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

anemoi/datasets/create/filter.py

@@ -0,0 +1,50 @@
+# (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 abc import ABC
+from abc import abstractmethod
+from typing import Any
+
+import earthkit.data as ekd
+
+
+class Filter(ABC):
+    """A base class for filters."""
+
+    def __init__(self, context: Any, *args: Any, **kwargs: Any) -> None:
+        """Initialise the filter.
+
+        Parameters
+        ----------
+        context : Any
+            The context in which the filter is created.
+        *args : tuple
+            Positional arguments.
+        **kwargs : dict
+            Keyword arguments.
+        """
+
+        self.context = context
+
+    @abstractmethod
+    def execute(self, data: ekd.FieldList) -> ekd.FieldList:
+        """Execute the filter.
+
+        Parameters
+        ----------
+        data : ekd.FieldList
+            The input data.
+
+        Returns
+        -------
+        ekd.FieldList
+            The output data.
+        """
+
+        pass

anemoi/datasets/create/filters/__init__.py

@@ -0,0 +1,33 @@
+# (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.
+#
+
+from typing import Any
+
+from anemoi.utils.registry import Registry
+
+filter_registry = Registry(__name__)
+
+
+def create_filter(context: Any, config: Any) -> Any:
+    """Create a filter based on the provided configuration.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the filter is created.
+    config : Any
+        The configuration for the filter.
+
+    Returns
+    -------
+    Any
+        The created filter.
+    """
+    return filter_registry.from_config(config, context)

anemoi/datasets/create/filters/empty.py

@@ -0,0 +1,37 @@
+# (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.
+
+
+from typing import Any
+
+import earthkit.data as ekd
+from anemoi.transform.fields import new_empty_fieldlist
+
+from .legacy import legacy_filter
+
+
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, **kwargs: Any) -> ekd.FieldList:
+    """Create a pipeline that returns an empty result.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    Any
+        An empty result.
+    """
+    return new_empty_fieldlist()

anemoi/datasets/create/filters/legacy.py

@@ -0,0 +1,93 @@
+# (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 inspect
+import logging
+import os
+from typing import Any
+from typing import Callable
+
+from ..filter import Filter
+from . import filter_registry
+
+LOG = logging.getLogger(__name__)
+
+
+class LegacyFilter(Filter):
+    """A legacy filter class.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the filter is created.
+    *args : tuple
+        Positional arguments.
+    **kwargs : dict
+        Keyword arguments.
+    """
+
+    def __init__(self, context: Any, *args: Any, **kwargs: Any) -> None:
+        super().__init__(context, *args, **kwargs)
+        self.args = args
+        self.kwargs = kwargs
+
+
+class legacy_filter:
+    """A decorator class for legacy filters.
+
+    Parameters
+    ----------
+    name : str
+        The name of the legacy filter.
+    """
+
+    def __init__(self, name: str) -> None:
+        name, _ = os.path.splitext(os.path.basename(name))
+        self.name = name
+
+    def __call__(self, execute: Callable) -> Callable:
+        """Call method to wrap the execute function.
+
+        Parameters
+        ----------
+        execute : Callable
+            The execute function to be wrapped.
+
+        Returns
+        -------
+        Callable
+            The wrapped execute function.
+        """
+        this = self
+        name = f"Legacy{self.name.title()}Filter"
+        source = ".".join([execute.__module__, execute.__name__])
+
+        def execute_wrapper(self, input) -> Any:
+            """Wrapper method to call the execute function."""
+            try:
+                return execute(self.context, input, *self.args, **self.kwargs)
+            except TypeError:
+                LOG.error(f"Error executing filter {this.name} from {source}")
+                LOG.error(f"Function signature is: {inspect.signature(execute)}")
+                LOG.error(f"Arguments are: {self.args=}, {self.kwargs=}")
+                raise
+
+        klass = type(
+            name,
+            (LegacyFilter,),
+            {
+                "execute": execute_wrapper,
+                "_source": source,
+            },
+        )
+
+        filter_registry.register(self.name)(klass)
+
+        return execute

anemoi/datasets/create/filters/noop.py

@@ -0,0 +1,37 @@
+# (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.
+
+from typing import Any
+
+import earthkit.data as ekd
+
+from .legacy import legacy_filter
+
+
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, *args: Any, **kwargs: Any) -> ekd.FieldList:
+    """No operation filter that returns the input as is.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : ekd.FieldList
+        List of input fields.
+    *args : Any
+        Additional arguments.
+    **kwargs : Any
+        Additional keyword arguments.
+
+    Returns
+    -------
+    List[Any]
+        The input list of fields.
+    """
+    return input

