Merge remote-tracking branch 'origin/develop' into feature/regrid

Author: b8raoult

CHANGELOG.md

@@ -13,6 +13,7 @@ Keep it human-readable, your future self will thank you!
 ### Added
 
 - Add regrid filter
+- Added repeat-member #18
 
 ## [0.1.0](https://github.com/ecmwf/anemoi-utils/transform/0.0.5...HEAD/compare/0.0.8...0.1.0) - 2024-11-18
 

src/anemoi/transform/filters/lambda_filters.py

@@ -0,0 +1,114 @@
+# (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 importlib
+import typing as tp
+
+from earthkit.data.core.fieldlist import Field
+from earthkit.data.core.fieldlist import FieldList
+
+from anemoi.transform.filters import filter_registry
+from anemoi.transform.filters.base import SimpleFilter
+
+
+@filter_registry.register("earthkitfieldlambda")
+class EarthkitFieldLambdaFilter(SimpleFilter):
+    """A filter to apply an arbitrary function to individual fields."""
+
+    def __init__(
+        self,
+        fn: str | tp.Callable[[Field, tp.Any], Field],
+        param: str | list[str],
+        args: list = [],
+        kwargs: dict[str, tp.Any] = {},
+        backward_fn: str | tp.Callable[[Field, tp.Any], Field] | None = None,
+    ):
+        """Initialise the EarthkitFieldLambdaFilter.
+
+        Parameters
+        ----------
+        fn: callable or str
+            The lambda function as a callable with the general signature
+            `fn(*earthkit.data.Field, *args, **kwargs) -> earthkit.data.Field` or
+            a string path to the function, such as "package.module.function".
+        param: list or str
+            The parameter name or list of parameter names to apply the function to.
+        args: list
+            The list of arguments to pass to the lambda function.
+        kwargs: dict
+            The dictionary of keyword arguments to pass to the lambda function.
+        backward_fn (optional): callable, str or None
+            The backward lambda function as a callable with the general signature
+            `backward_fn(*earthkit.data.Field, *args, **kwargs) -> earthkit.data.Field` or
+            a string path to the function, such as "package.module.function".
+
+        Examples
+        --------
+        >>> from anemoi.transform.filters.lambda_filters import EarthkitFieldLambdaFilter
+        >>> import earthkit.data as ekd
+        >>> fields = ekd.from_source(
+        ...        "mars",{"param": ["2t"],
+        ...                "levtype": "sfc",
+        ...                "dates": ["2023-11-17 00:00:00"]})
+        >>> kelvin_to_celsius = EarthkitFieldLambdaFilter(
+        ...     fn=lambda x, s: x.clone(values=x.values - s),
+        ...     param="2t",
+        ...     args=[273.15],
+        ... )
+        >>> fields = kelvin_to_celsius.forward(fields)
+        """
+        if not isinstance(args, list):
+            raise ValueError("Expected 'args' to be a list. " f"Got {args} instead.")
+        if not isinstance(kwargs, dict):
+            raise ValueError("Expected 'kwargs' to be a dictionary. " f"Got {kwargs} instead.")
+
+        self.fn = self._import_fn(fn) if isinstance(fn, str) else fn
+
+        if isinstance(backward_fn, str):
+            self.backward_fn = self._import_fn(backward_fn)
+        else:
+            self.backward_fn = backward_fn
+
+        self.param = param if isinstance(param, list) else [param]
+        self.args = args
+        self.kwargs = kwargs
+
+    def forward(self, data: FieldList) -> FieldList:
+        return self._transform(data, self.forward_transform, *self.param)
+
+    def backward(self, data: FieldList) -> FieldList:
+        if self.backward_fn:
+            return self._transform(data, self.backward_transform, *self.param)
+        raise NotImplementedError(f"{self} is not reversible.")
+
+    def forward_transform(self, *fields: Field) -> tp.Iterator[Field]:
+        """Apply the lambda function to the field."""
+        yield self.fn(*fields, *self.args, **self.kwargs)
+
+    def backward_transform(self, *fields: Field) -> tp.Iterator[Field]:
+        """Apply the backward lambda function to the field."""
+        yield self.backward_fn(*fields, *self.args, **self.kwargs)
+
+    def _import_fn(self, fn: str) -> tp.Callable[..., Field]:
+        try:
+            module_name, fn_name = fn.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            return getattr(module, fn_name)
+        except Exception as e:
+            raise ValueError(f"Could not import function {fn}") from e
+
+    def __repr__(self):
+        out = f"{self.__class__.__name__}(fn={self.fn},"
+        if self.backward_fn:
+            out += f"backward_fn={self.backward_fn},"
+        out += f"param={self.param},"
+        out += f"args={self.args},"
+        out += f"kwargs={self.kwargs},"
+        out += ")"
+        return out

