wip; doesnt work
MODIFY
pangeo_forge/recipe.py
pangeo_forge/recipe.py
@@ -13,7 +13,7 @@ import xarray as xr
import zarr
from rechunker.types import MultiStagePipeline, ParallelPipelines, Stage
-from .patterns import ExplicitURLSequence
+from .patterns import ExplicitURLSequence, VariableSequencePattern
from .storage import AbstractTarget, UninitializedTarget
from .utils import chunked_iterable, fix_scalar_attr_encoding
@@ -117,7 +117,7 @@ class BaseRecipe(ABC):
@dataclass
-class NetCDFtoZarrSequentialRecipe(BaseRecipe):
+class NetCDFtoZarrRecipe(BaseRecipe):
"""This class represents a dataset composed of many individual NetCDF files.
The files are arraged in a sequence along a single dimension, called the
`sequence_dim`. Each file may contain multiple variables.
@@ -149,8 +149,6 @@ class NetCDFtoZarrSequentialRecipe(BaseRecipe):
`(ds: xr.Dataset) -> ds: xr.Dataset`.
"""
- input_urls: Iterable[str] = field(repr=False, default_factory=list)
- sequence_dim: str = ""
inputs_per_chunk: int = 1
nitems_per_input: int = 1
target: Optional[AbstractTarget] = field(default_factory=UninitializedTarget)
@@ -164,13 +162,6 @@ class NetCDFtoZarrSequentialRecipe(BaseRecipe):
process_input: Optional[Callable[[xr.Dataset, str], xr.Dataset]] = None
process_chunk: Optional[Callable[[xr.Dataset], xr.Dataset]] = None
- def __post_init__(self):
- input_pattern = ExplicitURLSequence(self.input_urls)
- self._inputs = {k: v for k, v in input_pattern}
- self._chunks_inputs = {
- k: v for k, v in enumerate(chunked_iterable(self._inputs, self.inputs_per_chunk))
- }
-
@property
def prepare_target(self) -> Callable:
def _prepare_target():
@@ -180,17 +171,25 @@ class NetCDFtoZarrSequentialRecipe(BaseRecipe):
logger.info("Found an existing dataset in target")
logger.debug(f"{ds}")
except (FileNotFoundError, IOError, zarr.errors.GroupNotFoundError):
- first_chunk_key = next(self.iter_chunks())
- # for input_url in self.inputs_for_chunk(first_chunk_key):
- # self.cache_input(input_url)
- ds = self.open_chunk(first_chunk_key).chunk()
+
+ init_dsets = []
+ for chunk_key in self._init_chunks:
+ chunk_ds = self.open_chunk(
+ chunk_key
+ ).chunk() # make sure data are not in memory
+ init_dsets.append(chunk_ds)
+ # TODO: create csutomizable option for this step
+ ds = xr.merge(init_dsets, compat="identical", join="exact")
# make sure the concat dim has a valid fill_value to avoid
# overruns when writing chunk
ds[self.sequence_dim].encoding = {"_FillValue": -1}
# actually not necessary if we use decode_times=False
+
self.initialize_target(ds)
+ # Regardless of whether there is an existing dataset or we are creating a new one,
+ # we need to expand the sequence_dim to hold the entire expected size of the data
self.expand_target_dim(self.sequence_dim, self.sequence_len())
return _prepare_target
@@ -341,3 +340,48 @@ class NetCDFtoZarrSequentialRecipe(BaseRecipe):
def iter_chunks(self):
for k in self._chunks_inputs:
yield k
+
+
+@dataclass
+class NetCDFtoZarrSequentialRecipe(NetCDFtoZarrRecipe):
+ input_urls: Iterable[str] = field(repr=False, default_factory=list)
+ sequence_dim: str = ""
+
+ def __post_init__(self):
+ input_pattern = ExplicitURLSequence(self.input_urls)
+ self._inputs = {k: v for k, v in input_pattern}
+ self._chunks_inputs = {
+ k: v for k, v in enumerate(chunked_iterable(self._inputs, self.inputs_per_chunk))
+ }
+ # just the first chunk is needed to initialize the recipe
+ self._init_chunks = [next(iter(self._chunks_inputs))]
+
+
+@dataclass
+class test_NetCDFtoZarrMultiVarSequentialRecipe(NetCDFtoZarrRecipe):
+ input_pattern: VariableSequencePattern
+ sequence_dim: str = ""
+
+ def __post_init__(self):
+ self._variables = self.input_pattern.keys["variable"]
+ self._inputs = {k: v for k, v in self.input_pattern}
+ # input keys are tuples like
+ # ("temp", 0)
+ # ("temp", 1)
+ # ...
+ # ("salt", n)
+ # chunk_keys will be similarly named, but the index will have a different meaning
+
+ chunks_inputs = {}
+ # generate this iteratively
+ for vname in self._variables:
+ # a hacky way to filter
+ sequence_keys = [k[1] for k in self._inputs if k[0] == vname]
+ for chunk_key, chunk_inputs in enumerate(
+ chunked_iterable(sequence_keys, self.inputs_per_chunk)
+ ):
+ chunks_inputs[(vname, chunk_key)] = (vname, chunk_inputs)
+ self._chunks_inputs = chunks_inputs
+
+ # should be the first chunk from each variable
+ self._init_chunks = [chunk_key for chunk_key in self._chunks_inpust if chunk_key[1] == 0]

