Skip to content
Snippets Groups Projects
Commit 3d0a938b authored by Siddhant Tibrewal's avatar Siddhant Tibrewal
Browse files

added some documentation to the create-dataset helper functions

parent 65f1b8fb
No related branches found
No related tags found
1 merge request!43Draft: Resolve "Migrate to Zarr3 and add example for a dataset using sharding"
#!/usr/bin/env python3
import numpy as np
import zarr
def add_coordinates(
dataset: zarr.Dataset,
coordinates: list[tuple[float, float]],
coord_names: tuple[str, str] = ("lon", "lat"),
) -> None:
"""
Add longitude and latitude coordinates to the specified Zarr dataset.
Parameters
----------
dataset : zarr.Dataset
The Zarr dataset where the coordinates will be added.
coordinates : list[tuple[float, float]]
A list of tuples containing the (longitude, latitude) values for each point.
coord_names : tuple[str, str], optional
The names to use for the longitude and latitude arrays. Defaults to ("lon", "lat").
Notes
-----
This function creates two new arrays in the dataset: `coord_names[0]` for longitude and `coord_names[1]` for latitude.
The `crs` array is also created, with its attributes set to indicate that it's a "point_cloud" coordinate reference system.
Example: add_coordinates(dataset, [(10.2, 45.3), (20.4, 50.5)])
"""
def add_coordinates(dataset, coordinates, coord_names=("lon", "lat")):
crs = dataset.create_array(name="crs", dtype=np.float32, shape=(1,))
crs.attrs["_ARRAY_DIMENSIONS"] = ("crs",)
crs.attrs["grid_mapping_name"] = "point_cloud"
......@@ -28,15 +53,59 @@ def add_coordinates(dataset, coordinates, coord_names=("lon", "lat")):
lat.attrs["units"] = "degree"
lat.attrs["standard_name"] = "grid_latitude"
def add_healpix_grid(dataset, order):
def add_healpix_grid(dataset: zarr.Dataset, order: int):
"""
Add a HealPix grid to the specified Zarr dataset.
Parameters
----------
dataset : zarr.Dataset
The Zarr dataset where the HealPix grid will be added to the crs.
order : int
The order of the HealPix grid. This corresponds to 2^order for the NSIDE.
Notes
-----
The HealPix grid is stored as a single array named "crs" in the dataset, with the healpix_nside and healpix_order attributes set
accordingly. No values are added to it
"""
crs = dataset.create_array(name="crs", dtype=np.float32, shape=(1,))
crs.attrs["_ARRAY_DIMENSIONS"] = ("crs",)
crs.attrs["grid_mapping_name"] = "healpix"
crs.attrs["healpix_nside"] = 2**order
crs.attrs["healpix_order"] = "nest"
def add_healpix_hierarchy(dataset, order, prefix="healpix_"):
for o in range(order + 1):
def add_healpix_hierarchy(
dataset: zarr.Dataset,
order: int,
nr_of_coarsenings: int = 4,
prefix: str = "healpix_",
) -> None:
"""
Add a hierarchical structure to the specified Zarr dataset for a given Healpix order.
This function creates a group hierarchy with each level representing a specific resolution of the Healpix grid.
The `add_healpix_grid` function is used to create the actual grid arrays within each group.
Parameters
----------
dataset : zarr.Dataset
The Zarr dataset where the hierarchy will be added.
order : int
The maximum level in the hierarchy.
nr_of_coarsenings : int
Number of coarsening aggregation levels needed
prefix : str, optional
The prefix to use for naming each group. Defaults to "healpix_".
Notes
-----
This function sets up a hierarchical structure with each level representing a specific resolution of the Healpix grid.
The `hiopy::parent` attribute is used to link each group to its parent in the hierarchy, allowing for efficient navigation.
"""
for o in range(order, order - nr_of_coarsenings, -1):
zg = dataset.create_group(name=f"{prefix}{o}")
add_healpix_grid(zg, o)
if o < order:
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment