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

anemoi/datasets/create/input/action.py

@@ -7,253 +7,311 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-import json
 import logging
-from copy import deepcopy
-from typing import Any
-from typing import Dict
-from typing import List
 
-from earthkit.data.core.order import build_remapping
-
-from ...dates.groups import GroupOfDates
-from .context import Context
-from .template import substitute
+from anemoi.datasets.dates import DatesProvider
 
 LOG = logging.getLogger(__name__)
 
 
 class Action:
-    """Represents an action to be performed within a given context.
-
-    Attributes
-    ----------
-    context : ActionContext
-        The context in which the action exists.
-    kwargs : Dict[str, Any]
-        Additional keyword arguments.
-    args : Any
-        Additional positional arguments.
-    action_path : List[str]
-        The action path.
+    """An "Action" represents a single operation described in the yaml configuration, e.g. a source, a filter,
+    pipe, join, etc.
+
+    See :ref:`operations` for more details.
+
     """
 
-    def __init__(
-        self, context: "ActionContext", action_path: List[str], /, *args: Any, **kwargs: Dict[str, Any]
-    ) -> None:
-        """Initialize an Action instance.
-
-        Parameters
-        ----------
-        context : ActionContext
-            The context in which the action exists.
-        action_path : List[str]
-            The action path.
-        args : Any
-            Additional positional arguments.
-        kwargs : Dict[str, Any]
-            Additional keyword arguments.
-        """
-        if "args" in kwargs and "kwargs" in kwargs:
-            """We have:
-               args = []
-               kwargs = {args: [...], kwargs: {...}}
-            move the content of kwargs to args and kwargs.
-            """
-            assert len(kwargs) == 2, (args, kwargs)
-            assert not args, (args, kwargs)
-            args = kwargs.pop("args")
-            kwargs = kwargs.pop("kwargs")
-
-        assert isinstance(context, ActionContext), type(context)
-        self.context = context
-        self.kwargs = kwargs
-        self.args = args
-        self.action_path = action_path
-
-    @classmethod
-    def _short_str(cls, x: str) -> str:
-        """Shorten the string representation if it exceeds 1000 characters.
-
-        Parameters
-        ----------
-        x : str
-            The string to shorten.
-
-        Returns
-        -------
-        str
-            The shortened string.
-        """
-        x = str(x)
-        if len(x) < 1000:
-            return x
-        return x[:1000] + "..."
-
-    def _repr(self, *args: Any, _indent_: str = "\n", _inline_: str = "", **kwargs: Any) -> str:
-        """Generate a string representation of the Action instance.
-
-        Parameters
-        ----------
-        args : Any
-            Additional positional arguments.
-        _indent_ : str, optional
-            The indentation string, by default "\n".
-        _inline_ : str, optional
-            The inline string, by default "".
-        kwargs : Any
-            Additional keyword arguments.
-
-        Returns
-        -------
-        str
-            The string representation.
-        """
-        more = ",".join([str(a)[:5000] for a in args])
-        more += ",".join([f"{k}={v}"[:5000] for k, v in kwargs.items()])
-
-        more = more[:5000]
-        txt = f"{self.__class__.__name__}: {_inline_}{_indent_}{more}"
-        if _indent_:
-            txt = txt.replace("\n", "\n  ")
-        return txt
-
-    def __repr__(self) -> str:
-        """Return the string representation of the Action instance.
-
-        Returns
-        -------
-        str
-            The string representation.
-        """
-        return self._repr()
-
-    def select(self, dates: object, **kwargs: Any) -> None:
-        """Select dates for the action.
-
-        Parameters
-        ----------
-        dates : object
-            The dates to select.
-        kwargs : Any
-            Additional keyword arguments.
-        """
-        self._raise_not_implemented()
-
-    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_select(self, group_of_dates: GroupOfDates) -> str:
-        """Trace the selection of a group of dates.
-
-        Parameters
-        ----------
-        group_of_dates : GroupOfDates
-            The group of dates to trace.
-
-        Returns
-        -------
-        str
-            The trace string.
-        """
-        return f"{self.__class__.__name__}({group_of_dates})"
-
-
-class ActionContext(Context):
-    """Represents the context in which an action is performed.
-
-    Attributes
-    ----------
-    order_by : str
-        The order by criteria.
-    flatten_grid : bool
-        Whether to flatten the grid.
-    remapping : Dict[str, Any]
-        The remapping configuration.
-    use_grib_paramid : bool
-        Whether to use GRIB parameter ID.
+    def __init__(self, config, *path):
+        self.config = config
+        self.path = path
+        assert path[0] in (
+            "input",
+            "data_sources",
+        ), f"{self.__class__.__name__}: path must start with 'input' or 'data_sources': {path}"
+
+
+class Concat(Action):
+    """The Concat contruct is used to concat different actions that are responsible
+    for delivery fields for different dates.
+
+    See :ref:`building-concat` for more details.
+
+    .. block-code:: yaml
+
+        input:
+            concat:
+                - dates:
+                    start: 2023-01-01
+                    end: 2023-01-31
+                    frequency: 1d
+                  action: # some action
+                     ...
+
+                - dates:
+                    start: 2023-02-01
+                    end: 2023-02-28
+                    frequency: 1d
+                  action: # some action
+
     """
 
