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

anemoi/datasets/create/sources/grib.py

@@ -0,0 +1,297 @@
+# (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 glob
+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
+from earthkit.data import from_source
+from earthkit.data.indexing.fieldlist import FieldArray
+from earthkit.data.utils.patterns import Pattern
+
+from .legacy import legacy_source
+
+LOG = logging.getLogger(__name__)
+
+
+def _load(context: Any, name: str, record: Dict[str, Any]) -> tuple:
+    """Load data from a given source.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    name : str
+        The name of the data source.
+    record : dict of str to Any
+        The record containing source information.
+
+    Returns
+    -------
+    tuple
+        A tuple containing the data as a numpy array and the UUID of the HGrid.
+    """
+    ds = None
+
+    param = record["param"]
+
+    if "path" in record:
+        context.info(f"Using {name} from {record['path']} (param={param})")
+        ds = from_source("file", record["path"])
+
+    if "url" in record:
+        context.info(f"Using {name} from {record['url']} (param={param})")
+        ds = from_source("url", record["url"])
+
+    ds = ds.sel(param=param)
+
+    assert len(ds) == 1, f"{name} {param}, expected one field, got {len(ds)}"
+    ds = ds[0]
+
+    return ds.to_numpy(flatten=True), ds.metadata("uuidOfHGrid")
+
+
+class Geography:
+    """This class retrieves the latitudes and longitudes of unstructured grids,
+    and checks if the fields are compatible with the grid.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    latitudes : dict of str to Any
+        Latitude information.
+    longitudes : dict of str to Any
+        Longitude information.
+    """
+
+    def __init__(self, context: Any, latitudes: Dict[str, Any], longitudes: Dict[str, Any]) -> None:
+        """Initialize the Geography class.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the function is executed.
+        latitudes : dict of str to Any
+            Latitude information.
+        longitudes : dict of str to Any
+            Longitude information.
+        """
+        latitudes, uuidOfHGrid_lat = _load(context, "latitudes", latitudes)
+        longitudes, uuidOfHGrid_lon = _load(context, "longitudes", longitudes)
+
+        assert (
+            uuidOfHGrid_lat == uuidOfHGrid_lon
+        ), f"uuidOfHGrid mismatch: lat={uuidOfHGrid_lat} != lon={uuidOfHGrid_lon}"
+
+        context.info(f"Latitudes: {len(latitudes)}, Longitudes: {len(longitudes)}")
+        assert len(latitudes) == len(longitudes)
+
+        self.uuidOfHGrid = uuidOfHGrid_lat
+        self.latitudes = latitudes
+        self.longitudes = longitudes
+        self.first = True
+
+    def check(self, field: Any) -> None:
+        """Check if the field is compatible with the grid.
+
+        Parameters
+        ----------
+        field : Any
+            The field to check.
+        """
+        if self.first:
+            # We only check the first field, for performance reasons
+            assert (
+                field.metadata("uuidOfHGrid") == self.uuidOfHGrid
+            ), f"uuidOfHGrid mismatch: {field.metadata('uuidOfHGrid')} != {self.uuidOfHGrid}"
+            self.first = False
+
+
+class AddGrid:
+    """An earth-kit.data.Field wrapper that adds grid information.
+
+    Parameters
+    ----------
+    field : Any
+        The field to wrap.
+    geography : Geography
+        The geography information.
+    """
+
+    def __init__(self, field: Any, geography: Geography) -> None:
+        """Initialize the AddGrid class.
+
+        Parameters
+        ----------
+        field : Any
+            The field to wrap.
+        geography : Geography
+            The geography information.
+        """
+        self._field = field
+
+        geography.check(field)
+
+        self._latitudes = geography.latitudes
+        self._longitudes = geography.longitudes
+
+    def __getattr__(self, name: str) -> Any:
+        """Get an attribute from the wrapped field.
+
+        Parameters
+        ----------
+        name : str
+            The name of the attribute.
+
+        Returns
+        -------
+        Any
+            The attribute value.
+        """
+        return getattr(self._field, name)
+
+    def __repr__(self) -> str:
+        """Get the string representation of the wrapped field.
+
+        Returns
+        -------
+        str
+            The string representation.
+        """
+        return repr(self._field)
+
+    def grid_points(self) -> tuple:
+        """Get the grid points (latitudes and longitudes).
+
+        Returns
+        -------
+        tuple
+            The latitudes and longitudes.
+        """
+        return self._latitudes, self._longitudes
+
+    @property
+    def resolution(self) -> str:
+        """Get the resolution of the grid."""
+        return "unknown"
+
+
+def check(ds: Any, paths: List[str], **kwargs: Any) -> None:
+    """Check if the dataset matches the expected number of fields.
+
+    Parameters
+    ----------
+    ds : Any
+        The dataset to check.
+    paths : list of str
+        List of paths to the GRIB files.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Raises
+    ------
+    ValueError
+        If the number of fields does not match the expected count.
+    """
+    count = 1
+    for k, v in kwargs.items():
+        if isinstance(v, (tuple, list)):
+            count *= len(v)
+
+    if len(ds) != count:
+        raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})")
+
+
+def _expand(paths: List[str]) -> Any:
+    """Expand the given paths using glob.
+
+    Parameters
+    ----------
+    paths : list of str
+        List of paths to expand.
+
+    Returns
+    -------
+    Any
+        The expanded paths.
+    """
+    for path in paths:
+        cnt = 0
+        for p in glob.glob(path):
+            yield p
+            cnt += 1
+        if cnt == 0:
+            yield path
+
+
+@legacy_source(__file__)
+def execute(
+    context: Any,
+    dates: List[Any],
+    path: Union[str, List[str]],
+    latitudes: Optional[Dict[str, Any]] = None,
+    longitudes: Optional[Dict[str, Any]] = None,
+    *args: Any,
+    **kwargs: Any,
+) -> ekd.FieldList:
+    """Execute the function to load data from GRIB files.
+
+    Args:
+        context (Any): The context in which the function is executed.
+        dates (List[Any]): List of dates.
+        path (Union[str, List[str]]): Path or list of paths to the GRIB files.
+        latitudes (Optional[Dict[str, Any]], optional): Latitude information. Defaults to None.
+        longitudes (Optional[Dict[str, Any]], optional): Longitude information. Defaults to None.
+        *args (Any): Additional arguments.
+        **kwargs (Any): Additional keyword arguments.
+
+    Returns
+    -------
+    Any
+        The loaded dataset.
+    """
+    given_paths = path if isinstance(path, list) else [path]
+
+    geography = None
+    if latitudes is not None and longitudes is not None:
+        geography = Geography(context, latitudes, longitudes)
+
+    ds = from_source("empty")
+    dates = [d.isoformat() for d in dates]
+
+    for path in given_paths:
+        paths = Pattern(path, ignore_missing_keys=True).substitute(*args, date=dates, **kwargs)
+
+        for name in ("grid", "area", "rotation", "frame", "resol", "bitmap"):
+            if name in kwargs:
+                raise ValueError(f"MARS interpolation parameter '{name}' not supported")
+
+        for path in _expand(paths):
+            context.trace("📁", "PATH", path)
+            s = from_source("file", path)
+            s = s.sel(valid_datetime=dates, **kwargs)
+            ds = ds + s
+
+    if kwargs and not context.partial_ok:
+        check(ds, given_paths, valid_datetime=dates, **kwargs)
+
+    if geography is not None:
+        ds = FieldArray([AddGrid(_, geography) for _ in ds])
+
+    if len(ds) == 0:
+        LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")
+
+    return ds

