anemoi-datasets 0.5.25__py3-none-any.whl → 0.5.27__py3-none-any.whl

anemoi/datasets/create/sources/xarray_support/metadata.py

@@ -12,8 +12,6 @@ import datetime
 import logging
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import Optional
 
 from anemoi.utils.dates import as_datetime
 from earthkit.data.core.geography import Geography
@@ -23,87 +21,6 @@ from earthkit.data.utils.projections import Projection
 LOG = logging.getLogger(__name__)
 
 
-class _MDMapping:
-    """A class to handle metadata mapping for variables.
-
-    Attributes
-    ----------
-    variable : Any
-        The variable to map.
-    time : Any
-        The time associated with the variable.
-    mapping : Dict[str, str]
-        A dictionary mapping keys to variable names.
-    """
-
-    def __init__(self, variable: Any) -> None:
-        """Initialize the _MDMapping class.
-
-        Parameters
-        ----------
-        variable : Any
-            The variable to map.
-        """
-        self.variable = variable
-        self.time = variable.time
-        self.mapping = dict()
-        # Aliases
-
-    def _from_user(self, key: str) -> str:
-        """Get the internal key corresponding to a user-provided key.
-
-        Parameters
-        ----------
-        key : str
-            The user-provided key.
-
-        Returns
-        -------
-        str
-            The internal key corresponding to the user-provided key.
-        """
-        return self.mapping.get(key, key)
-
-    def from_user(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
-        """Convert user-provided keys to internal keys.
-
-        Parameters
-        ----------
-        kwargs : Dict[str, Any]
-            A dictionary of user-provided keys and values.
-
-        Returns
-        -------
-        Dict[str, Any]
-            A dictionary with internal keys and original values.
-        """
-        return {self._from_user(k): v for k, v in kwargs.items()}
-
-    def __repr__(self) -> str:
-        """Return a string representation of the _MDMapping object.
-
-        Returns
-        -------
-        str
-            String representation of the _MDMapping object.
-        """
-        return f"MDMapping({self.mapping})"
-
-    def fill_time_metadata(self, field: Any, md: Dict[str, Any]) -> None:
-        """Fill the time metadata for a field.
-
-        Parameters
-        ----------
-        field : Any
-            The field to fill metadata for.
-        md : Dict[str, Any]
-            The metadata dictionary to update.
-        """
-        valid_datetime = self.variable.time.fill_time_metadata(field._md, md)
-        if valid_datetime is not None:
-            md["valid_datetime"] = as_datetime(valid_datetime).isoformat()
-
-
 class XArrayMetadata(RawMetadata):
     """A class to handle metadata for XArray fields.
 
@@ -129,10 +46,16 @@ class XArrayMetadata(RawMetadata):
         field : Any
             The field to extract metadata from.
         """
+        from .field import XArrayField
+
+        assert isinstance(field, XArrayField), type(field)
         self._field = field
         md = field._md.copy()
-        self._mapping = _MDMapping(field.owner)
-        self._mapping.fill_time_metadata(field, md)
+
+        valid_datetime = field.owner.time.fill_time_metadata(field._md, md)
+        if valid_datetime is not None:
+            md["valid_datetime"] = as_datetime(valid_datetime).isoformat()
+
         super().__init__(md)
 
     @cached_property
@@ -140,7 +63,7 @@ class XArrayMetadata(RawMetadata):
         """Get the geography information for the field."""
         return XArrayFieldGeography(self._field, self._field.owner.grid)
 
-    def as_namespace(self, namespace: Optional[str] = None) -> Dict[str, Any]:
+    def as_namespace(self, namespace: str | None = None) -> dict[str, Any]:
         """Get the metadata as a specific namespace.
 
         Parameters
@@ -162,7 +85,7 @@ class XArrayMetadata(RawMetadata):
         elif namespace == "mars":
             return self._as_mars()
 
-    def _as_mars(self) -> Dict[str, Any]:
+    def _as_mars(self) -> dict[str, Any]:
         """Get the metadata as MARS namespace.
 
         Returns
@@ -172,7 +95,7 @@ class XArrayMetadata(RawMetadata):
         """
         return {}
 
-    def _base_datetime(self) -> Optional[datetime.datetime]:
+    def _base_datetime(self) -> datetime.datetime | None:
         """Get the base datetime for the field.
 
         Returns
@@ -182,7 +105,7 @@ class XArrayMetadata(RawMetadata):
         """
         return self._field.forecast_reference_time
 
-    def _valid_datetime(self) -> Optional[datetime.datetime]:
+    def _valid_datetime(self) -> datetime.datetime | None:
         """Get the valid datetime for the field.
 
         Returns
@@ -192,38 +115,6 @@ class XArrayMetadata(RawMetadata):
         """
         return self._get("valid_datetime")
 
-    def get(self, key: str, astype: Optional[type] = None, **kwargs: Any) -> Any:
-        """Get a metadata value by key.
-
-        Parameters
-        ----------
-        key : str
-            The key to get the value for.
-        astype : Optional[type]
-            The type to cast the value to.
-        **kwargs : Any
-            Additional keyword arguments.
-
-        Returns
-        -------
-        Any
-            The value for the specified key, optionally cast to the specified type.
-        """
-
-        if key == "levelist":
-            # Special case for levelist, for compatibility with GRIB
-            if key not in self._d and "level" in self._d:
-                key = "level"
-
-        if key in self._d:
-            if astype is not None:
-                return astype(self._d[key])
-            return self._d[key]
-
-        key = self._mapping._from_user(key)
-
-        return super().get(key, astype=astype, **kwargs)
-
 
 class XArrayFieldGeography(Geography):
     """A class to handle geography information for XArray fields.
@@ -280,7 +171,7 @@ class XArrayFieldGeography(Geography):
         """
         raise NotImplementedError()
 
-    def latitudes(self, dtype: Optional[type] = None) -> Any:
+    def latitudes(self, dtype: type | None = None) -> Any:
         """Get the latitudes for the field.
 
         Parameters
@@ -298,7 +189,7 @@ class XArrayFieldGeography(Geography):
             return result.astype(dtype)
         return result
 
-    def longitudes(self, dtype: Optional[type] = None) -> Any:
+    def longitudes(self, dtype: type | None = None) -> Any:
         """Get the longitudes for the field.
 
         Parameters
@@ -316,7 +207,7 @@ class XArrayFieldGeography(Geography):
             return result.astype(dtype)
         return result
 
-    def resolution(self) -> Optional[Any]:
+    def resolution(self) -> Any | None:
         """Get the resolution for the field.
 
         Returns
@@ -327,7 +218,7 @@ class XArrayFieldGeography(Geography):
         # TODO: implement resolution
         return None
 
-    def mars_grid(self) -> Optional[Any]:
+    def mars_grid(self) -> Any | None:
         """Get the MARS grid for the field.
 
         Returns
@@ -338,7 +229,7 @@ class XArrayFieldGeography(Geography):
         # TODO: implement mars_grid
         return None
 
-    def mars_area(self) -> Optional[Any]:
+    def mars_area(self) -> Any | None:
         """Get the MARS area for the field.
 
         Returns
@@ -350,7 +241,7 @@ class XArrayFieldGeography(Geography):
         # return [self.north, self.west, self.south, self.east]
         return None
 
-    def x(self, dtype: Optional[type] = None) -> None:
+    def x(self, dtype: type | None = None) -> None:
         """Get the x-coordinates for the field.
 
         Parameters
@@ -365,7 +256,7 @@ class XArrayFieldGeography(Geography):
         """
         raise NotImplementedError()
 
-    def y(self, dtype: Optional[type] = None) -> None:
+    def y(self, dtype: type | None = None) -> None:
         """Get the y-coordinates for the field.
 
         Parameters

anemoi/datasets/create/sources/xarray_support/patch.py

@@ -10,15 +10,13 @@
 
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
 
 import xarray as xr
 
 LOG = logging.getLogger(__name__)
 
 
-def patch_attributes(ds: xr.Dataset, attributes: Dict[str, Dict[str, Any]]) -> Any:
+def patch_attributes(ds: xr.Dataset, attributes: dict[str, dict[str, Any]]) -> Any:
     """Patch the attributes of the dataset.
 
     Parameters
@@ -40,7 +38,7 @@ def patch_attributes(ds: xr.Dataset, attributes: Dict[str, Dict[str, Any]]) -> A
     return ds
 
 
-def patch_coordinates(ds: xr.Dataset, coordinates: List[str]) -> Any:
+def patch_coordinates(ds: xr.Dataset, coordinates: list[str]) -> Any:
     """Patch the coordinates of the dataset.
 
     Parameters
@@ -61,13 +59,54 @@ 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,
 }
 
 
-def patch_dataset(ds: xr.Dataset, patch: Dict[str, Dict[str, Any]]) -> Any:
+def patch_dataset(ds: xr.Dataset, patch: dict[str, dict[str, Any]]) -> Any:
     """Patch the dataset.
 
     Parameters
@@ -82,7 +121,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/create/sources/xarray_support/time.py

@@ -13,9 +13,6 @@ import logging
 from abc import ABC
 from abc import abstractmethod
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
 
 from anemoi.utils.dates import as_datetime
 
@@ -29,7 +26,7 @@ class Time(ABC):
     """Base class for different time representations."""
 
     @classmethod
-    def from_coordinates(cls, coordinates: List[Coordinate]) -> "Time":
+    def from_coordinates(cls, coordinates: list[Coordinate]) -> "Time":
         """Create a Time instance from a list of coordinates.
 
         Returns
@@ -85,7 +82,7 @@ class Time(ABC):
         raise NotImplementedError(f"{len(date_coordinate)=} {len(time_coordinate)=} {len(step_coordinate)=}")
 
     @abstractmethod
-    def select_valid_datetime(self, variable: Variable) -> Optional[str]:
+    def select_valid_datetime(self, variable: Variable) -> str | None:
         """Select the valid datetime for a given variable.
 
         Parameters
@@ -101,7 +98,7 @@ class Time(ABC):
         pass
 
     @abstractmethod
-    def fill_time_metadata(self, coords_values: Dict[str, Any], metadata: Dict[str, Any]) -> None:
+    def fill_time_metadata(self, coords_values: dict[str, Any], metadata: dict[str, Any]) -> None:
         """Fill metadata with time information.
 
         Args
@@ -118,7 +115,7 @@ class Time(ABC):
 class Constant(Time):
     """Represents a constant time."""
 
-    def fill_time_metadata(self, coords_values: Dict[str, Any], metadata: Dict[str, Any]) -> None:
+    def fill_time_metadata(self, coords_values: dict[str, Any], metadata: dict[str, Any]) -> None:
         """Fill metadata with time information.
 
         Parameters
@@ -154,7 +151,7 @@ class Analysis(Time):
         """
         self.time_coordinate_name = time_coordinate.variable.name
 
-    def fill_time_metadata(self, coords_values: Dict[str, Any], metadata: Dict[str, Any]) -> Any:
+    def fill_time_metadata(self, coords_values: dict[str, Any], metadata: dict[str, Any]) -> Any:
         """Fill metadata with time information.
 
         Parameters
@@ -197,7 +194,7 @@ class ForecastFromValidTimeAndStep(Time):
     """Represents a forecast time derived from valid time and step."""
 
     def __init__(
-        self, time_coordinate: Coordinate, step_coordinate: Coordinate, date_coordinate: Optional[Coordinate] = None
+        self, time_coordinate: Coordinate, step_coordinate: Coordinate, date_coordinate: Coordinate | None = None
     ) -> None:
         """Initialize ForecastFromValidTimeAndStep with time, step, and optional date coordinates.
 
@@ -214,7 +211,7 @@ class ForecastFromValidTimeAndStep(Time):
         self.step_coordinate_name = step_coordinate.variable.name
         self.date_coordinate_name = date_coordinate.variable.name if date_coordinate else None
 
-    def fill_time_metadata(self, coords_values: Dict[str, Any], metadata: Dict[str, Any]) -> Any:
+    def fill_time_metadata(self, coords_values: dict[str, Any], metadata: dict[str, Any]) -> Any:
         """Fill metadata with time information.
 
         Returns
@@ -285,7 +282,7 @@ class ForecastFromValidTimeAndBaseTime(Time):
         self.date_coordinate_name = date_coordinate.name
         self.time_coordinate_name = time_coordinate.name
 
-    def fill_time_metadata(self, coords_values: Dict[str, Any], metadata: Dict[str, Any]) -> Any:
+    def fill_time_metadata(self, coords_values: dict[str, Any], metadata: dict[str, Any]) -> Any:
         """Fill metadata with time information.
 
         Returns
@@ -346,7 +343,7 @@ class ForecastFromBaseTimeAndDate(Time):
         self.date_coordinate_name = date_coordinate.name
         self.step_coordinate_name = step_coordinate.name
 
-    def fill_time_metadata(self, coords_values: Dict[str, Any], metadata: Dict[str, Any]) -> Any:
+    def fill_time_metadata(self, coords_values: dict[str, Any], metadata: dict[str, Any]) -> Any:
         """Fill metadata with time information.
 
         Returns
@@ -375,7 +372,7 @@ class ForecastFromBaseTimeAndDate(Time):
 
         return date + step
 
-    def select_valid_datetime(self, variable: Variable) -> Optional[str]:
+    def select_valid_datetime(self, variable: Variable) -> str | None:
         """Select the valid datetime for a given variable.
 
         Parameters

anemoi/datasets/create/sources/xarray_support/variable.py

@@ -12,10 +12,7 @@ import logging
 import math
 from functools import cached_property
 from typing import Any
-from typing import Dict
-from typing import List
 from typing import Optional
-from typing import Tuple
 
 import numpy as np
 import xarray as xr
@@ -49,10 +46,10 @@ class Variable:
         *,
         ds: xr.Dataset,
         variable: xr.DataArray,
-        coordinates: List[Any],
+        coordinates: list[Any],
         grid: Any,
         time: Any,
-        metadata: Dict[str, Any],
+        metadata: dict[str, Any],
     ):
         """Initialize the Variable object.
 
@@ -82,8 +79,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
@@ -107,7 +108,7 @@ class Variable:
         return self.length
 
     @property
-    def grid_mapping(self) -> Optional[Dict[str, Any]]:
+    def grid_mapping(self) -> dict[str, Any] | None:
         """Return the grid mapping of the variable."""
         grid_mapping = self.variable.attrs.get("grid_mapping", None)
         if grid_mapping is None:
@@ -142,7 +143,7 @@ class Variable:
         str
             A string representation of the variable.
         """
-        return "Variable[name=%s,coordinates=%s,metadata=%s]" % (
+        return "Variable[name={},coordinates={},metadata={}]".format(
             self.variable.name,
             self.coordinates,
             self._metadata,
@@ -173,7 +174,7 @@ class Variable:
         kwargs = {k: v for k, v in zip(self.names, coords)}
         return XArrayField(self, self.variable.isel(kwargs))
 
-    def sel(self, missing: Dict[str, Any], **kwargs: Any) -> Optional["Variable"]:
+    def sel(self, missing: dict[str, Any], **kwargs: Any) -> Optional["Variable"]:
         """Select a subset of the variable based on the given coordinates.
 
         Parameters
@@ -217,23 +218,26 @@ class Variable:
                 LOG.warning(f"Could not find {k}={v} in {c}")
             return None
 
-        coordinates = [x.reduced(i) if c is x else x for x in self.coordinates]
+        if c.scalar and i == 0:
+            variable = self
+        else:
+            coordinates = [x.reduced(i) if c is x else x for x in self.coordinates]
 
-        metadata = self._metadata.copy()
-        metadata.update({k: v})
+            metadata = self._metadata.copy()
+            metadata.update({k: v})
 
-        variable = Variable(
-            ds=self.ds,
-            variable=self.variable.isel({k: i}),
-            coordinates=coordinates,
-            grid=self.grid,
-            time=self.time,
-            metadata=metadata,
-        )
+            variable = Variable(
+                ds=self.ds,
+                variable=self.variable.isel({k: i}),
+                coordinates=coordinates,
+                grid=self.grid,
+                time=self.time,
+                metadata=metadata,
+            )
 
         return variable.sel(missing, **kwargs)
 
-    def match(self, **kwargs: Any) -> Tuple[bool, Optional[Dict[str, Any]]]:
+    def match(self, **kwargs: Any) -> tuple[bool, dict[str, Any] | None]:
         """Match the variable based on the given metadata.
 
         Parameters
@@ -285,7 +289,7 @@ class FilteredVariable:
         self.kwargs = kwargs
 
     @cached_property
-    def fields(self) -> List["XArrayField"]:
+    def fields(self) -> list["XArrayField"]:
         """Filter the fields of a variable based on metadata."""
         return [
             field

anemoi/datasets/create/sources/xarray_zarr.py

@@ -8,7 +8,6 @@
 # nor does it submit to any jurisdiction.
 
 from typing import Any
-from typing import List
 
 import earthkit.data as ekd
 
@@ -17,7 +16,7 @@ from .xarray import load_many
 
 
 @legacy_source(__file__)
-def execute(context: Any, dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+def execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
     """Execute the data loading process.
 
     Parameters

anemoi/datasets/create/sources/zenodo.py

@@ -9,8 +9,6 @@
 
 
 from typing import Any
-from typing import Dict
-from typing import List
 
 import earthkit.data as ekd
 from earthkit.data.core.fieldlist import MultiFieldList
@@ -47,15 +45,15 @@ def execute(context: Any, dates: Any, record_id: str, file_key: str, *args: Any,
     """
     import requests
 
-    result: List[Any] = []
+    result: list[Any] = []
 
     URLPATTERN = "https://zenodo.org/api/records/{record_id}"
     url = URLPATTERN.format(record_id=record_id)
     r = requests.get(url)
     r.raise_for_status()
-    record: Dict[str, Any] = r.json()
+    record: dict[str, Any] = r.json()
 
-    urls: Dict[str, str] = {}
+    urls: dict[str, str] = {}
     for file in record["files"]:
         urls[file["key"]] = file["links"]["self"]
 

anemoi/datasets/create/statistics/__init__.py

@@ -17,9 +17,6 @@ import pickle
 import shutil
 import socket
 from typing import Any
-from typing import List
-from typing import Optional
-from typing import Union
 
 import numpy as np
 import tqdm
@@ -77,7 +74,7 @@ def default_statistics_dates(dates: list[datetime.datetime]) -> tuple[datetime.d
     return dates[0], end
 
 
-def to_datetime(date: Union[str, datetime.datetime]) -> np.datetime64:
+def to_datetime(date: str | datetime.datetime) -> np.datetime64:
     """Convert a date to numpy datetime64 format.
 
     Parameters
@@ -97,7 +94,7 @@ def to_datetime(date: Union[str, datetime.datetime]) -> np.datetime64:
     return date
 
 
-def to_datetimes(dates: list[Union[str, datetime.datetime]]) -> list[np.datetime64]:
+def to_datetimes(dates: list[str | datetime.datetime]) -> list[np.datetime64]:
     """Convert a list of dates to numpy datetime64 format.
 
     Parameters
@@ -221,7 +218,7 @@ def check_variance(
 
 
 def compute_statistics(
-    array: NDArray[Any], check_variables_names: Optional[List[str]] = None, allow_nans: bool = False
+    array: NDArray[Any], check_variables_names: list[str] | None = None, allow_nans: bool = False
 ) -> dict[str, np.ndarray]:
     """Compute statistics for a given array, provides minimum, maximum, sum, squares, count and has_nans as a dictionary.