anemoi/datasets/create/filters/orog_to_z.py

@@ -0,0 +1,58 @@
+# (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.
+
+from collections import defaultdict
+from typing import Any
+from typing import Dict
+
+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, orog: str, z: str = "z") -> ekd.FieldList:
+    """Convert orography [m] to z (geopotential height).
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : FieldList
+        List of input fields.
+    orog : str
+        Orography parameter.
+    z : str, optional
+        Geopotential height parameter. Defaults to "z".
+
+    Returns
+    -------
+    FieldList
+        List of fields with geopotential height.
+    """
+    result = []
+    processed_fields: Dict[tuple, Dict[str, Any]] = defaultdict(dict)
+
+    for f in input:
+        key = f.metadata(namespace="mars")
+        param = key.pop("param")
+        if param == orog:
+            key = tuple(key.items())
+
+            if param in processed_fields[key]:
+                raise ValueError(f"Duplicate field {param} for {key}")
+
+            output = f.to_numpy(flatten=True) * 9.80665
+            result.append(new_field_from_numpy(f, output, param=z))
+        else:
+            result.append(f)
+
+    return new_fieldlist_from_list(result)

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

@@ -7,21 +7,44 @@
 # 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
+from typing import Dict
+from typing import Tuple
 
-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.meteo import thermo
 
-from .single_level_specific_humidity_to_relative_humidity import NewDataField
+from .legacy import legacy_filter
+
 
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, t: str, rh: str, q: str = "q") -> ekd.FieldList:
+    """Convert relative humidity on pressure levels to specific humidity.
 
-def execute(context, input, t, rh, q="q"):
-    """Convert relative humidity on pressure levels to specific humidity"""
-    result = FieldArray()
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    t : str
+        Temperature parameter.
+    rh : str
+        Relative humidity parameter.
+    q : str, optional
+        Specific humidity parameter. Defaults to "q".
 
-    params = (t, rh)
-    pairs = defaultdict(dict)
+    Returns
+    -------
+    ekd.FieldList
+        Array of fields with specific humidity.
+    """
+    result = []
+    params: Tuple[str, str] = (t, rh)
+    pairs: Dict[Tuple[Any, ...], Dict[str, Any]] = defaultdict(dict)
 
     # Gather all necessary fields
     for f in input:
@@ -53,6 +76,6 @@ def execute(context, input, t, rh, q="q"):
 
         # actual conversion from rh --> q_v
         q_pl = thermo.specific_humidity_from_relative_humidity(t_pl, rh_pl, pressure)
-        result.append(NewDataField(values[rh], q_pl, q))
+        result.append(new_field_from_numpy(values[rh], q_pl, param=q))
 
-    return result
+    return new_fieldlist_from_list(result)

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

@@ -9,19 +9,43 @@
 
 
 from collections import defaultdict
+from typing import Any
+from typing import Dict
 
+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.data.indexing.fieldlist import FieldArray
 from earthkit.meteo import thermo
 
-from .single_level_specific_humidity_to_relative_humidity import NewDataField
+from .legacy import legacy_filter
 
 
-def execute(context, input, t, q, rh="r"):
-    """Convert specific humidity on pressure levels to relative humidity"""
-    result = FieldArray()
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, t: str, q: str, rh: str = "r") -> FieldArray:
+    """Convert specific humidity on pressure levels to relative humidity.
 
-    params = (t, q)
-    pairs = defaultdict(dict)
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    t : str
+        Temperature parameter.
+    q : str
+        Specific humidity parameter.
+    rh : str, optional
+        Relative humidity parameter. Defaults to "r".
+
+    Returns
+    -------
+    ekd.FieldList
+        Array of fields with relative humidity.
+    """
+    result = []
+    params: tuple[str, str] = (t, q)
+    pairs: Dict[tuple, Dict[str, Any]] = defaultdict(dict)
 
     # Gather all necessary fields
     for f in input:
@@ -53,6 +77,6 @@ def execute(context, input, t, q, rh="r"):
 
         # actual conversion from rh --> q_v
         rh_pl = thermo.relative_humidity_from_specific_humidity(t_pl, q_pl, pressure)
-        result.append(NewDataField(values[q], rh_pl, rh))
+        result.append(new_field_from_numpy(values[q], rh_pl, param=rh))
 
-    return result
+    return new_fieldlist_from_list(result)

anemoi/datasets/create/filters/rename.py

