Store analysis in verif files

Author: tnipen

bris/outputs/verif.py

@@ -225,11 +225,13 @@ def finalize(self) -> None:
         frts = self.intermediate.get_forecast_reference_times()
         frts_unix = utils.datetime_to_unixtime(frts).astype(np.double)
 
+        leadtimes_seconds = self.intermediate.pm.leadtimes.astype(np.float32)
+
         coords = {}
         coords["time"] = (["time"], frts_unix, cf.get_attributes("time"))
         coords["leadtime"] = (
             ["leadtime"],
-            self.intermediate.pm.leadtimes.astype(np.float32) / 3600,
+            leadtimes_seconds / 3600,
             {"units": "hour"},
         )
         assert len(self.obs_ids) == len(self.opoints.get_lats()), (
@@ -345,15 +347,14 @@ def finalize(self) -> None:
                         self.ds["x"] = (["time", "leadtime", "location", "quantile"], x)
 
         # Find which valid times we need observations for
-        frts_ut = utils.datetime_to_unixtime(frts)
-        a, b = np.meshgrid(frts_ut, np.array(self.intermediate.pm.leadtimes))
+        # Build a 2D array of valid times for (init_time, leadtime)
+        a, b = np.meshgrid(frts_unix, np.array(self.intermediate.pm.leadtimes))
         valid_times = a + b
         valid_times = valid_times.transpose()
         if len(valid_times) == 0:
             utils.LOGGER.warning("Could not finalize verif, no valid times")
             return
 
-        # valid_times = np.sort(np.unique(valid_times.flatten()))
         unique_valid_times = np.sort(np.unique(valid_times.flatten()))
 
         start_time = int(np.min(unique_valid_times))
@@ -387,6 +388,21 @@ def finalize(self) -> None:
 
         self.ds["obs"] = (["time", "leadtime", "location"], obs)
 
+        if int(leadtimes_seconds[0]) == 0:
+            # Store the analysis state from "fcst", where it is available. Note: this picks the first
+            # forecast leadtime, which isn't necesssarily going to be an analysis.
+            analysis = self.create_nan_array(fcst_shape)
+            analysis_valid_times = frts_unix
+
+            for _t, valid_time in enumerate(frts_unix):
+                Itimes, Ileadtimes = np.where(valid_times == valid_time)
+                for i in range(len(Itimes)):
+                    analysis[Itimes[i], Ileadtimes[i], :] = fcst[_t, 0, :]
+
+            self.ds["analysis"] = (["time", "leadtime", "location"], analysis)
+        else:
+            print("Not writing analysis to verif file. The first forecast leadtime isn't 0")
+
         if self.num_members > 1:
             crps = self.compute_crps(ens, obs, self.fair_crps)
             self.ds["ensemble_crps"] = (["time", "leadtime", "location"], crps)
@@ -399,6 +415,7 @@ def finalize(self) -> None:
 
         data_variables = [
             "obs",
+            "analysis",
             "fcst",
             "ensemble",
             "ensemble_mean",

tests/test_outputs_verif.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pytest
+import xarray as xr
 
 from bris.outputs import Verif
 from bris.predict_metadata import PredictMetadata
@@ -64,6 +65,7 @@ def test_1():
                     output.add_forecast(times, member, pred)
 
                 output.finalize()
+                check_expected_variable(ofilename)
 
         altitudes = np.arange(len(lats))
         pm = PredictMetadata(
@@ -90,6 +92,15 @@ def test_1():
                 output.add_forecast(times, member, pred)
 
             output.finalize()
+            check_expected_variable(ofilename)
+
+def check_expected_variable(filename):
+    expected_variables = ["analysis", "fcst", "obs", "ensemble", "ensemble_mean",
+            "ensemble_variance", "ensemble_crps"]
+
+    with xr.open_dataset(filename) as file:
+        for variable in expected_variables:
+            assert variable in file, variable
 
 
 def test_2():