added argument to allow chunks-per-shard while configuring the dataset

apps/hiopy/configure/configure.py

@@ -32,8 +32,8 @@ def _collect_groups(dataset):
 
 def add_height(dataset, name, n):
     for g in _collect_groups(dataset):
-        height = g.create_array(name, fill_value=None, dtype=np.int32, shape=np.arange(n).shape)
-        height.append(data=np.arange(n)) 
+        height = g.create_array(name, fill_value=None, dtype=np.int64, shape=np.arange(n).shape)
+        height[:] = np.arange(n)
         height.attrs["_ARRAY_DIMENSIONS"] = [name]
         height.attrs["axis"] = "Z"
         height.attrs["long_name"] = "generalized_height"
@@ -53,6 +53,7 @@ def add_variable(
     frac_mask=None,
     yac_name=None,
     attributes=None,
+    chunks_per_shard=None,
     **kwargs,
 ):
     for g in _collect_groups(dataset):
@@ -72,6 +73,7 @@ def add_variable(
             raise Exception("Unknown crs.")
 
         _attributes = attributes or {}
+
         if zaxis is None:
             shape = (*ntime, crs_len)
             _chunk_shape = np.minimum(chunk_shape, shape) if chunk_shape is not None else None
@@ -83,8 +85,18 @@ def add_variable(
             _attributes["_ARRAY_DIMENSIONS"] = (*taxis_tuple, zaxis, spatial_attr)
         _attributes["grid_mapping"] = "crs"
 
+        _shard_shape = None
+        if chunks_per_shard is not None:
+            _shard_shape = tuple(i * chunks_per_shard for i in _chunk_shape)
+
         v = g.create_array(
-            name, shape=shape, dtype=np.float32, fill_value=np.nan, chunks=_chunk_shape, **kwargs
+            name,
+            shape=shape,
+            dtype=np.float32,
+            fill_value=np.nan,
+            chunks=_chunk_shape,
+            shards=_shard_shape,
+            **kwargs,
         )
 
         # TODO: Use a generic name instead of hiopy such that it represents arbitrary grid too

apps/hiopy/configure/create_dataset.py

@@ -38,4 +38,3 @@ def add_healpix_hierarchy(dataset, order, prefix="healpix_"):
         add_healpix_grid(zg, o)
         if o < order:
             zg.attrs["hiopy::parent"] = f"{prefix}{o+1}"
-

apps/hiopy/worker.py

@@ -67,7 +67,7 @@ def main():
         assert len(args.datasets) == 1, "Loco only supports reading from 1 dataset"
         loco_store = zarr.MemoryStore()
         zarr.copy_store(args.datasets[0].store, loco_store)
-        zarr.convenience.consolidate_metadata(loco_store)
+        zarr.consolidate_metadata(loco_store)
         loco_server = LocoServer(loco_store, args.loco_host, args.loco_port)
         args.datasets = [zarr.open(store=loco_store)]