feat(training): use DCT in PlotSpectrum to compute spectra on regional domains

training/src/anemoi/training/diagnostics/callbacks/plot.py

@@ -478,7 +478,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)
@@ -812,8 +815,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,
         )
 
@@ -1219,7 +1222,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,
@@ -1337,7 +1345,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,

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,8 +222,34 @@ 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.
+
+    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)
+
+
 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,
@@ -247,6 +293,13 @@ 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"
+
+    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))
@@ -262,12 +315,18 @@ 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)
 
-    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))
@@ -296,16 +355,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,7 +385,7 @@ def plot_power_spectrum(
     return fig
 
 
-def compute_spectra(field: np.ndarray) -> np.ndarray:
+def _compute_spectra_sht(field: np.ndarray) -> np.ndarray:
     """Compute spectral variability of a field by wavenumber.
 
     Parameters
@@ -346,7 +405,7 @@ def compute_spectra(field: np.ndarray) -> np.ndarray:
     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 +423,87 @@ 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.
+
+    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
+
+    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[valid], bins).astype(int)
+
+    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, 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 +607,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)
@@ -903,7 +1050,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))
@@ -1131,8 +1281,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])