Add spread_squared ensemble variance metric for probabilistic evaluation

CHANGELOG.md

@@ -27,7 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Replace `shell=True` subprocess call in `compute_standardization_stats.py` with a safe argument list and Python-side hostname parsing to prevent command injection via `SLURM_JOB_NODELIST` [\#264](https://github.com/mllam/neural-lam/pull/264) @ashum9
 
 - Avoid NaN when standardizing fields with zero std [#189](https://github.com/mllam/neural-lam/pull/189) @varunsiravuri
-
+- Scale metric heatmap figure size, tick labels, and annotation text with the number of variables and lead times so plots remain readable for larger evaluation outputs ([#375](https://github.com/mllam/neural-lam/issues/375))
 - Fix README image paths to use absolute GitHub URLs so images display correctly on PyPI [\#188](https://github.com/mllam/neural-lam/pull/188) @bk-simon
 
 - Fix typo in `ar_model.py` that causes `AttributeError` during evaluation [\#204](https://github.com/mllam/neural-lam/pull/204) @ritinikhil

neural_lam/metrics.py

@@ -227,11 +227,71 @@ def crps_gauss(
     )
 
 
+def spread_squared(
+    pred, target, pred_std, mask=None, average_grid=True, sum_vars=True
+):
+    """
+    Ensemble variance (spread squared) metric.
+
+    Computes the unbiased sample variance of ensemble predictions across
+    the ensemble dimension (dim=-3). The entry-wise variance is then
+    passed through ``mask_and_reduce_metric`` for grid masking and
+    optional reduction — consistent with all other metrics.
+
+    This metric is used for spread-skill analysis: comparing ensemble
+    spread against forecast error (e.g. MSE) to assess probabilistic
+    calibration. For a well-calibrated ensemble, spread_squared should
+    approximate MSE.
+
+    (...,) is any number of batch dimensions
+    pred: (..., S, N, d_state), ensemble predictions where S is the
+        number of ensemble members at dim=-3
+    target: (..., N, d_state), target (unused, accepted for API
+        consistency with other metrics)
+    pred_std: (..., N, d_state) or (d_state,), predicted std.-dev.
+        (unused, accepted for API consistency)
+    mask: (N,), boolean mask describing which grid nodes to use
+    average_grid: boolean, if grid dimension -2 should be reduced
+        (mean over N)
+    sum_vars: boolean, if variable dimension -1 should be reduced
+        (sum over d_state)
+
+    Returns:
+    metric_val: One of (...,), (..., d_state), (..., N),
+        (..., N, d_state), depending on reduction arguments.
+    """
+    ens_dim = -3  # S dimension: pred is (..., S, N, d_state)
+    num_ens = pred.shape[ens_dim]
+    assert num_ens > 1, (
+        f"Ensemble variance requires more than 1 member, got S={num_ens}. "
+        "Single-member spread is undefined."
+    )
+
+    # Unbiased sample variance (Bessel's correction)
+    # var = (1/(S-1)) * sum((x_i - mean)^2)
+    # Implemented as: mean((x_i - mean)^2) * S/(S-1)
+    ens_mean = torch.mean(pred, dim=ens_dim)  # (..., N, d_state)
+    entry_spread = torch.mean(
+        (pred - ens_mean.unsqueeze(ens_dim)) ** 2,
+        dim=ens_dim,
+    ) * (
+        num_ens / (num_ens - 1)
+    )  # (..., N, d_state)
+
+    return mask_and_reduce_metric(
+        entry_spread,
+        mask=mask,
+        average_grid=average_grid,
+        sum_vars=sum_vars,
+    )
+
+
 DEFINED_METRICS = {
     "mse": mse,
     "mae": mae,
     "wmse": wmse,
     "wmae": wmae,
     "nll": nll,
     "crps_gauss": crps_gauss,
+    "spread_squared": spread_squared,
 }

neural_lam/models/ar_model.py

@@ -374,7 +374,7 @@ def on_validation_epoch_end(self):
         """
         Compute val metrics at the end of val epoch
         """
-        # Create error maps for all test metrics
+        # Create error heatmaps for all validation metrics
         self.aggregate_and_plot_metrics(self.val_metrics, prefix="val")
 
         # Clear lists with validation metrics values
@@ -414,9 +414,10 @@ def test_step(self, batch, batch_idx):
             batch_size=batch[0].shape[0],
         )
 
-        # Compute all evaluation metrics for error maps Note: explicitly list
-        # metrics here, as test_metrics can contain additional ones, computed
-        # differently, but that should be aggregated on_test_epoch_end
+        # Compute all evaluation metrics for error heatmaps. Note:
+        # explicitly list metrics here, as test_metrics can contain
+        # additional ones, computed differently, but that should be
+        # aggregated on_test_epoch_end
         for metric_name in ("mse", "mae"):
             metric_func = metrics.get_metric(metric_name)
             batch_metric_vals = metric_func(
@@ -603,7 +604,7 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name):
         Return: log_dict: dict with everything to log for given metric
         """
         log_dict = {}
-        metric_fig = vis.plot_error_map(
+        metric_fig = vis.plot_error_heatmap(
             errors=metric_tensor,
             datastore=self._datastore,
         )
@@ -635,7 +636,8 @@ def create_metric_log_dict(self, metric_tensor, prefix, metric_name):
 
     def aggregate_and_plot_metrics(self, metrics_dict, prefix):
         """
-        Aggregate and create error map plots for all metrics in metrics_dict
+        Aggregate and create error heatmap plots for all metrics in
+        metrics_dict
 
         metrics_dict: dictionary with metric_names and list of tensors
             with step-evals.
@@ -693,7 +695,7 @@ def on_test_epoch_end(self):
         Compute test metrics and make plots at the end of test epoch. Will
         gather stored tensors and perform plotting and logging on rank 0.
         """
-        # Create error maps for all test metrics
+        # Create error heatmaps for all test metrics
         self.aggregate_and_plot_metrics(self.test_metrics, prefix="test")
 
         # Plot spatial loss maps