anemoi-datasets 0.5.26__py3-none-any.whl → 0.5.28__py3-none-any.whl

anemoi/datasets/create/sources/mars.py

@@ -9,25 +9,21 @@
 
 import datetime
 import re
+from collections.abc import Generator
 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
 from earthkit.data.utils.availability import Availability
 
-from anemoi.datasets.create.utils import to_datetime_list
+from anemoi.datasets.create.sources import source_registry
 
-from .legacy import legacy_source
+from .legacy import LegacySource
 
 DEBUG = False
 
 
-def to_list(x: Union[list, tuple, Any]) -> list:
+def to_list(x: list | tuple | Any) -> list:
     """Converts the input to a list if it is not already a list or tuple.
 
     Parameters
@@ -46,8 +42,8 @@ def to_list(x: Union[list, tuple, Any]) -> list:
 
 
 def _date_to_datetime(
-    d: Union[datetime.datetime, list, tuple, str],
-) -> Union[datetime.datetime, List[datetime.datetime]]:
+    d: datetime.datetime | list | tuple | str,
+) -> datetime.datetime | list[datetime.datetime]:
     """Converts the input date(s) to datetime objects.
 
     Parameters
@@ -67,7 +63,7 @@ def _date_to_datetime(
     return datetime.datetime.fromisoformat(d)
 
 
-def expand_to_by(x: Union[str, int, list]) -> Union[str, int, list]:
+def expand_to_by(x: str | int | list) -> str | int | list:
     """Expands a range expression to a list of values.
 
     Parameters
@@ -97,7 +93,7 @@ def expand_to_by(x: Union[str, int, list]) -> Union[str, int, list]:
     return x
 
 
-def normalise_time_delta(t: Union[datetime.timedelta, str]) -> datetime.timedelta:
+def normalise_time_delta(t: datetime.timedelta | str) -> datetime.timedelta:
     """Normalizes a time delta string to a datetime.timedelta object.
 
     Parameters
@@ -120,7 +116,7 @@ def normalise_time_delta(t: Union[datetime.timedelta, str]) -> datetime.timedelt
     return t
 
 
-def _normalise_time(t: Union[int, str]) -> str:
+def _normalise_time(t: int | str) -> str:
     """Normalizes a time value to a string in HHMM format.
 
     Parameters
@@ -136,15 +132,15 @@ def _normalise_time(t: Union[int, str]) -> str:
     t = int(t)
     if t < 100:
         t * 100
-    return "{:04d}".format(t)
+    return f"{t:04d}"
 
 
 def _expand_mars_request(
-    request: Dict[str, Any],
+    request: dict[str, Any],
     date: datetime.datetime,
     request_already_using_valid_datetime: bool = False,
     date_key: str = "date",
-) -> List[Dict[str, Any]]:
+) -> list[dict[str, Any]]:
     """Expands a MARS request with the given date and other parameters.
 
     Parameters
@@ -222,11 +218,11 @@ def _expand_mars_request(
 
 
 def factorise_requests(
-    dates: List[datetime.datetime],
-    *requests: Dict[str, Any],
+    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]:
+) -> Generator[dict[str, Any], None, None]:
     """Factorizes the requests based on the given dates.
 
     Parameters
@@ -268,7 +264,7 @@ def factorise_requests(
         yield r
 
 
-def use_grib_paramid(r: Dict[str, Any]) -> Dict[str, Any]:
+def use_grib_paramid(r: dict[str, Any]) -> dict[str, Any]:
     """Converts the parameter short names to GRIB parameter IDs.
 
     Parameters
@@ -362,135 +358,111 @@ MARS_KEYS = [
 ]
 
 
-@legacy_source(__file__)
-def mars(
-    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]
-
-    for r in requests:
-        param = r.get("param", [])
-        if not isinstance(param, (list, tuple)):
-            param = [param]
-        # check for "Norway bug" where yaml transforms 'no' into False, etc.
-        for p in param:
-            if p is False:
-                raise ValueError(
-                    "'param' cannot be 'False'. If you wrote 'param: no' or 'param: off' in yaml, you may want to use quotes?"
-                )
-            if p is None:
-                raise ValueError(
-                    "'param' cannot be 'None'. If you wrote 'param: no' in yaml, you may want to use quotes?"
-                )
-            if p is True:
-                raise ValueError(
-                    "'param' cannot be 'True'. If you wrote 'param: on' in yaml, you may want to use quotes?"
-                )
-
-    if len(dates) == 0:  # When using `repeated_dates`
-        assert len(requests) == 1, requests
-        assert "date" in requests[0], requests[0]
-        if isinstance(requests[0]["date"], datetime.date):
-            requests[0]["date"] = requests[0]["date"].strftime("%Y%m%d")
-    else:
-        requests = factorise_requests(
-            dates,
-            *requests,
-            request_already_using_valid_datetime=request_already_using_valid_datetime,
-            date_key=date_key,
-        )
-
-    requests = list(requests)
-
-    ds = from_source("empty")
-    context.trace("✅", f"{[str(d) for d in dates]}")
-    context.trace("✅", f"Will run {len(requests)} requests")
-    for r in requests:
-        r = {k: v for k, v in r.items() if v != ("-",)}
-        context.trace("✅", f"mars {r}")
-
-    for r in requests:
-        r = {k: v for k, v in r.items() if v != ("-",)}
-
-        if context.use_grib_paramid and "param" in r:
-            r = use_grib_paramid(r)
-
-        for k, v in r.items():
-            if k not in MARS_KEYS:
-                raise ValueError(
-                    f"⚠️ Unknown key {k}={v} in MARS request. Did you mean '{did_you_mean(k, MARS_KEYS)}' ?"
-                )
-        try:
-            if use_cdsapi_dataset:
-                ds = ds + from_source("cds", use_cdsapi_dataset, r)
-            else:
-                ds = ds + from_source("mars", **r)
-        except Exception as e:
-            if "File is empty:" not in str(e):
-                raise
-    return ds
-
-
-execute = mars
-
-
-if __name__ == "__main__":
-    import yaml
-
-    config = yaml.safe_load(
+@source_registry.register("mars")
+class MarsSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: list[datetime.datetime],
+        *requests: dict[str, Any],
+        request_already_using_valid_datetime: bool = False,
+        date_key: str = "date",
+        use_cdsapi_dataset: str | None = 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.
         """
-    - class: ea
-      expver: '0001'
-      grid: 20.0/20.0
-      levtype: sfc
-      param: [2t]
-      # param: [10u, 10v, 2d, 2t, lsm, msl, sdor, skt, slor, sp, tcw, z]
-      number: [0, 1]
-
-    # - class: ea
-    #   expver: '0001'
-    #   grid: 20.0/20.0
-    #   levtype: pl
-    #   param: [q]
-    #   levelist: [1000, 850]
 
-    """
-    )
-    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 = to_datetime_list(dates)
+        if not requests:
+            requests = [kwargs]
+
+        for r in requests:
+            param = r.get("param", [])
+            if not isinstance(param, (list, tuple)):
+                param = [param]
+            # check for "Norway bug" where yaml transforms 'no' into False, etc.
+            for p in param:
+                if p is False:
+                    raise ValueError(
+                        "'param' cannot be 'False'. If you wrote 'param: no' or 'param: off' in yaml, you may want to use quotes?"
+                    )
+                if p is None:
+                    raise ValueError(
+                        "'param' cannot be 'None'. If you wrote 'param: no' in yaml, you may want to use quotes?"
+                    )
+                if p is True:
+                    raise ValueError(
+                        "'param' cannot be 'True'. If you wrote 'param: on' in yaml, you may want to use quotes?"
+                    )
+
+        if len(dates) == 0:  # When using `repeated_dates`
+            assert len(requests) == 1, requests
+            assert "date" in requests[0], requests[0]
+            if isinstance(requests[0]["date"], datetime.date):
+                requests[0]["date"] = requests[0]["date"].strftime("%Y%m%d")
+        else:
+            requests = factorise_requests(
+                dates,
+                *requests,
+                request_already_using_valid_datetime=request_already_using_valid_datetime,
+                date_key=date_key,
+            )
 
-    DEBUG = True
-    for f in mars(None, dates, *config):
-        print(f, f.to_numpy().mean())
+        requests = list(requests)
+
+        ds = from_source("empty")
+        context.trace("✅", f"{[str(d) for d in dates]}")
+        context.trace("✅", f"Will run {len(requests)} requests")
+        for r in requests:
+            r = {k: v for k, v in r.items() if v != ("-",)}
+            context.trace("✅", f"mars {r}")
+
+        for r in requests:
+            r = {k: v for k, v in r.items() if v != ("-",)}
+
+            if context.use_grib_paramid and "param" in r:
+                r = use_grib_paramid(r)
+
+            for k, v in r.items():
+                if k not in MARS_KEYS:
+                    raise ValueError(
+                        f"⚠️ Unknown key {k}={v} in MARS request. Did you mean '{did_you_mean(k, MARS_KEYS)}' ?"
+                    )
+            try:
+                if use_cdsapi_dataset:
+                    ds = ds + from_source("cds", use_cdsapi_dataset, r)
+                else:
+                    ds = ds + from_source("mars", **r)
+            except Exception as e:
+                if "File is empty:" not in str(e):
+                    raise
+        return ds
+
+
+# TODO: make clearer the interface between sources that use mars.
+# Currently some sources use mars as a function rather than through the registry,
+# e.g. accumulations, accumulations2, hindcasts, recentre, tendencies
+mars = MarsSource._execute

anemoi/datasets/create/sources/netcdf.py

@@ -9,34 +9,37 @@
 
 
 from typing import Any
-from typing import List
 
 import earthkit.data as ekd
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 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)
+@source_registry.register("netcdf")
+class NetCDFSource(LegacySource):
+
+    @staticmethod
+    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

@@ -9,35 +9,37 @@
 
 
 from typing import Any
-from typing import Dict
-from typing import List
 
 import earthkit.data as ekd
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 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)
+@source_registry.register("opendap")
+class OpenDAPSource(LegacySource):
+
+    @staticmethod
+    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/sources/patterns.py

@@ -9,15 +9,13 @@
 
 import datetime
 import glob
+from collections.abc import Generator
 from typing import Any
-from typing import Generator
-from typing import List
-from typing import Tuple
 
 from earthkit.data.utils.patterns import Pattern
 
 
-def _expand(paths: List[str]) -> Generator[str, None, None]:
+def _expand(paths: list[str]) -> Generator[str, None, None]:
     """Expand the given paths to include all matching file paths.
 
     Parameters
@@ -54,8 +52,8 @@ def _expand(paths: List[str]) -> Generator[str, None, None]:
 
 
 def iterate_patterns(
-    path: str, dates: List[datetime.datetime], **kwargs: Any
-) -> Generator[Tuple[str, List[str]], None, None]:
+    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

anemoi/datasets/create/sources/recentre.py

@@ -9,17 +9,15 @@
 
 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 . import source_registry
+from .legacy import LegacySource
 from .mars import mars
 
 
-def to_list(x: Union[list, tuple, str]) -> List:
+def to_list(x: list | tuple | str) -> list:
     """Converts the input to a list. If the input is a string, it splits it by '/'.
 
     Parameters
@@ -39,7 +37,7 @@ def to_list(x: Union[list, tuple, str]) -> List:
     return [x]
 
 
-def normalise_number(number: Union[list, tuple, str]) -> List[int]:
+def normalise_number(number: list | tuple | str) -> list[int]:
     """Normalises the input number to a list of integers.
 
     Parameters
@@ -63,7 +61,7 @@ def normalise_number(number: Union[list, tuple, str]) -> List[int]:
     return number
 
 
-def normalise_request(request: Dict) -> Dict:
+def normalise_request(request: dict) -> dict:
     """Normalises the request dictionary by converting certain fields to lists.
 
     Parameters
@@ -85,7 +83,7 @@ def normalise_request(request: Dict) -> Dict:
     return request
 
 
-def load_if_needed(context: Any, dates: Any, dict_or_dataset: Union[Dict, Any]) -> Any:
+def load_if_needed(context: Any, dates: Any, dict_or_dataset: dict | Any) -> Any:
     """Loads the dataset if the input is a dictionary, otherwise returns the input.
 
     Parameters
@@ -108,43 +106,43 @@ def load_if_needed(context: Any, dates: Any, dict_or_dataset: Union[Dict, Any])
     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
+@source_registry.register("recentre")
+class RecentreSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: Any,
+        members: dict | Any,
+        centre: 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)