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

anemoi/datasets/create/sources/constants.py

@@ -8,46 +8,45 @@
 # nor does it submit to any jurisdiction.
 
 from typing import Any
-from typing import Dict
-from typing import List
 
 from earthkit.data import from_source
 
-from .legacy import legacy_source
-
-
-@legacy_source(__file__)
-def constants(context: Any, dates: List[str], template: Dict[str, Any], param: str) -> Any:
-    """Deprecated function to retrieve constants data.
-
-    Parameters
-    ----------
-    context : Any
-        The context object for tracing.
-    dates : list of str
-        List of dates for which data is required.
-    template : dict of str to Any
-        Template dictionary for the data source.
-    param : str
-        Parameter to retrieve.
-
-    Returns
-    -------
-    Any
-        Data retrieved from the source.
-    """
-    from warnings import warn
-
-    warn(
-        "The source `constants` is deprecated, use `forcings` instead.",
-        DeprecationWarning,
-        stacklevel=2,
-    )
-    context.trace("✅", f"from_source(constants, {template}, {param}")
-    if len(template) == 0:
-        raise ValueError("Forcings template is empty.")
-
-    return from_source("forcings", source_or_dataset=template, date=dates, param=param)
-
-
-execute: Any = constants
+from . import source_registry
+from .legacy import LegacySource
+
+
+@source_registry.register("constants")
+class ConstantsSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], template: dict[str, Any], param: str) -> Any:
+        """Deprecated function to retrieve constants data.
+
+        Parameters
+        ----------
+        context : Any
+            The context object for tracing.
+        dates : list of str
+            List of dates for which data is required.
+        template : dict of str to Any
+            Template dictionary for the data source.
+        param : str
+            Parameter to retrieve.
+
+        Returns
+        -------
+        Any
+            Data retrieved from the source.
+        """
+        from warnings import warn
+
+        warn(
+            "The source `constants` is deprecated, use `forcings` instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        context.trace("✅", f"from_source(constants, {template}, {param}")
+        if len(template) == 0:
+            raise ValueError("Forcings template is empty.")
+
+        return from_source("forcings", source_or_dataset=template, date=list(dates), param=param)

anemoi/datasets/create/sources/empty.py

@@ -9,29 +9,32 @@
 
 
 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
 
 
-@legacy_source(__file__)
-def execute(context: Any, dates: List[str], **kwargs: Any) -> ekd.FieldList:
-    """Executes the loading of an empty data source.
+@source_registry.register("empty")
+class EmptySource(LegacySource):
 
-    Parameters
-    ----------
-    context : object
-        The context in which the function is executed.
-    dates : list
-        List of dates for which data is to be loaded.
-    **kwargs : dict
-        Additional keyword arguments.
+    @staticmethod
+    def _execute(context: Any, dates: list[str], **kwargs: Any) -> ekd.FieldList:
+        """Executes the loading of an empty data source.
 
-    Returns
-    -------
-    ekd.FieldList
-        Loaded empty data source.
-    """
-    return ekd.from_source("empty")
+        Parameters
+        ----------
+        context : object
+            The context in which the function is executed.
+        dates : list
+            List of dates for which data is to be loaded.
+        **kwargs : dict
+            Additional keyword arguments.
+
+        Returns
+        -------
+        ekd.FieldList
+            Loaded empty data source.
+        """
+        return ekd.from_source("empty")

anemoi/datasets/create/sources/fdb.py

@@ -0,0 +1,133 @@
+# (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 datetime import datetime
+from typing import Any
+
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_grid
+from anemoi.transform.fields import new_fieldlist_from_list
+from anemoi.transform.flavour import RuleBasedFlavour
+from anemoi.transform.grids import grid_registry
+
+from anemoi.datasets.create.typing import DateList
+
+from ..source import Source
+from . import source_registry
+
+
+@source_registry.register("fdb")
+class FdbSource(Source):
+    """FDB data source."""
+
+    emoji = "💽"
+
+    def __init__(
+        self,
+        context,
+        fdb_request: dict[str, Any],
+        fdb_config: dict | None = None,
+        fdb_userconfig: dict | None = None,
+        flavour: dict[str, Any] | None = None,
+        grid_definition: str | None = None,
+        **kwargs: dict[str, Any],
+    ):
+        """Initialise the FDB input.
+
+        Parameters
+        ----------
+        context : dict
+            The context.
+        fdb_request: dict
+            The FDB request parameters.
+        fdb_config : dict, optional
+            The FDB config to use.
+        fdb_userconfig : dict, optional
+            The FDB userconfig to use.
+        flavour: dict, optional
+            The flavour configuration, see `anemoi.transform.flavour.RuleBasedFlavour`.
+        grid_definition : str, optional
+            The grid definition to use, see `anemoi.transform.grids.grid_registry`.
+        kwargs : dict, optional
+            Additional keyword arguments.
+        """
+        super().__init__(context)
+        self.request = fdb_request.copy()
+        self.configs = {"config": fdb_config, "userconfig": fdb_userconfig}
+
+        self.flavour = RuleBasedFlavour(flavour) if flavour else None
+        if grid_definition is not None:
+            self.grid = grid_registry.from_config(grid_definition)
+        else:
+            self.grid = None
+
+        if "step" not in self.request:
+            self.request["step"] = 0
+
+        self.request["param"] = _shortname_to_paramid(fdb_request["param"], kwargs.pop("param_id_map", None))
+
+        # temporary workarounds for FDB use at MeteoSwiss (adoption is ongoing)
+        # thus not documented
+        self.offset_from_date = kwargs.pop("offset_from_date", None)
+
+    def execute(self, dates: DateList) -> ekd.FieldList:
+        """Execute the FDB source.
+
+        Parameters
+        ----------
+        dates : DateList
+            The input dates.
+
+        Returns
+        -------
+        ekd.FieldList
+            The output data.
+        """
+
+        requests = []
+        for date in dates:
+            time_request = _time_request_keys(date, self.offset_from_date)
+            requests.append(self.request | time_request)
+
+        # in some cases (e.g. repeated_dates 'constant' mode), we might have a fully
+        # defined request already and an empty dates list
+        requests = requests or [self.request]
+
+        fl = ekd.from_source("empty")
+        for request in requests:
+            fl += ekd.from_source("fdb", request, **self.configs, read_all=True)
+
+        if self.grid is not None:
+            fl = new_fieldlist_from_list([new_field_from_grid(f, self.grid) for f in fl])
+
+        if self.flavour:
+            fl = self.flavour.map(fl)
+
+        return fl
+
+
+def _time_request_keys(dt: datetime, offset_from_date: bool | None = None) -> str:
+    """Defines the time-related keys for the FDB request."""
+    out = {}
+    out["date"] = dt.strftime("%Y%m%d")
+    if offset_from_date:
+        out["time"] = "0000"
+        out["step"] = int((dt - dt.replace(hour=0, minute=0)).total_seconds() // 3600)
+    else:
+        out["time"] = dt.strftime("%H%M")
+    return out
+
+
+def _shortname_to_paramid(shortname: list[str], param_id_map: dict[str, int] | None = None) -> list[int]:
+    from anemoi.datasets.create.sources.mars import use_grib_paramid
+
+    """Convert a shortname to a parameter ID."""
+    if param_id_map is None:
+        return use_grib_paramid(shortname)
+    return [param_id_map[s] for s in shortname]

anemoi/datasets/create/sources/forcings.py

@@ -8,35 +8,35 @@
 # nor does it submit to any jurisdiction.
 
 from typing import Any
-from typing import List
 
 from earthkit.data import from_source
 
-from .legacy import legacy_source
-
-
-@legacy_source(__file__)
-def forcings(context: Any, dates: List[str], template: str, param: str) -> Any:
-    """Loads forcing data from a specified source.
-
-    Parameters
-    ----------
-    context : object
-        The context in which the function is executed.
-    dates : list
-        List of dates for which data is to be loaded.
-    template : FieldList
-        Template for the data source.
-    param : str
-        Parameter for the data source.
-
-    Returns
-    -------
-    object
-        Loaded forcing data.
-    """
-    context.trace("✅", f"from_source(forcings, {template}, {param}")
-    return from_source("forcings", source_or_dataset=template, date=dates, param=param)
-
-
-execute = forcings
+from . import source_registry
+from .legacy import LegacySource
+
+
+@source_registry.register("forcings")
+class ForcingsSource(LegacySource):
+
+    @staticmethod
+    def _execute(context: Any, dates: list[str], template: str, param: str) -> Any:
+        """Loads forcing data from a specified source.
+
+        Parameters
+        ----------
+        context : object
+            The context in which the function is executed.
+        dates : list
+            List of dates for which data is to be loaded.
+        template : FieldList
+            Template for the data source.
+        param : str
+            Parameter for the data source.
+
+        Returns
+        -------
+        object
+            Loaded forcing data.
+        """
+        context.trace("✅", f"from_source(forcings, {template}, {param}")
+        return from_source("forcings", source_or_dataset=template, date=list(dates), param=param)

anemoi/datasets/create/sources/grib.py

@@ -11,10 +11,6 @@
 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 anemoi.transform.fields import new_field_from_grid
@@ -24,12 +20,13 @@ from anemoi.transform.grids import grid_registry
 from earthkit.data import from_source
 from earthkit.data.utils.patterns import Pattern
 
-from .legacy import legacy_source
+from . import source_registry
+from .legacy import LegacySource
 
 LOG = logging.getLogger(__name__)
 
 
-def check(ds: Any, paths: List[str], **kwargs: Any) -> None:
+def check(ds: Any, paths: list[str], **kwargs: Any) -> None:
     """Check if the dataset matches the expected number of fields.
 
     Parameters
@@ -51,11 +48,19 @@ def check(ds: Any, paths: List[str], **kwargs: Any) -> None:
         if isinstance(v, (tuple, list)):
             count *= len(v)
 
+    # in the case of static data (e.g repeated dates) dates might be empty
+    if len(ds) != count and kwargs.get("dates", []) == []:
+        LOG.warning(
+            f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})"
+            f" Received empty dates - assuming this is static data."
+        )
+        return
+
     if len(ds) != count:
         raise ValueError(f"Expected {count} fields, got {len(ds)} (kwargs={kwargs}, paths={paths})")
 
 
-def _expand(paths: List[str]) -> Any:
+def _expand(paths: list[str]) -> Any:
     """Expand the given paths using glob.
 
     Parameters
@@ -77,74 +82,85 @@ def _expand(paths: List[str]) -> Any:
             yield path
 
 
-@legacy_source(__file__)
-def execute(
-    context: Any,
-    dates: List[Any],
-    path: Union[str, List[str]],
-    flavour: Optional[Union[str, Dict[str, Any]]] = None,
-    grid_definition: Optional[Dict[str, Any]] = None,
-    *args: Any,
-    **kwargs: Any,
-) -> ekd.FieldList:
-    """Executes the function to load data from GRIB files.
-
-    Parameters
-    ----------
-    context : Any
-        The context in which the function is executed.
-    dates : list of Any
-        List of dates.
-    path : str or list of str
-        Path or list of paths to the GRIB files.
-    flavour : str or dict of str to Any, optional
-        Flavour information, by default None.
-    grid_definition : dict of str to Any, optional
-        Grid definition configuration to create a Grid object, by default None.
-    *args : Any
-        Additional positional arguments.
-    **kwargs : Any
-        Additional keyword arguments.
-
-    Returns
-    -------
-    Any
-        The loaded dataset.
-    """
-    given_paths = path if isinstance(path, list) else [path]
-    if flavour is not None:
-        flavour = RuleBasedFlavour(flavour)
-
-    if grid_definition is not None:
-        grid = grid_registry.from_config(grid_definition)
-    else:
-        grid = None
-
-    ds = from_source("empty")
-    dates = [d.isoformat() for d in dates]
-
-    for path in given_paths:
-        paths = Pattern(path).substitute(*args, date=dates, allow_extra=True, **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)
-            if flavour is not None:
-                s = flavour.map(s)
-            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 grid is not None:
-        ds = new_fieldlist_from_list([new_field_from_grid(f, grid) for f in ds])
-
-    if len(ds) == 0:
-        LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")
-
-    return ds
+@source_registry.register("grib")
+class GribSource(LegacySource):
+
+    @staticmethod
+    def _execute(
+        context: Any,
+        dates: list[Any],
+        path: str | list[str],
+        flavour: str | dict[str, Any] | None = None,
+        grid_definition: dict[str, Any] | None = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> ekd.FieldList:
+        """Executes the function to load data from GRIB files.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the function is executed.
+        dates : list of Any
+            List of dates.
+        path : str or list of str
+            Path or list of paths to the GRIB files.
+        flavour : str or dict of str to Any, optional
+            Flavour information, by default None.
+        grid_definition : dict of str to Any, optional
+            Grid definition configuration to create a Grid object, by default None.
+        *args : Any
+            Additional positional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Any
+            The loaded dataset.
+        """
+        given_paths = path if isinstance(path, list) else [path]
+        if flavour is not None:
+            flavour = RuleBasedFlavour(flavour)
+
+        if grid_definition is not None:
+            grid = grid_registry.from_config(grid_definition)
+        else:
+            grid = None
+
+        ds = from_source("empty")
+        dates = [d.isoformat() for d in dates]
+
+        for path in given_paths:
+
+            # do not substitute if not needed
+            if "{" not in path:
+                paths = [path]
+            else:
+                paths = Pattern(path).substitute(*args, date=dates, allow_extra=True, **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)
+                if flavour is not None:
+                    s = flavour.map(s)
+                sel_kwargs = kwargs.copy()
+                if dates != []:
+                    sel_kwargs["valid_datetime"] = dates
+                s = s.sel(**sel_kwargs)
+                ds = ds + s
+
+        if kwargs and not context.partial_ok:
+            check(ds, given_paths, valid_datetime=dates, **kwargs)
+
+        if grid is not None:
+            ds = new_fieldlist_from_list([new_field_from_grid(f, grid) for f in ds])
+
+        if len(ds) == 0:
+            LOG.warning(f"No fields found for {dates} in {given_paths} (kwargs={kwargs})")
+
+        return ds