add consolidate and subset methods
MODIFY
pangeo_forge_recipes/chunk_grid.py
pangeo_forge_recipes/chunk_grid.py
Abstract representation of ND chunked arrays
"""
+from itertools import chain, groupby
from typing import Dict, FrozenSet, Tuple
import numpy as np
+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.
@@ -110,6 +113,24 @@ class ChunkAxis:
def __len__(self):
return self._chunk_bounds[-1].item()
+ def subset(self, factor):
+ """Return a copy with chunks decimated by a subset factor."""
+
+ new_chunks = tuple(chain(*(calc_subsets(c, factor) for c in self.chunks)))
+ return self.__class__(new_chunks)
+
+ def consolidate(self, factor):
+ """Return a copy with chunks consolidated by a subset factor."""
+
+ new_chunks = []
+
+ def grouper(val):
+ return val[0] // factor
+
+ for _, gobj in groupby(enumerate(self.chunks), grouper):
+ new_chunks.append(sum(f[1] for f in gobj))
+ return self.__class__(tuple(new_chunks))
+
@property
def nchunks(self):
return len(self.chunks)

MODIFY
tests/test_chunk_grid.py
tests/test_chunk_grid.py
@@ -45,6 +45,20 @@ def test_chunk_axis():
_ = ca.chunk_index_to_array_slice(3)
+def test_chunk_axis_subsets():
+ ca = ChunkAxis(chunks=(2, 4, 3))
+ cas = ca.subset(2)
+ assert cas.chunks == (1, 1, 2, 2, 1, 2)
+
+
+def test_chunk_axis_consolidate():
+ ca = ChunkAxis(chunks=(2, 4, 3, 4, 2))
+ cac = ca.consolidate(2)
+ assert cac.chunks == (6, 7, 2)
+ cad = ca.consolidate(3)
+ assert cad.chunks == (9, 6)
+
+
def test_chunk_grid():
cg = ChunkGrid({"x": (2, 4, 3), "time": (7, 8)})
assert cg.dims == {"x", "time"}

add consolidate and subset methods