add subset and consolidate to ChunkGrid
MODIFY
pangeo_forge_recipes/chunk_grid.py
pangeo_forge_recipes/chunk_grid.py
@@ -12,8 +12,6 @@ from .utils import calc_subsets
# Most of this is probably already in Dask and Zarr!
# However, it's useful to write up our own little model that does just what we need.
-ChunkIndex = Tuple[int, ...]
-
class ChunkGrid:
"""A ChunkGrid contains several named ChunkAxis.
@@ -66,6 +64,30 @@ class ChunkGrid:
def ndim(self):
return len(self._chunk_axes)
+ def consolidate(self, factors: Dict[str, int]):
+ """Return a new ChunkGrid with chunks consolidated by a given factor
+ along specifed dimensions."""
+
+ # doesn't seem like the kosher way to do this but /shrug
+ new = self.__class__({})
+ new._chunk_axes = {
+ name: ca.consolidate(factors[name]) if name in factors else ca
+ for name, ca in self._chunk_axes.items()
+ }
+ return new
+
+ def subset(self, factors: Dict[str, int]):
+ """Return a new ChunkGrid with chunks decimated by a given subset factor
+ along specifed dimensions."""
+
+ # doesn't seem like the kosher way to do this but /shrug
+ new = self.__class__({})
+ new._chunk_axes = {
+ name: ca.subset(factors[name]) if name in factors else ca
+ for name, ca in self._chunk_axes.items()
+ }
+ return new
+
def chunk_index_to_array_slice(self, chunk_indexes: Dict[str, int]) -> Dict[str, slice]:
"""Convert a single index from chunk space to a slice in array space
for each specified dimension."""
@@ -104,7 +126,7 @@ class ChunkAxis:
"""
def __init__(self, chunks: Tuple[int, ...]):
- self.chunks = chunks
+ self.chunks = tuple(chunks) # want this immutable
self._chunk_bounds = np.hstack([0, np.cumsum(self.chunks)])
def __eq__(self, other):

MODIFY
tests/test_chunk_grid.py
tests/test_chunk_grid.py
@@ -45,7 +45,7 @@ def test_chunk_axis():
_ = ca.chunk_index_to_array_slice(3)
-def test_chunk_axis_subsets():
+def test_chunk_axis_subset():
ca = ChunkAxis(chunks=(2, 4, 3))
cas = ca.subset(2)
assert cas.chunks == (1, 1, 2, 2, 1, 2)
@@ -89,3 +89,23 @@ def test_chunk_grid_from_uniform_grid():
cg1 = ChunkGrid({"x": (2, 2), "y": (3, 3, 3, 1)})
cg2 = ChunkGrid.from_uniform_grid({"x": (2, 4), "y": (3, 10)})
assert cg1 == cg2
+
+
+def test_chunk_grid_consolidate_subset():
+ cg = ChunkGrid({"x": (2, 4, 3), "time": (7, 8)})
+
+ assert cg.consolidate({}) == cg
+ cgc1 = cg.consolidate({"x": 2})
+ assert cg.shape == cgc1.shape
+ assert cgc1.nchunks == {"x": 2, "time": 2}
+ cgc2 = cg.consolidate({"x": 2, "time": 2})
+ assert cg.shape == cgc2.shape
+ assert cgc2.nchunks == {"x": 2, "time": 1}
+
+ assert cg.subset({}) == cg
+ cgs1 = cg.subset({"x": 2})
+ assert cg.shape == cgs1.shape
+ assert cgs1.nchunks == {"x": 6, "time": 2}
+ cgs2 = cg.subset({"x": 2, "time": 2})
+ assert cg.shape == cgs2.shape
+ assert cgs2.nchunks == {"x": 6, "time": 4}

add subset and consolidate to ChunkGrid