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

anemoi/datasets/create/input/data_sources.py

@@ -17,7 +17,7 @@ from ...dates.groups import GroupOfDates
 from .action import Action
 from .action import action_factory
 from .misc import _tidy
-from .result import Result
+from .result.field import Result
 
 LOG = logging.getLogger(__name__)
 
@@ -55,6 +55,7 @@ class DataSourcesAction(Action):
 
         self.sources = [action_factory(config, context, ["data_sources"] + [a_path]) for a_path, config in configs]
         self.input = action_factory(input, context, ["input"])
+        self.names = [a_path for a_path, config in configs]
 
     def select(self, group_of_dates: GroupOfDates) -> "DataSourcesResult":
         """Selects the data sources result for the given group of dates.

anemoi/datasets/create/input/misc.py

@@ -8,9 +8,6 @@
 # nor does it submit to any jurisdiction.
 
 import logging
-from collections.abc import Callable
-from functools import wraps
-from typing import Any
 
 from earthkit.data import FieldList
 from earthkit.data.core.fieldlist import MultiFieldList
@@ -18,74 +15,6 @@ from earthkit.data.core.fieldlist import MultiFieldList
 LOG = logging.getLogger(__name__)
 
 
-def parse_function_name(name: str) -> tuple[str, int | None]:
-    """Parses a function name to extract the base name and an optional time delta.
-
-    Parameters
-    ----------
-    name : str
-        The function name to parse.
-
-    Returns
-    -------
-    tuple of (str, int or None)
-        The base name and an optional time delta.
-    """
-    if name.endswith("h") and name[:-1].isdigit():
-
-        if "-" in name:
-            name, delta = name.split("-")
-            sign = -1
-
-        elif "+" in name:
-            name, delta = name.split("+")
-            sign = 1
-
-        else:
-            return name, None
-
-        assert delta[-1] == "h", (name, delta)
-        delta = sign * int(delta[:-1])
-        return name, delta
-
-    return name, None
-
-
-def assert_fieldlist(method: Callable[..., Any]) -> Callable[..., Any]:
-    """Decorator to assert that the result of a method is an instance of FieldList.
-
-    Parameters
-    ----------
-    method : Callable[..., Any]
-        The method to decorate.
-
-    Returns
-    -------
-    Callable[..., Any]
-        The decorated method.
-    """
-
-    @wraps(method)
-    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
-
-        result = method(self, *args, **kwargs)
-        assert isinstance(result, FieldList), type(result)
-        return result
-
-    return wrapper
-
-
-def assert_is_fieldlist(obj: object) -> None:
-    """Asserts that the given object is an instance of FieldList.
-
-    Parameters
-    ----------
-    obj : object
-        The object to check.
-    """
-    assert isinstance(obj, FieldList), type(obj)
-
-
 def _flatten(ds: MultiFieldList | FieldList) -> list:
     """Flattens a MultiFieldList or FieldList into a list of FieldList objects.
 

anemoi/datasets/create/input/repeated_dates.py

@@ -14,17 +14,9 @@ from collections.abc import Generator
 from typing import Any
 
 import numpy as np
-from anemoi.transform.fields import new_field_with_valid_datetime
-from anemoi.transform.fields import new_fieldlist_from_list
 from anemoi.utils.dates import as_datetime
 from anemoi.utils.dates import frequency_to_timedelta
 
-from .action import Action
-from .action import action_factory
-from .join import JoinResult
-from .result import Result
-from .trace import trace_select
-
 LOG = logging.getLogger(__name__)
 
 
@@ -276,109 +268,3 @@ class DateMapperConstant(DateMapper):
                 group_of_dates,
             )
         ]
