diff --git a/CHANGELOG.md b/CHANGELOG.md index d976ad4b9..f7f98685e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 7db2cca6d..a995a28d1 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -227,6 +227,65 @@ 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, @@ -234,4 +293,5 @@ def crps_gauss( "wmae": wmae, "nll": nll, "crps_gauss": crps_gauss, + "spread_squared": spread_squared, } diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py index f1bcb461d..bf62d5a6f 100644 --- a/neural_lam/models/ar_model.py +++ b/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 diff --git a/neural_lam/vis.py b/neural_lam/vis.py index 06f2f6d35..89dd62c18 100644 --- a/neural_lam/vis.py +++ b/neural_lam/vis.py @@ -1,5 +1,7 @@ +# Standard library +import warnings + # Third-party -import cartopy.crs as ccrs import cartopy.feature as cfeature import matplotlib import matplotlib.colors @@ -12,7 +14,7 @@ from . import utils from .datastore.base import BaseRegularGridDatastore -# Font sizes shared across all plot functions for visual consistency. +# Font sizes shared across map plot functions for visual consistency. _TITLE_SIZE = 13 # suptitle and per-axes titles _LABEL_SIZE = 11 # axis / colorbar labels _TICK_SIZE = 11 # tick labels @@ -67,8 +69,9 @@ def plot_on_axis( ------- matplotlib.collections.QuadMesh The mesh object created by pcolormesh. - """ + # Third-party + import cartopy.crs as ccrs ax.coastlines(resolution="50m") ax.add_feature(cfeature.BORDERS, linestyle="-", alpha=0.5) @@ -109,14 +112,12 @@ def plot_on_axis( ) if boundary_alpha is not None: - # Overlay boundary mask mask_da = datastore.boundary_mask mask_values = mask_da.values if mask_values.ndim == 2 and mask_values.shape[1] == 1: mask_values = mask_values[:, 0] mask_2d = mask_values.reshape(grid_shape) - # Create overlay: 1 where boundary, NaN where interior overlay = np.where(mask_2d == 1, 1.0, np.nan) ax.pcolormesh( @@ -129,7 +130,6 @@ def plot_on_axis( ) if crop_to_interior: - # Calculate extent of interior mask_da = datastore.boundary_mask mask_values = mask_da.values if mask_values.ndim == 2 and mask_values.shape[1] == 1: @@ -154,63 +154,193 @@ def plot_on_axis( return mesh +# Annotations become unreadable when cells are smaller than this (in points) +# or when the total number of cells exceeds a readable count. +_MIN_CELL_SIZE_FOR_ANNOTATIONS = 18 +_MAX_CELLS_FOR_ANNOTATIONS = 800 + + +def _compute_heatmap_layout(n_rows: int, n_cols: int) -> dict[str, float]: + """Choose figure and font sizes from the heatmap dimensions. + + Scaling coefficients were empirically tuned for readability on grids + ranging from ~5 to ~50 variables/lead-times. The figure grows + proportionally so that cells never shrink below a readable size. + """ + max_dim = max(n_rows, n_cols) + + # Size the figure so each cell gets ~0.8 x 0.5 inches; floor at 8 x 4.5. + fig_width = float(max(4.5 + 0.8 * n_cols, 8.0)) + fig_height = float(max(2.5 + 0.5 * n_rows, 4.5)) + + # Approximate cell size in points (72 pt/inch) to decide whether + # in-cell annotations will be legible. + cell_w_pt = (fig_width / max(n_cols, 1)) * 72 + cell_h_pt = (fig_height / max(n_rows, 1)) * 72 + show_annotations = ( + n_rows * n_cols <= _MAX_CELLS_FOR_ANNOTATIONS + and min(cell_w_pt, cell_h_pt) >= _MIN_CELL_SIZE_FOR_ANNOTATIONS + ) + + return { + "fig_width": fig_width, + "fig_height": fig_height, + "tick_label_size": float(np.clip(15.0 - 0.18 * max_dim, 7.0, 14.0)), + "annotation_size": float(np.clip(13.0 - 0.22 * max_dim, 5.0, 12.0)), + "title_size": float(np.clip(16.0 - 0.15 * max_dim, 9.0, 15.0)), + "x_tick_rotation": 45.0 if n_cols > 12 else 0.0, + "show_annotations": show_annotations, + } + + +def _get_heatmap_var_labels( + datastore: BaseRegularGridDatastore, n_vars: int +) -> list[str]: + """Build state-variable labels, padding defensively if metadata is short.""" + var_names = list(datastore.get_vars_names(category="state")) + var_units = list(datastore.get_vars_units(category="state")) + + if len(var_names) < n_vars: + var_names.extend( + [f"state_feature_{i}" for i in range(len(var_names), n_vars)] + ) + if len(var_units) < n_vars: + var_units.extend([""] * (n_vars - len(var_units))) + + labels = [] + for name, unit in zip(var_names[:n_vars], var_units[:n_vars]): + labels.append(f"{name} ({unit})" if unit else name) + + return labels + + @matplotlib.rc_context(utils.fractional_plot_bundle(1)) -def plot_error_map(errors, datastore: BaseRegularGridDatastore, title=None): +def plot_error_heatmap( + errors, + datastore: BaseRegularGridDatastore, + title=None, + vmin: float | None = None, + vmax: float | None = None, +): """ - Plot a heatmap of errors of different variables at different - predictions horizons - errors: (pred_steps, d_f) + Plot a heatmap of errors for state variables across forecast lead times. + + Parameters + ---------- + errors : torch.Tensor + Error values with shape `(pred_steps, d_f)`. + datastore : BaseRegularGridDatastore + Datastore providing step length and variable metadata. + title : str, optional + Optional title for the figure. + vmin : float, optional + Minimum value for the colour scale. Defaults to 0 when all errors + are non-negative, otherwise the global minimum of the error array. + vmax : float, optional + Maximum value for the colour scale. Defaults to the global maximum + of the error array. """ - errors_np = errors.T.cpu().numpy() # (d_f, pred_steps) + errors_np = errors.detach().cpu().numpy().T # (d_f, pred_steps) d_f, pred_steps = errors_np.shape step_length = datastore.step_length - # Normalize all errors to [0,1] for color map - max_errors = errors_np.max(axis=1) # d_f - errors_norm = errors_np / np.expand_dims(max_errors, axis=1) - time_step_int, time_step_unit = utils.get_integer_time(step_length) + layout = _compute_heatmap_layout(n_rows=d_f, n_cols=pred_steps) - fig, ax = plt.subplots(figsize=(15, 10)) + finite_errors = errors_np[np.isfinite(errors_np)] + if finite_errors.size == 0: + computed_vmin, computed_vmax = 0.0, 1.0 + else: + computed_vmin = float(finite_errors.min()) + computed_vmax = float(finite_errors.max()) + if computed_vmin >= 0.0: + computed_vmin = 0.0 + if np.isclose(computed_vmin, computed_vmax): + computed_vmax = computed_vmin + 1.0 - ax.imshow( - errors_norm, - cmap="OrRd", - vmin=0, - vmax=1.0, + final_vmin = vmin if vmin is not None else computed_vmin + final_vmax = vmax if vmax is not None else computed_vmax + + fig, ax = plt.subplots( + figsize=(layout["fig_width"], layout["fig_height"]), + constrained_layout=True, + ) + + im = ax.imshow( + errors_np, + cmap="viridis", + vmin=final_vmin, + vmax=final_vmax, interpolation="none", aspect="auto", - alpha=0.8, ) + cbar = fig.colorbar(im, ax=ax, pad=0.02) + cbar.ax.tick_params(labelsize=layout["tick_label_size"]) + cbar.ax.yaxis.get_offset_text().set_fontsize(layout["tick_label_size"]) + + if layout["show_annotations"]: + for (j, i), error in np.ndenumerate(errors_np): + formatted_error = ( + f"{error:.3g}" if abs(error) < 1.0e4 else f"{error:.2E}" + ) + normed = im.norm(error) + text_color = ( + "white" if (np.isfinite(normed) and normed < 0.45) else "black" + ) + ax.text( + i, + j, + formatted_error, + ha="center", + va="center", + usetex=False, + fontsize=layout["annotation_size"], + color=text_color, + ) - # ax and labels - for (j, i), error in np.ndenumerate(errors_np): - # Numbers > 9999 will be too large to fit - formatted_error = f"{error:.3f}" if error < 9999 else f"{error:.2E}" - ax.text(i, j, formatted_error, ha="center", va="center", usetex=False) - - # Ticks and labels ax.set_xticks(np.arange(pred_steps)) pred_hor_i = np.arange(pred_steps) + 1 pred_hor_h = time_step_int * pred_hor_i - ax.set_xticklabels(pred_hor_h, size=_TICK_SIZE) - ax.set_xlabel(f"Lead time ({time_step_unit[0]})", size=_LABEL_SIZE) + ax.set_xticklabels( + pred_hor_h, + size=layout["tick_label_size"], + rotation=layout["x_tick_rotation"], + ha="right" if layout["x_tick_rotation"] > 0 else "center", + ) + ax.set_xlabel( + f"Lead time ({time_step_unit[0]})", size=layout["tick_label_size"] + ) ax.set_yticks(np.arange(d_f)) - var_names = datastore.get_vars_names(category="state") - var_units = datastore.get_vars_units(category="state") - y_ticklabels = [ - _tex_safe(f"{name} ({unit})") - for name, unit in zip(var_names, var_units) - ] - ax.set_yticklabels(y_ticklabels, rotation=30, size=_TICK_SIZE) + ax.set_yticklabels( + _get_heatmap_var_labels(datastore=datastore, n_vars=d_f), + size=layout["tick_label_size"], + ) if title: - ax.set_title(title, size=_TITLE_SIZE) + ax.set_title(title, size=layout["title_size"]) return fig +def plot_error_map( + errors, + datastore: BaseRegularGridDatastore, + title=None, + vmin: float | None = None, + vmax: float | None = None, +): + """Deprecated: use :func:`plot_error_heatmap` instead.""" + warnings.warn( + "plot_error_map is deprecated, use plot_error_heatmap instead", + DeprecationWarning, + stacklevel=2, + ) + return plot_error_heatmap( + errors, datastore=datastore, title=title, vmin=vmin, vmax=vmax + ) + + @matplotlib.rc_context(utils.fractional_plot_bundle(1)) def plot_prediction( datastore: BaseRegularGridDatastore, diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 000000000..5a1c9ac48 --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,242 @@ +""" +Tests for spread_squared ensemble variance metric. + +Verifies: +1. Mathematical correctness against known analytical values +2. Unbiased estimation property (converges to true variance) +3. Guard against single-member ensembles (S=1) +4. Output shape matches the API contract +5. Consistency with torch.var (Bessel-corrected) +6. Correct behavior with mask and reduction flags + +Shape note: tests use (B, S, N, F) without an explicit T dimension. +This is intentional — the metric follows the (..., S, N, d_state) convention +where T (from Issue #335's (B, T, N, F) shape) is just another batch dim +folded into (...). In ar_model.py, B and T are already flattened before +metrics are called, so (B, S, N, F) accurately reflects real call shapes. +""" +import pytest +import torch + +from neural_lam.metrics import spread_squared + + +class TestSpreadSquared: + """Tests for the spread_squared (ensemble variance) metric.""" + + def test_known_value_two_members(self): + """ + For pred = [1.0, 3.0], unbiased variance = 2.0. + mean = 2.0, deviations = [-1, 1], sum of sq = 2, / (2-1) = 2.0 + """ + # Shape: (B=1, S=2, N=1, d_state=1) — T folded into B per (...) convention + pred = torch.tensor([[[[1.0]], [[3.0]]]]) + target = torch.zeros(1, 1, 1) # (B, N, d_state) — unused + pred_std = torch.ones(1, 1, 1) # unused + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + + # With average_grid=False, sum_vars=False → shape (B, N, d_state) + assert result.shape == (1, 1, 1), ( + f"Expected shape (1,1,1), got {result.shape}" + ) + assert torch.allclose(result, torch.tensor([[[2.0]]])), ( + f"Expected 2.0, got {result.item()}" + ) + + def test_known_value_three_members(self): + """ + For pred = [2.0, 4.0, 6.0], mean=4.0, + unbiased var = ((2-4)^2 + (4-4)^2 + (6-4)^2) / (3-1) = 4.0. + """ + # Shape: (B=1, S=3, N=1, d_state=1) — T folded into B per (...) convention + pred = torch.tensor([[[[2.0]], [[4.0]], [[6.0]]]]) + target = torch.zeros(1, 1, 1) + pred_std = torch.ones(1, 1, 1) + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + + assert torch.allclose(result, torch.tensor([[[4.0]]])), ( + f"Expected 4.0, got {result.item()}" + ) + + def test_unbiased_estimation_large_ensemble(self): + """ + For a large ensemble drawn from N(0,1), the unbiased sample + variance should converge to the true variance (1.0). + """ + torch.manual_seed(42) + S = 10000 + # Shape: (B=1, S=10000, N=1, d_state=1) — T folded into B per (...) convention + pred = torch.randn(1, S, 1, 1) + target = torch.zeros(1, 1, 1) + pred_std = torch.ones(1, 1, 1) + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + + assert torch.abs(result.squeeze() - 1.0) < 0.05, ( + f"Expected ~1.0 for N(0,1) variance, got {result.item()}" + ) + + def test_rejects_single_member(self): + """S=1 must raise AssertionError — single-member variance + is undefined.""" + pred = torch.randn(1, 1, 1, 1) # S=1 + target = torch.zeros(1, 1, 1) + pred_std = torch.ones(1, 1, 1) + + with pytest.raises(AssertionError, match="more than 1 member"): + spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + + def test_output_shape_no_reduction(self): + """ + With average_grid=False, sum_vars=False, output shape should + be (..., N, d_state) with ensemble dim reduced. + + T is a batch dimension folded into (...); pred is (B*T, S, N, F). + """ + B, S, N, F = 4, 8, 100, 3 + pred = torch.randn(B, S, N, F) # (..., S, N, d_state) + target = torch.randn(B, N, F) + pred_std = torch.ones(F) + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + + # S reduced, rest preserved: (B, N, F) + assert result.shape == (B, N, F), ( + f"Expected ({B},{N},{F}), got {result.shape}" + ) + + def test_output_shape_with_reduction(self): + """ + With average_grid=True, sum_vars=True (defaults), output + should reduce N and d_state dims. + + T is a batch dimension folded into (...); pred is (B*T, S, N, F). + """ + B, S, N, F = 2, 5, 50, 4 + pred = torch.randn(B, S, N, F) + target = torch.randn(B, N, F) + pred_std = torch.ones(F) + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=True, + sum_vars=True, + ) + + # S reduced, N averaged, F summed: (B,) + assert result.shape == (B,), ( + f"Expected ({B},), got {result.shape}" + ) + + def test_matches_torch_var(self): + """ + Verify that spread_squared (without grid/var reduction) matches + torch.var with Bessel's correction across the ensemble dimension. + """ + torch.manual_seed(123) + B, S, N, F = 2, 5, 10, 4 + # T folded into B per (...) convention + pred = torch.randn(B, S, N, F) + target = torch.randn(B, N, F) + pred_std = torch.ones(F) + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + # ens_dim=-3 means dim=1 for 4D input (B, S, N, F) + expected = torch.var(pred, dim=-3, unbiased=True) + + assert torch.allclose(result, expected, atol=1e-6), ( + "spread_squared does not match torch.var(unbiased=True)" + ) + + def test_zero_spread(self): + """ + If all ensemble members are identical, variance should be 0. + """ + # Shape: (B=1, S=3, N=1, d_state=1) — all members = 5.0 + pred = torch.full((1, 3, 1, 1), 5.0) + target = torch.zeros(1, 1, 1) + pred_std = torch.ones(1, 1, 1) + + result = spread_squared( + pred, + target, + pred_std, + mask=None, + average_grid=False, + sum_vars=False, + ) + + assert torch.allclose(result, torch.tensor([[[0.0]]])), ( + f"Expected 0.0 for identical members, got {result.item()}" + ) + + def test_with_mask(self): + """ + Verify that the boolean mask correctly filters grid nodes. + """ + # Shape: (B=1, S=3, N=4, d_state=1) + pred = torch.randn(1, 3, 4, 1) + target = torch.randn(1, 4, 1) + pred_std = torch.ones(1) + mask = torch.tensor([True, False, True, False]) # keep 2 of 4 nodes + + result = spread_squared( + pred, + target, + pred_std, + mask=mask, + average_grid=False, + sum_vars=False, + ) + + # N filtered from 4 to 2 + assert result.shape == (1, 2, 1), ( + f"Expected shape (1,2,1) with mask, got {result.shape}" + ) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 84b35a308..5cdee6123 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -1,195 +1,45 @@ # Standard library from datetime import timedelta from pathlib import Path -from typing import Iterator -from unittest.mock import patch # Third-party -import matplotlib.figure -import matplotlib.pyplot as plt -import numpy as np -import pytest -import torch -import xarray as xr -from cartopy import crs as ccrs +import matplotlib + +matplotlib.use("Agg") # non-interactive backend for headless test runs + +# Third-party +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 +import pytest # noqa: E402 +import torch # noqa: E402 # First-party -from neural_lam import config as nlconfig -from neural_lam import vis -from neural_lam.create_graph import create_graph_from_datastore -from neural_lam.models.graph_lam import GraphLAM -from neural_lam.weather_dataset import WeatherDataset -from tests.conftest import init_datastore_example -from tests.dummy_datastore import DummyDatastore +from neural_lam import config as nlconfig # noqa: E402 +from neural_lam import vis # noqa: E402 +from neural_lam.create_graph import create_graph_from_datastore # noqa: E402 +from neural_lam.models.graph_lam import GraphLAM # noqa: E402 +from neural_lam.weather_dataset import WeatherDataset # noqa: E402 +from tests.dummy_datastore import DummyDatastore # noqa: E402 # Create output directory for test figures TEST_OUTPUT_DIR = Path(__file__).parent / "test_outputs" / "plotting" TEST_OUTPUT_DIR.mkdir(parents=True, exist_ok=True) -@pytest.fixture(autouse=True) -def mock_cartopy_downloads(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Prevent cartopy from downloading Natural Earth map data during tests. - Monkeypatches the GeoAxes methods used in vis.plot_on_axis. - """ - # Third-party - from cartopy.mpl.geoaxes import GeoAxes - - monkeypatch.setattr(GeoAxes, "coastlines", lambda *args, **kwargs: None) - monkeypatch.setattr(GeoAxes, "add_feature", lambda *args, **kwargs: None) - - -@pytest.fixture(autouse=True) -def close_all_figures_after_test() -> Iterator[None]: - """Ensure test-created matplotlib figures are always cleaned up.""" - yield - plt.close("all") +class HeatmapDatastore: + """Minimal datastore stub for error-heatmap plotting tests.""" + def __init__(self, n_vars, step_length=timedelta(hours=1)): + self._n_vars = n_vars + self.step_length = step_length -def test_plot_prediction() -> None: - """Check prediction plot structure, titles and shared color scaling.""" - datastore = init_datastore_example("dummydata") - n_grid = datastore.num_grid_points - - da_pred = xr.DataArray(np.linspace(0.0, 1.0, n_grid)) - da_target = xr.DataArray(np.linspace(1.0, 2.0, n_grid)) - - expected_vmin = float(np.nanmin([da_pred.values, da_target.values])) - expected_vmax = float(np.nanmax([da_pred.values, da_target.values])) - - fig = vis.plot_prediction( - datastore=datastore, - da_prediction=da_pred, - da_target=da_target, - title="Test Prediction", - vrange=(expected_vmin, expected_vmax), - boundary_alpha=None, - crop_to_interior=False, - ) + def get_vars_names(self, category): + assert category == "state" + return [f"state_var_{i}" for i in range(self._n_vars)] - assert isinstance(fig, matplotlib.figure.Figure) - assert len(fig.axes) == 3 - - ground_truth_ax, prediction_ax, _ = fig.axes - assert ground_truth_ax.get_title() == "Ground Truth" - assert prediction_ax.get_title() == "Prediction" - assert fig._suptitle.get_text() == "Test Prediction" - - assert len(ground_truth_ax.collections) == 1 - assert len(prediction_ax.collections) == 1 - - assert ground_truth_ax.collections[0].norm.vmin == expected_vmin - assert ground_truth_ax.collections[0].norm.vmax == expected_vmax - assert prediction_ax.collections[0].norm.vmin == expected_vmin - assert prediction_ax.collections[0].norm.vmax == expected_vmax - - -def test_plot_error_map() -> None: - """Check error heatmap content, labels and annotations.""" - datastore = init_datastore_example("dummydata") - d_f = len(datastore.get_vars_names(category="state")) - pred_steps = 4 - - errors = torch.arange(1, pred_steps * d_f + 1, dtype=torch.float32).reshape( - pred_steps, d_f - ) - - fig = vis.plot_error_map( - errors=errors, - datastore=datastore, - title="Test Error Map", - ) - - assert isinstance(fig, matplotlib.figure.Figure) - assert len(fig.axes) == 1 - - ax = fig.axes[0] - assert len(ax.images) == 1 - assert ax.images[0].get_array().shape == (d_f, pred_steps) - assert ax.get_xlabel() == "Lead time (h)" - assert ax.get_title() == "Test Error Map" - - expected_x_ticklabels = [str(step) for step in range(1, pred_steps + 1)] - actual_x_ticklabels = [tick.get_text() for tick in ax.get_xticklabels()] - assert actual_x_ticklabels == expected_x_ticklabels - - var_names = datastore.get_vars_names(category="state") - var_units = datastore.get_vars_units(category="state") - expected_y_ticklabels = [ - f"{name} ({unit})" for name, unit in zip(var_names, var_units) - ] - actual_y_ticklabels = [tick.get_text() for tick in ax.get_yticklabels()] - assert actual_y_ticklabels == expected_y_ticklabels - - assert len(ax.texts) == pred_steps * d_f - - -def test_plot_spatial_error() -> None: - """Check that plot_spatial_error runs without error and returns a Figure.""" - datastore = init_datastore_example("dummydata") - n_grid = datastore.num_grid_points - - error = torch.linspace(0.0, 1.0, n_grid) - - fig = vis.plot_spatial_error( - error=error, - datastore=datastore, - title="Test Spatial Error", - boundary_alpha=None, - crop_to_interior=False, - ) - - assert isinstance(fig, matplotlib.figure.Figure) - # GeoAxes + colorbar axes - assert len(fig.axes) == 2 - assert fig.texts[0].get_text() == "Test Spatial Error" - - -def test_plot_spatial_error_crop_to_interior_changes_extent() -> None: - """Check interior cropping forwards interior lon/lat bounds to - set_extent.""" - datastore = init_datastore_example("dummydata") - n_grid = datastore.num_grid_points - grid_shape = (datastore.grid_shape_state.x, datastore.grid_shape_state.y) - - boundary_mask = np.ones(grid_shape, dtype=int) - boundary_mask[2:-2, 2:-2] = 0 - datastore.ds["boundary_mask"] = xr.DataArray( - boundary_mask.reshape(n_grid), dims=["grid_index"] - ) - datastore.__dict__.pop("boundary_mask", None) - - lats_lons = datastore.get_lat_lon("state") - lons = lats_lons[:, 0].reshape(grid_shape) - lats = lats_lons[:, 1].reshape(grid_shape) - interior = boundary_mask == 0 - - expected_min_lon = float(lons[interior].min()) - expected_max_lon = float(lons[interior].max()) - expected_min_lat = float(lats[interior].min()) - expected_max_lat = float(lats[interior].max()) - - error = torch.linspace(0.0, 1.0, n_grid) - with patch( - "cartopy.mpl.geoaxes.GeoAxes.set_extent", autospec=True - ) as set_extent_mock: - vis.plot_spatial_error( - error=error, - datastore=datastore, - boundary_alpha=None, - crop_to_interior=True, - ) - - assert set_extent_mock.call_count == 1 - called_extent = set_extent_mock.call_args.args[1] - called_crs = set_extent_mock.call_args.kwargs["crs"] - - assert called_extent[0] == pytest.approx(expected_min_lon) - assert called_extent[1] == pytest.approx(expected_max_lon) - assert called_extent[2] == pytest.approx(expected_min_lat) - assert called_extent[3] == pytest.approx(expected_max_lat) - assert isinstance(called_crs, ccrs.PlateCarree) + def get_vars_units(self, category): + assert category == "state" + return ["unit"] * self._n_vars @pytest.fixture @@ -359,3 +209,120 @@ def test_plot_examples_integration_saves_figure( assert fig is not None assert isinstance(fig, plt.Figure) assert output_path.exists() + + +def test_plot_error_heatmap_uses_global_color_scale(): + """Heatmap colors should encode absolute values across all variables.""" + errors = torch.tensor( + [ + [1.0, 100.0, 10.0], + [2.0, 80.0, 5.0], + [3.0, 60.0, 2.5], + ] + ) # (pred_steps, d_f) + datastore = HeatmapDatastore(n_vars=errors.shape[1]) + + fig = vis.plot_error_heatmap(errors, datastore=datastore) + ax = fig.axes[0] + image = ax.images[0] + + np.testing.assert_allclose(image.get_array(), errors.T.numpy()) + assert image.norm.vmin == 0.0 + assert image.norm.vmax == pytest.approx(errors.max().item()) + assert len(fig.axes) == 2 # main axis + colorbar axis + + plt.close(fig) + + +def test_plot_error_heatmap_adapts_figure_and_font_sizes(): + """Dense heatmaps should get more space and smaller text.""" + small_errors = torch.ones((4, 5)) + large_errors = torch.ones((20, 30)) + + small_fig = vis.plot_error_heatmap( + small_errors, datastore=HeatmapDatastore(n_vars=small_errors.shape[1]) + ) + large_fig = vis.plot_error_heatmap( + large_errors, datastore=HeatmapDatastore(n_vars=large_errors.shape[1]) + ) + + small_ax = small_fig.axes[0] + large_ax = large_fig.axes[0] + + assert large_fig.get_size_inches()[0] > small_fig.get_size_inches()[0] + assert large_fig.get_size_inches()[1] > small_fig.get_size_inches()[1] + assert ( + large_ax.get_yticklabels()[0].get_fontsize() + < small_ax.get_yticklabels()[0].get_fontsize() + ) + assert large_ax.texts[0].get_fontsize() < small_ax.texts[0].get_fontsize() + assert large_ax.get_xticklabels()[0].get_rotation() == 45.0 + + plt.close(small_fig) + plt.close(large_fig) + + +def test_plot_error_heatmap_skips_annotations_for_very_dense_grids(): + """Very dense heatmaps should omit in-cell text to stay readable.""" + dense_errors = torch.ones((40, 50)) + fig = vis.plot_error_heatmap( + dense_errors, datastore=HeatmapDatastore(n_vars=dense_errors.shape[1]) + ) + ax = fig.axes[0] + + # No text annotations should be drawn when cells are too small + assert len(ax.texts) == 0 + # Figure should still grow beyond the old 18-inch cap + assert fig.get_size_inches()[0] > 18.0 + + plt.close(fig) + + +def test_plot_error_map_deprecated_wrapper(): + """The old plot_error_map name should still work but emit a warning.""" + errors = torch.ones((3, 4)) + datastore = HeatmapDatastore(n_vars=errors.shape[1]) + + with pytest.warns(DeprecationWarning, match="plot_error_heatmap"): + fig = vis.plot_error_map(errors, datastore=datastore) + + assert isinstance(fig, plt.Figure) + plt.close(fig) + + +def test_plot_error_heatmap_respects_explicit_vmin_vmax(): + errors = torch.tensor([[1.0, 10.0], [2.0, 5.0]]) + datastore = HeatmapDatastore(n_vars=errors.shape[1]) + + fig = vis.plot_error_heatmap( + errors, datastore=datastore, vmin=0.0, vmax=200.0 + ) + ax = fig.axes[0] + image = ax.images[0] + + assert image.norm.vmin == 0.0 + assert image.norm.vmax == 200.0 + + plt.close(fig) + + +def test_plot_error_heatmap_cross_run_same_scale(): + errors1 = torch.tensor([[1.0, 5.0], [2.0, 10.0]]) + errors2 = torch.tensor([[50.0, 100.0], [60.0, 120.0]]) + datastore = HeatmapDatastore(n_vars=errors1.shape[1]) + + fig1 = vis.plot_error_heatmap( + errors1, datastore=datastore, vmin=0.0, vmax=150.0 + ) + fig2 = vis.plot_error_heatmap( + errors2, datastore=datastore, vmin=0.0, vmax=150.0 + ) + + image1 = fig1.axes[0].images[0] + image2 = fig2.axes[0].images[0] + + assert image1.norm.vmin == image2.norm.vmin == 0.0 + assert image1.norm.vmax == image2.norm.vmax == 150.0 + + plt.close(fig1) + plt.close(fig2)