fix sequence_len bug
MODIFY
pangeo_forge/patterns.py
pangeo_forge/patterns.py
@@ -3,7 +3,6 @@ Filename patterns.
"""
import itertools
-import string
from abc import ABC
from dataclasses import dataclass
from typing import Dict, Generator, Iterable, List, Tuple
@@ -20,11 +19,15 @@ class URLPattern(BasePattern):
keys: Dict[str, Iterable]
def __post_init__(self):
- fmt_string_keys = [
- t[1] for t in string.Formatter().parse(self.fmt_string) if t[1] is not None
- ]
- if not set(fmt_string_keys) == set(self.keys):
- raise KeyError("Specified keys don't match fmt_string")
+ pass
+ # TODO: make a better check
+ # this check does not work if you have to use the keys multiple times in the format string
+ # fmt_string_keys = [
+ # t[1] for t in string.Formatter().parse(self.fmt_string) if t[1] is not None
+ # ]
+ # e.g. '{variable}/{variable[0]}.nc'
+ # if not set(fmt_string_keys) == set(self.keys):
+ # raise KeyError("Specified keys don't match fmt_string")
def __iter__(self):
for keys in itertools.product(*self.keys.values()):

MODIFY
pangeo_forge/recipe.py
pangeo_forge/recipe.py
@@ -6,7 +6,7 @@ import logging
from abc import ABC, abstractmethod
from contextlib import contextmanager
from dataclasses import dataclass, field
-from typing import Callable, Hashable, Iterable, NoReturn, Optional
+from typing import Callable, Dict, Hashable, Iterable, NoReturn, Optional
import fsspec
import xarray as xr
@@ -110,6 +110,12 @@ class BaseRecipe(ABC):
pipelines.append(pipeline)
return pipelines
+ # https://stackoverflow.com/questions/59986413/achieving-multiple-inheritance-using-python-dataclasses
+ def __post_init__(self):
+ # just intercept the __post_init__ calls so they
+ # aren't relayed to `object`
+ pass
+
# Notes about dataclasses:
# - https://www.python.org/dev/peps/pep-0557/#inheritance
@@ -147,8 +153,10 @@ class NetCDFtoZarrRecipe(BaseRecipe):
`(ds: xr.Dataset, filename: str) -> ds: xr.Dataset`.
:param process_chunk: Function to call on each concatenated chunk, with signature
`(ds: xr.Dataset) -> ds: xr.Dataset`.
+ :param target_chunks: Desired chunk structure for the targret dataset.
"""
+ sequence_dim: str = ""
inputs_per_chunk: int = 1
nitems_per_input: int = 1
target: Optional[AbstractTarget] = field(default_factory=UninitializedTarget)
@@ -161,6 +169,23 @@ class NetCDFtoZarrRecipe(BaseRecipe):
fsspec_open_kwargs: dict = field(default_factory=dict)
process_input: Optional[Callable[[xr.Dataset, str], xr.Dataset]] = None
process_chunk: Optional[Callable[[xr.Dataset], xr.Dataset]] = None
+ target_chunks: Optional[Dict[str, int]] = None
+
+ def __post_init__(self):
+ if self.target_chunks:
+ sequence_dim_chunks = self.target_chunks.get(self.sequence_dim)
+ if sequence_dim_chunks:
+ if sequence_dim_chunks > self.nitems_per_input:
+ raise ValueError(
+ f"`sequence_dim_chunks` ({sequence_dim_chunks}) must be ",
+ f"<= `nitems_per_input` ({self.nitems_per_input})",
+ )
+ if self.nitems_per_input % sequence_dim_chunks > 0:
+ raise ValueError(
+ "`sequence_dim_chunks` ({sequence_dim_chunks}) must "
+ f"evenly divide `nitems_per_input` ({self.nitems_per_input})"
+ )
+ # TODO: more input validation
@property
def prepare_target(self) -> Callable:
@@ -170,6 +195,12 @@ class NetCDFtoZarrRecipe(BaseRecipe):
ds = self.open_target()
logger.info("Found an existing dataset in target")
logger.debug(f"{ds}")
+
+ if self.target_chunks:
+ # TODO: check that target_chunks id compatibile with the
+ # existing chunks
+ pass
+
except (FileNotFoundError, IOError, zarr.errors.GroupNotFoundError):
init_dsets = []
@@ -188,9 +219,15 @@ class NetCDFtoZarrRecipe(BaseRecipe):
# make sure the concat dim has a valid fill_value to avoid
# overruns when writing chunk
- ds[self.sequence_dim].encoding = {"_FillValue": -1}
+ # this is problematic because this value will always come out as NaN
+ # when we open up the target
+ # ds[self.sequence_dim].encoding = {"_FillValue": -1}
# actually not necessary if we use decode_times=False
+ if self.target_chunks:
+ # target_chunks has undergone some validation in __post_init__
+ ds = ds.chunk(self.target_chunks)
+
self.initialize_target(ds)
# Regardless of whether there is an existing dataset or we are creating a new one,
@@ -315,6 +352,11 @@ class NetCDFtoZarrRecipe(BaseRecipe):
shape[axis] = dimsize
arr.resize(shape)
+ # now explicity write the sequence coordinate to avoid missing data
+ # when reopening
+ if dim in zgroup:
+ zgroup[dim][:] = 0
+
def inputs_for_chunk(self, chunk_key):
return self._chunks_inputs[chunk_key]
@@ -334,10 +376,6 @@ class NetCDFtoZarrRecipe(BaseRecipe):
region_slice = slice(start, start + self.nitems_for_chunk(chunk_key))
return {self.sequence_dim: region_slice}
- def sequence_len(self):
- # tells the total size of dataset along the sequence dimension
- return sum([self.nitems_for_chunk(k) for k in self.iter_chunks()])
-
def sequence_chunks(self):
# chunking
return {self.sequence_dim: self.inputs_per_chunk * self.nitems_per_input}
@@ -350,9 +388,9 @@ class NetCDFtoZarrRecipe(BaseRecipe):
@dataclass
class NetCDFtoZarrSequentialRecipe(NetCDFtoZarrRecipe):
input_urls: Iterable[str] = field(repr=False, default_factory=list)
- sequence_dim: str = ""
def __post_init__(self):
+ super().__post_init__()
input_pattern = ExplicitURLSequence(self.input_urls)
self._inputs = {k: v for k, v in input_pattern}
self._chunks_inputs = {
@@ -364,13 +402,17 @@ class NetCDFtoZarrSequentialRecipe(NetCDFtoZarrRecipe):
def chunk_position(self, chunk_key):
return chunk_key
+ def sequence_len(self):
+ # tells the total size of dataset along the sequence dimension
+ return sum([self.nitems_for_chunk(k) for k in self.iter_chunks()])
+
@dataclass
class NetCDFtoZarrMultiVarSequentialRecipe(NetCDFtoZarrRecipe):
input_pattern: VariableSequencePattern = field(default_factory=VariableSequencePattern)
- sequence_dim: str = ""
def __post_init__(self):
+ super().__post_init__()
self._variables = self.input_pattern.keys["variable"]
self._inputs = {k: v for k, v in self.input_pattern}
# input keys are tuples like
@@ -398,3 +440,7 @@ class NetCDFtoZarrMultiVarSequentialRecipe(NetCDFtoZarrRecipe):
def chunk_position(self, chunk_key):
return chunk_key[1]
+
+ def sequence_len(self):
+ n_sequence = len(self.input_pattern.keys[self.input_pattern._sequence_key])
+ return self.nitems_per_input * n_sequence

MODIFY
tests/conftest.py
tests/conftest.py
@@ -27,6 +27,8 @@ def get_open_port():
def daily_xarray_dataset():
"""Return a synthetic random xarray dataset."""
np.random.seed(1)
+ # TODO: change nt to 11 in order to catch the edge case where
+ # items_per_input does not evenly divide the length of the sequence dimension
nt, ny, nx = 10, 18, 36
time = pd.date_range(start="2010-01-01", periods=nt, freq="D")
lon = (np.arange(nx) + 0.5) * 360 / nx

MODIFY
tests/test_recipe.py
tests/test_recipe.py
+from contextlib import nullcontext as does_not_raise
+
import aiohttp
import pytest
import xarray as xr
@@ -59,27 +61,62 @@ def test_NetCDFtoZarrSequentialRecipeHttpAuth(
"process_input, process_chunk",
[(None, None), (incr_date, None), (None, incr_date), (incr_date, incr_date)],
)
+@pytest.mark.parametrize("inputs_per_chunk", [1, 2])
+@pytest.mark.parametrize(
+ "target_chunks,chunk_expectation",
+ [
+ (None, does_not_raise()),
+ ({"lon": 12}, does_not_raise()),
+ ({"lon": 12, "time": 1}, does_not_raise()),
+ ({"lon": 12, "time": 3}, pytest.raises(ValueError)),
+ ],
+)
def test_NetCDFtoZarrSequentialRecipe(
- daily_xarray_dataset, netcdf_local_paths, tmp_target, tmp_cache, process_input, process_chunk
+ daily_xarray_dataset,
+ netcdf_local_paths,
+ tmp_target,
+ tmp_cache,
+ process_input,
+ process_chunk,
+ inputs_per_chunk,
+ target_chunks,
+ chunk_expectation,
):
# the same recipe is created as a fixture in conftest.py
# I left it here explicitly because it makes the test easier to read.
paths, items_per_file = netcdf_local_paths
+ with chunk_expectation as excinfo:
r = recipe.NetCDFtoZarrSequentialRecipe(
input_urls=paths,
sequence_dim="time",
- inputs_per_chunk=1,
+ inputs_per_chunk=inputs_per_chunk,
nitems_per_input=items_per_file,
target=tmp_target,
input_cache=tmp_cache,
process_input=process_input,
process_chunk=process_chunk,
+ target_chunks=target_chunks,
)
+ if excinfo:
+ # don't continue if we got an exception
+ return
_manually_execute_recipe(r)
- ds_target = xr.open_zarr(tmp_target.get_mapper(), consolidated=True).load()
+ ds_target = xr.open_zarr(tmp_target.get_mapper(), consolidated=True)
+
+ # chunk validation
+ sequence_chunks = ds_target.chunks["time"]
+ if target_chunks is None:
+ target_chunks = {}
+ seq_chunk_len = target_chunks.pop("time", None) or (items_per_file * inputs_per_chunk)
+ # we expect all chunks but the last to have the expected size
+ assert all([item == seq_chunk_len for item in sequence_chunks[:-1]])
+ for other_dim, chunk_len in target_chunks.items():
+ all([item == chunk_len for item in ds_target.chunks[other_dim][:-1]])
+
+ ds_target.load()
ds_expected = daily_xarray_dataset.compute()
if process_input is not None:
@@ -125,3 +162,8 @@ def test_NetCDFtoZarrMultiVarSequentialRecipe(
input_cache=tmp_cache,
)
_manually_execute_recipe(r)
+
+ ds_target = xr.open_zarr(tmp_target.get_mapper(), consolidated=True).compute()
+ print(ds_target)
+ print(daily_xarray_dataset)
+ assert ds_target.identical(daily_xarray_dataset)

fix sequence_len bug