[weatherbench2] Add GaussianVariance metric.

PiperOrigin-RevId: 592111655

Author: ilopezgp

weatherbench2/metrics.py

@@ -712,15 +712,15 @@ def _rank_da(da: xr.DataArray) -> np.ndarray:
 
 @dataclasses.dataclass
 class GaussianCRPS(Metric):
-  """The spread measure associated with CRPS, E|X - X'|."""
+  """The analytical formulation of CRPS for a Gaussian."""
 
   def compute_chunk(
       self,
       forecast: xr.Dataset,
       truth: xr.Dataset,
       region: t.Optional[Region] = None,
   ) -> xr.Dataset:
-    """CRPSSpread, averaged over space, for a time chunk of data."""
+    """GaussianCRPS, averaged over space, for a time chunk of data."""
     return _spatial_average(
         _pointwise_gaussian_crps(forecast, truth),
         region=region,
@@ -770,6 +770,33 @@ def _pointwise_gaussian_crps(
   return xr.Dataset(dataset, coords=forecast.coords)
 
 
+@dataclasses.dataclass
+class GaussianVariance(Metric):
+  """The variance of a Gaussian forecast."""
+
+  def compute_chunk(
+      self,
+      forecast: xr.Dataset,
+      truth: xr.Dataset,
+      region: t.Optional[Region] = None,
+  ) -> xr.Dataset:
+    """GaussianVariance, averaged over space, for a time chunk of data."""
+    del truth  # unused
+    var_list = []
+    dataset = {}
+    for var in forecast.keys():
+      if f"{var}_std" in forecast.keys():
+        var_list.append(var)
+    for var_name in var_list:
+      variance = forecast[f"{var_name}_std"] * forecast[f"{var_name}_std"]
+      dataset[var_name] = variance
+
+    return _spatial_average(
+        xr.Dataset(dataset, coords=forecast.coords),
+        region=region,
+    )
+
+
 @dataclasses.dataclass
 class EnsembleStddevSqrtBeforeTimeAvg(EnsembleMetric):
   """The standard deviation of an ensemble of forecasts.

weatherbench2/metrics_test.py

@@ -259,6 +259,28 @@ def test_gaussian_crps(self):
     np.testing.assert_allclose(result['2m_temperature'].values, expected)
 
 
+class GaussianVarianceTest(parameterized.TestCase):
+
+  def test_gaussian_variance(self):
+    forecast = schema.mock_forecast_data(
+        variables_3d=[],
+        variables_2d=['2m_temperature', '2m_temperature_std'],
+        time_start='2022-01-01',
+        time_stop='2022-01-02',
+        lead_stop='1 day',
+    )
+    truth = schema.mock_truth_data(
+        variables_3d=[],
+        variables_2d=['2m_temperature'],
+        time_start='2022-01-01',
+        time_stop='2022-01-20',
+    )
+    forecast['2m_temperature_std'] = forecast['2m_temperature_std'] + 1.0
+    result = metrics.GaussianVariance().compute(forecast, truth)
+    expected = np.array([1.0, 1.0])
+    np.testing.assert_allclose(result['2m_temperature'].values, expected)
+
+
 class RankHistogramTest(parameterized.TestCase):
 
   @parameterized.named_parameters(