anemoi-datasets 0.5.19__py3-none-any.whl → 0.5.20__py3-none-any.whl

anemoi/datasets/data/__init__.py

@@ -16,6 +16,7 @@ from typing import Set
 # from .dataset import Shape
 # from .dataset import TupleIndex
 from .misc import _open_dataset
+from .misc import _save_dataset
 from .misc import add_dataset_path
 from .misc import add_named_dataset
 
@@ -92,6 +93,21 @@ def open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
     return ds
 
 
+def save_dataset(recipe: dict, zarr_path: str, n_workers: int = 1) -> None:
+    """Open a dataset and save it to disk.
+
+    Parameters
+    ----------
+    recipe : dict
+        Recipe used with open_dataset (not a dataset creation recipe).
+    zarr_path : str
+        Path to store the obtained anemoi dataset to disk.
+    n_workers : int
+        Number of workers to use for parallel processing. If none, sequential processing will be performed.
+    """
+    _save_dataset(recipe, zarr_path, n_workers)
+
+
 def list_dataset_names(*args: Any, **kwargs: Any) -> list[str]:
     """List the names of datasets.
 

anemoi/datasets/data/complement.py

@@ -207,7 +207,7 @@ class ComplementNone(Complement):
 class ComplementNearest(Complement):
     """A class to complement a target dataset with variables from a source dataset using nearest neighbor interpolation."""
 
-    def __init__(self, target: Any, source: Any) -> None:
+    def __init__(self, target: Any, source: Any, max_distance: float = None) -> None:
         """Initializes the ComplementNearest class.
 
         Parameters
@@ -216,6 +216,8 @@ class ComplementNearest(Complement):
             The target dataset.
         source : Any
             The source dataset.
+        max_distance : float, optional
+            The maximum distance for nearest neighbor interpolation, default is None.
         """
         super().__init__(target, source)
 
@@ -224,6 +226,7 @@ class ComplementNearest(Complement):
             self._source.longitudes,
             self._target.latitudes,
             self._target.longitudes,
+            max_distance=max_distance,
         )
 
     def check_compatibility(self, d1: Dataset, d2: Dataset) -> None:

anemoi/datasets/data/dataset.py

@@ -237,6 +237,13 @@ class Dataset(ABC, Sized):
 
             return Statistics(self, open_dataset(statistics))._subset(**kwargs).mutate()
 
+        # Note: trim_edge should go before thinning
+        if "trim_edge" in kwargs:
+            from .masked import TrimEdge
+
+            edge = kwargs.pop("trim_edge")
+            return TrimEdge(self, edge)._subset(**kwargs).mutate()
+
         if "thinning" in kwargs:
             from .masked import Thinning
 
@@ -284,6 +291,13 @@ class Dataset(ABC, Sized):
             interpolate_frequency = kwargs.pop("interpolate_frequency")
             return InterpolateFrequency(self, interpolate_frequency)._subset(**kwargs).mutate()
 
+        if "interpolate_variables" in kwargs:
+            from .interpolate import InterpolateNearest
+
+            interpolate_variables = kwargs.pop("interpolate_variables")
+            max_distance = kwargs.pop("max_distance", None)
+            return InterpolateNearest(self, interpolate_variables, max_distance=max_distance)._subset(**kwargs).mutate()
+
         # Keep last
         if "shuffle" in kwargs:
             from .subset import Subset

anemoi/datasets/data/interpolate.py

@@ -13,7 +13,11 @@ import logging
 from functools import cached_property
 from typing import Any
 from typing import Dict
+from typing import List
+from typing import Optional
 from typing import Set
+from typing import Tuple
+from typing import Union
 
 import numpy as np
 from anemoi.utils.dates import frequency_to_timedelta
@@ -214,3 +218,75 @@ class InterpolateFrequency(Forwards):
         return {
             # "frequency": frequency_to_string(self._frequency),
         }
+
+
+class InterpolateNearest(Forwards):
+    def __init__(
+        self, dataset: Dataset, interpolate_variables: List[str], max_distance: Optional[float] = None
+    ) -> None:
+        """Initialize the InterpolateNearest class.
+
+        Parameters
+        ----------
+        dataset : Dataset
+            The dataset to be interpolated.
+        interpolate_variables : List[str]
+            The variables to be interpolated.
+        max_distance : Optional[float], optional
+            The maximum distance for nearest neighbor search, by default None.
+        """
+        from ..grids import nearest_grid_points
+
+        super().__init__(dataset)
+        self.vars = interpolate_variables
+        self.var_mask = {
+            v: np.where(~np.isnan(dataset[0, dataset.name_to_index[v]].squeeze()))[0] for v in interpolate_variables
+        }
+        self.dataset = dataset
+        self.ngp_var = {
+            v: nearest_grid_points(
+                dataset.latitudes[self.var_mask[v]],
+                dataset.longitudes[self.var_mask[v]],
+                dataset.latitudes,
+                dataset.longitudes,
+                max_distance=max_distance,
+            )
+            for v in interpolate_variables
+        }
+
+    def tree(self) -> Node:
+        return Node(self, [self.forward.tree()], interpolate_variables=self.vars)
+
+    @property
+    def shape(self) -> Shape:
+        return self.forward.shape
+
+    @property
+    def metadata(self) -> Dict[str, Any]:
+        return self.forward.metadata()
+
+    @staticmethod
+    def slice_len(slice_obj: slice) -> int:
+        # Compute the length of the slice
+        return max(0, (slice_obj.stop - slice_obj.start + (slice_obj.step or 1) - 1) // (slice_obj.step or 1))
+
+    @expand_list_indexing
+    def _get_tuple(self, index: TupleIndex) -> NDArray[Any]:
+        index, changes = index_to_slices(index, self.shape)
+        source_data = self.forward[index[0]]
+        target_data = source_data.copy()
+        vars_to_interpolate = [self.forward.name_to_index[v] for v in self.vars]
+        for v, i in zip(self.vars, vars_to_interpolate):
+            target_data[:, i, ...] = source_data[:, i][..., self.var_mask[v][self.ngp_var[v]]]
+        result = target_data[(slice(None),) + index[1:]]
+        return apply_index_to_slices_changes(result, changes)
+
+    def __getitem__(self, index: Union[int, slice, Tuple[Union[int, slice], ...]]) -> NDArray[Any]:
+        if isinstance(index, (int, slice)):
+            index = (index, slice(None), slice(None), slice(None))
+        return self._get_tuple(index)
+
+    def subclass_metadata_specific(self) -> Dict[str, Any]:
+        return {
+            "interpolate_variables": self.vars,
+        }

anemoi/datasets/data/masked.py

@@ -254,3 +254,80 @@ class Cropping(Masked):
             The metadata specific to the Cropping subclass.
         """
         return dict(area=self.area)
