Give better error message in weatherbench/metrics when climatology doesn't contain forecast keys. Previously, the KeyError raised by xr would indicate that "some_var_mean" was not found, but that "_mean" variable was likely just a "let's do a final attempt" at finding variables, and this confuses the user.

PiperOrigin-RevId: 590711498

Author: langmore

weatherbench2/metrics.py

@@ -335,8 +335,17 @@ def compute_chunk(
       time_dim = "time"
     try:
       climatology_chunk = self.climatology[list(forecast.keys())]
-    except KeyError:
+    except KeyError as e:
+      not_found = set(forecast.keys()).difference(self.climatology.data_vars)
       clim_var_dict = {key + "_mean": key for key in forecast.keys()}  # pytype: disable=unsupported-operands
+      not_found_means = set(clim_var_dict).difference(
+          self.climatology.data_vars
+      )
+      if not_found and not_found_means:
+        raise KeyError(
+            f"Did not find {not_found} forecast keys in climatology. Appending "
+            "'mean' did not help"
+        ) from e
       climatology_chunk = self.climatology[list(clim_var_dict.keys())].rename(
           clim_var_dict
       )