-    def __init__(self, /, order_by: str, flatten_grid: bool, remapping: Dict[str, Any], use_grib_paramid: bool) -> None:
-        """Initialize an ActionContext instance.
-
-        Parameters
-        ----------
-        order_by : str
-            The order by criteria.
-        flatten_grid : bool
-            Whether to flatten the grid.
-        remapping : Dict[str, Any]
-            The remapping configuration.
-        use_grib_paramid : bool
-            Whether to use GRIB parameter ID.
-        """
-        super().__init__()
-        self.order_by = order_by
-        self.flatten_grid = flatten_grid
-        self.remapping = build_remapping(remapping)
-        self.use_grib_paramid = use_grib_paramid
-
-
-def action_factory(config: Dict[str, Any], context: ActionContext, action_path: List[str]) -> Action:
-    """Factory function to create an Action instance based on the configuration.
-
-    Parameters
-    ----------
-    config : Dict[str, Any]
-        The action configuration.
-    context : ActionContext
-        The context in which the action exists.
-    action_path : List[str]
-        The action path.
-
-    Returns
-    -------
-    Action
-        The created Action instance.
+    def __init__(self, config, *path):
+        super().__init__(config, *path, "concat")
+
+        assert isinstance(config, list), f"Value must be a dict {list}"
+
+        self.choices = []
+
+        for i, item in enumerate(config):
+
+            dates = item["dates"]
+            filtering_dates = DatesProvider.from_config(**dates)
+            action = action_factory({k: v for k, v in item.items() if k != "dates"}, *self.path, str(i))
+            self.choices.append((filtering_dates, action))
+
+    def __repr__(self):
+        return f"Concat({self.choices})"
+
+    def __call__(self, context, argument):
+
+        results = context.empty_result()
+
+        for filtering_dates, action in self.choices:
+            dates = context.matching_dates(filtering_dates, argument)
+            if len(dates) == 0:
+                continue
+            results += action(context, dates)
+
+        return context.register(results, self.path)
+
+
+class Join(Action):
+    """Implement the join operation to combine results from multiple actions.
+
+    See :ref:`building-join` for more details.
+
+    .. block-code:: yaml
+
+        input:
+            join:
+                - grib:
+                     ...
+
+                - netcdf: # some other action
+                     ...
+
     """
-    from .concat import ConcatAction
-    from .data_sources import DataSourcesAction
-    from .function import FunctionAction
-    from .join import JoinAction
-    from .pipe import PipeAction
-    from .repeated_dates import RepeatedDatesAction
-
-    # from .data_sources import DataSourcesAction
-
-    assert isinstance(context, Context), (type, context)
-    if not isinstance(config, dict):
-        raise ValueError(f"Invalid input config {config}")
-    if len(config) != 1:
-        print(json.dumps(config, indent=2, default=str))
-        raise ValueError(f"Invalid input config. Expecting dict with only one key, got {list(config.keys())}")
-
-    config = deepcopy(config)
-    key = list(config.keys())[0]
-
-    if isinstance(config[key], list):
-        args, kwargs = config[key], {}
-    elif isinstance(config[key], dict):
-        args, kwargs = [], config[key]
-    else:
-        raise ValueError(f"Invalid input config {config[key]} ({type(config[key])}")
-
-    cls = {
-        "data_sources": DataSourcesAction,
-        "data-sources": DataSourcesAction,
-        "concat": ConcatAction,
-        "join": JoinAction,
-        "pipe": PipeAction,
-        "function": FunctionAction,
-        "repeated_dates": RepeatedDatesAction,
-        "repeated-dates": RepeatedDatesAction,
-    }.get(key)
-
-    if cls is None:
-        from ..sources import create_source
-
-        source = create_source(None, substitute(context, config))
-        return FunctionAction(context, action_path + [key], key, source)
-
-    return cls(context, action_path + [key], *args, **kwargs)
+
+    def __init__(self, config, *path):
+        super().__init__(config, *path, "join")
+
+        assert isinstance(config, list), f"Value of Join Action must be a list, got: {config}"
+
+        self.actions = [action_factory(item, *self.path, str(i)) for i, item in enumerate(config)]
+
+    def __repr__(self):
+        return f"Join({self.actions})"
+
+    def __call__(self, context, argument):
+        results = context.empty_result()
+
+        for action in self.actions:
+            results += action(context, argument)
+
+        return context.register(results, self.path)
+
+
+class Pipe(Action):
+    """Implement the pipe operation to chain results from a
+    source through multiple filters.
+
+    See :ref:`building-pipe` for more details.
+
+    .. block-code:: yaml
+
+        input:
+            pipe:
+                - grib:
+                     ...
+
+                - rename:
+                     ...
+
+    """
+
+    def __init__(self, config, *path):
+        assert isinstance(config, list), f"Value of Pipe Action must be a list, got {config}"
+        super().__init__(config, *path, "pipe")
+        self.actions = [action_factory(item, *self.path, str(i)) for i, item in enumerate(config)]
+
+    def __repr__(self):
+        return f"Pipe({self.actions})"
+
+    def __call__(self, context, argument):
+        result = context.empty_result()
+
+        for i, action in enumerate(self.actions):
+            if i == 0:
+                result = action(context, argument)
+            else:
+                result = action(context, result)
+
+        return context.register(result, self.path)
+
+
+class Function(Action):
+    """Base class for sources and filters."""
+
+    def __init__(self, config, *path):
+        super().__init__(config, *path, self.name)
+
+    def __call__(self, context, argument):
+
+        config = context.resolve(self.config)  # Substitute the ${} variables in the config
+
+        config["_type"] = self.name  # Find a better way to do this
+
+        source = self.create_object(context, config)
+
+        return context.register(self.call_object(context, source, argument), self.path)
+
+
+class DatasetSourceMixin:
+    """Mixin class for sources defined in anemoi-datasets"""
+
+    def create_object(self, context, config):
+        from anemoi.datasets.create.sources import create_source as create_datasets_source
+
+        return create_datasets_source(context, config)
+
+    def call_object(self, context, source, argument):
+        return source.execute(context.source_argument(argument))
+
+
+class TransformSourceMixin:
+    """Mixin class for sources defined in anemoi-transform"""
+
+    def create_object(self, context, config):
+        from anemoi.transform.sources import create_source as create_transform_source
+
+        return create_transform_source(context, config)
+
+
+class TransformFilterMixin:
+    """Mixin class for filters defined in anemoi-transform"""
+
+    def create_object(self, context, config):
+        from anemoi.transform.filters import create_filter as create_transform_filter
+
+        return create_transform_filter(context, config)
+
+    def call_object(self, context, filter, argument):
+        return filter.forward(context.filter_argument(argument))
+
+
+class FilterFunction(Function):
+    """Action to call a filter on the argument (e.g. rename, regrid, etc.)."""
+
+    def __call__(self, context, argument):
+        return self.call(context, argument, context.filter_argument)
+
+
+def _make_name(name, what):
+    name = name.replace("_", "-")
+    name = "".join(x.title() for x in name.split("-"))
+    return name + what.title()
+
+
+def new_source(name, mixin):
+    return type(
+        _make_name(name, "source"),
+        (Function, mixin),
+        {"name": name},
+    )
+
+
+def new_filter(name, mixin):
+    return type(
+        _make_name(name, "filter"),
+        (Function, mixin),
+        {"name": name},
+    )
+
+
+class DataSources(Action):
+    """Action to call a source (e.g. mars, netcdf, grib, etc.)."""
+
+    def __init__(self, config, *path):
+        super().__init__(config, *path)
+        assert isinstance(config, (dict, list)), f"Invalid config type: {type(config)}"
+        if isinstance(config, dict):
+            self.sources = {k: action_factory(v, *path, k) for k, v in config.items()}
+        else:
+            self.sources = {i: action_factory(v, *path, str(i)) for i, v in enumerate(config)}
+
+    def __call__(self, context, argument):
+        for name, source in self.sources.items():
+            context.register(source(context, argument), self.path + (name,))
+
+
+class Recipe(Action):
+    """Action that represent a recipe (i.e. a sequence of data_sources and input)."""
+
+    def __init__(self, input, data_sources):
+        self.input = input
+        self.data_sources = data_sources
+
+    def __call__(self, context, argument):
+        # Load data_sources
+        self.data_sources(context, argument)
+        return self.input(context, argument)
+
+
+KLASS = {
+    "concat": Concat,
+    "join": Join,
+    "pipe": Pipe,
+    "data-sources": DataSources,
+}
+
+LEN_KLASS = len(KLASS)
+
+
+def make(key, config, *path):
+
+    if LEN_KLASS == len(KLASS):
+
+        # Load pluggins
+        from anemoi.transform.filters import filter_registry as transform_filter_registry
+        from anemoi.transform.sources import source_registry as transform_source_registry
+
+        from anemoi.datasets.create.sources import source_registry as dataset_source_registry
+
+        # Register sources, local first
+        for name in dataset_source_registry.registered:
+            if name not in KLASS:
+                KLASS[name.replace("_", "-")] = new_source(name, DatasetSourceMixin)
+
+        for name in transform_source_registry.registered:
+            if name not in KLASS:
+                KLASS[name.replace("_", "-")] = new_source(name, TransformSourceMixin)
+
+        # Register filters
+        for name in transform_filter_registry.registered:
+            if name not in KLASS:
+                KLASS[name.replace("_", "-")] = new_filter(name, TransformFilterMixin)
+
+    return KLASS[key.replace("_", "-")](config, *path)
+
+
+def action_factory(data, *path):
+
+    assert len(path) > 0, f"Path must contain at least one element {path}"
+    assert path[0] in ("input", "data_sources")
+
+    assert isinstance(data, dict), f"Input data must be a dictionary, got {type(data)}"
+    assert len(data) == 1, f"Input data must contain exactly one key-value pair {data} {'.'.join(x for x in path)}"
+
+    key, value = next(iter(data.items()))
+    return make(key, value, *path)

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

@@ -0,0 +1,71 @@
+# (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
+from abc import abstractmethod
+from typing import Any
+
+LOG = logging.getLogger(__name__)
+
+
+class Context(ABC):
+    """Context for building input data."""
+
+    def __init__(self, /, argument: Any) -> None:
+        self.results = {}
+        self.cache = {}
+        self.argument = argument
+
+    def trace(self, emoji, *message) -> None:
+
+        print(f"{emoji}: {message}")
+
+    def register(self, data: Any, path: list[str]) -> Any:
+
+        if not path:
+            return data
+
+        assert path[0] in ("input", "data_sources"), path
+
+        LOG.info(f"Registering data at path: {path}")
+        self.results[tuple(path)] = data
+        return data
+
+    def resolve(self, config):
+        config = config.copy()
+
+        for key, value in list(config.items()):
+            if isinstance(value, str) and value.startswith("${") and value.endswith("}"):
+                path = tuple(value[2:-1].split("."))
+                if path in self.results:
+                    config[key] = self.results[path]
+                else:
+                    LOG.warning(f"Path not found {path}")
+                    for p in sorted(self.results):
+                        LOG.info(f"   Available paths: {p}")
+                    raise KeyError(f"Path {path} not found in results: {self.results.keys()}")
+
+        return config
+
+    def create_source(self, config: Any, *path) -> Any:
+        from anemoi.datasets.create.input.action import action_factory
+
+        if not isinstance(config, dict):
+            # It is already a result (e.g. ekd.FieldList), loaded from ${a.b.c}
+            # TODO: something more elegant
+            return lambda *args, **kwargs: config
+
+        return action_factory(config, *path)
+
+    @abstractmethod
+    def empty_result(self) -> Any: ...
+
+    @abstractmethod
+    def create_result(self, data: Any) -> Any: ...

anemoi/datasets/create/input/context/field.py

@@ -0,0 +1,54 @@
+# (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 typing import Any
+
+from earthkit.data.core.order import build_remapping
+
+from ..result.field import FieldResult
+from . import Context
+
+
+class FieldContext(Context):
+
+    def __init__(
+        self,
+        /,
+        argument: Any,
+        order_by: str,
+        flatten_grid: bool,
+        remapping: dict[str, Any],
+        use_grib_paramid: bool,
+    ) -> None:
+        super().__init__(argument)
+        self.order_by = order_by
+        self.flatten_grid = flatten_grid
+        self.remapping = build_remapping(remapping)
+        self.use_grib_paramid = use_grib_paramid
+        self.partial_ok = False
+
+    def empty_result(self) -> Any:
+        import earthkit.data as ekd
+
+        return ekd.from_source("empty")
+
+    def source_argument(self, argument: Any) -> Any:
+        return argument  # .dates
+
+    def filter_argument(self, argument: Any) -> Any:
+        return argument
+
+    def create_result(self, data):
+        return FieldResult(self, data)
+
+    def matching_dates(self, filtering_dates, group_of_dates: Any) -> Any:
+        from anemoi.datasets.dates.groups import GroupOfDates
+
+        return GroupOfDates(sorted(set(group_of_dates) & set(filtering_dates)), group_of_dates.provider)