anemoi/datasets/create/{functions/sources → sources}/hindcasts.py

@@ -8,22 +8,56 @@
 # nor does it submit to any jurisdiction.
 
 import logging
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Union
 
 from earthkit.data.core.fieldlist import MultiFieldList
 
-from anemoi.datasets.create.functions.sources.mars import mars
+from anemoi.datasets.create.sources.mars import mars
+
+from .legacy import legacy_source
 
 LOGGER = logging.getLogger(__name__)
 
 
-def _to_list(x):
+def _to_list(x: Union[list, tuple, Any]) -> List[Any]:
+    """Converts the input to a list if it is not already a list or tuple.
+
+    Parameters
+    ----------
+    x : Any
+        The input to convert.
+
+    Returns
+    -------
+    list
+        A list containing the input elements.
+    """
     if isinstance(x, (list, tuple)):
         return x
     return [x]
 
 
-def hindcasts(context, dates, **request):
-
+@legacy_source(__file__)
+def hindcasts(context: Any, dates: List[Any], **request: Dict[str, Any]) -> MultiFieldList:
+    """Generates hindcast requests based on the provided dates and request parameters.
+
+    Parameters
+    ----------
+    context : Any
+        The context containing the dates provider and trace method.
+    dates : List[Any]
+        A list of dates for which to generate hindcast requests.
+    request : Dict[str, Any]
+        Additional request parameters.
+
+    Returns
+    -------
+    MultiFieldList
+        A MultiFieldList containing the hindcast data.
+    """
     from anemoi.datasets.dates import HindcastsDates
 
     provider = context.dates_provider

anemoi/datasets/create/sources/legacy.py

