add ChunkGrid conflict test
MODIFY
pangeo_forge_recipes/chunk_grid.py
pangeo_forge_recipes/chunk_grid.py
@@ -114,6 +114,20 @@ class ChunkGrid:
for name, array_slice in array_slices.items()
}
+ def chunk_conflicts(self, chunk_index: Dict[str, int], other: ChunkGrid) -> Dict[str, Set[int]]:
+ """Figure out which _other_ chunks from this ChunkGrid might potentially
+ be in conflict with the specificied chunk index when writing to a
+ different ChunkGrid.
+
+ :param chunk_index: The index of the chunk we want to write
+ :param other: The other ChunkAxis
+ """
+
+ return {
+ name: self._chunk_axes[name].chunk_conflicts(idx, other._chunk_axes[name])
+ for name, idx in chunk_index.items()
+ }
+
class ChunkAxis:
"""A ChunkAxis has two index spaces.
@@ -191,6 +205,9 @@ class ChunkAxis:
"""Figure out which _other_ chunks from this ChunkAxis might potentially
be in conflict with the specificied chunk index when writing to a
different ChunkAxis.
+
+ :param chunk_index: The index of the chunk we want to write
+ :param other: The other ChunkAxis
"""
if len(other) != len(self):

MODIFY
tests/test_chunk_grid.py
tests/test_chunk_grid.py
@@ -124,3 +124,12 @@ def test_chunk_axis_conflicts():
with pytest.raises(ValueError):
_ = ca1.chunk_conflicts(0, ChunkAxis((14,)))
+
+
+def test_chunk_grid_conflicts():
+ cg1 = ChunkGrid({"x": (2, 4, 3, 4, 2), "y": (10, 10, 10)})
+ cg2 = ChunkGrid({"x": (5, 4, 6), "y": (11, 9, 10)})
+
+ assert cg1.chunk_conflicts({"x": 0}, cg2) == {"x": {1}}
+ assert cg1.chunk_conflicts({"x": 0, "y": 0}, cg2) == {"x": {1}, "y": {1}}
+ assert cg1.chunk_conflicts({"y": 2}, cg2) == {"y": set()}

add ChunkGrid conflict test