add utils.calc_subsets
MODIFY
pangeo_forge_recipes/utils.py
pangeo_forge_recipes/utils.py
-import itertools
import logging
from contextlib import contextmanager
from typing import List, Sequence, Tuple
@@ -9,16 +8,6 @@ from dask.distributed import Lock, get_client
logger = logging.getLogger(__name__)
-# https://alexwlchan.net/2018/12/iterating-in-fixed-size-chunks/
-def chunked_iterable(iterable, size):
- it = iter(iterable)
- while True:
- chunk = tuple(itertools.islice(it, size))
- if not chunk:
- break
- yield chunk
-
-
# only needed because of
# https://github.com/pydata/xarray/issues/4631
def fix_scalar_attr_encoding(ds):
@@ -38,6 +27,16 @@ def fix_scalar_attr_encoding(ds):
return ds
+def calc_subsets(sequence_len: int, n_segments: int) -> List[int]:
+ """Given a sequence length, return a list that divides it into n_segments
+ (possibly uneven) integer segments."""
+ if n_segments > sequence_len:
+ raise ValueError(f"Can't split len {sequence_len} into {n_segments} segments")
+ step = sequence_len // n_segments
+ remainder = sequence_len % n_segments
+ return (n_segments - 1) * [step] + [step + remainder]
+
+
def chunk_bounds_and_conflicts(
chunks: Sequence[int], zchunks: int
) -> Tuple[List[int], List[Tuple[int, ...]]]:

MODIFY
tests/test_utils.py
tests/test_utils.py
-from pangeo_forge_recipes.utils import chunk_bounds_and_conflicts
+import pytest
+
+from pangeo_forge_recipes.utils import calc_subsets, chunk_bounds_and_conflicts
def test_chunk_conflicts():
@@ -10,3 +12,10 @@ def test_chunk_conflicts():
[(), (1,), (1,), ()],
)
assert chunk_bounds_and_conflicts([9, 12, 5], zchunks) == ([0, 9, 21, 26], [(0,), (0, 2), (2,)])
+
+
+def test_calc_subsets():
+ with pytest.raises(ValueError):
+ _ = calc_subsets(4, 5)
+ assert calc_subsets(5, 5) == [1, 1, 1, 1, 1]
+ assert calc_subsets(6, 5) == [1, 1, 1, 1, 2]

add utils.calc_subsets