anemoi-datasets 0.5.16__py3-none-any.whl → 0.5.17__py3-none-any.whl

anemoi/datasets/create/filters/transform.py

@@ -0,0 +1,53 @@
+# (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 typing import Dict
+
+import earthkit.data as ekd
+
+from ..filter import Filter
+
+
+class TransformFilter(Filter):
+    """Calls filters from anemoi.transform.filters
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the filter is created.
+    name : str
+        The name of the filter.
+    config : Dict[str, Any]
+        The configuration for the filter.
+    """
+
+    def __init__(self, context: Any, name: str, config: Dict[str, Any]) -> None:
+
+        from anemoi.transform.filters import create_filter
+
+        self.name = name
+        self.transform_filter = create_filter(self, config)
+
+    def execute(self, context: Any, input: ekd.FieldList) -> ekd.FieldList:
+        """Execute the transformation filter.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the execution occurs.
+        input : ekd.FieldList
+            The input data to be transformed.
+
+        Returns
+        -------
+        ekd.FieldList
+            The transformed data.
+        """
+        return self.transform_filter.forward(input)

anemoi/datasets/create/{functions/filters → filters}/unrotate_winds.py

@@ -7,29 +7,38 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-
 from collections import defaultdict
+from typing import Any
 
-# import numpy as np
-from earthkit.data.indexing.fieldlist import FieldArray
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
 from earthkit.geo.rotate import unrotate_vector
 
+from .legacy import legacy_filter
 
-class NewDataField:
-    def __init__(self, field, data):
-        self.field = field
-        self.data = data
-
-    def to_numpy(self, *args, **kwargs):
-        return self.data
 
-    def __getattr__(self, name):
-        return getattr(self.field, name)
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, u: str, v: str) -> ekd.FieldList:
+    """Unrotate the wind components of a GRIB file.
 
+    Parameters
+    ----------
+    context : Any
+        The execution context.
+    input : List[Any]
+        The list of input fields.
+    u : str
+        The parameter name for the u-component of the wind.
+    v : str
+        The parameter name for the v-component of the wind.
 
-def execute(context, input, u, v):
-    """Unrotate the wind components of a GRIB file."""
-    result = FieldArray()
+    Returns
+    -------
+    ekd.FieldList
+        The resulting field array with unrotated wind components.
+    """
+    result = []
 
     wind_params = (u, v)
     wind_pairs = defaultdict(dict)
@@ -72,10 +81,10 @@ def execute(context, input, u, v):
             lon_unrotated=raw_longs,
         )
 
-        result.append(NewDataField(x, u_new))
-        result.append(NewDataField(y, v_new))
+        result.append(new_field_from_numpy(x, u_new))
+        result.append(new_field_from_numpy(y, v_new))
 
-    return result
+    return new_fieldlist_from_list(result)
 
 
 if __name__ == "__main__":

anemoi/datasets/create/filters/uv_to_speeddir.py

@@ -0,0 +1,94 @@
+# (C) Copyright 2024 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.
+
+"""This module provides a function to convert u and v wind components to wind speed and direction."""
+
+from collections import defaultdict
+from typing import Any
+from typing import List
+
+import earthkit.data as ekd
+import numpy as np
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
+from earthkit.meteo.wind.array import xy_to_polar
+
+from .legacy import legacy_filter
+
+
+@legacy_filter(__file__)
+def execute(
+    context: Any,
+    input: List[Any],
+    u_component: str,
+    v_component: str,
+    wind_speed: str,
+    wind_dir: str,
+    in_radians: bool = False,
+) -> ekd.FieldList:
+    """Converts u and v wind components to wind speed and direction.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields containing wind components.
+    u_component : str
+        The name of the u component field.
+    v_component : str
+        The name of the v component field.
+    wind_speed : str
+        The name of the wind speed field to be created.
+    wind_dir : str
+        The name of the wind direction field to be created.
+    in_radians : bool, optional
+        If True, the wind direction is returned in radians. Default is False.
+
+    Returns
+    -------
+    ekd.FieldList
+        A FieldArray containing the wind speed and direction fields.
+    """
+    result = []
+
+    wind_params = (u_component, v_component)
+    wind_pairs = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+
+        if param not in wind_params:
+            result.append(f)
+            continue
+
+        key = tuple(key.items())
+
+        if param in wind_pairs[key]:
+            raise ValueError(f"Duplicate wind component {param} for {key}")
+
+        wind_pairs[key][param] = f
+
+    for _, pairs in wind_pairs.items():
+        if len(pairs) != 2:
+            raise ValueError("Missing wind component")
+
+        u = pairs[u_component]
+        v = pairs[v_component]
+
+        # assert speed.grid_mapping == dir.grid_mapping
+        magnitude, direction = xy_to_polar(u.to_numpy(flatten=True), v.to_numpy(flatten=True))
+        if in_radians:
+            direction = np.deg2rad(direction)
+
+        result.append(new_field_from_numpy(u, magnitude, params=wind_speed))
+        result.append(new_field_from_numpy(v, direction, params=wind_dir))
+
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/{functions/filters → filters}/wz_to_w.py