@@ -0,0 +1,93 @@
+# (C) Copyright 2025- 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 inspect
+import logging
+import os
+from typing import Any
+from typing import Callable
+
+from ..source import Source
+from . import source_registry
+
+LOG = logging.getLogger(__name__)
+
+
+class LegacySource(Source):
+    """A legacy source class.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the source is created.
+    *args : tuple
+        Positional arguments.
+    **kwargs : dict
+        Keyword arguments.
+    """
+
+    def __init__(self, context: Any, *args: Any, **kwargs: Any) -> None:
+        super().__init__(context, *args, **kwargs)
+        self.args = args
+        self.kwargs = kwargs
+
+
+class legacy_source:
+    """A decorator class for legacy sources.
+
+    Parameters
+    ----------
+    name : str
+        The name of the legacy source.
+    """
+
+    def __init__(self, name: str) -> None:
+        name, _ = os.path.splitext(os.path.basename(name))
+        self.name = name
+
+    def __call__(self, execute: Callable) -> Callable:
+        """Call method to wrap the execute function.
+
+        Parameters
+        ----------
+        execute : function
+            The execute function to be wrapped.
+
+        Returns
+        -------
+        function
+            The wrapped execute function.
+        """
+        this = self
+        name = f"Legacy{self.name.title()}Source"
+        source = ".".join([execute.__module__, execute.__name__])
+
+        def execute_wrapper(self, dates) -> Any:
+            """Wrapper method to call the execute function."""
+            try:
+                return execute(self.context, dates, *self.args, **self.kwargs)
+            except TypeError:
+                LOG.error(f"Error executing source {this.name} from {source}")
+                LOG.error(f"Function signature is: {inspect.signature(execute)}")
+                LOG.error(f"Arguments are: {self.args=}, {self.kwargs=}")
+                raise
+
+        klass = type(
+            name,
+            (LegacySource,),
+            {
+                "execute": execute_wrapper,
+                "_source": source,
+            },
+        )
+
+        source_registry.register(self.name)(klass)
+
+        return execute

anemoi/datasets/create/{functions/sources → sources}/mars.py

@@ -9,6 +9,12 @@
 
 import datetime
 import re
+from typing import Any
+from typing import Dict
+from typing import Generator
+from typing import List
+from typing import Optional
+from typing import Union
 
 from anemoi.utils.humanize import did_you_mean
 from earthkit.data import from_source
@@ -16,16 +22,44 @@ from earthkit.data.utils.availability import Availability
 
 from anemoi.datasets.create.utils import to_datetime_list
 
+from .legacy import legacy_source
+
 DEBUG = False
 
 
-def to_list(x):
+def to_list(x: Union[list, tuple, Any]) -> list:
+    """Converts the input to a list if it is not already a list or tuple.
+
+    Parameters
+    ----------
+    x : Any
+        The input value to be converted.
+
+    Returns
+    -------
+    list
+        A list containing the input value(s).
+    """
     if isinstance(x, (list, tuple)):
         return x
     return [x]
 
 
-def _date_to_datetime(d):
+def _date_to_datetime(
+    d: Union[datetime.datetime, list, tuple, str],
+) -> Union[datetime.datetime, List[datetime.datetime]]:
+    """Converts the input date(s) to datetime objects.
+
+    Parameters
+    ----------
+    d : Union[datetime.datetime, list, tuple, str]
+        The input date(s) to be converted.
+
+    Returns
+    -------
+    Union[datetime.datetime, List[datetime.datetime]]
+        A datetime object or a list of datetime objects.
+    """
     if isinstance(d, datetime.datetime):
         return d
     if isinstance(d, (list, tuple)):
@@ -33,8 +67,19 @@ def _date_to_datetime(d):
     return datetime.datetime.fromisoformat(d)
 
 
-def expand_to_by(x):
+def expand_to_by(x: Union[str, int, list]) -> Union[str, int, list]:
+    """Expands a range expression to a list of values.
+
+    Parameters
+    ----------
+    x : Union[str, int, list]
+        The input range expression.
 
+    Returns
+    -------
+    Union[str, int, list]
+        A list of expanded values.
+    """
     if isinstance(x, (str, int)):
         return expand_to_by(str(x).split("/"))
 
@@ -52,7 +97,19 @@ def expand_to_by(x):
     return x
 
 
-def normalise_time_delta(t):
+def normalise_time_delta(t: Union[datetime.timedelta, str]) -> datetime.timedelta:
+    """Normalizes a time delta string to a datetime.timedelta object.
+
+    Parameters
+    ----------
+    t : Union[datetime.timedelta, str]
+        The input time delta string.
+
+    Returns
+    -------
+    datetime.timedelta
+        A normalized datetime.timedelta object.
+    """
     if isinstance(t, datetime.timedelta):
         assert t == datetime.timedelta(hours=t.hours), t
 
