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

anemoi/datasets/create/sources/mars.py

@@ -16,9 +16,10 @@ 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 anemoi.datasets.create.sources.accumulate import IntervalsDatesProvider
 
-from .legacy import legacy_source
+from .legacy import LegacySource
 
 DEBUG = False
 
@@ -145,7 +146,7 @@ def _expand_mars_request(
 
     Parameters
     ----------
-    request : Dict[str, Any]
+    request : dict[str, Any]
         The input MARS request.
     date : datetime.datetime
         The date to be used in the request.
@@ -156,7 +157,7 @@ def _expand_mars_request(
 
     Returns
     -------
-    List[Dict[str, Any]]
+    List[dict[str, Any]]
         A list of expanded MARS requests.
     """
     requests = []
@@ -164,23 +165,26 @@ def _expand_mars_request(
     user_step = to_list(expand_to_by(request.get("step", [0])))
     user_time = None
     user_date = None
-
     if not request_already_using_valid_datetime:
-        user_time = request.get("time")
+        user_time = request.get("user_time")
         if user_time is not None:
             user_time = to_list(user_time)
             user_time = [_normalise_time(t) for t in user_time]
 
         user_date = request.get(date_key)
         if user_date is not None:
-            assert isinstance(user_date, str), user_date
+            if isinstance(user_date, int):
+                user_date = str(user_date)
+            elif isinstance(user_date, datetime.datetime):
+                user_date = user_date.strftime("%Y%m%d")
+            else:
+                raise ValueError(f"Invalid type for {user_date}")
             user_date = re.compile("^{}$".format(user_date.replace("-", "").replace("?", ".")))
 
     for step in user_step:
         r = request.copy()
 
         if not request_already_using_valid_datetime:
-
             if isinstance(step, str) and "-" in step:
                 assert step.count("-") == 1, step
 
@@ -190,30 +194,27 @@ def _expand_mars_request(
             base = date - datetime.timedelta(hours=hours)
             r.update(
                 {
-                    date_key: base.strftime("%Y%m%d"),
+                    "date": base.strftime("%Y%m%d"),
                     "time": base.strftime("%H%M"),
                     "step": step,
                 }
             )
-
         for pproc in ("grid", "rotation", "frame", "area", "bitmap", "resol"):
             if pproc in r:
                 if isinstance(r[pproc], (list, tuple)):
                     r[pproc] = "/".join(str(x) for x in r[pproc])
 
         if user_date is not None:
-            if not user_date.match(r[date_key]):
+            if not user_date.match(r["date"]):
                 continue
 
         if user_time is not None:
-            # It time is provided by the user, we only keep the requests that match the time
+            # If time is provided by the user, we only keep the requests that match the time
             if r["time"] not in user_time:
                 continue
 
         requests.append(r)
 
-    # assert requests, requests
-
     return requests
 
 
@@ -222,6 +223,7 @@ def factorise_requests(
     *requests: dict[str, Any],
     request_already_using_valid_datetime: bool = False,
     date_key: str = "date",
+    no_date_here: bool = False,
 ) -> Generator[dict[str, Any], None, None]:
     """Factorizes the requests based on the given dates.
 
@@ -229,33 +231,42 @@ def factorise_requests(
     ----------
     dates : List[datetime.datetime]
         The list of dates to be used in the requests.
-    requests : Dict[str, Any]
+    requests : List[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.
+    no_date_here : bool, optional
+        Flag indicating if there is no date in the "dates" list.
 
     Returns
     -------
-    Generator[Dict[str, Any], None, None]
+    Generator[dict[str, Any], None, None]
         Factorized requests.
     """
-    updates = []
-    for req in requests:
-        # req = normalise_request(req)
+    if isinstance(requests, tuple) and len(requests) == 1 and "requests" in requests[0]:
+        requests = requests[0]["requests"]
 
-        for d in dates:
-            updates += _expand_mars_request(
+    updates = []
+    for d in sorted(dates):
+        for req in requests:
+            if not no_date_here and (
+                ("date" in req)
+                and ("time" in req)
+                and d.strftime("%Y%m%d%H%M") != (str(req["date"]) + str(req["time"]).zfill(4))
+            ):
+                continue
+            new_req = _expand_mars_request(
                 req,
                 date=d,
                 request_already_using_valid_datetime=request_already_using_valid_datetime,
-                date_key=date_key,
+                date_key="user_date",
             )
+            updates += new_req
 
     if not updates:
         return
-
     compressed = Availability(updates)
     for r in compressed.iterate():
         for k, v in r.items():
@@ -269,12 +280,12 @@ def use_grib_paramid(r: dict[str, Any]) -> dict[str, Any]:
 
     Parameters
     ----------
-    r : Dict[str, Any]
+    r : dict[str, Any]
         The input request containing parameter short names.
 
     Returns
     -------
-    Dict[str, Any]
+    dict[str, Any]
         The request with parameter IDs.
     """
     from anemoi.utils.grib import shortname_to_paramid
@@ -358,135 +369,129 @@ 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]
+        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 isinstance(dates, IntervalsDatesProvider):
+            # When using accumulate source
+            requests_ = []
+            for request in requests:
+                for d, interval in dates.intervals:
+                    context.trace("🌧️", "interval:", interval)
+                    _, r, _ = dates._adjust_request_to_interval(interval, request)
+                    context.trace("🌧️", "  adjusted request =", r)
+                    requests_.append(r)
+            requests = requests_
+            context.trace("🌧️", f"Total requests: {len(requests)}")
+            requests = factorise_requests(
+                ["no_date_here"],
+                *requests,
+                request_already_using_valid_datetime=True,
+                date_key=date_key,
+                no_date_here=True,
+            )
 
-    """
-    )
-    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)
+        elif 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]}, {len(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)

anemoi/datasets/create/sources/recentre.py

@@ -12,7 +12,8 @@ from typing import Any
 
 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
 
 
@@ -105,43 +106,43 @@ def load_if_needed(context: Any, dates: Any, dict_or_dataset: dict | Any) -> Any
     return dict_or_dataset
 
 
-@legacy_source(__file__)
-def recentre(
-    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)
-
-
-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)