-
-
-class DateMapperResult(Result):
-    """A Result implementation that updates the valid datetime of the datasource."""
-
-    def __init__(
-        self,
-        context: Any,
-        action_path: list[str],
-        group_of_dates: Any,
-        source_result: Any,
-        mapper: DateMapper,
-        original_group_of_dates: Any,
-    ) -> None:
-        """Initialize DateMapperResult.
-
-        Parameters
-        ----------
-        context : Any
-            The context.
-        action_path : list of str
-            The action path.
-        group_of_dates : Any
-            The group of dates.
-        source_result : Any
-            The source result.
-        mapper : DateMapper
-            The date mapper.
-        original_group_of_dates : Any
-            The original group of dates.
-        """
-        super().__init__(context, action_path, group_of_dates)
-
-        self.source_results: Any = source_result
-        self.mapper: DateMapper = mapper
-        self.original_group_of_dates: Any = original_group_of_dates
-
-    @property
-    def datasource(self) -> Any:
-        """Get the datasource with updated valid datetime."""
-        result: list = []
-
-        for field in self.source_results.datasource:
-            for date in self.original_group_of_dates:
-                result.append(new_field_with_valid_datetime(field, date))
-
-        if not result:
-            raise ValueError("repeated_dates: no input data found")
-
-        return new_fieldlist_from_list(result)
-
-
-class RepeatedDatesAction(Action):
-    """An Action implementation that selects and transforms a group of dates."""
-
-    def __init__(self, context: Any, action_path: list[str], source: Any, mode: str, **kwargs: Any) -> None:
-        """Initialize RepeatedDatesAction.
-
-        Args:
-            context (Any): The context.
-            action_path (List[str]): The action path.
-            source (Any): The data source.
-            mode (str): The mode for date mapping.
-            **kwargs (Any): Additional arguments.
-        """
-        super().__init__(context, action_path, source, mode, **kwargs)
-
-        self.source: Any = action_factory(source, context, action_path + ["source"])
-        self.mapper: DateMapper = DateMapper.from_mode(mode, self.source, kwargs)
-
-    @trace_select
-    def select(self, group_of_dates: Any) -> JoinResult:
-        """Select and transform the group of dates.
-
-        Args:
-            group_of_dates (Any): The group of dates to select.
-
-        Returns
-        -------
-        JoinResult
-            The result of the join operation.
-        """
-        results: list = []
-        for one_date_group, many_dates_group in self.mapper.transform(group_of_dates):
-            results.append(
-                DateMapperResult(
-                    self.context,
-                    self.action_path,
-                    one_date_group,
-                    self.source.select(one_date_group),
-                    self.mapper,
-                    many_dates_group,
-                )
-            )
-
-        return JoinResult(self.context, self.action_path, group_of_dates, results)
-
-    def __repr__(self) -> str:
-        """Get the string representation of the action.
-
-        Returns
-        -------
-        str
-            The string representation.
-        """
-        return f"MultiDateMatchAction({self.source}, {self.mapper})"

anemoi/datasets/create/input/result/__init__.py

@@ -0,0 +1,17 @@
+# (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.
+
+import logging
+from abc import ABC
+
+LOG = logging.getLogger(__name__)
+
+
+class Result(ABC):
+    pass

anemoi/datasets/create/input/{result.py → result/field.py}

@@ -22,9 +22,7 @@ from anemoi.utils.humanize import seconds_to_human
 from anemoi.utils.humanize import shorten_list
 from earthkit.data.core.order import build_remapping
 
-from .action import ActionContext
-from .trace import trace
-from .trace import trace_datasource
+from . import Result
 
 LOG = logging.getLogger(__name__)
 
@@ -278,40 +276,22 @@ def _data_request(data: Any) -> dict[str, Any]:
     return dict(param_level=params_levels, param_step=params_steps, area=area, grid=grid)
 
 
-class Result:
+class FieldResult(Result):
     """Class to represent the result of an action in the dataset creation process."""
 
     empty: bool = False
     _coords_already_built: bool = False
 
-    def __init__(self, context: ActionContext, action_path: list[str], dates: Any) -> None:
-        """Initialize a Result instance.
+    def __init__(self, context: Any, datasource: Any) -> None:
 
-        Parameters
-        ----------
-        context : ActionContext
-            The context in which the result exists.
-        action_path : list of str
-            The action path.
-        dates : Any
-            The dates associated with the result.
-        """
         from anemoi.datasets.dates.groups import GroupOfDates
 
-        assert isinstance(dates, GroupOfDates), dates
-
-        assert isinstance(context, ActionContext), type(context)
-        assert isinstance(action_path, list), action_path
-
         self.context: Any = context