@@ -63,14 +120,49 @@ def normalise_time_delta(t):
     return t
 
 
-def _normalise_time(t):
+def _normalise_time(t: Union[int, str]) -> str:
+    """Normalizes a time value to a string in HHMM format.
+
+    Parameters
+    ----------
+    t : Union[int, str]
+        The input time value.
+
+    Returns
+    -------
+    str
+        A string representing the normalized time.
+    """
     t = int(t)
     if t < 100:
         t * 100
     return "{:04d}".format(t)
 
 
-def _expand_mars_request(request, date, request_already_using_valid_datetime=False, date_key="date"):
+def _expand_mars_request(
+    request: Dict[str, Any],
+    date: datetime.datetime,
+    request_already_using_valid_datetime: bool = False,
+    date_key: str = "date",
+) -> List[Dict[str, Any]]:
+    """Expands a MARS request with the given date and other parameters.
+
+    Parameters
+    ----------
+    request : Dict[str, Any]
+        The input MARS request.
+    date : datetime.datetime
+        The date to be used in the request.
+    request_already_using_valid_datetime : bool, optional
+        Flag indicating if the request already uses valid datetime.
+    date_key : str, optional
+        The key for the date in the request.
+
+    Returns
+    -------
+    List[Dict[str, Any]]
+        A list of expanded MARS requests.
+    """
     requests = []
 
     user_step = to_list(expand_to_by(request.get("step", [0])))
@@ -130,11 +222,29 @@ def _expand_mars_request(request, date, request_already_using_valid_datetime=Fal
 
 
 def factorise_requests(
-    dates,
-    *requests,
-    request_already_using_valid_datetime=False,
-    date_key="date",
-):
+    dates: List[datetime.datetime],
+    *requests: Dict[str, Any],
+    request_already_using_valid_datetime: bool = False,
+    date_key: str = "date",
+) -> Generator[Dict[str, Any], None, None]:
+    """Factorizes the requests based on the given dates.
+
+    Parameters
+    ----------
+    dates : List[datetime.datetime]
+        The list of dates to be used in the requests.
+    requests : Dict[str, Any]
+        The input requests to be factorized.
+    request_already_using_valid_datetime : bool, optional
+        Flag indicating if the requests already use valid datetime.
+    date_key : str, optional
+        The key for the date in the requests.
+
+    Returns
+    -------
+    Generator[Dict[str, Any], None, None]
+        Factorized requests.
+    """
     updates = []
     for req in requests:
         # req = normalise_request(req)
@@ -158,7 +268,19 @@ def factorise_requests(
         yield r
 
 
-def use_grib_paramid(r):
+def use_grib_paramid(r: Dict[str, Any]) -> Dict[str, Any]:
+    """Converts the parameter short names to GRIB parameter IDs.
+
+    Parameters
+    ----------
+    r : Dict[str, Any]
+        The input request containing parameter short names.
+
+    Returns
+    -------
+    Dict[str, Any]
+        The request with parameter IDs.
+    """
     from anemoi.utils.grib import shortname_to_paramid
 
     params = r["param"]
@@ -240,15 +362,40 @@ MARS_KEYS = [
 ]
 
 
+@legacy_source(__file__)
 def mars(
-    context,
-    dates,
-    *requests,
-    request_already_using_valid_datetime=False,
-    date_key="date",
-    use_cdsapi_dataset=None,
-    **kwargs,
-):
+    context: Any,
+    dates: List[datetime.datetime],
+    *requests: Dict[str, Any],
+    request_already_using_valid_datetime: bool = False,
+    date_key: str = "date",
+    use_cdsapi_dataset: Optional[str] = None,
+    **kwargs: Any,
+) -> Any:
+    """Executes MARS requests based on the given context, dates, and other parameters.
+
+    Parameters
+    ----------
+    context : Any
+        The context for the requests.
+    dates : List[datetime.datetime]
+        The list of dates to be used in the requests.
+    requests : Dict[str, Any]
+        The input requests to be executed.
+    request_already_using_valid_datetime : bool, optional
+        Flag indicating if the requests already use valid datetime.
+    date_key : str, optional
+        The key for the date in the requests.
+    use_cdsapi_dataset : Optional[str], optional
+        The dataset to be used with CDS API.
+    kwargs : Any
+        Additional keyword arguments for the requests.
+
+    Returns
+    -------
+    Any
+        The resulting dataset.
+    """
 
     if not requests:
         requests = [kwargs]
@@ -318,6 +465,7 @@ def mars(
 
 execute = mars
 
+
 if __name__ == "__main__":
     import yaml