feat: ability for apply-mask to work on a single param (#251)

## Description
Alternative implementation of #249. Add the ability for the `apply_mask`
filter to work only on a single param.

***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/fields/apply_mask.py

@@ -84,11 +84,13 @@ class MaskVariable(SingleFieldFilter):
     # threshold: float | None,
     # threshold_operator: Literal["<", "<=", ">", ">=", "==", "!="],
     # rename: str | None,
+    # param: str | None,
     optional_inputs = {
         "mask_value": None,
         "threshold": None,
         "threshold_operator": ">",
         "rename": None,
+        "param": None,
     }
 
     def prepare_filter(self):
@@ -120,6 +122,11 @@ def prepare_filter(self):
         else:
             self.mask = mask == self.mask_value
 
+    def forward_select(self):
+        if self.param is not None:
+            return {"param": self.param}
+        return {}
+
     def forward_transform(self, field: ekd.Field) -> ekd.Field:
         """Apply the forward transformation to the field.
 

tests/field_filters/test_apply_mask.py

@@ -102,10 +102,43 @@ def test_apply_mask(source, ekd_from_source, mask_name, rename, threshold_option
             expected_values = input_field.to_numpy(flatten=True).copy()
             expected_values[expected_mask] = np.nan
             result = output_field.to_numpy(flatten=True)
-            np.array_equal(expected_values, result, equal_nan=True)
+            assert np.array_equal(expected_values, result, equal_nan=True)
             assert np.sum(np.isnan(result)) == expected_mask_count
 
 
+def test_apply_mask_only_single_param(source, ekd_from_source):
+    apply_mask = filter_registry.create(
+        "apply_mask",
+        path="mixed_floats",
+        threshold=0.5,
+        threshold_operator=">",
+        param="t",
+    )
+    ekd_from_source.assert_called_once_with("file", "mixed_floats")
+
+    pipeline = source | apply_mask
+
+    input_fields = collect_fields_by_param(source)
+    output_fields = collect_fields_by_param(pipeline)
+
+    expected_mask = MASK_VALUES["mixed_floats"].copy().flatten()
+    expected_mask = np.greater(expected_mask, 0.5)
+    expected_mask_count = np.sum(expected_mask)
+
+    for param in DATA_VALUES.keys():
+        assert param in output_fields
+        for input_field, output_field in zip(input_fields[param], output_fields[param]):
+            if param == "t":
+                # only mask t
+                expected_values = input_field.to_numpy(flatten=True).copy()
+                expected_values[expected_mask] = np.nan
+                result = output_field.to_numpy(flatten=True)
+                assert np.array_equal(expected_values, result, equal_nan=True)
+                assert np.sum(np.isnan(result)) == expected_mask_count
+            else:
+                assert np.array_equal(input_field.to_numpy(flatten=True), output_field.to_numpy(flatten=True))
+
+
 if __name__ == "__main__":
     """Run all test functions that start with 'test_'."""
     for name, obj in list(globals().items()):