feat: Add support for interpolator

src/anemoi/inference/input.py

@@ -113,13 +113,15 @@ def __repr__(self) -> str:
             return f"{self.__class__.__name__}({self.purpose})"
 
     @abstractmethod
-    def create_input_state(self, *, date: Date | None) -> State:
+    def create_input_state(self, *, date: Date | None, **kwargs) -> State:
         """Create the input state dictionary.
 
         Parameters
         ----------
         date : Optional[Date]
             The date for which to create the input state.
+        **kwargs : Any
+            Additional keyword arguments.
 
         Returns
         -------

src/anemoi/inference/inputs/cds.py

@@ -138,13 +138,15 @@ def __init__(
         self.dataset = dataset
         self.kwargs = kwargs
 
-    def create_input_state(self, *, date: Date | None) -> State:
+    def create_input_state(self, *, date: Date | None, **kwargs) -> State:
         """Create the input state for the given date.
 
         Parameters
         ----------
         date : Optional[Date]
             The date for which to create the input state.
+        **kwargs : Any
+            Additional keyword arguments.
 
         Returns
         -------
@@ -162,6 +164,7 @@ def create_input_state(self, *, date: Date | None) -> State:
             ),
             variables=self.variables,
             date=date,
+            **kwargs,
         )
 
     def retrieve(self, variables: list[str], dates: list[Date]) -> Any:

src/anemoi/inference/inputs/cutout.py

@@ -11,20 +11,31 @@
 import logging
 from collections import defaultdict
 from collections.abc import Iterable
+from collections.abc import Mapping
 
 import numpy as np
 
 from anemoi.inference.input import Input
 from anemoi.inference.inputs import create_input
+from anemoi.inference.inputs import input_registry
 from anemoi.inference.types import Date
 from anemoi.inference.types import ProcessorConfig
 from anemoi.inference.types import State
 
-from . import input_registry
-
 LOG = logging.getLogger(__name__)
 
 
