Improve metric heatmap readability and scaling (#375)

CHANGELOG.md

@@ -30,6 +30,11 @@ 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
+- Improve metric heatmaps by renaming `plot_error_map`, using a relative
+  cross-variable color scale while keeping original values in annotations, and
+  scaling figure size, tick labels, and annotation text with the number of
+  variables and lead times so larger evaluation outputs remain readable
+  ([#375](https://github.com/mllam/neural-lam/issues/375))
 - Replaces multiple `assert` statements used for runtime input validation with explicit `ValueError` [\#279](https://github.com/mllam/neural-lam/pull/279) @Sir-Sloth-The-Lazy
 
 - 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

neural_lam/models/ar_model.py

@@ -381,7 +381,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
@@ -421,9 +421,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(
@@ -610,7 +611,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,
         )
@@ -642,7 +643,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.
@@ -700,7 +702,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