Adapt for combine2
MODIFY
pangeo_forge_recipes/recipes/reference_hdf_zarr.py
pangeo_forge_recipes/recipes/reference_hdf_zarr.py
@@ -3,7 +3,7 @@ from __future__ import annotations
import json
import os
from dataclasses import dataclass, field
-from typing import Hashable, Iterable, Optional
+from typing import Callable, Hashable, Iterable, Optional
import fsspec
import yaml
@@ -39,13 +39,6 @@ def finalize(config: HDFReferenceRecipe):
assert config.storage_config.target is not None, "target is required"
assert config.storage_config.metadata is not None, "metadata_cache is required"
remote_protocol = fsspec.utils.get_protocol(next(config.file_pattern.items())[1])
- concat_args = config.xarray_concat_args.copy()
- if "dim" in concat_args:
- raise ValueError(
- "Please do not set 'dim' in xarray_concat_args. "
- "It is determined automatically from the File Pattern."
- )
- concat_args["dim"] = config.file_pattern.concat_dims[0] # there should only be one
files = list(
config.storage_config.metadata.getitems(
@@ -59,12 +52,16 @@ def finalize(config: HDFReferenceRecipe):
files,
remote_protocol=remote_protocol,
remote_options=config.netcdf_storage_options,
- xarray_open_kwargs=config.xarray_open_kwargs,
- xarray_concat_args=concat_args,
+ coo_dtypes=config.coo_dtypes,
+ coo_map=config.coo_map,
+ identical_dims=config.identical_dims,
+ concat_dims=config.file_pattern.concat_dims,
+ preprocess=config.preprocess,
+ postprocess=config.postprocess,
)
# mzz does not support directly writing to remote yet
# get dict version and push it
- out = mzz.translate(None, template_count=config.template_count)
+ out = mzz.translate()
fs = config.storage_config.target.fs
protocol = fs.protocol if isinstance(fs.protocol, str) else fs.protocol[0]
with config.storage_config.target.open(config.output_json_fname, mode="wt") as f:
@@ -133,12 +130,19 @@ class HDFReferenceRecipe(BaseRecipe, StorageMixin, FilePatternMixin):
inlined into the output reference file
:param output_storage_options: dict of kwargs for creating fsspec
instance when writing final output
- :param template_count: the number of occurrences of a URL before it
- gets made a template. ``None`` to disable templating
- :param xarray_open_kwargs: kwargs passed to xarray.open_dataset.
- Only used if `file_pattern` has more than one file.
- :param xarray_concat_args: kwargs passed to xarray.concat
- Only used if `file_pattern` has more than one file.
+ :param identical_dims: coordiate-like variables that are assumed to be the
+ same in every input, and so do not need any concatenation
+ :param coo_map: set of "selectors" defining how to fetch the dimension coordinates
+ of any given chunk for each of the concat dimes. By default, this is the variable
+ in the dataset with the same name as the given concat dim, except for "var",
+ where the default is the name of each input variable. See the doc of
+ MultiZarrToZarr for more details on possibilities.
+ :param coo_dtyes: optional coercion of coordinate values before write
+ Note that, if using cftime to read coordinate values, output will also be
+ encoded with cftime (i.e., ints and special attributes) unless you specify
+ an "M8[*]" as the output type.
+ :param preprocess: a function applied to each HDF file's references before combine
+ :param postprocess: a function applied to the global combined references before write
"""
_compiler = hdf_reference_recipe_compiler
@@ -149,15 +153,18 @@ class HDFReferenceRecipe(BaseRecipe, StorageMixin, FilePatternMixin):
# store_chunk, finalize_target. The strategy was used for NWM's 370k files.
# TODO: as written, this recipe is specific to HDF5 files. fsspec-reference-maker
- # also supports TIFF and grib2 (and more coming)
+ # also supports TIFF, FITS, netCDF3 and grib2 (and more coming)
output_json_fname: str = "reference.json"
output_intake_yaml_fname: str = "reference.yaml"
netcdf_storage_options: dict = field(default_factory=dict)
inline_threshold: int = 500
output_storage_options: dict = field(default_factory=dict)
- template_count: Optional[int] = 20
- xarray_open_kwargs: dict = field(default_factory=dict)
- xarray_concat_args: dict = field(default_factory=dict)
+ concat_dims: list = field(default_factory=list)
+ identical_dims: Optional[list] = field(default_factory=list)
+ coo_map: Optional[dict] = field(default_factory=dict)
+ coo_dtypes: Optional[dict] = field(default_factory=dict)
+ preprocess: Optional[Callable] = None
+ postprocess: Optional[Callable] = None
def __post_init__(self):
self._validate_file_pattern()

MODIFY
tests/recipe_tests/test_HDFReferenceRecipe.py
tests/recipe_tests/test_HDFReferenceRecipe.py
@@ -13,7 +13,7 @@ HDFReferenceRecipe = reference_hdf_zarr.HDFReferenceRecipe
@pytest.mark.parametrize("default_storage", [True, False])
-@pytest.mark.parametrize("with_intake", [True, False])
+@pytest.mark.parametrize("with_intake", [False, True])
def test_single(
netcdf_local_file_pattern_sequential, tmpdir, with_intake, default_storage, execute_recipe,
):
@@ -21,6 +21,7 @@ def test_single(
path = list(file_pattern.items())[0][1]
expected = xr.open_dataset(path, engine="h5netcdf")
recipe = HDFReferenceRecipe(file_pattern)
+ recipe.coo_map = {"time": "cf:time"} # ensure cftime encoding end-to-end
if default_storage:
out_target = recipe.storage_config.target
@@ -44,8 +45,9 @@ def test_single(
"reference", fo=out_target._full_path("reference.json"), remote_protocol="file"
)
m = fs.get_mapper("")
- ds = xr.open_dataset(m, engine="zarr", chunks={}, consolidated=False)
- assert (ds.foo == expected.foo).all()
+ ds = xr.open_dataset(m, engine="zarr", chunks={}, backend_kwargs=dict(consolidated=False))
+
+ assert (expected.foo.values == ds.foo.values[: 1 * len(expected.time)]).all()
@pytest.mark.parametrize("with_intake", [True, False])
@@ -54,6 +56,7 @@ def test_multi(netcdf_local_file_pattern_sequential, tmpdir, with_intake, execut
paths = [f for _, f in list(file_pattern.items())]
expected = xr.open_mfdataset(paths, engine="h5netcdf")
recipe = HDFReferenceRecipe(file_pattern)
+ recipe.coo_map = {"time": "cf:time"} # ensure cftime encoding end-to-end
# make sure assigning storage later works
out_target = FSSpecTarget(fs=fsspec.filesystem("file"), root_path=str(tmpdir))

Adapt for combine2