Option to compute mse over north/south hemisphere

Author: 14renus

geoarches/dataloaders/era5.py

@@ -100,11 +100,13 @@ def get_level_variable_indices(
 
 
 def get_headline_level_variable_indices(
-    pressure_levels=arches_default_pressure_levels, level_variables=arches_default_level_variables
+    pressure_levels=arches_default_pressure_levels,
+    level_variables=arches_default_level_variables,
+    headline_variables=("Z500", "T850", "Q700", "U850", "V850"),
 ):
     """Mapping for main level variables."""
     out = get_level_variable_indices(pressure_levels, level_variables)
-    return {k: v for k, v in out.items() if k in ("Z500", "T850", "Q700", "U850", "V850")}
+    return {k: v for k, v in out.items() if k in headline_variables}
 
 
 class Era5Dataset(XarrayDataset):

geoarches/evaluation/metric_registry.py

@@ -46,6 +46,13 @@ def instantiate_metric(metric_name: str, **extra_kwargs):
 register_metric(
     "era5_deterministic_metrics_with_spatial", Era5DeterministicMetrics, compute_per_gridpoint=True
 )
+register_metric(
+    "era5_deterministic_metrics_with_spatial_and_hemisphere",
+    Era5DeterministicMetrics,
+    compute_per_gridpoint=True,
+    compute_per_hemisphere=True,
+    headline_variables=("Z500", "Z850", "T850", "Q700", "U850", "V850"),
+)
 register_metric(
     "era5_ensemble_metrics",
     Era5EnsembleMetrics,

geoarches/metrics/deterministic_metrics.py

@@ -74,6 +74,7 @@ def __init__(
         data_shape: tuple,
         compute_lat_weights_fn: Callable[[int], torch.tensor] = compute_lat_weights_weatherbench,
         compute_per_gridpoint: bool = False,
+        compute_per_hemisphere: bool = False,
     ):
         """
         Args:
@@ -87,6 +88,8 @@ def __init__(
                 Default function assumes latitudes are ordered -90 to 90.
             compute_per_gridpoint: Whether to also compute mse and rmse per gridpoint (along with aggregated globally mse and rmse).
                 Default: only compute globally aggregated mse and rmse.
+            compute_per_hemisphere: Whether to also compute mse and rmse per north and sount hemisphere
+                (along with aggregated globally mse and rmse).
         """
         Metric.__init__(self)
         MetricBase.__init__(
@@ -97,6 +100,7 @@ def __init__(
             # rollout_iterations=rollout_iterations,
         )
         self.compute_per_gridpoint = compute_per_gridpoint
+        self.compute_per_hemisphere = compute_per_hemisphere
 
         # Call `self.add_state`for every internal state that is needed for the metrics computations.
         # `dist_reduce_fx` indicates the function that should be used to reduce.
@@ -108,6 +112,11 @@ def __init__(
             "rmse_before_time_avg", default=torch.zeros(data_shape), dist_reduce_fx="sum"
         )
 
+        if self.compute_per_gridpoint:
+            # Aggregated over north and south hemispheres.
+            self.add_state("mse_north", default=torch.zeros(data_shape), dist_reduce_fx="sum")
+            self.add_state("mse_south", default=torch.zeros(data_shape), dist_reduce_fx="sum")
+
         if self.compute_per_gridpoint:
             # Per gridpoint.
             self.add_state(
@@ -134,6 +143,12 @@ def update(self, targets: torch.Tensor, preds: torch.Tensor) -> None:
             targets, preds
         ).sqrt().sum(0)
 
+        if self.compute_per_hemisphere:
+            num_lats = targets.shape[-2]
+            equator_index = num_lats // 2
+            self.mse_north = self.wmse(targets, preds, lat_range=(0, equator_index)).sum(0)
+            self.mse_south = self.wmse(targets, preds, lat_range=(equator_index, num_lats)).sum(0)
+
         if self.compute_per_gridpoint:
             self.mse_per_gridpt = self.mse_per_gridpt + self.spatial_mse(targets, preds).sum(0)
 
@@ -149,6 +164,16 @@ def compute(self) -> Dict[str, torch.Tensor]:
             rmse=(self.mse / self.nsamples).sqrt(),
         )
 
+        if self.compute_per_hemisphere:
+            all_metrics.update(
+                dict(
+                    mse_north=self.mse_north / self.nsamples,
+                    rmse_north=(self.mse_north / self.nsamples).sqrt(),
+                    mse_south=self.mse_south / self.nsamples,
+                    rmse_south=(self.mse_south / self.nsamples).sqrt(),
+                )
+            )
+
         if self.compute_per_gridpoint:
             all_metrics.update(
                 dict(
@@ -182,6 +207,7 @@ def __init__(
         surface_variables=era5.arches_default_surface_variables,
         level_variables=era5.arches_default_level_variables,
         pressure_levels=era5.arches_default_pressure_levels,
+        headline_variables=("Z500", "T850", "Q700", "U850", "V850"),
         compute_lat_weights_fn: Callable[[int], torch.tensor] = compute_lat_weights_weatherbench,
         compute_per_gridpoint: bool = False,
         lead_time_hours: int = 24,
@@ -192,6 +218,7 @@ def __init__(
             surface_variables: Names of surface variables (to select quantiles).
             level_variables: Names of level variables (used to get `variable_indices`).
             pressure_levels: pressure levels in data (used to get `variable_indices`).
+            headline_variables: Short names of level variables to output (used to get 'variable_indices').
             level_data_shape: (var, lev) shape for level variables.
             num_level_variables: Number of level variables (used to compute data_shape).
             compute_per_gridpoint: Whether to also compute mse and rmse per gridpoint (along with aggregated globally mse and rmse).
@@ -228,7 +255,9 @@ def __init__(
                     compute_per_gridpoint=compute_per_gridpoint,
                 ),
                 variable_indices=add_timedelta_index(
-                    era5.get_headline_level_variable_indices(pressure_levels, level_variables),
+                    era5.get_headline_level_variable_indices(
+                        pressure_levels, level_variables, headline_variables
+                    ),
                     lead_time_hours=lead_time_hours,
                     rollout_iterations=rollout_iterations,
                 ),

geoarches/metrics/metric_base.py

@@ -62,14 +62,27 @@ def __init__(
         super().__init__()
         self.compute_lat_weights_fn = compute_lat_weights_fn
 
-    def wmse(self, x: torch.Tensor, y: torch.Tensor | int = 0):
+    def wmse(
+        self, x: torch.Tensor, y: torch.Tensor | int = 0, lat_range: tuple[int, int] | None = None
+    ):
         """Latitude weighted mse error.
 
         Args:
             x: preds with shape (..., lat, lon)
             y: targets with shape (..., lat, lon)
+            lat_range: Optional tuple of (min_lat, max_lat) to restrict the latitude range for the computation.
+                If None, uses the full latitude range.
         """
         lat_coeffs = self.compute_lat_weights_fn(latitude_resolution=x.shape[-2]).to(x.device)
+
+        if lat_range is not None:
+            start_lat, end_lat = lat_range
+            x = x[..., start_lat:end_lat, :]
+            lat_coeffs = lat_coeffs[start_lat:end_lat, :]
+
+            if not isinstance(y, int):
+                y = y[..., start_lat:end_lat, :]
+
         return (x - y).pow(2).mul(lat_coeffs).nanmean((-2, -1))
 
     def spatial_mse(self, x: torch.Tensor, y: torch.Tensor | int = 0):