+def contains_key(obj, key: str) -> bool:
+    """Recursively check if `key` exists anywhere in a nested config (dict/DotDict/lists)."""
+    if isinstance(obj, Mapping):
+        if key in obj:
+            return True
+        return any(contains_key(v, key) for v in obj.values())
+    if isinstance(obj, (list, tuple, set)):
+        return any(contains_key(v, key) for v in obj)
+    return False
+
+
 def _mask_and_combine_states(
     existing_state: State,
     new_state: State,
@@ -138,9 +149,12 @@ def __init__(
                 cfg = cfg.copy()
                 mask = cfg.pop("mask", f"{src}/cutout_mask")
 
-            self.sources[src] = create_input(
-                context, cfg, variables=variables, pre_processors=pre_processors, purpose=purpose
-            )
+            if contains_key(cfg, "pre_processors"):
+                self.sources[src] = create_input(context, cfg, variables=variables, purpose=purpose)
+            else:
+                self.sources[src] = create_input(
+                    context, cfg, variables=variables, purpose=purpose, pre_processors=pre_processors
+                )
 
             if isinstance(mask, str):
                 self.masks[src] = self.sources[src].checkpoint.load_supporting_array(mask)
@@ -151,13 +165,15 @@ def __repr__(self):
         """Return a string representation of the Cutout object."""
         return f"Cutout({self.sources})"
 
-    def create_input_state(self, *, date: Date | None) -> State:
+    def create_input_state(self, *, date: Date | None, **kwargs) -> State:
         """Create the input state for the given date.
 
         Parameters
         ----------
         date : Optional[Date]
             The date for which to create the input state.
+        **kwargs : dict
+            Additional keyword arguments for the source input state creation.
 
         Returns
         -------
@@ -173,7 +189,7 @@ def create_input_state(self, *, date: Date | None) -> State:
         combined_state = {}
 
         for source in self.sources.keys():
-            source_state = self.sources[source].create_input_state(date=date)
+            source_state = self.sources[source].create_input_state(date=date, **kwargs)
             source_mask = self.masks[source]
 
             # Create the mask front padded with zeros

src/anemoi/inference/inputs/dataset.py

@@ -104,13 +104,15 @@ def __repr__(self) -> str:
         """Return a string representation of the DatasetInput."""
         return f"DatasetInput({self.open_dataset_args}, {self.open_dataset_kwargs})"
 
-    def create_input_state(self, *, date: Date | None = None) -> State:
+    def create_input_state(self, *, date: Date | None = None, **kwargs) -> State:
         """Create the input state for the given date.
 
         Parameters
         ----------
         date : Optional[Any]
             The date for which to create the input state.
+        **kwargs : Any
+            Additional keyword arguments.
 
         Returns
         -------

src/anemoi/inference/inputs/ekd.py

@@ -111,8 +111,6 @@ def __init__(
         ----------
         context : Any
             The context in which the input is used.
-        pre_processors : Optional[List[ProcessorConfig]], default None
-            Pre-processors to apply to the input
         namer : Optional[Union[Callable[[Any, Dict[str, Any]], str], Dict[str, Any]]]
             Optional namer for the input.
         """
@@ -128,12 +126,12 @@ def __init__(
         assert callable(self._namer), type(self._namer)
 
     def _filter_and_sort(self, data: Any, *, dates: list[Any], title: str) -> Any:
-        """Filter and sort the data.
+        """Filter and sort the data (earthkit FieldList/FieldArray).
 
         Parameters
         ----------
         data : Any
-            The data to filter and sort.
+            The data to filter and sort (FieldList or FieldArray).
         dates : List[Any]
             The list of dates to select.
         title : str
@@ -142,7 +140,7 @@ def _filter_and_sort(self, data: Any, *, dates: list[Any], title: str) -> Any:
         Returns
         -------
         Any
-            The filtered and sorted data.
+            The filtered and sorted data (FieldArray).
         """
 
         def _name(field: Any, _: Any, original_metadata: dict[str, Any]) -> str:
@@ -166,12 +164,12 @@ def _name(field: Any, _: Any, original_metadata: dict[str, Any]) -> str:
         return data
 
     def _find_variable(self, data: Any, name: str, **kwargs: Any) -> Any:
-        """Find a variable in the data.
+        """Find a variable in the data (earthkit FieldList/FieldArray selection).
 
         Parameters
         ----------
         data : Any
-            The data to search.
+            The data to search (FieldList or FieldArray).
         name : str
             The name of the variable to find.
         **kwargs : Any
@@ -180,7 +178,7 @@ def _find_variable(self, data: Any, name: str, **kwargs: Any) -> Any:
         Returns
         -------
         Any
-            The selected variable.
+            The selected variable (FieldArray subset).
         """
 
         def _name(field: Any, _: Any, original_metadata: dict[str, Any]) -> str:
@@ -198,9 +196,17 @@ def _create_state(
         longitudes: FloatArray | None = None,
         dtype: DTypeLike = np.float32,
         flatten: bool = True,
+        ref_date_index: int = -1,
     ) -> State:
         """Create a state from an ekd.FieldList.
 
+        Notes
+        -----
+        - The `fields` argument must be an earthkit FieldList (or FieldArray-compatible).
+        - This method intentionally converts state["fields"] from a FieldList to
+          a Dict[str, np.ndarray] with shape (len(dates), n_points).
+        - Pre-processors are run while state["fields"] is still a FieldList.
+
         Parameters
         ----------
         fields : ekd.FieldList
@@ -215,30 +221,17 @@ def _create_state(
             The data type.
         flatten : bool
             Whether to flatten the data.
+        ref_date_index : int
+            The index of the reference date in the dates list.
 
         Returns
         -------
         State
-            The created input state.
+            The created input state with state["fields"] as Dict[str, np.ndarray].
         """
-        fields = self.pre_process(fields)
-
-        dates = sorted([to_datetime(d) for d in dates])
-        date_to_index = {d.isoformat(): i for i, d in enumerate(dates)}
-
-        state = dict(date=dates[-1], latitudes=latitudes, longitudes=longitudes, fields=dict())
-
-        if len(fields) == 0:
-            LOG.warning("No input fields found for dates %s (%s)", [d.isoformat() for d in dates], self)
-            return state
-
-        state_fields = state["fields"]
-
-        fields = self._filter_and_sort(fields, dates=dates, title="Create input state")
-
         if latitudes is None and longitudes is None:
             try:
-                state["latitudes"], state["longitudes"] = fields[0].grid_points()
+                latitudes, longitudes = fields[0].grid_points()
                 LOG.info(
                     "%s: using `latitudes` and `longitudes` from the first input field",
                     self.__class__.__name__,
@@ -251,8 +244,6 @@ def _create_state(
                 latitudes = self.checkpoint.latitudes
                 longitudes = self.checkpoint.longitudes
                 if latitudes is not None and longitudes is not None:
-                    state["latitudes"] = latitudes
-                    state["longitudes"] = longitudes
                     LOG.info(
                         "%s: using `latitudes` and `longitudes` found in the checkpoint.",
                         self.__class__.__name__,
@@ -264,6 +255,21 @@ def _create_state(
                     )
                     raise e
 
+        state = dict(date=dates[ref_date_index], latitudes=latitudes, longitudes=longitudes, fields=fields)
+
+        # allow hooks to operate on the FieldList before conversion to numpy
+        state = self.pre_process(state)
+
+        fields = state["fields"]
+        state_fields = {}
+
+        if len(fields) == 0:
+            raise ValueError("No input fields provided")
+
+        dates = sorted([to_datetime(d) for d in dates])
+        date_to_index = {d.isoformat(): i for i, d in enumerate(dates)}
+        fields = self._filter_and_sort(fields, dates=dates, title="Create input state")
+
         check = defaultdict(set)
 
         n_points = fields[0].to_numpy(dtype=dtype, flatten=flatten).size
@@ -295,7 +301,7 @@ def _create_state(
                 raise ValueError(f"Duplicate dates for {name}")
 
             check[name].add(date_idx)
-
+        state["fields"] = state_fields
         for name, idx in check.items():
             if len(idx) != len(dates):
                 LOG.error("Missing dates for %s: %s", name, idx)
@@ -326,6 +332,7 @@ def _create_input_state(
         longitudes: FloatArray | None = None,
         dtype: DTypeLike = np.float32,
         flatten: bool = True,
+        ref_date_index: int = -1,
     ) -> State:
         """Create the input state.
 
@@ -345,6 +352,8 @@ def _create_input_state(
             The data type.
         flatten : bool
             Whether to flatten the data.
+        ref_date_index : int
+            The index of the reference date in the dates list.
 
         Returns
         -------
@@ -366,6 +375,7 @@ def _create_input_state(
             longitudes=longitudes,
             dtype=dtype,
             flatten=flatten,
+            ref_date_index=ref_date_index,
         )
 
     def _load_forcings_state(self, fields: ekd.FieldList, *, dates: list[Date], current_state: State) -> State:
@@ -397,7 +407,7 @@ def _load_forcings_state(self, fields: ekd.FieldList, *, dates: list[Date], curr
     def set_private_attributes(self, state: State, fields: ekd.FieldList) -> None:  # type: ignore
         """Set private attributes to the state.
 
-        Provides geography information if available retrieved from the fields.
+        Provides geography information if available retrieved from the fields (FieldList/FieldArray).
         """
         geography_information = {}
 

src/anemoi/inference/inputs/empty.py

@@ -46,13 +46,15 @@ def __init__(self, context: Context, **kwargs: Any) -> None:
         super().__init__(context, **kwargs)
         assert self.variables in (None, []), "EmptyInput should not have variables"
 
-    def create_input_state(self, *, date: Date | None) -> State:
+    def create_input_state(self, *, date: Date | None, **kwargs) -> State:
         """Create an empty input state.
 
         Parameters
         ----------
         date : Date or None
             The date for the input state.
+        **kwargs : Any
+            Additional keyword arguments.
 
         Returns
         -------

src/anemoi/inference/inputs/fdb.py

@@ -55,11 +55,11 @@ def __init__(
         # NOTE: this is a temporary workaround for #191 thus not documented
         self.param_id_map = kwargs.pop("param_id_map", {})
 
-    def create_input_state(self, *, date: Date | None) -> State:
+    def create_input_state(self, *, date: Date | None, **kwargs) -> State:
         date = np.datetime64(date).astype(datetime.datetime)
         dates = [date + h for h in self.checkpoint.lagged]
         ds = self.retrieve(variables=self.variables, dates=dates)
-        res = self._create_input_state(ds, variables=None, date=date)
+        res = self._create_input_state(ds, variables=None, date=date, **kwargs)
         return res
 
     def load_forcings_state(self, *, dates: list[Date], current_state: State) -> State: