[weatherbench2] Add Threshold classes, implement GaussianQuantileThreshold in Gaussian metrics.

weatherbench2/metrics.py

@@ -22,6 +22,7 @@
 
 import numpy as np
 from scipy import stats
+from weatherbench2 import thresholds
 from weatherbench2.regions import Region
 import xarray as xr
 
@@ -826,45 +827,16 @@ class GaussianBrierScore(Metric):
     Spatially averaged Brier score for a Gaussian distribution.
   """
 
-  climatology: xr.Dataset
-  threshold: float
+  threshold: thresholds.Threshold
 
   def compute_chunk(
       self,
       forecast: xr.Dataset,
       truth: xr.Dataset,
       region: t.Optional[Region] = None,
   ) -> xr.Dataset:
-    if "init_time" in forecast.dims:
-      time_dim = "valid_time"
-    else:
-      time_dim = "time"
-    climatology_chunk = _get_climatology_chunk(self.climatology, truth)
-    clim_std_dict = {key + "_std": key for key in truth.keys()}  # pytype: disable=unsupported-operands
-    try:
-      climatology_std_chunk = self.climatology[
-          list(clim_std_dict.keys())
-      ].rename(clim_std_dict)
-    except KeyError as e:
-      not_found_stds = set(clim_std_dict).difference(self.climatology.data_vars)
-      raise KeyError(
-          f"Did not find {not_found_stds} forecast keys in climatology."
-      ) from e
 
-    if hasattr(forecast, "level"):
-      climatology_chunk = climatology_chunk.sel(level=forecast.level)
-      climatology_std_chunk = climatology_std_chunk.sel(level=forecast.level)
-
-    time_selection = dict(dayofyear=forecast[time_dim].dt.dayofyear)
-    if "hour" in set(climatology_chunk.coords):
-      time_selection["hour"] = forecast[time_dim].dt.hour
-
-    climatology_chunk = climatology_chunk.sel(time_selection).compute()
-    climatology_std_chunk = climatology_std_chunk.sel(time_selection).compute()
-    threshold = (
-        climatology_chunk
-        + stats.norm.ppf(self.threshold) * climatology_std_chunk
-    )
+    threshold = self.threshold.compute(truth)
     truth_probability = xr.where(truth > threshold, 1.0, 0.0)
 
     var_list = []
@@ -909,45 +881,16 @@ class GaussianIgnoranceScore(Metric):
     Spatially averaged ignorance score for a Gaussian distribution.
   """
 
-  climatology: xr.Dataset
-  threshold: float
+  threshold: thresholds.Threshold
 
   def compute_chunk(
       self,
       forecast: xr.Dataset,
       truth: xr.Dataset,
       region: t.Optional[Region] = None,
   ) -> xr.Dataset:
-    if "init_time" in forecast.dims:
-      time_dim = "valid_time"
-    else:
-      time_dim = "time"
-    climatology_chunk = _get_climatology_chunk(self.climatology, truth)
-    clim_std_dict = {key + "_std": key for key in truth.keys()}  # pytype: disable=unsupported-operands
-    try:
-      climatology_std_chunk = self.climatology[
-          list(clim_std_dict.keys())
-      ].rename(clim_std_dict)
-    except KeyError as e:
-      not_found_stds = set(clim_std_dict).difference(self.climatology.data_vars)
-      raise KeyError(
-          f"Did not find {not_found_stds} forecast keys in climatology."
-      ) from e
 
-    if hasattr(forecast, "level"):
-      climatology_chunk = climatology_chunk.sel(level=forecast.level)
-      climatology_std_chunk = climatology_std_chunk.sel(level=forecast.level)
-
-    time_selection = dict(dayofyear=forecast[time_dim].dt.dayofyear)
-    if "hour" in set(climatology_chunk.coords):
-      time_selection["hour"] = forecast[time_dim].dt.hour
-
-    climatology_chunk = climatology_chunk.sel(time_selection).compute()
-    climatology_std_chunk = climatology_std_chunk.sel(time_selection).compute()
-    threshold = (
-        climatology_chunk
-        + stats.norm.ppf(self.threshold) * climatology_std_chunk
-    )
+    threshold = self.threshold.compute(truth)
     truth_probability = xr.where(truth > threshold, 1.0, 0.0)
 
     var_list = []

weatherbench2/metrics_test.py

@@ -19,6 +19,7 @@
 from weatherbench2 import regions
 from weatherbench2 import schema
 from weatherbench2 import test_utils
+from weatherbench2 import thresholds
 from weatherbench2 import utils
 import xarray as xr
 
@@ -352,10 +353,10 @@ def test_gaussian_brier_score(self, error, expected):
         .rename({'2m_temperature': '2m_temperature_std'})
     )
     climatology = xr.merge([climatology_mean, climatology_std])
