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

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
@@ -26,6 +23,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
@@ -60,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
@@ -132,7 +130,11 @@ class CoordinateGuesser(ABC):
             units=units,
         )
 
-        d: Optional[Coordinate] = None
+        d: Coordinate | None = 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:
@@ -209,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
@@ -308,9 +310,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}'")
@@ -357,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)
@@ -375,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
@@ -393,7 +395,11 @@ class CoordinateGuesser(ABC):
         pass
 
     @abstractmethod
-    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LatitudeCoordinate]:
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
+        pass
+
+    @abstractmethod
+    def _is_latitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> LatitudeCoordinate | None:
         """Checks if the coordinate is a latitude.
 
         Parameters
@@ -411,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
@@ -429,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
@@ -447,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
@@ -465,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
@@ -483,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
@@ -501,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
@@ -519,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
@@ -550,7 +556,16 @@ class DefaultCoordinateGuesser(CoordinateGuesser):
         """
         super().__init__(ds)
 
-    def _is_longitude(self, c: xr.DataArray, attributes: CoordinateAttributes) -> Optional[LongitudeCoordinate]:
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
+        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) -> LongitudeCoordinate | None:
         """Checks if the coordinate is a longitude.
 
         Parameters
@@ -579,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
@@ -608,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
@@ -631,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
@@ -654,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
@@ -677,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
@@ -700,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
@@ -726,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
@@ -744,12 +759,18 @@ 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")
 
         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,12 +780,9 @@ 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]:
+    def _is_number(self, c: xr.DataArray, attributes: CoordinateAttributes) -> EnsembleCoordinate | None:
         """Checks if the coordinate is an ensemble coordinate.
 
         Parameters
@@ -788,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
@@ -801,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
@@ -836,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
@@ -856,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
@@ -876,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
@@ -896,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
@@ -916,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
@@ -936,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
@@ -956,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
@@ -976,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
@@ -1021,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
@@ -1040,3 +1058,23 @@ class FlavourCoordinateGuesser(CoordinateGuesser):
             return EnsembleCoordinate(c)
 
         return None
+
+    def _is_point(self, c: xr.DataArray, attributes: CoordinateAttributes) -> PointCoordinate | None:
+        """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/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()