src/anemoi/transform/filters/repeat_members.py

@@ -0,0 +1,85 @@
+# (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 logging
+
+from ..fields import new_field_from_numpy
+from ..fields import new_fieldlist_from_list
+from . import filter_registry
+from .base import Filter
+
+LOG = logging.getLogger(__name__)
+
+
+def make_list_int(value):
+    if isinstance(value, str):
+        if "/" not in value:
+            return [value]
+        bits = value.split("/")
+        if len(bits) == 3 and bits[1].lower() == "to":
+            value = list(range(int(bits[0]), int(bits[2]) + 1, 1))
+
+        elif len(bits) == 5 and bits[1].lower() == "to" and bits[3].lower() == "by":
+            value = list(range(int(bits[0]), int(bits[2]) + int(bits[4]), int(bits[4])))
+
+    if isinstance(value, list):
+        return value
+    if isinstance(value, tuple):
+        return value
+    if isinstance(value, int):
+        return [value]
+
+    raise ValueError(f"Cannot make list from {value}")
+
+
+@filter_registry.register("repeat_members")
+class RepeatMembers(Filter):
+    """The filter can be used to replicate non-ensembles fields into ensemble fields.
+
+    Args: (only one of the following)
+      numbers: A list of numbers (1-based) of the fields to replicate.
+      members: A list of 0-based indices of the fields to replicate.
+      count: The number of times to replicate the fields.
+    """
+
+    def __init__(
+        self,
+        numbers=None,  # 1-based
+        members=None,  # 0-based
+        count=None,
+    ):
+        if sum(x is not None for x in (members, count, numbers)) != 1:
+            raise ValueError("Exactly one of members, count or numbers must be given")
+
+        if numbers is not None:
+            numbers = make_list_int(numbers)
+            members = [n - 1 for n in numbers]
+
+        if count is not None:
+            members = list(range(count))
+
+        members = make_list_int(members)
+        self.members = members
+        assert isinstance(members, (tuple, list)), f"members must be a list or tuple, got {type(members)}"
+
+    def forward(self, data):
+        result = []
+        for f in data:
+            array = f.to_numpy()
+            for member in self.members:
+                number = member + 1
+                new_field = new_field_from_numpy(array, template=f, number=number)
+                result.append(new_field)
+
+        return new_fieldlist_from_list(result)
+
+    def backward(self, data):
+        # this could be implemented
+        raise NotImplementedError("RepeatMembers is not reversible")

tests/__init__.py

@@ -7,51 +7,87 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
+import sys
+from pathlib import Path
 
 import earthkit.data as ekd
+import numpy.testing as npt
 from pytest import approx
 
+from anemoi.transform.filters.lambda_filters import EarthkitFieldLambdaFilter
 from anemoi.transform.filters.rescale import Convert
 from anemoi.transform.filters.rescale import Rescale
 
+sys.path.append(Path(__file__).parents[1].as_posix())
 
-def test_rescale():
+
+def test_rescale(fieldlist):
+    fieldlist = fieldlist.sel(param="2t")
     # rescale from K to °C
-    temp = ekd.from_source("mars", {"param": "2t", "levtype": "sfc", "dates": ["2023-11-17 00:00:00"]})
-    fieldlist = temp.to_fieldlist()
     k_to_deg = Rescale(scale=1.0, offset=-273.15, param="2t")
     rescaled = k_to_deg.forward(fieldlist)
-    assert rescaled[0].values.min() == temp.values.min() - 273.15
-    assert rescaled[0].values.std() == approx(temp.values.std())
+
+    npt.assert_allclose(rescaled[0].to_numpy(), fieldlist[0].to_numpy() - 273.15)
     # and back
     rescaled_back = k_to_deg.backward(rescaled)
-    assert rescaled_back[0].values.min() == temp.values.min()
-    assert rescaled_back[0].values.std() == approx(temp.values.std())
-    # rescale from °C to F
-    deg_to_far = Rescale(scale=9 / 5, offset=32, param="2t")
-    rescaled_farheneit = deg_to_far.forward(rescaled)
-    assert rescaled_farheneit[0].values.min() == 9 / 5 * rescaled[0].values.min() + 32
-    assert rescaled_farheneit[0].values.std() == approx((9 / 5) * rescaled[0].values.std())
-    # rescale from F to K
-    rescaled_back = k_to_deg.backward(deg_to_far.backward(rescaled_farheneit))
-    assert rescaled_back[0].values.min() == temp.values.min()
-    assert rescaled_back[0].values.std() == approx(temp.values.std())
-
-
-def test_convert():
+    npt.assert_allclose(rescaled_back[0].to_numpy(), fieldlist[0].to_numpy())
+
+
+def test_convert(fieldlist):
     # rescale from K to °C
