feat(input): add FDB input class (#190)

Author: frazane

src/anemoi/inference/inputs/fdb.py

@@ -0,0 +1,88 @@
+# (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 datetime
+import logging
+from typing import Any
+from typing import List
+from typing import Optional
+
+import earthkit.data as ekd
+import numpy as np
+
+from ..types import Date
+from ..types import State
+from . import input_registry
+from .grib import GribInput
+
+LOG = logging.getLogger(__name__)
+
+
+@input_registry.register("fdb")
+class FDBInput(GribInput):
+    """Get input fields from FDB."""
+
+    trace_name = "fdb"
+
+    def __init__(
+        self,
+        context,
+        *,
+        namer=None,
+        fdb_config: dict | None = None,
+        fdb_userconfig: dict | None = None,
+        **kwargs: dict[str, Any],
+    ):
+        """Initialise the FDB input.
+
+        Parameters
+        ----------
+        context : dict
+            The context runner.
+        namer : optional
+            The namer to use for the input.
+        fdb_config : dict, optional
+            The FDB config to use.
+        fdb_userconfig : dict, optional
+            The FDB userconfig to use.
+        kwargs : dict, optional
+            Additional keyword arguments for the request to FDB.
+        """
+        super().__init__(context, namer=namer)
+        self.kwargs = kwargs
+        self.configs = {"config": fdb_config, "userconfig": fdb_userconfig}
+        # NOTE: this is a temporary workaround for #191 thus not documented
+        self.param_id_map = kwargs.pop("param_id_map", {})
+        self.variables = self.checkpoint.variables_from_input(include_forcings=False)
+
+    def create_input_state(self, *, date: Optional[Date]) -> 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)
+        return res
+
+    def load_forcings_state(self, *, variables: List[str], dates: List[Date], current_state: State) -> State:
+        ds = self.retrieve(variables=variables, dates=dates)
+        return self._load_forcings_state(ds, variables=variables, dates=dates, current_state=current_state)
+
+    def retrieve(self, variables: List[str], dates: List[Date]) -> Any:
+        requests = self.checkpoint.mars_requests(
+            variables=variables,
+            dates=dates,
+            use_grib_paramid=self.context.use_grib_paramid,
+            patch_request=self.context.patch_data_request,
+        )
+        requests = [self.kwargs | r for r in requests]
+        # NOTE: this is a temporary workaround for #191
+        for request in requests:
+            request["param"] = [self.param_id_map.get(p, p) for p in request["param"]]
+        sources = [ekd.from_source("fdb", request, stream=False, **self.configs) for request in requests]
+        ds = ekd.from_source("multi", sources)
+        return ds