feat(containers): if chunks undefined, use default values

radiocosmology · May 18, 2021 · a754475 · a754475
1 parent e61d252
commit a754475
Showing 1 changed file with 27 additions and 9 deletions.
diff --git a/draco/core/containers.py b/draco/core/containers.py
@@ -67,6 +67,16 @@
     COMPRESSION = None
     COMPRESSION_OPTS = None
 
+chunk_sizes_default = {
+    "freq": 64,
+    "input": 128,
+    "pix": 128,
+    "pol": 2,
+    "ra": 128,
+    "stack": 256,
+    "time": 128,
+}
+
 
 class ContainerBase(memh5.BasicCont):
     """A base class for pipeline containers.
@@ -279,24 +289,32 @@ def chunkify(self):
         if not set them to dataset_spec values.
         """
         for dset in self.dataset_spec:
-            if "chunks" in self.dataset_spec[dset] and self[dset].chunks is None:
-                # ensure chunks aren't larger than dataset shape
+            if self[dset].chunks is None:
                 chunks = ()
-                for i, l in enumerate(self[dset].shape):
-                    chunks += (min(self.dataset_spec[dset]["chunks"][i], l),)
-                self._data._storage_root[dset].chunks = chunks
+                if "chunks" in self.dataset_spec[dset]:
+                    # ensure chunks aren't larger than dataset shape
+                    for i, l in enumerate(self[dset].shape):
+                        chunks += (min(self.dataset_spec[dset]["chunks"][i], l),)
+                    self._data._storage_root[dset].chunks = chunks
+                elif "axes" in self.dataset_spec[dset]:
+                    # Use default values based on axes names, or fall back to use axis shape
+                    for i, l in enumerate(self[dset].shape):
+                        axis_name = self.dataset_spec[dset]["axes"][i]
+                        chunks += (chunk_sizes_default.get(axis_name, l),)
             if (
                 "compression" in self.dataset_spec[dset]
                 and self[dset].compression is None
             ):
-                self._data._storage_root[dset].compression = self.dataset_spec[dset]["compression"]
+                self._data._storage_root[dset].compression = self.dataset_spec[dset][
+                    "compression"
+                ]
             if (
                 "compression_opts" in self.dataset_spec[dset]
                 and self[dset].compression_opts is None
             ):
-                self._data._storage_root[dset].compression_opts = self.dataset_spec[dset][
-                    "compression_opts"
-                ]
+                self._data._storage_root[dset].compression_opts = self.dataset_spec[
+                    dset
+                ]["compression_opts"]
 
     @property
     def datasets(self):