@@ -0,0 +1,205 @@
+# (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.
+
+import re
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+import earthkit.data as ekd
+from earthkit.data.indexing.fieldlist import FieldArray
+
+from .legacy import legacy_filter
+
+
+class RenamedFieldMapping:
+    """Rename a field based on the value of another field.
+
+    Parameters
+    ----------
+    field : Any
+        The field to be renamed.
+    what : str
+        The name of the field that will be used to rename the field.
+    renaming : dict of dict of str
+        A dictionary mapping the values of 'what' to the new names.
+    """
+
+    def __init__(self, field: Any, what: str, renaming: Dict[str, Dict[str, str]]) -> None:
+        """Initialize a RenamedFieldMapping instance.
+
+        Parameters
+        ----------
+        field : Any
+            The field to be renamed.
+        what : str
+            The name of the field that will be used to rename the field.
+        renaming : dict of dict of str
+            A dictionary mapping the values of 'what' to the new names.
+        """
+        self.field = field
+        self.what = what
+        self.renaming = {}
+        for k, v in renaming.items():
+            self.renaming[k] = {str(a): str(b) for a, b in v.items()}
+
+    def metadata(self, key: Optional[str] = None, **kwargs: Any) -> Any:
+        """Get metadata from the original field, with the option to rename the parameter.
+
+        Parameters
+        ----------
+        key : str, optional
+            The metadata key.
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Any
+            The metadata value.
+        """
+        if key is None:
+            return self.field.metadata(**kwargs)
+
+        value = self.field.metadata(key, **kwargs)
+        if key == self.what:
+            return self.renaming.get(self.what, {}).get(value, value)
+
+        return value
+
+    def __getattr__(self, name: str) -> Any:
+        """Get an attribute from the original field.
+
+        Parameters
+        ----------
+        name : str
+            The name of the attribute.
+
+        Returns
+        -------
+        Any
+            The attribute value.
+        """
+        return getattr(self.field, name)
+
+    def __repr__(self) -> str:
+        """Get the string representation of the original field.
+
+        Returns
+        -------
+        str
+            The string representation of the original field.
+        """
+        return repr(self.field)
+
+
+class RenamedFieldFormat:
+    """Rename a field based on a format string.
+
+    Parameters
+    ----------
+    field : Any
+        The field to be renamed.
+    what : str
+        The name of the field that will be used to rename the field.
+    format : str
+        The format string for renaming.
+    """
+
+    def __init__(self, field: Any, what: str, format: str) -> None:
+        """Initialize a RenamedFieldFormat instance.
+
+        Parameters
+        ----------
+        field : Any
+            The field to be renamed.
+        what : str
+            The name of the field that will be used to rename the field.
+        format : str
+            The format string for renaming.
+        """
+        self.field = field
+        self.what = what
+        self.format = format
+        self.bits = re.findall(r"{(\w+)}", format)
+
+    def metadata(self, *args: Any, **kwargs: Any) -> Any:
+        """Get metadata from the original field, with the option to rename the parameter using a format string.
+
+        Parameters
+        ----------
+        *args : Any
+            Positional arguments.
+        **kwargs : Any
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Any
+            The metadata value.
+        """
+        value = self.field.metadata(*args, **kwargs)
+        if args:
+            assert len(args) == 1
+            if args[0] == self.what:
+                bits = {b: self.field.metadata(b, **kwargs) for b in self.bits}
+                return self.format.format(**bits)
+        return value
+
+    def __getattr__(self, name: str) -> Any:
+        """Get an attribute from the original field.
+
+        Parameters
+        ----------
+        name : str
+            The name of the attribute.
+
+        Returns
+        -------
+        Any
+            The attribute value.
+        """
+        return getattr(self.field, name)
+
+    def __repr__(self) -> str:
+        """Get the string representation of the original field.
+
+        Returns
+        -------
+        str
+            The string representation of the original field.
+        """
+        return repr(self.field)
+
+
+@legacy_filter(__file__)
+def execute(context: Any, input: ekd.FieldList, what: str = "param", **kwargs: Any) -> ekd.FieldList:
+    """Rename fields based on the value of another field or a format string.
+
+    Parameters
+    ----------
+    context : Any
+        The context in which the function is executed.
+    input : List[Any]
+        List of input fields.
+    what : str, optional
+        The field to be used for renaming. Defaults to "param".
+    **kwargs : Any
+        Additional keyword arguments for renaming.
+
+    Returns
+    -------
+    ekd.FieldList
+        Array of renamed fields.
+    """
+
+    if what in kwargs and isinstance(kwargs[what], str):
+        return FieldArray([RenamedFieldFormat(fs, what, kwargs[what]) for fs in input])
+
+    return FieldArray([RenamedFieldMapping(fs, what, kwargs) for fs in input])