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

anemoi/datasets/create/sources/hindcasts.py

@@ -14,7 +14,8 @@ from earthkit.data.core.fieldlist import MultiFieldList
 
 from anemoi.datasets.create.sources.mars import mars
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOGGER = logging.getLogger(__name__)
 
@@ -37,57 +38,57 @@ def _to_list(x: list | tuple | Any) -> list[Any]:
     return [x]
 
 
-@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
-    assert isinstance(provider, HindcastsDates)
-
-    context.trace("H️", f"hindcasts {len(dates)=}")
-
-    request["param"] = _to_list(request["param"])
-    request["step"] = _to_list(request.get("step", 0))
-    request["step"] = [int(_) for _ in request["step"]]
-
-    context.trace("H️", f"hindcast {request}")
-
-    requests = []
-    for d in dates:
-        r = request.copy()
-        hindcast = provider.mapping[d]
-        r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
-        r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
-        r["time"] = hindcast.refdate.strftime("%H")
-        r["step"] = hindcast.step
-        requests.append(r)
-
-    if len(requests) == 0:
-        return MultiFieldList([])
-
-    return mars(
-        context,
-        dates,
-        *requests,
-        date_key="hdate",
-        request_already_using_valid_datetime=True,
-    )
-
-
-execute = hindcasts
+@source_registry.register("hindcasts")
+class HindcastsSource(LegacySource):
+
+    @staticmethod
+    def _execute(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
+        assert isinstance(provider, HindcastsDates)
+
+        context.trace("H️", f"hindcasts {len(dates)=}")
+
+        request["param"] = _to_list(request["param"])
+        request["step"] = _to_list(request.get("step", 0))
+        request["step"] = [int(_) for _ in request["step"]]
+
+        context.trace("H️", f"hindcast {request}")
+
+        requests = []
+        for d in dates:
+            r = request.copy()
+            hindcast = provider.mapping[d]
+            r["hdate"] = hindcast.hdate.strftime("%Y-%m-%d")
+            r["date"] = hindcast.refdate.strftime("%Y-%m-%d")
+            r["time"] = hindcast.refdate.strftime("%H")
+            r["step"] = hindcast.step
+            requests.append(r)
+
+        if len(requests) == 0:
+            return MultiFieldList([])
+
+        return mars(
+            context,
+            dates,
+            *requests,
+            date_key="hdate",
+            request_already_using_valid_datetime=True,
+        )

anemoi/datasets/create/sources/legacy.py

@@ -8,16 +8,13 @@
 # nor does it submit to any jurisdiction.
 
 
-import inspect
 import logging
-import os
-from collections.abc import Callable
+from abc import abstractmethod
 from typing import Any
 
-from anemoi.datasets.create.input.template import resolve
+from anemoi.datasets.create.input.context import Context
 
 from ..source import Source
-from . import source_registry
 
 LOG = logging.getLogger(__name__)
 
@@ -27,7 +24,7 @@ class LegacySource(Source):
 
     Parameters
     ----------
-    context : Any
+    context : Context
         The context in which the source is created.
     *args : tuple
         Positional arguments.
@@ -35,64 +32,15 @@ class LegacySource(Source):
         Keyword arguments.
     """
 
-    def __init__(self, context: Any, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, context: Context, *args: Any, **kwargs: Any) -> None:
         super().__init__(context, *args, **kwargs)
         self.args = args
         self.kwargs = kwargs
 
+    @staticmethod
+    @abstractmethod
+    def _execute(context, *args, **kwargs):
+        pass
 
-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."""
-
-            args, kwargs = resolve(self.context, (self.args, self.kwargs))
-
-            try:
-                return execute(self.context, dates, *args, **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: {args=}, {kwargs=}")
-                raise
-
-        klass = type(
-            name,
-            (LegacySource,),
-            {
-                "execute": execute_wrapper,
-                "_source": source,
-            },
-        )
-
-        source_registry.register(self.name)(klass)
-
-        return execute
+    def execute(self, dates: Any) -> Any:
+        return self._execute(self.context, dates, *self.args, **self.kwargs)

anemoi/datasets/create/sources/mars.py

@@ -16,9 +16,9 @@ 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
 
@@ -358,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: 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.
-    """
-
-    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

@@ -12,30 +12,34 @@ from typing import Any
 
 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

@@ -12,30 +12,34 @@ from typing import Any
 
 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)