MODIFY
setup.cfg
setup.cfg
@@ -3,7 +3,7 @@ max-line-length = 100
[isort]
known_first_party=pangeo_forge
-known_third_party=aiohttp,fsspec,numpy,pandas,pkg_resources,prefect,pytest,rechunker,setuptools,sphinx_book_theme,xarray,zarr
+known_third_party=aiohttp,fsspec,numpy,pandas,pange_forge,pkg_resources,prefect,pytest,rechunker,setuptools,sphinx_book_theme,xarray,zarr
multi_line_output=3
include_trailing_comma=True
force_grid_wrap=0

MODIFY
tests/conftest.py
tests/conftest.py
@@ -10,6 +10,7 @@ import pytest
import xarray as xr
from pangeo_forge import recipe
+from pangeo_forge.patterns import VariableSequencePattern
from pangeo_forge.storage import CacheFSSpecTarget, FSSpecTarget, UninitializedTarget
@@ -60,12 +61,14 @@ def _split_up_files_by_day(ds, day_param):
def _split_up_files_by_variable_and_day(ds, day_param):
all_dsets = []
all_fnames = []
+ fnames_by_variable = {}
for varname in ds.data_vars:
var_dsets, fnames = _split_up_files_by_day(ds[[varname]], day_param)
fnames = [f"{varname}_{fname}" for fname in fnames]
all_dsets += var_dsets
all_fnames += fnames
- return all_dsets, all_fnames
+ fnames_by_variable[varname] = fnames
+ return all_dsets, all_fnames, fnames_by_variable
@pytest.fixture(scope="session", params=["D", "2D"])
@@ -85,12 +88,15 @@ def netcdf_local_paths(daily_xarray_dataset, tmpdir_factory, request):
def netcdf_local_paths_by_variable(daily_xarray_dataset, tmpdir_factory, request):
"""Return a list of paths pointing to netcdf files."""
tmp_path = tmpdir_factory.mktemp("netcdf_data")
- datasets, fnames = _split_up_files_by_variable_and_day(
+ datasets, fnames, fnames_by_variable = _split_up_files_by_variable_and_day(
daily_xarray_dataset.copy(), request.param
)
full_paths = [tmp_path.join(fname) for fname in fnames]
xr.save_mfdataset(datasets, [str(path) for path in full_paths])
- return full_paths
+ items_per_file = {"D": 1, "2D": 2}[request.param]
+ path_format = str(tmp_path) + "/{variable}_{n:03d}.nc"
+ print(path_format)
+ return full_paths, items_per_file, fnames_by_variable, path_format
# TODO: refactor to allow netcdf_local_paths_by_variable to be passed without
@@ -164,3 +170,22 @@ def netCDFtoZarr_sequential_recipe(daily_xarray_dataset, netcdf_local_paths, tmp
input_cache=tmp_cache,
)
return r, daily_xarray_dataset, tmp_target
+
+
+@pytest.fixture
+def netCDFtoZarr_sequential_multi_variable_recipe(
+ daily_xarray_dataset, netcdf_local_paths_by_variable, tmp_target, tmp_cache
+):
+ paths, items_per_file, fnames_by_variable, path_format = netcdf_local_paths_by_variable
+ pattern = VariableSequencePattern(
+ path_format, keys={"variable": ["foo", "bar"], "n": list(range(len(paths) / 2))}
+ )
+ r = recipe.NetCDFtoZarrMultiVarSequentialRecipe(
+ input_pattern=pattern,
+ sequence_dim="time",
+ inputs_per_chunk=1,
+ nitems_per_input=items_per_file,
+ target=tmp_target,
+ input_cache=tmp_cache,
+ )
+ return r, daily_xarray_dataset, tmp_target

MODIFY
tests/test_fixtures.py
tests/test_fixtures.py
@@ -15,7 +15,8 @@ def test_fixture_local_files(daily_xarray_dataset, netcdf_local_paths):
# TODO: this is quite repetetive of the test above. Replace with parametrization.
def test_fixture_local_files_by_variable(daily_xarray_dataset, netcdf_local_paths_by_variable):
- paths = [str(path) for path in netcdf_local_paths_by_variable]
+ paths, items_per_file, fnames_by_variable, path_format = netcdf_local_paths_by_variable
+ paths = [str(path) for path in paths]
ds = xr.open_mfdataset(paths, combine="by_coords", concat_dim="time")
assert ds.identical(daily_xarray_dataset)

MODIFY
tests/test_recipe.py
tests/test_recipe.py
import aiohttp
import pytest
import xarray as xr
+from pange_forge.patterns import VariableSequencePattern
from pangeo_forge import recipe
from pangeo_forge.storage import UninitializedTargetError
@@ -15,36 +16,7 @@ def incr_date(ds, filename=""):
return ds
-@pytest.mark.skip(reason="Removed this class for now")
-@pytest.mark.parametrize(
- "file_urls, files_per_chunk, expected_keys, expected_filenames",
- [
- (dummy_fnames, 1, [0, 1, 2], [("a.nc",), ("b.nc",), ("c.nc",)]),
- (dummy_fnames, 2, [0, 1], [("a.nc", "b.nc",), ("c.nc",),],), # noqa: E231
- ],
-)
-def test_sequence_recipe(file_urls, files_per_chunk, expected_keys, expected_filenames, tmp_target):
-
- r = recipe.SequenceRecipe(
- input_urls=file_urls,
- sequence_dim="time",
- inputs_per_chunk=files_per_chunk,
- target=tmp_target,
- chunk_preprocess_funcs=[],
- )
-
- assert r.sequence_len() == len(file_urls)
-
- chunk_keys = list(r.iter_chunks())
- assert chunk_keys == expected_keys
-
- for k, expected in zip(r.iter_chunks(), expected_filenames):
- fnames = r.inputs_for_chunk(k)
- assert fnames == expected
-
-
def _manually_execute_recipe(r):
- # this is the cannonical way to manually execute a recipe
for input_key in r.iter_inputs():
r.cache_input(input_key)
r.prepare_target()
@@ -105,7 +77,6 @@ def test_NetCDFtoZarrSequentialRecipe(
process_chunk=process_chunk,
)
- # this is the cannonical way to manually execute a recipe
_manually_execute_recipe(r)
ds_target = xr.open_zarr(tmp_target.get_mapper(), consolidated=True).load()
@@ -134,6 +105,23 @@ def test_NetCDFtoZarrSequentialRecipeNoTarget(
input_urls=paths, sequence_dim="time", inputs_per_chunk=1, nitems_per_input=items_per_file,
)
- # this is the cannonical way to manually execute a recipe
with pytest.raises(UninitializedTargetError):
r.cache_input(next(r.iter_inputs()))
+
+
+@pytest.fixture
+def test_NetCDFtoZarrMultiVarSequentialRecipe(
+ daily_xarray_dataset, netcdf_local_paths_by_variable, tmp_target, tmp_cache
+):
+ paths, items_per_file, fnames_by_variable, path_format = netcdf_local_paths_by_variable
+ pattern = VariableSequencePattern(
+ path_format, keys={"variable": ["foo", "bar"], "n": list(range(len(paths) / 2))}
+ )
+ _ = recipe.NetCDFtoZarrMultiVarSequentialRecipe(
+ input_pattern=pattern,
+ sequence_dim="time",
+ inputs_per_chunk=1,
+ nitems_per_input=items_per_file,
+ target=tmp_target,
+ input_cache=tmp_cache,
+ )

wip; doesnt work