anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__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
@@ -79,7 +77,7 @@ def patch_rename(ds: xr.Dataset, renames: dict[str, str]) -> Any:
     return ds.rename(renames)
 
 
-def patch_sort_coordinate(ds: xr.Dataset, sort_coordinates: List[str]) -> Any:
+def patch_sort_coordinate(ds: xr.Dataset, sort_coordinates: list[str]) -> Any:
     """Sort the coordinates of the dataset.
 
     Parameters
@@ -108,7 +106,7 @@ PATCHES = {
 }
 
 
-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

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.
 
@@ -111,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:
@@ -146,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,
@@ -177,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
@@ -221,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
@@ -289,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,34 +8,37 @@
 # nor does it submit to any jurisdiction.
 
 from typing import Any
-from typing import List
 
 import earthkit.data as ekd
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 from .xarray import load_many
 
 
-@legacy_source(__file__)
-def execute(context: Any, dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
-    """Execute the data loading process.
-
-    Parameters
-    ----------
-    context : Any
-        The context in which the execution occurs.
-    dates : List[str]
-        List of dates for which data is to be loaded.
-    url : str
-        The URL from which data is to be loaded.
-    *args : tuple
-        Additional positional arguments.
-    **kwargs : dict
-        Additional keyword arguments.
-
-    Returns
-    -------
-    ekd.FieldList
-        The loaded data.
-    """
-    return load_many("🇿", context, dates, url, *args, **kwargs)
+@source_registry.register("xarray_zarr")
+class XarrayZarrSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+        """Execute the data loading process.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the execution occurs.
+        dates : List[str]
+            List of dates for which data is to be loaded.
+        url : str
+            The URL from which data is to be loaded.
+        *args : tuple
+            Additional positional arguments.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns
+        -------
+        ekd.FieldList
+            The loaded data.
+        """
+        return load_many("🇿", context, dates, url, *args, **kwargs)

anemoi/datasets/create/sources/zenodo.py

@@ -9,61 +9,63 @@
 
 
 from typing import Any
-from typing import Dict
-from typing import List
 
 import earthkit.data as ekd
 from earthkit.data.core.fieldlist import MultiFieldList
 from earthkit.data.sources.url import download_and_cache
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 from .patterns import iterate_patterns
 from .xarray import load_one
 
 
-@legacy_source(__file__)
-def execute(context: Any, dates: Any, record_id: str, file_key: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
-    """Executes the download and processing of files from Zenodo.
+@source_registry.register("zenodo")
+class ZenodoSource(LegacySource):
 
-    Parameters
-    ----------
-    context : Any
-        The context in which the function is executed.
-    dates : Any
-        The dates for which the data is required.
-    record_id : str
-        The Zenodo record ID.
-    file_key : str
-        The key to identify the file.
-    *args : Any
-        Additional arguments.
-    **kwargs : Any
-        Additional keyword arguments.
+    @staticmethod
+    def _execute(context: Any, dates: Any, record_id: str, file_key: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+        """Executes the download and processing of files from Zenodo.
 
-    Returns
-    -------
-    MultiFieldList
-        A list of fields loaded from the downloaded files.
-    """
-    import requests
+        Parameters
+        ----------
+        context : Any
+            The context in which the function is executed.
+        dates : Any
+            The dates for which the data is required.
+        record_id : str
+            The Zenodo record ID.
+        file_key : str
+            The key to identify the file.
+        *args : Any
+            Additional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
 
-    result: List[Any] = []
+        Returns
+        -------
+        MultiFieldList
+            A list of fields loaded from the downloaded files.
+        """
+        import requests
 
-    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()
+        result: list[Any] = []
 
-    urls: Dict[str, str] = {}
-    for file in record["files"]:
-        urls[file["key"]] = file["links"]["self"]
+        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()
 
-    for url, dates in iterate_patterns(file_key, dates, **kwargs):
-        if url not in urls:
-            continue
+        urls: dict[str, str] = {}
+        for file in record["files"]:
+            urls[file["key"]] = file["links"]["self"]
 
-        path = download_and_cache(urls[url])
-        result.append(load_one("?", context, dates, path, options={}, flavour=None, **kwargs))
+        for url, dates in iterate_patterns(file_key, dates, **kwargs):
+            if url not in urls:
+                continue
 
-    return MultiFieldList(result)
+            path = download_and_cache(urls[url])
+            result.append(load_one("?", context, dates, path, options={}, flavour=None, **kwargs))
+
+        return MultiFieldList(result)