From 651db4c114438f9856c8c1f636587c74095124fa Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Tue, 12 May 2026 17:02:12 +0200 Subject: [PATCH 1/7] Add DCT when computing spectra on regional domains --- .../training/diagnostics/callbacks/plot.py | 51 ++++-- .../src/anemoi/training/diagnostics/plots.py | 147 ++++++++++++++--- .../diagnostics/test_plotting_callbacks.py | 151 ++++++++++++++++-- 3 files changed, 302 insertions(+), 47 deletions(-) diff --git a/training/src/anemoi/training/diagnostics/callbacks/plot.py b/training/src/anemoi/training/diagnostics/callbacks/plot.py index a712e468a7..a012a0712f 100644 --- a/training/src/anemoi/training/diagnostics/callbacks/plot.py +++ b/training/src/anemoi/training/diagnostics/callbacks/plot.py @@ -424,7 +424,10 @@ def get_edge_trainable_modules( provider_specs = ( ("encoder_graph_provider", (dataset_name, model._graph_name_hidden)), ("decoder_graph_provider", (model._graph_name_hidden, dataset_name)), - ("processor_graph_provider", (model._graph_name_hidden, model._graph_name_hidden)), + ( + "processor_graph_provider", + (model._graph_name_hidden, model._graph_name_hidden), + ), ) for provider_name, edge_key in provider_specs: provider = self._resolve_edge_provider(getattr(model, provider_name, None), dataset_name) @@ -735,8 +738,8 @@ def __init__( # Build focus mask self.focus_mask = build_spatial_mask( - node_attribute_name=focus_area.get("mask_attr_name", None) if focus_area is not None else None, - latlon_bbox=focus_area.get("latlon_bbox", None) if focus_area is not None else None, + node_attribute_name=(focus_area.get("mask_attr_name", None) if focus_area is not None else None), + latlon_bbox=(focus_area.get("latlon_bbox", None) if focus_area is not None else None), name=focus_area.get("name", None) if focus_area is not None else None, ) @@ -850,7 +853,12 @@ def __init__( Batch frequency to plot at, by default None """ del kwargs - super().__init__(config, dataset_names=dataset_names, every_n_batches=every_n_batches, focus_area=focus_area) + super().__init__( + config, + dataset_names=dataset_names, + every_n_batches=every_n_batches, + focus_area=focus_area, + ) self.sample_idx = sample_idx self.parameters = parameters @@ -904,7 +912,12 @@ def _plot( output_tensor, ) - for x, y_true, y_pred, tag_suffix in pl_module.plot_adapter.iter_plot_samples(data, output_tensor): + for ( + x, + y_true, + y_pred, + tag_suffix, + ) in pl_module.plot_adapter.iter_plot_samples(data, output_tensor): fig = plot_predicted_multilevel_flat_sample( plot_parameters_dict, self.per_sample, @@ -964,7 +977,12 @@ def __init__( every_n_batches : int | None, optional Override for batch frequency, by default None """ - super().__init__(config, dataset_names=dataset_names, every_n_batches=every_n_batches, focus_area=focus_area) + super().__init__( + config, + dataset_names=dataset_names, + every_n_batches=every_n_batches, + focus_area=focus_area, + ) self.sample_idx = sample_idx self.parameters = parameters self.min_delta = min_delta @@ -1008,7 +1026,12 @@ def _plot( for name in self.parameters } - for x, y_true, y_pred, tag_suffix in pl_module.plot_adapter.iter_plot_samples(data, output_tensor): + for ( + x, + y_true, + y_pred, + tag_suffix, + ) in pl_module.plot_adapter.iter_plot_samples(data, output_tensor): fig = plot_power_spectrum( plot_parameters_dict_spectrum, latlons, @@ -1065,7 +1088,12 @@ def __init__( Override for batch frequency, by default None """ - super().__init__(config, dataset_names=dataset_names, every_n_batches=every_n_batches, focus_area=focus_area) + super().__init__( + config, + dataset_names=dataset_names, + every_n_batches=every_n_batches, + focus_area=focus_area, + ) self.sample_idx = sample_idx self.parameters = parameters self.precip_and_related_fields = precip_and_related_fields @@ -1117,7 +1145,12 @@ def _plot( for name in self.parameters } - for x, y_true, y_pred, tag_suffix in pl_module.plot_adapter.iter_plot_samples(data, output_tensor): + for ( + x, + y_true, + y_pred, + tag_suffix, + ) in pl_module.plot_adapter.iter_plot_samples(data, output_tensor): fig = plot_histogram( plot_parameters_dict_histogram, diff --git a/training/src/anemoi/training/diagnostics/plots.py b/training/src/anemoi/training/diagnostics/plots.py index 4a209de5f7..578ddee0c9 100644 --- a/training/src/anemoi/training/diagnostics/plots.py +++ b/training/src/anemoi/training/diagnostics/plots.py @@ -9,6 +9,7 @@ import logging +from typing import Literal import datashader as dsh import matplotlib.cm as cm @@ -24,6 +25,7 @@ from matplotlib.colors import Normalize from matplotlib.colors import TwoSlopeNorm from matplotlib.figure import Figure +from scipy.fft import dctn from scipy.interpolate import griddata from torch import Tensor @@ -165,15 +167,33 @@ def _interpolate_field( yt_field = yt xt_field = xt if yt is None else None - yp_i = griddata((pc_lon, pc_lat), yp_field, (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0) + yp_i = griddata( + (pc_lon, pc_lat), + yp_field, + (grid_pc_lon, grid_pc_lat), + method=method, + fill_value=0.0, + ) yt_i = None xt_i = None if yt_field is not None: - yt_i = griddata((pc_lon, pc_lat), yt_field, (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0) + yt_i = griddata( + (pc_lon, pc_lat), + yt_field, + (grid_pc_lon, grid_pc_lat), + method=method, + fill_value=0.0, + ) elif xt_field is not None: - xt_i = griddata((pc_lon, pc_lat), xt_field, (grid_pc_lon, grid_pc_lat), method=method, fill_value=0.0) + xt_i = griddata( + (pc_lon, pc_lat), + xt_field, + (grid_pc_lon, grid_pc_lat), + method=method, + fill_value=0.0, + ) return yp_i, yt_i, xt_i @@ -202,6 +222,13 @@ def _apply_nan_mask( return yp_i, yt_i, xt_i +def _is_global_coverage(pc_lon: np.ndarray, pc_lat: np.ndarray) -> bool: + """Heuristic to detect near-global coverage on equirectangular coordinates.""" + lon_span = np.nanmax(pc_lon) - np.nanmin(pc_lon) + lat_span = np.nanmax(pc_lat) - np.nanmin(pc_lat) + return lon_span >= (2.0 * np.pi * 0.95) and lat_span >= (np.pi * 0.95) + + def plot_power_spectrum( parameters: dict[str, int], latlons: np.ndarray, @@ -247,6 +274,7 @@ def plot_power_spectrum( ax = [ax] pc_lon, pc_lat = Projection.equirectangular().project(latlons) + resolved_spectrum_method = "sht" if _is_global_coverage(pc_lon, pc_lat) else "dct" # Calculate delta_lat on the projected grid delta_lat = abs(np.diff(pc_lat)) @@ -296,16 +324,16 @@ def plot_power_spectrum( if nan_flag: yp_i, yt_i, xt_i = _apply_nan_mask(yp_i, yt_i, xt_i) - amplitude_p = np.array(compute_spectra(yp_i)) + amplitude_p = np.array(compute_spectra(yp_i, method=resolved_spectrum_method)) if yt is not None: - amplitude_t = np.array(compute_spectra(yt_i)) + amplitude_t = np.array(compute_spectra(yt_i, method=resolved_spectrum_method)) ax[plot_idx].loglog( np.arange(1, amplitude_t.shape[0]), amplitude_t[1 : (amplitude_t.shape[0])], label="Truth (data)", ) else: - amplitude_x = np.array(compute_spectra(xt_i)) + amplitude_x = np.array(compute_spectra(xt_i, method=resolved_spectrum_method)) ax[plot_idx].loglog( np.arange(1, amplitude_x.shape[0]), amplitude_x[1 : (amplitude_x.shape[0])], @@ -326,27 +354,15 @@ def plot_power_spectrum( return fig -def compute_spectra(field: np.ndarray) -> np.ndarray: - """Compute spectral variability of a field by wavenumber. - - Parameters - ---------- - field : np.ndarray - lat lon field to calculate the spectra of - - Returns - ------- - np.ndarray - spectra of field by wavenumber - - """ +def _compute_spectra_sht(field: np.ndarray) -> np.ndarray: + """Compute spectra using spherical harmonic decomposition.""" try: from pyshtools.expand import SHGLQ from pyshtools.expand import SHExpandGLQ except ImportError as e: error_msg = ( "pyshtools is required to compute spherical harmonic power spectra. " - "It can be installed with the `plotting` dependency. `pip install anemoi-training[plotting]`.", + "It can be installed with the `plotting` dependency. `pip install anemoi-training[plotting]`." ) raise ImportError(error_msg) from e @@ -364,6 +380,65 @@ def compute_spectra(field: np.ndarray) -> np.ndarray: return np.sum(coeff_amp, axis=0) +def _compute_spectra_dct(field: np.ndarray) -> np.ndarray: + """Compute radial power spectrum for regional domains with 2-D DCT.""" + field = np.array(field) + dct_coeffs = dctn(field, type=2, norm="ortho") + variance = dct_coeffs**2 + + ni, nj = field.shape + ki, kj = np.meshgrid(np.arange(ni), np.arange(nj), indexing="ij") + alpha = np.sqrt((ki / ni) ** 2 + (kj / nj) ** 2) + + n_min = min(ni, nj) + k_max = n_min - 1 + bins = np.arange(1, k_max + 2) / n_min + band_idx = np.digitize(alpha, bins).astype(int) + + raw_var_full = np.bincount(band_idx.ravel(), weights=variance.ravel(), minlength=k_max + 2) + raw_var = raw_var_full[1 : k_max + 1].copy() + + mode_count_full = np.bincount(band_idx.ravel(), minlength=k_max + 2) + mode_count = mode_count_full[1 : k_max + 1].copy() + + k_g = k_max // 2 + k_ = np.arange(1, k_g + 1, dtype=int) + + alpha_p = np.zeros(k_max + 3, dtype=float) + alpha_p[1 : k_max + 1] = raw_var + var = 0.5 * alpha_p[2 * k_ - 1] + alpha_p[2 * k_] + 0.5 * alpha_p[2 * k_ + 1] + + mode_counts = np.zeros(k_max + 3, dtype=float) + mode_counts[1 : k_max + 1] = mode_count + gathered_mode_count = 0.5 * mode_counts[2 * k_ - 1] + mode_counts[2 * k_] + 0.5 * mode_counts[2 * k_ + 1] + + return np.where(gathered_mode_count > 0, var / gathered_mode_count, 0.0) + + +def compute_spectra(field: np.ndarray, method: Literal["sht", "dct"] = "sht") -> np.ndarray: + """Compute spectral variability by wavenumber. + + Parameters + ---------- + field : np.ndarray + Lat-lon field on a regular 2-D grid. + method : {"sht", "dct"}, optional + Spectral transform. "sht" for global spherical harmonics, + "dct" for regional-domain DCT. + + Returns + ------- + np.ndarray + Power spectrum by wavenumber. + """ + if method == "sht": + return _compute_spectra_sht(field) + if method == "dct": + return _compute_spectra_dct(field) + msg = f"Unknown spectra method: {method}" + raise ValueError(msg) + + def plot_histogram( parameters: dict[str, int], x: np.ndarray, @@ -467,7 +542,14 @@ def plot_histogram( alpha=0.7, label="Input" if y_true is None else "Truth (data)", ) - ax[plot_idx].bar(bins_yp[:-1], hist_yp, width=np.diff(bins_yp), color="red", alpha=0.7, label="Predicted") + ax[plot_idx].bar( + bins_yp[:-1], + hist_yp, + width=np.diff(bins_yp), + color="red", + alpha=0.7, + label="Predicted", + ) ax[plot_idx].set_title(variable_name) ax[plot_idx].set_xlabel(variable_name) @@ -845,7 +927,10 @@ def single_plot( dy, dx = ymax - ymin, xmax - xmin ybuffer, xbuffer = dy * 0.05, dx * 0.05 if transform is not None: - ax.set_extent([xmin - xbuffer, xmax + xbuffer, ymin - ybuffer, ymax + ybuffer], crs=transform) + ax.set_extent( + [xmin - xbuffer, xmax + xbuffer, ymin - ybuffer, ymax + ybuffer], + crs=transform, + ) else: ax.set_xlim((xmin - xbuffer, xmax + xbuffer)) ax.set_ylim((ymin - ybuffer, ymax + ybuffer)) @@ -1073,8 +1158,20 @@ def plot_rank_histograms( for plot_idx, (_variable_idx, variable_name) in enumerate(parameters.items()): rh_ = rh[:, plot_idx] - ax[plot_idx].bar(np.arange(0, n_ens + 1), rh_ / rh_.sum(), linewidth=1, color="blue", width=0.7) - ax[plot_idx].hlines(rh_.mean() / rh_.sum(), xmin=-0.5, xmax=n_ens + 0.5, linestyles="--", colors="red") + ax[plot_idx].bar( + np.arange(0, n_ens + 1), + rh_ / rh_.sum(), + linewidth=1, + color="blue", + width=0.7, + ) + ax[plot_idx].hlines( + rh_.mean() / rh_.sum(), + xmin=-0.5, + xmax=n_ens + 0.5, + linestyles="--", + colors="red", + ) ax[plot_idx].set_title(f"{variable_name[0]} ranks") _hide_axes_ticks(ax[plot_idx]) diff --git a/training/tests/unit/diagnostics/test_plotting_callbacks.py b/training/tests/unit/diagnostics/test_plotting_callbacks.py index 1ff259eb2b..55f1433b2c 100644 --- a/training/tests/unit/diagnostics/test_plotting_callbacks.py +++ b/training/tests/unit/diagnostics/test_plotting_callbacks.py @@ -234,7 +234,9 @@ def _make_pl_module_forecaster( pl_module.model.model._graph_data = {"data": MagicMock()} pl_module.model.model._graph_data["data"].__getitem__ = lambda _self, _k: MagicMock() graph_data = pl_module.model.model._graph_data["data"] - graph_data.__getitem__ = lambda k: torch.zeros(nlatlon, 2) if k == pl_module.model.model._graph_name_data else None + graph_data.__getitem__ = lambda k: ( + torch.zeros(nlatlon, 2) if k == pl_module.model.model._graph_name_data else None + ) # Use no-op output_mask pl_module.output_mask = {"data": NoOutputMask()} @@ -322,7 +324,13 @@ def test_process_forecaster_output_shapes(): # data: one sample from input_tensor (4 time steps); shape (time_steps, n_ens, nlatlon, nvar) assert data.shape == (1 + total_targets + 1, n_ens, nlatlon, nvar), data.shape # output_tensor: (output_times, n_step_output, n_ens, nlatlon, nvar) after mask - assert output_tensor.shape == (output_times, n_step_output, n_ens, nlatlon, nvar), output_tensor.shape + assert output_tensor.shape == ( + output_times, + n_step_output, + n_ens, + nlatlon, + nvar, + ), output_tensor.shape def test_process_temporal_downscaler_output_shapes(): @@ -356,7 +364,13 @@ def test_process_temporal_downscaler_output_shapes(): data, output_tensor = callback.process(pl_module, "data", outputs, batch) assert data.shape == (1 + total_targets + 1, n_ens, nlatlon, nvar), data.shape - assert output_tensor.shape == (pl_module.task.num_output_timesteps, 1, n_ens, nlatlon, nvar), output_tensor.shape + assert output_tensor.shape == ( + pl_module.task.num_output_timesteps, + 1, + n_ens, + nlatlon, + nvar, + ), output_tensor.shape def test_process_temporal_downscaler_multi_out_squeeze(): @@ -391,7 +405,13 @@ def test_process_temporal_downscaler_multi_out_squeeze(): # output_tensor: (num_output_timesteps, 1, n_ens, nlatlon, nvar) - 5D assert output_tensor.ndim == 5, output_tensor.shape - assert output_tensor.shape == (pl_module.task.num_output_timesteps, 1, 1, nlatlon, nvar), output_tensor.shape + assert output_tensor.shape == ( + pl_module.task.num_output_timesteps, + 1, + 1, + nlatlon, + nvar, + ), output_tensor.shape # ---- PlotLoss ---- @@ -436,7 +456,24 @@ def test_plot_loss_sort_and_color_by_parameter_group_with_groups(): }, ) # >15 parameters to trigger the grouping branch (<=15 keeps each param as its own group) - parameter_names = ["tp", "sp", "p0", "p1", "p2", "p3", "p4", "p5", "u10", "v10", "w0", "w1", "w2", "w3", "w4", "w5"] + parameter_names = [ + "tp", + "sp", + "p0", + "p1", + "p2", + "p3", + "p4", + "p5", + "u10", + "v10", + "w0", + "w1", + "w2", + "w3", + "w4", + "w5", + ] sort_idx, colors, xticks, legend_patches = callback.sort_and_color_by_parameter_group(parameter_names) assert sort_idx.shape == (len(parameter_names),) @@ -483,7 +520,10 @@ def test_plot_loss_temporal_downscaler(): "anemoi.training.diagnostics.callbacks.plot.argsort_variablename_variablelevel", return_value=np.arange(nvar), ), - patch("anemoi.training.diagnostics.callbacks.plot.plot_loss", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_loss", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -543,7 +583,10 @@ def test_plot_loss_diffusion(): "anemoi.training.diagnostics.callbacks.plot.argsort_variablename_variablelevel", return_value=np.arange(nvar), ), - patch("anemoi.training.diagnostics.callbacks.plot.plot_loss", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_loss", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -605,7 +648,10 @@ def test_plot_loss_forecaster(): "anemoi.training.diagnostics.callbacks.plot.argsort_variablename_variablelevel", return_value=np.arange(nvar), ), - patch("anemoi.training.diagnostics.callbacks.plot.plot_loss", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_loss", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -653,7 +699,10 @@ def test_plot_spectrum_temporal_downscaler(): with ( patch.object(callback, "_output_figure") as mock_output_figure, - patch("anemoi.training.diagnostics.callbacks.plot.plot_power_spectrum", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_power_spectrum", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -700,7 +749,10 @@ def test_plot_spectrum_forecaster(): with ( patch.object(callback, "_output_figure") as mock_output_figure, - patch("anemoi.training.diagnostics.callbacks.plot.plot_power_spectrum", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_power_spectrum", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -748,7 +800,10 @@ def test_plot_histogram_temporal_downscaler(): with ( patch.object(callback, "_output_figure") as mock_output_figure, - patch("anemoi.training.diagnostics.callbacks.plot.plot_histogram", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_histogram", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -795,7 +850,10 @@ def test_plot_histogram_forecaster(): with ( patch.object(callback, "_output_figure") as mock_output_figure, - patch("anemoi.training.diagnostics.callbacks.plot.plot_histogram", return_value=MagicMock()), + patch( + "anemoi.training.diagnostics.callbacks.plot.plot_histogram", + return_value=MagicMock(), + ), ): callback._plot( trainer, @@ -832,7 +890,10 @@ def test_plots_plot_loss_returns_figure(): x = np.array([0.1, 0.2, 0.15, 0.25]) colors = np.array(["C0", "C1", "C2", "C3"]) xticks = {"a": 0, "b": 1, "c": 2, "d": 3} - legend_patches = [mpatches.Patch(color="C0", label="a"), mpatches.Patch(color="C1", label="b")] + legend_patches = [ + mpatches.Patch(color="C0", label="a"), + mpatches.Patch(color="C1", label="b"), + ] fig = plot_loss(x, colors, xticks=xticks, legend_patches=legend_patches) @@ -900,6 +961,70 @@ def test_plots_plot_power_spectrum_returns_figure(): plt.close(fig) +def test_plots_plot_power_spectrum_auto_selects_sht_for_global_like_domain(): + """plot_power_spectrum uses SHT backend when the domain is detected as global.""" + from unittest.mock import patch + + import matplotlib.pyplot as plt + + from anemoi.training.diagnostics.plots import plot_power_spectrum + + parameters = {0: ("t2m", False)} + lat = np.linspace(-90, 90, 4) + lon = np.linspace(-180, 180, 4) + lat_grid, lon_grid = np.meshgrid(lat, lon, indexing="ij") + latlons = np.stack([lat_grid.ravel(), lon_grid.ravel()], axis=1) + nlatlon = latlons.shape[0] + rng = np.random.default_rng() + x = rng.standard_normal((nlatlon, 1)).astype(np.float64) + y_true = rng.standard_normal((nlatlon, 1)).astype(np.float64) + y_pred = rng.standard_normal((nlatlon, 1)).astype(np.float64) + + with ( + patch("anemoi.training.diagnostics.plots._is_global_coverage", return_value=True), + patch("anemoi.training.diagnostics.plots.compute_spectra", return_value=np.ones(8)) as mock_compute, + ): + fig = plot_power_spectrum(parameters, latlons, x, y_true, y_pred, min_delta=0.01) + + assert fig is not None + assert mock_compute.call_count >= 2 + assert all(call.kwargs["method"] == "sht" for call in mock_compute.call_args_list) + fig.clear() + plt.close(fig) + + +def test_plots_plot_power_spectrum_auto_selects_dct_for_regional_domain(): + """plot_power_spectrum uses DCT backend when the domain is detected as regional.""" + from unittest.mock import patch + + import matplotlib.pyplot as plt + + from anemoi.training.diagnostics.plots import plot_power_spectrum + + parameters = {0: ("t2m", False)} + lat = np.linspace(40, 50, 4) + lon = np.linspace(0, 10, 4) + lat_grid, lon_grid = np.meshgrid(lat, lon, indexing="ij") + latlons = np.stack([lat_grid.ravel(), lon_grid.ravel()], axis=1) + nlatlon = latlons.shape[0] + rng = np.random.default_rng() + x = rng.standard_normal((nlatlon, 1)).astype(np.float64) + y_true = rng.standard_normal((nlatlon, 1)).astype(np.float64) + y_pred = rng.standard_normal((nlatlon, 1)).astype(np.float64) + + with ( + patch("anemoi.training.diagnostics.plots._is_global_coverage", return_value=False), + patch("anemoi.training.diagnostics.plots.compute_spectra", return_value=np.ones(8)) as mock_compute, + ): + fig = plot_power_spectrum(parameters, latlons, x, y_true, y_pred, min_delta=0.01) + + assert fig is not None + assert mock_compute.call_count >= 2 + assert all(call.kwargs["method"] == "dct" for call in mock_compute.call_args_list) + fig.clear() + plt.close(fig) + + def test_plots_plot_predicted_multilevel_flat_sample_returns_figure(): """plot_predicted_multilevel_flat_sample returns a Figure and runs without error.""" import matplotlib.pyplot as plt From 36610ddd04b7bd8126cd6e8bf1fee9f9a058e4c5 Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Wed, 13 May 2026 10:02:43 +0200 Subject: [PATCH 2/7] Enhance documentation --- .../src/anemoi/training/diagnostics/plots.py | 56 ++++++++++++++++++- 1 file changed, 53 insertions(+), 3 deletions(-) diff --git a/training/src/anemoi/training/diagnostics/plots.py b/training/src/anemoi/training/diagnostics/plots.py index 639132f311..1da00865c4 100644 --- a/training/src/anemoi/training/diagnostics/plots.py +++ b/training/src/anemoi/training/diagnostics/plots.py @@ -223,7 +223,26 @@ def _apply_nan_mask( def _is_global_coverage(pc_lon: np.ndarray, pc_lat: np.ndarray) -> bool: - """Heuristic to detect near-global coverage on equirectangular coordinates.""" + """Heuristic to detect near-global coverage on equirectangular coordinates. + + Determines whether the provided longitude and latitude coordinates span + the majority of the globe. Returns True if coverage includes at least 95% + of the full longitude range (2π radians) and 95% of the full latitude range + (π radians), indicating global or near-global spatial coverage. + + Parameters + ---------- + pc_lon : np.ndarray + Longitude coordinates in radians, shape (n_points,) + pc_lat : np.ndarray + Latitude coordinates in radians, shape (n_points,) + + Returns + ------- + bool + True if the coordinates span >= 95% of the globe in both dimensions, + False otherwise. + """ lon_span = np.nanmax(pc_lon) - np.nanmin(pc_lon) lat_span = np.nanmax(pc_lat) - np.nanmin(pc_lat) return lon_span >= (2.0 * np.pi * 0.95) and lat_span >= (np.pi * 0.95) @@ -355,7 +374,19 @@ def plot_power_spectrum( def _compute_spectra_sht(field: np.ndarray) -> np.ndarray: - """Compute spectra using spherical harmonic decomposition.""" + """Compute spectral variability of a field by wavenumber. + + Parameters + ---------- + field : np.ndarray + lat lon field to calculate the spectra of + + Returns + ------- + np.ndarray + spectra of field by wavenumber + + """ try: from pyshtools.expand import SHGLQ from pyshtools.expand import SHExpandGLQ @@ -381,7 +412,26 @@ def _compute_spectra_sht(field: np.ndarray) -> np.ndarray: def _compute_spectra_dct(field: np.ndarray) -> np.ndarray: - """Compute radial power spectrum for regional domains with 2-D DCT.""" + """Compute radial power spectrum for regional domains with 2-D DCT. + + Implements the spectral analysis method described in: + + Denis, B., J. Côté, and R. Laprise, 2002: + "Spectral Decomposition of Two-Dimensional Atmospheric Fields on + Limited-Area Domains Using the Discrete Cosine Transform (DCT)." + Mon. Wea. Rev., 130, 1812-1829. + https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2 + + Parameters + ---------- + field : np.ndarray + lat lon field on a rectangular grid (Ni x Nj) with isotropic grid spacing. + + Returns + ------- + np.ndarray + Radial power spectrum P(k) computed from the input field. + """ field = np.array(field) dct_coeffs = dctn(field, type=2, norm="ortho") variance = dct_coeffs**2 From 2580d8811114cff27a8c5b4486aafd35b8fa2efc Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Fri, 15 May 2026 13:53:41 +0200 Subject: [PATCH 3/7] Log spectrum computation method (SHT vs DCT) and coverage type --- training/src/anemoi/training/diagnostics/plots.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/training/src/anemoi/training/diagnostics/plots.py b/training/src/anemoi/training/diagnostics/plots.py index 1da00865c4..698575a92b 100644 --- a/training/src/anemoi/training/diagnostics/plots.py +++ b/training/src/anemoi/training/diagnostics/plots.py @@ -295,6 +295,12 @@ def plot_power_spectrum( pc_lon, pc_lat = Projection.equirectangular().project(latlons) resolved_spectrum_method = "sht" if _is_global_coverage(pc_lon, pc_lat) else "dct" + LOGGER.info( + "Using %s method for power spectrum computation (%s coverage detected)", + resolved_spectrum_method.upper(), + "global" if resolved_spectrum_method == "sht" else "regional", + ) + # Calculate delta_lat on the projected grid delta_lat = abs(np.diff(pc_lat)) non_zero_delta_lat = delta_lat[delta_lat != 0] From 07dc2f3fa2387620edf71b3e009527e29ae949ee Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Thu, 21 May 2026 12:18:53 +0200 Subject: [PATCH 4/7] Remove invalid generic annotation on enumerate --- training/src/anemoi/training/diagnostics/plots.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/training/src/anemoi/training/diagnostics/plots.py b/training/src/anemoi/training/diagnostics/plots.py index 698575a92b..62acf5a460 100644 --- a/training/src/anemoi/training/diagnostics/plots.py +++ b/training/src/anemoi/training/diagnostics/plots.py @@ -249,7 +249,7 @@ def _is_global_coverage(pc_lon: np.ndarray, pc_lat: np.ndarray) -> bool: def plot_power_spectrum( - parameters: dict[str, int], + parameters: dict[int, tuple[str, bool]], latlons: np.ndarray, x: np.ndarray, y_true: np.ndarray | None, @@ -320,7 +320,7 @@ def plot_power_spectrum( regular_pc_lat = np.linspace(pc_lat.min(), pc_lat.max(), n_pix_lat) grid_pc_lon, grid_pc_lat = np.meshgrid(regular_pc_lon, regular_pc_lat) - for plot_idx, (variable_idx, (variable_name, diagnostic_only)) in enumerate[tuple[str, int]](parameters.items()): + for plot_idx, (variable_idx, (variable_name, diagnostic_only)) in enumerate(parameters.items()): xt = (x if x.ndim == 1 else x[..., variable_idx]).reshape(-1) yt = ( (y_true.reshape(-1) if y_true.ndim == 1 else y_true[..., variable_idx].reshape(-1)) From b6171fb9e1e6be4979efea2cc9eebbc867fe631b Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Thu, 21 May 2026 12:26:49 +0200 Subject: [PATCH 5/7] Add spectrum method-specific grid resolution to avoid DCT distortion --- training/src/anemoi/training/diagnostics/plots.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/training/src/anemoi/training/diagnostics/plots.py b/training/src/anemoi/training/diagnostics/plots.py index 62acf5a460..bff72af445 100644 --- a/training/src/anemoi/training/diagnostics/plots.py +++ b/training/src/anemoi/training/diagnostics/plots.py @@ -315,7 +315,13 @@ def plot_power_spectrum( # Define a regular grid for interpolation n_pix_lat = int(np.floor(abs(pc_lat.max() - pc_lat.min()) / min_delta_lat)) - n_pix_lon = (n_pix_lat - 1) * 2 + 1 # 2*lmax + 1 + if resolved_spectrum_method == "sht": + n_pix_lon = 2 * n_pix_lat - 1 # for SHT, need 2*lmax+1 points in longitude + else: + n_pix_lon = int( + np.floor(abs(pc_lon.max() - pc_lon.min()) / min_delta_lat), + ) # for DCT, use same resolution in lon and lat + regular_pc_lon = np.linspace(pc_lon.min(), pc_lon.max(), n_pix_lon) regular_pc_lat = np.linspace(pc_lat.min(), pc_lat.max(), n_pix_lat) grid_pc_lon, grid_pc_lat = np.meshgrid(regular_pc_lon, regular_pc_lat) From 10cdead08cf65001f88961a9fab754967170932f Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Thu, 21 May 2026 15:07:05 +0200 Subject: [PATCH 6/7] Limit the values of alpha to follow Denis et al. --- training/src/anemoi/training/diagnostics/plots.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/training/src/anemoi/training/diagnostics/plots.py b/training/src/anemoi/training/diagnostics/plots.py index bff72af445..123838e2e2 100644 --- a/training/src/anemoi/training/diagnostics/plots.py +++ b/training/src/anemoi/training/diagnostics/plots.py @@ -450,17 +450,20 @@ def _compute_spectra_dct(field: np.ndarray) -> np.ndarray: ni, nj = field.shape ki, kj = np.meshgrid(np.arange(ni), np.arange(nj), indexing="ij") + + # Denis et al. (2002) retain only coefficients with 0 < alpha < 1. alpha = np.sqrt((ki / ni) ** 2 + (kj / nj) ** 2) + valid = (alpha > 0.0) & (alpha < 1.0) n_min = min(ni, nj) k_max = n_min - 1 bins = np.arange(1, k_max + 2) / n_min - band_idx = np.digitize(alpha, bins).astype(int) + band_idx = np.digitize(alpha[valid], bins).astype(int) - raw_var_full = np.bincount(band_idx.ravel(), weights=variance.ravel(), minlength=k_max + 2) + raw_var_full = np.bincount(band_idx, weights=variance[valid], minlength=k_max + 2) raw_var = raw_var_full[1 : k_max + 1].copy() - mode_count_full = np.bincount(band_idx.ravel(), minlength=k_max + 2) + mode_count_full = np.bincount(band_idx, minlength=k_max + 2) mode_count = mode_count_full[1 : k_max + 1].copy() k_g = k_max // 2 From 16f1efcbf3e48f2dafaffa4d186eac800768eeaf Mon Sep 17 00:00:00 2001 From: Christophe Bovalo Date: Fri, 22 May 2026 11:41:51 +0200 Subject: [PATCH 7/7] Add deterministic DCT spectrum regression test --- .../diagnostics/test_plotting_callbacks.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/training/tests/unit/diagnostics/test_plotting_callbacks.py b/training/tests/unit/diagnostics/test_plotting_callbacks.py index 21725d1da3..99b4a782dc 100644 --- a/training/tests/unit/diagnostics/test_plotting_callbacks.py +++ b/training/tests/unit/diagnostics/test_plotting_callbacks.py @@ -965,6 +965,27 @@ def test_plots_plot_power_spectrum_auto_selects_dct_for_regional_domain(): plt.close(fig) +def test_compute_spectra_dct_single_mode_has_peaked_spectrum(): + """compute_spectra(..., method='dct') concentrates energy in the expected radial bin.""" + from anemoi.training.diagnostics.plots import compute_spectra + + ni = nj = 32 + mode = 4 + i = np.arange(ni) + # Single DCT mode in the meridional direction, constant in zonal direction. + field = np.cos(np.pi * (i + 0.5) * mode / ni)[:, None] * np.ones((1, nj)) + + spectrum = compute_spectra(field, method="dct") + + assert spectrum.shape == (15,) + assert int(np.argmax(spectrum)) == 1 + dominant = float(np.max(spectrum)) + residual = float(np.sum(spectrum) - dominant) + # Keep tolerance loose while still ensuring strong concentration in one band. + assert dominant > 1.0 + assert residual < dominant * 1e-6 + + def test_plots_plot_predicted_multilevel_flat_sample_returns_figure(): """plot_predicted_multilevel_flat_sample returns a Figure and runs without error.""" import matplotlib.pyplot as plt