add prune function to patterns
MODIFY
pangeo_forge_recipes/patterns.py
pangeo_forge_recipes/patterns.py
@@ -47,6 +47,7 @@ class MergeDim:
Index = Tuple[int, ...]
+CombineDim = Union[MergeDim, ConcatDim]
class FilePattern:
@@ -77,7 +78,7 @@ class FilePattern:
coords = {cdim.name: (cdim.name, cdim.keys) for cdim in combine_dims}
return xr.DataArray(fnames_np, dims=list(coords), coords=coords) # type: ignore
- def __init__(self, format_function: Callable, *combine_dims: Union[MergeDim, ConcatDim]):
+ def __init__(self, format_function: Callable, *combine_dims: CombineDim):
self.__setstate__((format_function, combine_dims))
def __getstate__(self):
@@ -157,3 +158,29 @@ def pattern_from_file_sequence(file_list, concat_dim, nitems_per_file=None):
return file_list[kwargs[concat_dim]]
return FilePattern(format_function, concat)
+
+
+def prune_pattern(fp: FilePattern, nkeep: int = 2) -> FilePattern:
+ """
+ Create a smaller pattern from a full pattern.
+ Keeps all MergeDims but only the first `nkeep` items from each ConcatDim
+
+ :param fp: The original pattern.
+ :param nkeep: The number of items to keep from each ConcatDim sequence.
+ """
+
+ new_combine_dims = [] # type: List[CombineDim]
+ for cdim in fp.combine_dims:
+ if isinstance(cdim, MergeDim):
+ new_combine_dims.append(cdim)
+ elif isinstance(cdim, ConcatDim):
+ new_keys = cdim.keys[:nkeep]
+ # this feels like a fragile way to copy a DataClass but it works for now
+ new_cdim = ConcatDim(
+ name=cdim.name, keys=new_keys, nitems_per_file=cdim.nitems_per_file
+ )
+ new_combine_dims.append(new_cdim)
+ else: # pragma: no cover
+ assert "Should never happen"
+
+ return FilePattern(fp.format_function, *new_combine_dims)

MODIFY
tests/test_patterns.py
tests/test_patterns.py
import pytest
-from pangeo_forge_recipes.patterns import ConcatDim, FilePattern, MergeDim, pattern_from_file_sequence
+from pangeo_forge_recipes.patterns import (
+ ConcatDim,
+ FilePattern,
+ MergeDim,
+ pattern_from_file_sequence,
+ prune_pattern,
+)
def test_file_pattern_concat():
@@ -68,3 +74,17 @@ def test_file_pattern_concat_merge(pickle):
assert fp[key] == fname
fnames.append(fname)
assert list(fp.items()) == list(zip(expected_keys, fnames))
+
+
+@pytest.mark.parametrize("nkeep", [1, 2])
+def test_prune(nkeep):
+ concat = ConcatDim(name="time", keys=list(range(3)))
+ merge = MergeDim(name="variable", keys=["foo", "bar"])
+
+ def format_function(time, variable):
+ return f"T_{time}_V_{variable}"
+
+ fp = FilePattern(format_function, merge, concat)
+ fp_pruned = prune_pattern(fp, nkeep=nkeep)
+ assert fp_pruned.dims == {"variable": 2, "time": nkeep}
+ assert len(list(fp_pruned.items())) == 2 * nkeep

add prune function to patterns