anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/create/{functions/sources/xarray → sources/xarray_support}/grid.py

@@ -9,45 +9,83 @@
 
 
 import logging
+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
 
 LOG = logging.getLogger(__name__)
 
 
-class Grid:
+class Grid(ABC):
+    """Abstract base class for grid structures."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         pass
 
     @property
-    def latitudes(self):
+    def latitudes(self) -> Any:
+        """Get the latitudes of the grid."""
         return self.grid_points[0]
 
     @property
-    def longitudes(self):
+    def longitudes(self) -> Any:
+        """Get the longitudes of the grid."""
         return self.grid_points[1]
 
+    @property
+    @abstractmethod
+    def grid_points(self) -> Tuple[Any, Any]:
+        """Get the grid points."""
+        pass
+
 
 class LatLonGrid(Grid):
-    def __init__(self, lat, lon, variable_dims):
+    """Grid class for latitude and longitude coordinates."""
+
+    def __init__(self, lat: Any, lon: Any, variable_dims: Any) -> None:
+        """Initialize the LatLonGrid class.
+
+        Parameters
+        ----------
+        lat : Any
+            The latitudes.
+        lon : Any
+            The longitudes.
+        variable_dims : Any
+            The variable dimensions.
+        """
         super().__init__()
         self.lat = lat
         self.lon = lon
 
 
 class XYGrid(Grid):
-    def __init__(self, x, y):
+    """Grid class for x and y coordinates."""
+
+    def __init__(self, x: Any, y: Any) -> None:
+        """Initialize the XYGrid class.
+
+        Parameters
+        ----------
+        x : Any
+            The x-coordinates.
+        y : Any
+            The y-coordinates.
+        """
         self.x = x
         self.y = y
 
 
 class MeshedGrid(LatLonGrid):
+    """Grid class for meshed latitude and longitude coordinates."""
 
     @cached_property
-    def grid_points(self):
-
+    def grid_points(self) -> Tuple[Any, Any]:
+        """Get the grid points for the meshed grid."""
         lat, lon = np.meshgrid(
             self.lat.variable.values,
             self.lon.variable.values,
@@ -57,8 +95,20 @@ class MeshedGrid(LatLonGrid):
 
 
 class UnstructuredGrid(LatLonGrid):
-
-    def __init__(self, lat, lon, variable_dims):
+    """Grid class for unstructured latitude and longitude coordinates."""
+
+    def __init__(self, lat: Any, lon: Any, variable_dims: Any) -> None:
+        """Initialize the UnstructuredGrid class.
+
+        Parameters
+        ----------
+        lat : Any
+            The latitudes.
+        lon : Any
+            The longitudes.
+        variable_dims : Any
+            The variable dimensions.
+        """
         super().__init__(lat, lon, variable_dims)
         assert len(lat) == len(lon), (len(lat), len(lon))
         self.variable_dims = variable_dims
@@ -67,8 +117,8 @@ class UnstructuredGrid(LatLonGrid):
         assert set(self.variable_dims) == set(self.grid_dims), (self.variable_dims, self.grid_dims)
 
     @cached_property
-    def grid_points(self):
-
+    def grid_points(self) -> Tuple[Any, Any]:
+        """Get the grid points for the unstructured grid."""
         assert 1 <= len(self.variable_dims) <= 2
 
         if len(self.variable_dims) == 1:
@@ -88,11 +138,31 @@ class UnstructuredGrid(LatLonGrid):
 
 
 class ProjectionGrid(XYGrid):
-    def __init__(self, x, y, projection):
+    """Grid class for projected x and y coordinates."""
+
+    def __init__(self, x: Any, y: Any, projection: Any) -> None:
+        """Initialize the ProjectionGrid class.
+
+        Parameters
+        ----------
+        x : Any
+            The x-coordinates.
+        y : Any
+            The y-coordinates.
+        projection : Any
+            The projection information.
+        """
         super().__init__(x, y)
         self.projection = projection
 
-    def transformer(self):
+    def transformer(self) -> Any:
+        """Get the transformer for the projection.
+
+        Returns
+        -------
+        Any
+            The transformer.
+        """
         from pyproj import CRS
         from pyproj import Transformer
 
@@ -107,10 +177,11 @@ class ProjectionGrid(XYGrid):
 
 
 class MeshProjectionGrid(ProjectionGrid):
+    """Grid class for meshed projected coordinates."""
 
     @cached_property
-    def grid_points(self):
-
+    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")
         lon, lat = transformer.transform(xv, yv)
@@ -118,19 +189,9 @@ class MeshProjectionGrid(ProjectionGrid):
 
 
 class UnstructuredProjectionGrid(XYGrid):
-    @cached_property
-    def grid_points(self):
-        assert False, "Not implemented"
-
-        # lat, lon = transformer.transform(
-        #      self.y.variable.values.flatten(),
-        #     self.x.variable.values.flatten(),
-
-        # )
+    """Grid class for unstructured projected coordinates."""
 
-        # lat = lat[::len(lat)//100]
-        # lon = lon[::len(lon)//100]
-
-        # print(len(lat), len(lon))
-
-        # return np.meshgrid(lat, lon)
+    @cached_property
+    def grid_points(self) -> Tuple[Any, Any]:
+        """Get the grid points for the unstructured projection grid."""
+        raise NotImplementedError("UnstructuredProjectionGrid")

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

@@ -0,0 +1,395 @@
+# (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.
+
+
+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
+from earthkit.data.core.metadata import RawMetadata
+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.
+
+    Attributes
+    ----------
+    LS_KEYS : List[str]
+        List of keys for the metadata.
+    NAMESPACES : List[str]
+        List of namespaces for the metadata.
+    MARS_KEYS : List[str]
+        List of MARS keys for the metadata.
+    """
+
+    LS_KEYS = ["variable", "level", "valid_datetime", "units"]
+    NAMESPACES = ["default", "mars"]
+    MARS_KEYS = ["param", "step", "levelist", "levtype", "number", "date", "time"]
+
+    def __init__(self, field: Any) -> None:
+        """Initialize the XArrayMetadata class.
+
+        Parameters
+        ----------
+        field : Any
+            The field to extract metadata from.
+        """
+        self._field = field
+        md = field._md.copy()
+        self._mapping = _MDMapping(field.owner)
+        self._mapping.fill_time_metadata(field, md)
+        super().__init__(md)
+
+    @cached_property
+    def geography(self) -> "XArrayFieldGeography":
+        """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]:
+        """Get the metadata as a specific namespace.
+
+        Parameters
+        ----------
+        namespace : Optional[str]
+            The namespace to use.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata in the specified namespace.
+        """
+        if not isinstance(namespace, str) and namespace is not None:
+            raise TypeError("namespace must be a str or None")
+
+        if namespace == "default" or namespace == "" or namespace is None:
+            return dict(self)
+
+        elif namespace == "mars":
+            return self._as_mars()
+
+    def _as_mars(self) -> Dict[str, Any]:
+        """Get the metadata as MARS namespace.
+
+        Returns
+        -------
+        Dict[str, Any]
+            The metadata in the MARS namespace.
+        """
+        return {}
+
+    def _base_datetime(self) -> Optional[datetime.datetime]:
+        """Get the base datetime for the field.
+
+        Returns
+        -------
+        Optional[datetime.datetime]
+            The base datetime for the field.
+        """
+        return self._field.forecast_reference_time
+
+    def _valid_datetime(self) -> Optional[datetime.datetime]:
+        """Get the valid datetime for the field.
+
+        Returns
+        -------
+        Optional[datetime.datetime]
+            The valid datetime for the field.
+        """
+        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 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.
+
+    Attributes
+    ----------
+    _field : Any
+        The field to extract geography information from.
+    _grid : Any
+        The grid associated with the field.
+    """
+
+    def __init__(self, field: Any, grid: Any) -> None:
+        """Initialize the XArrayFieldGeography class.
+
+        Parameters
+        ----------
+        field : Any
+            The field to extract geography information from.
+        grid : Any
+            The grid associated with the field.
+        """
+        self._field = field
+        self._grid = grid
+
+    def _unique_grid_id(self) -> None:
+        """Get the unique grid ID.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented.
+        """
+        raise NotImplementedError()
+
+    def bounding_box(self) -> None:
+        """Get the bounding box for the field.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented.
+        """
+        raise NotImplementedError()
+        # return BoundingBox(north=self.north, south=self.south, east=self.east, west=self.west)
+
+    def gridspec(self) -> None:
+        """Get the grid specification for the field.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented.
+        """
+        raise NotImplementedError()
+
+    def latitudes(self, dtype: Optional[type] = None) -> Any:
+        """Get the latitudes for the field.
+
+        Parameters
+        ----------
+        dtype : Optional[type]
+            The type to cast the latitudes to.
+
+        Returns
+        -------
+        Any
+            The latitudes for the field.
+        """
+        result = self._grid.latitudes
+        if dtype is not None:
+            return result.astype(dtype)
+        return result
+
+    def longitudes(self, dtype: Optional[type] = None) -> Any:
+        """Get the longitudes for the field.
+
+        Parameters
+        ----------
+        dtype : Optional[type]
+            The type to cast the longitudes to.
+
+        Returns
+        -------
+        Any
+            The longitudes for the field.
+        """
+        result = self._grid.longitudes
+        if dtype is not None:
+            return result.astype(dtype)
+        return result
+
+    def resolution(self) -> Optional[Any]:
+        """Get the resolution for the field.
+
+        Returns
+        -------
+        Optional[Any]
+            The resolution for the field.
+        """
+        # TODO: implement resolution
+        return None
+
+    def mars_grid(self) -> Optional[Any]:
+        """Get the MARS grid for the field.
+
+        Returns
+        -------
+        Optional[Any]
+            The MARS grid for the field.
+        """
+        # TODO: implement mars_grid
+        return None
+
+    def mars_area(self) -> Optional[Any]:
+        """Get the MARS area for the field.
+
+        Returns
+        -------
+        Optional[Any]
+            The MARS area for the field.
+        """
+        # TODO: code me
+        # return [self.north, self.west, self.south, self.east]
+        return None
+
+    def x(self, dtype: Optional[type] = None) -> None:
+        """Get the x-coordinates for the field.
+
+        Parameters
+        ----------
+        dtype : Optional[type]
+            The type to cast the x-coordinates to.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented.
+        """
+        raise NotImplementedError()
+
+    def y(self, dtype: Optional[type] = None) -> None:
+        """Get the y-coordinates for the field.
+
+        Parameters
+        ----------
+        dtype : Optional[type]
+            The type to cast the y-coordinates to.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented.
+        """
+        raise NotImplementedError()
+
+    def shape(self) -> Any:
+        """Get the shape of the field.
+
+        Returns
+        -------
+        Any
+            The shape of the field.
+        """
+        return self._field.shape
+
+    def projection(self) -> Projection:
+        """Get the projection for the field.
+
+        Returns
+        -------
+        Projection
+            The projection for the field.
+        """
+        return Projection.from_cf_grid_mapping(**self._field.grid_mapping)

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

