🎨 Move filter registration to CLI and rename ExampleFilter to DoublingFilter

Author: jemrobinson

ice_station_zebra/config/base.yaml

@@ -1,7 +1,8 @@
 defaults:
   - datasets:
     - samp_sicsouth_osisaf_25k_2020_2020_24h_v1
-    - samp_weathersouth_era5_0p5_2020_2020_24h_v1
+    # - samp_weathersouth_era5_0p5_2020_2020_24h_v1
+    - nbsamp_weathersouth_era5_0p5_2020_2020_24h_v1
   - evaluate: default
   - loggers:
     - wandb

ice_station_zebra/config/datasets/nbsamp_weathersouth_era5_0p5_2020_2020_24h_v1.yaml

@@ -11,7 +11,11 @@ nbsamp-weathersouth-era5-0p5-2020-2020-24h-v1:
     frequency: 24h
 
   input:
-    anemoi-dataset:
-      join:
-      - dataset: ${base_path}/data/anemoi/samp-weathersouth-era5-0p5-2020-2020-24h-v1.zarr
-        select: [ 2t, z_250, z_500 , 10u, 10v, cos_julian_day, sin_julian_day ]
+    pipe:
+      - anemoi-dataset:
+          join:
+          - dataset: ${base_path}/data/anemoi/samp-weathersouth-era5-0p5-2020-2020-24h-v1.zarr
+            select: [ 2t, z_250, z_500 , 10u, 10v, cos_julian_day, sin_julian_day ]
+      - doubling_filter:
+          input_field: 2t
+          output_field: doubled_2t

ice_station_zebra/data_processors/cli.py

@@ -5,6 +5,7 @@
 
 from ice_station_zebra.cli import hydra_adaptor
 
+from .filters import register_filters
 from .zebra_data_processor_factory import ZebraDataProcessorFactory
 
 # Create the typer app
@@ -17,6 +18,7 @@
 @hydra_adaptor
 def create(config: DictConfig) -> None:
     """Create all datasets."""
+    register_filters()
     factory = ZebraDataProcessorFactory(config)
     for dataset in factory.datasets:
         logger.info("Working on %s.", dataset.name)

ice_station_zebra/data_processors/filters/__init__.py

@@ -1,5 +1,13 @@
-from .example_filter import ExampleFilter
+from anemoi.transform.filters import filter_registry
+
+from .doubling_filter import DoublingFilter
+
+
+# Register all filers with anemoi-transform
+def register_filters() -> None:
+    filter_registry.register("doubling_filter", DoublingFilter)
+
 
 __all__ = [
-    "ExampleFilter",
+    "DoublingFilter",
 ]

…ata_processors/filters/example_filter.py …ta_processors/filters/doubling_filter.pyice_station_zebra/data_processors/filters/example_filter.py renamed to ice_station_zebra/data_processors/filters/doubling_filter.py ice_station_zebra/data_processors/filters/example_filter.py renamed to ice_station_zebra/data_processors/filters/doubling_filter.py

@@ -1,29 +1,24 @@
 from collections.abc import Iterator
 
-import earthkit.data as ekd  # type: ignore[import-untyped]
+import earthkit.data as ekd
 from anemoi.transform.filters.matching import MatchingFieldsFilter, matching
 
 
-# @filter_registry.register("example_filter")
-class ExampleFilter(MatchingFieldsFilter):
-    @matching(
-        select="param",
-        forward=["input_field"],
-    )
+class DoublingFilter(MatchingFieldsFilter):
+    @matching(select="param", forward=["input_field"])
     def __init__(
         self,
         *,
         input_field: str = "msl",
+        output_field: str = "doubled_msl",
     ) -> None:
-        """Create an example filter."""
+        """An example filter that takes two arguments."""
         self.input_field = input_field
+        self.output_field = output_field
 
     def forward_transform(self, input_field: ekd.Field) -> Iterator[ekd.Field]:
-        """An example forward transform that doubles the input field."""
+        """An example forward transform that doubles the input field as a new field."""
         yield input_field
         yield self.new_field_from_numpy(
-            input_field.to_numpy() * 2, template=input_field, param="test_field"
+            input_field.to_numpy() * 2, template=input_field, param=self.output_field
         )
-
-    def backward_transform(self, input_field: ekd.Field) -> Iterator[ekd.Field]:
-        yield input_field

ice_station_zebra/data_processors/zebra_data_processor.py

@@ -4,13 +4,11 @@
 
 from anemoi.datasets.commands.create import Create
 from anemoi.datasets.commands.inspect import InspectZarr
-from anemoi.transform.filters import filter_registry
 from omegaconf import DictConfig, OmegaConf
 from zarr.errors import PathNotFoundError
 
 from ice_station_zebra.types import AnemoiCreateArgs, AnemoiInspectArgs
 
-from .filters.example_filter import ExampleFilter
 from .preprocessors import IPreprocessor
 
 logger = logging.getLogger(__name__)
@@ -48,7 +46,6 @@ def create(self) -> None:
 
     def download(self) -> None:
         """Download a single Anemoi dataset."""
-        filter_registry.register("example_filter", ExampleFilter)
         self.preprocessor.download(self.path_preprocessor)
         logger.info("Creating dataset %s at %s.", self.name, self.path_dataset)
         Create().run(

ice_station_zebra/data_processors/zebra_data_processor_factory.py

@@ -24,5 +24,5 @@ def __init__(self, config: DictConfig) -> None:
                 .get("type", "None")
             ]
             self.datasets.append(
-                ZebraDataProcessor(dataset_name, config, cls_preprocessor)
+                ZebraDataProcessor(dataset_name, config, cls_preprocessor)  # type: ignore[type-abstract]
             )

pyproject.toml

@@ -62,6 +62,7 @@ omit= [
 module = [
     "anemoi.datasets.*",
     "anemoi.transform.*",
+    "earthkit.*",
     "icenet.data.sic.*",
     "zarr.*",
 ]