anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__py3-none-any.whl

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

@@ -11,10 +11,6 @@
 import json
 import logging
 from typing import Any
-from typing import Dict
-from typing import List
-from typing import Optional
-from typing import Union
 
 import xarray as xr
 import yaml
@@ -33,7 +29,7 @@ LOG = logging.getLogger(__name__)
 class XarrayFieldList(FieldList):
     """A class to represent a list of fields from an xarray Dataset."""
 
-    def __init__(self, ds: xr.Dataset, variables: List[Variable]) -> None:
+    def __init__(self, ds: xr.Dataset, variables: list[Variable]) -> None:
         """Initialize the XarrayFieldList.
 
         Parameters
@@ -44,7 +40,7 @@ class XarrayFieldList(FieldList):
             The list of variables.
         """
         self.ds: xr.Dataset = ds
-        self.variables: List[Variable] = variables.copy()
+        self.variables: list[Variable] = variables.copy()
         self.total_length: int = sum(v.length for v in variables)
 
     def __repr__(self) -> str:
@@ -90,8 +86,8 @@ class XarrayFieldList(FieldList):
         cls,
         ds: xr.Dataset,
         *,
-        flavour: Optional[Union[str, Dict[str, Any]]] = None,
-        patch: Optional[Dict[str, Any]] = None,
+        flavour: str | dict[str, Any] | None = None,
+        patch: dict[str, Any] | None = None,
     ) -> "XarrayFieldList":
         """Create an XarrayFieldList from an xarray Dataset.
 
@@ -112,7 +108,7 @@ class XarrayFieldList(FieldList):
         if patch is not None:
             ds = patch_dataset(ds, patch)
 
-        variables: List[Variable] = []
+        variables: list[Variable] = []
 
         if isinstance(flavour, str):
             with open(flavour) as f:
@@ -121,9 +117,9 @@ class XarrayFieldList(FieldList):
                 else:
                     flavour = json.load(f)
 
-        if isinstance(flavour, Dict):
-            flavour_coords: List[str] = [coords["name"] for coords in flavour["rules"].values()]
-            ds_dims: List[str] = [dim for dim in ds._dims]
+        if isinstance(flavour, dict):
+            flavour_coords: list[str] = [coords["name"] for coords in flavour["rules"].values()]
+            ds_dims: list[str] = [dim for dim in ds._dims]
             for dim in ds_dims:
                 if dim in flavour_coords and dim not in ds._coord_names:
                     ds = ds.assign_coords({dim: ds[dim]})
@@ -154,7 +150,7 @@ class XarrayFieldList(FieldList):
                 continue
 
             variable = ds[name]
-            coordinates: List[Any] = []
+            coordinates: list[Any] = []
 
             for coord in variable.coords:
 
@@ -210,7 +206,7 @@ class XarrayFieldList(FieldList):
             So we get an extra chance to filter the fields by the metadata.
         """
 
-        variables: List[Variable] = []
+        variables: list[Variable] = []
         count: int = 0
 
         for v in self.variables:
@@ -223,7 +219,7 @@ class XarrayFieldList(FieldList):
 
             if match:
                 count += 1
-                missing: Dict[str, Any] = {}
+                missing: dict[str, Any] = {}
 
                 # Select from the variable's coordinates (time, level, number, ....)
                 # This may return a new variable with a isel() slice of the selection

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

@@ -11,11 +11,8 @@
 import logging
 from abc import ABC
 from abc import abstractmethod
+from collections.abc import Hashable
 from typing import Any
-from typing import Dict
-from typing import Hashable
-from typing import Optional
-from typing import Tuple
 
 import xarray as xr
 from anemoi.utils.config import DotDict
@@ -61,11 +58,11 @@ class CoordinateGuesser(ABC):
             The dataset to guess coordinates from.
         """
         self.ds = ds
-        self._coordinate_cache: Dict[Hashable, Coordinate] = {}
-        self._grid_cache: Dict[Hashable, Grid] = {}
+        self._coordinate_cache: dict[Hashable, Coordinate] = {}
+        self._grid_cache: dict[Hashable, Grid] = {}
 
     @classmethod
-    def from_flavour(cls, ds: xr.Dataset, flavour: Optional[Dict[str, Any]]) -> "CoordinateGuesser":
+    def from_flavour(cls, ds: xr.Dataset, flavour: dict[str, Any] | None) -> "CoordinateGuesser":
         """Creates a CoordinateGuesser from a flavour.
 
         Parameters
@@ -133,7 +130,7 @@ class CoordinateGuesser(ABC):
             units=units,
         )
 
-        d: Optional[Coordinate] = None
+        d: Coordinate | None = None
 
         d = self._is_point(coordinate, attributes)
         if d is not None:
@@ -214,7 +211,7 @@ class CoordinateGuesser(ABC):
 
         raise NotImplementedError(f"Cannot establish grid {coordinates}")
 
