refactor: rescale/convert filters to SingleFieldFilter (#226)

## Description
Converts rescale and convert filters to SingleFieldFilters, changes
tests to ensure non-matching fields are not processed, and breaks up the
class hierarchy to avoid instantiating non-leaf class .

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

Author: aaron-hopkinson

src/anemoi/transform/filters/rescale.py

@@ -6,80 +6,64 @@
 # 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.abc import Iterator
+from abc import ABC
+from abc import abstractmethod
+from typing import Callable
 
 import earthkit.data as ekd
 
+from anemoi.transform.filter import SingleFieldFilter
 from anemoi.transform.filters import filter_registry
-from anemoi.transform.filters.matching import MatchingFieldsFilter
-from anemoi.transform.filters.matching import matching
 
 
-class Rescale(MatchingFieldsFilter):
+class Rescaler:
+    def __init__(self, scale, offset):
+        self.scale = scale
+        self.offset = offset
+
+    def forward(self, x):
+        return x * self.scale + self.offset
+
+    def backward(self, x):
+        return (x - self.offset) / self.scale
+
+
+class RescaleMixin(ABC):
+    # inheriting classes should define required_inputs (which must include param)
+    # and must define self.rescaler in prepare_filter
+    param: str
+    rescaler: Rescaler
+    # intended to be inherited from SingleFieldFilter
+    new_field_from_numpy: Callable
+
+    @abstractmethod
+    def prepare_filter(self):
+        raise NotImplementedError("prepare_filter must be implemented by subclasses.")
+
+    def forward_select(self):
+        return {"param": self.param}
+
+    def forward_transform(self, param: ekd.Field) -> ekd.Field:
+        """Apply the forward transformation (x to ax+b)."""
+        rescaled = self.rescaler.forward(param.to_numpy())
+        return self.new_field_from_numpy(rescaled, template=param, param=self.param)
+
+    def backward_transform(self, param: ekd.Field) -> ekd.Field:
+        """Apply the backward transformation (ax+b to x)."""
+        descaled = self.rescaler.backward(param.to_numpy())
+        return self.new_field_from_numpy(descaled, template=param, param=self.param)
+
+
+class Rescale(RescaleMixin, SingleFieldFilter):
     """A filter to rescale a parameter from a scale and an offset, and back."""
 
-    @matching(
-        select="param",
-        forward=("param",),
-        backward=("param",),
-    )
-    def __init__(
-        self,
-        *,
-        scale: float,
-        offset: float,
-        param: str,
-    ) -> None:
-        """Parameters
-        -------------
-        scale : float
-            The scale factor.
-        offset : float
-            The offset value.
-        param : str
-            The parameter to be rescaled.
-        """
+    required_inputs = ("scale", "offset", "param")
 
-        self.scale = scale
-        self.offset = offset
-        self.param = param
-
-    def forward_transform(self, param: ekd.Field) -> Iterator[ekd.Field]:
-        """Apply the forward transformation (x to ax+b).
-
-        Parameters
-        ----------
-        param : ekd.Field
-            The input data to be transformed.
-
-        Returns
-        -------
-        Iterator[ekd.Field]
-            A generator yielding the transformed data.
-        """
-        rescaled = param.to_numpy() * self.scale + self.offset
-        yield self.new_field_from_numpy(rescaled, template=param, param=self.param)
-
-    def backward_transform(self, param: ekd.Field) -> Iterator[ekd.Field]:
-        """Apply the backward transformation (ax+b to x).
-
-        Parameters
-        ----------
-        param : ekd.Field
-            The input data to be transformed.
-
-        Returns
-        -------
-        Iterator[ekd.Field]
-            A generator yielding the transformed data.
-        """
-        descaled = (param.to_numpy() - self.offset) / self.scale
-        yield self.new_field_from_numpy(descaled, template=param, param=self.param)
-
-
-class Convert(Rescale):
+    def prepare_filter(self):
+        self.rescaler = Rescaler(self.scale, self.offset)
+
+
+class Convert(RescaleMixin, SingleFieldFilter):
     """A filter to convert a parameter in a given unit to another unit, and back.
 
     This filter uses :mod:`cfunits` (see the `cfunits documentation <https://ncas-cms.github.io/cfunits/>`_)
@@ -102,26 +86,19 @@ class Convert(Rescale):
 
     """
 
-    def __init__(self, *, unit_in: str, unit_out: str, param: str) -> None:
-        """Parameters
-        -------------
-        unit_in : str
-            The input unit.
-        unit_out : str
-            The output unit.
-        param : str
-            The parameter to be converted.
-        """
+    required_inputs = ("unit_in", "unit_out", "param")
+
+    def prepare_filter(self):
         from cfunits import Units
 
-        u0 = Units(unit_in)
-        u1 = Units(unit_out)
+        u0 = Units(self.unit_in)
+        u1 = Units(self.unit_out)
         x1, x2 = 0.0, 1.0
         y1, y2 = Units.conform([x1, x2], u0, u1)
-        a = (y2 - y1) / (x2 - x1)
-        b = y1 - a * x1
+        scale = (y2 - y1) / (x2 - x1)
+        offset = y1 - scale * x1
 
-        super().__init__(scale=a, offset=b, param=param)
+        self.rescaler = Rescaler(scale, offset)
 
 
 filter_registry.register("rescale", Rescale)

tests/test_rescale.py

@@ -11,10 +11,8 @@
 import numpy.testing as npt
 import pytest
 from anemoi.utils.testing import skip_if_offline
-from pytest import approx
 
-from anemoi.transform.filters.rescale import Convert
-from anemoi.transform.filters.rescale import Rescale
+from anemoi.transform.filters import filter_registry
 
 
 def skip_missing_udunits2():
@@ -40,15 +38,24 @@ def test_rescale(fieldlist: ekd.FieldList) -> None:
         The fieldlist to use for testing.
     """
 
-    fieldlist = fieldlist.sel(param="2t")
+    before_filter = {field.metadata("param"): field.to_numpy().copy() for field in fieldlist}
+
     # rescale from K to °C
-    k_to_deg = Rescale(scale=1.0, offset=-273.15, param="2t")
+    k_to_deg = filter_registry.create("rescale", scale=1.0, offset=-273.15, param="2t")
     rescaled = k_to_deg.forward(fieldlist)
+    after_forward = {field.metadata("param"): field.to_numpy().copy() for field in rescaled}
 
-    npt.assert_allclose(rescaled[0].to_numpy(), fieldlist[0].to_numpy() - 273.15)
     # and back
     rescaled_back = k_to_deg.backward(rescaled)
-    npt.assert_allclose(rescaled_back[0].to_numpy(), fieldlist[0].to_numpy())
+    after_backward = {field.metadata("param"): field.to_numpy().copy() for field in rescaled_back}
+
+    for param in ("2t", "sp"):
+        npt.assert_allclose(before_filter[param], after_backward[param])
+
+        if param == "2t":
+            npt.assert_allclose(before_filter[param] - 273.15, after_forward[param])
+        else:
+            npt.assert_allclose(before_filter[param], after_forward[param])
 
 
 @skip_missing_udunits2()
@@ -61,21 +68,20 @@ def test_convert(fieldlist: ekd.FieldList) -> None:
     fieldlist : ekd.FieldList
         The fieldlist to use for testing.
     """
+    before_filter = {field.metadata("param"): field.to_numpy().copy() for field in fieldlist}
     # rescale from K to °C
-    fieldlist = fieldlist.sel(param="2t")
-    k_to_deg = Convert(unit_in="K", unit_out="degC", param="2t")
+    k_to_deg = filter_registry.create("convert", unit_in="K", unit_out="degC", param="2t")
     rescaled = k_to_deg.forward(fieldlist)
-    assert rescaled[0].values.min() == fieldlist.values.min() - 273.15
-    assert rescaled[0].values.std() == approx(fieldlist.values.std())
+    after_forward = {field.metadata("param"): field.to_numpy().copy() for field in rescaled}
+
     # and back
     rescaled_back = k_to_deg.backward(rescaled)
-    assert rescaled_back[0].values.min() == fieldlist.values.min()
-    assert rescaled_back[0].values.std() == approx(fieldlist.values.std())
+    after_backward = {field.metadata("param"): field.to_numpy().copy() for field in rescaled_back}
 
+    for param in ("2t", "sp"):
+        npt.assert_allclose(before_filter[param], after_backward[param])
 
-if __name__ == "__main__":
-    """Run all test functions that start with 'test_'."""
-    for name, obj in list(globals().items()):
-        if name.startswith("test_") and callable(obj):
-            print(f"Running {name}...")
-            obj()
+        if param == "2t":
+            npt.assert_allclose(before_filter[param] - 273.15, after_forward[param])
+        else:
+            npt.assert_allclose(before_filter[param], after_forward[param])