From eb114b01946a4f7c6a805ef9702857ac0ec09985 Mon Sep 17 00:00:00 2001
From: Liam Gray <lgray@phas.ubc.ca>
Date: Mon, 7 Oct 2024 11:22:52 -0700
Subject: [PATCH] feat(containers): increase default chunk size for a few
 containers

---
 draco/core/containers.py | 45 +++++++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/draco/core/containers.py b/draco/core/containers.py
index 537b04fd..653677d0 100644
--- a/draco/core/containers.py
+++ b/draco/core/containers.py
@@ -1134,7 +1134,7 @@ class SiderealStream(
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": {
                 "weight_dataset": "vis_weight",
             },
@@ -1147,7 +1147,7 @@ class SiderealStream(
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": True,
         },
         "input_flags": {
@@ -1171,7 +1171,7 @@ class SiderealStream(
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (3, 64, 128, 128),
+            "chunks": (3, 32, 512, 2048),
             "truncate": True,
         },
         "nsample": {
@@ -1182,7 +1182,7 @@ class SiderealStream(
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
         },
         "effective_ra": {
             "axes": ["freq", "stack", "ra"],
@@ -1192,7 +1192,7 @@ class SiderealStream(
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": True,
         },
     }
@@ -1239,7 +1239,7 @@ class SiderealDirtyStream(FreqContainer, VisContainer, SiderealContainer):
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": True,
         },
         "mask": {
@@ -1255,6 +1255,9 @@ class SiderealDirtyStream(FreqContainer, VisContainer, SiderealContainer):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (32, 1, 2048),
             "truncate": False,
         },
         "modes": {
@@ -1263,6 +1266,8 @@ class SiderealDirtyStream(FreqContainer, VisContainer, SiderealContainer):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
             "truncate": False,
         },
     }
@@ -1316,7 +1321,7 @@ class SiderealBandCovariance(FreqContainer, VisContainer, SiderealContainer):
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 1, 128),
+            "chunks": (32, 512, 1, 2048),
             "truncate": True,
         },
     }
@@ -1564,7 +1569,7 @@ class TimeStream(FreqContainer, VisContainer, TODContainer):
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": {
                 "weight_dataset": "vis_weight",
             },
@@ -1577,7 +1582,7 @@ class TimeStream(FreqContainer, VisContainer, TODContainer):
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": True,
         },
         "input_flags": {
@@ -1621,7 +1626,7 @@ class FactorizedTimeStream(TimeStream):
             "distributed_axis": "freq",
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
-            "chunks": (64, 128, 128),
+            "chunks": (32, 512, 2048),
             "truncate": True,
         },
         "vis_weight": {
@@ -1630,6 +1635,9 @@ class FactorizedTimeStream(TimeStream):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (512, 2048),
             "truncate": False,
         },
         "modes": {
@@ -1638,6 +1646,9 @@ class FactorizedTimeStream(TimeStream):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
+            "compression": COMPRESSION,
+            "compression_opts": COMPRESSION_OPTS,
+            "chunks": (512, 512),
             "truncate": False,
         },
     }
@@ -2046,7 +2057,7 @@ class VisGridStream(FreqContainer, SiderealContainer, VisBase):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
-            "chunks": (1, 64, 1, 64, 128),
+            "chunks": (1, 32, 1, 256, 2048),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
             "truncate": {
@@ -2059,7 +2070,7 @@ class VisGridStream(FreqContainer, SiderealContainer, VisBase):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
-            "chunks": (1, 64, 1, 64, 128),
+            "chunks": (1, 32, 1, 256, 2048),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
             "truncate": True,
@@ -2069,7 +2080,7 @@ class VisGridStream(FreqContainer, SiderealContainer, VisBase):
             "dtype": np.int32,
             "initialise": True,
             "distributed": False,
-            "chunks": (1, 64, 1, 64, 128),
+            "chunks": (1, 1, 256, 2048),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
         },
@@ -2284,7 +2295,7 @@ class RingMap(FreqContainer, SiderealContainer, DataWeightContainer):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
-            "chunks": (1, 1, 64, 128, 128),
+            "chunks": (1, 1, 32, 2048, 512),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
             "truncate": {
@@ -2297,7 +2308,7 @@ class RingMap(FreqContainer, SiderealContainer, DataWeightContainer):
             "initialise": True,
             "distributed": True,
             "distributed_axis": "freq",
-            "chunks": (1, 64, 128, 128),
+            "chunks": (1, 32, 2048, 512),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
             "truncate": True,
@@ -2308,7 +2319,7 @@ class RingMap(FreqContainer, SiderealContainer, DataWeightContainer):
             "initialise": False,
             "distributed": True,
             "distributed_axis": "freq",
-            "chunks": (1, 1, 64, 128, 128),
+            "chunks": (1, 1, 32, 2048, 512),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
             "truncate": True,
@@ -2319,7 +2330,7 @@ class RingMap(FreqContainer, SiderealContainer, DataWeightContainer):
             "initialise": False,
             "distributed": True,
             "distributed_axis": "freq",
-            "chunks": (4, 512, 512),
+            "chunks": (4, 512, 2048),
             "compression": COMPRESSION,
             "compression_opts": COMPRESSION_OPTS,
             "truncate": True,