complete ChunkGrid
MODIFY
pangeo_forge_recipes/chunk_grid.py
pangeo_forge_recipes/chunk_grid.py
@@ -14,6 +14,10 @@ ChunkIndex = Tuple[int, ...]
class ChunkGrid:
+ """A ChunkGrid contains several named ChunkAxis.
+ The order of the axes does not matter.
+ """
+
def __init__(self, chunks):
self._chunk_axes = {name: ChunkAxis(axis_chunks) for name, axis_chunks in chunks.items()}
@@ -33,13 +37,23 @@ class ChunkGrid:
def ndim(self):
return len(self._chunk_axes)
- def chunk_index_to_slice(self, chunk_index: ChunkIndex) -> Tuple[slice, ...]:
- return tuple(
- [axis.chunk_index_to_slice(idx) for axis, idx in zip(self._chunk_axes, chunk_index)]
- )
+ def chunk_index_to_array_slice(self, chunk_indexes: Dict[str, int]) -> Dict[str, slice]:
+ return {
+ name: self._chunk_axes[name].chunk_index_to_array_slice(chunk_index)
+ for name, chunk_index in chunk_indexes.items()
+ }
- def slice_to_chunk_indexes(self, slices: Tuple[slice]) -> ChunkIndex:
- return tuple([axis.slice_to_chunk_index(sl) for axis, sl in zip(self._chunk_axes, slices)])
+ def array_index_to_chunk_index(self, array_indexes: Dict[str, int]) -> Dict[str, int]:
+ return {
+ name: self._chunk_axes[name].array_index_to_chunk_index(array_index)
+ for name, array_index in array_indexes.items()
+ }
+
+ def array_slice_to_chunk_slice(self, array_slices: Dict[str, slice]) -> Dict[str, slice]:
+ return {
+ name: self._chunk_axes[name].array_slice_to_chunk_slice(array_slice)
+ for name, array_slice in array_slices.items()
+ }
@dataclass
@@ -50,7 +64,6 @@ class ChunkAxis:
Chunk index space describes chunk positions.
A ChunkAxis helps translate between these two spaces.
-
"""
chunks: Tuple[int, ...]
@@ -66,6 +79,8 @@ class ChunkAxis:
return len(self.chunks)
def chunk_index_to_array_slice(self, chunk_index: int) -> slice:
+ if chunk_index < 0 or chunk_index >= self.nchunks:
+ raise IndexError("chunk_index out of range")
return slice(self._chunk_bounds[chunk_index], self._chunk_bounds[chunk_index + 1])
def array_index_to_chunk_index(self, array_index: int) -> int:

MODIFY
tests/test_chunk_grid.py
tests/test_chunk_grid.py
@@ -36,6 +36,14 @@ def test_chunk_axis():
assert ca.array_slice_to_chunk_slice(slice(2, 5)) == slice(1, 2)
assert ca.array_slice_to_chunk_slice(slice(6, 7)) == slice(2, 3)
+ with pytest.raises(IndexError):
+ _ = ca.chunk_index_to_array_slice(-1)
+ assert ca.chunk_index_to_array_slice(0) == slice(0, 2)
+ assert ca.chunk_index_to_array_slice(1) == slice(2, 6)
+ assert ca.chunk_index_to_array_slice(2) == slice(6, 9)
+ with pytest.raises(IndexError):
+ _ = ca.chunk_index_to_array_slice(3)
+
def test_chunk_grid():
cg = ChunkGrid({"x": (2, 4, 3), "time": (7, 8)})
@@ -43,3 +51,21 @@ def test_chunk_grid():
assert cg.shape == {"x": 9, "time": 15}
assert cg.nchunks == {"x": 3, "time": 2}
assert cg.ndim == 2
+
+ assert cg.array_index_to_chunk_index({"x": 2}) == {"x": 1}
+ assert cg.array_index_to_chunk_index({"time": 10}) == {"time": 1}
+ assert cg.array_index_to_chunk_index({"x": 7, "time": 10}) == {"x": 2, "time": 1}
+
+ assert cg.array_slice_to_chunk_slice({"x": slice(0, 9)}) == {"x": slice(0, 3)}
+ assert cg.array_slice_to_chunk_slice({"time": slice(0, 15)}) == {"time": slice(0, 2)}
+ assert cg.array_slice_to_chunk_slice({"x": slice(0, 9), "time": slice(0, 15)}) == {
+ "x": slice(0, 3),
+ "time": slice(0, 2),
+ }
+
+ assert cg.chunk_index_to_array_slice({"x": 1}) == {"x": slice(2, 6)}
+ assert cg.chunk_index_to_array_slice({"time": 1}) == {"time": slice(7, 15)}
+ assert cg.chunk_index_to_array_slice({"x": 1, "time": 1}) == {
+ "x": slice(2, 6),
+ "time": slice(7, 15),
+ }

complete ChunkGrid