+
+
+class TrimEdge(Masked):
+    """A class that removes the boundary of a dataset."""
+
+    def __init__(self, forward, edge):
+        if isinstance(edge, int):
+            self.edge = [edge] * 4
+        elif isinstance(edge, (list, tuple)) and len(edge) == 4:
+            self.edge = edge
+        else:
+            raise ValueError("'edge' must be an integer or a list of 4 integers")
+
+        for e in self.edge:
+            if not isinstance(e, int) or e < 0:
+                raise ValueError("'edge' must be integer(s) 0 or greater")
+
+        shape = forward.field_shape
+        if len(shape) != 2:
+            raise ValueError("TrimEdge only works on regular grids")
+
+        if self.edge[0] + self.edge[1] >= shape[0]:
+            raise ValueError("Too much triming of the first grid dimension, resulting in an empty dataset")
+        if self.edge[2] + self.edge[3] >= shape[1]:
+            raise ValueError("Too much triming of the second grid dimension, resulting in an empty dataset")
+
+        mask = np.full(shape, True, dtype=bool)
+        mask[0 : self.edge[0], :] = False
+        mask[:, 0 : self.edge[2]] = False
+        if self.edge[1] != 0:
+            mask[-self.edge[1] :, :] = False
+        if self.edge[3] != 0:
+            mask[:, -self.edge[3] :] = False
+
+        mask = mask.flatten()
+
+        super().__init__(forward, mask)
+
+    def mutate(self) -> Dataset:
+        """Mutate the dataset.
+
+        Returns
+        -------
+        Dataset
+            The mutated dataset.
+        """
+        if self.edge is None:
+            return self.forward.mutate()
+        return super().mutate()
+
+    def tree(self) -> Node:
+        """Get the tree representation of the dataset.
+
+        Returns
+        -------
+        Node
+            The tree representation of the dataset.
+        """
+        return Node(self, [self.forward.tree()], edge=self.edge)
+
+    def forwards_subclass_metadata_specific(self) -> Dict[str, Any]:
+        """Get the metadata specific to the TrimEdge subclass.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata specific to the TrimEdge subclass.
+        """
+        return dict(edge=self.edge)
+
+    @property
+    def field_shape(self) -> Shape:
+        """Returns the field shape of the dataset."""
+        x, y = self.forward.field_shape
+        x -= self.edge[0] + self.edge[1]
+        y -= self.edge[2] + self.edge[3]
+        return x, y

anemoi/datasets/data/misc.py

@@ -539,3 +539,162 @@ def _open_dataset(*args: Any, **kwargs: Any) -> "Dataset":
         return dataset._subset(**kwargs)
 
     return sets[0]._subset(**kwargs)