@@ -7,37 +7,39 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-
 from collections import defaultdict
-
-from earthkit.data.indexing.fieldlist import FieldArray
-
-
-class NewDataField:
-    def __init__(self, field, data, new_name):
-        self.field = field
-        self.data = data
-        self.new_name = new_name
-
-    def to_numpy(self, *args, **kwargs):
-        return self.data
-
-    def metadata(self, key=None, **kwargs):
-        if key is None:
-            return self.field.metadata(**kwargs)
-
-        value = self.field.metadata(key, **kwargs)
-        if key == "param":
-            return self.new_name
-        return value
-
-    def __getattr__(self, name):
-        return getattr(self.field, name)
-
-
-def execute(context, input, wz, t, w="w"):
-    """Convert geometric vertical velocity (m/s) to vertical velocity (Pa / s)"""
-    result = FieldArray()
+from typing import Any
+
+import earthkit.data as ekd
+from anemoi.transform.fields import new_field_from_numpy
+from anemoi.transform.fields import new_fieldlist_from_list
+
+from .legacy import legacy_filter
+
+
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, wz: str, t: str, w: str = "w") -> ekd.FieldList:
+    """Convert geometric vertical velocity (m/s) to vertical velocity (Pa / s).
+
+    Parameters
+    ----------
+    context : Any
+        The context for the execution.
+    input : List[Any]
+        The list of input fields.
+    wz : str
+        The parameter name for geometric vertical velocity.
+    t : str
+        The parameter name for temperature.
+    w : str, optional
+        The parameter name for vertical velocity. Defaults to "w".
+
+    Returns
+    -------
+    ekd.FieldList
+        The resulting FieldArray with converted vertical velocity fields.
+    """
+    result = []
 
     params = (wz, t)
     pairs = defaultdict(dict)
@@ -67,12 +69,28 @@ def execute(context, input, wz, t, w="w"):
         pressure = keys[4][1] * 100  # TODO: REMOVE HARDCODED INDICES
 
         w_pl = wz_to_w(wz_pl, t_pl, pressure)
-        result.append(NewDataField(values[wz], w_pl, w))
+        result.append(new_field_from_numpy(values[wz], w_pl, param=w))
+
+    return new_fieldlist_from_list(result)
+
 
-    return result
+def wz_to_w(wz: Any, t: Any, pressure: float) -> Any:
+    """Convert geometric vertical velocity (m/s) to vertical velocity (Pa / s).
 
+    Parameters
+    ----------
+    wz : Any
+        The geometric vertical velocity data.
+    t : Any
+        The temperature data.
+    pressure : float
+        The pressure value.
 
-def wz_to_w(wz, t, pressure):
+    Returns
+    -------
+    Any
+        The vertical velocity data in Pa / s.
+    """
     g = 9.81
     Rd = 287.058
 

anemoi/datasets/create/input/__init__.py

@@ -9,14 +9,37 @@
 
 import logging
 from copy import deepcopy
+from typing import Any
+from typing import Union
+
+from anemoi.datasets.dates.groups import GroupOfDates
 
 from .trace import trace_select
 
 LOG = logging.getLogger(__name__)
 
 