-        self.group_of_dates: Any = dates
-        self.action_path: list[str] = action_path
-
-    @property
-    @trace_datasource
-    def datasource(self) -> Any:
-        """Retrieve the data source for the result."""
-        self._raise_not_implemented()
+        self.datasource = datasource
+        self.group_of_dates = context.argument
+        assert isinstance(
+            self.group_of_dates, GroupOfDates
+        ), f"Expected group_of_dates to be a GroupOfDates, got {type(self.group_of_dates)}: {self.group_of_dates}"
 
     @property
     def data_request(self) -> dict[str, Any]:
@@ -326,7 +306,7 @@ class Result:
         Any
             The data cube.
         """
-        trace("🧊", f"getting cube from {self.__class__.__name__}")
+
         ds: Any = self.datasource
 
         remapping: Any = self.context.remapping
@@ -519,66 +499,6 @@ class Result:
         print()
         exit(1)
 
-    def _repr(self, *args: Any, _indent_: str = "\n", **kwargs: Any) -> str:
-        """Return the string representation of the Result instance.
-
-        Parameters
-        ----------
-        args : Any
-            Additional positional arguments.
-        _indent_ : str
-            Indentation string.
-        kwargs : Any
-            Additional keyword arguments.
-
-        Returns
-        -------
-        str
-            The string representation.
-        """
-        more: str = ",".join([str(a)[:5000] for a in args])
-        more += ",".join([f"{k}={v}"[:5000] for k, v in kwargs.items()])
-
-        dates: str = " no-dates"
-        if self.group_of_dates is not None:
-            dates = f" {len(self.group_of_dates)} dates"
-            dates += " ("
-            dates += "/".join(d.strftime("%Y-%m-%dT%H:%M") for d in self.group_of_dates)
-            if len(dates) > 100:
-                dates = dates[:100] + "..."
-            dates += ")"
-
-        more = more[:5000]
-        txt: str = f"{self.__class__.__name__}:{dates}{_indent_}{more}"
-        if _indent_:
-            txt = txt.replace("\n", "\n  ")
-        return txt
-
-    def __repr__(self) -> str:
-        """Return the string representation of the Result instance."""
-        return self._repr()
-
-    def _raise_not_implemented(self) -> None:
-        """Raise a NotImplementedError indicating the method is not implemented."""
-        raise NotImplementedError(f"Not implemented in {self.__class__.__name__}")
-
-    def _trace_datasource(self, *args: Any, **kwargs: Any) -> str:
-        """Trace the data source for the result.
-
-        Parameters
-        ----------
-        args : Any
-            Additional positional arguments.
-        kwargs : Any
-            Additional keyword arguments.
-
-        Returns
-        -------
-        str
-            The trace string.
-        """
-        return f"{self.__class__.__name__}({self.group_of_dates})"
-
     def build_coords(self) -> None:
         """Build the coordinates for the result."""
         if self._coords_already_built:

anemoi/datasets/create/sources/accumulations.py

@@ -20,12 +20,13 @@ from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
 from numpy.typing import NDArray
 
-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
 from .mars import mars
 
 LOG = logging.getLogger(__name__)
+MISSING_VALUE = 1e-38
 
 
 def _member(field: Any) -> int:
@@ -168,6 +169,7 @@ class Accumulation:
         # are used to store the end step
 
         edition = template.metadata("edition")
+        assert np.all(self.values != MISSING_VALUE)
 
         if edition == 1 and self.endStep > 254:
             self.out.write(
@@ -176,6 +178,7 @@ class Accumulation:
                 stepType="instant",
                 step=self.endStep,
                 check_nans=True,
+                missing_value=MISSING_VALUE,
             )
         else:
             self.out.write(
@@ -185,6 +188,7 @@ class Accumulation:
                 startStep=self.startStep,
                 endStep=self.endStep,
                 check_nans=True,
+                missing_value=MISSING_VALUE,
             )
         self.values = None
         self.done = True
@@ -205,9 +209,6 @@ class Accumulation:
         if step not in self.steps:
             return
 
-        if not np.all(values >= 0):
-            warnings.warn(f"Negative values for {field}: {np.nanmin(values)} {np.nanmax(values)}")
-
         assert not self.done, (self.key, step)
         assert step not in self.seen, (self.key, step)
 
@@ -966,97 +967,76 @@ def _scda(request: dict[str, Any]) -> dict[str, Any]:
     return request
 
 
-@legacy_source(__file__)
-def accumulations(
-    context: Any, dates: list[datetime.datetime], use_cdsapi_dataset: str | None = None, **request: Any
-) -> Any:
-    """Computes accumulations based on the provided context, dates, and request parameters.
-
-    Parameters
-    ----------
-    context : Any
-        Context for the computation.
-    dates : List[datetime.datetime]
-        List of dates.
-    use_cdsapi_dataset : Optional[str], optional
-        CDSAPI dataset to use. Defaults to None.
-    **request : Any
-        Additional request parameters.
-
-    Returns
-    -------
-    Any
-        The computed accumulations.
-    """
-
-    if (
-        request.get("class") == "ea"
-        and request.get("stream", "oper") == "oper"
-        and request.get("accumulation_period") == 24
-    ):
-        from .accumulations2 import accumulations as accumulations2
-
-        LOG.warning(
-            "🧪️ Experimental features: Using accumulations2, because class=ea stream=oper and accumulation_period=24"
-        )
-        return accumulations2(context, dates, **request)
-
-    _to_list(request["param"])
-    class_ = request.get("class", "od")
-    stream = request.get("stream", "oper")
-
-    user_accumulation_period = request.pop("accumulation_period", 6)
-    accumulations_reset_frequency = request.pop("accumulations_reset_frequency", None)
-    user_date = request.pop("date", None)
-
-    # If `data_accumulation_period` is not set, this means that the accumulations are from the start
-    # of the forecast.
-
-    KWARGS = {
-        ("od", "oper"): dict(patch=_scda),
-        ("od", "elda"): dict(base_times=(6, 18)),
-        ("od", "enfo"): dict(base_times=(0, 6, 12, 18)),
-        ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
-        ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
-        ("rr", "oper"): dict(base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
-        ("l5", "oper"): dict(data_accumulation_period=1, base_times=(0,)),
-    }
-
-    kwargs = KWARGS.get((class_, stream), {})
-
-    context.trace("🌧️", f"accumulations {request} {user_accumulation_period} {kwargs}")
-
-    return _compute_accumulations(
-        context,
-        dates,
-        request,
-        user_accumulation_period=user_accumulation_period,
-        accumulations_reset_frequency=accumulations_reset_frequency,
-        use_cdsapi_dataset=use_cdsapi_dataset,
-        user_date=user_date,
-        **kwargs,
-    )
+@source_registry.register("accumulations")
+class AccumulationsSource(LegacySource):
 
+    @staticmethod
+    def _execute(
+        context: Any, dates: list[datetime.datetime], use_cdsapi_dataset: str | None = None, **request: Any
+    ) -> Any:
+        """Computes accumulations based on the provided context, dates, and request parameters.
 