-    def _check_dims(self, variable: Any, x_or_lon: Any, y_or_lat: Any) -> Tuple[Any, bool]:
+    def _check_dims(self, variable: Any, x_or_lon: Any, y_or_lat: Any) -> tuple[Any, bool]:
         """Checks the dimensions of the variable against the coordinates.
 
         Parameters
@@ -362,7 +359,7 @@ class CoordinateGuesser(ABC):
                 grid_mapping = self.ds.attrs["crs"]
                 LOG.warning(f"Using CRS {grid_mapping} from global attributes")
 
-        grid: Optional[Grid] = None
+        grid: Grid | None = None
         if grid_mapping is not None:
 
             grid_mapping = dict(self.ds[grid_mapping].attrs)
@@ -380,7 +377,7 @@ class CoordinateGuesser(ABC):
         raise NotImplementedError(f"Unstructured grid {x.name} {y.name}")
 
     @abstractmethod
-    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
+    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LongitudeCoordinate | None:
         """Checks if the coordinate is a longitude.
 
         Parameters
@@ -398,11 +395,11 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[PointCoordinate]:
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
         pass
 
     @abstractmethod
-    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
+    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LatitudeCoordinate | None:
         """Checks if the coordinate is a latitude.
 
         Parameters
@@ -420,7 +417,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[XCoordinate]:
+    def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> XCoordinate | None:
         """Checks if the coordinate is an x coordinate.
 
         Parameters
@@ -438,7 +435,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[YCoordinate]:
+    def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> YCoordinate | None:
         """Checks if the coordinate is a y coordinate.
 
         Parameters
@@ -456,7 +453,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[TimeCoordinate]:
+    def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> TimeCoordinate | None:
         """Checks if the coordinate is a time coordinate.
 
         Parameters
@@ -474,7 +471,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[DateCoordinate]:
+    def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> DateCoordinate | None:
         """Checks if the coordinate is a date coordinate.
 
         Parameters
@@ -492,7 +489,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[StepCoordinate]:
+    def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> StepCoordinate | None:
         """Checks if the coordinate is a step coordinate.
 
         Parameters
@@ -510,7 +507,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LevelCoordinate]:
+    def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LevelCoordinate | None:
         """Checks if the coordinate is a level coordinate.
 
         Parameters
@@ -528,7 +525,7 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
+    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> EnsembleCoordinate | None:
         """Checks if the coordinate is an ensemble coordinate.
 
         Parameters
@@ -559,7 +556,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         """
         super().__init__(ds)
 
-    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[PointCoordinate]:
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
         if attributes.standard_name in ["cell", "station", "poi", "point"]:
             return PointCoordinate(c)
 
@@ -568,7 +565,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
+    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LongitudeCoordinate | None:
         """Checks if the coordinate is a longitude.
 
         Parameters
@@ -597,7 +594,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
+    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LatitudeCoordinate | None:
         """Checks if the coordinate is a latitude.
 
         Parameters
@@ -626,7 +623,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[XCoordinate]:
+    def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> XCoordinate | None:
         """Checks if the coordinate is an x coordinate.
 
         Parameters
@@ -649,7 +646,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[YCoordinate]:
+    def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> YCoordinate | None:
         """Checks if the coordinate is a y coordinate.
 
         Parameters
@@ -672,7 +669,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[TimeCoordinate]:
+    def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> TimeCoordinate | None:
         """Checks if the coordinate is a time coordinate.
 
         Parameters
@@ -695,7 +692,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[DateCoordinate]:
+    def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> DateCoordinate | None:
         """Checks if the coordinate is a date coordinate.
 
         Parameters
@@ -718,7 +715,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[StepCoordinate]:
+    def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> StepCoordinate | None:
         """Checks if the coordinate is a step coordinate.
 
         Parameters
@@ -744,7 +741,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LevelCoordinate]:
+    def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LevelCoordinate | None:
         """Checks if the coordinate is a level coordinate.
 
         Parameters
@@ -762,6 +759,9 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         if attributes.standard_name == "atmosphere_hybrid_sigma_pressure_coordinate":
             return LevelCoordinate(c, "ml")
 
+        if attributes.standard_name == "model_level_number":
+            return LevelCoordinate(c, "ml")
+
         if attributes.long_name == "height" and attributes.units == "m":
             return LevelCoordinate(c, "height")
 
@@ -782,7 +782,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
+    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> EnsembleCoordinate | None:
         """Checks if the coordinate is an ensemble coordinate.
 
         Parameters
@@ -806,7 +806,7 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
 class FlavourCoordinateGuesser(CoordinateGuesser):
     """Implementation of CoordinateGuesser that uses a flavour for guessing."""
 
-    def __init__(self, ds: xr.Dataset, flavour: Dict[str, Any]) -> None:
+    def __init__(self, ds: xr.Dataset, flavour: dict[str, Any]) -> None:
         """Initializes the FlavourCoordinateGuesser.
 
         Parameters
@@ -819,7 +819,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
         super().__init__(ds)
         self.flavour = flavour
 
-    def _match(self, c: xr.DataArray, key: str, attributes: CoordinateAttributes) -> Optional[Dict[str, Any]]:
+    def _match(self, c: xr.DataArray, key: str, attributes: CoordinateAttributes) -> dict[str, Any] | None:
         """Matches the coordinate against the flavour rules.
 
         Parameters
@@ -854,7 +854,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
+    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LongitudeCoordinate | None:
         """Checks if the coordinate is a longitude using the flavour rules.
 
         Parameters
@@ -874,7 +874,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
+    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LatitudeCoordinate | None:
         """Checks if the coordinate is a latitude using the flavour rules.
 
         Parameters
@@ -894,7 +894,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[XCoordinate]:
+    def _is_x(self, c: xr.DataArray, attributes: CoordinateAttributes) -> XCoordinate | None:
         """Checks if the coordinate is an x coordinate using the flavour rules.
 
         Parameters
@@ -914,7 +914,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[YCoordinate]:
+    def _is_y(self, c: xr.DataArray, attributes: CoordinateAttributes) -> YCoordinate | None:
         """Checks if the coordinate is a y coordinate using the flavour rules.
 
         Parameters
@@ -934,7 +934,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[TimeCoordinate]:
+    def _is_time(self, c: xr.DataArray, attributes: CoordinateAttributes) -> TimeCoordinate | None:
         """Checks if the coordinate is a time coordinate using the flavour rules.
 
         Parameters
@@ -954,7 +954,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[StepCoordinate]:
+    def _is_step(self, c: xr.DataArray, attributes: CoordinateAttributes) -> StepCoordinate | None:
         """Checks if the coordinate is a step coordinate using the flavour rules.
 
         Parameters
@@ -974,7 +974,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[DateCoordinate]:
+    def _is_date(self, c: xr.DataArray, attributes: CoordinateAttributes) -> DateCoordinate | None:
         """Checks if the coordinate is a date coordinate using the flavour rules.
 
         Parameters
@@ -994,7 +994,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LevelCoordinate]:
+    def _is_level(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LevelCoordinate | None:
         """Checks if the coordinate is a level coordinate using the flavour rules.
 
         Parameters
@@ -1039,7 +1039,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         raise NotImplementedError(f"levtype for {c=}")
 
-    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[EnsembleCoordinate]:
+    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> EnsembleCoordinate | None:
         """Checks if the coordinate is an ensemble coordinate using the flavour rules.
 
         Parameters
@@ -1059,7 +1059,7 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
 
         return None
 
-    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[PointCoordinate]:
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
         """Checks if the coordinate is a point coordinate using the flavour rules.
 
         Parameters

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

@@ -13,7 +13,6 @@ from abc import ABC
 from abc import abstractmethod
 from functools import cached_property
 from typing import Any
-from typing import Tuple
 
 import numpy as np
 
@@ -38,7 +37,7 @@ class Grid(ABC):
 
     @property
     @abstractmethod
-    def grid_points(self) -> Tuple[Any, Any]:
+    def grid_points(self) -> tuple[Any, Any]:
         """Get the grid points."""
         pass
 
@@ -85,7 +84,7 @@ class MeshedGrid(LatLonGrid):
     """Grid class for meshed latitude and longitude coordinates."""
 
     @cached_property
-    def grid_points(self) -> Tuple[Any, Any]:
+    def grid_points(self) -> tuple[Any, Any]:
         """Get the grid points for the meshed grid."""
 
         if self.variable_dims == (self.lon.variable.name, self.lat.variable.name):
@@ -128,7 +127,7 @@ class UnstructuredGrid(LatLonGrid):
         assert set(self.variable_dims) == set(self.grid_dims), (self.variable_dims, self.grid_dims)
 
     @cached_property
-    def grid_points(self) -> Tuple[Any, Any]:
+    def grid_points(self) -> tuple[Any, Any]:
         """Get the grid points for the unstructured grid."""
         assert 1 <= len(self.variable_dims) <= 2
 
@@ -191,7 +190,7 @@ class MeshProjectionGrid(ProjectionGrid):
     """Grid class for meshed projected coordinates."""
 
     @cached_property
-    def grid_points(self) -> Tuple[Any, Any]:
+    def grid_points(self) -> tuple[Any, Any]:
         """Get the grid points for the mesh projection grid."""
         transformer = self.transformer()
         xv, yv = np.meshgrid(self.x.variable.values, self.y.variable.values)  # , indexing="ij")
@@ -199,10 +198,17 @@ class MeshProjectionGrid(ProjectionGrid):
         return lat.flatten(), lon.flatten()
 
 
-class UnstructuredProjectionGrid(XYGrid):
+class UnstructuredProjectionGrid(ProjectionGrid):
     """Grid class for unstructured projected coordinates."""
 
     @cached_property
-    def grid_points(self) -> Tuple[Any, Any]:
-        """Get the grid points for the unstructured projection grid."""
-        raise NotImplementedError("UnstructuredProjectionGrid")
+    def grid_points(self) -> tuple[Any, Any]:
+        """Get the grid points for the unstructured grid."""
+        if self.projection == "epsg:4326":
+            # WGS84, no transformation needed
+            return self.y.variable.values.flatten(), self.x.variable.values.flatten()
+        transformer = self.transformer()
+        xv, yv = np.meshgrid(self.x.variable.values, self.y.variable.values)  # , indexing="ij")
+        lon, lat = transformer.transform(xv, yv)
+
+        return lat.flatten(), lon.flatten()