@@ -0,0 +1,91 @@
+# (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.
+
+
+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:
+    """Patch the attributes of the dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    attributes : Dict[str, Dict[str, Any]]
+        The attributes to patch.
+
+    Returns
+    -------
+    Any
+        The patched dataset.
+    """
+    for name, value in attributes.items():
+        variable = ds[name]
+        variable.attrs.update(value)
+
+    return ds
+
+
+def patch_coordinates(ds: xr.Dataset, coordinates: List[str]) -> Any:
+    """Patch the coordinates of the dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    coordinates : List[str]
+        The coordinates to patch.
+
+    Returns
+    -------
+    Any
+        The patched dataset.
+    """
+    for name in coordinates:
+        ds = ds.assign_coords({name: ds[name]})
+
+    return ds
+
+
+PATCHES = {
+    "attributes": patch_attributes,
+    "coordinates": patch_coordinates,
+}
+
+
+def patch_dataset(ds: xr.Dataset, patch: Dict[str, Dict[str, Any]]) -> Any:
+    """Patch the dataset.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to patch.
+    patch : Dict[str, Dict[str, Any]]
+        The patch to apply.
+
+    Returns
+    -------
+    Any
+        The patched dataset.
+    """
+    for what, values in patch.items():
+        if what not in PATCHES:
+            raise ValueError(f"Unknown patch type {what!r}")
+
+        ds = PATCHES[what](ds, values)
+
+    return ds