-execute = accumulations
-
-if __name__ == "__main__":
-    import yaml
+        Parameters
+        ----------
+        context : Any
+            Context for the computation.
+        dates : List[datetime.datetime]
+            List of dates.
+        use_cdsapi_dataset : Optional[str], optional
+            CDSAPI dataset to use. Defaults to None.
+        **request : Any
+            Additional request parameters.
 
-    config = yaml.safe_load(
+        Returns
+        -------
+        Any
+            The computed accumulations.
         """
-      class: ea
-      expver: '0001'
-      grid: 20./20.
-      levtype: sfc
-#      number: [0, 1]
-#      stream: enda
-      param: [cp, tp]
-#      accumulation_period: 6h
-    """
-    )
-    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)
 
-    for f in accumulations(None, dates, **config):
-        print(f, f.to_numpy().mean())
+        if (
+            request.get("class") == "ea"
+            and request.get("stream", "oper") == "oper"
+            and request.get("accumulation_period") == 24
+        ):
+            from .accumulations2 import Accumulations2Source
+
+            LOG.warning(
+                "🧪️ Experimental features: Using accumulations2, because class=ea stream=oper and accumulation_period=24"
+            )
+            return Accumulations2Source._execute(context, dates, **request)
+
+        _to_list(request["param"])
+        class_ = request.get("class", "od")
+        stream = request.get("stream", "oper")
+
+        user_accumulation_period = request.pop("accumulation_period", 6)
+        accumulations_reset_frequency = request.pop("accumulations_reset_frequency", None)
+        user_date = request.pop("date", None)
+
+        # If `data_accumulation_period` is not set, this means that the accumulations are from the start
+        # of the forecast.
+
+        KWARGS = {
+            ("od", "oper"): dict(patch=_scda),
+            ("od", "elda"): dict(base_times=(6, 18)),
+            ("od", "enfo"): dict(base_times=(0, 6, 12, 18)),
+            ("ea", "oper"): dict(data_accumulation_period=1, base_times=(6, 18)),
+            ("ea", "enda"): dict(data_accumulation_period=3, base_times=(6, 18)),
+            ("rr", "oper"): dict(base_times=(0, 3, 6, 9, 12, 15, 18, 21)),
+            ("l5", "oper"): dict(data_accumulation_period=1, base_times=(0,)),
+        }
+
+        kwargs = KWARGS.get((class_, stream), {})
+
+        context.trace("🌧️", f"accumulations {request} {user_accumulation_period} {kwargs}")
+
+        return _compute_accumulations(
+            context,
+            dates,
+            request,
+            user_accumulation_period=user_accumulation_period,
+            accumulations_reset_frequency=accumulations_reset_frequency,
+            use_cdsapi_dataset=use_cdsapi_dataset,
+            user_date=user_date,
+            **kwargs,
+        )