-    threshold = 0.8
-    result = metrics.GaussianBrierScore(climatology, threshold).compute(
-        forecast, truth
+    threshold = thresholds.GaussianQuantileThreshold(
+        climatology=climatology, quantile=0.8
     )
+    result = metrics.GaussianBrierScore(threshold).compute(forecast, truth)
     expected_arr = np.array([expected, expected])
     np.testing.assert_allclose(
         result['2m_temperature'].values, expected_arr, rtol=1e-4
@@ -394,10 +395,10 @@ def test_gaussian_ignorance_score(self, error, expected):
         .rename({'2m_temperature': '2m_temperature_std'})
     )
     climatology = xr.merge([climatology_mean, climatology_std])
-    threshold = 0.8
-    result = metrics.GaussianIgnoranceScore(climatology, threshold).compute(
-        forecast, truth
+    threshold = thresholds.GaussianQuantileThreshold(
+        climatology=climatology, quantile=0.8
     )
+    result = metrics.GaussianIgnoranceScore(threshold).compute(forecast, truth)
     print('Ignorance score result is', result)
     expected_arr = np.array([expected, expected])
     np.testing.assert_allclose(

weatherbench2/thresholds.py

@@ -0,0 +1,123 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Thresholding classes for discrete probabilistic metrics."""
+
+from collections import abc
+import dataclasses
+import typing
+
+from scipy import stats
+import xarray as xr
+
+
+def _get_climatology_mean(
+    climatology: xr.Dataset, variables: abc.Sequence[str]
+) -> xr.Dataset:
+  """Returns the climatological mean of the given variables."""
+  try:
+    climatology_mean = climatology[variables]
+  except KeyError as e:
+    not_found = set(variables).difference(climatology.data_vars)
+    clim_var_dict = {var + "_mean": var for var in variables}  # pytype: disable=unsupported-operands
+    not_found_means = set(clim_var_dict).difference(climatology.data_vars)
+    if not_found and not_found_means:
+      raise KeyError(
+          f"Did not find {not_found} keys in climatology. Appending "
+          "'mean' did not help."
+      ) from e
+    climatology_mean = climatology[list(clim_var_dict.keys())].rename(
+        clim_var_dict
+    )
+  return typing.cast(xr.Dataset, climatology_mean)
+
+
+def _get_climatology_std(
+    climatology: xr.Dataset, variables: abc.Sequence[str]
+) -> xr.Dataset:
+  """Returns the climatological standard deviation of the given variables."""
+  clim_std_dict = {key + "_std": key for key in variables}  # pytype: disable=unsupported-operands
+  try:
+    climatology_std = climatology[list(clim_std_dict.keys())].rename(
+        clim_std_dict
+    )
+    return typing.cast(xr.Dataset, climatology_std)
+  except KeyError as e:
+    not_found_stds = set(clim_std_dict).difference(climatology.data_vars)
+    raise KeyError(f"Did not find {not_found_stds} keys in climatology.") from e
+
+
+@dataclasses.dataclass
+class Threshold:
+  """Threshold for discrete probabilistic metric evaluation."""
+
+  def compute(self, truth: xr.Dataset) -> xr.Dataset:
+    """Computes the threshold for each true label variable.
+
+    Args:
+      truth: A dataset of label variables.
+
+    Returns:
+      A dataset containing thresholds for all label variables.
+    """
+    raise NotImplementedError
+
+
+@dataclasses.dataclass
+class QuantileThreshold(Threshold):
+  """Quantile threshold for discrete probabilistic metrics."""
+
+  def compute(self, truth: xr.Dataset) -> xr.Dataset:
+    """Computes the threshold for each label variable."""
+
+    raise NotImplementedError
+
+
+@dataclasses.dataclass
+class GaussianQuantileThreshold(Threshold):
+  """Gaussian quantile threshold for discrete probabilistic metrics.
+
+  Attributes:
+    climatology: Dataset containing the mean and standard deviation of the
+      climatological distribution.
+    quantile: The quantile to be evaluated given a Gaussian approximation of the
+      climatological distribution.
+  """
+
+  climatology: xr.Dataset
+  quantile: float
+
+  def compute(self, truth: xr.Dataset) -> xr.Dataset:
+    """Computes the threshold for each label variable."""
+    if "time" in truth.dims:
+      time_dim = "time"
+    else:
+      time_dim = "valid_time"
+
+    climatology_chunk = self.climatology
+    if hasattr(truth, "level"):
+      climatology_chunk = climatology_chunk.sel(level=truth.level)
+
+    time_selection = dict(dayofyear=truth["time"].dt.dayofyear)
+    if "hour" in set(climatology_chunk.coords):
+      time_selection["hour"] = truth[time_dim].dt.hour
+    climatology_chunk = climatology_chunk.sel(time_selection).compute()
+
+    variables = [str(key) for key in truth.keys()]
+    climatology_mean = _get_climatology_mean(climatology_chunk, variables)
+    climatology_std = _get_climatology_std(climatology_chunk, variables)
+    threshold = (
+        climatology_mean + stats.norm.ppf(self.quantile) * climatology_std
+    )
+    return threshold