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

anemoi/datasets/create/sources/netcdf.py

@@ -0,0 +1,42 @@
+# (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 typing import Any
+from typing import List
+
+import earthkit.data as ekd
+
+from .legacy import legacy_source
+from .xarray import load_many
+
+
+@legacy_source(__file__)
+def execute(context: Any, dates: List[str], path: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+    """Execute the loading of multiple NetCDF files.
+
+    Parameters
+    ----------
+    context : object
+        The context in which the function is executed.
+    dates : list
+        List of dates for which data is to be loaded.
+    path : str
+        Path to the directory containing the NetCDF files.
+    *args : tuple
+        Additional positional arguments.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    object
+        The loaded data.
+    """
+    return load_many("📁", context, dates, path, *args, **kwargs)

anemoi/datasets/create/sources/opendap.py

@@ -0,0 +1,43 @@
+# (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 typing import Any
+from typing import Dict
+from typing import List
+
+import earthkit.data as ekd
+
+from .legacy import legacy_source
+from .xarray import load_many
+
+
+@legacy_source(__file__)
+def execute(context: Dict[str, Any], dates: List[str], url: str, *args: Any, **kwargs: Any) -> ekd.FieldList:
+    """Execute the data loading process from an OpenDAP source.
+
+    Parameters
+    ----------
+    context : dict
+        The context in which the function is executed.
+    dates : list
+        List of dates for which data is to be loaded.
+    url : str
+        The URL of the OpenDAP source.
+    *args : tuple
+        Additional positional arguments.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    xarray.Dataset
+        The loaded dataset.
+    """
+    return load_many("🌐", context, dates, url, *args, **kwargs)

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

@@ -7,16 +7,29 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import datetime
 import glob
-import logging
+from typing import Any
+from typing import Generator
+from typing import List
+from typing import Tuple
 
 from earthkit.data.utils.patterns import Pattern
 
-LOG = logging.getLogger(__name__)
 
+def _expand(paths: List[str]) -> Generator[str, None, None]:
+    """Expand the given paths to include all matching file paths.
 
-def _expand(paths):
+    Parameters
+    ----------
+    paths : List[str]
+        List of paths to expand.
 
+    Returns
+    -------
+    Generator[str]
+        Expanded file paths.
+    """
     if not isinstance(paths, list):
         paths = [paths]
 
@@ -40,7 +53,25 @@ def _expand(paths):
             yield path
 
 
-def iterate_patterns(path, dates, **kwargs):
+def iterate_patterns(
+    path: str, dates: List[datetime.datetime], **kwargs: Any
+) -> Generator[Tuple[str, List[str]], None, None]:
+    """Iterate over patterns and expand them with given dates and additional keyword arguments.
+
+    Parameters
+    ----------
+    path : str
+        The pattern path to iterate over.
+    dates : List[datetime.datetime]
+        List of datetime objects to substitute in the pattern.
+    **kwargs : Any
+        Additional keyword arguments to substitute in the pattern.
+
+    Returns
+    -------
+    Generator[Tuple[str, List[str]]]
+        The expanded path and list of ISO formatted dates.
+    """
     given_paths = path if isinstance(path, list) else [path]
 
     dates = [d.isoformat() for d in dates]

anemoi/datasets/create/sources/recentre.py

@@ -0,0 +1,150 @@
+# (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 copy import deepcopy
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Union
+
+from anemoi.datasets.compute.recentre import recentre as _recentre
+
+from .legacy import legacy_source
+from .mars import mars
+
+
+def to_list(x: Union[list, tuple, str]) -> List:
+    """Converts the input to a list. If the input is a string, it splits it by '/'.
+
+    Parameters
+    ----------
+    x : Union[list, tuple, str]
+        The input to convert.
+
+    Returns
+    -------
+    list
+        The converted list.
+    """
+    if isinstance(x, (list, tuple)):
+        return x
+    if isinstance(x, str):
+        return x.split("/")
+    return [x]
+
+
+def normalise_number(number: Union[list, tuple, str]) -> List[int]:
+    """Normalises the input number to a list of integers.
+
+    Parameters
+    ----------
+    number : Union[list, tuple, str]
+        The number to normalise.
+
+    Returns
+    -------
+    list
+        The normalised list of integers.
+    """
+    number = to_list(number)
+
+    if len(number) > 4 and (number[1] == "to" and number[3] == "by"):
+        return list(range(int(number[0]), int(number[2]) + 1, int(number[4])))
+
+    if len(number) > 2 and number[1] == "to":
+        return list(range(int(number[0]), int(number[2]) + 1))
+
+    return number
+
+
+def normalise_request(request: Dict) -> Dict:
+    """Normalises the request dictionary by converting certain fields to lists.
+
+    Parameters
+    ----------
+    request : dict
+        The request dictionary to normalise.
+
+    Returns
+    -------
+    dict
+        The normalised request dictionary.
+    """
+    request = deepcopy(request)
+    if "number" in request:
+        request["number"] = normalise_number(request["number"])
+    if "time" in request:
+        request["time"] = to_list(request["time"])
+    request["param"] = to_list(request["param"])
+    return request
+
+
+def load_if_needed(context: Any, dates: Any, dict_or_dataset: Union[Dict, Any]) -> Any:
+    """Loads the dataset if the input is a dictionary, otherwise returns the input.
+
+    Parameters
+    ----------
+    context : Any
+        The context for loading the dataset.
+    dates : Any
+        The dates for loading the dataset.
+    dict_or_dataset : Union[dict, Any]
+        The input dictionary or dataset.
+
+    Returns
+    -------
+    Any
+        The loaded dataset or the original input.
+    """
+    if isinstance(dict_or_dataset, dict):
+        dict_or_dataset = normalise_request(dict_or_dataset)
+        dict_or_dataset = mars(context, dates, dict_or_dataset)
+    return dict_or_dataset
+
+
+@legacy_source(__file__)
+def recentre(
+    context: Any,
+    dates: Any,
+    members: Union[Dict, Any],
+    centre: Union[Dict, Any],
+    alpha: float = 1.0,
+    remapping: Dict = {},
+    patches: Dict = {},
+) -> Any:
+    """Recentres the members dataset using the centre dataset.
+
+    Parameters
+    ----------
+    context : Any
+        The context for recentering.
+    dates : Any
+        The dates for recentering.
+    members : Union[dict, Any]
+        The members dataset or request dictionary.
+    centre : Union[dict, Any]
+        The centre dataset or request dictionary.
+    alpha : float, optional
+        The alpha value for recentering. Defaults to 1.0.
+    remapping : dict, optional
+        The remapping dictionary. Defaults to {}.
+    patches : dict, optional
+        The patches dictionary. Defaults to {}.
+
+    Returns
+    -------
+    Any
+        The recentred dataset.
+    """
+    members = load_if_needed(context, dates, members)
+    centre = load_if_needed(context, dates, centre)
+    return _recentre(members=members, centre=centre, alpha=alpha)
+
+
+execute = recentre

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

@@ -7,14 +7,37 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+from datetime import datetime
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+
 from earthkit.data import from_source
 
 from anemoi.datasets.create.utils import to_datetime_list
 
-DEBUG = True
+from .legacy import legacy_source
+
 
+@legacy_source(__file__)
+def source(context: Optional[Any], dates: List[datetime], **kwargs: Any) -> Any:
+    """Generates a source based on the provided context, dates, and additional keyword arguments.
 
-def source(context, dates, **kwargs):
+    Parameters
+    ----------
+    context : Optional[Any]
+        The context in which the source is generated.
+    dates : List[datetime]
+        A list of datetime objects representing the dates.
+    **kwargs : Any
+        Additional keyword arguments for the source generation.
+
+    Returns
+    -------
+    Any
+        The generated source.
+    """
     name = kwargs.pop("name")
     context.trace("✅", f"from_source({name}, {dates}, {kwargs}")
     if kwargs["date"] == "$from_dates":
@@ -29,7 +52,7 @@ execute = source
 if __name__ == "__main__":
     import yaml
 
-    config = yaml.safe_load(
+    config: Dict[str, Any] = yaml.safe_load(
         """
       name: mars
       class: ea
@@ -42,9 +65,8 @@ if __name__ == "__main__":
       time: $from_dates
     """
     )
-    dates = yaml.safe_load("[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]")
+    dates: List[str] = yaml.safe_load("[2022-12-30 18:00, 2022-12-31 00:00, 2022-12-31 06:00, 2022-12-31 12:00]")
     dates = to_datetime_list(dates)
 
-    DEBUG = True
     for f in source(None, dates, **config):
         print(f, f.to_numpy().mean())

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

@@ -9,21 +9,50 @@
 
 import datetime
 from collections import defaultdict
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Tuple
 
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
 
-from anemoi.datasets.create.functions import assert_is_fieldlist
 from anemoi.datasets.create.utils import to_datetime_list
 
+from .legacy import legacy_source
 
-def _date_to_datetime(d):
+
+def _date_to_datetime(d: Any) -> Any:
+    """Converts a date string or a list/tuple of date strings to datetime objects.
+
+    Parameters
+    ----------
+    d : Any
+        A date string or a list/tuple of date strings.
+
+    Returns
+    -------
+    Any
+        A datetime object or a list/tuple of datetime objects.
+    """
     if isinstance(d, (list, tuple)):
         return [_date_to_datetime(x) for x in d]
     return datetime.datetime.fromisoformat(d)
 
 
-def normalise_time_delta(t):
+def normalise_time_delta(t: Any) -> datetime.timedelta:
+    """Normalizes a time delta string to a datetime.timedelta object.
+
+    Parameters
+    ----------
+    t : Any
+        A time delta string ending with 'h' or a datetime.timedelta object.
+
+    Returns
+    -------
+    datetime.timedelta
+        A normalized datetime.timedelta object.
+    """
     if isinstance(t, datetime.timedelta):
         assert t == datetime.timedelta(hours=t.hours), t
 
@@ -34,7 +63,19 @@ def normalise_time_delta(t):
     return t
 
 
-def group_by_field(ds):
+def group_by_field(ds: Any) -> Dict[Tuple, List[Any]]:
+    """Groups fields by their metadata excluding 'date', 'time', and 'step'.
+
+    Parameters
+    ----------
+    ds : Any
+        A dataset object.
+
+    Returns
+    -------
+    Dict[Tuple, List[Any]]
+        A dictionary where keys are tuples of metadata items and values are lists of fields.
+    """
     d = defaultdict(list)
     for field in ds.order_by("valid_datetime"):
         m = field.metadata(namespace="mars")
@@ -45,7 +86,24 @@ def group_by_field(ds):
     return d
 
 
-def tendencies(dates, time_increment, **kwargs):
+@legacy_source(__file__)
+def tendencies(dates: List[datetime.datetime], time_increment: Any, **kwargs: Any) -> Any:
+    """Computes tendencies for the given dates and time increment.
+
+    Parameters
+    ----------
+    dates : List[datetime.datetime]
+        A list of datetime objects.
+    time_increment : Any
+        A time increment string ending with 'h' or a datetime.timedelta object.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    Any
+        A dataset object with computed tendencies.
+    """
     print("✅", kwargs)
     time_increment = normalise_time_delta(time_increment)
 
@@ -53,7 +111,7 @@ def tendencies(dates, time_increment, **kwargs):
     all_dates = sorted(list(set(dates + shifted_dates)))
 
     # from .mars import execute as mars
-    from anemoi.datasets.create.functions.mars import execute as mars
+    from anemoi.datasets.create.mars import execute as mars
 
     ds = mars(dates=all_dates, **kwargs)
 
@@ -107,7 +165,6 @@ def tendencies(dates, time_increment, **kwargs):
     from earthkit.data import from_source
 
     ds = from_source("file", path)
-    assert_is_fieldlist(ds)
     # save a reference to the tmp file so it is deleted
     # only when the dataset is not used anymore
     ds._tmp = tmp

anemoi/datasets/create/sources/xarray.py

@@ -0,0 +1,92 @@
+# (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.
+
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+import earthkit.data as ekd
+
+from anemoi.datasets.create.typing import DateList
+
+from ..source import Source
+from .xarray_support import XarrayFieldList
+from .xarray_support import load_many
+from .xarray_support import load_one
+
+__all__ = ["load_many", "load_one", "XarrayFieldList"]
+
+
+class XarraySourceBase(Source):
+    """An Xarray base data source, intended to be subclassed."""
+
+    emoji = "✖️"  # For tracing
+
+    options: Optional[Dict[str, Any]] = None
+    flavour: Optional[Dict[str, Any]] = None
+    patch: Optional[Dict[str, Any]] = None
+
+    path_or_url: Optional[str] = None
+
+    def __init__(self, context: Any, path: str = None, url: str = None, *args: Any, **kwargs: Any):
+        """Initialise the source.
+
+        Parameters
+        ----------
+        context : Any
+            The context for the data source.
+        *args : Any
+            Additional positional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+        """
+        super().__init__(context, *args, **kwargs)
+
+        if path is not None and url is not None:
+            raise ValueError("Cannot specify both path and url")
+
+        if path is not None:
+            self.path_or_url = path
+        else:
+            self.path_or_url = url
+
+        self.args = args
+        self.kwargs = kwargs
+
+    def execute(self, dates: DateList) -> ekd.FieldList:
+        """Execute the data loading process for the given dates.
+
+        Parameters
+        ----------
+        dates : DateList
+            List of dates for which data needs to be loaded.
+
+        Returns
+        -------
+        ekd.FieldList
+            The loaded data fields.
+        """
+
+        # For now, just a simple wrapper around load_many
+        # TODO: move the implementation here
+
+        return load_many(
+            self.emoji,
+            self.context,
+            dates,
+            pattern=self.path_or_url,
+            options=self.options,
+            flavour=self.flavour,
+            patch=self.patch,
+            **self.kwargs,
+        )
+
+
+class XarraySource(XarraySourceBase):
+    pass

anemoi/datasets/create/sources/xarray_kerchunk.py

@@ -0,0 +1,36 @@
+# (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 . import source_registry
+from .xarray import XarraySourceBase
+
+
+@source_registry.register("xarray_kerchunk")
+class XarrayKerchunkSource(XarraySourceBase):
+    """An Xarray data source that uses the `kerchunk` engine."""
+
+    emoji = "🧱"
+
+    def __init__(self, context, json, *args, **kwargs: dict):
+        super().__init__(context, *args, **kwargs)
+
+        self.path_or_url = "reference://"
+
+        self.options = {
+            "engine": "zarr",
+            "backend_kwargs": {
+                "consolidated": False,
+                "storage_options": {
+                    "fo": json,
+                    "remote_protocol": "s3",
+                    "remote_options": {"anon": True},
+                },
+            },
+        }

anemoi/datasets/create/sources/xarray_support/README.md

@@ -0,0 +1 @@
+The code under this directory will be migrated to earthkit-data in the future