+
+
+def append_to_zarr(new_data: np.ndarray, new_dates: np.ndarray, zarr_path: str) -> None:
+    """Append data from a subset (for one date) to the Zarr store.
+
+    Parameters
+    ----------
+    new_data : np.ndarray
+        The new data to append.
+    new_dates : np.ndarray
+        The new dates to append.
+    zarr_path : str
+        The path to the Zarr store.
+
+    Notes
+    -----
+    - "dates" dataset is created with chunks equal to len(big_dataset.dates).
+    - "data" dataset is created with chunk size 1 along the first (time) dimension.
+    """
+    print("Appending data for", new_dates, flush=True)
+    # Re-open the zarr store to avoid root object accumulating memory.
+    root = zarr.open(zarr_path, mode="a")
+    # Convert new dates to strings (using str) regardless of input dtype.
+    new_dates = np.array(new_dates, dtype="datetime64[ns]")
+    dates_ds = root["dates"]
+    old_len = dates_ds.shape[0]
+    dates_ds.resize((old_len + len(new_dates),))
+    dates_ds[old_len:] = new_dates
+    # Append to "data" dataset.
+    data_ds = root["data"]
+    old_shape = data_ds.shape  # (time, n_vars, ensembles, gridpoints)
+    new_shape = (old_shape[0] + len(new_dates),) + old_shape[1:]
+    data_ds.resize(new_shape)
+    data_ds[old_shape[0] :] = new_data
+
+
+def process_date(date: Any, big_dataset: Any) -> Tuple[np.ndarray, np.ndarray]:
+    """Open the subset corresponding to the given date and return (date, subset).
+
+    Parameters
+    ----------
+    date : Any
+        The date to process.
+    big_dataset : Any
+        The dataset to process.
+
+    Returns
+    -------
+    Tuple[np.ndarray, np.ndarray]
+        The subset and the date.
+    """
+    print("Processing:", date, flush=True)
+    subset = _open_dataset(big_dataset, start=date, end=date).mutate()
+    s = subset[:]
+    date = subset.dates
+    return s, date
+
+
+def initialize_zarr_store(root: Any, big_dataset: Any, recipe: Dict[str, Any]) -> None:
+    """Initialize the Zarr store with the given dataset and recipe.
+
+    Parameters
+    ----------
+    root : Any
+        The root of the Zarr store.
+    big_dataset : Any
+        The dataset to initialize the store with.
+    recipe : Dict[str, Any]
+        The recipe for initializing the store.
+    """
+    ensembles = big_dataset.shape[1]
+    # Create or append to "dates" dataset.
+    if "dates" not in root:
+        full_length = len(big_dataset.dates)
+        root.create_dataset("dates", data=np.array([], dtype="datetime64[s]"), chunks=(full_length,))
+
+    if "data" not in root:
+        dims = (1, len(big_dataset.variables), ensembles, big_dataset.grids[0])
+        root.create_dataset(
+            "data",
+            shape=dims,
+            dtype=np.float64,
+            chunks=dims,
+        )
+
+    for k, v in big_dataset.statistics.items():
+        if k not in root:
+            root.create_dataset(
+                k,
+                data=v,
+                compressor=None,
+            )
+
+    # Create spatial coordinate datasets if missing.
+    if "latitudes" not in root or "longitudes" not in root:
+        root.create_dataset("latitudes", data=big_dataset.latitudes, compressor=None)
+        root.create_dataset("longitudes", data=big_dataset.longitudes, compressor=None)
+
+    # Set store-wide attributes if not already set.
+    if "frequency" not in root.attrs:
+        root.attrs["frequency"] = "10m"
+        root.attrs["resolution"] = "1km"
+        root.attrs["name_to_index"] = {k: i for i, k in enumerate(big_dataset.variables)}
+        root.attrs["ensemble_dimension"] = 1
+        root.attrs["field_shape"] = big_dataset.field_shape
+        root.attrs["flatten_grid"] = True
+        root.attrs["recipe"] = recipe
+
+
+def _save_dataset(recipe: Dict[str, Any], zarr_path: str, n_workers: int = 1) -> None:
+    """Incrementally create (or update) a Zarr store from an Anemoi dataset.
+
+    Parameters
+    ----------
+    recipe : Dict[str, Any]
+        The recipe for creating the dataset.
+    zarr_path : str
+        The path to the Zarr store.
+    n_workers : int, optional
+        The number of worker processes to use, by default 1.
+
+    Notes
+    -----
+    Worker processes extract data for each date in parallel, but all writes
+    to the store happen sequentially in the main process (i.e. single-writer).
+
+    The "dates" dataset is created with chunking equal to the full length of
+    big_dataset.dates, while "data" is chunked with 1 in the time dimension.
+    """
+    from concurrent.futures import ProcessPoolExecutor
+
+    full_ds = _open_dataset(recipe).mutate()
+    print("Opened full dataset.", flush=True)
+
+    # Use ProcessPoolExecutor for parallel data extraction.
+    # Workers return (date, subset) tuples.
+    root = zarr.open(zarr_path, mode="a")
+    initialize_zarr_store(root, full_ds, recipe)
+    print("Zarr store initialized.", flush=True)
+
+    existing_dates = np.array(sorted(root["dates"]), dtype="datetime64[s]")
+    all_dates = full_ds.dates
+    # To resume creation of the Zarr store in case the job is interrupted.
+    dates_to_process = np.array(sorted(set(all_dates).difference(existing_dates)), dtype="datetime64[s]")
+
+    use_pool = False
+
+    if use_pool:
+        with ProcessPoolExecutor(n_workers) as pool:
+            futures = [pool.submit(process_date, date, full_ds) for date in dates_to_process]
+            for future in futures:
+                subset, date = future.result()
+                # All appends happen sequentially here to
+                #  avoid dates being added in a random order.
+                append_to_zarr(subset, date, zarr_path)
+    else:
+        for date in dates_to_process:
+            subset, date = process_date(date, full_ds)
+            append_to_zarr(subset, date, zarr_path)

anemoi/datasets/grids.py

@@ -604,6 +604,7 @@ def nearest_grid_points(
     source_longitudes: NDArray[Any],
     target_latitudes: NDArray[Any],
     target_longitudes: NDArray[Any],
+    max_distance: float = None,
 ) -> NDArray[Any]:
     """Find the nearest grid points from source to target coordinates.
 
@@ -617,6 +618,9 @@ def nearest_grid_points(
         Target latitude coordinates.
     target_longitudes : NDArray[Any]
         Target longitude coordinates.
+    max_distance: float, optional
+        Maximum distance between nearest point and point to interpolate. Defaults to None.
+        For example, 1e-3 is 1 km.
 
     Returns
     -------
@@ -632,8 +636,10 @@ def nearest_grid_points(
 
     target_xyz = latlon_to_xyz(target_latitudes, target_longitudes)
     target_points = np.array(target_xyz).transpose()
-
-    _, indices = cKDTree(source_points).query(target_points, k=1)
+    if max_distance is None:
+        _, indices = cKDTree(source_points).query(target_points, k=1)
+    else:
+        _, indices = cKDTree(source_points).query(target_points, k=1, distance_upper_bound=max_distance)
     return indices
 
 

{anemoi_datasets-0.5.19.dist-info → anemoi_datasets-0.5.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-datasets
-Version: 0.5.19
+Version: 0.5.20
 Summary: A package to hold various functions to support training of ML models on ECMWF data.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License:                                  Apache License
@@ -205,6 +205,7 @@ License:                                  Apache License
            See the License for the specific language governing permissions and
            limitations under the License.
         
+Project-URL: Changelog, https://github.com/ecmwf/anemoi-datasets/CHANGELOG.md
 Project-URL: Documentation, https://anemoi-datasets.readthedocs.io/
 Project-URL: Homepage, https://github.com/ecmwf/anemoi-datasets/
 Project-URL: Issues, https://github.com/ecmwf/anemoi-datasets/issues
@@ -224,8 +225,8 @@ Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 License-File: LICENSE
-Requires-Dist: anemoi-transform>=0.1.4
-Requires-Dist: anemoi-utils[provenance]>=0.4.16
+Requires-Dist: anemoi-transform>=0.1.6
+Requires-Dist: anemoi-utils[provenance]>=0.4.18
 Requires-Dist: cfunits
 Requires-Dist: numpy
 Requires-Dist: pyyaml
@@ -233,8 +234,12 @@ Requires-Dist: semantic-version
 Requires-Dist: tqdm
 Requires-Dist: zarr<=2.18.4
 Provides-Extra: all
-Requires-Dist: anemoi-datasets[create,remote,xarray]; extra == "all"
+Requires-Dist: anemoi-datasets[comparelam,create,remote,xarray]; extra == "all"
+Provides-Extra: comparelam
+Requires-Dist: prettytable; extra == "comparelam"
+Requires-Dist: termcolor; extra == "comparelam"
 Provides-Extra: create
+Requires-Dist: cachetools; extra == "create"
 Requires-Dist: earthkit-data[mars]>=0.12.4; extra == "create"
 Requires-Dist: earthkit-geo>=0.3; extra == "create"
 Requires-Dist: earthkit-meteo>=0.3; extra == "create"
@@ -244,6 +249,7 @@ Requires-Dist: pyproj>=3; extra == "create"
 Provides-Extra: dev
 Requires-Dist: anemoi-datasets[all,docs,tests]; extra == "dev"
 Provides-Extra: docs
+Requires-Dist: anemoi-datasets[all]; extra == "docs"
 Requires-Dist: nbsphinx; extra == "docs"
 Requires-Dist: pandoc; extra == "docs"
 Requires-Dist: sphinx; extra == "docs"

{anemoi_datasets-0.5.19.dist-info → anemoi_datasets-0.5.20.dist-info}/RECORD

@@ -1,15 +1,17 @@
 anemoi/datasets/__init__.py,sha256=i_wsAT3ezEYF7o5dpqGrpoG4wmLS-QIBug18uJbSYMs,1065
 anemoi/datasets/__main__.py,sha256=ErwAqE3rBc7OaNO2JRsEOhWpB8ldjAt7BFSuRhbnlqQ,936
-anemoi/datasets/_version.py,sha256=c9HpQ99YdGTjwGSu9DP-RLCHGpdS8jfmCirVxrPm0i4,513
-anemoi/datasets/grids.py,sha256=ALvRRMvu0GaDCnNlOO-cRCfbpywA-1w_wzSylqpqgNY,17795
+anemoi/datasets/_version.py,sha256=0Ib1b9ijC26N78gsH6oTzFbmTu3c8De1ac8TlFGaEZU,513
+anemoi/datasets/grids.py,sha256=Hhj1aOXHvDjmI46M_UlLSjCs1qYqxH-uqd_kapDSdbU,18134
 anemoi/datasets/testing.py,sha256=fy_JzavUwLlK_2rtXAT-UGUyo5gjyQW2y826zf334Wg,2645
 anemoi/datasets/commands/__init__.py,sha256=O5W3yHZywRoAqmRUioAr3zMCh0hGVV18wZYGvc00ioM,698
 anemoi/datasets/commands/cleanup.py,sha256=FX082xkHKCSd8d-FUN5zDBSiKA-QYQEeUZ6dCUD-Ob8,1816
+anemoi/datasets/commands/compare-lam.py,sha256=F5GYRsKOtdhDePhifgf1TCj5L2T8EVIA2N8AdO-AKKY,14857
 anemoi/datasets/commands/compare.py,sha256=jzhjbbt1U-YANTVRBhrwSh2CcYgk4qX2IiTMJtcn82s,3678
 anemoi/datasets/commands/copy.py,sha256=UlvW9YIlP7jwKAY7TikVAfhkrjaQ9Kkxqfx8jEut-Jg,16010
 anemoi/datasets/commands/create.py,sha256=5BXdPZMO-ZULBnEVgyeSS-IMy4p84HSyFVG855gqj3k,6598
 anemoi/datasets/commands/finalise-additions.py,sha256=2LqU7ke3i-yRQbjkgldX6e2QlyE-tKqp0b6QOhJF19g,1985
 anemoi/datasets/commands/finalise.py,sha256=-YtN9wFFDrM_i_V9YHoXZsajF3eAax-73Zsi4uHAFCI,1709
+anemoi/datasets/commands/grib-index.py,sha256=1aO4Lm_7ZpRF_lRRF-OEMZ37ByYPZfUAIFW7jOsIf7Q,3161
 anemoi/datasets/commands/init-additions.py,sha256=2vMom5L38UvLLopzP2z-R_Fq31fU2uMvKXoDq5d8oI4,1931
 anemoi/datasets/commands/init.py,sha256=5IKyJ_hJA4lLIbpT88XtcGzXccHLSGwSoqVSvVJGxPg,2852
 anemoi/datasets/commands/inspect.py,sha256=dGx_3IQTNaV3iGfwhHp_G1fIecBWeRbOLU9ckrZ7epg,26578
@@ -39,8 +41,8 @@ anemoi/datasets/create/filters/empty.py,sha256=Dw1kUnAlFt6b5ds0kmrw9Gak09XjSqF8m
 anemoi/datasets/create/filters/legacy.py,sha256=6JY6uX7m-8NZjoZ1sqs0EAqT-uorvnZ-eFOwMU3LmRU,2536
 anemoi/datasets/create/filters/noop.py,sha256=WHl-k3NojGJMX4iNYxQ6Ln21pM8ERP4z8pQ5zLRDvXs,1019
 anemoi/datasets/create/filters/orog_to_z.py,sha256=vnZ1hD9LXoOfHCIbzkurMuBl_NSfXSiiHS2yZt8ndeQ,1784
-anemoi/datasets/create/filters/pressure_level_relative_humidity_to_specific_humidity.py,sha256=2xWyqIZm8AFqS9IDPylOB3UfFu9aqe47aZ8qAbmpNUo,2558
-anemoi/datasets/create/filters/pressure_level_specific_humidity_to_relative_humidity.py,sha256=L7cZvPZ1zNHbfbHQw-jn1TZEfO0i0JX5O7AoytuLy1Q,2575
+anemoi/datasets/create/filters/pressure_level_relative_humidity_to_specific_humidity.py,sha256=dBAQFNAc3GEZ_HwyDrcctFaKZQYddh2ldnDA2XSfSRg,2646
+anemoi/datasets/create/filters/pressure_level_specific_humidity_to_relative_humidity.py,sha256=51t6kbNZsXK87H0HVR9j0a54siokID47ve9r1a8rOLE,2663
 anemoi/datasets/create/filters/rename.py,sha256=pKi3CU6fvox2sPH7szXdA79NjdIcSz59IB7HsiS_9Co,5779
 anemoi/datasets/create/filters/rotate_winds.py,sha256=fVyAbypO_EsENHjQCujbEXp2gUEb97sMoG0s4YiPXfc,3102
 anemoi/datasets/create/filters/single_level_dewpoint_to_relative_humidity.py,sha256=hCS3yiN9nZf-P6shQmBm5Or9rMOwU1fTwHw_qFIjT9s,2378
@@ -52,9 +54,9 @@ anemoi/datasets/create/filters/sum.py,sha256=aGT6JkdHJ3i2SKzklqiyJ4ZFV3bVMYhHOSo
 anemoi/datasets/create/filters/transform.py,sha256=gIDLvaJlnn3Nc6P29aPOvNYM6yBWcIGrR2e_1bM6_Nw,1418
 anemoi/datasets/create/filters/unrotate_winds.py,sha256=3AJf0crnVVySLlXLIdfEUxRRlQeKgheUuD-UCrSrgo8,2798
 anemoi/datasets/create/filters/uv_to_speeddir.py,sha256=Zdc34AG5Bsz-Z7JGuznyRJr6F-BnWKXPiI3mjmOpbek,2883
-anemoi/datasets/create/filters/wz_to_w.py,sha256=42AQvTHk-ISyHlwvXfU3yiTGiDkfrs5kiKNkyqqtQpg,2725
+anemoi/datasets/create/filters/wz_to_w.py,sha256=slOiX5RibG48Zrkss8Qjpb-8ZTnvSvmKlk1Hy45_wzU,2812
 anemoi/datasets/create/input/__init__.py,sha256=XeURpmbReQvpELltGFKzg3oZFXWRdUxW9SK3K662SBQ,3364
-anemoi/datasets/create/input/action.py,sha256=xXLqVsoygxyaROiXc7TW9DCEOzVh1YgPDAqUpcOb9fs,7619
+anemoi/datasets/create/input/action.py,sha256=WmRPe5ZYQz8vxAtOr6hLYGLzikcldyps6dG3BwHiBp8,7709
 anemoi/datasets/create/input/concat.py,sha256=bU8SWfBVfK8bRAmmN4UO9zpIGxwQvRUk9_vwrKPOTE4,5355
 anemoi/datasets/create/input/context.py,sha256=qrLccxMe9UkyQxsNuR6JSK7oLzZq21dt38AxZ9kYzsY,2714
 anemoi/datasets/create/input/data_sources.py,sha256=4xUUShM0pCXIZVPJW_cSNMUwCO_wLx996MLFpTLChm0,4385
@@ -65,18 +67,20 @@ anemoi/datasets/create/input/join.py,sha256=RAdgE4lVcC71_J47dNa1weJuWdTXSQIvo06G
 anemoi/datasets/create/input/misc.py,sha256=FVaH_ym52RZI_fnLSMM_dKTQmWTrInucP780E3gGqvw,3357
 anemoi/datasets/create/input/pipe.py,sha256=-tCz161IwXoI8pl1hilA9T_j5eHSr-sgbijFLp9HHNc,2083
 anemoi/datasets/create/input/repeated_dates.py,sha256=HaPzDCNHQBY1VVp6gvd3drwjWjYpSBh-GLgHqBRJTz0,12012
-anemoi/datasets/create/input/result.py,sha256=wZHquaytAatwqGZifHzJDsysed_cEwSr9X2ND6Yq9o0,24218
+anemoi/datasets/create/input/result.py,sha256=wX5iAAtm50ezLn2QsIDIQnGEGmgE08lWltRwnkFWSLw,24220
 anemoi/datasets/create/input/step.py,sha256=WcR9NgRvUKF60Fo5veLvRCAQMrOd55x1gOEAmd2t2r4,5948
 anemoi/datasets/create/input/template.py,sha256=Iycw9VmfA0WEIDP_Of8bp-8HsV0EUfwbnm0WjxiO4GA,4092
 anemoi/datasets/create/input/trace.py,sha256=dakPYMmwKq6s17Scww1CN-xYBD3btJTGeDknOhAcnEM,3320
 anemoi/datasets/create/sources/__init__.py,sha256=XNiiGaC6NbxnGfl6glPw-gTJASi3vsGKwVlfkMqYGk4,950
 anemoi/datasets/create/sources/accumulations.py,sha256=ZA8F8RJPMHok5RpIHH4x-txwiSll8zuWwqJ3rn95JHk,20295
 anemoi/datasets/create/sources/accumulations2.py,sha256=iBORRrH0N7r3gMWm3mCkJ6XmB-dO_lEckHPwvmk9fu0,20673
+anemoi/datasets/create/sources/anemoi_dataset.py,sha256=2xJJTmKlv87F_2ECMKeehaeW7_oWLlDcLt8C_Prp1RI,2017
 anemoi/datasets/create/sources/constants.py,sha256=5O6d9tEuAmVjl5vNkNfmkaAjKXFlw1UjeueTsF1GZCI,1528
 anemoi/datasets/create/sources/eccc_fstd.py,sha256=8HK38f444HcWMvBhooP0XqTfMXYoCbN_8G9RI_Ne5rc,659
 anemoi/datasets/create/sources/empty.py,sha256=5mVIVRUwnBfE3zp-bvNA_imXCSpyds-4mewcI8HXAiY,1020
 anemoi/datasets/create/sources/forcings.py,sha256=877OZoXUoJncQ2_AAGSijwWqM-4kJJdxdIa6SFvZBUw,1216
-anemoi/datasets/create/sources/grib.py,sha256=zFBFWNFDVPCMSDRheNuaLZ7EaInjDt9OTJwVOPj9j-U,8371
+anemoi/datasets/create/sources/grib.py,sha256=nGjrKTQZB8GQCvfX7SgfdmcVJVSXF9kpb2GmSX7H7bM,8731
+anemoi/datasets/create/sources/grib_index.py,sha256=Pnm0RLga9lpD4MqVaZr7IqXMBlw1DtTIWZRfz7fq30Q,19026
 anemoi/datasets/create/sources/hindcasts.py,sha256=_4880rgd4AsRxlDXVi6dkh8mlKXrz2i27btVlmlMFjY,2611
 anemoi/datasets/create/sources/legacy.py,sha256=RJce-9TwmUUCFbgC8A3Dp61nSBfB8_lWti8WNoOMIcU,2652
 anemoi/datasets/create/sources/mars.py,sha256=tesQz7Ne6SLBChE_cNJU6Sxr6e0LXFlUKQ8gCdRiCMw,13155
@@ -91,11 +95,11 @@ anemoi/datasets/create/sources/xarray_kerchunk.py,sha256=vdFaFzze8VLjYUgIX8Lc39E
 anemoi/datasets/create/sources/xarray_zarr.py,sha256=McY-vgXmUbGAkBViAfYwBUeVmGUU-Qr8UW-jUGu5-9s,1209
 anemoi/datasets/create/sources/zenodo.py,sha256=KEetFEk5GzGFpoos8rbBQBTa2XElWG7oTYjfZXgbu0Q,2065
 anemoi/datasets/create/sources/xarray_support/README.md,sha256=56olM9Jh0vI0_bU9GI-IqbBcz4DZXWONqvdzN_VeAFE,78
-anemoi/datasets/create/sources/xarray_support/__init__.py,sha256=CLq5gcHhP5XGzr_LD6D4VPYqPUECmlYztzVC842KMU0,5563
+anemoi/datasets/create/sources/xarray_support/__init__.py,sha256=8Dv7KQW8O3VvHOsSbqYdjaaomYIhXIKgSGatnNEweNU,5564
 anemoi/datasets/create/sources/xarray_support/coordinates.py,sha256=rPEuijS77mQ9V9tpN7wjg-w9rBxj7bZf_c30lLgSscE,11029
 anemoi/datasets/create/sources/xarray_support/field.py,sha256=YRxx6kh1qO2qQ6I_VyR51h3dwNiiFM7CNwQNfpp-p-E,6375
-anemoi/datasets/create/sources/xarray_support/fieldlist.py,sha256=CG8ecTXCr37pNiykoxR6Sb21Xxsz6AS5K5-KE4qMEmo,8243
-anemoi/datasets/create/sources/xarray_support/flavour.py,sha256=GYodfpKfTBBWiyXytRrin6NK07ltlyz0UF7x4gQ3Fok,31836
+anemoi/datasets/create/sources/xarray_support/fieldlist.py,sha256=UyUljq2Ax-PpQ-bvG4Dsi_lkZucuPgCy120EadDeUMU,8271
+anemoi/datasets/create/sources/xarray_support/flavour.py,sha256=UyfzBjYMNfugMCq-r5Ie3qDuorLwaalPi_0oZHckZcg,32073
 anemoi/datasets/create/sources/xarray_support/grid.py,sha256=lsE8bQwBH9pflzvsJ89Z6ExYPdHJd54xorMNzL2gTd0,6181
 anemoi/datasets/create/sources/xarray_support/metadata.py,sha256=WRO86l-ZB7iJ7pG5Vz9kVv5h1MokfF0fuy0bNSNBRIc,10687
 anemoi/datasets/create/sources/xarray_support/patch.py,sha256=Snk8bz7gp0HrG0MrY5hrXu7VC0tKgtoiWXByi2sBYJc,2037
@@ -103,10 +107,10 @@ anemoi/datasets/create/sources/xarray_support/time.py,sha256=Y_lZTUOXWJH4jcSgyL4
 anemoi/datasets/create/sources/xarray_support/variable.py,sha256=fcazws9vuizmx55JCXwbkwffg4WxJllPrEg2US1VysE,9163
 anemoi/datasets/create/statistics/__init__.py,sha256=_BuPcuUrwQAEcMQVds93EV9M5ys2ao8jCWKV4OVoSSA,18291
 anemoi/datasets/create/statistics/summary.py,sha256=JdtChTmsr1Y958_nka36HltTbeZkawuGbprbfZD7Ux8,4790
-anemoi/datasets/data/__init__.py,sha256=dLzKYFX0eCi7urHA9t530SqZ_GYxTUyQeEcXYV8lZho,2521
-anemoi/datasets/data/complement.py,sha256=C55ZyWO8uM-bGbZkpuh80z95XtQjIr_NBnsxiKDWWtE,9643
+anemoi/datasets/data/__init__.py,sha256=wzhk_7VQImge12Xkg99xuiFOC7DAjBW1mu446y0Iq60,3057
+anemoi/datasets/data/complement.py,sha256=QcbcAW3HaG1g4QBrEzH6Lzk8YEdnK9hJP4OpxW_oU2E,9835
 anemoi/datasets/data/concat.py,sha256=eY5rujcdal00BJCv00mKSlxp0FKVvPQd7uqrBnL9fj4,8996
-anemoi/datasets/data/dataset.py,sha256=Dz74L_RihBzHJyHqlCKcXHBa0J_PkW3YYFofhv-Rh-4,30694
+anemoi/datasets/data/dataset.py,sha256=d-LZPCczG2-ZAua29FGP3_QWzwdnWPHhKpG1dHFQKio,31290
 anemoi/datasets/data/debug.css,sha256=z2X_ZDSnZ9C3pyZPWnQiEyAxuMxUaxJxET4oaCImTAQ,211
 anemoi/datasets/data/debug.py,sha256=hVa1jAQ-TK7CoKJNyyUC0eZPobFG-FpkVXEaO_3B-MA,10796
 anemoi/datasets/data/ensemble.py,sha256=-36kMjuT2y5jUeSnjCRTCyE4um6DLAADBVSKSTkHZZg,5352
@@ -114,11 +118,11 @@ anemoi/datasets/data/fill_missing.py,sha256=ceONpzD-PWLMTtG4WOw6USw-Cd1O55VYzfpA
 anemoi/datasets/data/forwards.py,sha256=d0LL0AORQApGmuPIfMWevOOZuePlgZ6GldHUKmqnTtg,19732
 anemoi/datasets/data/grids.py,sha256=vTAfGq3SaTU4tQzzfeRpFAKlmfL-JVvpjP4e3nGWO3s,22045
 anemoi/datasets/data/indexing.py,sha256=DasVd1j0FB0iTw6eqvhiLka4ztf2zJcI5NgWxmtxzCw,7526
-anemoi/datasets/data/interpolate.py,sha256=D6yBMlY-PTngbMC0Mx9cQn-wL_Oe5M8JNBQ0RaIj1Hc,6482
+anemoi/datasets/data/interpolate.py,sha256=-kSYwdjKH7zJtfITdbqdH6KyOFGVZDyHg4TaFk9shEI,9279
 anemoi/datasets/data/join.py,sha256=ZEHOsCecKBkKKH-vki404Sm7r7cV368ECO7PXPpay3s,9212
-anemoi/datasets/data/masked.py,sha256=xYhvdv4z11QVdZE1d4yGxSJ7VKixQcJy6iviZ_s54bo,7802
+anemoi/datasets/data/masked.py,sha256=giOvHLcGbLf6mZPqZjAxQd1kvydmkepDFh2EqchXLTQ,10213
 anemoi/datasets/data/merge.py,sha256=SvQhJHf-C-Kn7hEjFqomienk-epPPjMtoccRNCJpMtw,8733
-anemoi/datasets/data/misc.py,sha256=WUXqN7IdIdqxM8fIaOcr9FE7UbgvoKHtwsfqPm1cMpU,15253
+anemoi/datasets/data/misc.py,sha256=l8ZJwEjUf9CoPn-RhnilVDVTGypkVIzJ5xV77h8zzkc,20951
 anemoi/datasets/data/missing.py,sha256=ogfVDponbs0bGHMxps32Fj_fq4gT26R70yEMco5gdK8,12593
 anemoi/datasets/data/rescale.py,sha256=nGfJ5tWCncMJ7NMXkLbmt6z0ELrD6FxpbjJreQ3W91g,7004
 anemoi/datasets/data/select.py,sha256=Xs6uOzJL0CoOGeWA_E5_ukr8Jav2kXbZ41vhk7Vr8PE,8277
@@ -130,9 +134,9 @@ anemoi/datasets/data/xy.py,sha256=-jWzYismrK3eI3YCKIBpU1BCmraRncmVn0_2IUY--lk,75
 anemoi/datasets/dates/__init__.py,sha256=pEArHDQ7w5E0WC8Vvf9ypyKSdm6gnhoN9TmooITB7C4,13617
 anemoi/datasets/dates/groups.py,sha256=IOveL6IyTXZwEdXZEnRAnpu9pINY95VN7LzcpLfJ09E,10105
 anemoi/datasets/utils/__init__.py,sha256=hCW0QcLHJmE-C1r38P27_ZOvCLNewex5iQEtZqx2ckI,393
-anemoi_datasets-0.5.19.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
-anemoi_datasets-0.5.19.dist-info/METADATA,sha256=-iNWmeuYT_FuUpsKNhTMb79nXFJhh7sERHF2fW1XJGM,15727
-anemoi_datasets-0.5.19.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-anemoi_datasets-0.5.19.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
-anemoi_datasets-0.5.19.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
-anemoi_datasets-0.5.19.dist-info/RECORD,,
+anemoi_datasets-0.5.20.dist-info/licenses/LICENSE,sha256=8HznKF1Vi2IvfLsKNE5A2iVyiri3pRjRPvPC9kxs6qk,11354
+anemoi_datasets-0.5.20.dist-info/METADATA,sha256=iGcUiyvv_JjdtsXo_oq_WgwweD1nPE0xAM8n5d0qeQU,16039
+anemoi_datasets-0.5.20.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+anemoi_datasets-0.5.20.dist-info/entry_points.txt,sha256=yR-o-4uiPEA_GLBL81SkMYnUoxq3CAV3hHulQiRtGG0,66
+anemoi_datasets-0.5.20.dist-info/top_level.txt,sha256=DYn8VPs-fNwr7fNH9XIBqeXIwiYYd2E2k5-dUFFqUz0,7
+anemoi_datasets-0.5.20.dist-info/RECORD,,