+class Context:
+    """Context for building input data."""
+
+    pass
+
+
 class InputBuilder:
-    def __init__(self, config, data_sources, **kwargs):
+    """Builder class for creating input data from configuration and data sources."""
+
+    def __init__(self, config: dict, data_sources: Union[dict, list], **kwargs: Any) -> None:
+        """Initialize the InputBuilder.
+
+        Parameters
+        ----------
+        config : dict
+            Configuration dictionary.
+        data_sources : Union[dict, list]
+            Data sources.
+        **kwargs : Any
+            Additional keyword arguments.
+        """
         self.kwargs = kwargs
 
         config = deepcopy(config)
@@ -31,7 +54,19 @@ class InputBuilder:
         self.action_path = ["input"]
 
     @trace_select
-    def select(self, group_of_dates):
+    def select(self, group_of_dates: GroupOfDates) -> Any:
+        """Select data based on the group of dates.
+
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            Group of dates to select data for.
+
+        Returns
+        -------
+        Any
+            Selected data.
+        """
         from .action import ActionContext
         from .action import action_factory
 
@@ -40,7 +75,14 @@ class InputBuilder:
         action = action_factory(self.config, context, self.action_path)
         return action.select(group_of_dates)
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Return a string representation of the InputBuilder.
+
+        Returns
+        -------
+        str
+            String representation.
+        """
         from .action import ActionContext
         from .action import action_factory
 
@@ -48,8 +90,37 @@ class InputBuilder:
         a = action_factory(self.config, context, self.action_path)
         return repr(a)
 
-    def _trace_select(self, group_of_dates):
+    def _trace_select(self, group_of_dates: GroupOfDates) -> str:
+        """Trace the select operation.
+
+        Parameters
+        ----------
+        group_of_dates : GroupOfDates
+            Group of dates to select data for.
+
+        Returns
+        -------
+        str
+            Trace string.
+        """
         return f"InputBuilder({group_of_dates})"
 
 
-build_input = InputBuilder
+def build_input(config: dict, data_sources: Union[dict, list], **kwargs: Any) -> InputBuilder:
+    """Build an InputBuilder instance.
+
+    Parameters
+    ----------
+    config : dict
+        Configuration dictionary.
+    data_sources : Union[dict, list]
+        Data sources.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    InputBuilder
+        An instance of InputBuilder.
+    """
+    return InputBuilder(config, data_sources, **kwargs)

anemoi/datasets/create/input/action.py

@@ -9,17 +9,49 @@
 
 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 .misc import is_function
 
 LOG = logging.getLogger(__name__)
 
 
 class Action:
-    def __init__(self, context, action_path, /, *args, **kwargs):
+    """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.
+    """
+
+    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 = []
@@ -38,13 +70,43 @@ class Action:
         self.action_path = action_path
 
     @classmethod
-    def _short_str(cls, x):
+    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, _indent_="\n", _inline_="", **kwargs):
+    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()])
 
@@ -54,18 +116,77 @@ class Action:
             txt = txt.replace("\n", "\n  ")
         return txt
 
-    def select(self, dates, **kwargs):
+    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):
+    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):
+    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):
-    def __init__(self, /, order_by, flatten_grid, remapping, use_grib_paramid):
+    """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, /, 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
@@ -73,8 +194,23 @@ class ActionContext(Context):
         self.use_grib_paramid = use_grib_paramid
 
 
-def action_factory(config, context, action_path):
+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.
+    """
     from .concat import ConcatAction
     from .data_sources import DataSourcesAction
     from .function import FunctionAction
@@ -110,9 +246,9 @@ def action_factory(config, context, action_path):
     }.get(key)
 
     if cls is None:
-        if not is_function(key, "sources"):
-            raise ValueError(f"Unknown action '{key}' in {config}")
-        cls = FunctionAction
-        args = [key] + args
+        from ..sources import create_source
+
+        source = create_source(None, config)
+        return FunctionAction(context, action_path + [key], key, source)
 
     return cls(context, action_path + [key], *args, **kwargs)