add chunk conflict finder
MODIFY
pangeo_forge/utils.py
pangeo_forge/utils.py
import itertools
+from typing import Iterable, List, Tuple
import numpy as np
@@ -30,3 +31,29 @@ def fix_scalar_attr_encoding(ds):
ds[v].attrs.update(_fixed_attrs(ds[v].attrs))
ds[v].encoding.update(_fixed_attrs(ds[v].encoding))
return ds
+
+
+def calc_chunk_conflicts(chunks: Iterable[int], zchunks: int) -> List[Tuple[int]]:
+ n_chunks = len(chunks)
+
+ chunk_bounds = np.hstack([0, np.cumsum(chunks)])
+ chunk_overlap = []
+ for start, stop in zip(chunk_bounds[:-1], chunk_bounds[1:]):
+ chunk_start = start // zchunks
+ chunk_stop = (stop - 1) // zchunks
+ chunk_overlap.append((chunk_start, chunk_stop))
+
+ chunk_conflicts = []
+ for n, chunk_pair in enumerate(chunk_overlap):
+ conflicts = set()
+ if n > 0:
+ prev_pair = chunk_overlap[n - 1]
+ if prev_pair[1] == chunk_pair[0]:
+ conflicts.add(chunk_pair[0])
+ if n < (n_chunks - 1):
+ next_pair = chunk_overlap[n + 1]
+ if next_pair[0] == chunk_pair[1]:
+ conflicts.add(chunk_pair[1])
+ chunk_conflicts.append(tuple(conflicts))
+
+ return chunk_conflicts

MODIFY
tests/test_utils.py
tests/test_utils.py
import pytest
-import pangeo_forge.utils
+from pangeo_forge import utils
@pytest.mark.parametrize(
@@ -13,5 +13,13 @@ import pangeo_forge.utils
],
)
def test_chunked_iterable(iterable, size, expected):
- actual = list(pangeo_forge.utils.chunked_iterable(iterable, size))
+ actual = list(utils.chunked_iterable(iterable, size))
assert actual == expected
+
+
+def test_chunk_conflicts():
+ zchunks = 10
+ assert utils.calc_chunk_conflicts([10, 10], zchunks) == [(), ()]
+ assert utils.calc_chunk_conflicts([9, 10], zchunks) == [(0,), (0,)]
+ assert utils.calc_chunk_conflicts([10, 9, 11, 10], zchunks) == [(), (1,), (1,), ()]
+ assert utils.calc_chunk_conflicts([9, 12, 5], zchunks) == [(0,), (0, 2), (2,)]

add chunk conflict finder