Improved output of initial state

CHANGELOG.md

@@ -11,14 +11,15 @@ Keep it human-readable, your future self will thank you!
 ## [Unreleased]
 
 ### Added
-
+- Add initial state output in netcdf format
 - Fix: Enable inference when no constant forcings are used
 - Add anemoi-transform link to documentation
 - Add support for unstructured grids
 - Add CONTRIBUTORS.md file (#36)
 
 ### Changed
-
+- Raw output of initial state contains only values at initial time
+- Changed default naming of raw output
 - Add cos_solar_zenith_angle to list of known forcings
 - Add missing classes in checkpoint handling
 - Rename Condition to State [#24](https://github.com/ecmwf/anemoi-inference/pull/24)

src/anemoi/inference/output.py

@@ -6,6 +6,7 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 #
+import copy
 from abc import ABC
 from abc import abstractmethod
 
@@ -28,5 +29,12 @@ def write_initial_state(self, state):
     def write_state(self, state):
         pass
 
+    def reduce(self, state):
+        """Creates new state which is projection of original state on the last step in the multi-steps dimension."""
+        reduced_state = copy.deepcopy(state)
+        for field, values in state["fields"].items():
+            reduced_state["fields"][field] = values[-1, :]
+        return reduced_state
+
     def close(self):
         pass

src/anemoi/inference/outputs/netcdf.py

@@ -53,12 +53,18 @@ def _init(self, state):
 
         values = len(state["latitudes"])
 
-        self.values_dim = self.ncfile.createDimension("values", values)
+        time = 0
+        self.reference_date = state["date"]
+        if hasattr(self.context, "time_step") and hasattr(self.context, "lead_time"):
+            time = self.context.lead_time // self.context.time_step
+        if hasattr(self.context, "reference_date"):
+            self.reference_date = self.context.reference_date
 
-        self.time_dim = self.ncfile.createDimension("time", self.context.lead_time // self.context.time_step)
+        self.values_dim = self.ncfile.createDimension("values", values)
+        self.time_dim = self.ncfile.createDimension("time", time)
         self.time_var = self.ncfile.createVariable("time", "i4", ("time",), **compression)
 
-        self.time_var.units = "seconds since {0}".format(self.context.reference_date)
+        self.time_var.units = "seconds since {0}".format(self.reference_date)
         self.time_var.long_name = "time"
         self.time_var.calendar = "gregorian"
 
@@ -99,12 +105,13 @@ def _init(self, state):
         return self.ncfile
 
     def write_initial_state(self, state):
-        LOG.warning("NetCDFOutput: Writing of initial state is not supported.")
+        reduced_state = self.reduce(state)
+        self.write_state(reduced_state)
 
     def write_state(self, state):
         self._init(state)
 
-        step = state["date"] - self.context.reference_date
+        step = state["date"] - self.reference_date
         self.time_var[self.n] = step.total_seconds()
 
         for name, value in state["fields"].items():

src/anemoi/inference/outputs/raw.py

@@ -24,20 +24,30 @@
 class RawOutput(Output):
     """_summary_"""
 
-    def __init__(self, context, path):
+    def __init__(
+        self,
+        context,
+        path,
+        template="{date}.npz",
+        strftime="%Y%m%d%H%M%S",
+    ):
         super().__init__(context)
         self.path = path
+        self.template = template
+        self.strftime = strftime
 
     def __repr__(self):
         return f"RawOutput({self.path})"
 
     def write_initial_state(self, state):
-        self.write_state(state)
+        reduced_state = self.reduce(state)
+        self.write_state(reduced_state)
 
     def write_state(self, state):
         os.makedirs(self.path, exist_ok=True)
-        fn_state = f"{self.path}/{state['date'].strftime('%Y%m%d_%H')}"
+        date = state["date"].strftime(self.strftime)
+        fn_state = f"{self.path}/{self.template.format(date=date)}"
         restate = {f"field_{key}": val for key, val in state["fields"].items()}
-        restate["longitudes"] = state["longitudes"]
-        restate["latitudes"] = state["latitudes"]
+        for key in ["date", "longitudes", "latitudes"]:
+            restate[key] = np.array(state[key], dtype=str)
         np.savez_compressed(fn_state, **restate)