anemoi-datasets 0.5.25__tar.gz → 0.5.26__tar.gz

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/.github/pull_request_template.md

@@ -11,3 +11,5 @@
 <!-- Include any additional information, caveats, or considerations that the reviewer should be aware of. -->
 
 ***As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting  (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/***
+
+By opening this pull request, I affirm that all authors agree to the [Contributor License Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

anemoi_datasets-0.5.26/.release-please-manifest.json

@@ -0,0 +1,3 @@
+{
+    ".": "0.5.26"
+}

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/CHANGELOG.md

@@ -8,6 +8,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Please add your functional changes to the appropriate section in the PR.
 Keep it human-readable, your future self will thank you!
 
+## [0.5.26](https://github.com/ecmwf/anemoi-datasets/compare/0.5.25...0.5.26) (2025-07-11)
+
+
+### Features
+
+* Add planetary planetary source ([#368](https://github.com/ecmwf/anemoi-datasets/issues/368)) ([c85337f](https://github.com/ecmwf/anemoi-datasets/commit/c85337f4b08381dc1e2e8ba71134c998c6258c2d))
+* Extend xarray coordinate guesser to include pressure as a long_name ([#335](https://github.com/ecmwf/anemoi-datasets/issues/335)) ([bcbffaa](https://github.com/ecmwf/anemoi-datasets/commit/bcbffaa0b0ad011db5b1fd5879f9829f0d2e68af))
+* Extented complement nearest functionality to support k>1 ([#371](https://github.com/ecmwf/anemoi-datasets/issues/371)) ([a2c5f92](https://github.com/ecmwf/anemoi-datasets/commit/a2c5f921f57fce60bd3fe728c21dc5e5b7d3a4a1))
+* Point data ([#358](https://github.com/ecmwf/anemoi-datasets/issues/358)) ([78c4f85](https://github.com/ecmwf/anemoi-datasets/commit/78c4f85bf9cd4d6cc7f17268c8c526da32158220))
+* Relaxed compatibility check for horizontal grid. ([#339](https://github.com/ecmwf/anemoi-datasets/issues/339)) ([a7dc859](https://github.com/ecmwf/anemoi-datasets/commit/a7dc859b6ec2fb0bad2ed7b7b8ee3c76034a258b))
+* Sort xarray coordinate with patch ([#356](https://github.com/ecmwf/anemoi-datasets/issues/356)) ([638cbdb](https://github.com/ecmwf/anemoi-datasets/commit/638cbdb013b94b87748654ae338e0d02e69c1900))
+* Use pytest-xdist to run tests in parallel. ([#346](https://github.com/ecmwf/anemoi-datasets/issues/346)) ([1d656ee](https://github.com/ecmwf/anemoi-datasets/commit/1d656ee62c51ee0cbe043943581cef16b2f7d7db))
+
+
+### Bug Fixes
+
+* Deprecated filters in anemoi datasets ([#350](https://github.com/ecmwf/anemoi-datasets/issues/350)) ([d08c6f5](https://github.com/ecmwf/anemoi-datasets/commit/d08c6f55c3d983c930480bd777e74fd5a5bea6eb))
+* **grib-index:** Match against filename, not full path ([#374](https://github.com/ecmwf/anemoi-datasets/issues/374)) ([e71a952](https://github.com/ecmwf/anemoi-datasets/commit/e71a95221493a703515d7a14a17f3dead06c8430))
+* Missing method in xarray flavour ([#369](https://github.com/ecmwf/anemoi-datasets/issues/369)) ([d0b9588](https://github.com/ecmwf/anemoi-datasets/commit/d0b9588f0bff626ae751a74cf4cbbc373a2f67f4))
+* Remove trailing slash from zarr paths.  ([#362](https://github.com/ecmwf/anemoi-datasets/issues/362)) ([9a9bfe5](https://github.com/ecmwf/anemoi-datasets/commit/9a9bfe50406d299f44f7cce7e3befcabef6fefaf))
+
+
+### Documentation
+
+* Fix pre-commit ([#367](https://github.com/ecmwf/anemoi-datasets/issues/367)) ([6e6ebe2](https://github.com/ecmwf/anemoi-datasets/commit/6e6ebe2e81fc735a2af422b8c222b0a2f229d8dc))
+* Zarr dataset creation ([#365](https://github.com/ecmwf/anemoi-datasets/issues/365)) ([622c261](https://github.com/ecmwf/anemoi-datasets/commit/622c26156de44ae4eaa3359fef836741b1d81d67))
+
 ## [0.5.25](https://github.com/ecmwf/anemoi-datasets/compare/0.5.24...0.5.25) (2025-06-11)
 
 

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: anemoi-datasets
-Version: 0.5.25
+Version: 0.5.26
 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
@@ -226,7 +226,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: anemoi-transform>=0.1.10
-Requires-Dist: anemoi-utils[provenance]>=0.4.21
+Requires-Dist: anemoi-utils[provenance]>=0.4.26
 Requires-Dist: cfunits
 Requires-Dist: numcodecs<0.16
 Requires-Dist: numpy
@@ -262,6 +262,7 @@ Requires-Dist: requests; extra == "remote"
 Provides-Extra: tests
 Requires-Dist: anemoi-datasets[xarray]; extra == "tests"
 Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-xdist; extra == "tests"
 Provides-Extra: xarray
 Requires-Dist: adlfs; extra == "xarray"
 Requires-Dist: gcsfs; extra == "xarray"

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/docs/datasets/building/sources/anemoi-dataset.rst

@@ -17,3 +17,9 @@ An anemoi-dataset can be a source for a dataset:
 The parameters are the same as those used in the ``open_dataset``
 function, which allows you to subset and combine datasets. See
 :ref:`opening-datasets` for more information.
+
+In particular, this is how local zarr datasets created with anemoi in a
+can be used as a source, contrary to :ref:`xarray-zarr` :
+
+.. literalinclude:: yaml/anemoi-zarr-dataset.yaml
+   :language: yaml

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/docs/datasets/building/sources/xarray-zarr.rst

@@ -1,3 +1,5 @@
+.. _xarray-zarr:
+
 #############
  xarray-zarr
 #############
@@ -17,4 +19,8 @@ it is necessary to use the :ref:`join <building-join>` operation to join
 separate lists containing 2D variables and 3D variables. If all vertical
 levels are desired, then it is acceptable to specify a single source.
 
+Also, an ``xarray-zarr`` source uses the ``url`` keyword, and cannot be
+used for accessing local datasets. For using local zarr datasets as
+sources, use instead :ref:`anemoi-dataset_source`.
+
 See :ref:`create-cf-data` for more information.

anemoi_datasets-0.5.26/docs/datasets/building/sources/yaml/anemoi-zarr-dataset.yaml

@@ -0,0 +1,3 @@
+input:
+  anemoi-dataset:
+    dataset: path/to/dataset.zarr

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/docs/datasets/using/combining.rst

@@ -235,13 +235,15 @@ variables of `dataset1` and return the result.
        source=dataset2,
        what="variables",
        interpolate="nearest",
+       k=1,
    )
 
 Currently ``what`` can only be ``variables`` and can be omitted.
 
 The value for ``interpolate`` can be one of ``none`` (default) or
 ``nearest``. In the case of ``none``, the grids of the two datasets must
-match.
+match. In case of ``interpolate``, an additional parameter ``k`` can be
+set to specify the number of nearest neighbors to use.
 
 This feature was originally designed to be used in conjunction with
 ``cutout``, where `dataset1` is the lam, and `dataset2` is the global

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/docs/howtos/create/02-cf-data.rst

@@ -46,9 +46,18 @@ can contain patterns. See :ref:`file-pattern` for more information.
  Zarr
 ******
 
+For using remote hosted zarr datasets as sources, use
+:ref:`xarray-zarr`.
+
 .. literalinclude:: yaml/zarr1.yaml
    :language: yaml
 
+For using local zarr datasets (such as anemoi-generated datasets), use
+:ref:`anemoi-dataset_source`.
+
+.. literalinclude:: yaml/zarr2.yaml
+   :language: yaml
+
 *********************************************
  Handling data that is not 100% CF-compliant
 *********************************************

anemoi_datasets-0.5.26/docs/howtos/create/yaml/zarr2.yaml

@@ -0,0 +1,8 @@
+dates:
+  start: 2023-01-01T00:00:00
+  end: 2023-01-02T18:00:00
+  frequency: 6h
+
+input:
+  anemoi-dataset:
+    dataset: /path/to/input.zarr

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/docs/howtos/usage/code/cutout-complement1.py

@@ -14,4 +14,5 @@ ds = open_dataset(
     },
     source="global-dataset",
     interpolation="nearest",
+    k=1,
 )

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/docs/howtos/usage/yaml/cutout-complement1.yaml

@@ -9,3 +9,4 @@ dataset:
          adjust: dates
    source: global-dataset
    interpolation: nearest
+   k: 1

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/pyproject.toml

@@ -50,7 +50,7 @@ dynamic = [
 ]
 dependencies = [
   "anemoi-transform>=0.1.10",
-  "anemoi-utils[provenance]>=0.4.21",
+  "anemoi-utils[provenance]>=0.4.26",
   "cfunits",
   "numcodecs<0.16",                   # Until we move to zarr3
   "numpy",
@@ -100,6 +100,7 @@ optional-dependencies.remote = [
 optional-dependencies.tests = [
   "anemoi-datasets[xarray]",
   "pytest",
+  "pytest-xdist",
 ]
 
 optional-dependencies.xarray = [
@@ -131,6 +132,13 @@ version_file = "src/anemoi/datasets/_version.py"
 [tool.isort]
 profile = "black"
 
+[tool.pytest.ini_options]
+testpaths = "tests"
+addopts = [
+  "--numprocesses=auto",
+  "--strict-config",
+]
+
 [tool.mypy]
 strict = false
 exclude = [

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.25'
-__version_tuple__ = version_tuple = (0, 5, 25)
+__version__ = version = '0.5.26'
+__version_tuple__ = version_tuple = (0, 5, 26)

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/commands/grib-index.py

@@ -81,7 +81,7 @@ class GribIndexCmd(Command):
             bool
                 True if the path matches, False otherwise.
             """
-            return fnmatch.fnmatch(path, args.match)
+            return fnmatch.fnmatch(os.path.basename(path), args.match)
 
         from anemoi.datasets.create.sources.grib_index import GribIndex
 

anemoi_datasets-0.5.25/src/anemoi/datasets/create/filters/transform.py → anemoi_datasets-0.5.26/src/anemoi/datasets/create/filter.py

@@ -1,4 +1,4 @@
-# (C) Copyright 2025  Anemoi contributors.
+# (C) Copyright 2025- Anemoi contributors.
 #
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
@@ -12,10 +12,8 @@ from typing import Dict
 
 import earthkit.data as ekd
 
-from ..filter import Filter
 
-
-class TransformFilter(Filter):
+class TransformFilter:
     """Calls filters from anemoi.transform.filters
 
     Parameters
@@ -29,7 +27,6 @@ class TransformFilter(Filter):
     """
 
     def __init__(self, context: Any, name: str, config: Dict[str, Any]) -> None:
-
         from anemoi.transform.filters import create_filter
 
         self.name = name

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/input/step.py

@@ -8,7 +8,6 @@
 # nor does it submit to any jurisdiction.
 
 import logging
-import warnings
 from copy import deepcopy
 from typing import Any
 from typing import Dict
@@ -165,24 +164,11 @@ def step_factory(config: Dict[str, Any], context: ActionContext, action_path: Li
     if cls is not None:
         return cls(context, action_path, previous_step, *args, **kwargs)
 
-    # Try filters from datasets filter registry
+    # Try filters from transform filter registry
     from anemoi.transform.filters import filter_registry as transform_filter_registry
 
-    from ..filters import create_filter as create_datasets_filter
-    from ..filters import filter_registry as datasets_filter_registry
-
-    if datasets_filter_registry.is_registered(key):
-
-        if transform_filter_registry.is_registered(key):
-            warnings.warn(f"Filter `{key}` is registered in both datasets and transform filter registries")
-
-        filter = create_datasets_filter(None, config)
-        return FunctionStepAction(context, action_path + [key], previous_step, key, filter)
-
-    # Use filters from transform registry
-
     if transform_filter_registry.is_registered(key):
-        from ..filters.transform import TransformFilter
+        from ..filter import TransformFilter
 
         return FunctionStepAction(
             context, action_path + [key], previous_step, key, TransformFilter(context, key, config)

anemoi_datasets-0.5.26/src/anemoi/datasets/create/sources/planetary_computer.py

@@ -0,0 +1,44 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+from . import source_registry
+from .xarray import XarraySourceBase
+
+
+@source_registry.register("planetary_computer")
+class PlanetaryComputerSource(XarraySourceBase):
+    """An Xarray data source for the planetary_computer."""
+
+    emoji = "🪐"
+
+    def __init__(self, context, data_catalog_id, version="v1", *args, **kwargs: dict):
+
+        import planetary_computer
+        import pystac_client
+
+        self.data_catalog_id = data_catalog_id
+        self.flavour = kwargs.pop("flavour", None)
+        self.patch = kwargs.pop("patch", None)
+        self.options = kwargs.pop("options", {})
+
+        catalog = pystac_client.Client.open(
+            f"https://planetarycomputer.microsoft.com/api/stac/{version}/",
+            modifier=planetary_computer.sign_inplace,
+        )
+        collection = catalog.get_collection(self.data_catalog_id)
+
+        asset = collection.assets["zarr-abfs"]
+
+        if "xarray:storage_options" in asset.extra_fields:
+            self.options["storage_options"] = asset.extra_fields["xarray:storage_options"]
+
+        self.options.update(asset.extra_fields["xarray:open_kwargs"])
+
+        super().__init__(context, url=asset.href, *args, **kwargs)

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/sources/xarray_support/__init__.py

@@ -20,7 +20,6 @@ import xarray as xr
 from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.sources.patterns import iterate_patterns
-from anemoi.datasets.data.stores import name_to_zarr_store
 
 from ..legacy import legacy_source
 from .fieldlist import XarrayFieldList
@@ -89,37 +88,22 @@ def load_one(
         The loaded dataset.
     """
 
-    """
-    We manage the S3 client ourselves, bypassing fsspec and s3fs layers, because sometimes something on the stack
-    zarr/fsspec/s3fs/boto3 (?) seem to flags files as missing when they actually are not (maybe when S3 reports some sort of
-    connection error). In that case,  Zarr will silently fill the chunks that could not be downloaded with NaNs.
-    See https://github.com/pydata/xarray/issues/8842
-
-    We have seen this bug triggered when we run many clients in parallel, for example, when we create a new dataset using `xarray-zarr`.
-    """
-
     if options is None:
         options = {}
 
     context.trace(emoji, dataset, options, kwargs)
 
-    if isinstance(dataset, str) and ".zarr" in dataset:
-        data = xr.open_zarr(name_to_zarr_store(dataset), **options)
-    elif "planetarycomputer" in dataset:
-        store = name_to_zarr_store(dataset)
-        if "store" in store:
-            data = xr.open_zarr(**store)
-        if "filename_or_obj" in store:
-            data = xr.open_dataset(**store)
-    else:
-        data = xr.open_dataset(dataset, **options)
+    if isinstance(dataset, str) and dataset.endswith(".zarr"):
+        # If the dataset is a zarr store, we need to use the zarr engine
+        options["engine"] = "zarr"
+
+    data = xr.open_dataset(dataset, **options)
 
     fs = XarrayFieldList.from_xarray(data, flavour=flavour, patch=patch)
 
     if len(dates) == 0:
         result = fs.sel(**kwargs)
     else:
-        print("dates", dates, kwargs)
         result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 
     if len(result) == 0:
@@ -130,7 +114,7 @@ def load_one(
             a = ["valid_datetime", k.metadata("valid_datetime", default=None)]
             for n in kwargs.keys():
                 a.extend([n, k.metadata(n, default=None)])
-            print([str(x) for x in a])
+            LOG.warning(f"{[str(x) for x in a]}")
 
             if i > 16:
                 break

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/sources/xarray_support/coordinates.py

@@ -95,6 +95,7 @@ class Coordinate:
     is_member = False
     is_x = False
     is_y = False
+    is_point = False
 
     def __init__(self, variable: xr.DataArray) -> None:
         """Initialize the coordinate.
@@ -390,6 +391,13 @@ class EnsembleCoordinate(Coordinate):
         return value
 
 
+class PointCoordinate(Coordinate):
+    """Coordinate class for point data."""
+
+    is_point = True
+    mars_names = ("point",)
+
+
 class LongitudeCoordinate(Coordinate):
     """Coordinate class for longitude."""
 

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/sources/xarray_support/field.py

@@ -87,13 +87,10 @@ class XArrayField(Field):
                 coordinate = owner.by_name[coord_name]
                 self._md[coord_name] = coordinate.normalise(extract_single_value(coord_value))
 
-        # print(values.ndim, values.shape, selection.dims)
         # By now, the only dimensions should be latitude and longitude
         self._shape = tuple(list(self.selection.shape)[-2:])
         if math.prod(self._shape) != math.prod(self.selection.shape):
-            print(self.selection.ndim, self.selection.shape)
-            print(self.selection)
-            raise ValueError("Invalid shape for selection")
+            raise ValueError(f"Invalid shape for selection {self._shape=}, {self.selection.shape=} {self.selection=}")
 
     @property
     def shape(self) -> Tuple[int, int]:

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/sources/xarray_support/flavour.py

@@ -26,6 +26,7 @@ from .coordinates import EnsembleCoordinate
 from .coordinates import LatitudeCoordinate
 from .coordinates import LevelCoordinate
 from .coordinates import LongitudeCoordinate
+from .coordinates import PointCoordinate
 from .coordinates import ScalarCoordinate
 from .coordinates import StepCoordinate
 from .coordinates import TimeCoordinate
@@ -134,6 +135,10 @@ class CoordinateGuesser(ABC):
 
         d: Optional[Coordinate] = None
 
+        d = self._is_point(coordinate, attributes)
+        if d is not None:
+            return d
+
         d = self._is_longitude(coordinate, attributes)
         if d is not None:
             return d
@@ -308,9 +313,9 @@ class CoordinateGuesser(ABC):
             return self._grid_cache[(x.name, y.name, dim_vars)]
 
         grid_mapping = variable.attrs.get("grid_mapping", None)
-        if grid_mapping is not None:
-            print(f"grid_mapping: {grid_mapping}")
-            print(self.ds[grid_mapping])
+        # if grid_mapping is not None:
+        #     print(f"grid_mapping: {grid_mapping}")
+        #     print(self.ds[grid_mapping])
 
         if grid_mapping is None:
             LOG.warning(f"No 'grid_mapping' attribute provided for '{variable.name}'")
@@ -392,6 +397,10 @@ class CoordinateGuesser(ABC):
         """
         pass
 
+    @abstractmethod
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[PointCoordinate]:
+        pass
+
     @abstractmethod
     def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
         """Checks if the coordinate is a latitude.
@@ -550,6 +559,15 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         """
         super().__init__(ds)
 
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[PointCoordinate]:
+        if attributes.standard_name in ["cell", "station", "poi", "point"]:
+            return PointCoordinate(c)
+
+        if attributes.name in ["cell", "station", "poi", "point"]:  # WeatherBench
+            return PointCoordinate(c)
+
+        return None
+
     def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
         """Checks if the coordinate is a longitude.
 
@@ -750,6 +768,9 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         if attributes.standard_name == "air_pressure" and attributes.units == "hPa":
             return LevelCoordinate(c, "pl")
 
+        if attributes.long_name == "pressure" and attributes.units in ["hPa", "Pa"]:
+            return LevelCoordinate(c, "pl")
+
         if attributes.name == "level":
             return LevelCoordinate(c, "pl")
 
@@ -759,9 +780,6 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         if attributes.standard_name == "depth":
             return LevelCoordinate(c, "depth")
 
-        if attributes.name == "vertical" and attributes.units == "hPa":
-            return LevelCoordinate(c, "pl")
-
         return None
 
     def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
@@ -1040,3 +1058,23 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
             return EnsembleCoordinate(c)
 
         return None
+
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[PointCoordinate]:
+        """Checks if the coordinate is a point coordinate using the flavour rules.
+
+        Parameters
+        ----------
+        c : xr.DataArray
+            The coordinate to check.
+        attributes : CoordinateAttributes
+            The attributes of the coordinate.
+
+        Returns
+        -------
+        Optional[PointCoordinate]
+            The StepCoorPointCoordinateinate if matched, else None.
+        """
+        if self._match(c, "point", attributes):
+            return PointCoordinate(c)
+
+        return None

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/sources/xarray_support/patch.py

@@ -61,9 +61,50 @@ def patch_coordinates(ds: xr.Dataset, coordinates: List[str]) -> Any:
     return ds
 
 
+def patch_rename(ds: xr.Dataset, renames: dict[str, str]) -> Any:
+    """Rename variables in the dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    renames : dict[str, str]
+        Mapping from old variable names to new variable names.
+
+    Returns
+    -------
+    Any
+        The patched dataset.
+    """
+    return ds.rename(renames)
+
+
+def patch_sort_coordinate(ds: xr.Dataset, sort_coordinates: List[str]) -> Any:
+    """Sort the coordinates of the dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    sort_coordinates : List[str]
+        The coordinates to sort.
+
+    Returns
+    -------
+    Any
+        The patched dataset.
+    """
+
+    for name in sort_coordinates:
+        ds = ds.sortby(name)
+    return ds
+
+
 PATCHES = {
     "attributes": patch_attributes,
     "coordinates": patch_coordinates,
+    "rename": patch_rename,
+    "sort_coordinates": patch_sort_coordinate,
 }
 
 
@@ -82,7 +123,9 @@ def patch_dataset(ds: xr.Dataset, patch: Dict[str, Dict[str, Any]]) -> Any:
     Any
         The patched dataset.
     """
-    for what, values in patch.items():
+
+    ORDER = ["coordinates", "attributes", "rename", "sort_coordinates"]
+    for what, values in sorted(patch.items(), key=lambda x: ORDER.index(x[0])):
         if what not in PATCHES:
             raise ValueError(f"Unknown patch type {what!r}")
 

{anemoi_datasets-0.5.25 → anemoi_datasets-0.5.26}/src/anemoi/datasets/create/sources/xarray_support/variable.py

@@ -82,8 +82,12 @@ class Variable:
 
         self.time = time
 
-        self.shape = tuple(len(c.variable) for c in coordinates if c.is_dim and not c.scalar and not c.is_grid)
-        self.names = {c.variable.name: c for c in coordinates if c.is_dim and not c.scalar and not c.is_grid}
+        self.shape = tuple(
+            len(c.variable) for c in coordinates if c.is_dim and not c.scalar and not c.is_grid and not c.is_point
+        )
+        self.names = {
+            c.variable.name: c for c in coordinates if c.is_dim and not c.scalar and not c.is_grid and not c.is_point
+        }
         self.by_name = {c.variable.name: c for c in coordinates}
 
         # We need that alias for the time dimension