From b3d73c09cf3c8656182f924221851f955541fbae Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Fri, 14 Aug 2026 09:52:57 +0200 Subject: [PATCH 01/15] layer fix for backward compatibility --- src/weathergen/model/layers.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/weathergen/model/layers.py b/src/weathergen/model/layers.py index 3605269659..f5942a3f21 100644 --- a/src/weathergen/model/layers.py +++ b/src/weathergen/model/layers.py @@ -103,13 +103,16 @@ def __init__( self.noise_conditioning = LinearNormConditioning(dim_in) self.noise_conditioning = LinearNormConditioning(dim_in) elif dim_aux is not None: - self.lnorm = AdaLayerNorm(dim_in, dim_aux, norm_eps=norm_eps) + # Not aliased to self.lnorm: forward() picks this up via the self.layers loop + # below (isinstance(layer, AdaLayerNorm)), and a second attribute name here would + # duplicate this module's params under both "layers.0.*" and "lnorm.*" in the + # state_dict, breaking strict loading of checkpoints saved before this norm lived + # in self.layers only. + self.layers.append(AdaLayerNorm(dim_in, dim_aux, norm_eps=norm_eps)) else: - self.lnorm = norm(dim_in, eps=norm_eps) - - # TODO: The below should be consolidated – implementing in layer list for backward compatibility - if not is_dit: - self.layers.append(self.lnorm) + # See comment above: kept out of self.lnorm for the same reason, forward() calls + # it via the self.layers loop. + self.layers.append(norm(dim_in, eps=norm_eps)) if self.mlp_type == "swiglu": self.layers.append(torch.nn.Linear(dim_in, 2 * dim_hidden)) From cbe94b8fd7119aae9ad385bd362dd272778bed34 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Thu, 30 Jul 2026 18:32:26 +0200 Subject: [PATCH 02/15] Latent-RMSE + inference diagnostics onto develop-ssl-diffusion-v1 Squash of the mh/...-latent-rmse-and-diagnostics-ssl feature commits (1785ef3d..4356ee61) resolved onto develop-ssl-diffusion-v1. Adds: - Latent-rollout RMSE diagnostic: record_latent_rollout side channel + isolated batch.latent_rmse_source truth encode (utils/latent_rmse.py). - Per-ODE-step map & spectrum diagnostics (inference_diagnostics.py, inference_spectra.py, evaluate/scores/psd.py; diag_* config keys). - Harmonized diffusion sampling diagnostics (per-element rmse panel, 4-panel). Conflict resolutions vs develop-ssl: encoder 3-value return threaded through the chunked truth-latent encode; develop's LATENT_CONDITIONING_TOKENS key kept; _get_data_windows 3-return reconciled with teacher_time_offset windows; fe_diffusion_num_inference_steps config key; MapPanelPlotter adapted to develop's eval Plotter API (norm built locally). 28 diagnostics/spectra tests pass. Co-Authored-By: Claude Opus 4.8 --- ...config_diffusion_d2048_forecast_concat.yml | 10 + src/weathergen/datasets/batch.py | 23 + .../datasets/multi_stream_data_sampler.py | 101 +++- src/weathergen/model/diffusion.py | 143 +++-- src/weathergen/model/inference_diagnostics.py | 558 ++++++++++++++++++ src/weathergen/model/inference_spectra.py | 188 ++++++ src/weathergen/model/model.py | 84 +++ src/weathergen/train/trainer.py | 93 ++- src/weathergen/utils/latent_rmse.py | 152 +++++ tests/test_inference_diagnostics.py | 227 +++++++ tests/test_inference_spectra.py | 148 +++++ tests/test_sht_roundtrip.py | 95 +++ 12 files changed, 1751 insertions(+), 71 deletions(-) create mode 100644 src/weathergen/model/inference_diagnostics.py create mode 100644 src/weathergen/model/inference_spectra.py create mode 100644 src/weathergen/utils/latent_rmse.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/config/config_diffusion_d2048_forecast_concat.yml b/config/config_diffusion_d2048_forecast_concat.yml index b402b583db..c0a5526c32 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 diff --git a/src/weathergen/datasets/batch.py b/src/weathergen/datasets/batch.py index fdfddfd70b..564bd4d38b 100644 --- a/src/weathergen/datasets/batch.py +++ b/src/weathergen/datasets/batch.py @@ -355,6 +355,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""" @@ -364,6 +381,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 @@ -374,6 +394,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 13b7da1684..3e4f7fec90 100644 --- a/src/weathergen/datasets/multi_stream_data_sampler.py +++ b/src/weathergen/datasets/multi_stream_data_sampler.py @@ -137,6 +137,18 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage): ) self.teacher_time_offset = 0 + # 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) + if self.latent_rollout_rmse: + assert self.output_offset == 0, ( + f"latent_rollout_rmse requires forecast.offset == 0, got {self.output_offset}" + ) + + # initialise fsm, but can change for future mini_epochs self.batch_size = get_batch_size_from_config(mode_cfg) self.num_workers = cf.data_loading.num_workers self.shuffle = mode_cfg.shuffle @@ -462,6 +474,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, @@ -631,7 +680,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): """ @@ -704,6 +773,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_name, stream_ds in self.streams_datasets.items(): @@ -720,14 +792,14 @@ 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.readers ) # When teacher_time_offset > 0, load a separate set of data windows # shifted forward in time for the teacher (target) samples. if self.teacher_time_offset > 0: - (input_data_target, output_data_target) = self._get_data_windows( + (input_data_target, output_data_target, _) = self._get_data_windows( idx + self.teacher_time_offset, num_forecast_steps, i_max, stream_ds.readers ) else: @@ -749,6 +821,12 @@ def _get_batch(self, idx: int, num_forecast_steps: int): input_tokens_target = input_tokens output_tokens_target = output_tokens + 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 tidx = source_to_target[sidx].item() @@ -828,10 +906,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"]: diff --git a/src/weathergen/model/diffusion.py b/src/weathergen/model/diffusion.py index dbae5b9130..7812b6ad02 100644 --- a/src/weathergen/model/diffusion.py +++ b/src/weathergen/model/diffusion.py @@ -171,6 +171,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 @@ -185,7 +188,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. @@ -204,7 +207,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_inference_steps``, which + itself defaults to 10 — preserving the historical hardcoded value bit-identically. + Set ``fe_diffusion_num_inference_steps`` in the config (or via ``--options``) to override. Returns: torch.Tensor: Model output (denoised prediction during training, @@ -243,9 +249,9 @@ 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 - # Allow the number of ODE denoising steps to be set from the config. - # Falls back to the `num_steps` argument default when not configured. - num_steps = self.cf.get("fe_diffusion_num_inference_steps", None) or num_steps + # Number of ODE denoising steps from config (develop's key); fall back to the + # historical default of 10 when neither the config nor the arg provides one. + num_steps = self.cf.get("fe_diffusion_num_inference_steps", None) or num_steps or 10 return self.inference_forward( fstep=fstep, num_steps=num_steps, @@ -583,7 +589,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": [], @@ -596,6 +603,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( @@ -615,6 +633,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 @@ -636,12 +657,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) @@ -652,77 +686,64 @@ def _plot_sampling_diagnostics(self, track: dict, num_steps: int) -> None: import matplotlib matplotlib.use("Agg") + import os 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..04249fdcf1 --- /dev/null +++ b/src/weathergen/model/inference_diagnostics.py @@ -0,0 +1,558 @@ +# (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): + import matplotlib as mpl + + opts = self._parse_map_kwargs(dict(map_kwargs or {}), self.stream) + # The eval Plotter builds the norm inside scatter_plot (not in _parse_map_kwargs) and + # no longer returns an "extra"/"norm" key, so replicate that here for _render_scatter. + if opts.get("levels") is not None: + norm = mpl.colors.BoundaryNorm(opts["levels"], opts["cmap"].N, extend="both") + else: + norm = mpl.colors.Normalize( + vmin=opts.get("vmin"), vmax=opts.get("vmax"), clip=False + ) + extra = opts.get("extra", {}) + 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, norm, opts["cmap"], marker_size, opts["marker"], 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 e90576b165..d73b304737 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 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) + 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 @@ -381,6 +410,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 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 @@ -862,6 +895,17 @@ def forward( continue if self.forecast_engine: + # 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, @@ -947,6 +991,15 @@ def forward( output.add_physical_prediction( step, sname, (torch.cat(list(pred_tuple), dim=0),) ) + # 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. @@ -973,6 +1026,20 @@ def forward( 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: """ @@ -1009,6 +1076,23 @@ def _get_initial_conditions( return source_samples, tokens, posteriors, intermediates + 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 c8b0e6c76a..56906d86e4 100644 --- a/src/weathergen/train/trainer.py +++ b/src/weathergen/train/trainer.py @@ -27,6 +27,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 import ModelOutput from weathergen.model.model_interface import ( @@ -52,6 +53,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 @@ -73,6 +75,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`` @@ -821,9 +833,33 @@ 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 + + # 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 + # 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) + 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.record_latent_rollout = latent_rmse is not None and _noise_idx == 0 if noise_level is None: loss_suffix = "" @@ -836,6 +872,7 @@ def validate(self, mini_epoch, mode_cfg, batch_size): stage_suffix = f"_eta{eta_str}" if len(noise_levels) > 1 else "" dataset_val_iter = iter(self.data_loader_validation) + num_samples_write = mode_cfg.get("output", {}).get("num_samples", 0) * batch_size with torch.no_grad(): # print progress bar but only in interactive mode, i.e. when without ddp @@ -852,6 +889,9 @@ def validate(self, mini_epoch, mode_cfg, batch_size): chunk_size = mode_cfg.forecast.get("chunk_size", len(output_idxs)) chunks = self._get_forecast_step_chunks(output_idxs, chunk_size) batch.to_device_for_chunked_inference(self.device, chunks[-1]) + if ode_diag is not None: + # Self-disables for every batch after the first. + ode_diag.set_batch(bidx) # evaluate model with torch.autocast( @@ -889,14 +929,48 @@ def validate(self, mini_epoch, mode_cfg, batch_size): if is_diffusion: _expand_targets_to_match_preds(preds, targets_and_auxs) - # Write output for diffusion — _process_validation_chunks - # skips writing when is_diffusion=True because chunked IO - # is incompatible with the ODE trajectory fstep layout. - # Do it here instead, mirroring the non-diffusion path. - num_samples_write = ( - mode_cfg.get("output", {}).get("num_samples", 0) * batch_size + # 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] ) - if is_diffusion and _noise_idx == 0 and bidx < num_samples_write: + + _ = self.loss_calculator_val.compute_loss( + preds=preds, + targets_and_aux=targets_and_auxs, + 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 is_diffusion and _noise_idx == 0 and bidx < num_samples_write: denormalize_data_fct = ( (lambda x0, x1: x1) if mode_cfg.get("output", {}).get("normalized_samples", False) @@ -951,6 +1025,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.record_latent_rollout = False + 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..57d2e799cd --- /dev/null +++ b/src/weathergen/utils/latent_rmse.py @@ -0,0 +1,152 @@ +# (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 json +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" + + 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=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}") 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", + ) From f5cdc6448589e625caa86a02e1ebcbb49349f629 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Sat, 1 Aug 2026 18:18:33 +0200 Subject: [PATCH 03/15] fixes to latent rmse --- .../config_diffusion_d2048_forecast_n320.yml | 20 +- ...iffusion_d2048_forecast_n320_inference.yml | 367 ++++++++++++++++++ .../datasets/multi_stream_data_sampler.py | 4 +- src/weathergen/model/inference_diagnostics.py | 7 +- src/weathergen/model/inference_spectra.py | 43 +- src/weathergen/model/model.py | 1 + src/weathergen/utils/latent_rmse.py | 2 +- 7 files changed, 435 insertions(+), 9 deletions(-) create mode 100644 config/config_diffusion_d2048_forecast_n320_inference.yml diff --git a/config/config_diffusion_d2048_forecast_n320.yml b/config/config_diffusion_d2048_forecast_n320.yml index 648eba47f1..62132d7d91 100644 --- a/config/config_diffusion_d2048_forecast_n320.yml +++ b/config/config_diffusion_d2048_forecast_n320.yml @@ -73,13 +73,23 @@ with_step_conditioning: True # False frequency_embedding_dim: 256 embedding_dim: 512 noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" -sigma_min: &sigma_min 0.02 -sigma_max: &sigma_max 10 +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 sigma_data: 1.0 rho: 7 p_mean: 3.5 # only required if noise_distribution is log_normal p_std: 1.2 # only required if noise_distribution is log_normal +# --- 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) + healpix_level: 5 # Use 2D RoPE instead of traditional global positional encoding @@ -94,7 +104,8 @@ use_xsa: False with_mixed_precision: True with_flash_attention: True compile_model: False -with_fsdp: False +with_fsdp: True +# ddp_static_graph: True attention_dtype: bf16 mixed_precision_dtype: bf16 mlp_norm_eps: 1e-5 @@ -153,7 +164,8 @@ freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_to # load_chkpt: {'run_id': 'wd0u4he8', 'epoch': -1} # d2048 hl5, p_mean=3.5, forecasting # load_chkpt: {'run_id': 'iy534ge3', 'epoch': -1} # d2048 hl5 pre-training with z, 64 epochs # load_chkpt: {'run_id': 'btnf4goa', 'epoch': -1} # d2048 hl5 pre-training with z, 128 epochs -load_chkpt: {'run_id': 'c7d5iw3q', 'epoch': -1} # n320, 82 epochs +# load_chkpt: {'run_id': 'c7d5iw3q', 'epoch': -1} # n320, 82 epochs +load_chkpt: {'run_id': 'y9futkc8', 'epoch': -1} # n320 6h, 82 epochs norm_type: "LayerNorm" diff --git a/config/config_diffusion_d2048_forecast_n320_inference.yml b/config/config_diffusion_d2048_forecast_n320_inference.yml new file mode 100644 index 0000000000..79553055f0 --- /dev/null +++ b/config/config_diffusion_d2048_forecast_n320_inference.yml @@ -0,0 +1,367 @@ +# (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. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 2048 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: true +ae_local_num_queries: 1 +ae_local_queries_per_cell: false + +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: true +ae_adapter_with_residual: true +ae_adapter_dropout_rate: 0.1 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 4 +ae_global_num_heads: 32 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: true +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 2 +ae_global_trailing_layer_norm: false + +ae_aggregation_num_blocks: 0 +ae_aggregation_num_heads: 32 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: true +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 2 + +decoder_type: "PerceiverIOCoordConditioning" +pred_adapter_kv: false +pred_self_attention: true +pred_dyadic_dims: false +pred_mlp_adaln: true +num_class_tokens: 0 +num_register_tokens: 0 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2048 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +fe_impute_latent_noise_std: 0.0 # 1e-4 +fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False +# Diffusion related parameters +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 # only required if noise_distribution is log_normal +p_std: 1.2 # only required if noise_distribution is log_normal + +# --- 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) + +healpix_level: 5 + +# 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 +rope_2D: False +# mlp_type: swiglu +# use_xsa: True +mlp_type: mlp +use_xsa: False + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +# ddp_static_graph: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-5 +norm_eps: 1e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + + +#below for FE training only +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for DECODER training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for FE and DECODER training +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +# load_chkpt: {'run_id': 't0bdz7qn', 'epoch': -1} # multi-var d2048 hl5, sigma_data=1.7 +# load_chkpt: {'run_id': 'dcl584vo', 'epoch': -1} # z500 d2048 hl5, sigma_data=159.08 +# load_chkpt: {'run_id': 'tvkicam9', 'epoch': -1} # z500 d2048 hl3 enc-lnorm, sigma_data=1.0 +# load_chkpt: {'run_id': 'q9grso75', 'epoch': -1} # z500 d2048 hl3, sigma_data=39.2936 +# load_chkpt: {'run_id': 'qxivdyqz', 'epoch': -1} # z500 d2048 hl5 enc-lnorm, sigma_data=1.0 +# load_chkpt: {'run_id': 'h8x1qgz3', 'epoch': -1} # z500 d128 hl5, sigma_data=12.93 +# load_chkpt: {'run_id': '', 'epoch': -1} # z500 d128 hl5 enc-lnorm, sigma_data=1.0 +# load_chkpt: {'run_id': 'wvpb76ai', 'epoch': -1} # multi-var d2048 hl3 enc-lnorm, sigma_data=1.0 +# load_chkpt: {'run_id': 'ae4wlc5m', 'epoch': -1} # multi-var d2048 hl3, sigma_data=2.7047 +# load_chkpt: {'run_id': 'r45iwyns', 'epoch': -1} # multi-var d512 hl3, sigma_data=1.1785 +# load_chkpt: {'run_id': 'ydka6uql', 'epoch': -1} # multi-var d512 hl4, sigma_data=0.827 +# load_chkpt: {'run_id': 'lwjkb3y4', 'epoch': -1} # multi-var d512 hl5, sigma_data=0.5789 +# load_chkpt: {'run_id': 'v8kd6xc1', 'epoch': -1} # multi-var d512 hl5 nopos, sigma_data=0.6481 +# load_chkpt: {'run_id': 'lwjkb3y4', 'epoch': -1} # multi-var d512 hl5 enc-lnorm, sigma_data=1.0 +# load_chkpt: {'run_id': 'y1gu5md8', 'epoch': -1} # multi-var d512 hl5, sigma_dqta=1.0, diffusion-full-pipeline +# load_chkpt: {'run_id': 'mal6u4gc', 'epoch': -1} # multi-var d512 hl5, sigma_dqta=1.0, geoinfos 64 epochs, diffusion-full-pipeline +# load_chkpt: {'run_id': 'zrpncqb0', 'epoch': -1} # multi-var d512 hl5, sigma_dqta=1.0, geoinfos 196 epochs, diffusion-full-pipeline +# load_chkpt: {'run_id': 'm6fs8wvj', 'epoch': -1} # multi-var d512 hl5, sigma_data=1.0, swiglu xsa geoinfos, diffusion-full-pipeline +# load_chkpt: {'run_id': 'cgxt9imf', 'epoch': -1} # diffusion model d512 to fine-tune decoder, p_mean=0.5, SwiGLU+XSA+geoinfos, based on m6fs8wvj backbone +# load_chkpt: {'run_id': 'wo5mf2z4', 'epoch': -1} # diffusion model d512 to fine-tune decoder, p_mean=1.5, SwiGLU+XSA+geoinfos, based on m6fs8wvj backbone +# load_chkpt: {'run_id': 'zf6wnmpe', 'epoch': -1} # multi-var d2048 hl5, sigma_data=1.832 +# load_chkpt: {'run_id': 'mivw6jda', 'epoch': -1} # multi-var d2048 hl5 enc-lnorm, sigma_data=1.0 +# load_chkpt: {'run_id': 'l3rxe29i', 'epoch': -1} # multi-var d2048 hl5, sigma_data=1.0, swiglu xsa geoinfos, diffusion-full-pipeline +# load_chkpt: {'run_id': 'riyz96d4', 'epoch': -1} # diffusion model d2048 to fine-tune decoder, p_mean=0.5, based on l3rxe29i backbone +# load_chkpt: {'run_id': 'bokn5d2w', 'epoch': -1} # diffusion model d2048 to fine-tune decoder, p_mean=1.5, based on l3rxe29i backbone +# load_chkpt: {'run_id': 'uwyv1zdh', 'epoch': -1} # diffusion model d2048 to fine-tune decoder, p_mean=2.0, based on l3rxe29i backbone +# load_chkpt: {'run_id': 'dvslhdp3', 'epoch': -1} # diffusion model d2048 to fine-tune decoder, p_mean=2.5, based on l3rxe29i backbone +# load_chkpt: {'run_id': 'mtlfdgvh', 'epoch': -1} # diffusion model d2048 to fine-tune decoder, p_mean=3.0, based on l3rxe29i backbone +# load_chkpt: {'run_id': 'jnxvjtur', 'epoch': -1} # d2048 hl5, p_mean=-2.5, forecasting +# load_chkpt: {'run_id': 'ulw4wid5', 'epoch': -1} # d2048 hl5, p_mean=-2.0, forecasting +# load_chkpt: {'run_id': 'lxpbthna', 'epoch': -1} # d2048 hl5, p_mean=-1.5, forecasting +# load_chkpt: {'run_id': 'n1e97r8z', 'epoch': -1} # d2048 hl5, p_mean=-1.2, forecasting +# load_chkpt: {'run_id': 'halvwj21', 'epoch': -1} # d2048 hl5, p_mean=-0.5, forecasting +# load_chkpt: {'run_id': 'dlsp3k31', 'epoch': -1} # d2048 hl5, p_mean=0.5, forecasting +# load_chkpt: {'run_id': 'ug7huxi2', 'epoch': -1} # d2048 hl5, p_mean=1.5, forecasting +# load_chkpt: {'run_id': 'i3y5fhda', 'epoch': -1} # d2048 hl5, p_mean=2.5, forecasting +# load_chkpt: {'run_id': 'wd0u4he8', 'epoch': -1} # d2048 hl5, p_mean=3.5, forecasting +# load_chkpt: {'run_id': 'iy534ge3', 'epoch': -1} # d2048 hl5 pre-training with z, 64 epochs +# load_chkpt: {'run_id': 'btnf4goa', 'epoch': -1} # d2048 hl5 pre-training with z, 128 epochs +# load_chkpt: {'run_id': 'c7d5iw3q', 'epoch': -1} # n320, 82 epochs +load_chkpt: {'run_id': 'f9qbdiwi', 'epoch': -1} # n320 6h, 82 epochs + + +norm_type: "LayerNorm" + +##################################### + +# streams_directory: "./config/streams/era5_1deg_forecasting_d2048/" +# streams_directory: "./config/streams/era5_1deg_forecasting_z500/" +streams_directory: "./config/streams/era5_n320_forecasting/" +streams: ??? + +# type of zarr_store +zarr_store: "zip" # "zarr" for LocalStore, "zip" for ZipStore + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + desc: "" + run_id: ??? + run_history: [] + +# logging frequency in the training loop (in number of batches) +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 250 + log_grad_norms: False + +# parameters for data loading +data_loading : + + num_workers: 12 + rng_seed: ??? + repeat_data_in_mini_epoch : False + + +# config for training +training_config: + + # training_mode: "masking", "student_teacher", "latent_loss" + training_mode: ["masking","student_teacher"] # ["student_teacher", "physical_loss"] + + num_mini_epochs: 128 + samples_per_mini_epoch: 4096 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T18:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 #5e-5 + lr_max: 1e-5 #1e-4 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 512 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.975 + beta2 : 0.9875 + eps : 2e-08 + + losses : { + "physical": { + type: LossPhysical, + weight: 0.0, + loss_fcts: { + "mse": {}, + }, + target_and_aux_calc: "Physical", + }, + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + # masking strategy: "random", "healpix", "forecast" + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: {diffusion_rn: True}, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + + +# validation config; full validation config is merge of training and validation config +validation_config: + + # Noise levels (eta values in standard normal space) at which to evaluate the + # diffusion model during validation. sigma = exp(eta * p_std + p_mean). + # Each value produces a separate validation pass with independently logged metrics. + validation_noise_levels: [1.0, 2.0, 3.0, 4.0] + + samples_per_mini_epoch: 256 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: False + + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py index 3e4f7fec90..75aced2c7b 100644 --- a/src/weathergen/datasets/multi_stream_data_sampler.py +++ b/src/weathergen/datasets/multi_stream_data_sampler.py @@ -507,7 +507,9 @@ def _build_latent_rmse_stream_data( (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) + stream_data.add_source( + self._stage, step, rdata, source_cells_lens, source_cells, rdata.is_spoof + ) return stream_data diff --git a/src/weathergen/model/inference_diagnostics.py b/src/weathergen/model/inference_diagnostics.py index 04249fdcf1..084d263940 100644 --- a/src/weathergen/model/inference_diagnostics.py +++ b/src/weathergen/model/inference_diagnostics.py @@ -449,7 +449,12 @@ def maybe_create(cf, model, denormalize: Callable[[str, torch.Tensor], torch.Ten from weathergen.common.config import get_path_run stream_name = cf.get("diag_stream", "ERA5") - streams = {s["name"]: s for s in cf.streams} + # cf.streams is a dict keyed by stream name on this branch (develop-ssl); older code passed a + # list of stream dicts. Support both so the diagnostics work regardless of config shape. + if hasattr(cf.streams, "items"): + streams = {name: info for name, info in cf.streams.items()} + else: + 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 diff --git a/src/weathergen/model/inference_spectra.py b/src/weathergen/model/inference_spectra.py index 0823a9917f..13134f3f66 100644 --- a/src/weathergen/model/inference_spectra.py +++ b/src/weathergen/model/inference_spectra.py @@ -143,6 +143,27 @@ def canonical_grid_order(lats: np.typing.NDArray, lons: np.typing.NDArray) -> np return np.lexsort((np.asarray(lons) % 360.0, -np.asarray(lats))) +def _detect_reduced_gaussian(lats: np.typing.NDArray, n_points: int) -> str | None: + """Return ``"reduced"`` if the points are the ECMWF reduced Gaussian grid ``sht_psd`` supports. + + ``detect_grid_type`` (evaluate pkg) only knows octahedral and regular grids; ERA5's native grid + is a classic reduced Gaussian (N320 = 542080 points / 640 latitudes) that matches neither. + ``sht_psd``'s ``grid_type="reduced"`` path is N320-specific, so gate on nlat == 640 and confirm + the exact point count via anemoi's grid table (already the dependency that path relies on). + """ + nlat = len(np.unique(lats)) + if nlat != 640: # sht_psd's reduced path is hard-wired to N320 + return None + try: + from anemoi.transform.grids.named import lookup + + expected = len(np.asarray(lookup("N320")["latitudes"])) + except Exception as exc: # noqa: BLE001 - anemoi missing / offline / unknown grid + logger.warning(f"Reduced Gaussian (N320) PSD unavailable: {exc}") + return None + return "reduced" if n_points == expected else None + + def physical_psd( values: np.typing.NDArray, lats: np.typing.NDArray, @@ -159,9 +180,27 @@ def physical_psd( lats = np.asarray(lats) lons = np.asarray(lons) - grid_type = detect_grid_type(lats, lons, values.shape[-1]) + # detect_grid_type warns "PSD via SHT skipped" for any grid it doesn't know -- including ERA5's + # reduced Gaussian, which we DO handle below -- so silence that one call to avoid a misleading + # log; messaging is re-emitted here only if both detectors fail. + psd_logger = logging.getLogger("weathergen.evaluate.scores.psd") + _prev_level = psd_logger.level + psd_logger.setLevel(logging.ERROR) + try: + grid_type = detect_grid_type(lats, lons, values.shape[-1]) + finally: + psd_logger.setLevel(_prev_level) if grid_type is None: - return None + # ERA5's native reduced Gaussian N320 is recognised by neither octahedral nor regular + # detection; sht_psd handles it via grid_type="reduced". Recognise it here instead of + # silently skipping the physical spectrum. + grid_type = _detect_reduced_gaussian(lats, values.shape[-1]) + if grid_type is None: + logger.warning( + f"Physical PSD skipped: {values.shape[-1]} points / {len(np.unique(lats))} " + "latitudes is not a recognised global grid (octahedral, regular, or N320)." + ) + return None nlat = len(np.unique(lats)) expected = sum(_octahedral_lons_per_lat(nlat)) if grid_type == "octahedral" else None diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index d73b304737..123c5914bd 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -9,6 +9,7 @@ # granted to it by virtue of its status as an intergovernmental organisation # nor does it submit to any jurisdiction. +import copy import logging import math import typing diff --git a/src/weathergen/utils/latent_rmse.py b/src/weathergen/utils/latent_rmse.py index 57d2e799cd..0e7b5eb773 100644 --- a/src/weathergen/utils/latent_rmse.py +++ b/src/weathergen/utils/latent_rmse.py @@ -135,7 +135,7 @@ def plot(self, output_dir) -> None: # 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) + self._write_json(plotter.out_plot_dir_lines / 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.""" From ae02d42aab786261b964b9f2667b7ffd54ebde86 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Fri, 7 Aug 2026 17:28:57 +0200 Subject: [PATCH 04/15] add claude skill /weathergen-inference-diagnostics --- .../weathergen-inference-diagnostics/SKILL.md | 116 ++++++++++++ .../references/decoder-overlay.md | 132 ++++++++++++++ .../scripts/make_decoder_overlay.py | 172 ++++++++++++++++++ 3 files changed, 420 insertions(+) create mode 100644 .claude/skills/weathergen-inference-diagnostics/SKILL.md create mode 100644 .claude/skills/weathergen-inference-diagnostics/references/decoder-overlay.md create mode 100755 .claude/skills/weathergen-inference-diagnostics/scripts/make_decoder_overlay.py diff --git a/.claude/skills/weathergen-inference-diagnostics/SKILL.md b/.claude/skills/weathergen-inference-diagnostics/SKILL.md new file mode 100644 index 0000000000..833907cbfe --- /dev/null +++ b/.claude/skills/weathergen-inference-diagnostics/SKILL.md @@ -0,0 +1,116 @@ +--- +name: weathergen-inference-diagnostics +description: Run sampling/ODE diagnostics (per-step maps and spectra along the denoising trajectory) or the latent-RMSE-vs-lead-time curve on an already-trained WeatherGenerator run, and attach a backbone decoder to a run that has none. Use when asked to diagnose a diffusion or flow-matching sampler, to plot latent RMSE during rollout, or when inference fails with "assert len(outputs_physical) == 1", "Empty preds but non-empty targets", or a request for physical output from a model trained with a latent-only loss. +--- + +# Inference diagnostics on a trained run + +Both diagnostics run inside an ordinary `uv run inference` job on an existing run — never +retraining, never a separate `evaluate` step. Every flag below is passed via `--options`. + +Both features must exist in the *checked-out working copy*. Job launchers that submit a snapshot of +the run's original training code will not have them — run these interactively. + +## Which one is being asked for + +| Ask | Section | Needs a physical decoder? | +|---|---|---| +| latent RMSE vs lead time, rollout error curve | A | no | +| maps / spectra / `x0_hat` / trajectory inspection | B | **yes** | + +Start from A when the run has no decoder — it is the diagnostic that works unconditionally. + +## A. Latent-RMSE rollout curve + +RMSE between the rolled-out latent and the encoded truth latent, per lead step. Pure latent space, +so it works on latent-only runs as-is. + +```bash +uv run inference --from-run-id --options \ + test_config.start_date= test_config.end_date= \ + test_config.samples_per_mini_epoch=1 test_config.output.num_samples=0 \ + test_config.latent_rollout_rmse=True \ + training_config.forecast.num_steps=16 diffusion_rollout=True \ + fe_diffusion_num_ensemble_members=1 fe_diffusion_num_steps=10 \ + 'validation_config.validation_noise_levels=[]' \ + data_loading.num_workers=0 +``` + +For the flow-matching engine replace the two `fe_diffusion_*` flags with +`fe_flow_num_ensemble_members= fe_flow_sampler=ode fm_num_steps=10 fm_sde_sigma=0.0`. + +- Take ``/`` from the run's own `validation_config` in `models//model_.json`. +- `output.num_samples=0` skips the zarr write: it is not needed for the curve, it costs hundreds of + MB per sample, and it is the code path that requires a decoder. +- Asserts `forecast.offset == 0`. +- Output: `results//line_plots/compare_rmse_global__latent.png` + plus a JSON sidecar of the values. +- 1 sample / 1 member is the smoke test. Scale `samples_per_mini_epoch` for a real average; several + ensemble members can occupy an entire large GPU, and fragmentation OOM is fixed with + `export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` (without it the process may hang at a + `(Pdb++)` post-mortem prompt instead of exiting). +- Interpret against the climatology floor of **that** latent space (√2·σ_anom of its own encoder). + A floor measured on a different encoder does not transfer. + +## B. Sampling (ODE) diagnostics + +Decodes intermediate sampler states and renders, per ODE step, `x_t`, `x0_hat = D_t(x_t)`, +`decode(z)` and the ground truth, as maps and power spectra. `decode(z)` vs truth is the control +that separates decoder error from sampler error. + +```bash +uv run inference --from-run-id --options \ + test_config.start_date= test_config.end_date= \ + test_config.samples_per_mini_epoch=1 test_config.output.num_samples=1 \ + training_config.forecast.num_steps=1 'validation_config.validation_noise_levels=[]' \ + diffusion_rollout=True fe_diffusion_num_ensemble_members=1 fe_diffusion_num_steps=10 \ + diag_ode_maps=True diag_stream= 'diag_channels=["", ""]' \ + diag_latent_channels=128 data_loading.num_workers=0 +``` + +| flag | meaning | default | +|---|---|---| +| `diag_ode_maps` | master switch | `False` | +| `diag_stream` | stream to decode | `ERA5` | +| `diag_channels` | physical channels to plot | `["2t", "q_850"]` | +| `diag_ode_every_n_steps` | record every n-th ODE step (two decoder passes each) | `1` | +| `diag_latent_channels` | latent channels in the latent spectra | `128` | + +`diag_stream` must be a stream of the run's config, and each `diag_channels` entry must appear in +that stream's `val_target_channels`: + +```bash +python3 -c " +import json +d = json.load(open('models//model_.json')) +print(list(d['streams'])) +print(d['streams']['']['val_target_channels']) +" +``` + +Output: `results//plots/ode_diagnostics/{maps,spectra}/`. + +Prerequisites, and what happens when they fail: a non-diffusion run, a stage other than +`inference`, a forecast engine without a `.diagnostics` hook, or a `diag_stream` absent from +`cf.streams` each log a warning and silently disable the diagnostic. A missing or duplicated +`LossPhysical` term asserts instead. + +## Failure modes + +Both of these mean the run has no usable physical decoder, not that a flag is wrong: + +- `AssertionError` at `assert len(outputs_physical) == 1` (`utils/validation_io.py`) — the active + test config has no `LossPhysical` term. +- `AssertionError: Empty preds but non-empty targets` — a `LossPhysical` term exists but the model + built no decoder, so nothing was predicted while targets were still loaded. + +A run trained with a latent-only loss has no decoder weights in its checkpoint and builds none at +load time: `Model.__init__` creates `embed_target_coords` / `target_token_engines` / `pred_heads` +only when `LossPhysical` appears in `training_config.losses` or `validation_config.losses` +(`test_config` is never inspected). Adding the loss alone therefore yields a *randomly initialised* +decoder, which is worse than none — the plots look plausible and mean nothing. + +To decode anyway, borrow the decoder of the backbone the run was initialised from: read +`references/decoder-overlay.md` and generate the overlay with +`scripts/make_decoder_overlay.py`. If no suitable backbone exists, say so and offer section A +instead rather than producing physical plots from an untrained decoder. diff --git a/.claude/skills/weathergen-inference-diagnostics/references/decoder-overlay.md b/.claude/skills/weathergen-inference-diagnostics/references/decoder-overlay.md new file mode 100644 index 0000000000..5d13a27848 --- /dev/null +++ b/.claude/skills/weathergen-inference-diagnostics/references/decoder-overlay.md @@ -0,0 +1,132 @@ +# Attaching a backbone decoder to a run that has none + +A run trained with a latent-only loss carries encoder and forecast-engine weights but no decoder. +An *overlay config*, passed with `--config`, adds the missing pieces without touching the run's +stored config: + +```bash +uv run inference --from-run-id \ + --config config/inference_decoder_overlay_.yml \ + --options +``` + +`scripts/make_decoder_overlay.py` writes that file. Read the three blocks below to check its output +or to write one by hand. + +## Block 1 — the physical loss (always) + +```yaml +validation_config: + losses: + physical: + type: LossPhysical + weight: 0.0 # computed and logged, but not added to the total + target_and_aux_calc: Physical + loss_fcts: + mse: {} +``` + +It must sit in `validation_config` (or `training_config`): that is what makes the decoder modules +get built at all. `test_config` inherits from `validation_config`, so this also satisfies the +"exactly one `LossPhysical` term" that the diagnostics and `write_output` require. Weight `0.0` +keeps the term out of the combined loss while still computing and logging it. + +## Block 2 — the decoder weights (always) + +```yaml +load_decoder_chkpt: {run_id: , mini_epoch: -1} +``` + +This overlays *only* `embed_target_coords.*`, `target_token_engines.*` and `pred_heads.*` on top of +the primary checkpoint; encoder and forecast engine are left alone. The backbone is normally the +run named in the model's `load_chkpt`: + +```bash +python3 -c " +import json +print(json.load(open('models//model_.json'))['load_chkpt']) +" +``` + +Confirm that checkpoint actually carries a decoder and read off its shapes: + +```bash +uv run python -c " +import torch +sd = torch.load('models//_latest.chkpt', + map_location='meta', mmap=True, weights_only=True) +for k, v in sd.items(): + if k.startswith(('embed_target_coords', 'pred_heads')): + print(k, tuple(v.shape)) +" +``` + +Two lines come back, e.g. `embed_target_coords..linear.weight (512, C)` and +`pred_heads..pred_heads.0.0.weight (T, 512)`. They fix everything the stream block must match: + +- **``** — the decoder is keyed by stream *name*. The overlay must provide a stream with exactly + this name. +- **`T`** — number of target channels of that stream. +- **`C`** — input width of the target-coordinate embedding, which is `geoinfo_size + 105` + (`get_targets_coords_size`: `geoinfo + 5*(3*5) + 3*8 + 6`). So `C` implies the number of geoinfo + channels the stream must declare. + +A mismatch in `T` or `C` is a hard `size mismatch` at load; a mismatch in the *name* is silent — +`load_state_dict(strict=False)` drops the weights and you decode with random ones. + +## Block 3 — the decoded stream (when the model's own config lacks it) + +Copy the stream **verbatim from the backbone's model JSON**; do not hand-write it: + +```bash +python3 -c " +import json, yaml +d = json.load(open('models//model_.json')) +print(yaml.safe_dump({'streams': {'': d['streams']['']}}, sort_keys=False)) +" +``` + +Keep the derived `train_source_channels` / `train_target_channels` / `val_source_channels` / +`val_target_channels` / `target_channel_weights` lists that the JSON carries. They are computed +only when streams are read from a `streams_directory`, which does not happen for an overlay config +— and with `*_target_channels` missing, `is_stream_forcing` sees zero target channels, classifies +the stream as forcing, and builds no decoder for it at all, without an error. +`val_target_channels` is also the list that `diag_channels` names must come from. + +Then switch off reconstruction for the model's own stream(s), so no second, randomly initialised +decoder is created: + +```yaml +streams: + : + reconstruct: false +``` + +Needed for any stream that declares target channels. Its geoinfo count typically differs from the +decoded stream's, so its decoder could not take the checkpoint weights in any case. + +## Verifying before you burn GPU time + +```bash +uv run python -c " +import weathergen.common.config as config +from pathlib import Path +from weathergen.utils.utils import is_stream_reconstructed +cf = config.load_merge_configs(None, '', -1, None, + Path('config/inference_decoder_overlay_.yml'), {}) +print('streams:', list(cf.streams)) +print('reconstructed:', {k: is_stream_reconstructed(v) for k, v in cf.streams.items()}) +print('val losses:', {k: v.type for k, v in cf.validation_config.losses.items()}) +print('load_decoder_chkpt:', cf.load_decoder_chkpt) +" +``` + +Expect exactly one stream reconstructed — the decoded one — one `LossPhysical` term, and the +backbone in `load_decoder_chkpt`. During the run, check the log for +`Loading decoder weights from id=...` and for any `Missing keys` naming decoder modules. + +## Caveat to state in any write-up + +A backbone decoder was trained on encoded ground-truth latents, never finetuned on +generated ones, so its error is part of every physical number produced this way. The +`decode(z)`-vs-truth panels of the ODE diagnostics are the control for exactly this. diff --git a/.claude/skills/weathergen-inference-diagnostics/scripts/make_decoder_overlay.py b/.claude/skills/weathergen-inference-diagnostics/scripts/make_decoder_overlay.py new file mode 100755 index 0000000000..1cf89cdea4 --- /dev/null +++ b/.claude/skills/weathergen-inference-diagnostics/scripts/make_decoder_overlay.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""Generate an inference overlay config that attaches a backbone's physical decoder to a +run that has none (typically a run trained with a latent-only loss). + +Writes the three blocks documented in ../references/decoder-overlay.md: + + 1. a zero-weight LossPhysical in validation_config -> makes the decoder modules exist + 2. load_decoder_chkpt -> supplies the decoder weights + 3. the backbone's decoded stream, verbatim -> makes the shapes and the name match + plus reconstruct: false on the model's own target-carrying streams + +Run it from the repo root (so that `models/` resolves), under the project environment because +the checkpoint inspection needs torch: + + uv run python .claude/skills/weathergen-inference-diagnostics/scripts/make_decoder_overlay.py \ + --model [--backbone ] [--decoded-stream ] + +With no --backbone it uses the model's own `load_chkpt.run_id`. With no --decoded-stream it takes +the stream the backbone's decoder weights are keyed to, and refuses if that is ambiguous. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import yaml + +DECODER_PREFIXES = ("embed_target_coords", "target_token_engines", "pred_heads") +# get_targets_coords_size(): geoinfo + 5*(3*5) + 3*8 + 6 +COORD_SIZE_OFFSET = 5 * (3 * 5) + 3 * 8 + 6 + + +def die(msg: str) -> None: + sys.exit(f"error: {msg}") + + +def run_config(models_dir: Path, run_id: str) -> dict: + for name in (f"model_{run_id}.json", f"model_{run_id}_latest.json"): + path = models_dir / run_id / name + if path.is_file(): + return json.loads(path.read_text()) + die(f"no model config for {run_id!r} under {models_dir / run_id}") + + +def checkpoint_path(models_dir: Path, run_id: str, mini_epoch: int) -> Path: + tag = "latest" if mini_epoch in (-1, None) else f"chkpt{mini_epoch:05d}" + path = models_dir / run_id / f"{run_id}_{tag}.chkpt" + if not path.is_file(): + die(f"no checkpoint {path}") + return path + + +def decoder_shapes(path: Path) -> dict[str, dict[str, tuple]]: + """Map stream name -> {parameter name: shape} for every decoder parameter.""" + import torch # local: only this step needs the project environment + + state = torch.load(path, map_location="meta", mmap=True, weights_only=True) + found: dict[str, dict[str, tuple]] = {} + for key, value in state.items(): + bare = key[len("module.") :] if key.startswith("module.") else key + if bare.startswith(DECODER_PREFIXES): + stream = bare.split(".")[1] + found.setdefault(stream, {})[bare] = tuple(value.shape) + return found + + +def check_shapes(stream_cfg: dict, shapes: dict[str, tuple], stream: str) -> list[str]: + """Compare the checkpoint's decoder shapes against the stream config it will be built from.""" + notes = [] + n_targets = len(stream_cfg.get("val_target_channels") or stream_cfg.get("target") or []) + n_geoinfo = len(stream_cfg.get("geoinfo_channels") or []) + + for name, shape in shapes.items(): + if name.startswith("pred_heads") and name.endswith("0.0.weight"): + if shape[0] != n_targets: + notes.append( + f"pred head emits {shape[0]} channels but stream {stream!r} declares " + f"{n_targets} target channels" + ) + if name.startswith("embed_target_coords") and name.endswith("weight"): + expected = n_geoinfo + COORD_SIZE_OFFSET + if shape[-1] != expected: + notes.append( + f"coord embedding expects {shape[-1]} inputs but stream {stream!r} implies " + f"{expected} ({n_geoinfo} geoinfo channels + {COORD_SIZE_OFFSET})" + ) + return notes + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--model", required=True, help="run id of the model to run inference with") + ap.add_argument("--backbone", help="run id providing the decoder (default: model's load_chkpt)") + ap.add_argument("--decoded-stream", help="stream the decoder is keyed to (default: detected)") + ap.add_argument("--mini-epoch", type=int, default=-1, help="backbone checkpoint to read") + ap.add_argument("--models-dir", type=Path, default=Path("models")) + ap.add_argument("-o", "--output", type=Path, help="default: config/inference_decoder_overlay_.yml") + args = ap.parse_args() + + model_cfg = run_config(args.models_dir, args.model) + + backbone = args.backbone or (model_cfg.get("load_chkpt") or {}).get("run_id") + if not backbone: + die(f"{args.model} has no load_chkpt; pass --backbone explicitly") + backbone_cfg = run_config(args.models_dir, backbone) + + shapes = decoder_shapes(checkpoint_path(args.models_dir, backbone, args.mini_epoch)) + if not shapes: + die(f"backbone {backbone!r} has no decoder weights ({'/'.join(DECODER_PREFIXES)})") + + stream = args.decoded_stream + if stream is None: + if len(shapes) > 1: + die(f"backbone {backbone!r} has decoders for {sorted(shapes)}; pass --decoded-stream") + stream = next(iter(shapes)) + elif stream not in shapes: + die(f"backbone {backbone!r} has no decoder for stream {stream!r} (has {sorted(shapes)})") + + if stream not in backbone_cfg.get("streams", {}): + die(f"stream {stream!r} is in the checkpoint but not in the backbone's model config") + stream_cfg = backbone_cfg["streams"][stream] + + for note in check_shapes(stream_cfg, shapes[stream], stream): + print(f"warning: {note}", file=sys.stderr) + + # the model's own streams keep their targets but must not get a second, untrained decoder + streams: dict[str, dict] = {} + for name, cfg in model_cfg.get("streams", {}).items(): + if name == stream: + continue + if cfg.get("train_target_channels") or cfg.get("val_target_channels") or cfg.get("target"): + streams[name] = {"reconstruct": False} + streams[stream] = stream_cfg # verbatim, including the derived channel lists + + overlay = { + "streams": streams, + "validation_config": { + "losses": { + "physical": { + "type": "LossPhysical", + "weight": 0.0, + "target_and_aux_calc": "Physical", + "loss_fcts": {"mse": {}}, + } + } + }, + "load_decoder_chkpt": {"run_id": backbone, "mini_epoch": args.mini_epoch}, + } + + out = args.output or Path("config") / f"inference_decoder_overlay_{backbone}.yml" + out.parent.mkdir(parents=True, exist_ok=True) + header = ( + f"# Decoder overlay: run {args.model} with the physical decoder of {backbone}.\n" + f"# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md.\n" + f"# Decoded stream: {stream}. The block below is copied verbatim from the backbone's\n" + f"# model config, derived channel lists included: nothing recomputes them for an overlay.\n" + ) + with out.open("w") as fh: + fh.write(header) + yaml.safe_dump(overlay, fh, sort_keys=False, width=100) + + print(f"wrote {out}") + print(f" decoded stream : {stream}") + print(f" not reconstructed: {[k for k in streams if k != stream] or '(none)'}") + print(f"\nuse with:\n uv run inference --from-run-id {args.model} --config {out} --options ...") + + +if __name__ == "__main__": + main() From 687508644fd42bbe105099edd4d7a65cb213ac3d Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Tue, 11 Aug 2026 10:40:36 +0200 Subject: [PATCH 05/15] some configs --- config/inference_decoder_overlay_zs4o66ch.yml | 344 ++++++++++++++++ config/zs4o66ch_reconstructed_config.yml | 385 ++++++++++++++++++ 2 files changed, 729 insertions(+) create mode 100644 config/inference_decoder_overlay_zs4o66ch.yml create mode 100644 config/zs4o66ch_reconstructed_config.yml diff --git a/config/inference_decoder_overlay_zs4o66ch.yml b/config/inference_decoder_overlay_zs4o66ch.yml new file mode 100644 index 0000000000..19ebf7bc8d --- /dev/null +++ b/config/inference_decoder_overlay_zs4o66ch.yml @@ -0,0 +1,344 @@ +# Decoder overlay: run tb4l9yu8 with the physical decoder of zs4o66ch. +# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md. +# Decoded stream: ERA5. The block below is copied verbatim from the backbone's +# model config, derived channel lists included: nothing recomputes them for an overlay. +streams: + ERA5: + type: anemoi + # Not in the backbone's block: the anemoi reader only subsamples when the stream declares a + # frequency (data_reader_anemoi.py:69), so the 1-hourly zarr returns all 6 hourly slices of + # each 6-hour window -- 6 x 542080 points with lat/lon repeated. The physical PSD then matches + # no known grid and every physical spectrum panel comes out "unavailable". 06:00:00 matches + # the run's forecast time_step (and ERA5_in's own frequency) and yields one N320 slice. + frequency: 06:00:00 + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target_exclude: + - z + - w_10 + - w_50 + - w_100 + - w_150 + - w_200 + - w_250 + - w_300 + - w_400 + - w_500 + - w_600 + - w_700 + - w_850 + - w_925 + - w_1000 + - slor + - sdor + - tcw + - cp + - tp + - q_50 + - q_100 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + reconstruct: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 +validation_config: + # tb4l9yu8 trains with training_mode: [student_teacher] only. The data sampler collects + # target_coords / target_values solely when "masking" is in the mode + # (multi_stream_data_sampler.py:759), so without this the batch carries no ERA5 targets, + # predict_decoders skips the stream and the diagnostics die on KeyError: 'ERA5'. + # The source side is unchanged: "masking" only *adds* target_coords/target_values. + training_mode: + - masking + - student_teacher + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: zs4o66ch + mini_epoch: -1 diff --git a/config/zs4o66ch_reconstructed_config.yml b/config/zs4o66ch_reconstructed_config.yml new file mode 100644 index 0000000000..5b2c1dd6ad --- /dev/null +++ b/config/zs4o66ch_reconstructed_config.yml @@ -0,0 +1,385 @@ +# (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. + + +# ============================================================================= +# RECONSTRUCTED training config for run zs4o66ch (reconstructed 2026-08-10) +# +# Source of truth: models/zs4o66ch/model_zs4o66ch.json (the checkpoint's own +# resolved config). Reconstructed by diffing that JSON against every config in +# config/ -- it matches config_jepa_multi_data_all_years.yml exactly except for +# the single ema_ramp_up_ratio deviation marked DEVIATION below. Every other +# key, including all 10 stream blocks, is identical. +# +# What zs4o66ch is: a masked autoencoder + JEPA SSL model. fe_num_blocks: 0 and +# forecast.num_steps: 0 -- it has NO forecasting engine and was never trained on +# a lead time. Its ERA5 decoder was trained purely as a 20%-keep spatial +# inpainter (model_input rate 0.2, `subset`-correspondence with an 80% +# contiguous cropping_healpix target crop) under MSE, with reconstruction +# competing against JEPA (w 1.0), context loss (w 0.4) and deep-SSL. +# +# NOT reconstructable from the checkpoint (runtime state, left as ??? / defaults): +# general.run_id / rank / world_size, data_loading.rng_seed (was 1785849062), +# data_path_* and path_shared_* (recorded as JSC/ECMWF paths: /e/data1/..., +# /e/scratch/weatherai/... -- NOT the Santis paths this checkout resolves to), +# mlflow.tracking_uri, world_size 8 (2 nodes x 4). +# +# Training history: stage was `train_continue` with from_run_id == zs4o66ch, +# i.e. six in-place continuation segments (run_history 6620, 13740, 21380, +# 28500, 35620, 36120), ending at istep 36120 -- short of the 64 x 8192 the +# num_mini_epochs/samples_per_mini_epoch below imply for a single leg. +# +# Numeric literals in scientific notation were rewritten 1e-6 -> 1.0e-6 etc. +# YAML 1.1 parses the bare form as a *string*; the checkpoint holds floats, so +# these forms are the faithful ones. No value changed. +# ============================================================================= + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 1024 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: True + +ae_local_num_queries: 1 +ae_local_queries_per_cell: False +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: True +ae_adapter_with_residual: True +ae_adapter_dropout_rate: 0.1 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 3 +ae_global_num_heads: 16 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: True +# TODO: switching to < 1 triggers triton-related issues. +# See https://github.com/ecmwf/WeatherGenerator/issues/1050 +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 4 +ae_global_trailing_layer_norm: True + +ae_aggregation_num_blocks: 5 +ae_aggregation_num_heads: 16 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: True +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 4 + +decoder_type: PerceiverIOCoordConditioning # CrossAttentionAdaNormConditioning +pred_adapter_kv: False +pred_self_attention: True +pred_dyadic_dims: False +pred_mlp_adaln: True +num_class_tokens: 0 +num_register_tokens: 64 + +# noise before predictor in JEPA +noise_pre_predictor_std: 5.0e-4 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +fe_num_blocks: 0 +fe_diffusion_model: null +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_layer_norm_after_blocks: [7] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth layer +fe_impute_latent_noise_std: 0.0 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False + +healpix_level: 5 +# 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 +rope_2D: True +# mlp_type: swiglu +# use_xsa: True + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: False +ddp_find_unused_parameters: False +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1.0e-5 +norm_eps: 1.0e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + +freeze_modules: "" + +norm_type: "LayerNorm" +qk_norm_type: "RMSNorm" +# type of zarr_store +zarr_store: "zip" # "zarr" for LocalStore, "zip" for ZipStore + +##################################### + +streams_directory: "./config/streams/era5_georing_avhrr_synop_lowres_sjepa/" +streams: ??? + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + run_id: ??? + run_history: [] + +# logging frequency in the training loop (in number of batches) +train_logging: + terminal: 100 + metrics: 20 + checkpoint: 500 + log_grad_norms: False + + loss_spike_detection: + enabled: False + skip_batch: False # also makes the skip explicit instead of relying on GradScaler + loss_threshold: 0.01 + + # Collapse monitoring for SSL training (JEPA/DINO/iBOT) + # Detects representation collapse via various metrics + # For forecasting, supports sequences of latents with per-step and aggregate metrics + collapse_monitoring: + enabled: False + compute_frequency: 100 # batches between metric computations + log_frequency: 100 # batches between metric logging + metrics: + effective_rank: + enabled: true + tensor_source: "both" # "student", "teacher", or "both" + sample_size: 2048 # max samples for SVD (0 = no sampling) + # For forecasting sequences: "all" (per-step + aggregates), + # "aggregate_only" (mean/min/max/degradation), "per_step_only" + forecast_aggregation: "all" + singular_values: + enabled: true + tensor_source: "both" + sample_size: 2048 + forecast_aggregation: "all" + dimension_variance: + enabled: true + tensor_source: "both" # cheap to compute, good early indicator + forecast_aggregation: "all" + prototype_entropy: + enabled: true # only applies to DINO + ema_beta: + enabled: true + +# parameters for data loading +data_loading : + + num_workers: 10 + rng_seed: ??? + + +# config for training +training_config: + + # training_mode: "masking", "student_teacher", "latent_loss" + training_mode: ["masking", "student_teacher"] + + # Deep self-supervision (V-JEPA 2.1 style): compute SSL loss at multiple encoder depths. + # When enabled, intermediate encoder representations are tapped and used as additional + # targets/predictions for the SSL loss. Disabled by default (empty tap_after). + deep_ssl: + enabled: true + tap_after: # named tap points (final output always included implicitly) + - local2global # after local-to-global adapter, before global attention + - "global:0" # after 1st global layer (0-indexed) + - "global:1" # after 2nd global layer + fusion_hidden_factor: 4 # MLP hidden dim = factor * dim_embed + level_weights: [0.25, 0.25, 0.25, 0.25] # per-level loss weights (len = num taps + 1) + + # Context loss (V-JEPA 2.1): L1 on context (visible) tokens between student and teacher. + # Complements the JEPA loss which targets masked tokens. + context_loss: + enabled: true + weight: 0.4 + warmup_steps: 1000 # linear warmup from 0 to weight over this many steps + + num_mini_epochs: 64 + samples_per_mini_epoch: 8192 + shuffle: True + + start_date: 1980-01-01T00:00 + end_date: 2022-12-31T00:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + window_offset_prediction : 0 + + learning_rate_scheduling : + lr_start: 1.0e-6 + lr_max: 1.0e-4 + lr_final_decay: 1.0e-6 + lr_final: 0.0 + num_steps_warmup: 512 + num_steps_cooldown: 2048 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 0.5 + weight_decay: 0.05 + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.9875 # at B=8 beta1 =0.9 + beta2 : 0.994 # at B=8 beta1 approx 0.95 + eps : 2.0e-8 + + losses : { + "physical": { + enabled: True, + type: LossPhysical, + loss_fcts: { "mse": { + target_source_correspondence: {0 : {0 : "subset"} }, + }, }, + }, + "student-teacher": { + enabled: True, + type: LossLatentSSLStudentTeacher, + weight: 1.0, + loss_fcts : { + "JEPA": { + 'weight': 1, "loss_extra_args": {}, "out_dim": 2048, "head": transformer, + "num_blocks": 6, "num_heads": 24, "with_qk_lnorm": True, "intermediate_dim": 1536, + "dropout_rate": 0.1, + target_source_correspondence: {0 : {0 : "subset"} }, + }, + }, + target_and_aux_calc: { "EMATeacher" : + # DEVIATION from config_jepa_multi_data_all_years.yml (null there). + # zs4o66ch's checkpoint records 0.1 -- presumably set via --options at launch. + { ema_ramp_up_ratio : 0.1, + ema_halflife_in_thousands: 300000, + model_param_overrides : { + training_config: { losses: { student-teacher:{ loss_fcts :{JEPA: {head: identity} }}}} + }, + } + } + } + } + + model_input: { + "random_easy" : { + # masking strategy: "random", "forecast" + masking_strategy: "random", + num_samples: 1, + num_steps_input: 1, + masking_strategy_config : { + diffusion_rn : True, + rate : 0.2, + rate_sampling: False + }, + }, + } + + target_input: { + "random_easy_target" : { + masking_strategy: "cropping_healpix", + num_samples: 1, + masking_strategy_config : { rate : 0.8, hl_mask: 3, rate_sampling: False }, + }, + } + + forecast : + time_step: 00:00:00 + num_steps: 0 + policy: null + + +# validation config; full validation config is merge of training and validation config +validation_config: + + samples_per_mini_epoch: 256 + shuffle: False + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T00:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : False + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1.0e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: False + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null + From e5dd0d423242e70c3a52b16d92d2328b3376139e Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Fri, 14 Aug 2026 11:37:39 +0200 Subject: [PATCH 06/15] fix model.py after merge --- src/weathergen/model/model.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/weathergen/model/model.py b/src/weathergen/model/model.py index 123c5914bd..4c080795c0 100644 --- a/src/weathergen/model/model.py +++ b/src/weathergen/model/model.py @@ -904,7 +904,9 @@ def forward( # 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) + lambda toks, step=step: self.decode_tokens( + model_params, step, toks, source_samples + ) ) # apply forecasting engine @@ -1082,7 +1084,7 @@ def decode_tokens( model_params: ModelParams, step: int, tokens: torch.Tensor, - batch: ModelBatch, + batch: BatchSamples, ) -> dict: """Decode arbitrary latent tokens to physical space, outside the forward's bookkeeping. @@ -1091,7 +1093,7 @@ def decode_tokens( Returns ``{stream_name: (pred_per_batch_item, ...)}``. """ return self.predict_decoders( - model_params, step, tokens, batch, ModelOutput(1), out_step=0 + model_params, step, tokens, batch, ModelOutput([0], 0, batch), out_step=0 ).physical[0] def predict_latent( From 842aaed9bab7b548162a49bf9a9cebee391fbfad Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Fri, 14 Aug 2026 11:56:52 +0200 Subject: [PATCH 07/15] merge commit --- pyproject.toml | 35 +++++++++++------------------------ 1 file changed, 11 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 00065b4e3b..8cd6880012 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -239,32 +239,19 @@ explicit = true [tool.uv.sources] weathergen-common = { workspace = true } weathergen-evaluate = { workspace = true } -weathergen-metrics = { workspace = true } -weathergen-readers-extra = { workspace = true } - - -flash-attn = [ - { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, - { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.8.3+cu12torch2.9cxx11abiFALSE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, - # There is no official support for flash-attn on macos -] - torch = [ -# Explicit pin for GPU - { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, - { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, -# Explicit pin for CPU. - { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, - { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, - - # No acceleration for now. - # TODO: explore pytorch + metal backend - { index = "pytorch-cpu", marker = "sys_platform == 'darwin'", extra="gpu"}, + { index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, + { index = "pytorch-cpu", marker = "sys_platform == 'macosx'"}, +] +# This URL was evaluated this way: +# uv run ~/WeatherGenerator-private/hpc/hpc2020/ecmwf/get-flash-atten.sh +# get-flash-atten.sh only supports x86_64 (Dao-AILab has no official aarch64 wheels for +# v2.7.4.post1); the aarch64 build below is a community build from cathalobrien/get-flash-attn, +# matched to the same torch (2.6), CUDA (12), Python (3.12) and ABI (FALSE) as the x86_64 wheel. +flash-attn = [ + { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.7.4+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, + { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, ] -# TODO: switch back to this branch, but this is a larger update and the tag name is incorrect -# anemoi-datasets = { git = "https://github.com/ecmwf/anemoi-datasets", tag = "v0.0.2-special-zarr3"} -anemoi-datasets = { git = "https://github.com/tjhunter/anemoi-datasets", tag = "0.0.1"} - [tool.pytest.ini_options] log_cli = true From 0ea7fc7139b7692ca3f2c4bf74e9eb87ab20c052 Mon Sep 17 00:00:00 2001 From: iluise Date: Fri, 14 Aug 2026 15:34:16 +0200 Subject: [PATCH 08/15] run fast eval without target --- .../evaluate/io/data/io_orchestration.py | 31 ++++- .../weathergen/evaluate/io/data/io_workers.py | 123 +++++++++++++++++- .../weathergen/evaluate/io/wegen_reader.py | 3 + .../src/weathergen/evaluate/run_evaluation.py | 11 +- 4 files changed, 154 insertions(+), 14 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py b/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py index 4131bf20f0..5561c58638 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py +++ b/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py @@ -83,6 +83,7 @@ class IOState: offset: np.timedelta64 | None = ( None # fallback offset in hours for init_time when source_interval is missing ) + anemoi_target_cfg: dict | None = None # if set, read targets from anemoi dataset # --------------------------------------------------------------------------- @@ -254,6 +255,7 @@ def _build_io_state( n_io_workers: int, ens_select: EnsembleSelect, rank: str = "", + inference_cfg: dict | None = None, ) -> IOState: """Resolve all I/O parameters that are shared between the two impl paths.""" zarr_path = str(fname_zarr) @@ -278,6 +280,19 @@ def _build_io_state( if isinstance(regrid_opts, bool) and regrid_opts: regrid_opts = {"target_grid": [1.5, 1.5]} + # ---- Resolve anemoi target config from inference config ---- + anemoi_target_cfg = None + if inference_cfg: + stream_info = inference_cfg.get("streams", {}).get(stream, {}) + if stream_info.get("type") in ("anemoi", "anemoi_operan") and stream_info.get("filenames"): + data_path = inference_cfg.get("data_path_anemoi", "") + filename = str(Path(data_path) / stream_info["filenames"][0]) + anemoi_target_cfg = { + "filename": filename, + "channels": stream_info.get("val_target_channels", []), + } + _logger.info(f"Anemoi target source: {filename}") + return IOState( run_id=run_id, zarr_path=zarr_path, @@ -295,10 +310,11 @@ def _build_io_state( coords=coords, lat=lat, lon=lon, - n_workers=n_io_workers, + n_workers=min(n_io_workers, 20) if anemoi_target_cfg is not None else n_io_workers, rank=rank, offset=offset, regrid_opts=regrid_opts, + anemoi_target_cfg=anemoi_target_cfg, ) @@ -315,6 +331,7 @@ def _parallel_read( backend: str, label: str, regrid_opts: dict, + anemoi_target_cfg: dict | None = None, ) -> tuple[list, bool]: """Dispatch _read_sample over samples, with parallel→sequential fallback. @@ -333,6 +350,7 @@ def _parallel_read( read_coords=need_coords, is_gridded=is_gridded, regrid_opts=regrid_opts, + anemoi_target_cfg=anemoi_target_cfg, ) calls = [delayed(_read_sample)(sample=s, **kwargs) for s in samples] @@ -550,6 +568,7 @@ def get_data_dirstore(state: IOState) -> ReaderOutput: backend=state.backend, label=f"RUN {state.run_id} [rank {state.rank}] - {state.stream} fstep {fs}", regrid_opts=state.regrid_opts, + anemoi_target_cfg=state.anemoi_target_cfg, ) # If _parallel_read fell back to sequential, honour that for the rest if fell_back: @@ -583,7 +602,7 @@ def get_data_dirstore(state: IOState) -> ReaderOutput: del results - if n_workers > 1: + if n_workers > 1 and state.backend == "loky": get_reusable_executor().shutdown(wait=True) _logger.info( @@ -623,6 +642,7 @@ def get_data_zipstore(state: IOState) -> ReaderOutput: read_coords=not state.is_gridded, is_gridded=state.is_gridded, regrid_opts=state.regrid_opts, + anemoi_target_cfg=state.anemoi_target_cfg, ) calls = [ delayed(_read_sample)(sample=s, fsteps=[fs], **kwargs) @@ -637,6 +657,11 @@ def get_data_zipstore(state: IOState) -> ReaderOutput: verbose=5, ) + _logger.info( + f"RUN {state.run_id} [rank {state.rank}] - {state.stream}: " + f"dispatch_parallel returned {len(flat_results)} results. Assembling..." + ) + # --- Re-group: flat_results[sample_idx * n_fsteps + fstep_idx] -------- n_fsteps = len(state.fsteps) # Gather per-sample results in the same shape as get_data_dirstore expects @@ -694,7 +719,7 @@ def get_data_zipstore(state: IOState) -> ReaderOutput: del flat_results - if state.n_workers > 1: + if state.n_workers > 1 and state.backend == "loky": get_reusable_executor().shutdown(wait=True) _logger.info( diff --git a/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py b/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py index 643246d422..2b9b2dc7c3 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py +++ b/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py @@ -15,7 +15,9 @@ import contextlib import logging +import threading +import anemoi.datasets import numpy as np import zarr from numpy.typing import NDArray @@ -24,6 +26,101 @@ _logger = logging.getLogger(__name__) +# Module-level cache so that threading workers reuse the same anemoi dataset +# handle across tasks dispatched to the same process. +_anemoi_cache: dict[str, tuple] = {} +_anemoi_lock = threading.Lock() + + +def _open_anemoi_dataset(anemoi_cfg: dict) -> tuple: + """Open the anemoi dataset once and precompute target channel indices. + + Uses a module-level cache keyed by filename so that threading workers + reuse the same handle. A lock ensures only one thread opens the dataset. + + Returns ``(ds, target_idx, ds_dates)`` for reuse across forecast steps. + """ + filename = anemoi_cfg["filename"] + if filename in _anemoi_cache: + return _anemoi_cache[filename] + + with _anemoi_lock: + # Double-check after acquiring lock + if filename in _anemoi_cache: + return _anemoi_cache[filename] + + target_channels = anemoi_cfg["channels"] + + _logger.info(f"Opening anemoi dataset {filename} (lazy)") + ds = anemoi.datasets.open_dataset(filename) + target_idx = [ds.variables.index(ch) for ch in target_channels] + ds_dates = ds.dates.astype("datetime64[s]") + result = (ds, target_idx, ds_dates) + _anemoi_cache[filename] = result + return result + + +def _read_anemoi_target( + ds, + target_idx: list[int], + ds_dates: np.ndarray, + times: np.ndarray, + channel_idxs: list[int] | None, +) -> np.ndarray: + """Read target data from a pre-opened anemoi dataset for the given valid times. + + Parameters + ---------- + ds : anemoi dataset handle + Already-opened dataset (from ``_open_anemoi_dataset``). + target_idx : list[int] + Indices of val_target_channels within ``ds.variables``. + ds_dates : np.ndarray + ``ds.dates`` cast to ``datetime64[s]`` (precomputed). + times : np.ndarray + Valid times (datetime64) to extract from the dataset. + channel_idxs : list[int] | None + Channel indices to select (applied after reading). + + Returns + ------- + np.ndarray + Target data array, shape ``(n_points, n_channels)``, float32. + """ + unique_times = np.unique(times) + + # Read one time-slice at a time (typically ≤6 per fstep) + all_data = [] + for ut in unique_times: + ut_s = np.datetime64(ut, "s") + matches = np.where(ds_dates == ut_s)[0] + + if len(matches) == 0: + raise ValueError( + f"Time {ut} not found in anemoi dataset. " + f"Available range: {ds_dates[0]} .. {ds_dates[-1]}" + ) + + idx = int(matches[0]) + # ds[idx] → shape (n_variables, n_ens, n_gridpoints) + data_slice = np.asarray(ds[idx]) + # Select target channels and squeeze ensemble dim → (n_gridpoints, n_target_channels) + data_slice = data_slice[target_idx, 0, :].T + all_data.append(data_slice) + + # For gridded data every grid point shares the same unique time, + # so each slice already has the right shape. + if len(all_data) == 1: + target_data = all_data[0] + else: + target_data = np.concatenate(all_data, axis=0) + + # Apply the same early channel selection as the zarr path + if channel_idxs is not None: + target_data = target_data[:, channel_idxs] + + return target_data.astype(np.float32) + def _compute_early_channel_selection( read_channels: list[str], @@ -92,6 +189,7 @@ def _read_sample( read_coords: bool = False, is_gridded: bool = True, regrid_opts: dict | None = None, + anemoi_target_cfg: dict | None = None, ) -> tuple[list[NDArray], list[NDArray], list[NDArray], dict]: """ Read all forecast steps for one sample via direct zarr array access. @@ -172,22 +270,35 @@ def _read_sample( if source_interval: break - for fs in fsteps: + # Open anemoi dataset once for all fsteps (if configured) + anemoi_handle = None + if anemoi_target_cfg is not None: + anemoi_handle = _open_anemoi_dataset(anemoi_target_cfg) + + for fi, fs in enumerate(fsteps): base = f"{sample}/{stream}/{fs}" # Direct array access — bypasses OutputDataset/as_xarray/dask entirely pred_data = np.asarray(ds[f"{base}/prediction/data"]) - target_data = np.asarray(ds[f"{base}/target/data"]) times_data = np.asarray(ds[f"{base}/prediction/times"]) - # Select channels by index + if anemoi_handle is not None: + # Read target from anemoi dataset + anemoi_ds, target_idx, ds_dates = anemoi_handle + target_data = _read_anemoi_target( + anemoi_ds, target_idx, ds_dates, times_data, channel_idxs + ) + else: + target_data = np.asarray(ds[f"{base}/target/data"]) + if channel_idxs is not None: + target_data = target_data[:, channel_idxs] + + # Select channels by index for prediction if channel_idxs is not None: pred_data = ( pred_data[:, channel_idxs] if pred_data.ndim == 2 else pred_data[:, channel_idxs, :] ) - target_data = target_data[:, channel_idxs] - - # Handle sub-steps (gridded data with multiple valid_times per fstep). + # Handle sub-steps (gridded data with multiple valid_times per fstep). # For scatter/observation data each observation has its own timestamp, # so splitting by unique time would create one tiny array per obs — # thousands of them — causing the assembly code to hang. diff --git a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py index 6589813305..ffb1ccd36e 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py +++ b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py @@ -616,6 +616,9 @@ def get_data( self._num_io_workers, ens_select, rank=rank_file.stem.split("rank")[-1], + inference_cfg=self.inference_cfg + if self.eval_cfg.get("type") == "anemoi-target" + else None, ) get_data_fn = get_data_zipstore if state.is_zip else get_data_dirstore result = get_data_fn(state) diff --git a/packages/evaluate/src/weathergen/evaluate/run_evaluation.py b/packages/evaluate/src/weathergen/evaluate/run_evaluation.py index e6af331012..a273124356 100755 --- a/packages/evaluate/src/weathergen/evaluate/run_evaluation.py +++ b/packages/evaluate/src/weathergen/evaluate/run_evaluation.py @@ -162,7 +162,7 @@ def get_reader( region: str | None = None, metric: dict[str, object] | None = None, ): - if reader_type == "zarr": + if reader_type == "zarr" or reader_type == "anemoi-target": reader = WeatherGenZarrReader(run, run_id, private_paths) elif reader_type == "csv": reader = CsvReader(run, run_id, private_paths) @@ -220,11 +220,12 @@ def _process_stream( _logger.info(f"Stream {stream} not found for run {run_id}. Skipping.") return run_id, stream, {}, {} - needs_plotting = stream_dict.get("plotting") and type_ == "zarr" + is_zarr_like = type_ in ("zarr", "anemoi-target") + needs_plotting = stream_dict.get("plotting") and is_zarr_like needs_scoring = stream_dict.get("evaluation", False) output_data = None - if (needs_plotting or needs_scoring) and type_ == "zarr": + if (needs_plotting or needs_scoring) and is_zarr_like: available_data = reader.check_availability(stream, mode="evaluation") output_data = None @@ -247,9 +248,9 @@ def _process_stream( if not needs_scoring: return run_id, stream, {}, {} - plot_score_maps = plot_score_options.get("plot_score_maps", False) and type_ == "zarr" + plot_score_maps = plot_score_options.get("plot_score_maps", False) and is_zarr_like plot_score_init_time_series = ( - plot_score_options.get("plot_score_init_time_series", False) and type_ == "zarr" + plot_score_options.get("plot_score_init_time_series", False) and is_zarr_like ) stream_loaded_scores, recomputable_metrics = reader.load_scores(stream, regions, metrics) From f34c0187c662e6b2ef90723af27bd471aeec16b4 Mon Sep 17 00:00:00 2001 From: iluise Date: Fri, 14 Aug 2026 16:47:08 +0200 Subject: [PATCH 09/15] fix streams problem (only for eval) --- packages/common/src/weathergen/common/io.py | 14 ++++++++++++-- .../src/weathergen/evaluate/io/data/io_workers.py | 7 +++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/packages/common/src/weathergen/common/io.py b/packages/common/src/weathergen/common/io.py index 685cd5d8e7..0973e72315 100644 --- a/packages/common/src/weathergen/common/io.py +++ b/packages/common/src/weathergen/common/io.py @@ -499,9 +499,19 @@ def forecast_offset(self) -> int: def example_key(self) -> ItemKey: try: sample, example_sample = next(self.data_root.groups()) + # Find the first stream that has prediction/target data (not just source) + for stream, example_stream in example_sample.groups(): + fstep_keys = sorted(example_stream.group_keys()) + for fk in fstep_keys: + fstep_group = example_stream[fk] + child_names = set(fstep_group.group_keys()) + if "prediction" in child_names or "target" in child_names: + return ItemKey(sample, int(fk), stream) + # Fallback: use first stream / first fstep even without prediction/target stream, example_stream = next(example_sample.groups()) - fstep = 0 - except StopIteration as e: + fstep_keys = sorted(example_stream.group_keys()) + fstep = int(fstep_keys[0]) if fstep_keys else 0 + except (StopIteration, IndexError) as e: msg = f"Data store at: {self._store_path} is empty." raise FileNotFoundError(msg) from e diff --git a/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py b/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py index 2b9b2dc7c3..61946379b2 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py +++ b/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py @@ -89,6 +89,13 @@ def _read_anemoi_target( """ unique_times = np.unique(times) + # Handle empty times (some fsteps have no data) + if len(unique_times) == 0: + n_ch = len(target_idx) + if channel_idxs is not None: + n_ch = len(channel_idxs) + return np.empty((0, n_ch), dtype=np.float32) + # Read one time-slice at a time (typically ≤6 per fstep) all_data = [] for ut in unique_times: From 780a9a24df9707cac6417dd9e44fecedc9d96fd1 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Fri, 14 Aug 2026 17:07:00 +0200 Subject: [PATCH 10/15] changes to pyproject and spectral plotting (changes evaluate package!) --- packages/common/pyproject.toml | 2 +- packages/evaluate/pyproject.toml | 2 +- .../evaluate/plotting/line_plots.py | 94 +++++++++++++++++++ .../evaluate/plotting/plot_utils.py | 18 ++++ .../src/weathergen/evaluate/scores/psd.py | 17 +++- packages/metrics/pyproject.toml | 2 +- packages/readers_extra/pyproject.toml | 2 +- pyproject.toml | 38 +++++--- 8 files changed, 157 insertions(+), 18 deletions(-) diff --git a/packages/common/pyproject.toml b/packages/common/pyproject.toml index 58a5ad407c..148dc99ddf 100644 --- a/packages/common/pyproject.toml +++ b/packages/common/pyproject.toml @@ -21,7 +21,7 @@ dependencies = [ dev = [ "pytest>=8.3", "pytest-mock>=3.14.1", - "ruff==0.9.7", + "ruff==0.12.2", "pyrefly==0.36.0", ] diff --git a/packages/evaluate/pyproject.toml b/packages/evaluate/pyproject.toml index be7ea34950..403649227e 100644 --- a/packages/evaluate/pyproject.toml +++ b/packages/evaluate/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ dev = [ "pytest>=8.3", "pytest-mock>=3.14.1", - "ruff==0.9.7", + "ruff==0.12.2", "pyrefly==0.36.0", ] diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py b/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py index 862c82d427..ecc665090c 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py @@ -18,6 +18,7 @@ import numpy as np import seaborn as sns import xarray as xr +from matplotlib.lines import Line2D from weathergen.evaluate.plotting.plot_utils import ( align_labels, @@ -812,3 +813,96 @@ def psd_plot( _logger.debug(f"Saving PSD summary plot to {fname}") fig.savefig(str(fname), bbox_inches="tight", dpi=self.dpi_val) plt.close(fig) + + def psd_evolution_plot( + self, + per_fstep_datasets: dict[int, dict], + tag: str = "", + variable: str = "", + label: str = "", + ) -> None: + """Overlay PSD spectra (and their pred/target ratio) across forecast lead-time steps in + one two-panel plot, colour-coded by step. + + Mirrors the diffusion-diagnostics "evolution" plot style (viridis colormap + colorbar + keyed to the swept step), applied here to forecast lead-time steps instead of diffusion + denoising steps, and mirrors ``psd_plot``'s two-panel (spectra + ratio) layout. In the + top panel, solid lines are predictions and dashed lines are targets; both share the + same per-step colour so prediction/target can be compared at a given lead time while + also seeing how each evolves across lead time. The bottom panel shows the pred/target + ratio per step with the same per-step colour, where lead-time drift is expected to be + most visible. + + Parameters + ---------- + per_fstep_datasets : dict[int, dict] + Maps forecast step -> dict with keys ``frequencies``, ``psd_target``, + ``psd_prediction``, ``psd_method`` (as produced by ``_extract_psd_attrs``). + tag : str + Filename tag. + variable : str + Channel name, used in the title. + label : str + Run label, used in the title. + """ + out_dir = Path(self.out_plot_dir_psd) + out_dir.mkdir(parents=True, exist_ok=True) + + fsteps = sorted(per_fstep_datasets) + if len(fsteps) < 2: + return + + fig, (ax_spec, ax_ratio) = plt.subplots( + 2, + 1, + figsize=self.fig_size or (7, 7), + gridspec_kw={"height_ratios": [2, 1], "hspace": 0.08}, + ) + cmap = plt.get_cmap("viridis") + n = max(len(fsteps) - 1, 1) + + for i, fstep in enumerate(fsteps): + ds = per_fstep_datasets[fstep] + freq = np.asarray(ds["frequencies"]) + pred = np.asarray(ds["psd_prediction"]) + tar = np.asarray(ds["psd_target"]) + c = cmap(i / n) + + ax_spec.loglog(freq, pred, color=c, lw=1.0) + ax_spec.loglog(freq, tar, color=c, lw=1.0, ls="--", alpha=0.6) + + with np.errstate(divide="ignore", invalid="ignore"): + ratio = np.where(tar > 0, pred / tar, np.nan) + ax_ratio.semilogx(freq, ratio, color=c, lw=1.0) + + sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(fsteps[0], fsteps[-1])) + fig.colorbar(sm, ax=[ax_spec, ax_ratio], label="Forecast step") + + ax_spec.legend( + handles=[ + Line2D([], [], color="gray", lw=1.0, ls="-", label="Prediction"), + Line2D([], [], color="gray", lw=1.0, ls="--", alpha=0.6, label="Target"), + ], + frameon=False, + fontsize=8, + ) + psd_method = next(iter(per_fstep_datasets.values())).get("psd_method", "sht") + title_parts = [f"PSD evolution ({psd_method})"] + if variable: + title_parts.append(variable) + if label: + title_parts.append(label) + ax_spec.set_title(" – ".join(title_parts)) + ax_spec.set_ylabel("Power") + ax_spec.grid(True, which="both", ls="--", alpha=0.4) + + ax_ratio.axhline(1.0, ls="--", color="gray", lw=0.8) + ax_ratio.set_ylabel("Pred / Target") + ax_ratio.set_xlabel("Frequency (1/deg)") + ax_ratio.set_ylim(0, 2) + ax_ratio.grid(True, which="both", ls="--", alpha=0.4) + + fname = out_dir / f"{tag or 'psd_evolution'}.{self.image_format}" + _logger.debug(f"Saving PSD evolution plot to {fname}") + fig.savefig(str(fname), bbox_inches="tight", dpi=self.dpi_val) + plt.close(fig) diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py index bf9bd421e4..86547e6078 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py @@ -843,10 +843,12 @@ def psd_plot_metric_region( label = runs[run_id].get("label", run_id) + per_fstep_datasets = {} for fstep in attr_fsteps: psd_datasets = _extract_psd_attrs(data_ch, fstep, ch) if psd_datasets is None: continue + per_fstep_datasets[fstep] = psd_datasets[0] method_tag = psd_datasets[0].get("psd_method", "sht") name = create_filename( @@ -861,6 +863,22 @@ def psd_plot_metric_region( variable=ch, forecast_step=str(fstep), ) + + if len(per_fstep_datasets) >= 2: + method_tag = next(iter(per_fstep_datasets.values())).get( + "psd_method", "sht" + ) + evo_name = create_filename( + prefix=[metric, method_tag, region], + middle=[run_id], + suffix=[stream, ch, "evolution"], + ) + plotter.psd_evolution_plot( + per_fstep_datasets, + tag=evo_name, + variable=ch, + label=label, + ) _logger.info(f"PSD plots saved successfully into: {plotter.out_plot_dir_psd}") diff --git a/packages/evaluate/src/weathergen/evaluate/scores/psd.py b/packages/evaluate/src/weathergen/evaluate/scores/psd.py index 4021321797..5bc22265f9 100644 --- a/packages/evaluate/src/weathergen/evaluate/scores/psd.py +++ b/packages/evaluate/src/weathergen/evaluate/scores/psd.py @@ -328,7 +328,7 @@ def detect_grid_type( Returns ------- str | None - ``"octahedral"``, ``"regular"``, or ``None`` if detection fails. + ``"octahedral"``, ``"regular"``, ``"reduced"``, or ``None`` if detection fails. """ unique_lats = np.unique(lats) nlat = len(unique_lats) @@ -358,6 +358,19 @@ def detect_grid_type( _logger.debug(f"Detected regular lat-lon grid (nlat={nlat}).") return "regular" + # Check reduced Gaussian (ECMWF N320, e.g. ERA5's native grid) + if nlat == 640: # sht_psd's "reduced" path is hard-wired to N320 + try: + from anemoi.transform.grids.named import lookup + + expected_reduced = len(np.asarray(lookup("N320")["latitudes"])) + except Exception as exc: # noqa: BLE001 - anemoi missing / offline / unknown grid + _logger.warning(f"Reduced Gaussian (N320) grid check unavailable: {exc}") + expected_reduced = None + if expected_reduced is not None and n_points == expected_reduced: + _logger.debug(f"Detected reduced Gaussian (N320) grid (nlat={nlat}).") + return "reduced" + # 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: @@ -712,7 +725,7 @@ def compute_psd_score( lat_range : tuple[float, float] Latitude bounds for fft method. grid_type : str | None - Pre-detected grid type (``"octahedral"``, ``"regular"``). + Pre-detected grid type (``"octahedral"``, ``"regular"``, ``"reduced"``). When ``None``, the grid type is auto-detected from lats/lons. Pass a pre-computed value to avoid repeated detection across channels. diff --git a/packages/metrics/pyproject.toml b/packages/metrics/pyproject.toml index 404c96ebd9..f1e3cf9b4a 100644 --- a/packages/metrics/pyproject.toml +++ b/packages/metrics/pyproject.toml @@ -13,7 +13,7 @@ dependencies = [ dev = [ "pytest>=8.3", "pytest-mock>=3.14.1", - "ruff==0.9.7", + "ruff==0.12.2", "pyrefly==0.36.0", ] diff --git a/packages/readers_extra/pyproject.toml b/packages/readers_extra/pyproject.toml index 0a322926bb..55fcd21f50 100644 --- a/packages/readers_extra/pyproject.toml +++ b/packages/readers_extra/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ dev = [ "pytest>=8.3", "pytest-mock>=3.14.1", - "ruff==0.9.7", + "ruff==0.12.2", "pyrefly==0.36.0", ] diff --git a/pyproject.toml b/pyproject.toml index 8cd6880012..e97ac5f22b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,7 +65,8 @@ dev = [ "jupytext>=1.17.2", "pytest>=8.3", "pytest-mock>=3.14.1", - "ruff==0.9.7", + # If you change the version, also change it in scripts/actions.sh + "ruff==0.12.2", "tensorboard>=2.20.0", "pdbpp>=0.11.7", "pylint==4.0.3", @@ -239,20 +240,33 @@ explicit = true [tool.uv.sources] weathergen-common = { workspace = true } weathergen-evaluate = { workspace = true } -torch = [ - { index = "pytorch-cu124", marker = "sys_platform == 'linux' or sys_platform == 'win32'" }, - { index = "pytorch-cpu", marker = "sys_platform == 'macosx'"}, -] -# This URL was evaluated this way: -# uv run ~/WeatherGenerator-private/hpc/hpc2020/ecmwf/get-flash-atten.sh -# get-flash-atten.sh only supports x86_64 (Dao-AILab has no official aarch64 wheels for -# v2.7.4.post1); the aarch64 build below is a community build from cathalobrien/get-flash-attn, -# matched to the same torch (2.6), CUDA (12), Python (3.12) and ABI (FALSE) as the x86_64 wheel. +weathergen-metrics = { workspace = true } +weathergen-readers-extra = { workspace = true } + + flash-attn = [ - { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.7.4+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, - { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, + { url = "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.9cxx11abiTRUE-cp312-cp312-linux_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" }, + { url = "https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.8.3+cu12torch2.9cxx11abiFALSE-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, + # There is no official support for flash-attn on macos ] +torch = [ +# Explicit pin for GPU + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_aarch64.whl", marker = 'sys_platform == "linux" and platform_machine == "aarch64"', extra="gpu" }, + { url = "https://download-r2.pytorch.org/whl/cu126/torch-2.9.1+cu126-cp312-cp312-manylinux_2_28_x86_64.whl", marker = 'sys_platform == "linux" and platform_machine == "x86_64"', extra="gpu" }, +# Explicit pin for CPU. + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_aarch64.whl",marker = "sys_platform == 'linux' and platform_machine == 'aarch64'",extra = "cpu" }, + { url = "https://download-r2.pytorch.org/whl/cpu/torch-2.9.1+cpu-cp312-cp312-manylinux_2_28_x86_64.whl",marker = "sys_platform == 'linux' and platform_machine == 'x86_64'",extra = "cpu" }, + + # No acceleration for now. + # TODO: explore pytorch + metal backend + { index = "pytorch-cpu", marker = "sys_platform == 'darwin'", extra="gpu"}, +] +# TODO: switch back to this branch, but this is a larger update and the tag name is incorrect +# anemoi-datasets = { git = "https://github.com/ecmwf/anemoi-datasets", tag = "v0.0.2-special-zarr3"} +anemoi-datasets = { git = "https://github.com/tjhunter/anemoi-datasets", tag = "0.0.1"} + + [tool.pytest.ini_options] log_cli = true log_cli_level = "INFO" From ba9b896754cc72f6ca7ecff22859aec6bcbf1125 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Mon, 17 Aug 2026 10:34:25 +0200 Subject: [PATCH 11/15] committing configs --- ...onfig_diffusion_forecast_JEPA_oqkez8qs.yml | 323 +++++++++ ...onfig_diffusion_forecast_JEPA_zs4066ch.yml | 323 +++++++++ .../diffusion/config_diffusion_hsq7hsb6.yml | 349 ++++++++++ .../diffusion/config_diffusion_oliqngmf.yml | 351 ++++++++++ .../diffusion/config_diffusion_skkbxfya.yml | 627 ++++++++++++++++++ .../config_diffusion_xxoplye6_resolved.yml | 248 +++++++ config/evaluate/eval_config_diffusion.yml | 19 +- config/inference_decoder_overlay_dy0jlrmw.yml | 388 +++++++++++ config/inference_decoder_overlay_fqxmlc0v.yml | 330 +++++++++ ...ence_decoder_overlay_ni41n7gy_mapw62ot.yml | 388 +++++++++++ ...ence_decoder_overlay_ni41n7gy_skkbxfya.yml | 389 +++++++++++ config/inference_decoder_overlay_oliqngmf.yml | 410 ++++++++++++ config/inference_decoder_overlay_oqkez8qs.yml | 330 +++++++++ config/inference_decoder_overlay_xxoplye6.yml | 330 +++++++++ ...ence_decoder_overlay_xxoplye6_mapw62ot.yml | 388 +++++++++++ ...ence_decoder_overlay_xxoplye6_skkbxfya.yml | 330 +++++++++ config/runs_plot_train.yml | 60 +- 17 files changed, 5552 insertions(+), 31 deletions(-) create mode 100644 config/diffusion/config_diffusion_forecast_JEPA_oqkez8qs.yml create mode 100644 config/diffusion/config_diffusion_forecast_JEPA_zs4066ch.yml create mode 100644 config/diffusion/config_diffusion_hsq7hsb6.yml create mode 100644 config/diffusion/config_diffusion_oliqngmf.yml create mode 100644 config/diffusion/config_diffusion_skkbxfya.yml create mode 100644 config/diffusion/config_diffusion_xxoplye6_resolved.yml create mode 100644 config/inference_decoder_overlay_dy0jlrmw.yml create mode 100644 config/inference_decoder_overlay_fqxmlc0v.yml create mode 100644 config/inference_decoder_overlay_ni41n7gy_mapw62ot.yml create mode 100644 config/inference_decoder_overlay_ni41n7gy_skkbxfya.yml create mode 100644 config/inference_decoder_overlay_oliqngmf.yml create mode 100644 config/inference_decoder_overlay_oqkez8qs.yml create mode 100644 config/inference_decoder_overlay_xxoplye6.yml create mode 100644 config/inference_decoder_overlay_xxoplye6_mapw62ot.yml create mode 100644 config/inference_decoder_overlay_xxoplye6_skkbxfya.yml diff --git a/config/diffusion/config_diffusion_forecast_JEPA_oqkez8qs.yml b/config/diffusion/config_diffusion_forecast_JEPA_oqkez8qs.yml new file mode 100644 index 0000000000..8ad1f2487c --- /dev/null +++ b/config/diffusion/config_diffusion_forecast_JEPA_oqkez8qs.yml @@ -0,0 +1,323 @@ +# (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. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 1024 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: true +ae_local_num_queries: 1 +ae_local_queries_per_cell: false + +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: true +ae_adapter_with_residual: true +ae_adapter_dropout_rate: 0.1 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 3 +ae_global_num_heads: 16 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: true +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 4 +ae_global_trailing_layer_norm: true + +ae_aggregation_num_blocks: 5 +ae_aggregation_num_heads: 16 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: true +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 4 + +decoder_type: "PerceiverIOCoordConditioning" +pred_adapter_kv: false +pred_self_attention: true +pred_dyadic_dims: false +pred_mlp_adaln: true +num_class_tokens: 0 +num_register_tokens: 64 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +diffusion: {noise_stream: ERA5_in} +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2048 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +fe_impute_latent_noise_std: 0.0 # 1e-4 +fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False +# Diffusion related parameters +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 # only required if noise_distribution is log_normal +p_std: 1.2 # only required if noise_distribution is log_normal + +healpix_level: 5 + +# 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 +rope_2D: True +mlp_type: mlp +use_xsa: False +# mlp_type: swiglu +# use_xsa: True + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-5 +norm_eps: 1e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + + +#below for FE training only +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for DECODER training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for FE and DECODER training +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +load_chkpt: {'run_id': 'oqkez8qs', 'mini_epoch': 84} # n320 + obs sjepa encoder backbone, fine-tuned from zs4o66ch (matches models/oqkez8qs); note: config key is "mini_epoch", not "epoch" (see model_interface.py load_chkpt handling) + + +norm_type: "LayerNorm" +qk_norm_type: "RMSNorm" + +##################################### + +# streams_directory: "./config/streams/era5_1deg_forecasting_d2048/" +# streams_directory: "./config/streams/era5_1deg_forecasting_z500/" +streams_directory: "./config/streams/era5_n320_forecasting/" +streams: ??? + +# type of zarr_store +zarr_store: "zip" # "zarr" for LocalStore, "zip" for ZipStore + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + desc: "" + run_id: ??? + run_history: [] + +# logging frequency in the training loop (in number of batches) +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 250 + log_grad_norms: False + +# parameters for data loading +data_loading : + + num_workers: 8 + rng_seed: ??? + repeat_data_in_mini_epoch : False + + +# config for training +training_config: + + # training_mode: "masking", "student_teacher", "latent_loss" + training_mode: ["student_teacher"] # ["student_teacher", "physical_loss"] + + num_mini_epochs: 128 + samples_per_mini_epoch: 4096 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T18:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 #5e-5 + lr_max: 1e-5 #1e-4 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 512 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.975 + beta2 : 0.9875 + eps : 2e-08 + + losses : { + # "physical": { + # type: LossPhysical, + # weight: 0.0, + # loss_fcts: { + # "mse": {}, + # }, + # target_and_aux_calc: "Physical", + # }, + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + # masking strategy: "random", "healpix", "forecast" + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + + +# validation config; full validation config is merge of training and validation config +validation_config: + + # Noise levels (eta values in standard normal space) at which to evaluate the + # diffusion model during validation. sigma = exp(eta * p_std + p_mean). + # Each value produces a separate validation pass with independently logged metrics. + validation_noise_levels: [1.0, 2.0, 3.0, 4.0] + + samples_per_mini_epoch: 256 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: True + + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null \ No newline at end of file diff --git a/config/diffusion/config_diffusion_forecast_JEPA_zs4066ch.yml b/config/diffusion/config_diffusion_forecast_JEPA_zs4066ch.yml new file mode 100644 index 0000000000..a035939170 --- /dev/null +++ b/config/diffusion/config_diffusion_forecast_JEPA_zs4066ch.yml @@ -0,0 +1,323 @@ +# (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. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 1024 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: true +ae_local_num_queries: 1 +ae_local_queries_per_cell: false + +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: true +ae_adapter_with_residual: true +ae_adapter_dropout_rate: 0.1 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 3 +ae_global_num_heads: 16 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: true +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 4 +ae_global_trailing_layer_norm: true + +ae_aggregation_num_blocks: 5 +ae_aggregation_num_heads: 16 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: true +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 4 + +decoder_type: "PerceiverIOCoordConditioning" +pred_adapter_kv: false +pred_self_attention: true +pred_dyadic_dims: false +pred_mlp_adaln: true +num_class_tokens: 0 +num_register_tokens: 64 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +diffusion: {noise_stream: ERA5_in} +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2048 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +fe_impute_latent_noise_std: 0.0 # 1e-4 +fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False +# Diffusion related parameters +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 # only required if noise_distribution is log_normal +p_std: 1.2 # only required if noise_distribution is log_normal + +healpix_level: 5 + +# 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 +rope_2D: True +mlp_type: mlp +use_xsa: False +# mlp_type: swiglu +# use_xsa: True + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-5 +norm_eps: 1e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + + +#below for FE training only +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for DECODER training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for FE and DECODER training +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +load_chkpt: {'run_id': 'zs4o66ch', 'epoch': 32} # n320 + obs sjepa encoder backbone (matches models/zs4o66ch) + + +norm_type: "LayerNorm" +qk_norm_type: "RMSNorm" + +##################################### + +# streams_directory: "./config/streams/era5_1deg_forecasting_d2048/" +# streams_directory: "./config/streams/era5_1deg_forecasting_z500/" +streams_directory: "./config/streams/era5_n320_forecasting/" +streams: ??? + +# type of zarr_store +zarr_store: "zip" # "zarr" for LocalStore, "zip" for ZipStore + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + desc: "" + run_id: ??? + run_history: [] + +# logging frequency in the training loop (in number of batches) +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 250 + log_grad_norms: False + +# parameters for data loading +data_loading : + + num_workers: 8 + rng_seed: ??? + repeat_data_in_mini_epoch : False + + +# config for training +training_config: + + # training_mode: "masking", "student_teacher", "latent_loss" + training_mode: ["student_teacher"] # ["student_teacher", "physical_loss"] + + num_mini_epochs: 128 + samples_per_mini_epoch: 4096 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T18:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 #5e-5 + lr_max: 1e-5 #1e-4 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 512 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.975 + beta2 : 0.9875 + eps : 2e-08 + + losses : { + # "physical": { + # type: LossPhysical, + # weight: 0.0, + # loss_fcts: { + # "mse": {}, + # }, + # target_and_aux_calc: "Physical", + # }, + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + # masking strategy: "random", "healpix", "forecast" + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + + +# validation config; full validation config is merge of training and validation config +validation_config: + + # Noise levels (eta values in standard normal space) at which to evaluate the + # diffusion model during validation. sigma = exp(eta * p_std + p_mean). + # Each value produces a separate validation pass with independently logged metrics. + validation_noise_levels: [1.0, 2.0, 3.0, 4.0] + + samples_per_mini_epoch: 256 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: True + + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null \ No newline at end of file diff --git a/config/diffusion/config_diffusion_hsq7hsb6.yml b/config/diffusion/config_diffusion_hsq7hsb6.yml new file mode 100644 index 0000000000..ee336e1693 --- /dev/null +++ b/config/diffusion/config_diffusion_hsq7hsb6.yml @@ -0,0 +1,349 @@ +# (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. +# +# Diffusion fine-tune config for hsq7hsb6, built from +# /e/project1/e-ext-2025e01-128/karlbauer1/repos/weathergen-configs/diffusion/config_diffusion_skkbxfya.yml +# (the same pattern skkbxfya used to fine-tune a plain forecast checkpoint into a +# diffusion model), with two changes required for hsq7hsb6 specifically: +# - streams_directory points at the FROZEN weathergen-configs snapshot copied +# into od08us1u's own launch dir (od08us1u is the first run in hsq7hsb6's +# run_history lineage: od08us1u -> ccf5i0sc -> hsq7hsb6), NOT lessig's live +# weathergen-configs repo. Verified against the fully-resolved 'streams' dict +# in model_hsq7hsb6.json: the frozen od08us1u snapshot matches on all 10 +# streams, while lessig's live copy has since dropped 'cos_sza' from +# geoinfo_channels for AVHRR/METOP_ABC_IASI (4 vs the resolved 5), which +# caused a checkpoint-loading size mismatch on encoder.embed_engine.embeds.* +# the first time this config was launched (run rvalsiui, job 1353959). +# - rope_2D is set to False to match hsq7hsb6's actual resolved config +# (model_hsq7hsb6.json), since freeze_modules below freezes encoder/ +# ae_local_engine/ae_global_engine, which must match the positional-encoding +# scheme the checkpoint was trained with. All other architecture fields +# (ae_local_*, ae_global_*, ae_aggregation_*, decoder_type, embed_*, norm_type, +# etc.) already match hsq7hsb6's resolved config, so they are unchanged. +# fe_* / diffusion-specific fields are left at skkbxfya's template values, since +# the forecast engine is not in freeze_modules and is newly trained regardless of +# hsq7hsb6's original (non-diffusion) fe_num_blocks: 16. +# +# fe_diffusion_latent_dim (2560), ae_local_max_tokens_per_cell (128), and +# training_config.learning_rate_scheduling.lr_max (1e-4) are set here to match +# values otherwise passed via --options at launch: launch-slurm.py's MLflow +# upload logs each --config file and the --options-derived config_command_line.yaml +# as separate param batches against the same run_id, and MLflow params are +# immutable, so a --config file literal that differs from a same-key --options +# override crashes the submission (INVALID_PARAMETER_VALUE) before the SLURM job +# is ever created. Keeping the file's values equal to the launch-time overrides +# avoids that crash. training_config.learning_rate_scheduling.num_steps_cooldown +# already matches the launch-time override (32768) so needed no change. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 2048 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: True + +ae_local_num_queries: 1 +ae_local_queries_per_cell: False +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: True +ae_adapter_with_residual: True +ae_adapter_dropout_rate: 0.1 +ae_local_max_tokens_per_cell: 128 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 4 +ae_global_num_heads: 32 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: True +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 2 +ae_global_trailing_layer_norm: False + +ae_aggregation_num_blocks: 0 +ae_aggregation_num_heads: 32 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: True +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 2 + +decoder_type: PerceiverIOCoordConditioning +pred_adapter_kv: False +pred_self_attention: True +pred_dyadic_dims: False +pred_mlp_adaln: True +num_class_tokens: 0 +num_register_tokens: 0 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +diffusion: {noise_stream: ERA5} +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2560 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +fe_impute_latent_noise_std: 0.0 # 1e-4 +fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False +# Diffusion related parameters +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 # only required if noise_distribution is log_normal +p_std: 1.2 # only required if noise_distribution is log_normal + +healpix_level: 5 + +# 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 +# hsq7hsb6 was trained with rope_2D: False (see model_hsq7hsb6.json) -- must match +# since ae_local_engine/ae_global_engine/encoder are frozen below. +rope_2D: False +mlp_type: mlp +use_xsa: False +# mlp_type: swiglu +# use_xsa: True + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-5 +norm_eps: 1e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + +#below for FE training only +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for DECODER training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for FE and DECODER training +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +load_chkpt: {'run_id': 'hsq7hsb6', 'mini_epoch': -1} # hsq7hsb6 2-step/12h forecast fine-tune checkpoint + +norm_type: "LayerNorm" +# qk_norm_type: "RMSNorm" + +##################################### + +streams_directory: "/e/scratch/weatherai/slurm/slurm_weathergen_od08us1u_dir/WeatherGenerator/weathergen-configs/streams_preprint/era5_georing_avhrr_synop_lowres_linear/" +streams: ??? + +zarr_store: "zip" + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + desc: "" + run_id: ??? + run_history: [] + +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 250 + log_grad_norms: False + +data_loading: + num_workers: 10 + num_workers_validation: 4 + rng_seed: ??? + repeat_data_in_mini_epoch: False + +training_config: + training_mode: ["masking", "student_teacher"] + + num_mini_epochs: 128 + samples_per_mini_epoch: 4096 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T00:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 #5e-5 + lr_max: 1e-4 #1e-5 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 32768 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.975 + beta2 : 0.9875 + eps : 2e-08 + + losses : { + # "physical": { + # type: LossPhysical, + # weight: 0.0, + # loss_fcts: { + # "mse": {}, + # }, + # target_and_aux_calc: "Physical", + # }, + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + # masking strategy: "random", "healpix", "forecast" + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + # chunk_size: 4 + + +validation_config: + + # Noise levels (eta values in standard normal space) at which to evaluate the + # diffusion model during validation. sigma = exp(eta * p_std + p_mean). + # Each value produces a separate validation pass with independently logged metrics. + validation_noise_levels: [1.0, 4.0] + + samples_per_mini_epoch: 64 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: False + + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null diff --git a/config/diffusion/config_diffusion_oliqngmf.yml b/config/diffusion/config_diffusion_oliqngmf.yml new file mode 100644 index 0000000000..b4ef577367 --- /dev/null +++ b/config/diffusion/config_diffusion_oliqngmf.yml @@ -0,0 +1,351 @@ +# (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. +# +# Diffusion fine-tune config for oliqngmf, built from config_diffusion_hsq7hsb6.yml +# (same skkbxfya-derived pattern used to turn a plain forecast/masking checkpoint +# into a diffusion model), with two changes required for oliqngmf specifically: +# - streams_directory points at the FROZEN weathergen-configs snapshot copied +# into oliqngmf's own launch dir (oliqngmf's run_history lineage contains only +# itself, so its own launch dir is the earliest available snapshot — same role +# od08us1u played for hsq7hsb6). Verified against the fully-resolved 'streams' +# dict in model_oliqngmf.json: `diff -rq` between this frozen snapshot and +# lessig's live weathergen-configs repo shows zero differences (unlike the +# hsq7hsb6 case, where the live copy had since dropped 'cos_sza' from +# geoinfo_channels and caused a checkpoint-loading size mismatch), so no drift +# risk either way — the frozen path is used to stay consistent with the +# established pattern. +# - rope_2D is set to True to match oliqngmf's actual resolved config +# (model_oliqngmf.json: rope_2D=True, vs hsq7hsb6's False), since freeze_modules +# below freezes encoder/ae_local_engine/ae_global_engine, which must match the +# positional-encoding scheme the checkpoint was trained with. All other +# architecture fields (ae_local_*, ae_global_*, ae_aggregation_*, decoder_type, +# embed_*, norm_type, pred_*, healpix_level, etc.) were checked field-by-field +# against model_oliqngmf.json and match hsq7hsb6's resolved config exactly, so +# they are unchanged from the hsq7hsb6 template. +# fe_* / diffusion-specific fields are left at the template values, since the +# forecast engine is not in freeze_modules and is newly trained regardless of +# oliqngmf's original (non-diffusion) fe_num_blocks: 16. +# +# fe_diffusion_latent_dim (2560), ae_local_max_tokens_per_cell (128), and +# training_config.learning_rate_scheduling.lr_max (1e-4) are set here to match +# values otherwise passed via --options at launch: launch-slurm.py's MLflow +# upload logs each --config file and the --options-derived config_command_line.yaml +# as separate param batches against the same run_id, and MLflow params are +# immutable, so a --config file literal that differs from a same-key --options +# override crashes the submission (INVALID_PARAMETER_VALUE) before the SLURM job +# is ever created. Keeping the file's values equal to the launch-time overrides +# avoids that crash. training_config.learning_rate_scheduling.num_steps_cooldown +# already matches the launch-time override (32768) so needed no change. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 2048 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: True + +ae_local_num_queries: 1 +ae_local_queries_per_cell: False +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: True +ae_adapter_with_residual: True +ae_adapter_dropout_rate: 0.1 +ae_local_max_tokens_per_cell: 128 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 4 +ae_global_num_heads: 32 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: True +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 2 +ae_global_trailing_layer_norm: False + +ae_aggregation_num_blocks: 0 +ae_aggregation_num_heads: 32 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: True +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 2 + +decoder_type: PerceiverIOCoordConditioning +pred_adapter_kv: False +pred_self_attention: True +pred_dyadic_dims: False +pred_mlp_adaln: True +num_class_tokens: 0 +num_register_tokens: 0 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +diffusion: {noise_stream: ERA5} +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2560 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +fe_impute_latent_noise_std: 0.0 # 1e-4 +fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False +# Diffusion related parameters +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 # only required if noise_distribution is log_normal +p_std: 1.2 # only required if noise_distribution is log_normal + +healpix_level: 5 + +# 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 +# oliqngmf was trained with rope_2D: True (see model_oliqngmf.json) -- must match +# since ae_local_engine/ae_global_engine/encoder are frozen below. +rope_2D: True +mlp_type: mlp +use_xsa: False +# mlp_type: swiglu +# use_xsa: True + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-5 +norm_eps: 1e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + +#below for FE training only +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for DECODER training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for FE and DECODER training +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +load_chkpt: {'run_id': 'oliqngmf', 'mini_epoch': -1} # oliqngmf forecast/masking backbone checkpoint + +norm_type: "LayerNorm" +# qk_norm_type: "RMSNorm" + +##################################### + +streams_directory: "/e/scratch/weatherai/slurm/slurm_weathergen_oliqngmf_dir/WeatherGenerator/weathergen-configs/streams_preprint/era5_georing_avhrr_synop_lowres_linear/" +streams: ??? + +zarr_store: "zip" + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + desc: "" + run_id: ??? + run_history: [] + +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 250 + log_grad_norms: False + +data_loading: + num_workers: 10 + num_workers_validation: 4 + rng_seed: ??? + repeat_data_in_mini_epoch: False + +training_config: + training_mode: ["masking", "student_teacher"] + + num_mini_epochs: 128 + samples_per_mini_epoch: 4096 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T00:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 #5e-5 + lr_max: 1e-4 #1e-5 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 32768 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.975 + beta2 : 0.9875 + eps : 2e-08 + + losses : { + # "physical": { + # type: LossPhysical, + # weight: 0.0, + # loss_fcts: { + # "mse": {}, + # }, + # target_and_aux_calc: "Physical", + # }, + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + # masking strategy: "random", "healpix", "forecast" + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + # chunk_size: 4 + + +validation_config: + + # Noise levels (eta values in standard normal space) at which to evaluate the + # diffusion model during validation. sigma = exp(eta * p_std + p_mean). + # Each value produces a separate validation pass with independently logged metrics. + validation_noise_levels: [1.0, 4.0] + + samples_per_mini_epoch: 64 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: False + + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null diff --git a/config/diffusion/config_diffusion_skkbxfya.yml b/config/diffusion/config_diffusion_skkbxfya.yml new file mode 100644 index 0000000000..97793760e3 --- /dev/null +++ b/config/diffusion/config_diffusion_skkbxfya.yml @@ -0,0 +1,627 @@ +# (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. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 2048 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: True + +ae_local_num_queries: 1 +ae_local_queries_per_cell: False +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: True +ae_adapter_with_residual: True +ae_adapter_dropout_rate: 0.1 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 4 +ae_global_num_heads: 32 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: True +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 2 +ae_global_trailing_layer_norm: False + +ae_aggregation_num_blocks: 0 +ae_aggregation_num_heads: 32 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: True +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 2 + +decoder_type: PerceiverIOCoordConditioning +pred_adapter_kv: False +pred_self_attention: True +pred_dyadic_dims: False +pred_mlp_adaln: True +num_class_tokens: 0 +num_register_tokens: 0 + +# number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# one is training an auto-encoder +diffusion: {noise_stream: ERA5} +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2048 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +fe_impute_latent_noise_std: 0.0 # 1e-4 +fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# currently fixed to 1.0 (due to limitations with flex_attention and triton) +forecast_att_dense_rate: 1.0 +with_step_conditioning: True # False +# Diffusion related parameters +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 # only required if noise_distribution is log_normal +p_std: 1.2 # only required if noise_distribution is log_normal + +healpix_level: 5 + +# 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 +rope_2D: True +mlp_type: mlp +use_xsa: False +# mlp_type: swiglu +# use_xsa: True + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-5 +norm_eps: 1e-4 + +latent_noise_kl_weight: 0.0 # 1e-5 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + +#below for FE training only +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for DECODER training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +#below for FE and DECODER training +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +load_chkpt: {'run_id': 'skkbxfya', 'epoch': -1} # Christian's 8-node model + +norm_type: "LayerNorm" +# qk_norm_type: "RMSNorm" + +##################################### + +streams_directory: "/e/project1/e-ext-2025e01-128/karlbauer1/repos/weathergen-configs/diffusion/streams/era5_georing_avhrr_synop_lowres_linear/" +streams: ??? + +zarr_store: "zip" + +general: + + # mutable parameters + istep: 0 + rank: ??? + world_size: ??? + + # local_rank, + # with_ddp, + # data_path_*, + # model_path, + # run_path, + # path_shared_ + + multiprocessing_method: "fork" + + desc: "" + run_id: ??? + run_history: [] + +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 250 + log_grad_norms: False + +data_loading: + num_workers: 10 + num_workers_validation: 4 + rng_seed: ??? + repeat_data_in_mini_epoch: False + +training_config: + training_mode: ["masking", "student_teacher"] + + num_mini_epochs: 128 + samples_per_mini_epoch: 4096 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T00:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 #5e-5 + lr_max: 1e-5 #1e-4 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 32768 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "sqrt" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + # parameters are scaled by number of DDP workers + beta1 : 0.975 + beta2 : 0.9875 + eps : 2e-08 + + losses : { + # "physical": { + # type: LossPhysical, + # weight: 0.0, + # loss_fcts: { + # "mse": {}, + # }, + # target_and_aux_calc: "Physical", + # }, + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + # masking strategy: "random", "healpix", "forecast" + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + # chunk_size: 4 + + +validation_config: + + # Noise levels (eta values in standard normal space) at which to evaluate the + # diffusion model during validation. sigma = exp(eta * p_std + p_mean). + # Each value produces a separate validation pass with independently logged metrics. + validation_noise_levels: [1.0, 4.0] + + samples_per_mini_epoch: 64 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + # whether to track the exponential moving average of weights for validation + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + # parameters for validation samples that are written to disk + output : { + # number of samples that are written + num_samples: 0, + # write samples in normalized model space + normalized_samples: False, + # output streams to write; default all + streams: null, + } + + # run validation before training starts (mainly for model development) + validate_before_training: False + + +# test config; full test config is merge of validation and test config +# test config is used by default when running inference + +# Tags for experiment tracking +# These tags will be logged in MLFlow along with completed runs for train, eval, val +# The tags are free-form, with the following rules: +# - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# - tags should not duplicate existing config entries. +# - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +wgtags: + # The name of the organization of the person running the experiment. + # This may be autofilled in the future. Expected values are lowercase strings + # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" + org: null + # The Github issue corresponding to this run (number such as 1234) + # Github issues are the central point when running experiment and contain + # links to hedgedocs, code branches, pull requests etc. + # It is recommended to associate a run with a Github issue. + issue: null + # The name of the experiment. This is a distinctive codename for the experiment campaign being run. + # This is expected to be the primary tag for comparing experiments in MLFlow, along with the + # issue number. + # Expected values are lowercase strings with no spaces, just underscores: + # Examples: "rollout_ablation_grid" + exp: null + # *** Experiment-specific tags *** + # All extra tags (including lists, dictionaries, etc.) are treated + # as strings by mlflow, so treat all extra tags as simple string key: value pairs. + grid: null + + + +# # (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. + +# embed_orientation: "channels" +# embed_unembed_mode: "block" +# embed_dropout_rate: 0.1 + +# ae_local_dim_embed: 2048 +# ae_local_num_blocks: 2 +# ae_local_num_heads: 16 +# ae_local_dropout_rate: 0.1 +# ae_local_with_qk_lnorm: True + +# ae_local_num_queries: 1 +# ae_local_queries_per_cell: False +# ae_adapter_num_heads: 16 +# ae_adapter_embed: 128 +# ae_adapter_with_qk_lnorm: True +# ae_adapter_with_residual: True +# ae_adapter_dropout_rate: 0.1 + +# ae_global_dim_embed: 2048 +# ae_global_num_blocks: 4 +# ae_global_num_heads: 32 +# ae_global_dropout_rate: 0.1 +# ae_global_with_qk_lnorm: True +# ae_global_att_dense_rate: 1.0 +# ae_global_block_factor: 64 +# ae_global_mlp_hidden_factor: 2 +# ae_global_trailing_layer_norm: False + +# ae_aggregation_num_blocks: 0 +# ae_aggregation_num_heads: 32 +# ae_aggregation_dropout_rate: 0.1 +# ae_aggregation_with_qk_lnorm: True +# ae_aggregation_att_dense_rate: 1.0 +# ae_aggregation_block_factor: 64 +# ae_aggregation_mlp_hidden_factor: 2 + +# decoder_type: PerceiverIOCoordConditioning +# pred_adapter_kv: False +# pred_self_attention: True +# pred_dyadic_dims: False +# pred_mlp_adaln: True +# num_class_tokens: 0 +# num_register_tokens: 0 + +# # number of steps offset applied to first target window; if set to zero and forecast_steps=0 then +# # one is training an auto-encoder +# diffusion: {noise_stream: ERA5} +# fe_num_blocks: 12 +# fe_num_heads: 16 +# fe_dropout_rate: 0.1 +# fe_with_qk_lnorm: True +# fe_diffusion_latent_dim: 2048 # set to an int > ae_global_dim_embed to inflate the diffusion latent space; null = use ae_global_dim_embed +# fe_diffusion_model: True +# fe_diffusion_predict_residual: False +# fe_diffusion_classifier_free_guidance_prob: 0.0 +# fe_diffusion_model_conditioning: "forecast" # options: "date_time", "time", "forecast" +# fe_diffusion_model_conditioning_type: "concatenate" # options: "cross_attn", "concatenate", ..., "ada_ln" +# fe_layer_norm_after_blocks: [11] # Index starts at 0. Thus, [3] adds a LayerNorm after the fourth block +# fe_impute_latent_noise_std: 0.0 # 1e-4 +# fe_impute_latent_diffusion_noise_std: 1e-3 # 1e-4 +# # currently fixed to 1.0 (due to limitations with flex_attention and triton) +# forecast_att_dense_rate: 1.0 +# with_step_conditioning: True # False +# # Diffusion related parameters +# frequency_embedding_dim: 256 +# embedding_dim: 512 +# noise_distribution: &noise_distribution "log_uniform" # "log_uniform" or "log_normal" +# sigma_min: &sigma_min 0.4 +# sigma_max: &sigma_max 50 +# sigma_data: 1.0 +# rho: 7 +# p_mean: 3.5 # only required if noise_distribution is log_normal +# p_std: 1.2 # only required if noise_distribution is log_normal + +# healpix_level: 5 + +# # 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 +# rope_2D: True +# mlp_type: mlp +# use_xsa: False +# # mlp_type: swiglu +# # use_xsa: True + +# with_mixed_precision: True +# with_flash_attention: True +# compile_model: False +# with_fsdp: True +# attention_dtype: bf16 +# mixed_precision_dtype: bf16 +# mlp_norm_eps: 1e-5 +# norm_eps: 1e-4 + +# latent_noise_kl_weight: 0.0 # 1e-5 +# latent_noise_gamma: 2.0 +# latent_noise_saturate_encodings: 5 +# latent_noise_use_additive_noise: False +# latent_noise_deterministic_latents: True + +# #below for FE training only +# freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +# #below for DECODER training only +# # freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*noise_embedder.*|.*fe.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" +# #below for FE and DECODER training +# # freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +# load_chkpt: {'run_id': 'skkbxfya', 'epoch': -1} # Christian's 8-node model + +# norm_type: "LayerNorm" +# # qk_norm_type: "RMSNorm" + +# ##################################### + +# streams_directory: "/e/project1/e-ext-2025e01-128/karlbauer1/repos/weathergen-configs/diffusion/streams/era5_georing_avhrr_synop_lowres_linear/" +# streams: ??? + +# zarr_store: "zip" + +# general: + +# # mutable parameters +# istep: 0 +# rank: ??? +# world_size: ??? + +# # local_rank, +# # with_ddp, +# # data_path_*, +# # model_path, +# # run_path, +# # path_shared_ + +# multiprocessing_method: "fork" + +# desc: "" +# run_id: ??? +# run_history: [] + +# train_logging: +# terminal: 10 +# metrics: 20 +# checkpoint: 250 +# log_grad_norms: False + +# data_loading: +# num_workers: 10 +# num_workers_validation: 6 +# rng_seed: ??? +# repeat_data_in_mini_epoch: False + +# training_config: +# training_mode: ["masking", "student_teacher"] + +# num_mini_epochs: 128 +# samples_per_mini_epoch: 4096 +# shuffle: True + +# start_date: 1979-01-01T00:00 +# end_date: 2022-12-31T00:00 + +# time_window_step: 06:00:00 +# time_window_len: 06:00:00 + +# learning_rate_scheduling : +# lr_start: 1e-6 #5e-5 +# lr_max: 1e-5 #1e-4 +# lr_final_decay: 1e-6 +# lr_final: 0.0 +# num_steps_warmup: 64 +# num_steps_cooldown: 32768 +# policy_warmup: "cosine" +# policy_decay: "constant" +# policy_cooldown: "linear" +# parallel_scaling_policy: "sqrt" + +# optimizer: +# grad_clip: 1.0 +# weight_decay: 0.1 +# log_grad_norms: False +# adamw : +# # parameters are scaled by number of DDP workers +# beta1 : 0.975 +# beta2 : 0.9875 +# eps : 2e-08 + +# losses : { +# # "physical": { +# # type: LossPhysical, +# # weight: 0.0, +# # loss_fcts: { +# # "mse": {}, +# # }, +# # target_and_aux_calc: "Physical", +# # }, +# "latent_diff": { +# type: LossLatentDiffusion, +# weight: 1.0, +# target_and_aux_calc: DiffusionLatentTargetEncoder, +# loss_fcts: { "mse": { }, }, +# } +# } + +# model_input: { +# "forecasting" : { +# # masking strategy: "random", "healpix", "forecast" +# masking_strategy: "forecast", +# masking_strategy_config: { +# diffusion_rn: True, +# noise_distribution: *noise_distribution, +# sigma_min: *sigma_min, +# sigma_max: *sigma_max, +# }, +# num_steps_input: 2, +# num_samples: 1, +# } +# } + +# target_input: { +# "forecasting" : { +# masking_strategy: "forecast", +# masking_strategy_config: { +# diffusion_rn: True, +# noise_distribution: *noise_distribution, +# sigma_min: *sigma_min, +# sigma_max: *sigma_max, +# }, +# num_steps_input: 1, +# num_samples: 1, +# } +# } + +# forecast : +# time_step: 06:00:00 +# num_steps: 1 +# offset: 0 +# policy: "fixed" +# # chunk_size: 4 + + +# validation_config: + +# # Noise levels (eta values in standard normal space) at which to evaluate the +# # diffusion model during validation. sigma = exp(eta * p_std + p_mean). +# # Each value produces a separate validation pass with independently logged metrics. +# validation_noise_levels: [1.0, 2.0, 3.0, 4.0] + +# samples_per_mini_epoch: 256 +# shuffle: True + +# start_date: 2023-10-01T00:00 +# end_date: 2023-12-31T18:00 + +# # whether to track the exponential moving average of weights for validation +# validate_with_ema: +# enabled : True +# ema_ramp_up_ratio: 0.09 +# ema_halflife_in_thousands: 1e-3 + +# # parameters for validation samples that are written to disk +# output : { +# # number of samples that are written +# num_samples: 0, +# # write samples in normalized model space +# normalized_samples: False, +# # output streams to write; default all +# streams: null, +# } + +# # run validation before training starts (mainly for model development) +# validate_before_training: False + + +# # test config; full test config is merge of validation and test config +# # test config is used by default when running inference + +# # Tags for experiment tracking +# # These tags will be logged in MLFlow along with completed runs for train, eval, val +# # The tags are free-form, with the following rules: +# # - tags should be primitive types (strings, numbers, booleans). NO lists or dictionaries +# # - tags should not duplicate existing config entries. +# # - try to reuse existing tags where possible. MLFlow does not like having too many unique tags +# # - do not use long strings in values (less than 20 characters is a good rule of thumb, we may enforce this in the future) +# wgtags: +# # The name of the organization of the person running the experiment. +# # This may be autofilled in the future. Expected values are lowercase strings +# # e.g. "ecmwf", "cmcc", "metnor", "jsc", "escience" +# org: null +# # The Github issue corresponding to this run (number such as 1234) +# # Github issues are the central point when running experiment and contain +# # links to hedgedocs, code branches, pull requests etc. +# # It is recommended to associate a run with a Github issue. +# issue: null +# # The name of the experiment. This is a distinctive codename for the experiment campaign being run. +# # This is expected to be the primary tag for comparing experiments in MLFlow, along with the +# # issue number. +# # Expected values are lowercase strings with no spaces, just underscores: +# # Examples: "rollout_ablation_grid" +# exp: null +# # *** Experiment-specific tags *** +# # All extra tags (including lists, dictionaries, etc.) are treated +# # as strings by mlflow, so treat all extra tags as simple string key: value pairs. +# grid: null diff --git a/config/diffusion/config_diffusion_xxoplye6_resolved.yml b/config/diffusion/config_diffusion_xxoplye6_resolved.yml new file mode 100644 index 0000000000..4ba1e8d4b6 --- /dev/null +++ b/config/diffusion/config_diffusion_xxoplye6_resolved.yml @@ -0,0 +1,248 @@ +# (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. +# +# Reconstructed from the resolved runtime config in +# /e/scratch/weatherai/shared_work/models/xxoplye6/model_xxoplye6.json +# (snapshot at general.istep=2980), not from the base-config file on disk. + +embed_orientation: "channels" +embed_unembed_mode: "block" +embed_dropout_rate: 0.1 + +ae_local_dim_embed: 2048 +ae_local_num_blocks: 2 +ae_local_num_heads: 16 +ae_local_dropout_rate: 0.1 +ae_local_with_qk_lnorm: True + +ae_local_num_queries: 1 +ae_local_queries_per_cell: False +ae_adapter_num_heads: 16 +ae_adapter_embed: 128 +ae_adapter_with_qk_lnorm: True +ae_adapter_with_residual: True +ae_adapter_dropout_rate: 0.1 +ae_local_max_tokens_per_cell: 128 + +ae_global_dim_embed: 2048 +ae_global_num_blocks: 4 +ae_global_num_heads: 32 +ae_global_dropout_rate: 0.1 +ae_global_with_qk_lnorm: True +ae_global_att_dense_rate: 1.0 +ae_global_block_factor: 64 +ae_global_mlp_hidden_factor: 2 +ae_global_trailing_layer_norm: False + +ae_aggregation_num_blocks: 0 +ae_aggregation_num_heads: 32 +ae_aggregation_dropout_rate: 0.1 +ae_aggregation_with_qk_lnorm: True +ae_aggregation_att_dense_rate: 1.0 +ae_aggregation_block_factor: 64 +ae_aggregation_mlp_hidden_factor: 2 + +decoder_type: PerceiverIOCoordConditioning +pred_adapter_kv: False +pred_self_attention: True +pred_dyadic_dims: False +pred_mlp_adaln: True +num_class_tokens: 0 +num_register_tokens: 0 + +diffusion: {noise_stream: ERA5} +fe_num_blocks: 12 +fe_num_heads: 16 +fe_dropout_rate: 0.1 +fe_with_qk_lnorm: True +fe_diffusion_latent_dim: 2560 +fe_diffusion_model: True +fe_diffusion_predict_residual: False +fe_diffusion_classifier_free_guidance_prob: 0.0 +fe_diffusion_model_conditioning: "forecast" +fe_diffusion_model_conditioning_type: "concatenate" +fe_layer_norm_after_blocks: [11] +fe_impute_latent_noise_std: 0.0 +fe_impute_latent_diffusion_noise_std: 1e-3 +forecast_att_dense_rate: 1.0 +with_step_conditioning: True +frequency_embedding_dim: 256 +embedding_dim: 512 +noise_distribution: &noise_distribution "log_uniform" +sigma_min: &sigma_min 0.4 +sigma_max: &sigma_max 50 +sigma_data: 1.0 +rho: 7 +p_mean: 3.5 +p_std: 1.2 + +healpix_level: 5 + +rope_2D: True +mlp_type: mlp +use_xsa: False + +with_mixed_precision: True +with_flash_attention: True +compile_model: False +with_fsdp: True +attention_dtype: bf16 +mixed_precision_dtype: bf16 +mlp_norm_eps: 1e-05 +norm_eps: 1e-04 + +latent_noise_kl_weight: 0.0 +latent_noise_gamma: 2.0 +latent_noise_saturate_encodings: 5 +latent_noise_use_additive_noise: False +latent_noise_deterministic_latents: True + +freeze_modules: ".*latent_pre_norm.*|.*latent_heads.*|.*pred_heads.*|.*target_token_engines.*|.*embed_target_coords.*|.*encoder.*|.*StreamEmbedder_ERA5.*|.*embed_engine.*|.*ae_local_engine.*|.*ae_local_global_engine.*|.*ae_global_engine.*" + +load_chkpt: {'run_id': 'skkbxfya', 'epoch': -1} + +norm_type: "LayerNorm" + +##################################### + +streams_directory: "/e/project1/e-ext-2025e01-128/karlbauer1/repos/weathergen-configs/diffusion/streams/era5_georing_avhrr_synop_lowres_linear/" +streams: ??? + +zarr_store: "zip" + +general: + istep: 2980 + rank: ??? + world_size: ??? + + multiprocessing_method: "fork" + + desc: "" + run_id: xxoplye6 + run_history: [[xxoplye6, 2400], [xxoplye6, 2980]] + +train_logging: + terminal: 10 + metrics: 20 + checkpoint: 50 + log_grad_norms: False + +data_loading: + num_workers: 10 + num_workers_validation: 6 + rng_seed: 1786564175 + repeat_data_in_mini_epoch: False + +training_config: + training_mode: ["masking", "student_teacher"] + + num_mini_epochs: 64 + samples_per_mini_epoch: 16384 + shuffle: True + + start_date: 1979-01-01T00:00 + end_date: 2022-12-31T00:00 + + time_window_step: 06:00:00 + time_window_len: 06:00:00 + + learning_rate_scheduling : + lr_start: 1e-6 + lr_max: 2e-4 + lr_final_decay: 1e-6 + lr_final: 0.0 + num_steps_warmup: 64 + num_steps_cooldown: 32768 + policy_warmup: "cosine" + policy_decay: "constant" + policy_cooldown: "linear" + parallel_scaling_policy: "constant" + + optimizer: + grad_clip: 1.0 + weight_decay: 0.1 + log_grad_norms: False + adamw : + beta1 : 0.99 + beta2 : 0.999 + eps : 2e-08 + + losses : { + "latent_diff": { + type: LossLatentDiffusion, + weight: 1.0, + target_and_aux_calc: DiffusionLatentTargetEncoder, + loss_fcts: { "mse": { }, }, + } + } + + model_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 2, + num_samples: 1, + } + } + + target_input: { + "forecasting" : { + masking_strategy: "forecast", + masking_strategy_config: { + diffusion_rn: True, + noise_distribution: *noise_distribution, + sigma_min: *sigma_min, + sigma_max: *sigma_max, + }, + num_steps_input: 1, + num_samples: 1, + } + } + + forecast : + time_step: 06:00:00 + num_steps: 1 + offset: 0 + policy: "fixed" + + +validation_config: + + validation_noise_levels: [1.0, 2.0, 3.0, 4.0] + + samples_per_mini_epoch: 256 + shuffle: True + + start_date: 2023-10-01T00:00 + end_date: 2023-12-31T18:00 + + validate_with_ema: + enabled : True + ema_ramp_up_ratio: 0.09 + ema_halflife_in_thousands: 1e-3 + + output : { + num_samples: 0, + normalized_samples: False, + streams: null, + } + + validate_before_training: False + + +wgtags: + org: null + issue: null + exp: null + grid: null diff --git a/config/evaluate/eval_config_diffusion.yml b/config/evaluate/eval_config_diffusion.yml index 15a567af25..c59c1a010a 100644 --- a/config/evaluate/eval_config_diffusion.yml +++ b/config/evaluate/eval_config_diffusion.yml @@ -18,8 +18,8 @@ # vmax: 40 evaluation: - metrics : ["rmse", "mae"] - regions: ["global", "nhem"] + metrics : ["psd"] + regions: ["global"] summary_plots : true ratio_plots : false heat_maps : false @@ -37,13 +37,13 @@ evaluation: default_streams: ERA5: - channels: ["2t", "10u"] #, "10v", "z_500", "t_850", "u_850", "v_850", "q_850", ] + channels: ["2t", "q_850"] #, "10v", "z_500", "t_850", "u_850", "v_850", "q_850", ] evaluation: forecast_step: "all" sample: "all" ensemble: "all" #supported: "all", "mean", [0,1,2] plotting: - sample: [1] + sample: [0] forecast_step: [1] #supported: "all", [1,2,3,...], "1-50" (equivalent of [1,2,3,...50]) ensemble: [0] #supported: "all", "mean", [0,1,2] plot_maps: true @@ -53,9 +53,14 @@ default_streams: run_ids : - kuia5xr0: - label: "debugging model g0vdqua7" - results_base_dir : "../results/" + ri8gqyup: + label: "forecast finetuned deterministic (skkbxfya -> dy0jlrmw)" + + # kvp6fds7: + # label: "forecast finetuned deterministic (skkbxfya -> mapw62ot)" + iy3niftw: + label: "decoder finetuned diffusion (skkbxfya -> ni41n7gy -> mapw62ot)" + # results_base_dir : "../results/" #NEW: if "streams" is not specified, the default streams are used diff --git a/config/inference_decoder_overlay_dy0jlrmw.yml b/config/inference_decoder_overlay_dy0jlrmw.yml new file mode 100644 index 0000000000..9b5dbd0fcf --- /dev/null +++ b/config/inference_decoder_overlay_dy0jlrmw.yml @@ -0,0 +1,388 @@ +# Decoder overlay: run ni41n7gy with the physical decoder of mapw62ot. +# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md. +# Decoded stream: ERA5. The block below is copied verbatim from the backbone's +# model config, derived channel lists included: nothing recomputes them for an overlay. +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target: + - 10u + - 10v + - 2d + - 2t + - lcc + - mcc + - hcc + - tcc + - msl + - skt + - sp + - ssrd + - strd + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 271040 + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +# load_decoder_chkpt: +# run_id: dy0jlrmw +# mini_epoch: -1 diff --git a/config/inference_decoder_overlay_fqxmlc0v.yml b/config/inference_decoder_overlay_fqxmlc0v.yml new file mode 100644 index 0000000000..eaddcfd943 --- /dev/null +++ b/config/inference_decoder_overlay_fqxmlc0v.yml @@ -0,0 +1,330 @@ +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target_exclude: + - z + - w_10 + - w_50 + - w_100 + - w_150 + - w_200 + - w_250 + - w_300 + - w_400 + - w_500 + - w_600 + - w_700 + - w_850 + - w_925 + - w_1000 + - slor + - sdor + - tcw + - cp + - tp + - q_50 + - q_100 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + reconstruct: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + training_mode: + - masking + - student_teacher + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: oqkez8qs + mini_epoch: 84 diff --git a/config/inference_decoder_overlay_ni41n7gy_mapw62ot.yml b/config/inference_decoder_overlay_ni41n7gy_mapw62ot.yml new file mode 100644 index 0000000000..aed3571042 --- /dev/null +++ b/config/inference_decoder_overlay_ni41n7gy_mapw62ot.yml @@ -0,0 +1,388 @@ +# Decoder overlay: run ni41n7gy with the physical decoder of mapw62ot. +# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md. +# Decoded stream: ERA5. The block below is copied verbatim from the backbone's +# model config, derived channel lists included: nothing recomputes them for an overlay. +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target: + - 10u + - 10v + - 2d + - 2t + - lcc + - mcc + - hcc + - tcc + - msl + - skt + - sp + - ssrd + - strd + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 271040 + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: mapw62ot + mini_epoch: -1 diff --git a/config/inference_decoder_overlay_ni41n7gy_skkbxfya.yml b/config/inference_decoder_overlay_ni41n7gy_skkbxfya.yml new file mode 100644 index 0000000000..7e8f68d99c --- /dev/null +++ b/config/inference_decoder_overlay_ni41n7gy_skkbxfya.yml @@ -0,0 +1,389 @@ +# Decoder overlay: run ni41n7gy with the physical decoder of skkbxfya. +# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md. +# Decoded stream: ERA5. The block below is copied verbatim from the backbone's +# model config, derived channel lists included: nothing recomputes them for an overlay. +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target: + - 10u + - 10v + - 2d + - 2t + - lcc + - mcc + - hcc + - tcc + - msl + - skt + - sp + - ssrd + - strd + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: skkbxfya + mini_epoch: 84 diff --git a/config/inference_decoder_overlay_oliqngmf.yml b/config/inference_decoder_overlay_oliqngmf.yml new file mode 100644 index 0000000000..38ddbd432b --- /dev/null +++ b/config/inference_decoder_overlay_oliqngmf.yml @@ -0,0 +1,410 @@ +# Decoder overlay: run huwssly8 with the physical decoder of oliqngmf. +# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md. +# Decoded stream: ERA5. The block below is copied verbatim from the backbone's +# model config, derived channel lists included: nothing recomputes them for an overlay. +streams: + METOP_ABC_IASI: + reconstruct: false + AVHRR: + reconstruct: false + METEOSAT_SEVIRI_IR: + reconstruct: false + GOES_ABI_IR: + reconstruct: false + HIMAWARI_AHI_IR: + reconstruct: false + GOES_ABI_VIS: + reconstruct: false + HIMAWARI_AHI_VIS: + reconstruct: false + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + # Injected manually (not in oliqngmf's own resolved config, same gotcha hit on hsq7hsb6/ + # skkbxfya): without this, inference produces hourly (6x) rather than 6-hourly physical + # output, and spectra come out "not available" ("Physical PSD skipped"). + frequency: 06:00:00 + source: [] + target: + - 10u + - 10v + - 2d + - 2t + - lcc + - mcc + - hcc + - tcc + - msl + - skt + - sp + - ssrd + - strd + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + data_paths: + - /e/data1/slmet/ml_training + - /p/project1/hclimrep/koldunovn1 + - /p/scratch/hclimrep/shared/weather_generator_data/icon_art + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 +validation_config: + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: oliqngmf + mini_epoch: -1 diff --git a/config/inference_decoder_overlay_oqkez8qs.yml b/config/inference_decoder_overlay_oqkez8qs.yml new file mode 100644 index 0000000000..eaddcfd943 --- /dev/null +++ b/config/inference_decoder_overlay_oqkez8qs.yml @@ -0,0 +1,330 @@ +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target_exclude: + - z + - w_10 + - w_50 + - w_100 + - w_150 + - w_200 + - w_250 + - w_300 + - w_400 + - w_500 + - w_600 + - w_700 + - w_850 + - w_925 + - w_1000 + - slor + - sdor + - tcw + - cp + - tp + - q_50 + - q_100 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + reconstruct: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + training_mode: + - masking + - student_teacher + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: oqkez8qs + mini_epoch: 84 diff --git a/config/inference_decoder_overlay_xxoplye6.yml b/config/inference_decoder_overlay_xxoplye6.yml new file mode 100644 index 0000000000..5679b1ba08 --- /dev/null +++ b/config/inference_decoder_overlay_xxoplye6.yml @@ -0,0 +1,330 @@ +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target_exclude: + - z + - w_10 + - w_50 + - w_100 + - w_150 + - w_200 + - w_250 + - w_300 + - w_400 + - w_500 + - w_600 + - w_700 + - w_850 + - w_925 + - w_1000 + - slor + - sdor + - tcw + - cp + - tp + - q_50 + - q_100 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + reconstruct: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + training_mode: + - masking + - student_teacher + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +# load_decoder_chkpt: +# run_id: skkbxfya +# mini_epoch: 84 diff --git a/config/inference_decoder_overlay_xxoplye6_mapw62ot.yml b/config/inference_decoder_overlay_xxoplye6_mapw62ot.yml new file mode 100644 index 0000000000..2f8c9d5cd5 --- /dev/null +++ b/config/inference_decoder_overlay_xxoplye6_mapw62ot.yml @@ -0,0 +1,388 @@ +# Decoder overlay: run xxoplye6 with the physical decoder of mapw62ot. +# Generated by make_decoder_overlay.py -- see the skill's references/decoder-overlay.md. +# Decoded stream: ERA5. The block below is copied verbatim from the backbone's +# model config, derived channel lists included: nothing recomputes them for an overlay. +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target: + - 10u + - 10v + - 2d + - 2t + - lcc + - mcc + - hcc + - tcc + - msl + - skt + - sp + - ssrd + - strd + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 271040 + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: mapw62ot + mini_epoch: -1 diff --git a/config/inference_decoder_overlay_xxoplye6_skkbxfya.yml b/config/inference_decoder_overlay_xxoplye6_skkbxfya.yml new file mode 100644 index 0000000000..341892d5b9 --- /dev/null +++ b/config/inference_decoder_overlay_xxoplye6_skkbxfya.yml @@ -0,0 +1,330 @@ +streams: + ERA5: + type: anemoi + filenames: + - aifs-ea-an-oper-0001-mars-n320-1979-2024-1h-v2-with-era51.zarr + stream_id: 1 + source: [] + target_exclude: + - z + - w_10 + - w_50 + - w_100 + - w_150 + - w_200 + - w_250 + - w_300 + - w_400 + - w_500 + - w_600 + - w_700 + - w_850 + - w_925 + - w_1000 + - slor + - sdor + - tcw + - cp + - tp + - q_50 + - q_100 + geoinfo_channels: + - z + - lsm + - slor + - sdor + - insolation + - cos_local_time + - sin_local_time + - cos_julian_day + - sin_julian_day + - cos_latitude + - sin_latitude + - cos_longitude + - sin_longitude + loss_weight: 1.0 + location_weight: cosine_latitude + token_size: 32 + tokenize_spacetime: false + max_num_targets: 542080 + diagnostic: true + reconstruct: true + embed: + net: transformer + num_tokens: 1 + num_heads: 4 + dim_embed: 512 + num_blocks: 2 + embed_target_coords: + net: linear + dim_embed: 512 + target_readout: + num_layers: 4 + num_heads: 4 + pred_head: + ens_size: 1 + num_layers: 1 + name: ERA5 + train_source_channels: [] + train_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + target_channel_weights: + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + - 1.0 + val_source_channels: [] + val_target_channels: + - 10u + - 10v + - 2d + - 2t + - hcc + - lcc + - mcc + - msl + - q_10 + - q_1000 + - q_150 + - q_200 + - q_250 + - q_300 + - q_400 + - q_500 + - q_600 + - q_700 + - q_850 + - q_925 + - skt + - sp + - ssrd + - strd + - t_10 + - t_100 + - t_1000 + - t_150 + - t_200 + - t_250 + - t_300 + - t_400 + - t_50 + - t_500 + - t_600 + - t_700 + - t_850 + - t_925 + - tcc + - u_10 + - u_100 + - u_1000 + - u_150 + - u_200 + - u_250 + - u_300 + - u_400 + - u_50 + - u_500 + - u_600 + - u_700 + - u_850 + - u_925 + - v_10 + - v_100 + - v_1000 + - v_150 + - v_200 + - v_250 + - v_300 + - v_400 + - v_50 + - v_500 + - v_600 + - v_700 + - v_850 + - v_925 + - z_10 + - z_100 + - z_1000 + - z_150 + - z_200 + - z_250 + - z_300 + - z_400 + - z_50 + - z_500 + - z_600 + - z_700 + - z_850 + - z_925 + frequency: 06:00:00 +validation_config: + training_mode: + - masking + - student_teacher + losses: + physical: + type: LossPhysical + weight: 0.0 + target_and_aux_calc: Physical + loss_fcts: + mse: {} +load_decoder_chkpt: + run_id: skkbxfya + mini_epoch: 84 diff --git a/config/runs_plot_train.yml b/config/runs_plot_train.yml index 04791a3a54..96f166605e 100644 --- a/config/runs_plot_train.yml +++ b/config/runs_plot_train.yml @@ -3,30 +3,42 @@ train : plot : - perrjey5 : - slurm_id : 0 - description: "N320-obs, with absolute coords, larger decoder, time_window_step=6h" - c71eo6pu : - slurm_id : 0 - description: "perrjey5 contd : N320-obs, with absolute coords, larger decoder, time_window_step=6h" - xc1n4mw3 : - slurm_id : 0 - description: "develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml" - kdl0orun : - slurm_id : 0 - description: "develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml, streams_prepint, linear" - e4327u4h : - slurm_id : 0 - description: "develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml, streams_prepint, transformer" - d5nrjup6 : - slurm_id : 0 - description: "develop" - iqbifydo : - slurm_id : 0 - description: "develop-ssl-diffusion-v1" - fsvwafcl : - slurm_id : 0 - description: "develop-ssl-diffusion-v1, FSDP" + huwssly8: + slurm_id : 0 + description: "new MTM run diff" + fqxmlc0v: + slurm_id : 0 + description: "old MTM run diff" + oliqngmf : + slurm_id : 0 + description: "new MTM run" + hsq7hsb6: + slurm_id : 0 + description: "old MTM run" + # perrjey5 : + # slurm_id : 0 + # description: "N320-obs, with absolute coords, larger decoder, time_window_step=6h" + # c71eo6pu : + # slurm_id : 0 + # description: "perrjey5 contd : N320-obs, with absolute coords, larger decoder, time_window_step=6h" + # xc1n4mw3 : + # slurm_id : 0 + # description: "develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml" + # kdl0orun : + # slurm_id : 0 + # description: "develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml, streams_prepint, linear" + # e4327u4h : + # slurm_id : 0 + # description: "develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml, streams_prepint, transformer" + # d5nrjup6 : + # slurm_id : 0 + # description: "develop" + # iqbifydo : + # slurm_id : 0 + # description: "develop-ssl-diffusion-v1" + # fsvwafcl : + # slurm_id : 0 + # description: "develop-ssl-diffusion-v1, FSDP" # # develop-ssl-diffusion-v1, config_forecasting_n320_obs.yml From 8a9587e904902500ed7eff46aeea01738b92ce4e Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Mon, 17 Aug 2026 17:25:01 +0200 Subject: [PATCH 12/15] add functionality for spectra overlay --- config/evaluate/eval_config_diffusion.yml | 6 +- .../evaluate/plotting/plot_utils.py | 55 ++++-- .../src/weathergen/evaluate/scores/score.py | 182 ++++++++++++++++++ 3 files changed, 224 insertions(+), 19 deletions(-) diff --git a/config/evaluate/eval_config_diffusion.yml b/config/evaluate/eval_config_diffusion.yml index c59c1a010a..9b324ea03c 100644 --- a/config/evaluate/eval_config_diffusion.yml +++ b/config/evaluate/eval_config_diffusion.yml @@ -18,7 +18,7 @@ # vmax: 40 evaluation: - metrics : ["psd"] + metrics : ["psd","rmse"] regions: ["global"] summary_plots : true ratio_plots : false @@ -53,12 +53,12 @@ default_streams: run_ids : - ri8gqyup: + st6x4i29: label: "forecast finetuned deterministic (skkbxfya -> dy0jlrmw)" # kvp6fds7: # label: "forecast finetuned deterministic (skkbxfya -> mapw62ot)" - iy3niftw: + m3ezilrq: label: "decoder finetuned diffusion (skkbxfya -> ni41n7gy -> mapw62ot)" # results_base_dir : "../results/" #NEW: if "streams" is not specified, the default streams are used diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py index 7e52591c77..d8fcda5047 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py @@ -800,12 +800,18 @@ def psd_plot_metric_region( PSD curves (frequencies, target PSD, prediction PSD) are stored in ``score.attrs`` by ``Scores.calc_psd`` and read back here. + + For a given forecast step, all runs are overlaid on one plot. Evolution + plots (across forecast steps) remain one-per-run. """ streams_set = collect_streams(runs) channels_set = collect_channels(scores_dict, metric, region, runs) for stream in streams_set: for ch in channels_set: + # First pass: gather each run's per-fstep PSD data for this stream/channel. + run_fstep_datasets: dict[str, dict] = {} + run_labels: dict[str, str] = {} for run_id, data in scores_dict[metric][region].get(stream, {}).items(): if ch not in np.atleast_1d(data.channel.values): continue @@ -819,8 +825,6 @@ def psd_plot_metric_region( _logger.warning(f"PSD attrs missing for {run_id}/{stream}/{ch}. Skipping.") continue - label = runs[run_id].get("label", run_id) - per_fstep_datasets = {} for fstep in attr_fsteps: psd_datasets = _extract_psd_attrs(data_ch, fstep, ch) @@ -828,20 +832,39 @@ def psd_plot_metric_region( continue per_fstep_datasets[fstep] = psd_datasets[0] - method_tag = psd_datasets[0].get("psd_method", "sht") - name = create_filename( - prefix=[metric, method_tag, region], - middle=[run_id], - suffix=[stream, ch, f"fstep{fstep}"], - ) - plotter.psd_plot( - psd_datasets, - [label], - tag=name, - variable=ch, - forecast_step=str(fstep), - ) + if not per_fstep_datasets: + continue + + run_fstep_datasets[run_id] = per_fstep_datasets + run_labels[run_id] = runs[run_id].get("label", run_id) + + if not run_fstep_datasets: + continue + + # Second pass: one combined plot per forecast step, overlaying every run + # that has data for it. + all_fsteps = sorted({fstep for d in run_fstep_datasets.values() for fstep in d}) + for fstep in all_fsteps: + run_ids = [rid for rid, d in run_fstep_datasets.items() if fstep in d] + psd_datasets = [run_fstep_datasets[rid][fstep] for rid in run_ids] + labels = [run_labels[rid] for rid in run_ids] + + method_tag = psd_datasets[0].get("psd_method", "sht") + name = create_filename( + prefix=[metric, method_tag, region], + middle=run_ids, + suffix=[stream, ch, f"fstep{fstep}"], + ) + plotter.psd_plot( + psd_datasets, + labels, + tag=name, + variable=ch, + forecast_step=str(fstep), + ) + # Third pass: per-run evolution plot (unchanged — one run per plot). + for run_id, per_fstep_datasets in run_fstep_datasets.items(): if len(per_fstep_datasets) >= 2: method_tag = next(iter(per_fstep_datasets.values())).get( "psd_method", "sht" @@ -855,7 +878,7 @@ def psd_plot_metric_region( per_fstep_datasets, tag=evo_name, variable=ch, - label=label, + label=run_labels[run_id], ) _logger.info(f"PSD plots saved successfully into: {plotter.out_plot_dir_psd}") diff --git a/packages/evaluate/src/weathergen/evaluate/scores/score.py b/packages/evaluate/src/weathergen/evaluate/scores/score.py index afee96d9a3..daff06334f 100755 --- a/packages/evaluate/src/weathergen/evaluate/scores/score.py +++ b/packages/evaluate/src/weathergen/evaluate/scores/score.py @@ -16,6 +16,7 @@ import xarray as xr from scipy.spatial import cKDTree +from weathergen.evaluate.scores.psd import compute_psd_score, detect_grid_type from weathergen.evaluate.scores.score_utils import to_list # from common.io import MockIO @@ -198,6 +199,7 @@ def __init__( "seeps": self.calc_seeps, "qq_analysis": self.calc_quantiles, "nse": self.calc_nse, + "psd": self.calc_psd, } self.prob_metrics_dict = { "ssr": self.calc_ssr, @@ -1807,3 +1809,183 @@ def calc_quantiles( _logger.info(f"Q-Q analysis completed with {len(overall_qq_score.attrs)} attributes") return overall_qq_score + + def calc_psd( + self, + p: xr.DataArray, + gt: xr.DataArray, + psd_method: str = "sht", + psd_sht_truncation: int | None = None, + lat_range: tuple[float, float] = (-60.0, 60.0), + ) -> xr.DataArray: + """Compute power spectral density for prediction and ground truth. + + Returns a scalar summary score (log-spectral MSE) and stores the full + PSD curves in ``.attrs`` for plotting downstream. + + Parameters + ---------- + p: xr.DataArray + Forecast data array + gt: xr.DataArray + Ground truth data array + psd_method: str + Method to compute the PSD. Options: 'sht' (spherical harmonic transform), + 'fft' (2D Fourier transform) + psd_sht_truncation: int | None + Maximum spherical harmonic degree for truncation. If None, no truncation is applied. + lat_range: tuple[float, float] + Latitude range (min, max) to include in PSD calculation. Default is (-60, + 60) degrees. + + Returns + ------- + xr.DataArray + Power spectral density score (log-spectral MSE) averaged over aggregation dimensions. + + """ + if self._agg_dims is None: + raise ValueError("Cannot calculate PSD without aggregation dimensions.") + if len(self._agg_dims) != 1: + raise ValueError( + f"PSD expects exactly one spatial aggregation dimension, " + f"got agg_dims={self._agg_dims}." + ) + spatial_dim = self._agg_dims[0] + if spatial_dim not in gt.dims: + raise ValueError( + f"Spatial dimension '{spatial_dim}' not found in dims {list(gt.dims)}." + ) + + # PSD requires a spatial dimension with lat/lon coords (e.g. "ipoint"). + # If the aggregation dim is "sample" or "ens" (e.g. from score map pipeline), + # PSD is not applicable — return NaN gracefully. + if spatial_dim in ("sample", "ens"): + _logger.debug(f"PSD: aggregation dim is '{spatial_dim}' (not spatial). Skipping.") + return xr.DataArray(np.nan) + + n_points = gt.sizes[spatial_dim] + nlat, lats, lons = self._get_psd_grid_info(gt, spatial_dim) + + if psd_method == "fft" and (lats is None or lons is None): + raise ValueError(f"PSD method 'fft' requires lat/lon coords on '{spatial_dim}'.") + + # Detect grid type once for the entire stream (avoid repeated detection per channel) + grid_type = None + if psd_method == "sht" and lats is not None and lons is not None: + grid_type = detect_grid_type(lats, lons, n_points) + + psd_kwargs = dict( + lats=lats, + lons=lons, + nlat=nlat, + n_points=n_points, + psd_method=psd_method, + psd_sht_truncation=psd_sht_truncation, + lat_range=lat_range, + grid_type=grid_type, + ) + + # Dims to preserve (e.g. channel) vs batch dims (sample, ens) + other_dims = [d for d in gt.dims if d != spatial_dim] + preserve_dims = [d for d in other_dims if d not in ("sample", "ens")] + + if not preserve_dims: + gt_np, p_np = self._stack_for_psd(gt, p, spatial_dim, n_points) + slice_score, slice_attrs = compute_psd_score(gt=gt_np, p=p_np, **psd_kwargs) + score = xr.DataArray(slice_score) + score.attrs.update(slice_attrs) + score.attrs["psd_method"] = psd_method + return score + + # Iterate over preserved dims (typically per channel) + shape = tuple(gt.sizes[d] for d in preserve_dims) + score_values = np.empty(shape) + all_attrs: dict = {} + + for idx in np.ndindex(*shape): + sel = dict(zip(preserve_dims, idx, strict=False)) + gt_slice = gt.isel(**sel) + p_slice = p.isel(**sel) + gt_np, p_np = self._stack_for_psd(gt_slice, p_slice, spatial_dim, n_points) + + slice_score, slice_attrs = compute_psd_score(gt=gt_np, p=p_np, **psd_kwargs) + score_values[idx] = slice_score + + key = "_".join( + str(gt.coords[d].values[i]) if d in gt.coords else str(i) for d, i in sel.items() + ) + for k, v in slice_attrs.items(): + all_attrs[f"{key}/{k}"] = v + + coords = {d: gt.coords[d] for d in preserve_dims if d in gt.coords} + score = xr.DataArray(score_values, dims=preserve_dims, coords=coords) + all_attrs["psd_method"] = psd_method + all_attrs["preserve_dims"] = preserve_dims + score.attrs.update(all_attrs) + return score + + @staticmethod + def _get_psd_grid_info( + gt: xr.DataArray, spatial_dim: str + ) -> tuple[int | None, np.typing.NDArray | None, np.typing.NDArray | None]: + """ + Extract nlat, lats, lons from ground-truth coords. + + Parameters + ---------- + gt: xr.DataArray + Ground truth data array with lat/lon coordinates. + spatial_dim: str + Name of the spatial dimension along which to compute the PSD. + Returns + ------- + nlat: int | None + Number of latitude points, or None if lat/lon coords are not found. + lats: np.typing.NDArray | None + Latitude values, or None if lat/lon coords are not found. + lons: np.typing.NDArray | None + Longitude values, or None if lat/lon coords are not found. + + """ + if "lat" in gt.coords and "lon" in gt.coords: + if gt.coords["lat"].dims == (spatial_dim,) and gt.coords["lon"].dims == (spatial_dim,): + lats = gt.coords["lat"].values + lons = gt.coords["lon"].values + return len(np.unique(lats)), lats, lons + raise ValueError(f"PSD requires lat/lon coords on spatial dimension '{spatial_dim}'.") + + @staticmethod + def _stack_for_psd( + gt: xr.DataArray, p: xr.DataArray, spatial_dim: str, n_points: int + ) -> tuple[np.typing.NDArray, np.typing.NDArray]: + """ + Reshape data to (n_batch, n_points) for PSD computation. + + Parameters + ---------- + gt: xr.DataArray + Ground truth data array. + p: xr.DataArray + Forecast data array. + spatial_dim: str + Name of the spatial dimension along which to compute the PSD. + n_points: int + Number of points along the spatial dimension. + Returns + ------- + gt_np: np.typing.NDArray + Reshaped ground truth data of shape (n_batch, n_points). + p_np: np.typing.NDArray + Reshaped forecast data of shape (n_batch, n_points). + """ + non_spatial = [d for d in gt.dims if d != spatial_dim] + if non_spatial: + gt_np = gt.transpose(*non_spatial, spatial_dim).values.reshape(-1, n_points) + p_np = p.transpose( + *[d for d in p.dims if d != spatial_dim], spatial_dim + ).values.reshape(-1, n_points) + else: + gt_np = gt.values.reshape(1, -1) + p_np = p.values.reshape(1, -1) + return gt_np, p_np From 58b809c7257c13df49b1c63c07a08a978d3e1381 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Tue, 18 Aug 2026 18:09:25 +0200 Subject: [PATCH 13/15] cusotm plotting upgrade --- .../evaluate/plotting/line_plots.py | 270 +++++++++++++++--- .../evaluate/plotting/plot_orchestration.py | 3 + .../evaluate/plotting/plot_utils.py | 154 ++++++++-- 3 files changed, 366 insertions(+), 61 deletions(-) diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py b/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py index d132f90c29..f6607a11f0 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py @@ -18,7 +18,10 @@ import numpy as np import seaborn as sns import xarray as xr +from matplotlib.colors import TwoSlopeNorm from matplotlib.lines import Line2D +from numpy.typing import NDArray +from PIL import Image from weathergen.evaluate.plotting.plot_utils import ( align_labels, @@ -50,6 +53,9 @@ def __init__(self, plotter_cfg: dict, output_basedir: str | Path): - "std": plot mean +/- standard deviation - "minmax": plot min-max range - "members": plot individual ensemble members + - fps: frames per second for the PSD forecast-step animation (default 2) + - psd_show_ratio: if False, PSD plots are single-panel (spectra only) + - psd_animation: if False, no per-run PSD gif over forecast steps is written output_basedir: Base directory under which the plots will be saved. Expected scheme `/`. @@ -62,6 +68,9 @@ def __init__(self, plotter_cfg: dict, output_basedir: str | Path): self.add_grid = plotter_cfg.get("add_grid") self.plot_ensemble = plotter_cfg.get("plot_ensemble", False) self.baseline = plotter_cfg.get("baseline") + self.fps = plotter_cfg.get("fps", 2) + self.psd_show_ratio = plotter_cfg.get("psd_show_ratio", True) + self.psd_animate = plotter_cfg.get("psd_animation", True) self.out_plot_dir_lines = Path(output_basedir) / "line_plots" self.out_plot_dir_ratio = Path(output_basedir) / "ratio_plots" self.out_plot_dir_psd = Path(output_basedir) / "psd_plots" @@ -708,6 +717,45 @@ def heat_map( # PSD summary plot # ------------------------------------------------------------------ + def _psd_axes( + self, default_size: tuple[float, float] + ) -> tuple[plt.Figure, plt.Axes, plt.Axes | None]: + """Create the PSD figure: spectra panel + pred/target ratio panel, or spectra only. + + Returns ``(fig, ax_spec, ax_ratio)`` where ``ax_ratio`` is ``None`` when + ``self.psd_show_ratio`` is False. + """ + size = self.fig_size or default_size + if self.psd_show_ratio: + fig, (ax_spec, ax_ratio) = plt.subplots( + 2, + 1, + figsize=size, + gridspec_kw={"height_ratios": [2, 1], "hspace": 0.08}, + ) + return fig, ax_spec, ax_ratio + # Single panel: drop the space the ratio panel used to occupy. + fig, ax_spec = plt.subplots(1, 1, figsize=(size[0], size[1] * 0.7)) + return fig, ax_spec, None + + def _finish_psd_axes( + self, + ax_spec: plt.Axes, + ax_ratio: plt.Axes | None, + ratio_ylabel: str = "Pred / Target", + ) -> None: + """Apply the shared x-label / ratio-panel styling, ratio panel permitting.""" + xlabel = "Frequency (1/deg)" + if ax_ratio is None: + # x-label moves up onto the spectra panel when there is no ratio panel below it. + ax_spec.set_xlabel(xlabel) + return + ax_ratio.axhline(1.0, ls="--", color="gray", lw=0.8) + ax_ratio.set_ylabel(ratio_ylabel) + ax_ratio.set_xlabel(xlabel) + ax_ratio.set_ylim(0, 2) + ax_ratio.grid(True, which="both", ls="--", alpha=0.4) + def psd_plot( self, psd_datasets: list[dict], @@ -715,7 +763,7 @@ def psd_plot( tag: str = "", variable: str = "", forecast_step: str = "", - ) -> None: + ) -> Path: """Create a PSD summary plot overlaying multiple runs. Each entry in *psd_datasets* is a dict with keys @@ -730,6 +778,11 @@ def psd_plot( Human-readable label for each run. tag : str Filename tag. + + Returns + ------- + Path + Path of the written figure (used to assemble forecast-step animations). """ out_dir = Path(self.out_plot_dir_psd) out_dir.mkdir(parents=True, exist_ok=True) @@ -738,12 +791,7 @@ def psd_plot( freq = np.asarray(psd_datasets[0]["frequencies"]) tar_psd = np.asarray(psd_datasets[0]["psd_target"]) - fig, (ax_spec, ax_ratio) = plt.subplots( - 2, - 1, - figsize=self.fig_size or (10, 8), - gridspec_kw={"height_ratios": [2, 1], "hspace": 0.08}, - ) + fig, ax_spec, ax_ratio = self._psd_axes((10, 8)) # Upper panel: log-log spectra ax_spec.loglog(freq, tar_psd, color="black", lw=1.5, label="Target") @@ -768,24 +816,22 @@ def psd_plot( ax_spec.legend(frameon=False, fontsize=7) ax_spec.grid(True, which="both", ls="--", alpha=0.4) - # Lower panel: ratio (pred / target) - for i, (ds, label) in enumerate(zip(psd_datasets, labels, strict=False)): - c = colors[i % len(colors)] - pred = np.asarray(ds["psd_prediction"]) - with np.errstate(divide="ignore", invalid="ignore"): - ratio = np.where(tar_psd > 0, pred / tar_psd, np.nan) - ax_ratio.semilogx(freq, ratio, color=c, lw=1.2, label=label) - ax_ratio.axhline(1.0, ls="--", color="gray", lw=0.8) - ax_ratio.set_ylabel("Pred / Target") - ax_ratio.set_xlabel("Frequency (1/deg)") - ax_ratio.set_ylim(0, 2) - ax_ratio.grid(True, which="both", ls="--", alpha=0.4) + # Lower panel: ratio (pred / target) - only when enabled. + if ax_ratio is not None: + for i, (ds, label) in enumerate(zip(psd_datasets, labels, strict=False)): + c = colors[i % len(colors)] + pred = np.asarray(ds["psd_prediction"]) + with np.errstate(divide="ignore", invalid="ignore"): + ratio = np.where(tar_psd > 0, pred / tar_psd, np.nan) + ax_ratio.semilogx(freq, ratio, color=c, lw=1.2, label=label) + self._finish_psd_axes(ax_spec, ax_ratio) name = tag or "psd" fname = out_dir / f"{name}.{self.image_format}" _logger.debug(f"Saving PSD summary plot to {fname}") fig.savefig(str(fname), bbox_inches="tight", dpi=self.dpi_val) plt.close(fig) + return fname def psd_evolution_plot( self, @@ -799,12 +845,12 @@ def psd_evolution_plot( Mirrors the diffusion-diagnostics "evolution" plot style (viridis colormap + colorbar keyed to the swept step), applied here to forecast lead-time steps instead of diffusion - denoising steps, and mirrors ``psd_plot``'s two-panel (spectra + ratio) layout. In the - top panel, solid lines are predictions and dashed lines are targets; both share the - same per-step colour so prediction/target can be compared at a given lead time while - also seeing how each evolves across lead time. The bottom panel shows the pred/target - ratio per step with the same per-step colour, where lead-time drift is expected to be - most visible. + denoising steps, and mirrors ``psd_plot``'s two-panel (spectra + ratio) layout. Solid, + colour-coded lines are the per-step predictions. Because the target spectrum is nearly + lead-time invariant, it is averaged over all forecast steps and drawn once in grey + rather than once per step; the ratio panel divides every step's prediction by that same + averaged target, so it matches what the upper panel shows. The ratio panel is omitted + when ``psd_show_ratio`` is False. Parameters ---------- @@ -825,36 +871,48 @@ def psd_evolution_plot( if len(fsteps) < 2: return - fig, (ax_spec, ax_ratio) = plt.subplots( - 2, - 1, - figsize=self.fig_size or (7, 7), - gridspec_kw={"height_ratios": [2, 1], "hspace": 0.08}, - ) + fig, ax_spec, ax_ratio = self._psd_axes((7, 7)) cmap = plt.get_cmap("viridis") n = max(len(fsteps) - 1, 1) + # The target spectrum barely changes with lead time: average it across forecast + # steps and draw it once, instead of one dashed per-step curve. + ref_freq = np.asarray(per_fstep_datasets[fsteps[0]]["frequencies"]) + targets = [np.asarray(per_fstep_datasets[f]["psd_target"]) for f in fsteps] + if all(t.shape == targets[0].shape for t in targets): + tar_mean = np.nanmean(np.vstack(targets), axis=0) + else: + _logger.warning( + f"PSD evolution ({label} / {variable}): target spectra differ in length across " + "forecast steps; falling back to the first step's target as reference." + ) + tar_mean = targets[0] + for i, fstep in enumerate(fsteps): ds = per_fstep_datasets[fstep] freq = np.asarray(ds["frequencies"]) pred = np.asarray(ds["psd_prediction"]) - tar = np.asarray(ds["psd_target"]) c = cmap(i / n) ax_spec.loglog(freq, pred, color=c, lw=1.0) - ax_spec.loglog(freq, tar, color=c, lw=1.0, ls="--", alpha=0.6) - with np.errstate(divide="ignore", invalid="ignore"): - ratio = np.where(tar > 0, pred / tar, np.nan) - ax_ratio.semilogx(freq, ratio, color=c, lw=1.0) + if ax_ratio is not None and pred.shape == tar_mean.shape: + with np.errstate(divide="ignore", invalid="ignore"): + ratio = np.where(tar_mean > 0, pred / tar_mean, np.nan) + ax_ratio.semilogx(freq, ratio, color=c, lw=1.0) + + # Single, step-independent grey target on top of the prediction bundle. + ax_spec.loglog(ref_freq, tar_mean, color="0.45", lw=1.8, ls="--", zorder=5) sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(fsteps[0], fsteps[-1])) - fig.colorbar(sm, ax=[ax_spec, ax_ratio], label="Forecast step") + fig.colorbar( + sm, ax=[ax_spec] if ax_ratio is None else [ax_spec, ax_ratio], label="Forecast step" + ) ax_spec.legend( handles=[ - Line2D([], [], color="gray", lw=1.0, ls="-", label="Prediction"), - Line2D([], [], color="gray", lw=1.0, ls="--", alpha=0.6, label="Target"), + Line2D([], [], color=cmap(0.5), lw=1.0, ls="-", label="Prediction (per step)"), + Line2D([], [], color="0.45", lw=1.8, ls="--", label="Target (mean over steps)"), ], frameon=False, fontsize=8, @@ -869,13 +927,137 @@ def psd_evolution_plot( ax_spec.set_ylabel("Power") ax_spec.grid(True, which="both", ls="--", alpha=0.4) - ax_ratio.axhline(1.0, ls="--", color="gray", lw=0.8) - ax_ratio.set_ylabel("Pred / Target") - ax_ratio.set_xlabel("Frequency (1/deg)") - ax_ratio.set_ylim(0, 2) - ax_ratio.grid(True, which="both", ls="--", alpha=0.4) + self._finish_psd_axes(ax_spec, ax_ratio, ratio_ylabel="Pred / Target (mean)") fname = out_dir / f"{tag or 'psd_evolution'}.{self.image_format}" _logger.debug(f"Saving PSD evolution plot to {fname}") fig.savefig(str(fname), bbox_inches="tight", dpi=self.dpi_val) plt.close(fig) + + def psd_gif(self, frame_paths: list[str | Path], tag: str = "") -> Path | None: + """Assemble already-written per-forecast-step PSD frames into an animated gif. + + Frames are used in the order given (no filename globbing), with the same Pillow + ``save_all``/``append_images`` pattern used for the map/histogram animations in + ``plot_orchestration.py``. + + Parameters + ---------- + frame_paths : list[str | Path] + Ordered frame paths, as returned by :meth:`psd_plot`. + tag : str + Filename tag (extension ``.gif`` is appended). + + Returns + ------- + Path | None + Path of the gif, or None when fewer than two frames were available. + """ + paths = [Path(p) for p in frame_paths if p is not None and Path(p).exists()] + if len(paths) < 2: + _logger.debug(f"PSD animation '{tag}' skipped: fewer than two frames.") + return None + + out_dir = Path(self.out_plot_dir_psd) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / f"{tag or 'psd_animation'}.gif" + duration_ms = int(1000 / self.fps) if self.fps and self.fps > 0 else 400 + + images = [Image.open(p) for p in paths] + try: + # savefig(bbox_inches="tight") can produce slightly different canvas sizes per + # frame (e.g. "step 1" vs "step 10" label width); normalize so Pillow doesn't + # paste subsequent frames onto the first frame's canvas misaligned. + base_size = images[0].size + frames = [im if im.size == base_size else im.resize(base_size) for im in images] + frames[0].save( + out_path, + save_all=True, + append_images=frames[1:], + duration=duration_ms, + loop=0, + ) + finally: + for im in images: + im.close() + _logger.debug(f"Saved PSD animation to {out_path}") + return out_path + + def psd_gap_heatmap( + self, + freq: NDArray, + fsteps: list[int], + grid: NDArray, + vmin: float, + vmax: float, + tag: str = "", + variable: str = "", + label: str = "", + psd_method: str = "sht", + ) -> Path: + """Render a precomputed frequency x forecast-step gap grid as a heatmap. + + Same quantity as the ratio panels (log(prediction/target) = log(Pred/Target)), but as + one dense grid instead of overlaid per-step curves, so frequency-band-specific drift + over lead time is visible directly. Positive = over-prediction, negative = + under-prediction, matching this codebase's ``calc_bias`` convention (``p - gt``, + score.py) and the ratio panels' ``pred / target`` (ratio > 1 = over-prediction). + ``vmin``/``vmax`` are supplied by the caller (``psd_plot_metric_region`` computes them + once across every run being compared, so all runs for a given stream/channel share one + colour scale) rather than fixed here, since a fixed bound can badly over- or under-shoot + the data depending on the run. The colormap is centred on 0 (perfect prediction) via + ``TwoSlopeNorm`` even when ``vmin``/``vmax`` aren't symmetric around it. + + The frequency axis for the SHT method (the only one actually exercised by this eval + config) is an integer spherical-harmonic wavenumber starting at 0, not a continuous + FFT-style frequency — plotted on a **linear** axis. A log axis previously mishandled the + zero wavenumber and produced wildly wrong, data-independent limits. + + Parameters + ---------- + freq : np.ndarray + Frequency/wavenumber axis, shape ``(n_freq,)``. + fsteps : list[int] + Forecast steps included in ``grid``'s rows, same order. + grid : np.ndarray + ``log(prediction) - log(target)``, shape ``(len(fsteps), len(freq))``. + vmin, vmax : float + Colour scale bounds (already extended to include 0 by the caller). + tag : str + Filename tag. + variable : str + Channel name, used in the title. + label : str + Run label, used in the title. + psd_method : str + Used in the title only. + + Returns + ------- + Path + Path of the written figure. + """ + out_dir = Path(self.out_plot_dir_psd) + out_dir.mkdir(parents=True, exist_ok=True) + + fig, ax = plt.subplots(figsize=self.fig_size or (8, 6)) + norm = TwoSlopeNorm(vcenter=0.0, vmin=vmin, vmax=vmax) + mesh = ax.pcolormesh(freq, fsteps, grid, cmap="coolwarm", norm=norm, shading="nearest") + ax.set_xlabel("Frequency (1/deg)") + ax.set_ylabel("Forecast step") + + title_parts = [f"PSD gap: log(pred) - log(target) ({psd_method})"] + if variable: + title_parts.append(variable) + if label: + title_parts.append(label) + ax.set_title(" - ".join(title_parts)) + + cbar = fig.colorbar(mesh, ax=ax) + cbar.set_label("log(prediction) - log(target)") + + fname = out_dir / f"{tag or 'psd_gap_heatmap'}.{self.image_format}" + _logger.debug(f"Saving PSD gap heatmap to {fname}") + fig.savefig(str(fname), bbox_inches="tight", dpi=self.dpi_val) + plt.close(fig) + return fname diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py b/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py index d349d68600..b4579f3294 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py @@ -1114,10 +1114,13 @@ def plot_summary(cfg: dict, scores_dict: dict, summary_dir: Path): "image_format": plt_opt.get("image_format", "png"), "dpi_val": plt_opt.get("dpi_val", 300), "fig_size": plt_opt.get("fig_size", (8, 10)), + "fps": plt_opt.get("fps", 2), "log_scale": eval_opt.get("log_scale", False), "add_grid": eval_opt.get("add_grid", False), "plot_ensemble": eval_opt.get("plot_ensemble", False), "baseline": eval_opt.get("baseline", None), + "psd_show_ratio": eval_opt.get("psd_show_ratio", True), + "psd_animation": eval_opt.get("psd_animation", True), } plotter = LinePlots(plot_cfg, summary_dir) diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py index d8fcda5047..60dbc8c22e 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py @@ -75,7 +75,7 @@ def compute_offsets(n, spacing=0.11): Returns ------- - np.ndarray + NDArray Array of length *n* with offsets centred at zero. """ idx = np.arange(n) @@ -205,7 +205,7 @@ def compute_marker_size(marker_size_base: float, scale: bool, lat: NDArray) -> f Returns ------- - float or np.ndarray + float or NDArray Scalar when *scale* is ``False``; array matching *lat* otherwise. """ if not scale: @@ -260,7 +260,7 @@ def auto_marker_size( Returns ------- - float or np.ndarray + float or NDArray Marker size(s) ready to pass to ``ax.scatter(s=...)``. """ if n_points >= density_threshold: @@ -277,12 +277,12 @@ def _flatten_or_average(arr: NDArray) -> NDArray: Parameters ---------- - arr : np.ndarray + arr : NDArray Input array, possibly multi-dimensional. Returns ------- - np.ndarray + NDArray Flattened 1D array, averaged across extra dimensions if needed. """ if arr.ndim > 1: @@ -789,6 +789,50 @@ def _extract_psd_attrs(data_ch: xr.DataArray, fstep: int, ch: str) -> list[dict] return None +def _compute_psd_gap_grid( + per_fstep_datasets: dict[int, dict], label: str = "", variable: str = "" +) -> tuple[NDArray, list[int], NDArray] | None: + """Build a (n_fsteps, n_freq) grid of log(prediction) - log(target) for one run. + + Positive = over-prediction, negative = under-prediction — matching this codebase's existing + ``calc_bias`` convention (``p - gt``, score.py) and the PSD ratio panels' ``pred / target`` + (ratio > 1 = over-prediction), rather than the mathematically-arbitrary opposite sign. + + Forecast steps whose frequency array doesn't match the first usable step's shape are + dropped (with a warning) rather than raising, since PSD's frequency grid can in principle + differ across steps depending on NaN-masking. + + Returns + ------- + tuple[NDArray, list[int], NDArray] | None + ``(freq, used_fsteps, grid)``, or None if fewer than two forecast steps were usable. + """ + fsteps = sorted(per_fstep_datasets) + ref_freq = np.asarray(per_fstep_datasets[fsteps[0]]["frequencies"]) + rows, used_fsteps = [], [] + for fstep in fsteps: + ds = per_fstep_datasets[fstep] + freq = np.asarray(ds["frequencies"]) + if freq.shape != ref_freq.shape: + _logger.warning( + f"PSD gap heatmap ({label} / {variable}): fstep {fstep} has a differently " + "shaped frequency grid; excluding it from the heatmap." + ) + continue + tar = np.asarray(ds["psd_target"]) + pred = np.asarray(ds["psd_prediction"]) + with np.errstate(divide="ignore", invalid="ignore"): + gap = np.log(np.where(pred > 0, pred, np.nan)) - np.log( + np.where(tar > 0, tar, np.nan) + ) + rows.append(gap) + used_fsteps.append(fstep) + + if len(rows) < 2: + return None + return ref_freq, used_fsteps, np.vstack(rows) + + def psd_plot_metric_region( metric: str, region: str, @@ -801,8 +845,10 @@ def psd_plot_metric_region( PSD curves (frequencies, target PSD, prediction PSD) are stored in ``score.attrs`` by ``Scores.calc_psd`` and read back here. - For a given forecast step, all runs are overlaid on one plot. Evolution - plots (across forecast steps) remain one-per-run. + For a given forecast step, all runs are overlaid on one plot. Evolution plots (across + forecast steps) remain one-per-run, accompanied by a per-run frequency x forecast-step + gap heatmap and a per-run, per-forecast-step animated gif (unless disabled via + ``plotter.psd_animate``). """ streams_set = collect_streams(runs) channels_set = collect_channels(scores_dict, metric, region, runs) @@ -863,23 +909,97 @@ def psd_plot_metric_region( forecast_step=str(fstep), ) - # Third pass: per-run evolution plot (unchanged — one run per plot). + # Third pass: per-run evolution plot, per-run gap heatmap (colour scale shared + # across every run being compared for this stream/channel, computed up front), + # plus (if enabled) per-run per-fstep frames assembled into an animated gif. + run_gap_grids: dict[str, tuple[NDArray, list[int], NDArray]] = {} for run_id, per_fstep_datasets in run_fstep_datasets.items(): - if len(per_fstep_datasets) >= 2: - method_tag = next(iter(per_fstep_datasets.values())).get( - "psd_method", "sht" - ) - evo_name = create_filename( + result = _compute_psd_gap_grid( + per_fstep_datasets, label=run_labels[run_id], variable=ch + ) + if result is not None: + run_gap_grids[run_id] = result + + if run_gap_grids: + all_gap_vals = np.concatenate( + [grid.ravel() for _, _, grid in run_gap_grids.values()] + ) + gap_min = float(np.nanmin(all_gap_vals)) + gap_max = float(np.nanmax(all_gap_vals)) + # Extend to include 0 (perfect prediction) so TwoSlopeNorm's centre is always + # valid, even if every run happens to be biased the same direction. TwoSlopeNorm + # requires vmin < vcenter(=0) < vmax strictly, so nudge either bound off zero + # with an epsilon scaled to the data when it would otherwise land exactly (or + # only) on one side. + eps = max(abs(gap_min), abs(gap_max), 1e-9) * 1e-3 + gap_vmin = min(gap_min, -eps) + gap_vmax = max(gap_max, eps) + + for run_id, per_fstep_datasets in run_fstep_datasets.items(): + if len(per_fstep_datasets) < 2: + continue + + method_tag = next(iter(per_fstep_datasets.values())).get("psd_method", "sht") + evo_name = create_filename( + prefix=[metric, method_tag, region], + middle=[run_id], + suffix=[stream, ch, "evolution"], + ) + plotter.psd_evolution_plot( + per_fstep_datasets, + tag=evo_name, + variable=ch, + label=run_labels[run_id], + ) + + if run_id in run_gap_grids: + freq, used_fsteps, grid = run_gap_grids[run_id] + heatmap_name = create_filename( prefix=[metric, method_tag, region], middle=[run_id], - suffix=[stream, ch, "evolution"], + suffix=[stream, ch, "gap_heatmap"], ) - plotter.psd_evolution_plot( - per_fstep_datasets, - tag=evo_name, + plotter.psd_gap_heatmap( + freq, + used_fsteps, + grid, + vmin=gap_vmin, + vmax=gap_vmax, + tag=heatmap_name, variable=ch, label=run_labels[run_id], + psd_method=method_tag, ) + + if not getattr(plotter, "psd_animate", True): + continue + + # Single-run frames (one per forecast step) -> animated gif. The literal + # "frame" token keeps these filenames distinct from pass 2's combined-run + # per-fstep plots, which degenerate to the same middle=[run_id] when only + # one run is being compared. + frame_paths = [] + for fstep in sorted(per_fstep_datasets): + frame_name = create_filename( + prefix=[metric, method_tag, region], + middle=[run_id], + suffix=[stream, ch, "frame", f"fstep{int(fstep):03d}"], + ) + frame_paths.append( + plotter.psd_plot( + [per_fstep_datasets[fstep]], + [run_labels[run_id]], + tag=frame_name, + variable=ch, + forecast_step=str(fstep), + ) + ) + anim_name = create_filename( + prefix=[metric, method_tag, region], + middle=[run_id], + suffix=[stream, ch, "animation"], + ) + plotter.psd_gif(frame_paths, tag=anim_name) _logger.info(f"PSD plots saved successfully into: {plotter.out_plot_dir_psd}") From ccaea9eb35c26c981b8a8ab09726d9811e3b2146 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Tue, 18 Aug 2026 18:12:35 +0200 Subject: [PATCH 14/15] eval config --- config/evaluate/eval_config_diffusion.yml | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/config/evaluate/eval_config_diffusion.yml b/config/evaluate/eval_config_diffusion.yml index 9b324ea03c..5012680dc6 100644 --- a/config/evaluate/eval_config_diffusion.yml +++ b/config/evaluate/eval_config_diffusion.yml @@ -29,6 +29,8 @@ evaluation: print_summary: false #print out score values on screen. it can be verbose log_scale: false add_grid: false + psd_show_ratio: true # false -> PSD plots (incl. gif frames) are spectra-only, no ratio panel + psd_animation: true # false -> skip the per-run PSD gif over forecast steps score_cards: false bar_plots: false num_processes: 0 #options: int, "auto", 0 means no parallelism (default) @@ -44,7 +46,7 @@ default_streams: ensemble: "all" #supported: "all", "mean", [0,1,2] plotting: sample: [0] - forecast_step: [1] #supported: "all", [1,2,3,...], "1-50" (equivalent of [1,2,3,...50]) + forecast_step: all #supported: "all", [1,2,3,...], "1-50" (equivalent of [1,2,3,...50]) ensemble: [0] #supported: "all", "mean", [0,1,2] plot_maps: true plot_target: true @@ -53,15 +55,21 @@ default_streams: run_ids : - st6x4i29: + + ku6i70v3: label: "forecast finetuned deterministic (skkbxfya -> dy0jlrmw)" - # kvp6fds7: # label: "forecast finetuned deterministic (skkbxfya -> mapw62ot)" - m3ezilrq: - label: "decoder finetuned diffusion (skkbxfya -> ni41n7gy -> mapw62ot)" + vfjgla68: + label: "decoder finetuned diffusion (skkbxfya -> ni41n7gy -> ft2z04v2)" # results_base_dir : "../results/" #NEW: if "streams" is not specified, the default streams are used - + # a7agq09c: + # label: "decoder finetuned diffusion (skkbxfya -> xxoplye6 -> hzsc0c1m)" + # ryl57hzv: + # label: "decoder finetuned diffusion (fx276yn3 (t-jepa) -> qadjp1xy -> y8rbtuey)" + # esxycoev: + # label: "decoder finetuned diffusion (cw6a4szu (t-jepa) -> c3nvsqu8 -> v9cqd62s)" + # pj9vb67u: + # label: "decoder finetuned diffusion (oqkez8qs (s-jepa) -> r10vvlcm -> gfac2ap5)" - \ No newline at end of file From c9df516d761c35d0fe1bed9a2a58d73234e86a51 Mon Sep 17 00:00:00 2001 From: moritzhauschulz Date: Wed, 19 Aug 2026 09:49:27 +0200 Subject: [PATCH 15/15] custom eval plots --- .../src/weathergen/evaluate/plotting/line_plots.py | 9 +++++++++ .../weathergen/evaluate/plotting/plot_orchestration.py | 9 +++++++++ .../src/weathergen/evaluate/plotting/plot_utils.py | 5 +++++ 3 files changed, 23 insertions(+) diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py b/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py index f6607a11f0..5c2193a884 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/line_plots.py @@ -24,6 +24,7 @@ from PIL import Image from weathergen.evaluate.plotting.plot_utils import ( + RASTER_ANIMATION_FORMATS, align_labels, channel_sort_key, clean_label, @@ -958,6 +959,14 @@ def psd_gif(self, frame_paths: list[str | Path], tag: str = "") -> Path | None: _logger.debug(f"PSD animation '{tag}' skipped: fewer than two frames.") return None + if self.image_format.lower() not in RASTER_ANIMATION_FORMATS: + _logger.warning( + f"PSD animation '{tag}' skipped: image_format={self.image_format!r} is not a " + f"raster format Pillow can read as frames (supported: " + f"{sorted(RASTER_ANIMATION_FORMATS)})." + ) + return None + out_dir = Path(self.out_plot_dir_psd) out_dir.mkdir(parents=True, exist_ok=True) out_path = out_dir / f"{tag or 'psd_animation'}.gif" diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py b/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py index b4579f3294..8afd72aca2 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/plot_orchestration.py @@ -32,6 +32,7 @@ group_by_init_hour, ) from weathergen.evaluate.plotting.plot_utils import ( + RASTER_ANIMATION_FORMATS, bar_plot_metric_region, heat_maps_metric_region, plot_metric_region, @@ -469,6 +470,14 @@ def _build_single_animation( if not output_dir.is_dir(): return [] + if image_format.lower() not in RASTER_ANIMATION_FORMATS: + _logger.warning( + f"Skipping animation for {run_id} {tag} {stream} {var}: image_format=" + f"{image_format!r} is not a raster format Pillow/imageio can read as frames " + f"(supported: {sorted(RASTER_ANIMATION_FORMATS)})." + ) + return [] + region_part = region if region else "" if sample is not None: head = "_".join(filter(None, [prefix, run_id, tag, str(sample)])) diff --git a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py index 60dbc8c22e..fc489aec91 100644 --- a/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py +++ b/packages/evaluate/src/weathergen/evaluate/plotting/plot_utils.py @@ -18,6 +18,11 @@ _logger = logging.getLogger(__name__) +# Animations (gif/mp4) are stitched from per-frame static plots with Pillow or imageio, both of +# which require raster frames. Vector/document formats like "pdf"/"svg"/"eps" are valid choices +# for the static plots themselves (via image_format) but can't be decoded as animation frames. +RASTER_ANIMATION_FORMATS = {"png", "jpg", "jpeg", "bmp", "tiff", "tif", "webp"} + # Shared helpers def calculate_average_over_dim(