added ChunkGrid.from_uniform_grid
MODIFY
pangeo_forge_recipes/chunk_grid.py
pangeo_forge_recipes/chunk_grid.py
Abstract representation of ND chunked arrays
"""
-from dataclasses import dataclass
from typing import Dict, FrozenSet, Tuple
import numpy as np
@@ -16,11 +15,33 @@ ChunkIndex = Tuple[int, ...]
class ChunkGrid:
"""A ChunkGrid contains several named ChunkAxis.
The order of the axes does not matter.
+
+ :param chunks: Dictionary mapping dimension names to chunks in each dimension.
"""
- def __init__(self, chunks):
+ def __init__(self, chunks: Dict[str, Tuple[int, ...]]):
self._chunk_axes = {name: ChunkAxis(axis_chunks) for name, axis_chunks in chunks.items()}
+ def __eq__(self, other):
+ if self.dims != other.dims:
+ return False
+ for name in self._chunk_axes:
+ if self._chunk_axes[name] != other._chunk_axes[name]:
+ return False
+ return True
+
+ @classmethod
+ def from_uniform_grid(cls, chunksize_and_dimsize: Dict[str, Tuple[int, int]]):
+ all_chunks = {}
+ for name, (chunksize, dimsize) in chunksize_and_dimsize.items():
+ assert dimsize > 0
+ assert chunksize > 0 and chunksize <= dimsize
+ chunks = (dimsize // chunksize) * (chunksize,)
+ if dimsize % chunksize > 0:
+ chunks = chunks + (dimsize % chunksize,)
+ all_chunks[name] = chunks
+ return cls(all_chunks)
+
@property
def dims(self) -> FrozenSet[str]:
return frozenset(self._chunk_axes)
@@ -38,25 +59,31 @@ class ChunkGrid:
return len(self._chunk_axes)
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."""
+
return {
name: self._chunk_axes[name].chunk_index_to_array_slice(chunk_index)
for name, chunk_index in chunk_indexes.items()
}
def array_index_to_chunk_index(self, array_indexes: Dict[str, int]) -> Dict[str, int]:
+ """Figure out which chunk a single array-space index comes from
+ for each specified dimension."""
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]:
+ """Find all chunks that intersect with a given array-space slice
+ for each specified dimension."""
return {
name: self._chunk_axes[name].array_slice_to_chunk_slice(array_slice)
for name, array_slice in array_slices.items()
}
-@dataclass
class ChunkAxis:
"""A ChunkAxis has two index spaces.
@@ -64,13 +91,17 @@ class ChunkAxis:
Chunk index space describes chunk positions.
A ChunkAxis helps translate between these two spaces.
- """
- chunks: Tuple[int, ...]
+ :param chunks: The explicit size of each chunk
+ """
- def __post_init__(self):
+ def __init__(self, chunks: Tuple[int, ...]):
+ self.chunks = chunks
self._chunk_bounds = np.hstack([0, np.cumsum(self.chunks)])
+ def __eq__(self, other):
+ return self.chunks == other.chunks
+
def __len__(self):
return self._chunk_bounds[-1].item()
@@ -79,17 +110,22 @@ class ChunkAxis:
return len(self.chunks)
def chunk_index_to_array_slice(self, chunk_index: int) -> slice:
+ """Convert a single index from chunk space to a slice in array space."""
+
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:
+ """Figure out which chunk a single array-space index comes from."""
+
if array_index < 0 or array_index >= len(self):
raise IndexError("Index out of range")
return self._chunk_bounds.searchsorted(array_index, side="right") - 1
def array_slice_to_chunk_slice(self, sl: slice) -> slice:
- """Find all chunks that intersect with a given slice."""
+ """Find all chunks that intersect with a given array-space slice."""
+
if sl.step != 1 and sl.step is not None:
raise IndexError("Only works with step=1 or None")
if sl.start < 0:

MODIFY
tests/test_chunk_grid.py
tests/test_chunk_grid.py
@@ -69,3 +69,9 @@ def test_chunk_grid():
"x": slice(2, 6),
"time": slice(7, 15),
}
+
+
+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

added ChunkGrid.from_uniform_grid