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

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

@@ -7,19 +7,41 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import datetime
 import logging
-
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+import earthkit.data as ekd
+import xarray as xr
 from earthkit.data.core.fieldlist import MultiFieldList
 
+from anemoi.datasets.create.sources.patterns import iterate_patterns
 from anemoi.datasets.data.stores import name_to_zarr_store
 
-from .. import iterate_patterns
+from ..legacy import legacy_source
 from .fieldlist import XarrayFieldList
 
 LOG = logging.getLogger(__name__)
 
 
-def check(what, ds, paths, **kwargs):
+def check(what: str, ds: xr.Dataset, paths: List[str], **kwargs: Any) -> None:
+    """Checks if the dataset has the expected number of fields.
+
+    Parameters
+    ----------
+    what : str
+        Description of what is being checked.
+    ds : xr.Dataset
+        The dataset to check.
+    paths : List[str]
+        List of paths.
+    **kwargs : Any
+        Additional keyword arguments.
+    """
     count = 1
     for k, v in kwargs.items():
         if isinstance(v, (tuple, list)):
@@ -29,8 +51,43 @@ def check(what, ds, paths, **kwargs):
         raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, {what}s={paths})")
 
 
-def load_one(emoji, context, dates, dataset, *, options={}, flavour=None, patch=None, **kwargs):
-    import xarray as xr
+def load_one(
+    emoji: str,
+    context: Any,
+    dates: List[str],
+    dataset: Union[str, xr.Dataset],
+    *,
+    options: Optional[Dict[str, Any]] = None,
+    flavour: Optional[str] = None,
+    patch: Optional[Any] = None,
+    **kwargs: Any,
+) -> ekd.FieldList:
+    """Loads a single dataset.
+
+    Parameters
+    ----------
+    emoji : str
+        Emoji for tracing.
+    context : Any
+        Context object.
+    dates : List[str]
+        List of dates.
+    dataset : Union[str, xr.Dataset]
+        The dataset to load.
+    options : Dict[str, Any], optional
+        Additional options for loading the dataset.
+    flavour : Optional[str], optional
+        Flavour of the dataset.
+    patch : Optional[Any], optional
+        Patch for the dataset.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    MultiFieldList
+        The loaded dataset.
+    """
 
     """
     We manage the S3 client ourselve, bypassing fsspec and s3fs layers, because sometimes something on the stack
@@ -41,6 +98,9 @@ def load_one(emoji, context, dates, dataset, *, options={}, flavour=None, patch=
     We have seen this bug triggered when we run many clients in parallel, for example, when we create a new dataset using `xarray-zarr`.
     """
 
+    if options is None:
+        options = {}
+
     context.trace(emoji, dataset, options, kwargs)
 
     if isinstance(dataset, str) and ".zarr" in dataset:
@@ -59,6 +119,7 @@ def load_one(emoji, context, dates, dataset, *, options={}, flavour=None, patch=
     if len(dates) == 0:
         result = fs.sel(**kwargs)
     else:
+        print("dates", dates, kwargs)
         result = MultiFieldList([fs.sel(valid_datetime=date, **kwargs) for date in dates])
 
     if len(result) == 0:
@@ -79,8 +140,27 @@ def load_one(emoji, context, dates, dataset, *, options={}, flavour=None, patch=
     return result
 
 
-def load_many(emoji, context, dates, pattern, **kwargs):
-
+def load_many(emoji: str, context: Any, dates: List[datetime.datetime], pattern: str, **kwargs: Any) -> ekd.FieldList:
+    """Loads multiple datasets.
+
+    Parameters
+    ----------
+    emoji : str
+        Emoji for tracing.
+    context : Any
+        Context object.
+    dates : List[str]
+        List of dates.
+    pattern : str
+        Pattern for loading datasets.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    MultiFieldList
+        The loaded datasets.
+    """
     result = []
 
     for path, dates in iterate_patterns(pattern, dates, **kwargs):
@@ -89,5 +169,26 @@ def load_many(emoji, context, dates, pattern, **kwargs):
     return MultiFieldList(result)
 
 
-def execute(context, dates, url, *args, **kwargs):
+@legacy_source("xarray")
+def execute(context: Any, dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+    """Executes the loading of datasets.
+
+    Parameters
+    ----------
+    context : Any
+        Context object.
+    dates : List[str]
+        List of dates.
+    url : str
+        URL pattern for loading datasets.
+    *args : Any
+        Additional arguments.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    ekd.FieldList
+        The loaded datasets.
+    """
     return load_many("🌐", context, dates, url, *args, **kwargs)

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

@@ -0,0 +1,442 @@
+# (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.
+
+
+from __future__ import annotations
+
+import datetime
+import logging
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+import numpy as np
+import xarray as xr
+from earthkit.data.utils.dates import to_datetime
+
+LOG = logging.getLogger(__name__)
+
+
+def is_scalar(variable: Any) -> bool:
+    """Check if the variable is scalar.
+
+    Parameters
+    ----------
+    variable : Any
+        The variable to check.
+
+    Returns
+    -------
+    bool
+        True if the variable is scalar, False otherwise.
+    """
+    shape = variable.shape
+    if shape == (1,):
+        return True
+    if len(shape) == 0:
+        return True
+    return False
+
+
+def extract_single_value(variable: Any) -> Any:
+    """Extract a single value from the variable.
+
+    Parameters
+    ----------
+    variable : Any
+        The variable to extract the value from.
+
+    Returns
+    -------
+    Any
+        The extracted single value.
+    """
+    shape = variable.shape
+    if np.issubdtype(variable.values.dtype, np.datetime64):
+        if len(shape) == 0:
+            return to_datetime(variable.values)  # Convert to python datetime
+        if shape == (1,):
+            return to_datetime(variable.values[0])
+        assert False, (shape, variable.values[:2])
+
+    if np.issubdtype(variable.values.dtype, np.timedelta64):
+        if len(shape) == 0:
+            # Convert to python timedelta64
+            return datetime.timedelta(seconds=variable.values.astype("timedelta64[s]").astype(int).item())
+        assert False, (shape, variable.values)
+
+    if shape == (1,):
+        return variable.values[0]
+
+    if len(shape) == 0:
+        return variable.values.item()
+
+    assert False, (shape, variable.values)
+
+
+class Coordinate:
+    """Base class for coordinates."""
+
+    is_grid = False
+    is_dim = True
+    is_lat = False
+    is_lon = False
+    is_time = False
+    is_step = False
+    is_date = False
+    is_member = False
+    is_x = False
+    is_y = False
+
+    def __init__(self, variable: xr.DataArray) -> None:
+        """Initialize the coordinate.
+
+        Parameters
+        ----------
+        variable : Any
+            The variable representing the coordinate.
+        """
+        self.variable = variable
+        self.scalar = is_scalar(variable)
+        self.kwargs: Dict[str, Any] = {}  # Used when creating a new coordinate (reduced method)
+
+    def __len__(self) -> int:
+        """Get the length of the coordinate.
+
+        Returns
+        -------
+        int
+            The length of the coordinate.
+        """
+        return 1 if self.scalar else len(self.variable)
+
+    def __repr__(self) -> str:
+        """Get the string representation of the coordinate.
+
+        Returns
+        -------
+        str
+            The string representation of the coordinate.
+        """
+        return "%s[name=%s,values=%s,shape=%s]" % (
+            self.__class__.__name__,
+            self.variable.name,
+            self.variable.values if self.scalar else len(self),
+            self.variable.shape,
+        )
+
+    def reduced(self, i: int) -> Coordinate:
+        """Create a new coordinate with a single value.
+
+        Parameters
+        ----------
+        i : int
+            The index of the value to select.
+
+        Returns
+        -------
+        Coordinate
+            A new coordinate with the selected value.
+        """
+        return self.__class__(
+            self.variable.isel({self.variable.dims[0]: i}),
+            **self.kwargs,
+        )
+
+    def index(self, value: Union[Any, list, tuple]) -> Optional[Union[int, list]]:
+        """Return the index of the value in the coordinate.
+
+        Parameters
+        ----------
+        value : Union[Any, list, tuple]
+            The value to search for.
+
+        Returns
+        -------
+        Optional[Union[int, list]]
+            The index or indices of the value in the coordinate, or None if not found.
+        """
+        if isinstance(value, (list, tuple)):
+            if len(value) == 1:
+                return self._index_single(value)
+            else:
+                return self._index_multiple(value)
+        return self._index_single(value)
+
+    def _index_single(self, value: Any) -> Optional[int]:
+        """Return the index of a single value in the coordinate.
+
+        Parameters
+        ----------
+        value : Any
+            The value to search for.
+
+        Returns
+        -------
+        Optional[int]
+            The index of the value in the coordinate, or None if not found.
+        """
+        values = self.variable.values
+
+        # Check if dimension is 0D
+        if not isinstance(values, (list, np.ndarray)):
+            values = [values]
+
+        # Assume the array is sorted
+        index = np.searchsorted(values, value)
+
+        if index < len(values) and values[index] == value:
+            return index
+
+        # If not found, we need to check if the value is in the array
+
+        index = np.where(values == value)[0]
+        if len(index) > 0:
+            return index[0]
+
+        return None
+
+    def _index_multiple(self, value: list) -> Optional[list]:
+        """Return the indices of multiple values in the coordinate.
+
+        Parameters
+        ----------
+        value : list
+            The values to search for.
+
+        Returns
+        -------
+        Optional[list]
+            The indices of the values in the coordinate, or None if not found.
+        """
+        values = self.variable.values
+
+        # Check if dimension is 0D
+        if not isinstance(values, (list, np.ndarray)):
+            values = [values]
+
+        # Assume the array is sorted
+
+        index = np.searchsorted(values, value)
+        index = index[index < len(values)]
+
+        if np.all(values[index] == value):
+            return index
+
+        # If not found, we need to check if the value is in the array
+
+        index = np.where(np.isin(values, value))[0]
+
+        # We could also return incomplete matches
+        if len(index) == len(value):
+            return index
+
+        return None
+
+    @property
+    def name(self) -> str:
+        """Get the name of the coordinate."""
+        return self.variable.name
+
+    def normalise(self, value: Any) -> Any:
+        """Normalize the value for the coordinate.
+
+        Parameters
+        ----------
+        value : Any
+            The value to normalize.
+
+        Returns
+        -------
+        Any
+            The normalized value.
+        """
+        # Subclasses to format values that will be added to the field metadata
+        return value
+
+    @property
+    def single_value(self) -> Any:
+        """Get the single value of the coordinate."""
+        return extract_single_value(self.variable)
+
+
+class TimeCoordinate(Coordinate):
+    """Coordinate class for time."""
+
+    is_time = True
+    mars_names = ("valid_datetime",)
+
+    def index(self, time: datetime.datetime) -> Optional[int]:
+        """Return the index of the time in the coordinate.
+
+        Parameters
+        ----------
+        time : datetime.datetime
+            The time to search for.
+
+        Returns
+        -------
+        Optional[int]
+            The index of the time in the coordinate, or None if not found.
+        """
+        return super().index(np.datetime64(time))
+
+
+class DateCoordinate(Coordinate):
+    """Coordinate class for date."""
+
+    is_date = True
+    mars_names = ("date",)
+
+    def index(self, date: datetime.datetime) -> Optional[int]:
+        """Return the index of the date in the coordinate.
+
+        Parameters
+        ----------
+        date : datetime.datetime
+            The date to search for.
+
+        Returns
+        -------
+        Optional[int]
+            The index of the date in the coordinate, or None if not found.
+        """
+        return super().index(np.datetime64(date))
+
+
+class StepCoordinate(Coordinate):
+    """Coordinate class for step."""
+
+    is_step = True
+    mars_names = ("step",)
+
+
+class LevelCoordinate(Coordinate):
+    """Coordinate class for level.
+
+    Parameters
+    ----------
+    variable : Any
+        The variable representing the coordinate.
+    levtype : str
+        The type of level.
+    """
+
+    mars_names = ("level", "levelist")
+
+    def __init__(self, variable: Any, levtype: str) -> None:
+        """Initialize the level coordinate.
+
+        Parameters
+        ----------
+        variable : Any
+            The variable representing the coordinate.
+        levtype : str
+            The type of level.
+        """
+        super().__init__(variable)
+        self.levtype = levtype
+        # kwargs is used when creating a new coordinate (reduced method)
+        self.kwargs = {"levtype": levtype}
+
+    def normalise(self, value: Any) -> Any:
+        """Normalize the value for the level coordinate.
+
+        Parameters
+        ----------
+        value : Any
+            The value to normalize.
+
+        Returns
+        -------
+        Any
+            The normalized value.
+        """
+        # Some netcdf have pressue levels in float
+        if int(value) == value:
+            return int(value)
+        return value
+
+
+class EnsembleCoordinate(Coordinate):
+    """Coordinate class for ensemble."""
+
+    is_member = True
+    mars_names = ("number",)
+
+    def normalise(self, value: Any) -> Any:
+        """Normalize the value for the ensemble coordinate.
+
+        Parameters
+        ----------
+        value : Any
+            The value to normalize.
+
+        Returns
+        -------
+        Any
+            The normalized value.
+        """
+        if int(value) == value:
+            return int(value)
+        return value
+
+
+class LongitudeCoordinate(Coordinate):
+    """Coordinate class for longitude."""
+
+    is_grid = True
+    is_lon = True
+    mars_names = ("longitude",)
+
+
+class LatitudeCoordinate(Coordinate):
+    """Coordinate class for latitude."""
+
+    is_grid = True
+    is_lat = True
+    mars_names = ("latitude",)
+
+
+class XCoordinate(Coordinate):
+    """Coordinate class for X."""
+
+    is_grid = True
+    is_x = True
+    mars_names = ("x",)
+
+
+class YCoordinate(Coordinate):
+    """Coordinate class for Y."""
+
+    is_grid = True
+    is_y = True
+    mars_names = ("y",)
+
+
+class ScalarCoordinate(Coordinate):
+    """Coordinate class for scalar."""
+
+    is_grid = False
+
+    @property
+    def mars_names(self) -> Tuple[str, ...]:
+        """Get the MARS names for the coordinate."""
+        return (self.variable.name,)
+
+
+class UnsupportedCoordinate(Coordinate):
+    """Coordinate class for unsupported coordinates."""
+
+    @property
+    def mars_names(self) -> tuple:
+        """Get the MARS names for the coordinate."""
+        return (self.variable.name,)