Merge branch 'ecmwf:develop' into feature/lambda-filter

Author: frazane

CHANGELOG.md

@@ -8,7 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 Please add your functional changes to the appropriate section in the PR.
 Keep it human-readable, your future self will thank you!
 
-## [Unreleased](https://github.com/ecmwf/anemoi-utils/transform/0.0.5...HEAD/compare/0.0.8...HEAD)
+## [Unreleased](https://github.com/ecmwf/anemoi-utils/transform/0.0.5...HEAD/compare/0.1.0...HEAD)
+
+- 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
 
 ## [0.0.8](https://github.com/ecmwf/anemoi-utils/transform/0.0.5...HEAD/compare/0.0.5...0.0.8) - 2024-10-26
 

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/test_repeat_members.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")