anemoi/datasets/create/sources/accumulations2.py

@@ -18,10 +18,10 @@ import numpy as np
 from earthkit.data.core.temporary import temp_file
 from earthkit.data.readers.grib.output import new_grib_output
 
+from anemoi.datasets.create.sources import source_registry
 from anemoi.datasets.create.sources.mars import mars
-from anemoi.datasets.create.utils import to_datetime_list
 
-from .legacy import legacy_source
+from .legacy import LegacySource
 
 LOG = logging.getLogger(__name__)
 
@@ -599,49 +599,20 @@ def _scda(request: dict[str, Any]) -> dict[str, Any]:
     return request
 
 
-@legacy_source(__file__)
-def accumulations(context, dates, **request):
-    _to_list(request["param"])
-    user_accumulation_period = request.pop("accumulation_period", 6)
-    user_accumulation_period = datetime.timedelta(hours=user_accumulation_period)
+@source_registry.register("accumulations2")
+class Accumulations2Source(LegacySource):
 
-    context.trace("🌧️", f"accumulations {request} {user_accumulation_period}")
+    @staticmethod
+    def _execute(context, dates, **request):
+        _to_list(request["param"])
+        user_accumulation_period = request.pop("accumulation_period", 6)
+        user_accumulation_period = datetime.timedelta(hours=user_accumulation_period)
 
-    return _compute_accumulations(
-        context,
-        dates,
-        request,
-        user_accumulation_period=user_accumulation_period,
-    )
-
-
-execute = accumulations
-
-if __name__ == "__main__":
-    import yaml
-
-    config = yaml.safe_load(
-        """
-      class: ea
-      expver: '0001'
-      grid: 20./20.
-      levtype: sfc
-#      number: [0, 1]
-#      stream: enda
-      param: [cp, tp]
-#      accumulation_period: 6h
-      accumulation_period: 2
-    """
-    )
-    dates = yaml.safe_load("[2022-12-31 00:00, 2022-12-31 06:00]")
-    # 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)
-
-    class Context:
-        use_grib_paramid = True
+        context.trace("🌧️", f"accumulations {request} {user_accumulation_period}")
 
-        def trace(self, *args):
-            print(*args)
-
-    for f in accumulations(Context, dates, **config):
-        print(f, f.to_numpy().mean())
+        return _compute_accumulations(
+            context,
+            dates,
+            request,
+            user_accumulation_period=user_accumulation_period,
+        )