-    temp = ekd.from_source("mars", {"param": "2t", "levtype": "sfc", "dates": ["2023-11-17 00:00:00"]})
-    fieldlist = temp.to_fieldlist()
+    fieldlist = fieldlist.sel(param="2t")
     k_to_deg = Convert(unit_in="K", unit_out="degC", param="2t")
     rescaled = k_to_deg.forward(fieldlist)
-    assert rescaled[0].values.min() == temp.values.min() - 273.15
-    assert rescaled[0].values.std() == approx(temp.values.std())
+    assert rescaled[0].values.min() == fieldlist.values.min() - 273.15
+    assert rescaled[0].values.std() == approx(fieldlist.values.std())
     # and back
     rescaled_back = k_to_deg.backward(rescaled)
-    assert rescaled_back[0].values.min() == temp.values.min()
-    assert rescaled_back[0].values.std() == approx(temp.values.std())
+    assert rescaled_back[0].values.min() == fieldlist.values.min()
+    assert rescaled_back[0].values.std() == approx(fieldlist.values.std())
+
+
+# used in the test below
+def _do_something(field, a):
+    return field.clone(values=field.values * a)
+
+
+def test_singlefieldlambda(fieldlist):
+
+    fieldlist = fieldlist.sel(param="sp")
+
+    def undo_something(field, a):
+        return field.clone(values=field.values / a)
+
+    something = EarthkitFieldLambdaFilter(
+        fn="tests.test_filters._do_something",
+        param="sp",
+        args=[10],
+        backward_fn=undo_something,
+    )
+
+    transformed = something.forward(fieldlist)
+    npt.assert_allclose(transformed[0].to_numpy(), fieldlist[0].to_numpy() * 10)
+
+    untransformed = something.backward(transformed)
+    npt.assert_allclose(untransformed[0].to_numpy(), fieldlist[0].to_numpy())
 
 
 if __name__ == "__main__":
-    test_rescale()
-    test_convert()
+
+    fieldlist = ekd.from_source(
+        "mars",
+        {
+            "param": ["2t", "sp"],
+            "levtype": "sfc",
+            "dates": ["2023-11-17 00:00:00"],
+        },
+    )
+
+    test_rescale(fieldlist)
+    try:
+        test_convert(fieldlist)
+    except FileNotFoundError:
+        print("Skipping test_convert because of missing UNIDATA UDUNITS2 library, " "required by cfunits.")
+    test_singlefieldlambda(fieldlist)
+
+    print("All tests passed.")

tests/test_filters.py

@@ -0,0 +1,71 @@
+# (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 earthkit.data as ekd
+import numpy as np
+
+from anemoi.transform.filters.repeat_members import RepeatMembers
+
+
+def _get_template():
+    temp = ekd.from_source("mars", {"param": "2t", "levtype": "sfc", "dates": ["2023-11-17 00:00:00"]})
+    fieldlist = temp.to_fieldlist()
+    return fieldlist, fieldlist[0].values, fieldlist[0].metadata
+
+
+def test_repeat_members_using_numbers_1():
+    fieldlist, values, metadata = _get_template()
+
+    repeat = RepeatMembers(numbers=[1, 2, 3])
+    repeated = repeat.forward(fieldlist)
+    assert len(repeated) == 3
+    for i, f in enumerate(repeated):
+        assert f.values.shape == values.shape
+        assert np.all(f.values == values)
+        assert f.metadata("number") == i + 1
+        assert f.metadata("name") == metadata("name")
+
+
+def test_repeat_members_using_numbers_2():
+    fieldlist, values, metadata = _get_template()
+
+    repeat = RepeatMembers(numbers="1/to/3")
+    repeated = repeat.forward(fieldlist)
+    assert len(repeated) == 3
+    for i, f in enumerate(repeated):
+        assert f.values.shape == values.shape
+        assert np.all(f.values == values)
+        assert f.metadata("number") == i + 1
+        assert f.metadata("name") == metadata("name")
+
+
+def test_repeat_members_using_members():
+    fieldlist, values, metadata = _get_template()
+
+    repeat = RepeatMembers(members=[0, 1, 2])
+    repeated = repeat.forward(fieldlist)
+    assert len(repeated) == 3
+    for i, f in enumerate(repeated):
+        assert f.values.shape == values.shape
+        assert np.all(f.values == values)
+        assert f.metadata("number") == i + 1
+        assert f.metadata("name") == metadata("name")
+
+
+def test_repeat_members_using_count():
+    fieldlist, values, metadata = _get_template()
+
+    repeat = RepeatMembers(count=3)
+    repeated = repeat.forward(fieldlist)
+    assert len(repeated) == 3
+    for i, f in enumerate(repeated):
+        assert f.values.shape == values.shape
+        assert np.all(f.values == values)
+        assert f.metadata("number") == i + 1
+        assert f.metadata("name") == metadata("name")