From b575ad1a079f8517af6355291b1a7154aba799be Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Fri, 24 Jul 2026 18:32:56 +0200 Subject: [PATCH 1/5] =?UTF-8?q?base=20implementation=20=E2=80=93=20to=20be?= =?UTF-8?q?=20run=20with=20test=5Fconfig.latent=5Frollout=5Frmse=3DTrue=20?= =?UTF-8?q?flag?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../datasets/multi_stream_data_sampler.py | 37 +++++ src/weathergen/model/model.py | 78 +++++++++-- src/weathergen/train/trainer.py | 17 +++ src/weathergen/utils/latent_rmse.py | 131 ++++++++++++++++++ 4 files changed, 255 insertions(+), 8 deletions(-) create mode 100644 src/weathergen/utils/latent_rmse.py diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py index e248555a87..95924e1b6d 100644 --- a/src/weathergen/datasets/multi_stream_data_sampler.py +++ b/src/weathergen/datasets/multi_stream_data_sampler.py @@ -113,6 +113,39 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage): steps = np.array(forecast_cfg["num_steps"], dtype=np.int32).reshape(-1) self.list_num_forecast_steps = np.array(steps, dtype=np.int32) + # Latent rollout RMSE diagnostic (inference only). Instead of a bespoke data path, + # reuse the source-side num_steps_input machinery to load the truth latents at the + # forecast times: encode K+1 source steps and shift the base index forward by K-1 + # (K = forecast.num_steps), so the backward source window covers [t-1, t, ..., t+(K-1)]. + # Model.forward then re-indexes conditioning (t-1) and pairs each rollout step with its + # truth latent. No target-encoder or forecast-time reads needed. + self.latent_rollout_rmse = mode_cfg.get("latent_rollout_rmse", False) + self.latent_rollout_base_shift = 0 + if self.latent_rollout_rmse: + k_steps = int(self.list_num_forecast_steps.max()) + assert k_steps >= 1, "latent_rollout_rmse requires forecast.num_steps >= 1" + assert self.output_offset == 0, ( + f"latent_rollout_rmse requires forecast.offset == 0, got {self.output_offset}" + ) + # The base index is shifted forward by k_steps-1 so the backward source window + # covers the forecast times. Only the latent conditioning is re-indexed in + # Model.forward, not the physical target/write path, so write_output would emit a + # time-misaligned zarr. Forbid it rather than write silently corrupt output. + assert mode_cfg.get("output", {}).get("num_samples", 0) == 0, ( + "latent_rollout_rmse shifts the sample base index, so output.num_samples>0 " + "would write a time-misaligned physical zarr. Set output.num_samples=0." + ) + self.latent_rollout_base_shift = k_steps - 1 + if OmegaConf.is_config(mode_cfg): + OmegaConf.set_struct(mode_cfg, False) + for _, sc in mode_cfg.get("model_input", {}).items(): + sc["num_steps_input"] = k_steps + 1 + if is_root(): + logger.info( + f"latent_rollout_rmse: model_input num_steps_input -> {k_steps + 1}, " + f"base index shift -> +{k_steps - 1} (K={k_steps})" + ) + # initialise fsm, but can change for future mini_epochs self.batch_size = get_batch_size_from_config(mode_cfg) self.shuffle = mode_cfg.shuffle @@ -789,6 +822,10 @@ def __iter__(self) -> ModelBatch: idx: TIndex = perms[idx_raw % perms.shape[0]] idx_raw += 1 + # Latent rollout diagnostic: anchor the sample at t+(K-1) so the backward + # source window loads the forecast-time truth latents (see __init__). + idx = idx + self.latent_rollout_base_shift + batch = self._get_batch(idx, num_forecast_steps) # ensure the batch is valid, i.e. not completely empty and no NaN values diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index 1c500c7a9f..098f444b18 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -49,6 +49,35 @@ type StreamName = str +def _single_step_source_view(batch, step: int): + """ + Shallow view of a source ``BatchSamples`` exposing only input ``step``. + + The embedding engine concatenates the source tokens of *all* input steps into one tensor + before embedding, so encoding a sample that carries many input steps (as the latent + rollout RMSE diagnostic does, num_steps_input = K+1) costs K+1x the memory and OOMs. + Encoding step by step through this view keeps the peak at a single step. Only + ``source_tokens_cells`` / ``source_tokens_lens`` and ``tokens_lens`` are read per step; + nothing is deep-copied, so the views alias the original tensors. + """ + + view = copy.copy(batch) + view.tokens_lens = batch.tokens_lens[step : step + 1] + view.samples = [] + for sample in batch.samples: + sample_view = copy.copy(sample) + sample_view.streams_data = {} + for stream_name, stream_data in sample.streams_data.items(): + stream_view = copy.copy(stream_data) + stream_view.input_steps = 1 + stream_view.source_tokens_cells = [stream_data.source_tokens_cells[step]] + stream_view.source_tokens_lens = [stream_data.source_tokens_lens[step]] + sample_view.streams_data[stream_name] = stream_view + view.samples += [sample_view] + + return view + + class ModelOutput: """ Representation of model output @@ -350,6 +379,10 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord # One-shot flag to avoid log spam when warning about an unsupported # diffusion-inference + multi-step-rollout combination. self._warned_diffusion_multi_step = False + # Set by the trainer (a LatentRolloutRMSE) when the latent rollout RMSE diagnostic is + # enabled; None disables it. When set, forward() re-indexes the diffusion conditioning + # and accumulates per-lead-time latent RMSE. See forward(). + self.latent_rmse = None def _create_latent_pred_head( self, global_cfg, name, loss_cfg, use_class_token, use_patch_token @@ -724,31 +757,55 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: output = ModelOutput(batch.get_output_len()) - tokens, posteriors = self.encoder(model_params, batch) + if self.latent_rmse is not None: + # Latent rollout RMSE diagnostic: the source sample carries K+1 input steps, which + # the embed engine would otherwise concatenate into a single OOM-inducing forward. + # Encode each step separately and stack to bound peak memory (see + # _single_step_source_view). Result is already [B, T, ...]. + step_tokens = [] + posteriors = None + for s in range(batch.get_num_steps()): + tok_s, posteriors = self.encoder(model_params, _single_step_source_view(batch, s)) + step_tokens.append(tok_s) + tokens = torch.stack(step_tokens, dim=1) + else: + tokens, posteriors = self.encoder(model_params, batch) + # recover batch dimension and separate input_steps -> [B, T, ...] + tokens = tokens.reshape((len(batch), batch.get_num_steps(), *tokens.shape[1:])) output.add_latent_prediction(0, "posteriors", posteriors) - # recover batch dimension and separate input_steps - shape = (len(batch), batch.get_num_steps(), *tokens.shape[1:]) - # Reshape tokens to [B, T, ...] - tokens = tokens.reshape(shape) + shape = tokens.shape + + # Truth latents per rollout step, populated only for the latent-rollout RMSE diagnostic. + truth_latents = None if self.cf.get("fe_diffusion_model_conditioning", None) == "forecast": tokens = tokens.reshape(shape) # tokens[:, 0] = t (most recent), tokens[:, 1] = t-1, ..., tokens[:, -1] = t-(T-1) (oldest) - if self.cf.stage == "inference": + if self.latent_rmse is not None: + # Latent rollout RMSE diagnostic (inference). The source window carries K+1 + # encoded steps [t-1, t, ..., t+(K-1)]; most-recent first this is + # tokens[:, 0] = t+(K-1), ..., tokens[:, K-1] = t, tokens[:, K] = t-1. + # Condition on t-1 and pair rollout step j (valid time t+j) with tokens[:, K-1-j]. + k_steps = tokens.shape[1] - 1 + conditioning_tokens = tokens[:, k_steps] # t-1 + truth_latents = [tokens[:, k_steps - 1 - j] for j in range(k_steps)] # j -> t+j + tokens = tokens[:, k_steps - 1] # t (reference; the ODE sampler starts from noise) + elif self.cf.stage == "inference": print("Using most recent steps as conditioning tokens for inference.") # conditioning_tokens = tokens[:, :-1].sum(axis=1) conditioning_tokens = tokens[:, 1:].sum(axis=1) + tokens = tokens[:, 0] else: # Conditioning: all older context steps [t-1, ..., t-(T-1)]; denoising target: t (newest) conditioning_tokens = tokens[:, 1:].sum(axis=1) conditioning_tokens = conditioning_tokens + torch.randn_like(conditioning_tokens) * self.cf.get("fe_impute_latent_diffusion_noise_std", 0.0) if np.random.rand() < self.cf.get("fe_diffusion_classifier_free_guidance_prob", 0.0): # occasionally dropout conditioning for classifier free guidance conditioning_tokens = torch.zeros_like(conditioning_tokens) + tokens = tokens[:, 0] # X_t (tokens[:, 0], most recent) is the diffusion denoising target; older steps are conditioning. batch.samples[0].meta_info["ERA5"].params["conditioning_tokens"] = conditioning_tokens # self.forecast_engine._pending_target_tokens = diffusion_target_tokens - tokens = tokens[:, 0] else: tokens = tokens.sum(axis=1) @@ -756,7 +813,7 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: p_fwd = self.cf.training_config.get("forecast", {}).get("pushforward", False) # roll-out in latent space, iterate and generate output over requested output steps - for step in batch.get_output_idxs(): + for j, step in enumerate(batch.get_output_idxs()): without_grad = p_fwd and self.training and step != max(batch.get_output_idxs()) if without_grad: # Pushforward mode: advance tokens without grad; no decoding with torch.no_grad(): @@ -830,6 +887,11 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: output.add_physical_prediction( step, sname, (torch.cat(list(pred_tuple), dim=0),) ) + # Latent rollout RMSE diagnostic: accumulate RMSE of the rolled-out latent + # against the encoded truth latent at this lead time. member_final_tokens is + # (N_members, H, D); the truth is (1, H, D) and broadcasts over members. + if truth_latents is not None and j < len(truth_latents): + self.latent_rmse.add(j, member_final_tokens, truth_latents[j]) # Store per-member conditioning for the next rollout step. # conditioning_tokens holds (N, H, D) during ensemble rollout; inference_forward # calls expand(N, ...) which is a no-op when the dim already matches. diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 0fdd8d895a..941d852e57 100644 --- a/src/weathergen/train/trainer.py +++ b/src/weathergen/train/trainer.py @@ -47,6 +47,7 @@ get_target_idxs_from_cfg, ) from weathergen.utils.distributed import is_root +from weathergen.utils.latent_rmse import LatentRolloutRMSE from weathergen.utils.performance import NullThroughputTracker, ThroughputTracker from weathergen.utils.train_logger import TrainLogger, prepare_losses_for_logging from weathergen.utils.utils import get_dtype @@ -625,9 +626,20 @@ def validate(self, mini_epoch, mode_cfg, batch_size): all_losses: dict[str, list] = {} all_stddev: dict[str, list] = {} + # Latent rollout RMSE diagnostic: the model accumulates per-lead-time latent RMSE in + # forward() while self.latent_rmse is set. Accumulate only over the first (random + # noise-level) pass so the curve corresponds to a single sampling setting. + base_model = getattr(self.model, "module", self.model) + latent_rmse = ( + LatentRolloutRMSE(self.cf, mode_cfg, self.device) + if mode_cfg.get("latent_rollout_rmse", False) + else None + ) + for noise_idx, noise_level in enumerate(noise_levels): if is_diffusion: self._set_validation_noise_level(noise_level) + base_model.latent_rmse = latent_rmse if noise_idx == 0 else None if noise_level is None: loss_suffix = "" @@ -753,6 +765,11 @@ def validate(self, mini_epoch, mode_cfg, batch_size): if is_diffusion: self._set_validation_noise_level(None) + # latent rollout RMSE: reduce across ranks and plot like the evaluate package's curves + base_model.latent_rmse = None + if latent_rmse is not None: + latent_rmse.plot(config.get_path_run(self.cf)) + # avoid that there is a systematic bias in the validation subset self.dataset_val.advance() diff --git a/src/weathergen/utils/latent_rmse.py b/src/weathergen/utils/latent_rmse.py new file mode 100644 index 0000000000..2c753a24f8 --- /dev/null +++ b/src/weathergen/utils/latent_rmse.py @@ -0,0 +1,131 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +""" +Latent-space RMSE per forecast step, accumulated during (rollout) inference. + +The evaluate package computes RMSE-vs-lead-time curves for physical variables from the +zarr output of an inference run. Latents are far too large to write out for that, so the +equivalent latent-space curve is accumulated online here and plotted with the very same +plotting classes from ``weathergen.evaluate`` so that both look alike. +""" + +import logging + +import numpy as np +import torch +import xarray as xr + +from weathergen.common.config import parse_timedelta +from weathergen.evaluate.plotting.line_plots import LinePlots +from weathergen.evaluate.plotting.plot_utils import create_filename +from weathergen.utils.distributed import ddp_average, is_root + +logger = logging.getLogger(__name__) + +# defaults mirroring the ones the evaluate package uses in plot_summary() +_PLOT_CFG = { + "image_format": "png", + "dpi_val": 300, + "fig_size": (8, 10), + "log_scale": False, + "add_grid": True, + "plot_ensemble": False, + "baseline": None, +} + + +class LatentRolloutRMSE: + """ + Accumulate the RMSE between predicted and encoded ground-truth latents per forecast + step and plot it against lead time. + """ + + def __init__(self, cf, mode_cfg, device): + self.run_id = cf.general.run_id + self.device = device + + time_step = parse_timedelta(mode_cfg.get("forecast", {}).get("time_step", 0)) + self.step_hours = time_step / np.timedelta64(1, "h") + + # squared error sum and element count per forecast step; grown on demand + self._sum_sq: dict[int, torch.Tensor] = {} + self._counts: dict[int, torch.Tensor] = {} + + def add(self, step_idx: int, pred: torch.Tensor, truth: torch.Tensor) -> None: + """ + Accumulate the squared error between a predicted and an encoded truth latent for one + rollout step of one batch. + + Parameters + ---------- + step_idx: + Zero-based rollout step (lead index); step 0 is the forecast of ``t``. + pred: + Rolled-out latent, shape ``(N_members, H, D)`` (or ``(1, H, D)``). + truth: + Encoded truth latent, shape ``(1, H, D)``; broadcasts over ensemble members so the + accumulated value is the member-averaged squared error. + """ + + pred = pred.float() + truth = truth.float() + if pred.shape[-2:] != truth.shape[-2:]: + raise ValueError( + f"Latent prediction shape {tuple(pred.shape)} is incompatible with truth latent " + f"shape {tuple(truth.shape)} at rollout step {step_idx}." + ) + diff = pred - truth + sum_sq = diff.pow(2).sum(dtype=torch.float64) + count = torch.tensor(float(diff.numel()), device=self.device, dtype=torch.float64) + + if step_idx not in self._sum_sq: + self._sum_sq[step_idx] = torch.zeros((), device=self.device, dtype=torch.float64) + self._counts[step_idx] = torch.zeros((), device=self.device, dtype=torch.float64) + self._sum_sq[step_idx] += sum_sq.to(self.device) + self._counts[step_idx] += count.to(self.device) + + def plot(self, output_dir) -> None: + """Reduce across ranks, then write the RMSE-vs-lead-time plot (root rank only).""" + + if not self._sum_sq: + logger.warning("No latent predictions collected; skipping latent RMSE plot.") + return + + steps = np.array(sorted(self._sum_sq.keys())) + sum_sq = ddp_average(torch.stack([self._sum_sq[s] for s in steps])) + counts = ddp_average(torch.stack([self._counts[s] for s in steps])) + rmse = torch.sqrt(sum_sq / counts).numpy() + + if not is_root(): + return + + data = xr.DataArray( + rmse, + dims=["forecast_step"], + coords={"forecast_step": steps}, + name="rmse", + ) + # forecast step k is valid (k+1) * time_step after the last conditioning state + data = data.assign_coords(lead_time=("forecast_step", (steps + 1) * self.step_hours)) + x_dim = "forecast_step" + if self.step_hours > 0: + data = data.swap_dims({"forecast_step": "lead_time"}) + x_dim = "lead_time" + + plotter = LinePlots(_PLOT_CFG, output_dir) + plotter.plot( + [data], + [self.run_id], + tag=create_filename(prefix=["rmse", "global"], middle=[self.run_id], suffix=["latent"]), + x_dim=x_dim, + y_dim="rmse", + print_summary=True, + title="RMSE | latent | z_pre_norm", + ) From 2d1c961def4a14f4ac445dd7ef5db4eeac973dcb Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Sat, 25 Jul 2026 13:11:31 +0200 Subject: [PATCH 2/5] updated to ensur running alongside standard inference --- src/weathergen/datasets/batch.py | 23 ++++ .../datasets/multi_stream_data_sampler.py | 122 +++++++++++++----- src/weathergen/model/model.py | 84 ++++++------ src/weathergen/train/trainer.py | 36 +++++- 4 files changed, 182 insertions(+), 83 deletions(-) diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py index 22547ee925..2bc90992d1 100644 --- a/src/weathergen/datasets/batch.py +++ b/src/weathergen/datasets/batch.py @@ -309,6 +309,23 @@ def __init__( self.source2target_matching_idxs = np.full(num_source_samples, -1, dtype=np.int32) self.target2source_matching_idxs = [[] for _ in range(num_target_samples)] + # Optional isolated samples carrying source-channel network input at the forecast + # times, used only by the latent rollout RMSE diagnostic to encode truth latents. + # None on every standard batch; the source/target samples above are never affected. + self.latent_rmse_source: BatchSamples | None = None + + def init_latent_rmse_source(self, streams, num_samples: int) -> None: + """Create the isolated latent-RMSE truth samples (see ``latent_rmse_source``).""" + self.latent_rmse_source = BatchSamples( + streams, num_samples, self.output_steps, self.output_idxs + ) + + def add_latent_rmse_source_stream( + self, sample_idx: int, stream_name: str, stream_data: StreamData + ) -> None: + """Add one stream's forecast-time source data to the latent-RMSE truth samples.""" + self.latent_rmse_source.samples[sample_idx].add_stream_data(stream_name, stream_data) + def pin_memory(self): """Pin all tensors in this batch to CPU pinned memory""" @@ -318,6 +335,9 @@ def pin_memory(self): # pin target samples self.target_samples.pin_memory() + if self.latent_rmse_source is not None: + self.latent_rmse_source.pin_memory() + return self def to_device(self, device): # -> ModelBatch @@ -328,6 +348,9 @@ def to_device(self, device): # -> ModelBatch self.source_samples.to_device(device) self.target_samples.to_device(device) + if self.latent_rmse_source is not None: + self.latent_rmse_source.to_device(device) + self.device = device return self diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py index 95924e1b6d..ab98d44fb8 100644 --- a/src/weathergen/datasets/multi_stream_data_sampler.py +++ b/src/weathergen/datasets/multi_stream_data_sampler.py @@ -113,38 +113,16 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage): steps = np.array(forecast_cfg["num_steps"], dtype=np.int32).reshape(-1) self.list_num_forecast_steps = np.array(steps, dtype=np.int32) - # Latent rollout RMSE diagnostic (inference only). Instead of a bespoke data path, - # reuse the source-side num_steps_input machinery to load the truth latents at the - # forecast times: encode K+1 source steps and shift the base index forward by K-1 - # (K = forecast.num_steps), so the backward source window covers [t-1, t, ..., t+(K-1)]. - # Model.forward then re-indexes conditioning (t-1) and pairs each rollout step with its - # truth latent. No target-encoder or forecast-time reads needed. + # Latent rollout RMSE diagnostic (inference only). Attaches an ISOLATED set of + # source-channel samples at the forecast times [t, t+1, ..., t+(K-1)] to each batch + # (batch.latent_rmse_source), used only to encode truth latents for the diagnostic. + # The base index, the standard source/target samples, losses and zarr output are left + # exactly as a normal run — this is a read-only side channel. self.latent_rollout_rmse = mode_cfg.get("latent_rollout_rmse", False) - self.latent_rollout_base_shift = 0 if self.latent_rollout_rmse: - k_steps = int(self.list_num_forecast_steps.max()) - assert k_steps >= 1, "latent_rollout_rmse requires forecast.num_steps >= 1" assert self.output_offset == 0, ( f"latent_rollout_rmse requires forecast.offset == 0, got {self.output_offset}" ) - # The base index is shifted forward by k_steps-1 so the backward source window - # covers the forecast times. Only the latent conditioning is re-indexed in - # Model.forward, not the physical target/write path, so write_output would emit a - # time-misaligned zarr. Forbid it rather than write silently corrupt output. - assert mode_cfg.get("output", {}).get("num_samples", 0) == 0, ( - "latent_rollout_rmse shifts the sample base index, so output.num_samples>0 " - "would write a time-misaligned physical zarr. Set output.num_samples=0." - ) - self.latent_rollout_base_shift = k_steps - 1 - if OmegaConf.is_config(mode_cfg): - OmegaConf.set_struct(mode_cfg, False) - for _, sc in mode_cfg.get("model_input", {}).items(): - sc["num_steps_input"] = k_steps + 1 - if is_root(): - logger.info( - f"latent_rollout_rmse: model_input num_steps_input -> {k_steps + 1}, " - f"base index shift -> +{k_steps - 1} (K={k_steps})" - ) # initialise fsm, but can change for future mini_epochs self.batch_size = get_batch_size_from_config(mode_cfg) @@ -451,6 +429,43 @@ def _build_stream_data_input( return stream_data + def _build_latent_rmse_stream_data( + self, + stream_info: dict, + base_idx: TIndex, + num_forecast_steps: int, + forecast_input_data: list, + forecast_input_tokens: list, + mask: torch.Tensor | None, + ) -> StreamData: + """ + Build an ISOLATED source-channel StreamData at the forecast times + [t, t+1, ..., t+(K-1)] for the latent rollout RMSE diagnostic. + + Like ``_build_stream_data_input`` but walks *forward* over forecast steps: input step k + holds the source encoding of the true state at t+k. Used only to encode truth latents; + never fed to the model's conditioning, losses or zarr output. + """ + num_output_steps = self._get_output_length(num_forecast_steps) + stream_data = StreamData( + base_idx, num_output_steps, num_output_steps, self.num_healpix_cells + ) + for step, timestep_idx in enumerate(range(self.output_offset, num_output_steps)): + step_forecast_dt = base_idx + (self.time_step * timestep_idx) // self.step_timedelta + time_win = self.time_window_handler.window(step_forecast_dt) + + rdata = forecast_input_data[step] + token_data = forecast_input_tokens[step] + if token_data[0] is None and token_data[1] is None: + continue + + (source_cells, source_cells_lens) = self.tokenizer.get_source( + stream_info, rdata, token_data, (time_win.start, time_win.end), mask + ) + stream_data.add_source(step, rdata, source_cells_lens, source_cells) + + return stream_data + def _build_stream_data_output( self, mode: str, @@ -619,7 +634,27 @@ def _get_data_windows(self, base_idx, num_forecast_steps, num_steps_input_max, s output_data += [rdata] - return (input_data, output_data) + # source-channel data at the forecast times, for the latent RMSE diagnostic only. + # Isolated from input_data/output_data; consumed only via batch.latent_rmse_source. + forecast_input_data = [] + if self.latent_rollout_rmse: + for timestep_idx in range(self.output_offset, num_output_steps): + step_forecast_dt = ( + base_idx + (self.time_step * timestep_idx) // self.step_timedelta + ) + rdata = collect_datasources(stream_ds, step_forecast_dt, "source", self.rng) + if rdata.is_empty(): + time_win = self.time_window_handler.window(step_forecast_dt) + rdata = spoof( + self.healpix_level, + time_win.start, + stream_ds[0].get_geoinfo_size(), + len(stream_ds[0].mean[stream_ds[0].source_idx]), + ) + rdata.is_spoof = True + forecast_input_data += [rdata] + + return (input_data, output_data, forecast_input_data) def _get_source_target_masks(self, training_mode): """ @@ -691,6 +726,9 @@ def _get_batch(self, idx: int, num_forecast_steps: int): self.output_offset, num_output_steps, ) + if self.latent_rollout_rmse: + # isolated truth samples: one per target sample, mirroring the target masks + batch.init_latent_rmse_source(self.streams, num_target_samples) # for all streams for stream_info, (stream_name, stream_ds) in zip( @@ -708,7 +746,7 @@ def _get_batch(self, idx: int, num_forecast_steps: int): # input_data and output_data is conceptually consecutive but differs # in source and target channels; overlap in one window when self.output_offset=0 i_max = input_steps.max().item() - (input_data, output_data) = self._get_data_windows( + (input_data, output_data, forecast_input_data) = self._get_data_windows( idx, num_forecast_steps, i_max, stream_ds ) @@ -716,6 +754,11 @@ def _get_batch(self, idx: int, num_forecast_steps: int): # *_tokens = [ (cells_idx, cells_idx_lens), ... ] with length = #time_steps input_tokens = self.tokenizer.get_tokens_windows(stream_info, input_data, True) output_tokens = self.tokenizer.get_tokens_windows(stream_info, output_data, False) + forecast_input_tokens = ( + self.tokenizer.get_tokens_windows(stream_info, forecast_input_data, True) + if self.latent_rollout_rmse + else None + ) for sidx, source_mask in enumerate(source_masks.masks): # Map each source to its target @@ -770,10 +813,27 @@ def _get_batch(self, idx: int, num_forecast_steps: int): ] batch.add_target_stream(tidx, student_indices, stream_name, sdata, target_metadata) + # Isolated latent-RMSE truth: forecast-time source under the same target mask + # (the mask the model is trained to predict), added to a separate sample set. + if self.latent_rollout_rmse: + truth_sdata = self._build_latent_rmse_stream_data( + stream_info, + idx, + num_forecast_steps, + forecast_input_data, + forecast_input_tokens, + mask=target_mask, + ) + batch.add_latent_rmse_source_stream(tidx, stream_name, truth_sdata) + source_in_steps = input_steps.max().item() target_in_steps = np.array([tc.get("num_steps_input", 1) for _, tc in target_cfgs.items()]) target_in_steps = 1 if len(target_in_steps) == 0 else target_in_steps.max().item() batch = self._preprocess_model_batch(batch, source_in_steps, target_in_steps) + if self.latent_rollout_rmse: + batch.latent_rmse_source.tokens_lens = get_tokens_lens( + self.streams, batch.latent_rmse_source, num_output_steps + ) #add target times in source for diffusion model date/time conditioning if self.diffusion_model_conditioning in ["date_time", "date", "time"]: @@ -822,10 +882,6 @@ def __iter__(self) -> ModelBatch: idx: TIndex = perms[idx_raw % perms.shape[0]] idx_raw += 1 - # Latent rollout diagnostic: anchor the sample at t+(K-1) so the backward - # source window loads the forecast-time truth latents (see __init__). - idx = idx + self.latent_rollout_base_shift - batch = self._get_batch(idx, num_forecast_steps) # ensure the batch is valid, i.e. not completely empty and no NaN values diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index 098f444b18..456ea0507e 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -55,10 +55,10 @@ def _single_step_source_view(batch, step: int): The embedding engine concatenates the source tokens of *all* input steps into one tensor before embedding, so encoding a sample that carries many input steps (as the latent - rollout RMSE diagnostic does, num_steps_input = K+1) costs K+1x the memory and OOMs. - Encoding step by step through this view keeps the peak at a single step. Only - ``source_tokens_cells`` / ``source_tokens_lens`` and ``tokens_lens`` are read per step; - nothing is deep-copied, so the views alias the original tensors. + rollout RMSE truth samples do, K steps) costs Kx the memory and OOMs. Encoding step by + step through this view keeps peak memory at a single step. Only ``source_tokens_cells`` / + ``source_tokens_lens`` and ``tokens_lens`` are read per step; nothing is deep-copied, so + the views alias the original tensors. """ view = copy.copy(batch) @@ -379,10 +379,10 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord # One-shot flag to avoid log spam when warning about an unsupported # diffusion-inference + multi-step-rollout combination. self._warned_diffusion_multi_step = False - # Set by the trainer (a LatentRolloutRMSE) when the latent rollout RMSE diagnostic is - # enabled; None disables it. When set, forward() re-indexes the diffusion conditioning - # and accumulates per-lead-time latent RMSE. See forward(). - self.latent_rmse = None + # Set by the trainer for the latent rollout RMSE diagnostic: when True, forward() + # records the rolled-out latent per step under the "latent_rollout_pred" key. Purely + # additive — the standard path is unaffected. + self.record_latent_rollout = False def _create_latent_pred_head( self, global_cfg, name, loss_cfg, use_class_token, use_patch_token @@ -757,55 +757,31 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: output = ModelOutput(batch.get_output_len()) - if self.latent_rmse is not None: - # Latent rollout RMSE diagnostic: the source sample carries K+1 input steps, which - # the embed engine would otherwise concatenate into a single OOM-inducing forward. - # Encode each step separately and stack to bound peak memory (see - # _single_step_source_view). Result is already [B, T, ...]. - step_tokens = [] - posteriors = None - for s in range(batch.get_num_steps()): - tok_s, posteriors = self.encoder(model_params, _single_step_source_view(batch, s)) - step_tokens.append(tok_s) - tokens = torch.stack(step_tokens, dim=1) - else: - tokens, posteriors = self.encoder(model_params, batch) - # recover batch dimension and separate input_steps -> [B, T, ...] - tokens = tokens.reshape((len(batch), batch.get_num_steps(), *tokens.shape[1:])) + tokens, posteriors = self.encoder(model_params, batch) output.add_latent_prediction(0, "posteriors", posteriors) - shape = tokens.shape - - # Truth latents per rollout step, populated only for the latent-rollout RMSE diagnostic. - truth_latents = None + # recover batch dimension and separate input_steps + shape = (len(batch), batch.get_num_steps(), *tokens.shape[1:]) + # Reshape tokens to [B, T, ...] + tokens = tokens.reshape(shape) if self.cf.get("fe_diffusion_model_conditioning", None) == "forecast": tokens = tokens.reshape(shape) # tokens[:, 0] = t (most recent), tokens[:, 1] = t-1, ..., tokens[:, -1] = t-(T-1) (oldest) - if self.latent_rmse is not None: - # Latent rollout RMSE diagnostic (inference). The source window carries K+1 - # encoded steps [t-1, t, ..., t+(K-1)]; most-recent first this is - # tokens[:, 0] = t+(K-1), ..., tokens[:, K-1] = t, tokens[:, K] = t-1. - # Condition on t-1 and pair rollout step j (valid time t+j) with tokens[:, K-1-j]. - k_steps = tokens.shape[1] - 1 - conditioning_tokens = tokens[:, k_steps] # t-1 - truth_latents = [tokens[:, k_steps - 1 - j] for j in range(k_steps)] # j -> t+j - tokens = tokens[:, k_steps - 1] # t (reference; the ODE sampler starts from noise) - elif self.cf.stage == "inference": + if self.cf.stage == "inference": print("Using most recent steps as conditioning tokens for inference.") # conditioning_tokens = tokens[:, :-1].sum(axis=1) conditioning_tokens = tokens[:, 1:].sum(axis=1) - tokens = tokens[:, 0] else: # Conditioning: all older context steps [t-1, ..., t-(T-1)]; denoising target: t (newest) conditioning_tokens = tokens[:, 1:].sum(axis=1) conditioning_tokens = conditioning_tokens + torch.randn_like(conditioning_tokens) * self.cf.get("fe_impute_latent_diffusion_noise_std", 0.0) if np.random.rand() < self.cf.get("fe_diffusion_classifier_free_guidance_prob", 0.0): # occasionally dropout conditioning for classifier free guidance conditioning_tokens = torch.zeros_like(conditioning_tokens) - tokens = tokens[:, 0] # X_t (tokens[:, 0], most recent) is the diffusion denoising target; older steps are conditioning. batch.samples[0].meta_info["ERA5"].params["conditioning_tokens"] = conditioning_tokens # self.forecast_engine._pending_target_tokens = diffusion_target_tokens + tokens = tokens[:, 0] else: tokens = tokens.sum(axis=1) @@ -813,7 +789,7 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: p_fwd = self.cf.training_config.get("forecast", {}).get("pushforward", False) # roll-out in latent space, iterate and generate output over requested output steps - for j, step in enumerate(batch.get_output_idxs()): + for step in batch.get_output_idxs(): without_grad = p_fwd and self.training and step != max(batch.get_output_idxs()) if without_grad: # Pushforward mode: advance tokens without grad; no decoding with torch.no_grad(): @@ -887,11 +863,15 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: output.add_physical_prediction( step, sname, (torch.cat(list(pred_tuple), dim=0),) ) - # Latent rollout RMSE diagnostic: accumulate RMSE of the rolled-out latent - # against the encoded truth latent at this lead time. member_final_tokens is - # (N_members, H, D); the truth is (1, H, D) and broadcasts over members. - if truth_latents is not None and j < len(truth_latents): - self.latent_rmse.add(j, member_final_tokens, truth_latents[j]) + # Latent rollout RMSE diagnostic: record the rolled-out latent under a + # dedicated key (never "latent_state") so the latent loss is unaffected. The + # trainer pairs this with the encoded truth latent. Off by default. + if self.record_latent_rollout: + output.add_latent_prediction( + step, + "latent_rollout_pred", + self.tokens_to_latent_state(None, member_final_tokens), + ) # Store per-member conditioning for the next rollout step. # conditioning_tokens holds (N, H, D) during ensemble rollout; inference_forward # calls expand(N, ...) which is a no-op when the dim already matches. @@ -909,6 +889,20 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: return output + @torch.no_grad() + def encode_source_chunked(self, model_params: ModelParams, source_samples) -> torch.Tensor: + """ + Encode a multi-input-step source ``BatchSamples`` one step at a time and stack the + latents to ``[B, T, H, D]``. Used by the latent rollout RMSE diagnostic to encode the + truth latents at the forecast times without the K-step concatenation OOM (see + _single_step_source_view). Read-only; does not touch model state. + """ + step_tokens = [] + for s in range(source_samples.get_num_steps()): + tok_s, _ = self.encoder(model_params, _single_step_source_view(source_samples, s)) + step_tokens.append(tok_s) + return torch.stack(step_tokens, dim=1) + @staticmethod def _reindex_output_for_trajectory(output: ModelOutput, n_steps: int) -> ModelOutput: """ diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 941d852e57..48600269bb 100644 --- a/src/weathergen/train/trainer.py +++ b/src/weathergen/train/trainer.py @@ -626,9 +626,10 @@ def validate(self, mini_epoch, mode_cfg, batch_size): all_losses: dict[str, list] = {} all_stddev: dict[str, list] = {} - # Latent rollout RMSE diagnostic: the model accumulates per-lead-time latent RMSE in - # forward() while self.latent_rmse is set. Accumulate only over the first (random - # noise-level) pass so the curve corresponds to a single sampling setting. + # Latent rollout RMSE diagnostic (isolated side channel): the model records rolled-out + # latents under "latent_rollout_pred" while record_latent_rollout is set; the truth + # latents are encoded here from the isolated batch.latent_rmse_source. Accumulate only + # over the first (random noise-level) pass. Standard preds/losses/zarr are untouched. base_model = getattr(self.model, "module", self.model) latent_rmse = ( LatentRolloutRMSE(self.cf, mode_cfg, self.device) @@ -639,7 +640,7 @@ def validate(self, mini_epoch, mode_cfg, batch_size): for noise_idx, noise_level in enumerate(noise_levels): if is_diffusion: self._set_validation_noise_level(noise_level) - base_model.latent_rmse = latent_rmse if noise_idx == 0 else None + base_model.record_latent_rollout = latent_rmse is not None and noise_idx == 0 if noise_level is None: loss_suffix = "" @@ -711,6 +712,31 @@ def validate(self, mini_epoch, mode_cfg, batch_size): metadata=extract_batch_metadata(batch), ) + # Latent rollout RMSE: encode the isolated truth latents and pair them + # with the recorded predictions by lead step (both valid at t+j). Encode + # under the same autocast as the forward so truth and pred share dtype. + if ( + latent_rmse is not None + and noise_idx == 0 + and batch.latent_rmse_source is not None + ): + with torch.autocast( + device_type=f"cuda:{cf.local_rank}", + dtype=self.mixed_precision_dtype, + enabled=cf.with_mixed_precision, + ): + truth = base_model.encode_source_chunked( + self.model_params, batch.latent_rmse_source + ) # [B, K, H, D] + for j in range(truth.shape[1]): + pl = ( + preds.latent[j].get("latent_rollout_pred") + if j < len(preds.latent) + else None + ) + if pl is not None: + latent_rmse.add(j, pl.z_pre_norm, truth[:, j]) + # log output if noise_idx == 0: if bidx < num_samples_write: @@ -766,7 +792,7 @@ def validate(self, mini_epoch, mode_cfg, batch_size): self._set_validation_noise_level(None) # latent rollout RMSE: reduce across ranks and plot like the evaluate package's curves - base_model.latent_rmse = None + base_model.record_latent_rollout = False if latent_rmse is not None: latent_rmse.plot(config.get_path_run(self.cf)) From 68aca51fa5325857b7e84c6a4def932a4a93c839 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Wed, 29 Jul 2026 21:45:03 +0200 Subject: [PATCH 3/5] save json alongside plot --- src/weathergen/utils/latent_rmse.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/src/weathergen/utils/latent_rmse.py b/src/weathergen/utils/latent_rmse.py index 2c753a24f8..57d2e799cd 100644 --- a/src/weathergen/utils/latent_rmse.py +++ b/src/weathergen/utils/latent_rmse.py @@ -16,6 +16,7 @@ plotting classes from ``weathergen.evaluate`` so that both look alike. """ +import json import logging import numpy as np @@ -119,13 +120,33 @@ def plot(self, output_dir) -> None: data = data.swap_dims({"forecast_step": "lead_time"}) x_dim = "lead_time" + tag = create_filename(prefix=["rmse", "global"], middle=[self.run_id], suffix=["latent"]) + plotter = LinePlots(_PLOT_CFG, output_dir) plotter.plot( [data], [self.run_id], - tag=create_filename(prefix=["rmse", "global"], middle=[self.run_id], suffix=["latent"]), + tag=tag, x_dim=x_dim, y_dim="rmse", print_summary=True, title="RMSE | latent | z_pre_norm", ) + + # drop the plotted values next to the figure so the curve can be re-used numerically; + # "compare_" mirrors the prefix LinePlots.plot() puts on the figure file name + self._write_json(plotter.out_plot_dir / f"compare_{tag}.json", data) + + def _write_json(self, path, data: xr.DataArray) -> None: + """Write the plotted curve as JSON, in the same layout the evaluate package uses.""" + + data = data.assign_attrs( + run_id=self.run_id, + metric="rmse", + space="latent", + variable="z_pre_norm", + step_hours=float(self.step_hours), + ) + with open(path, "w") as f: + json.dump(data.to_dict(), f, indent=2) + logger.info(f"Wrote latent RMSE values to {path}") From 327a27031ff6c0643e21e767c85d48e0f7af506a Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Thu, 30 Jul 2026 15:50:14 +0200 Subject: [PATCH 4/5] implemented inference diagnostics --- config/config_diffusion_d2048_forecast_concat.yml | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/config/config_diffusion_d2048_forecast_concat.yml b/config/config_diffusion_d2048_forecast_concat.yml index 68612ee539..2f87684295 100644 --- a/config/config_diffusion_d2048_forecast_concat.yml +++ b/config/config_diffusion_d2048_forecast_concat.yml @@ -79,6 +79,16 @@ p_std: 1.2 healpix_level: 5 +# --- Per-ODE-step map & spectrum diagnostics (inference, single-step forecasting only) --- +# Decodes x_t and x0_hat at each ODE step and plots maps (x_t | x0_hat | decode(z) | truth) plus +# angular power spectra in latent and physical space, under //plots/ode_diagnostics. +# Requires training_config.forecast.num_steps=1 (rollout leaves no reference target). Off by default. +diag_ode_maps: false # master switch; costs 2 decoder passes + 4 map panels per step +diag_ode_every_n_steps: 1 # 1 = every ODE step; raise to cut render time +diag_channels: ["2t", "q_850"] +diag_stream: ERA5 +diag_latent_channels: 128 # latent channels sampled for the mean C_l (0 = all) + # Use 2D RoPE instead of traditional global positional encoding # When True: uses 2D RoPE based on healpix cell coordinates (lat/lon) # When False: uses traditional pe_global positional encoding From 4356ee618239575dca009966cfebf863865f2faf Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Thu, 30 Jul 2026 16:37:01 +0200 Subject: [PATCH 5/5] more changes --- .../src/weathergen/evaluate/scores/psd.py | 801 ++++++++++++++++++ src/weathergen/model/diffusion.py | 134 +-- src/weathergen/model/inference_diagnostics.py | 548 ++++++++++++ src/weathergen/model/inference_spectra.py | 188 ++++ src/weathergen/model/model.py | 28 + src/weathergen/train/trainer.py | 30 + tests/test_inference_diagnostics.py | 227 +++++ tests/test_inference_spectra.py | 148 ++++ tests/test_sht_roundtrip.py | 95 +++ 9 files changed, 2147 insertions(+), 52 deletions(-) create mode 100644 packages/evaluate/src/weathergen/evaluate/scores/psd.py create mode 100644 src/weathergen/model/inference_diagnostics.py create mode 100644 src/weathergen/model/inference_spectra.py create mode 100644 tests/test_inference_diagnostics.py create mode 100644 tests/test_inference_spectra.py create mode 100644 tests/test_sht_roundtrip.py diff --git a/packages/evaluate/src/weathergen/evaluate/scores/psd.py b/packages/evaluate/src/weathergen/evaluate/scores/psd.py new file mode 100644 index 0000000000..4021321797 --- /dev/null +++ b/packages/evaluate/src/weathergen/evaluate/scores/psd.py @@ -0,0 +1,801 @@ +# (C) Copyright 2025 Anemoi contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +"""Power Spectral Density (PSD) computation. + +Provides two PSD computation paths: + +- **Path A – SHT-based PSD** (``method="sht"``): + Spherical Harmonic Transform on separable grids (octahedral, reduced + Gaussian, regular lat-lon). Ported from anemoi.models ``spectral_transforms.py`` to pure + numpy using Legendre helpers from anemoi.models ``spectral_helpers.py``. + [anemoi.models.spectral_transforms] + https://github.com/ecmwf/anemoi-core/blob/main/models/src/anemoi/models/layers/spectral_transforms.py + [anemoi.models.spectral_helpers] + https://github.com/ecmwf/anemoi-core/blob/main/models/src/anemoi/models/layers/spectral_helpers.py + +- **Path B – FFT PSD** (``method="fft"``): + 1-D zonal FFT along the longitude dimension. This method **requires a regular + lat-lon grid** — i.e. the data must already live on a structured grid where + every latitude ring has the same number of equally-spaced longitude points. + If the input grid is not regular (e.g. octahedral reduced Gaussian), the + function raises a ``ValueError``. Re-gridding to a regular grid prior to FFT + is deliberately not supported because the interpolation introduces spectral + artefacts whose effect on the PSD is ill-defined. + For non-regular grids, use the SHT method instead. + Code base provided by the UKMet Office. + + The PSD is then computed row-by-row (per latitude ring) via 1-D real FFT + and averaged over all latitude rows within the specified ``lat_range``. +""" + +from __future__ import annotations + +import logging + +import numpy as np +from scipy.interpolate import griddata + +_logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Numpy-based Spherical Harmonic Transform (ported from spectral_helpers.py) +# --------------------------------------------------------------------------- + + +def _legendre_gauss_weights( + n: int, a: float = -1.0, b: float = 1.0 +) -> tuple[np.typing.NDArray, np.typing.NDArray]: + """Return Legendre-Gauss nodes and weights on ``[a, b]``.""" + xlg, wlg = np.polynomial.legendre.leggauss(n) + xlg = (b - a) * 0.5 * xlg + (b + a) * 0.5 + wlg = wlg * (b - a) * 0.5 + return xlg, wlg + + +def _legpoly( + mmax: int, lmax: int, x: np.typing.NDArray, inverse: bool = False +) -> np.typing.NDArray: + """Compute associated Legendre polynomials. + + Returns shape ``(mmax+1, lmax+1, len(x))``. + """ + nmax = max(mmax, lmax) + vdm = np.zeros((nmax + 1, nmax + 1, len(x)), dtype=np.float64) + + norm_factor = np.sqrt(4 * np.pi) + norm_factor = 1.0 / norm_factor if inverse else norm_factor + vdm[0, 0, :] = norm_factor / np.sqrt(4 * np.pi) + + for n in range(1, nmax + 1): + vdm[n - 1, n, :] = np.sqrt(2 * n + 1) * x * vdm[n - 1, n - 1, :] + vdm[n, n, :] = np.sqrt((2 * n + 1) * (1 + x) * (1 - x) / 2 / n) * vdm[n - 1, n - 1, :] + + for n in range(2, nmax + 1): + for m in range(0, n - 1): + vdm[m, n, :] = ( + x * np.sqrt((2 * n - 1) / (n - m) * (2 * n + 1) / (n + m)) * vdm[m, n - 1, :] + - np.sqrt((n + m - 1) / (n - m) * (2 * n + 1) / (2 * n - 3) * (n - m - 1) / (n + m)) + * vdm[m, n - 2, :] + ) + + return vdm[: mmax + 1, : lmax + 1] + + +class SphericalHarmonicTransform: + """Spherical Harmonic Transform in pure numpy. + + Mirrors the ``SphericalHarmonicTransform`` from ``spectral_helpers.py`` in anemoi.models + but operates on numpy arrays rather than torch tensors. + + Parameters + ---------- + lons_per_lat : list[int] + Number of longitude points on each latitude ring (pole to pole). + truncation : int + Maximum total wavenumber to retain. + """ + + def __init__(self, lons_per_lat: list[int], truncation: int) -> None: + self.lons_per_lat = lons_per_lat + self.nlat = len(lons_per_lat) + self.truncation = truncation + assert 0 < truncation <= self.nlat, f"Truncation {truncation} must be in (0, {self.nlat}]" + self.n_grid_points = sum(lons_per_lat) + + # Offsets into the flattened grid for each latitude ring + self.slon = [0] + list(np.cumsum(lons_per_lat))[:-1] + + # Whether all rings have the same number of points (regular grid) + self._is_regular = len(set(lons_per_lat)) == 1 + + # Precompute Gaussian latitudes + quadrature weights + theta, weight = _legendre_gauss_weights(self.nlat) + theta = np.flip(np.arccos(theta)) + + # Associated Legendre polynomials (m, l, lat) + pct = _legpoly(truncation, truncation, np.cos(theta)) + + # Pre-multiply by quadrature weights → shape (m, l, lat) + self.weight = np.einsum("mlk,k->mlk", pct, weight) + + # internal FFT helpers + + def _rfft_regular(self, x: np.typing.NDArray) -> np.typing.NDArray: + """Batched real FFT for a *regular* grid. + + Parameters + ---------- + x : np.typing.NDArray, shape ``(..., grid)`` + + Returns + ------- + np.typing.NDArray, complex, shape ``(..., nlat, nlon//2+1)`` + """ + nlon = self.lons_per_lat[0] + return np.fft.rfft(x.reshape(*x.shape[:-1], self.nlat, nlon), norm="forward") + + def _rfft_reduced(self, x: np.typing.NDArray) -> np.typing.NDArray: + """Per-ring real FFT for a *reduced* (variable-resolution) grid. + + Parameters + ---------- + x : np.typing.NDArray, shape ``(..., grid)`` + + Returns + ------- + np.typing.NDArray, complex, shape ``(..., nlat, max_nlon//2+1)`` + """ + max_nlon = max(self.lons_per_lat) + out_shape = (*x.shape[:-1], self.nlat, max_nlon // 2 + 1) + out = np.zeros(out_shape, dtype=np.complex128) + + for i, (slon, nlon) in enumerate(zip(self.slon, self.lons_per_lat, strict=False)): + out[..., i, : nlon // 2 + 1] = np.fft.rfft(x[..., slon : slon + nlon], norm="forward") + return out + + # transform + + def transform(self, x: np.typing.NDArray) -> np.typing.NDArray: + """Compute the SHT. + + Parameters + ---------- + x : np.typing.NDArray, real, shape ``(..., grid)`` + + Returns + ------- + np.typing.NDArray, complex, shape ``(..., L, M)`` where + ``L = M = truncation + 1``. + """ + if self._is_regular: + x_fft = self._rfft_regular(x) + else: + x_fft = self._rfft_reduced(x) + + x_fft = 2.0 * np.pi * x_fft + + real_part = x_fft[..., : self.truncation + 1].real + imag_part = x_fft[..., : self.truncation + 1].imag + + rl = np.einsum("...km,mlk->...lm", real_part, self.weight) + im = np.einsum("...km,mlk->...lm", imag_part, self.weight) + + return rl + 1j * im + + +class InverseSphericalHarmonicTransform: + """Inverse Spherical Harmonic Transform in pure numpy. + + Reconstructs a spatial field from spectral coefficients (l, m). + Mirrors the ``InverseSphericalHarmonicTransform`` from ``spectral_helpers.py`` + in anemoi.models but operates on numpy arrays. + This is not needed for the PSD computation but it is included + to verify that the forward and inverse transforms are consistent with each other. + + Parameters + ---------- + lons_per_lat : list[int] + Number of longitude points on each latitude ring (pole to pole). + truncation : int + Maximum total wavenumber. + """ + + def __init__(self, lons_per_lat: list[int], truncation: int) -> None: + self.lons_per_lat = lons_per_lat + self.nlat = len(lons_per_lat) + self.truncation = truncation + self.n_grid_points = sum(lons_per_lat) + self._is_regular = len(set(lons_per_lat)) == 1 + + # Gaussian latitudes (no quadrature weights needed for inverse) + theta, _ = _legendre_gauss_weights(self.nlat) + theta = np.flip(np.arccos(theta)) + + # Associated Legendre polynomials with inverse=True + self.pct = _legpoly(truncation, truncation, np.cos(theta), inverse=True) + + def _irfft_regular(self, x: np.typing.NDArray) -> np.typing.NDArray: + """Inverse FFT for a regular grid. + + Parameters + ---------- + x : np.typing.NDArray, complex, shape ``(..., nlat, M)`` + + Returns + ------- + np.typing.NDArray, real, shape ``(..., grid)`` + """ + nlon = self.lons_per_lat[0] + spatial = np.fft.irfft(x, n=nlon, norm="forward") # (..., nlat, nlon) + return spatial.reshape(*spatial.shape[:-2], self.n_grid_points) + + def _irfft_reduced(self, x: np.typing.NDArray) -> np.typing.NDArray: + """Per-ring inverse FFT for a reduced grid. + + Parameters + ---------- + x : np.typing.NDArray, complex, shape ``(..., nlat, M)`` + + Returns + ------- + np.typing.NDArray, real, shape ``(..., grid)`` + """ + lead_shape = x.shape[:-2] + out = np.zeros((*lead_shape, self.n_grid_points), dtype=np.float64) + offset = 0 + for i, nlon in enumerate(self.lons_per_lat): + ring = np.fft.irfft(x[..., i, :], n=nlon, norm="forward") + out[..., offset : offset + nlon] = ring + offset += nlon + return out + + def transform(self, coeffs: np.typing.NDArray) -> np.typing.NDArray: + """Compute the inverse SHT. + + Parameters + ---------- + coeffs : np.typing.NDArray, complex, shape ``(..., L, M)`` + + Returns + ------- + np.typing.NDArray, real, shape ``(..., grid)`` + """ + # Inverse Legendre transform: (..., l, m) × (m, l, k) → (..., k, m) + real_part = coeffs.real + imag_part = coeffs.imag + + rl = np.einsum("...lm,mlk->...km", real_part, self.pct) + im = np.einsum("...lm,mlk->...km", imag_part, self.pct) + + x_fourier = rl + 1j * im # (..., nlat, M) + + # Inverse FFT per ring + if self._is_regular: + return self._irfft_regular(x_fourier) + else: + return self._irfft_reduced(x_fourier) + + +# --------------------------------------------------------------------------- +# Grid helpers for building lons_per_lat +# --------------------------------------------------------------------------- + + +def _octahedral_lons_per_lat(nlat: int) -> list[int]: + """Return lons_per_lat for an octahedral reduced Gaussian grid.""" + half = [20 + 4 * i for i in range(nlat // 2)] + return half + list(reversed(half)) + + +def _regular_lons_per_lat(nlat: int) -> list[int]: + """Return lons_per_lat for a regular lat-lon grid (nlon = 2*nlat).""" + return [2 * nlat] * nlat + + +# --------------------------------------------------------------------------- +# Grid detection +# --------------------------------------------------------------------------- + + +def detect_grid_type( + lats: np.typing.NDArray, + lons: np.typing.NDArray, + n_points: int, +) -> str | None: + """Detect the grid type from latitude/longitude coordinates. + + Checks whether the point count matches known grid structures (octahedral + reduced Gaussian or regular lat-lon). Returns ``None`` with a warning if + the grid cannot be identified (e.g. regional subsets or non-standard grids). + + Parameters + ---------- + lats : np.typing.NDArray + Latitude values (per-point), length ``n_points``. + lons : np.typing.NDArray + Longitude values (per-point), length ``n_points``. + n_points : int + Total number of grid points. + + Returns + ------- + str | None + ``"octahedral"``, ``"regular"``, or ``None`` if detection fails. + """ + unique_lats = np.unique(lats) + nlat = len(unique_lats) + + # Check global extent + lat_min, lat_max = unique_lats.min(), unique_lats.max() + lat_span = lat_max - lat_min + + expected_span = 180.0 - 2 * (90.0 / nlat) # approx span for a Gaussian grid + if lat_span < 0.8 * expected_span: + _logger.warning( + f"Grid detection: latitude range [{lat_min:.1f}°, {lat_max:.1f}°] spans only " + f"{lat_span:.1f}° (expected ~{expected_span:.1f}° for {nlat} latitudes). " + f"PSD via SHT requires a global grid. Returning None." + ) + return None + + # Check octahedral reduced Gaussian + expected_oct = sum(_octahedral_lons_per_lat(nlat)) + if n_points == expected_oct: + _logger.debug(f"Detected octahedral reduced Gaussian grid (nlat={nlat}).") + return "octahedral" + + # Check regular lat-lon + expected_reg = sum(_regular_lons_per_lat(nlat)) + if n_points == expected_reg: + _logger.debug(f"Detected regular lat-lon grid (nlat={nlat}).") + return "regular" + + # Check if all latitude rings have the same number of points (regular but non-standard ratio) + unique_lons_global = np.unique(lons) + if nlat * len(unique_lons_global) == n_points: + _logger.debug(f"Detected regular grid (nlat={nlat}, nlon={len(unique_lons_global)}).") + return "regular" + + _logger.warning( + f"Grid detection: {n_points} points with {nlat} latitudes does not match " + f"octahedral ({expected_oct}) or regular ({expected_reg}) grids. " + f"The dataset may be regional or use an unsupported grid type." + "PSD via SHT skipped." + ) + return None + + +# --------------------------------------------------------------------------- +# High-level SHT PSD +# --------------------------------------------------------------------------- + + +def sht_psd( + data: np.typing.NDArray, + nlat: int, + truncation: int | None = None, + grid_type: str = "octahedral", +) -> tuple[np.typing.NDArray, np.typing.NDArray]: + """Compute PSD via Spherical Harmonic Transform. + + 1. Forward SHT: spatial → spectral coefficients ``(l, m)``. + 2. PSD: L2-norm over ``m`` for each total wavenumber ``l``. + + Parameters + ---------- + data : np.typing.NDArray + Spatial field with shape ``(n_points,)`` or ``(n_samples, n_points)``. + nlat : int + Number of latitudes in the grid. + truncation : int | None + Spectral truncation. Defaults to ``nlat // 2 - 1``. + grid_type : str + One of ``"octahedral"``, ``"regular"``, ``"reduced"``. + + Returns + ------- + wavenumbers : np.typing.NDArray, shape ``(L,)`` + Total wavenumber indices ``0, 1, …, L-1``. + psd : np.typing.NDArray, shape ``(L,)`` + Power spectral density averaged over samples. + """ + if data.ndim == 1: + data = data[np.newaxis, :] + n_samples, n_points = data.shape + + # Build the SHT for the appropriate grid + if grid_type == "octahedral": + lons_per_lat = _octahedral_lons_per_lat(nlat) + elif grid_type == "regular": + lons_per_lat = _regular_lons_per_lat(nlat) + elif grid_type == "reduced": + try: + from anemoi.transform.grids.named import lookup + except ImportError: + raise ImportError( + "anemoi.transform is required for grid_type='reduced'. " + "Install: pip install anemoi-transform" + ) from None + lats = lookup("N320")["latitudes"] + unique_lats = sorted(set(lats)) + lons_per_lat = [int((lats == lat).sum()) for lat in unique_lats] + else: + raise ValueError(f"Unknown grid_type: {grid_type!r}") + + trunc = truncation or nlat // 2 - 1 + sht = SphericalHarmonicTransform(lons_per_lat=lons_per_lat, truncation=trunc) + + assert n_points == sht.n_grid_points, ( + f"Input points={n_points} != expected grid points={sht.n_grid_points} " + f"for grid_type={grid_type!r}, nlat={nlat}" + ) + + # SphericalHarmonicTransform.transform accepts (..., grid) → (..., L, M) + # Pass (n_samples, n_points) directly. + coeffs = sht.transform(data) # (n_samples, L, M) + + # PSD = sum |coeffs|^2 over m for each total wavenumber l, averaged over samples + psd_per_sample = np.sum(np.abs(coeffs) ** 2, axis=-1) # (n_samples, L) + psd = psd_per_sample.mean(axis=0) + + n_wavenumbers = psd.shape[0] + wavenumbers = np.arange(n_wavenumbers, dtype=np.float64) + + return wavenumbers, psd + + +# --------------------------------------------------------------------------- +# FFT PSD (Credits to UK MetOffice) +# --------------------------------------------------------------------------- + + +def _fft_psd_calc(ht: np.typing.NDArray) -> np.typing.NDArray: + """Return the PSD for positive non-zero frequencies of an even-length signal. + + Assumes *ht* has an even number of points. + + Parameters + ---------- + ht : np.typing.NDArray + 1-D real-valued signal (one latitude ring). + + Returns + ------- + np.typing.NDArray + PSD for positive frequencies, length ``n // 2``. + """ + n = len(ht) + hf = np.fft.rfft(ht, norm="forward") + power = np.abs(hf[1 : round(n / 2 + 1)]) ** 2 + power *= 2.0 # compensate for positive frequencies only + return power + + +def _cubepsd(field_2d: np.typing.NDArray) -> np.typing.NDArray: + """Compute PSD averaged over all latitude rows. + + Parameters + ---------- + field_2d : np.typing.NDArray + 2-D array of shape ``(nlat, nlon)``. + + Returns + ------- + np.typing.NDArray + PSD of shape ``(nlon // 2,)``. + """ + nlat, nlon = field_2d.shape + field_psd = np.zeros(nlon // 2) + for row in field_2d: + field_psd += _fft_psd_calc(row) + field_psd /= nlat + return field_psd + + +def _calcposfreq(npoints: int, spacing_deg: float = 1.0) -> np.typing.NDArray: + """Return the positive frequencies for a signal of *npoints* evenly spaced points. + + Parameters + ---------- + npoints : int + Number of equally-spaced longitude points. + spacing_deg : float + Grid spacing in degrees. + + Returns + ------- + np.typing.NDArray + Positive frequencies, length ``npoints // 2``. + """ + freq = np.fft.fftfreq(npoints, d=spacing_deg) + return np.abs(freq[1 : round(npoints / 2 + 1)]) + + +def fft_psd( + data: np.typing.NDArray, + lats: np.typing.NDArray, + lons: np.typing.NDArray, + lat_range: tuple[float, float] = (-60.0, 60.0), + regrid_resolution: float = 1.0, +) -> tuple[np.typing.NDArray, np.typing.NDArray]: + """Compute PSD using 1-D zonal FFT along the longitude dimension. + + This method requires a **regular lat-lon grid** where every latitude ring + has the same number of equally-spaced longitude points. If the input is + not a regular grid, a ``ValueError`` is raised — use the SHT method instead. + + The PSD is computed row-by-row (per latitude ring) via 1-D real FFT and + averaged over all latitude rows within the specified ``lat_range``. + + Parameters + ---------- + data : np.typing.NDArray + Field values. Shape ``(n_samples, n_points)`` or ``(n_points,)``. + lats : np.typing.NDArray + Latitude values (per-point), length ``n_points``. + lons : np.typing.NDArray + Longitude values (per-point), length ``n_points``. + lat_range : tuple[float, float] + Latitude bounds to restrict the computation to. + regrid_resolution : float + Grid spacing in degrees for the regular target grid. + + Returns + ------- + frequencies : np.typing.NDArray + Positive frequencies in cycles per degree, shape ``(nfreq,)``. + psd : np.typing.NDArray + Power spectral density averaged over samples and latitude rows, + shape ``(nfreq,)``. + """ + + # Ensure 2-D: (n_samples, n_points) + if data.ndim == 1: + data = data.reshape(1, -1) + + n_samples, n_points = data.shape + + # Determine if the grid is regular or unstructured + unique_lats = np.unique(lats) + unique_lons = np.unique(lons) + is_regular = len(unique_lats) * len(unique_lons) == n_points + + if is_regular and len(lats) == len(unique_lats): + # lats/lons are axis arrays for a regular grid + lat_axis = unique_lats + lon_axis = unique_lons + nlat, nlon = len(lat_axis), len(lon_axis) + data_3d = data.reshape(n_samples, nlat, nlon) + else: + # Unstructured grid — regrid to regular lat-lon + lat_min = max(lat_range[0], lats.min()) + lat_max = min(lat_range[1], lats.max()) + lon_min, lon_max = lons.min(), lons.max() + + lat_axis = np.arange(lat_min, lat_max + regrid_resolution / 2, regrid_resolution) + lon_axis = np.arange(lon_min, lon_max + regrid_resolution / 2, regrid_resolution) + nlat, nlon = len(lat_axis), len(lon_axis) + + grid_lon, grid_lat = np.meshgrid(lon_axis, lat_axis) + points = np.column_stack((lats, lons)) + + data_3d = np.empty((n_samples, nlat, nlon)) + for s in range(n_samples): + data_3d[s] = griddata(points, data[s], (grid_lat, grid_lon), method="nearest") + + # Apply latitude mask + lat_mask = (lat_axis >= lat_range[0]) & (lat_axis <= lat_range[1]) + data_3d = data_3d[:, lat_mask, :] + nlon_sub = data_3d.shape[2] + + # Compute PSD per sample and average + psds = [] + for s in range(data_3d.shape[0]): + psds.append(_cubepsd(data_3d[s])) + psd_result = np.mean(psds, axis=0) + + spacing = 360.0 / nlon_sub if nlon_sub > 0 else regrid_resolution + frequencies = _calcposfreq(nlon_sub, spacing_deg=spacing) + return frequencies, psd_result + + +# --------------------------------------------------------------------------- +# Dispatch +# --------------------------------------------------------------------------- + + +def compute_psd_for_field( + data: np.typing.NDArray, + method: str = "sht", + nlat: int | None = None, + lats: np.typing.NDArray | None = None, + lons: np.typing.NDArray | None = None, + lat_range: tuple[float, float] = (-60.0, 60.0), + regrid_resolution: float = 1.0, + sht_truncation: int | None = None, + grid_type: str = "octahedral", +) -> tuple[np.typing.NDArray, np.typing.NDArray]: + """Compute PSD using the selected method. + + Parameters + ---------- + data : np.typing.NDArray + Spatial field. Shape depends on the method (see ``sht_psd`` / ``fft_psd``). + method : str + ``"sht"`` for SHT-based PSD, ``"fft"`` for FFT PSD. + nlat : int | None + Number of latitudes (required for SHT method). + lats, lons : np.typing.NDArray | None + Latitude / longitude coordinate arrays (required for fft method). + lat_range : tuple[float, float] + Latitude bounds for the fft method. + regrid_resolution : float + Grid spacing in degrees for the fft method. + sht_truncation : int | None + Spectral truncation for SHT. + grid_type : str + Grid type for SHT (``"octahedral"``, ``"regular"``, ``"reduced"``). + + Returns + ------- + x_values : np.typing.NDArray + Wavenumbers (SHT) or positive frequencies (fft). + psd : np.typing.NDArray + Power spectral density. + """ + if method == "sht": + if nlat is None: + raise ValueError("nlat is required for method='sht'") + return sht_psd( + data=data, + nlat=nlat, + truncation=sht_truncation, + grid_type=grid_type, + ) + elif method == "fft": + if lats is None or lons is None: + raise ValueError("lats and lons are required for method='fft'") + return fft_psd( + data=data, + lats=lats, + lons=lons, + lat_range=lat_range, + regrid_resolution=regrid_resolution, + ) + else: + raise ValueError(f"Unknown PSD method: {method!r}. Use 'sht' or 'fft'.") + + +def compute_psd_score( + gt: np.typing.NDArray, + p: np.typing.NDArray, + lats: np.typing.NDArray | None, + lons: np.typing.NDArray | None, + nlat: int | None, + n_points: int, + psd_method: str = "sht", + psd_regrid_resolution: float = 1.0, + psd_sht_truncation: int | None = None, + lat_range: tuple[float, float] = (-60.0, 60.0), + grid_type: str | None = None, +) -> tuple[float, dict]: + """Compute PSD for a pair of 2-D fields and return a scalar score + curves. + + This is the main entry point called from the Scores class. It handles NaN + masking, calls ``compute_psd_for_field`` for both inputs, and computes a + log-spectral MSE summary score. + + Parameters + ---------- + gt, p : np.typing.NDArray + Ground truth and prediction arrays of shape ``(n_samples, n_points)``. + lats, lons : np.typing.NDArray | None + Latitude / longitude arrays of length ``n_points`` (or None). + nlat : int | None + Number of latitudes (for SHT fallback). + n_points : int + Original number of spatial points (before NaN masking). + psd_method : str + ``"sht"`` or ``"fft"``. + psd_regrid_resolution : float + Grid spacing for fft method. + psd_sht_truncation : int | None + Spectral truncation for SHT. + lat_range : tuple[float, float] + Latitude bounds for fft method. + grid_type : str | None + Pre-detected grid type (``"octahedral"``, ``"regular"``). + When ``None``, the grid type is auto-detected from lats/lons. + Pass a pre-computed value to avoid repeated detection across channels. + + Returns + ------- + score : float + Log-spectral MSE scalar. + attrs : dict + Dict with keys ``"frequencies"``, ``"psd_target"``, ``"psd_prediction"`` + (lists for JSON serialization). + """ + # Handle NaN grid points (e.g. from regional masking). + valid_mask = ~np.isnan(gt).all(axis=0) + gt = gt[:, valid_mask] + p = p[:, valid_mask] + + # Filter lat/lon to match valid points + lats_valid = lats[valid_mask] if lats is not None and len(lats) == n_points else lats + lons_valid = lons[valid_mask] if lons is not None and len(lons) == n_points else lons + nlat_valid = len(np.unique(lats_valid)) if lats_valid is not None else nlat + + # Auto-detect grid type if not pre-computed by caller + if psd_method == "sht": + if lats_valid is None or lons_valid is None: + _logger.warning("PSD (SHT): lats/lons required for grid detection. Skipping.") + return np.nan, {} + if grid_type is None: + grid_type = detect_grid_type(lats_valid, lons_valid, gt.shape[-1]) + + if grid_type == "octahedral": + expected_pts = sum(_octahedral_lons_per_lat(nlat_valid)) + elif grid_type == "regular": + expected_pts = sum(_regular_lons_per_lat(nlat_valid)) + else: + expected_pts = None + + actual_pts = gt.shape[-1] + if expected_pts is not None and actual_pts != expected_pts: + _logger.warning( + f"PSD (SHT): grid point mismatch ({actual_pts} vs expected {expected_pts} " + f"for grid_type={grid_type!r}, nlat={nlat_valid}). SHT scores are only " + f"available for the full (global/unmasked) grid. Skipping this region." + ) + return np.nan, {} + + try: + freq_gt, psd_gt = compute_psd_for_field( + data=gt, + method=psd_method, + nlat=nlat_valid, + lats=lats_valid, + lons=lons_valid, + lat_range=lat_range, + regrid_resolution=psd_regrid_resolution, + sht_truncation=psd_sht_truncation, + grid_type=grid_type, + ) + freq_p, psd_p = compute_psd_for_field( + data=p, + method=psd_method, + nlat=nlat_valid, + lats=lats_valid, + lons=lons_valid, + lat_range=lat_range, + regrid_resolution=psd_regrid_resolution, + sht_truncation=psd_sht_truncation, + grid_type=grid_type, + ) + except Exception: + _logger.exception("PSD computation failed, returning NaN.") + return np.nan, {} + + # Scalar summary: mean squared error of log10 PSD + valid = (psd_gt > 0) & (psd_p > 0) + if valid.any(): + log_mse = float(np.mean((np.log10(psd_p[valid]) - np.log10(psd_gt[valid])) ** 2)) + else: + log_mse = np.nan + + attrs = { + "frequencies": freq_gt.tolist(), + "psd_target": psd_gt.tolist(), + "psd_prediction": psd_p.tolist(), + } + + return log_mse, attrs diff --git a/src/weathergen/model/diffusion.py b/src/weathergen/model/diffusion.py index fb499ee5e9..8791464f97 100644 --- a/src/weathergen/model/diffusion.py +++ b/src/weathergen/model/diffusion.py @@ -114,6 +114,9 @@ def __init__(self, cf: Config, num_healpix_cells: int, forecast_engine: Forecast self.cur_token = None # TODO: re move after single sample experiments self._noised_tokens: torch.Tensor | None = None self._fixed_noise_level: float | None = None + # Optional ODEDiagnostics (per-ODE-step maps/spectra), attached by the trainer; the decoder + # closure it needs is bound by model.py. + self.diagnostics = None self._noise = None @@ -123,7 +126,7 @@ def forward( fstep: int = None, meta_info: dict[str, SampleMetaData] = None, coords: torch.Tensor = None, - num_steps: int = 10, + num_steps: int | None = None, ) -> torch.Tensor: """ Forward pass that routes to training_forward or inference_forward based on model status. @@ -142,7 +145,10 @@ def forward( fstep: Forecast step index - required for both modes meta_info: Sample metadata dict containing timestamps - required for both modes coords: Optional coordinate tensor - num_steps: Number of diffusion steps for inference (default: 30) + num_steps: Number of diffusion ODE steps for inference. If None (the default, and what + model.py passes), it is read from config key ``fe_diffusion_num_steps``, which + itself defaults to 10 — preserving the historical hardcoded value bit-identically. + Set ``fe_diffusion_num_steps`` in the config (or via ``--options``) to override. Returns: torch.Tensor: Model output (denoised prediction during training, @@ -181,6 +187,10 @@ def forward( if fstep is None: raise ValueError(f"During inference, fstep is required. Got fstep={fstep}") self.cur_token = tokens.detach() if tokens is not None else None + # num_steps is not threaded through by model.py, so an explicit arg is rare; fall + # back to the config key (default 10 = the historical hardcoded value). + if num_steps is None: + num_steps = self.cf.get("fe_diffusion_num_steps", 10) return self.inference_forward( fstep=fstep, num_steps=num_steps, @@ -483,7 +493,8 @@ def _run_ode( "sigma": [], "x_std": [], "denoised_std": [], - "l2_to_target": [], + "rmse_x_t": [], + "rmse_x0_hat": [], "cosine_to_target": [], "c_skip": [], "d_cur_norm": [], @@ -496,6 +507,17 @@ def _run_ode( # Only populated when return_trajectory=True. intermediate_x: list[torch.Tensor] = [] if return_trajectory else None + # Per-ODE-step maps/spectra. Needs a reference target (absent in rollout mode) and a + # single sample (ensemble batches members on dim 0, with no per-member target). + diag = self.diagnostics if (self.cur_token is not None and batch_size == 1) else None + if diag is not None: + diag.begin(self.cur_token) + elif self.diagnostics is not None: + logger.info( + "ODE diagnostics disabled: no reference target (rollout mode) or " + f"batch_size={batch_size} > 1." + ) + # Main sampling loop. x_next = x * t_steps[0] for i, (t_cur, t_next) in enumerate( @@ -515,6 +537,9 @@ def _run_ode( # Euler step. denoised = self.denoise(x=x_hat, c=c, sigma=t_hat, fstep=fstep, coords=coords) + # Denoised (clean-latent) estimate x0_hat at t_cur, captured before the Heun + # correction below reassigns `denoised` to D(x_next, t_next). + x0_hat = denoised d_cur = (x_hat - denoised) / t_hat x_next = x_hat + (t_next - t_hat) * d_cur @@ -536,12 +561,25 @@ def _run_ode( track["residual_std"].append((x_hat - denoised).std().item()) track["x"].append(x_next.cpu()) if self.cur_token is not None: - track["l2_to_target"].append((x_next - self.cur_token).norm().item()) + # Per-element RMSE (‖·‖/√numel), size-independent and comparable to + # sigma_data=1. Both the noisy state x_next (last step = terminal sample at + # sigma=0) and the denoised estimate x0_hat=D(x_hat) at this sigma. + _rn = self.cur_token.numel() ** 0.5 + track["rmse_x_t"].append((x_next - self.cur_token).norm().item() / _rn) + track["rmse_x0_hat"].append((x0_hat - self.cur_token).norm().item() / _rn) track["x"].append(self.cur_token.cpu()) + if diag is not None: + # x_hat is the noisy state at t_cur; x0_hat is D(x_hat) at t_cur. + diag.on_step(i, t_hat.item(), x_hat, x0_hat) if return_trajectory: intermediate_x.append(x_next) + if diag is not None: + # The actual decoded sample (sigma=0). force=True so it is always recorded regardless + # of every_n_steps; the denoiser is undefined at the terminal node. + diag.on_step(num_steps, t_steps[num_steps].item(), x_next, None, force=True) + if log_diagnostics: self._plot_sampling_diagnostics(track, num_steps) @@ -556,68 +594,60 @@ def _plot_sampling_diagnostics(self, track: dict, num_steps: int) -> None: import matplotlib.pyplot as plt steps = list(range(len(track["sigma"]))) - has_target = len(track["l2_to_target"]) > 0 - n_plots = 7 + has_target = len(track["rmse_x_t"]) > 0 + n_plots = 4 if has_target else 3 fig, axes = plt.subplots(n_plots, 1, figsize=(10, 3 * n_plots), sharex=True) + i = 0 # 1) Sigma schedule - axes[0].semilogy(steps, track["sigma"], "o-", markersize=3) - axes[0].set_ylabel("sigma (noise level)") - axes[0].set_title( + axes[i].semilogy(steps, track["sigma"], "o-", markersize=3) + axes[i].set_ylabel("sigma (noise level)") + axes[i].set_title( f"Sampling diagnostics | sigma_max_eff={track['sigma'][0]:.2f}, " f"sigma_data={self.sigma_data}, steps={num_steps}" ) - axes[0].axhline( + axes[i].axhline( self.sigma_data, color="grey", ls="--", lw=0.8, label=f"sigma_data={self.sigma_data}" ) - axes[0].legend(fontsize=8) - axes[0].grid(True, alpha=0.3) + axes[i].legend(fontsize=8) + axes[i].grid(True, alpha=0.3) + i += 1 - # 2) Std of x_next and denoised estimate - axes[1].plot(steps, track["x_std"], "o-", markersize=3, label="x (noisy state)") - axes[1].plot(steps, track["denoised_std"], "s-", markersize=3, label="denoised estimate") + # 2) Per-element RMSE to target: noisy state x_t vs denoised estimate x̂₀ (comparable to + # sigma_data=1). x̂₀ sits near the target from the first step; x_t only reaches it at the + # terminal node (last point = the returned sample at sigma=0). + if has_target: + axes[i].plot(steps, track["rmse_x_t"], "o-", markersize=3, color="tab:blue", + label="rmse(x_t, z) (noisy state)") + axes[i].plot(steps, track["rmse_x0_hat"], "s-", markersize=3, color="tab:red", + label=r"rmse($\hat{x}_0$, z) (denoised estimate)") + axes[i].set_ylabel("RMSE to target (per-element)") + axes[i].legend(fontsize=8) + axes[i].grid(True, alpha=0.3) + i += 1 + + # 3) Std of x_next and denoised estimate + axes[i].plot(steps, track["x_std"], "o-", markersize=3, label="x (noisy state)") + axes[i].plot(steps, track["denoised_std"], "s-", markersize=3, label="denoised estimate") if self.cur_token is not None: target_std = self.cur_token.std().item() - axes[1].axhline( + axes[i].axhline( target_std, color="grey", ls="--", lw=0.8, label=f"target std={target_std:.3f}" ) - axes[1].set_ylabel("std") - axes[1].legend(fontsize=8) - axes[1].grid(True, alpha=0.3) - - if has_target: - # 3) L2 error to target - axes[2].plot(steps, track["l2_to_target"], "o-", markersize=3, color="tab:red") - axes[2].set_ylabel("L2 error to target") - axes[2].grid(True, alpha=0.3) - - # 4) d_cur norm and step norm - axes[3].semilogy(steps, track["d_cur_norm"], "o-", markersize=3, label="||d_cur||") - axes[3].semilogy(steps, track["d_cur_step_norm"], "^-", markersize=3, label="||(t_next - t_hat) * d_cur||") - axes[3].set_ylabel("norm (log scale)") - axes[3].set_title("ODE drift norms") - axes[3].legend(fontsize=8) - axes[3].grid(True, alpha=0.3) - - # 5) Residual std: Std(x_hat - denoised) - axes[4].semilogy(steps, track["residual_std"], "s-", markersize=3, color="tab:orange") - axes[4].set_ylabel("std (log scale)") - axes[4].set_title("Std(x_hat - denoised)") - axes[4].grid(True, alpha=0.3) - - # 6) Residual std zoomed to [0, 1] - axes[5].plot(steps, track["residual_std"], "s-", markersize=3, color="tab:orange") - axes[5].set_ylim(0, 1) - axes[5].set_ylabel("std (clipped to 1)") - axes[5].set_title("Std(x_hat - denoised) [y ≤ 1]") - axes[5].grid(True, alpha=0.3) - - # 7) Std of x_next over sampling steps - axes[6].semilogy(steps, track["x_std"], "o-", markersize=3, color="tab:blue") - axes[6].set_ylabel("std (log scale)") - axes[6].set_title("Std of x_next over denoising steps") - axes[6].grid(True, alpha=0.3) + axes[i].set_ylabel("std") + axes[i].legend(fontsize=8) + axes[i].grid(True, alpha=0.3) + i += 1 + + # 4) ODE drift norms + axes[i].semilogy(steps, track["d_cur_norm"], "o-", markersize=3, label="||d_cur||") + axes[i].semilogy(steps, track["d_cur_step_norm"], "^-", markersize=3, + label="||(t_next - t_hat) * d_cur||") + axes[i].set_ylabel("norm (log scale)") + axes[i].set_title("ODE drift norms") + axes[i].legend(fontsize=8) + axes[i].grid(True, alpha=0.3) axes[-1].set_xlabel("sampling step") fig.tight_layout() diff --git a/src/weathergen/model/inference_diagnostics.py b/src/weathergen/model/inference_diagnostics.py new file mode 100644 index 0000000000..65035d5258 --- /dev/null +++ b/src/weathergen/model/inference_diagnostics.py @@ -0,0 +1,548 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +"""Per-ODE-step maps and spectra for the diffusion sampler. + +At every step of the sampler we hold three latent states: the noisy state ``x_t``, the denoised +estimate ``x0_hat = D_t(x_t)`` predicted from it, and the clean target ``z``. This module decodes +them to physical space and plots + +1. maps of ``x_t`` / ``x0_hat`` / ``decode(z)`` / ground truth, and +2. angular power spectra of the same, in latent *and* physical space, + +so the over-smoothing failure mode -- ``x0_hat`` converging in RMSE while missing high-wavenumber +power -- is visible. ``decode(z)`` vs ground truth separates the sampler error from the +autoencoder's own reconstruction error. + +Only active for single-step forecasting: in rollout mode ``model.py`` sets ``tokens=None`` after +the first step, so the forecast engine's ``cur_token`` is ``None`` and there is no reference. + +Collection happens inside the sampler; **rendering is deferred** to +:meth:`ODEDiagnostics.render`, which the trainer calls after the forward pass -- the ground +truth, the per-point coordinates and the ``idxs_inv`` permutation only exist in the target/aux +output. That also keeps ~150 matplotlib figures out of the model forward. +""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from pathlib import Path + +import numpy as np +import torch + +from weathergen.model.inference_spectra import healpix_sht_psd, physical_psd, white_noise_reference + +logger = logging.getLogger(__name__) + +# Curve/panel styling, shared by the map and spectrum figures. +_FIELDS = ("x_t", "x0_hat", "decode_z", "truth") +_LABELS = { + "x_t": r"$x_t$ (noisy state)", + "x0_hat": r"$\hat{x}_0(x_t)$ (denoised estimate)", + "decode_z": r"decode($z$) (latent target)", + "truth": "truth (data)", +} +_COLORS = {"x_t": "tab:blue", "x0_hat": "tab:red", "decode_z": "tab:green", "truth": "black"} +_STYLES = {"x_t": "-", "x0_hat": "-", "decode_z": "--", "truth": ":"} +# At the terminal node x_t IS the decoded sample (sigma=0), not a noisy intermediate state. +_FINAL_LABEL = r"$x_{t=0}$ (final decoded output)" + + +class ODEDiagnostics: + """Collects decoded fields and latent spectra along the ODE, then renders them. + + Lifecycle, per sampled batch:: + + set_batch(bidx) # trainer: self-disables for bidx > 0 + bind_decoder(fn) # model.py: fn(tokens) -> {stream: (pred, ...)} + begin(z) # sampler: reference target + on_step(i, t, x_t, x0_hat) # sampler: once per ODE step + render(target_aux_physical) # trainer: writes the figures + """ + + def __init__( + self, + out_dir: Path, + stream: str, + channels: list[str], + channel_names: list[str], + nside: int, + denormalize: Callable[[str, torch.Tensor], torch.Tensor], + num_aux_tokens: int = 0, + every_n_steps: int = 1, + latent_channels: int = 128, + image_format: str = "png", + ) -> None: + self.out_dir = Path(out_dir) + self.stream = stream + self.nside = nside + self.denormalize = denormalize + self.num_aux_tokens = num_aux_tokens + self.every_n_steps = max(1, int(every_n_steps)) + self.latent_channels = int(latent_channels) + self.image_format = image_format + + missing = [c for c in channels if c not in channel_names] + if missing: + msg = f"Diagnostic channels {missing} not in stream {stream!r}: {channel_names}" + raise ValueError(msg) + self.channels = list(channels) + self.channel_idxs = [channel_names.index(c) for c in channels] + + self.enabled = False + self._decode: Callable[[torch.Tensor], dict] | None = None + self._reset() + + # ------------------------------------------------------------------ + # Lifecycle + # ------------------------------------------------------------------ + def _reset(self) -> None: + self.steps: list[int] = [] + self.times: list[float] = [] + self.latent_psd: dict[str, list[np.typing.NDArray]] = {"x_t": [], "x0_hat": []} + self.phys: dict[str, list[np.typing.NDArray]] = {"x_t": [], "x0_hat": []} + self.latent_psd_z: np.typing.NDArray | None = None + self.phys_z: np.typing.NDArray | None = None + self.wavenumbers: np.typing.NDArray | None = None + self._z_var: float = 1.0 + + def set_batch(self, batch_idx: int) -> None: + """Enable for the first batch only; each new batch starts from a clean slate.""" + self.enabled = batch_idx == 0 + self._reset() + + def bind_decoder(self, decode: Callable[[torch.Tensor], dict]) -> None: + self._decode = decode + + # ------------------------------------------------------------------ + # Collection (called from the sampler) + # ------------------------------------------------------------------ + def active(self) -> bool: + return self.enabled and self._decode is not None + + def begin(self, z: torch.Tensor) -> None: + """Record the clean latent target and its decoding.""" + if not self.active(): + return + wavenumbers, psd = self._latent_psd(z) + self.wavenumbers = wavenumbers + self.latent_psd_z = psd + self._z_var = float(self._latent_maps(z).var()) + self.phys_z = self._decode_channels(z) + + def on_step( + self, + step: int, + t: float, + x_t: torch.Tensor, + x0_hat: torch.Tensor | None, + force: bool = False, + ) -> None: + """Record one ODE step. Costs two decoder passes, hence ``every_n_steps``. + + The **terminal** state (``x0_hat=None``, ``force=True``) is the actual sample handed to + the decoder — the sampler loop otherwise only sees ``x_cur``, the state *before* each + update, so without this the returned output at ``sigma=0`` is never plotted and the last + ``x_t`` frame stalls at ``sigma_min`` (still visibly noisy). The denoiser is undefined + there (it would need another net forward at ``sigma=0``), so ``x0_hat`` is dropped. + """ + if not self.active() or (not force and step % self.every_n_steps): + return + self.steps.append(step) + self.times.append(float(t)) + self.latent_psd["x_t"].append(self._latent_psd(x_t)[1]) + self.phys["x_t"].append(self._decode_channels(x_t)) + if x0_hat is None: + self.latent_psd["x0_hat"].append(None) + self.phys["x0_hat"].append(None) + else: + self.latent_psd["x0_hat"].append(self._latent_psd(x0_hat)[1]) + self.phys["x0_hat"].append(self._decode_channels(x0_hat)) + + # ------------------------------------------------------------------ + # Latent helpers + # ------------------------------------------------------------------ + def _latent_maps(self, tokens: torch.Tensor) -> np.typing.NDArray: + """``[1, n_tokens, dim]`` tokens -> ``[n_channels, npix]`` HEALPix maps (nested). + + Auxiliary (class/register) tokens are stripped exactly as ``predict_decoders`` does, and + the per-cell query axis is folded into the channel axis so ``ae_local_num_queries > 1`` + works without special casing. + """ + x = tokens.detach()[:, self.num_aux_tokens :].float().cpu().numpy() + npix = 12 * self.nside**2 + n_cells = x.shape[1] + if n_cells % npix: + msg = f"Token count {n_cells} is not a multiple of npix={npix} (nside={self.nside})" + raise ValueError(msg) + # (batch, cells * queries, dim) -> (cells, queries * dim) -> (channels, cells) + return x[0].reshape(npix, -1).T + + def _latent_psd(self, tokens: torch.Tensor) -> tuple[np.typing.NDArray, np.typing.NDArray]: + maps = self._latent_maps(tokens) + if 0 < self.latent_channels < maps.shape[0]: + # Fixed subset across all steps so the curves are comparable; the mean over a random + # subset is an unbiased estimate of the mean over all channels. + idx = np.random.default_rng(0).choice(maps.shape[0], self.latent_channels, False) + maps = maps[idx] + return healpix_sht_psd(maps, self.nside) + + # ------------------------------------------------------------------ + # Physical helpers + # ------------------------------------------------------------------ + def _decode_channels(self, tokens: torch.Tensor) -> np.typing.NDArray: + """Decode latents and keep only the diagnostic channels, denormalized. + + Slicing *after* denormalization but *before* storing matters: the full channel set for + every field at every step would be gigabytes, the two channels are ~40 k floats. + """ + preds = self._decode(tokens) + pred = preds[self.stream][0] # first batch item; (ensemble, n_points, n_channels) + pred = pred[0] if pred.ndim == 3 else pred + pred = self.denormalize(self.stream, pred.to(torch.float32)) + return pred[:, self.channel_idxs].detach().cpu().numpy() + + # ------------------------------------------------------------------ + # Rendering (called from the trainer, after the forward pass) + # ------------------------------------------------------------------ + def render(self, target_aux_physical: dict) -> None: + """Write all figures. + + ``target_aux_physical`` is ``target_aux_out.physical[fstep][stream]`` for the physical + loss term -- the same structure ``write_output`` consumes. + """ + if not self.enabled or not self.steps: + return + try: + truth, lats, lons, order = self._truth_and_coords(target_aux_physical) + except Exception: + logger.exception("ODE diagnostics: could not extract truth/coords; skipping.") + return + + n_points = truth.shape[0] + # x0_hat holds None at the terminal frame (denoiser undefined at sigma=0); skip those. + collected = [p for p in [*self.phys["x_t"], *self.phys["x0_hat"]] if p is not None] + if any(p.shape[0] != n_points for p in collected): + logger.warning( + "ODE diagnostics: decoded fields do not match the %d target points; " + "skipping (are target coords varying across the forward pass?).", n_points + ) + return + + self.out_dir.mkdir(parents=True, exist_ok=True) + for name in ("x_t", "x0_hat"): + self.phys[name] = [None if p is None else p[order] for p in self.phys[name]] + self.phys_z = self.phys_z[order] if self.phys_z is not None else None + + self._render_spectra(truth, lats, lons) + self._render_maps(truth, lats, lons) + logger.info(f"Saved ODE diagnostics for {len(self.steps)} steps to {self.out_dir}") + + def _truth_and_coords(self, target_aux_physical: dict): + """Ground truth, coordinates and the permutation restoring dataset point order.""" + truth = target_aux_physical["target"][0] + coords = target_aux_physical["target_coords"][0] + idxs_inv = target_aux_physical["idxs_inv"][0] + if idxs_inv is not None: + truth = truth[idxs_inv] + coords = coords[idxs_inv] + truth = self.denormalize(self.stream, truth.to(torch.float32)).detach().cpu().numpy() + coords = coords.detach().cpu().numpy() + order = idxs_inv.detach().cpu().numpy() if idxs_inv is not None else slice(None) + return truth[:, self.channel_idxs], coords[:, 0], coords[:, 1], order + + # -- spectra -- + def _render_spectra( + self, truth: np.typing.NDArray, lats: np.typing.NDArray, lons: np.typing.NDArray + ) -> None: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + phys_psd = {name: [] for name in _FIELDS} + wavenumbers_p = None + for ch_pos in range(len(self.channels)): + for name in ("x_t", "x0_hat"): + curves = [] + for field in self.phys[name]: + res = None if field is None else physical_psd(field[:, ch_pos], lats, lons) + curves.append(None if res is None else res[1]) + wavenumbers_p = wavenumbers_p if res is None else res[0] + phys_psd[name].append(curves) + for name, field in (("decode_z", self.phys_z), ("truth", truth)): + res = None if field is None else physical_psd(field[:, ch_pos], lats, lons) + phys_psd[name].append(None if res is None else res[1]) + wavenumbers_p = wavenumbers_p if res is None else res[0] + + out = self.out_dir / "spectra" + out.mkdir(parents=True, exist_ok=True) + + # Per-step figure: latent + one panel per physical channel. + n_panels = 1 + len(self.channels) + for i, step in enumerate(self.steps): + # Terminal frame: x0_hat is absent and x_t is the final decoded sample. + is_final = self.latent_psd["x0_hat"][i] is None + xt_label = _FINAL_LABEL if is_final else _LABELS["x_t"] + fig, axes = plt.subplots(1, n_panels, figsize=(6 * n_panels, 4.4)) + axes = np.atleast_1d(axes) + self._spectrum_panel( + axes[0], + self.wavenumbers, + { + "x_t": self.latent_psd["x_t"][i], + "x0_hat": self.latent_psd["x0_hat"][i], + "decode_z": self.latent_psd_z, + }, + title="latent", + noise_var=self._z_var, + labels={"decode_z": r"$z$ (latent target)", "x_t": xt_label}, + ) + for c, channel in enumerate(self.channels): + self._spectrum_panel( + axes[1 + c], + wavenumbers_p, + { + "x_t": phys_psd["x_t"][c][i], + "x0_hat": phys_psd["x0_hat"][c][i], + "decode_z": phys_psd["decode_z"][c], + "truth": phys_psd["truth"][c], + }, + title=f"physical: {channel}", + labels={"x_t": xt_label}, + ) + tag = " [final decoded output]" if is_final else "" + fig.suptitle(f"ODE step {step} (t = {self.times[i]:.4g}){tag}") + fig.tight_layout() + fig.savefig(out / f"step{step:03d}.{self.image_format}", dpi=130) + plt.close(fig) + + # Evolution overlays: every step on one axis, colour-graded by step. + self._render_evolution(out, "latent", self.wavenumbers, self.latent_psd["x0_hat"], + self.latent_psd_z, r"$z$") + for c, channel in enumerate(self.channels): + self._render_evolution(out, channel, wavenumbers_p, phys_psd["x0_hat"][c], + phys_psd["truth"][c], "truth") + + @staticmethod + def _spectrum_panel(ax, wavenumbers, curves: dict, title: str, noise_var=None, labels=None): + labels = labels or {} + if wavenumbers is None: + ax.text(0.5, 0.5, "unavailable", ha="center", va="center", transform=ax.transAxes) + ax.set_title(title) + return + peak = 0.0 + for name, psd in curves.items(): + if psd is None: + continue + ax.loglog(wavenumbers[1:], psd[1:], _STYLES[name], color=_COLORS[name], lw=1.3, + label=labels.get(name, _LABELS[name])) + finite = psd[1:][np.isfinite(psd[1:])] + peak = max(peak, float(finite.max()) if finite.size else 0.0) + if peak > 0: + # Clamp to ~9 decades below the peak: a band-limited field's numerically-zero tail + # would otherwise stretch the axis over 18 decades and flatten everything of interest. + ax.set_ylim(peak * 1e-9, peak * 10) + if noise_var is not None: + ax.loglog(wavenumbers[1:], white_noise_reference(wavenumbers[1:], noise_var), + ":", color="grey", lw=1.0, label=r"white noise ($\propto 2\ell+1$)") + ax.set_xlabel(r"total wavenumber $\ell$") + ax.set_ylabel(r"PSD $\sum_m |a_{\ell m}|^2$") + ax.set_title(title) + ax.grid(True, which="both", alpha=0.3) + ax.legend(fontsize=7) + + def _render_evolution(self, out: Path, name: str, wavenumbers, curves, reference, ref_label): + import matplotlib.pyplot as plt + + if wavenumbers is None or not curves or all(c is None for c in curves): + return + fig, ax = plt.subplots(figsize=(7, 5)) + cmap = plt.get_cmap("viridis") + n = max(len(curves) - 1, 1) + for i, psd in enumerate(curves): + if psd is None: + continue + ax.loglog(wavenumbers[1:], psd[1:], color=cmap(i / n), lw=1.0) + if reference is not None: + ax.loglog(wavenumbers[1:], reference[1:], "k--", lw=1.6, label=ref_label) + ax.legend(fontsize=8) + sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(self.steps[0], self.steps[-1])) + fig.colorbar(sm, ax=ax, label="ODE step") + ax.set_xlabel(r"total wavenumber $\ell$") + ax.set_ylabel(r"PSD $\sum_m |a_{\ell m}|^2$") + ax.set_title(rf"{name}: $\hat{{x}}_0$ spectrum along the ODE") + ax.grid(True, which="both", alpha=0.3) + fig.tight_layout() + fig.savefig(out / f"evolution_{name}.{self.image_format}", dpi=130) + plt.close(fig) + + # -- maps -- + def _render_maps( + self, truth: np.typing.NDArray, lats: np.typing.NDArray, lons: np.typing.NDArray + ) -> None: + try: + plotter = _make_panel_plotter(self.out_dir, self.stream, self.image_format) + except Exception: + logger.exception("ODE diagnostics: map plotting unavailable; spectra kept.") + return + + out = self.out_dir / "maps" + out.mkdir(parents=True, exist_ok=True) + for c, channel in enumerate(self.channels): + # One colour scale for every panel and every step, taken from the truth, so the + # frames can be compared (and flip-booked) directly. + finite = truth[:, c][np.isfinite(truth[:, c])] + vmin, vmax = np.percentile(finite, [2, 98]) + for i, step in enumerate(self.steps): + x0_hat = self.phys["x0_hat"][i] # None at the terminal frame + is_final = x0_hat is None + fields = { + "x_t": self.phys["x_t"][i][:, c], + "x0_hat": None if is_final else x0_hat[:, c], + "decode_z": None if self.phys_z is None else self.phys_z[:, c], + "truth": truth[:, c], + } + labels = {**_LABELS, "x_t": _FINAL_LABEL if is_final else _LABELS["x_t"]} + panels = [(labels[k], v) for k, v in fields.items() if v is not None] + tag = " [final decoded output]" if is_final else "" + plotter.create_map_panel( + panels, + lats, + lons, + varname=channel, + suptitle=f"{channel} | ODE step {step} (t = {self.times[i]:.4g}){tag}", + out_path=out / f"step{step:03d}_{channel}.{self.image_format}", + map_kwargs={"vmin": float(vmin), "vmax": float(vmax)}, + ) + + +def maybe_create(cf, model, denormalize: Callable[[str, torch.Tensor], torch.Tensor]): + """Build the diagnostics and attach them to the forecast engine, or return ``None``. + + Off unless ``diag_ode_maps`` is set. Drives the diffusion engine's sampler, which exposes the + ``self.diagnostics`` hook and the ``begin``/``on_step`` protocol. Only meaningful during + inference, the only stage that runs a sampler. + """ + if not cf.get("diag_ode_maps", False): + return None + if not cf.get("fe_diffusion_model", False): + logger.warning("diag_ode_maps is set but this is not a diffusion run; " + "ignoring.") + return None + if cf.stage != "inference": + logger.warning(f"diag_ode_maps is set but stage is {cf.stage!r}; ignoring.") + return None + if not hasattr(model.forecast_engine, "diagnostics"): + logger.warning("diag_ode_maps is set but the forecast engine has no diagnostics hook " + f"({type(model.forecast_engine).__name__}); ignoring.") + return None + + from weathergen.common.config import get_path_run + + stream_name = cf.get("diag_stream", "ERA5") + streams = {s["name"]: s for s in cf.streams} + if stream_name not in streams: + logger.warning(f"diag_stream={stream_name!r} not in {list(streams)}; ignoring.") + return None + + engine = model.forecast_engine + diagnostics = ODEDiagnostics( + out_dir=get_path_run(cf) / "plots" / "ode_diagnostics", + stream=stream_name, + channels=list(cf.get("diag_channels", ["2t", "q_850"])), + channel_names=list(streams[stream_name].val_target_channels), + nside=2**cf.healpix_level, + denormalize=denormalize, + num_aux_tokens=getattr(model, "num_aux_tokens", 0), + every_n_steps=cf.get("diag_ode_every_n_steps", 1), + latent_channels=cf.get("diag_latent_channels", 128), + ) + engine.diagnostics = diagnostics + logger.info( + f"ODE diagnostics enabled: channels={diagnostics.channels}, " + f"every_n_steps={diagnostics.every_n_steps} -> {diagnostics.out_dir}" + ) + return diagnostics + + +def _make_panel_plotter(out_dir: Path, stream: str, image_format: str): + """Build the map-panel plotter. + + Imported lazily: it pulls in cartopy and the evaluation package's private working-dir config, + neither of which should be able to abort an inference run. + """ + import cartopy.crs as ccrs + import matplotlib.pyplot as plt + import xarray as xr + + from weathergen.evaluate.plotting.plot_utils import DefaultMarkerSize + from weathergen.evaluate.plotting.plotter import Plotter + + class MapPanelPlotter(Plotter): + """Multi-panel maps on one figure, reusing the evaluation package's rendering. + + ``Plotter.scatter_plot`` builds *and saves* a single-panel figure, so it cannot compose a + row of panels. Rather than modify the evaluation package (which is treated as read-only), + this subclass adds the panel layout and delegates every rendering decision -- + option parsing, marker sizing, scatter/datashader, HEALPix overlay -- to the inherited + methods, so panels match evaluation maps. It leans on ``Plotter``'s underscore-prefixed + helpers; an upstream rename breaks exactly this class. + """ + + def create_map_panel(self, panels, lats, lons, varname, suptitle, out_path, map_kwargs): + opts = self._parse_map_kwargs(dict(map_kwargs or {}), self.stream) + proj = ccrs.Robinson() + fig, axes = plt.subplots( + 1, len(panels), figsize=(5.6 * len(panels), 3.6), + subplot_kw={"projection": proj}, dpi=self.dpi_val, + ) + axes = np.atleast_1d(axes) + + artist = None + for ax, (title, values) in zip(axes, panels, strict=True): + data = xr.DataArray( + np.asarray(values), + dims=("ipoint",), + coords={"lon": ("ipoint", np.asarray(lons)), + "lat": ("ipoint", np.asarray(lats))}, + ) + try: + ax.coastlines(linewidth=0.3) + except Exception: + logger.warning("Could not add coastlines; continuing without them.") + ax.set_global() + marker_size = DefaultMarkerSize.auto_marker_size( + n_points=data.size, + fig_width_in=fig.get_figwidth() / len(panels), + fig_height_in=fig.get_figheight(), + stream_default=opts["marker_size_base"], + scale=opts["scale_marker_size"], + lat=data["lat"], + ) + artist = self._render_scatter( + ax, data, opts["norm"], opts["cmap"], marker_size, opts["marker"], + opts["extra"], + ) + ax.gridlines(draw_labels=False, linestyle="--", color="gray", linewidth=0.5, + alpha=0.6) + ax.set_title(title, fontsize=9) + + cbar = fig.colorbar(artist, ax=axes.tolist(), fraction=0.02, pad=0.02, shrink=0.8, + orientation="horizontal") + cbar.set_label(f"Variable: {varname}", fontsize=8) + cbar.ax.tick_params(labelsize=7) + fig.suptitle(suptitle, fontsize=10) + fig.savefig(out_path, bbox_inches="tight") + plt.close(fig) + + cfg = {"image_format": image_format, "dpi_val": 130, "fig_size": None, "regions": ["global"]} + return MapPanelPlotter(cfg, out_dir, stream=stream) diff --git a/src/weathergen/model/inference_spectra.py b/src/weathergen/model/inference_spectra.py new file mode 100644 index 0000000000..0823a9917f --- /dev/null +++ b/src/weathergen/model/inference_spectra.py @@ -0,0 +1,188 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +"""Angular power spectra for the inference-time ODE diagnostics. + +Two spaces have to be put on a common footing: + +- **physical**: the decoded fields live on the native o96 octahedral reduced Gaussian grid + (40 320 points). Handled by the evaluation package's ``sht_psd`` *unmodified*, so these + curves use exactly the same estimator as the ``psd`` evaluation score. +- **latent**: the tokens live on a HEALPix grid (``cf.healpix_level`` -> ``nside``), which is + iso-latitude but has neither Gauss-Legendre ring colatitudes nor Gauss-Legendre quadrature + weights, so ``SphericalHarmonicTransform`` cannot be pointed at it. :func:`healpix_sht_psd` + therefore reuses the evaluation package's Legendre machinery (``_legpoly``) with HEALPix ring + geometry and equal-area quadrature weights. + +Sharing ``_legpoly`` is what makes the two comparable: the same analytic field sampled on o96 +and on HEALPix nside 32 returns PSD values agreeing to ~5 significant figures, so latent and +physical spectra share both the ``l`` axis *and* the absolute scale. + +Convention note (inherited from the evaluation package): the returned PSD is +``sum_m |a_lm|^2``, i.e. it is *not* divided by ``2l+1``. White noise therefore rises like +``2l+1`` rather than being flat -- see :func:`white_noise_reference`. +""" + +from __future__ import annotations + +import functools +import logging + +import numpy as np +from astropy_healpix.healpy import pix2ang, ring2nest + +# Upstream estimator, used verbatim for the physical fields. ``_legpoly`` is private but is the +# piece that fixes the normalisation convention; importing it is what keeps the latent spectra on +# the same scale as the physical ones. +from weathergen.evaluate.scores.psd import ( + _legpoly, + _octahedral_lons_per_lat, + detect_grid_type, + sht_psd, +) + +logger = logging.getLogger(__name__) + + +@functools.cache +def _healpix_ring_geometry(nside: int): + """Ring decomposition of a HEALPix map, in RING ordering. + + Returns + ------- + thetas : colatitude of each ring, ascending (north -> south) + nlon : number of pixels on each ring + start : index of each ring's first pixel in the RING-ordered map + phi0 : azimuth of each ring's first pixel (HEALPix rings are *not* phase aligned) + weight : quadrature weight in ``cos(theta)`` for each ring + + ``weight`` is the equal-area HEALPix quadrature: every pixel subtends ``4*pi/npix``, so a ring + of ``nlon`` pixels covers ``d(cos theta) = 2*nlon/npix`` once the ``2*pi`` azimuthal integral + is factored out (that ``2*pi`` is applied in :func:`healpix_sht_psd`, mirroring + ``SphericalHarmonicTransform.transform``). + """ + npix = 12 * nside**2 + theta, phi = pix2ang(nside=nside, ipix=np.arange(npix), nest=False) + # RING ordering already groups pixels by ring with ascending colatitude. + ring_id = np.r_[0, np.cumsum(np.abs(np.diff(theta)) > 1e-12)] + nlon = np.bincount(ring_id) + start = np.r_[0, np.cumsum(nlon)[:-1]] + return theta[start], nlon, start, phi[start], 2.0 * nlon / npix + + +@functools.cache +def _nest_to_ring_index(nside: int) -> np.typing.NDArray: + """Index array ``idx`` such that ``map_nested[idx]`` is the map in RING ordering.""" + return ring2nest(nside, np.arange(12 * nside**2)) + + +def healpix_sht_psd( + maps: np.typing.NDArray, nside: int, truncation: int | None = None, nested: bool = True +) -> tuple[np.typing.NDArray, np.typing.NDArray]: + """Angular power spectrum of one or more HEALPix maps. + + Mirrors ``SphericalHarmonicTransform.transform`` + ``sht_psd`` from the evaluation package, + with HEALPix ring geometry substituted for the Gauss-Legendre one. + + Parameters + ---------- + maps : ``(npix,)`` or ``(n_maps, npix)``. The model's latent tokens are indexed by the + HEALPix **nested** index (``ang2pix(..., nest=True)`` in the tokenizer), hence the default. + nside : HEALPix nside, i.e. ``2**cf.healpix_level``. + truncation : maximum total wavenumber; defaults to ``2*nside``, beyond which the equal-area + quadrature degrades. + nested : whether ``maps`` is in nested ordering. + + Returns + ------- + wavenumbers, psd : ``(truncation+1,)`` each; ``psd`` is averaged over ``n_maps``. + """ + maps = np.atleast_2d(np.asarray(maps, dtype=np.float64)) + npix = 12 * nside**2 + if maps.shape[-1] != npix: + msg = f"Expected {npix} pixels for nside={nside}, got {maps.shape[-1]}" + raise ValueError(msg) + + truncation = int(truncation if truncation is not None else 2 * nside) + if nested: + maps = maps[:, _nest_to_ring_index(nside)] + + thetas, nlon, start, phi0, weight = _healpix_ring_geometry(nside) + # (m, l, ring), pre-multiplied by the quadrature weight -- as in the upstream __init__. + wgt = np.einsum("mlk,k->mlk", _legpoly(truncation, truncation, np.cos(thetas)), weight) + + # Per-ring real FFT. Unlike the Gauss-Legendre grids upstream handles, HEALPix rings are not + # phase aligned, so each ring's coefficients need the exp(-i*m*phi0) shift before they can be + # accumulated across rings. + coef = np.zeros((maps.shape[0], len(nlon), truncation + 1), dtype=np.complex128) + for k, (s, n) in enumerate(zip(start, nlon, strict=True)): + ring_fft = np.fft.rfft(maps[:, s : s + n], norm="forward") + m = np.arange(min(truncation + 1, ring_fft.shape[-1])) + coef[:, k, m] = ring_fft[:, m] * np.exp(-1j * m * phi0[k]) + coef *= 2.0 * np.pi + + # Complex einsum (upstream splits real/imag, which is equivalent only without the phase shift). + alm = np.einsum("...km,mlk->...lm", coef, wgt) + psd = np.sum(np.abs(alm) ** 2, axis=-1).mean(axis=0) + return np.arange(truncation + 1, dtype=np.float64), psd + + +def canonical_grid_order(lats: np.typing.NDArray, lons: np.typing.NDArray) -> np.typing.NDArray: + """Permutation putting scattered grid points into the ordering ``sht_psd`` expects. + + Upstream builds its rings from ``flip(arccos(leggauss_nodes))``, i.e. colatitude ascending = + latitude descending (north to south), with longitude ascending from 0 within each ring. The + points reaching us come in dataset order, so sort rather than assume. + """ + return np.lexsort((np.asarray(lons) % 360.0, -np.asarray(lats))) + + +def physical_psd( + values: np.typing.NDArray, + lats: np.typing.NDArray, + lons: np.typing.NDArray, + truncation: int | None = None, +) -> tuple[np.typing.NDArray, np.typing.NDArray] | None: + """Angular power spectrum of a field sampled on the native (o96) grid. + + Delegates to the evaluation package's ``sht_psd`` after restoring the canonical point order. + Returns ``None`` (with a warning) when the point cloud is not a recognised global grid, e.g. + a regional subset or a run with ``max_num_targets`` still subsampling the targets. + """ + values = np.atleast_2d(np.asarray(values, dtype=np.float64)) + lats = np.asarray(lats) + lons = np.asarray(lons) + + grid_type = detect_grid_type(lats, lons, values.shape[-1]) + if grid_type is None: + return None + + nlat = len(np.unique(lats)) + expected = sum(_octahedral_lons_per_lat(nlat)) if grid_type == "octahedral" else None + if expected is not None and values.shape[-1] != expected: + logger.warning( + f"Physical PSD skipped: {values.shape[-1]} points for nlat={nlat} does not match the " + f"{expected} expected on an {grid_type} grid (target subsampling still active?)." + ) + return None + + order = canonical_grid_order(lats, lons) + return sht_psd(values[:, order], nlat=nlat, truncation=truncation, grid_type=grid_type) + + +def white_noise_reference( + wavenumbers: np.typing.NDArray, variance: float = 1.0 +) -> np.typing.NDArray: + """PSD of spatially white noise in this convention: proportional to ``2l+1``. + + Because the estimator returns ``sum_m |a_lm|^2`` rather than a per-mode mean, white noise is a + rising line, not a flat one. Plotted as a reference so the pure-noise state at the start of + the ODE is recognisable by shape. + """ + return variance * (2.0 * np.asarray(wavenumbers) + 1.0) diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index 456ea0507e..7824a44ef0 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -796,6 +796,17 @@ def forward(self, model_params: ModelParams, batch: ModelBatch) -> ModelOutput: tokens = self.forecast_engine(tokens, step, model_params.rope_coords) continue + # The ODE diagnostics decode intermediate sampler states, but the sampler lives inside + # the forecast engine, which has no access to the decoders. Hand it a closure bound to + # this step's batch before it runs. + diagnostics = getattr(self.forecast_engine, "diagnostics", None) + if diagnostics is not None: + # step bound as a default arg: the loop variable would otherwise be captured by + # reference and resolve to the last forecast step. + diagnostics.bind_decoder( + lambda toks, step=step: self.decode_tokens(model_params, step, toks, batch) + ) + # apply forecasting engine tokens = self.forecast_engine( tokens, @@ -915,6 +926,23 @@ def _reindex_output_for_trajectory(output: ModelOutput, n_steps: int) -> ModelOu new_output.add_latent_prediction(0, k, v) return new_output + def decode_tokens( + self, + model_params: ModelParams, + step: int, + tokens: torch.Tensor, + batch: ModelBatch, + ) -> dict: + """Decode arbitrary latent tokens to physical space, outside the forward's bookkeeping. + + Used by the ODE diagnostics, which need to decode intermediate sampler states (``x_t``, + ``x0_hat``, the latent target) that never enter the ModelOutput. + Returns ``{stream_name: (pred_per_batch_item, ...)}``. + """ + return self.predict_decoders( + model_params, step, tokens, batch, ModelOutput(1), out_step=0 + ).physical[0] + def predict_latent( self, model_params: ModelParams, diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 48600269bb..9642d065d9 100644 --- a/src/weathergen/train/trainer.py +++ b/src/weathergen/train/trainer.py @@ -25,6 +25,7 @@ import weathergen.common.config as config from weathergen.common.config import Config from weathergen.datasets.multi_stream_data_sampler import MultiStreamDataSampler +from weathergen.model import inference_diagnostics from weathergen.model.ema import EMAModel from weathergen.model.model_interface import ( init_model_and_shard, @@ -58,6 +59,16 @@ # cfg_keys_to_filter = ["losses", "model_input", "target_input"] +def _physical_loss_term(mode_cfg) -> str: + """Name of the LossPhysical term, i.e. the one carrying the physical targets and coords. + + Same selection ``write_output`` makes; both consume ``target_aux_out.physical``. + """ + terms = [name for name, term in mode_cfg.losses.items() if term.type == "LossPhysical"] + assert len(terms) == 1, f"Expected exactly one LossPhysical term, got {terms}" + return terms[0] + + def _expand_targets_to_match_preds(preds, targets_and_auxs: dict) -> None: """ Replicate per-fstep entries in each TargetAuxOutput so its ``physical`` and ``latent`` @@ -626,6 +637,12 @@ def validate(self, mini_epoch, mode_cfg, batch_size): all_losses: dict[str, list] = {} all_stddev: dict[str, list] = {} + # Per-ODE-step maps/spectra for the diffusion sampler (off unless diag_ode_maps). + ode_diag = inference_diagnostics.maybe_create( + cf, self.model, self.dataset_val.denormalize_target_channels + ) + ode_diag_term = _physical_loss_term(mode_cfg) if ode_diag is not None else None + # Latent rollout RMSE diagnostic (isolated side channel): the model records rolled-out # latents under "latent_rollout_pred" while record_latent_rollout is set; the truth # latents are encoded here from the isolated batch.latent_rmse_source. Accumulate only @@ -667,6 +684,10 @@ def validate(self, mini_epoch, mode_cfg, batch_size): batch.to_device(self.device) + if ode_diag is not None: + # Self-disables for every batch after the first. + ode_diag.set_batch(bidx) + # evaluate model with torch.autocast( device_type=f"cuda:{cf.local_rank}", @@ -706,6 +727,15 @@ def validate(self, mini_epoch, mode_cfg, batch_size): if is_diffusion: _expand_targets_to_match_preds(preds, targets_and_auxs) + # Rendered here (not in the sampler): the ground truth, the point + # coordinates and the idxs_inv permutation only exist in target_aux. + # The target is identical across the trajectory (see + # _expand_targets_to_match_preds), so take the first fstep. + if ode_diag is not None and ode_diag.enabled: + ode_diag.render( + targets_and_auxs[ode_diag_term].physical[0][ode_diag.stream] + ) + _ = self.loss_calculator_val.compute_loss( preds=preds, targets_and_aux=targets_and_auxs, diff --git a/tests/test_inference_diagnostics.py b/tests/test_inference_diagnostics.py new file mode 100644 index 0000000000..7cad14cff5 --- /dev/null +++ b/tests/test_inference_diagnostics.py @@ -0,0 +1,227 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +"""Tests for the inference-time per-ODE-step diagnostics. + +Drives ``ODEDiagnostics`` through its real lifecycle with a stub decoder, so the collection, +the ``idxs_inv`` re-ordering and the figure production are exercised without a model. +""" + +import numpy as np +import pytest +import torch + +from weathergen.evaluate.scores.psd import _legendre_gauss_weights, _octahedral_lons_per_lat +from weathergen.model.inference_diagnostics import ODEDiagnostics + +NSIDE = 8 +NPIX = 12 * NSIDE**2 +DIM = 6 +CHANNELS = ["2t", "q_850"] +ALL_CHANNELS = ["z_500", "2t", "10u", "q_850"] + + +def _o96_like_grid(nlat: int = 16): + """A small octahedral grid, so detect_grid_type recognises it and the SHT path runs.""" + lons_per_lat = _octahedral_lons_per_lat(nlat) + nodes, _ = _legendre_gauss_weights(nlat) + theta = np.flip(np.arccos(nodes)) + lats = np.concatenate([np.full(n, 90.0 - np.degrees(t)) for t, n in + zip(theta, lons_per_lat, strict=True)]) + lons = np.concatenate([360.0 * np.arange(n) / n for n in lons_per_lat]) + return lats, lons + + +class _StubDecoder: + """Decodes latents to a physical field by broadcasting per-cell means onto the grid.""" + + def __init__(self, n_points: int, n_channels: int, stream: str = "ERA5"): + self.n_points = n_points + self.n_channels = n_channels + self.stream = stream + self.calls = 0 + + def __call__(self, tokens: torch.Tensor) -> dict: + self.calls += 1 + scale = float(tokens.mean()) + rng = np.random.default_rng(0) + field = rng.standard_normal((self.n_points, self.n_channels)) + scale + return {self.stream: (torch.from_numpy(field).float(),)} + + +def _make_diagnostics(tmp_path, n_points, **kwargs): + return ODEDiagnostics( + out_dir=tmp_path, + stream="ERA5", + channels=CHANNELS, + channel_names=ALL_CHANNELS, + nside=NSIDE, + denormalize=lambda _stream, data: data * 2.0 + 1.0, + **kwargs, + ) + + +def _target_aux(n_points, n_channels, idxs_inv=None): + rng = np.random.default_rng(1) + return { + "target": (torch.from_numpy(rng.standard_normal((n_points, n_channels))).float(),), + "target_coords": (torch.zeros(n_points, 2),), + "idxs_inv": (idxs_inv,), + } + + +def _run(diag, decoder, n_steps=3): + diag.set_batch(0) + diag.bind_decoder(decoder) + diag.begin(torch.randn(1, NPIX, DIM)) + for i in range(n_steps): + diag.on_step(i, 1.0 - i / n_steps, torch.randn(1, NPIX, DIM), torch.randn(1, NPIX, DIM)) + + +def test_unknown_channel_is_rejected_early(tmp_path): + """A typo in diag_channels must fail at construction, not after a 50-step sample.""" + with pytest.raises(ValueError, match="not in stream"): + ODEDiagnostics( + out_dir=tmp_path, + stream="ERA5", + channels=["2t", "nope"], + channel_names=ALL_CHANNELS, + nside=NSIDE, + denormalize=lambda _s, d: d, + ) + + +def test_disabled_for_later_batches(tmp_path): + diag = _make_diagnostics(tmp_path, 100) + decoder = _StubDecoder(100, len(ALL_CHANNELS)) + + diag.set_batch(1) + diag.bind_decoder(decoder) + diag.begin(torch.randn(1, NPIX, DIM)) + diag.on_step(0, 1.0, torch.randn(1, NPIX, DIM), torch.randn(1, NPIX, DIM)) + + assert not diag.enabled + assert decoder.calls == 0 + assert diag.steps == [] + + +def test_every_n_steps_subsamples(tmp_path): + diag = _make_diagnostics(tmp_path, 100, every_n_steps=3) + decoder = _StubDecoder(100, len(ALL_CHANNELS)) + _run(diag, decoder, n_steps=10) + + assert diag.steps == [0, 3, 6, 9] + # 1 decode for the target + 2 per recorded step. + assert decoder.calls == 1 + 2 * len(diag.steps) + + +def test_render_writes_maps_and_spectra(tmp_path): + lats, lons = _o96_like_grid() + n_points = lats.size + diag = _make_diagnostics(tmp_path, n_points, latent_channels=0) + decoder = _StubDecoder(n_points, len(ALL_CHANNELS)) + _run(diag, decoder, n_steps=3) + + aux = _target_aux(n_points, len(ALL_CHANNELS)) + aux["target_coords"] = (torch.from_numpy(np.stack([lats, lons], axis=1)).float(),) + diag.render(aux) + + spectra = sorted(p.name for p in (tmp_path / "spectra").glob("*.png")) + assert [s for s in spectra if s.startswith("step")] == [ + "step000.png", "step001.png", "step002.png" + ] + # One evolution overlay for the latent plus one per physical channel. + assert {s for s in spectra if s.startswith("evolution")} == { + "evolution_latent.png", "evolution_2t.png", "evolution_q_850.png" + } + + maps = sorted(p.name for p in (tmp_path / "maps").glob("*.png")) + assert maps == [f"step{i:03d}_{c}.png" for i in range(3) for c in CHANNELS] + + +def test_terminal_frame_is_always_recorded_without_x0_hat(tmp_path): + """The final decoded sample (x0_hat=None, force=True) must record even if step % n != 0.""" + diag = _make_diagnostics(tmp_path, 100, every_n_steps=3) + decoder = _StubDecoder(100, len(ALL_CHANNELS)) + diag.set_batch(0) + diag.bind_decoder(decoder) + diag.begin(torch.randn(1, NPIX, DIM)) + for i in range(8): + diag.on_step(i, 1.0 - i / 8, torch.randn(1, NPIX, DIM), torch.randn(1, NPIX, DIM)) + # step 8 is not a multiple of 3, but force=True records it anyway. + diag.on_step(8, 0.0, torch.randn(1, NPIX, DIM), None, force=True) + + assert diag.steps == [0, 3, 6, 8] + assert diag.phys["x0_hat"][-1] is None + assert diag.latent_psd["x0_hat"][-1] is None + assert diag.phys["x_t"][-1] is not None # the output field is still recorded + + +def test_terminal_frame_renders_with_three_map_panels(tmp_path): + """At the terminal frame the x0_hat panel is dropped, leaving x_t | decode(z) | truth.""" + lats, lons = _o96_like_grid() + n_points = lats.size + diag = _make_diagnostics(tmp_path, n_points, latent_channels=0) + decoder = _StubDecoder(n_points, len(ALL_CHANNELS)) + diag.set_batch(0) + diag.bind_decoder(decoder) + diag.begin(torch.randn(1, NPIX, DIM)) + diag.on_step(0, 1.0, torch.randn(1, NPIX, DIM), torch.randn(1, NPIX, DIM)) + diag.on_step(1, 0.0, torch.randn(1, NPIX, DIM), None, force=True) + + aux = _target_aux(n_points, len(ALL_CHANNELS)) + aux["target_coords"] = (torch.from_numpy(np.stack([lats, lons], axis=1)).float(),) + diag.render(aux) # must not raise on the None (terminal) x0_hat frame + + # Both frames produced a figure for each channel, including the terminal one. + maps = sorted(p.name for p in (tmp_path / "maps").glob("*.png")) + assert {"step000_2t.png", "step000_q_850.png", + "step001_2t.png", "step001_q_850.png"} <= set(maps) + + +def test_render_applies_idxs_inv_to_predictions(tmp_path): + """Predictions and targets must be permuted identically, as write_output does.""" + lats, lons = _o96_like_grid() + n_points = lats.size + diag = _make_diagnostics(tmp_path, n_points, latent_channels=0) + _run(diag, _StubDecoder(n_points, len(ALL_CHANNELS)), n_steps=1) + + before = diag.phys["x_t"][0].copy() + perm = torch.from_numpy(np.random.default_rng(3).permutation(n_points)) + aux = _target_aux(n_points, len(ALL_CHANNELS), idxs_inv=perm) + aux["target_coords"] = (torch.from_numpy(np.stack([lats, lons], axis=1)).float(),) + diag.render(aux) + + np.testing.assert_array_equal(diag.phys["x_t"][0], before[perm.numpy()]) + + +def test_render_is_a_noop_without_collected_steps(tmp_path): + diag = _make_diagnostics(tmp_path, 50) + diag.set_batch(0) + diag.render(_target_aux(50, len(ALL_CHANNELS))) + + assert not (tmp_path / "spectra").exists() + assert not (tmp_path / "maps").exists() + + +def test_latent_channel_subset_is_stable_across_steps(tmp_path): + """The same channels must be used at every step, or the curves are not comparable.""" + diag = _make_diagnostics(tmp_path, 50, latent_channels=2) + tokens = torch.randn(1, NPIX, DIM) + + first = diag._latent_psd(tokens)[1] + second = diag._latent_psd(tokens)[1] + + np.testing.assert_allclose(first, second) + + +def test_token_count_mismatch_is_reported(tmp_path): + diag = _make_diagnostics(tmp_path, 50) + with pytest.raises(ValueError, match="not a multiple of npix"): + diag._latent_maps(torch.randn(1, NPIX + 3, DIM)) diff --git a/tests/test_inference_spectra.py b/tests/test_inference_spectra.py new file mode 100644 index 0000000000..89fc081a16 --- /dev/null +++ b/tests/test_inference_spectra.py @@ -0,0 +1,148 @@ +# (C) Copyright 2025 WeatherGenerator contributors. +# +# This software is licensed under the terms of the Apache Licence Version 2.0 +# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. +# +# In applying this licence, ECMWF does not waive the privileges and immunities +# granted to it by virtue of its status as an intergovernmental organisation +# nor does it submit to any jurisdiction. + +"""Tests for the inference-time ODE spectral diagnostics.""" + +import numpy as np +import pytest +from astropy_healpix.healpy import pix2ang +from scipy.special import sph_harm_y + +from weathergen.evaluate.scores.psd import _legendre_gauss_weights, _octahedral_lons_per_lat +from weathergen.model.inference_spectra import ( + canonical_grid_order, + healpix_sht_psd, + physical_psd, + white_noise_reference, +) + +NSIDE = 32 +NPIX = 12 * NSIDE**2 +TRUNC = 2 * NSIDE + + +def _healpix_angles(nested: bool): + return pix2ang(nside=NSIDE, ipix=np.arange(NPIX), nest=nested) + + +def _real_ylm(ell: int, m: int, theta, phi): + y = sph_harm_y(ell, m, theta, phi) + return y.real if m == 0 else np.sqrt(2) * y.real + + +@pytest.mark.parametrize(("ell", "m"), [(3, 2), (10, 0), (20, 7)]) +def test_single_harmonic_is_isolated(ell: int, m: int) -> None: + """A pure Y_lm must put essentially all its power at that l.""" + theta, phi = _healpix_angles(nested=True) + _, psd = healpix_sht_psd(_real_ylm(ell, m, theta, phi), NSIDE, TRUNC) + + assert int(np.argmax(psd)) == ell + off_peak = (psd.sum() - psd[ell]) / psd.sum() + assert off_peak < 1e-3 + + +def test_nested_and_ring_orderings_agree() -> None: + """The nest->ring reindex must be applied; otherwise the spectrum is silently scrambled.""" + theta_n, phi_n = _healpix_angles(nested=True) + theta_r, phi_r = _healpix_angles(nested=False) + + _, psd_nested = healpix_sht_psd(_real_ylm(7, 3, theta_n, phi_n), NSIDE, TRUNC, nested=True) + _, psd_ring = healpix_sht_psd(_real_ylm(7, 3, theta_r, phi_r), NSIDE, TRUNC, nested=False) + + np.testing.assert_allclose(psd_nested, psd_ring, rtol=1e-10, atol=1e-14) + + +def test_ignoring_the_reindex_is_detectably_wrong() -> None: + """Guard the guard: reading a nested map as if it were ring order must change the answer.""" + theta, phi = _healpix_angles(nested=True) + field = _real_ylm(7, 3, theta, phi) + + _, correct = healpix_sht_psd(field, NSIDE, TRUNC, nested=True) + _, scrambled = healpix_sht_psd(field, NSIDE, TRUNC, nested=False) + + assert int(np.argmax(correct)) == 7 + assert not np.allclose(correct, scrambled, rtol=1e-3) + + +def test_white_noise_follows_the_2l_plus_1_reference() -> None: + """In this convention (sum over m, no 1/(2l+1)) white noise rises like 2l+1, not flat.""" + rng = np.random.default_rng(0) + ell, psd = healpix_sht_psd(rng.standard_normal((64, NPIX)), NSIDE, TRUNC) + + ratio = psd[1:] / white_noise_reference(ell[1:]) + # Flat ratio => the measured spectrum has the 2l+1 shape. + assert ratio.std() / ratio.mean() < 0.15 + + +def _o96_grid(): + nlat = 192 + lons_per_lat = _octahedral_lons_per_lat(nlat) + nodes, _ = _legendre_gauss_weights(nlat) + theta = np.flip(np.arccos(nodes)) + theta_pts = np.concatenate([np.full(n, t) for t, n in zip(theta, lons_per_lat, strict=True)]) + phi_pts = np.concatenate([2 * np.pi * np.arange(n) / n for n in lons_per_lat]) + return nlat, theta_pts, phi_pts + + +def test_latent_and_physical_share_one_normalisation() -> None: + """The point of reusing _legpoly: one analytic field, two grids, same PSD. + + Without this, the latent and physical panels would silently use different y-scales. + """ + modes = [(4, 1, 1.0), (11, 5, 0.6), (25, 3, 0.3)] + + def field(theta, phi): + return sum(amp * _real_ylm(ell, m, theta, phi) for ell, m, amp in modes) + + theta_hp, phi_hp = _healpix_angles(nested=True) + _, psd_hp = healpix_sht_psd(field(theta_hp, phi_hp), NSIDE, TRUNC) + + nlat, theta_pts, phi_pts = _o96_grid() + lats = 90.0 - np.degrees(theta_pts) + lons = np.degrees(phi_pts) + result = physical_psd(field(theta_pts, phi_pts), lats, lons, truncation=TRUNC) + assert result is not None + _, psd_o96 = result + + for ell, _, _ in modes: + assert psd_hp[ell] == pytest.approx(psd_o96[ell], rel=1e-3) + assert psd_hp.sum() == pytest.approx(psd_o96.sum(), rel=1e-3) + + +def test_physical_psd_is_order_independent() -> None: + """Points arrive in dataset order, so the estimator must sort them itself.""" + nlat, theta_pts, phi_pts = _o96_grid() + lats = 90.0 - np.degrees(theta_pts) + lons = np.degrees(phi_pts) + values = _real_ylm(6, 2, theta_pts, phi_pts) + + rng = np.random.default_rng(1) + shuffle = rng.permutation(values.size) + + _, psd = physical_psd(values, lats, lons, truncation=TRUNC) + _, psd_shuffled = physical_psd(values[shuffle], lats[shuffle], lons[shuffle], truncation=TRUNC) + + np.testing.assert_allclose(psd, psd_shuffled, rtol=1e-10, atol=1e-14) + + +def test_canonical_order_is_north_to_south_then_east() -> None: + lats = np.array([-10.0, 45.0, 45.0, 80.0]) + lons = np.array([0.0, 200.0, 10.0, 5.0]) + np.testing.assert_array_equal(canonical_grid_order(lats, lons), [3, 2, 1, 0]) + + +def test_subsampled_grid_is_refused_not_guessed() -> None: + """With max_num_targets still active the point cloud is not a grid; must return None.""" + _, theta_pts, phi_pts = _o96_grid() + rng = np.random.default_rng(2) + keep = rng.choice(theta_pts.size, size=20000, replace=False) + lats = 90.0 - np.degrees(theta_pts[keep]) + lons = np.degrees(phi_pts[keep]) + + assert physical_psd(_real_ylm(6, 2, theta_pts[keep], phi_pts[keep]), lats, lons) is None diff --git a/tests/test_sht_roundtrip.py b/tests/test_sht_roundtrip.py new file mode 100644 index 0000000000..3fae345dcd --- /dev/null +++ b/tests/test_sht_roundtrip.py @@ -0,0 +1,95 @@ +"""Test that SHT forward → inverse is (approximately) the identity.""" + +import numpy as np +import pytest + +from weathergen.evaluate.scores.psd import ( + InverseSphericalHarmonicTransform, + SphericalHarmonicTransform, + _octahedral_lons_per_lat, + _regular_lons_per_lat, +) + + +@pytest.mark.parametrize("grid_type,nlat", [ + ("regular", 32), + ("regular", 64), + ("octahedral", 32), + ("octahedral", 64), +]) +def test_sht_roundtrip_identity(grid_type: str, nlat: int) -> None: + """Applying SHT then inverse SHT on random noise recovers the original field.""" + rng = np.random.default_rng(42) + + if grid_type == "regular": + lons_per_lat = _regular_lons_per_lat(nlat) + else: + lons_per_lat = _octahedral_lons_per_lat(nlat) + + n_grid_points = sum(lons_per_lat) + truncation = nlat // 2 - 1 + + sht = SphericalHarmonicTransform(lons_per_lat=lons_per_lat, truncation=truncation) + isht = InverseSphericalHarmonicTransform(lons_per_lat=lons_per_lat, truncation=truncation) + + # Random spatial field + x = rng.standard_normal(n_grid_points) + + # Forward → inverse + coeffs = sht.transform(x) + x_reconstructed = isht.transform(coeffs) + + # The reconstruction is approximate due to truncation, but should be close + # for smooth-enough fields. For a bandlimited signal it should be exact. + # Use a generous tolerance since truncation discards high-frequency content. + assert x_reconstructed.shape == x.shape, ( + f"Shape mismatch: {x_reconstructed.shape} vs {x.shape}" + ) + + # Check correlation is positive — truncation discards high-frequency content + # so white noise won't be perfectly recovered, but the low-frequency part should match. + corr = np.corrcoef(x.ravel(), x_reconstructed.ravel())[0, 1] + assert corr > 0.30, f"Correlation too low: {corr:.4f}" + + # More importantly: verify the energy is preserved for the retained modes + # by checking that the relative L2 error is bounded + rel_error = np.linalg.norm(x - x_reconstructed) / np.linalg.norm(x) + assert rel_error < 1.0, f"Relative L2 error too large: {rel_error:.4f}" + + +@pytest.mark.parametrize("grid_type,nlat", [ + ("regular", 32), + ("regular", 64), + ("octahedral", 32), + ("octahedral", 64), +]) +def test_sht_roundtrip_bandlimited(grid_type: str, nlat: int) -> None: + """For a bandlimited signal, SHT → inverse SHT should be near-exact.""" + if grid_type == "regular": + lons_per_lat = _regular_lons_per_lat(nlat) + else: + lons_per_lat = _octahedral_lons_per_lat(nlat) + + n_grid_points = sum(lons_per_lat) + truncation = nlat // 2 - 1 + + sht = SphericalHarmonicTransform(lons_per_lat=lons_per_lat, truncation=truncation) + isht = InverseSphericalHarmonicTransform(lons_per_lat=lons_per_lat, truncation=truncation) + + # Create a bandlimited signal by doing inverse SHT on random coefficients + rng = np.random.default_rng(123) + L = truncation + 1 + random_coeffs = rng.standard_normal((L, L)) + 1j * rng.standard_normal((L, L)) + # Make it physically meaningful: zero out upper triangle (m > l) + for l in range(L): + random_coeffs[l, l + 1:] = 0.0 + + # Inverse → forward → inverse should give back the same spatial field + x_bandlimited = isht.transform(random_coeffs) + coeffs_recovered = sht.transform(x_bandlimited) + x_roundtrip = isht.transform(coeffs_recovered) + + np.testing.assert_allclose( + x_roundtrip, x_bandlimited, rtol=1e-6, atol=1e-10, + err_msg="Roundtrip on bandlimited signal should be near-exact", + )