Software companion
Band-Weighted Temporal Sequence Diffusion for Probabilistic ENSO Trajectories
Overview
The earlier diffusion head folds the forecast-lead axis into the batch axis, so each month is denoised independently and a member has no reason to be temporally coherent. The head here keeps the lead axis and denoises the whole trajectory jointly, alternating spatial residual blocks with temporal convolution and temporal attention so that a member is a path rather than fourteen unrelated maps.
Architecture alone did not get there, and the code records why. Two further mechanisms are separable switches rather than defaults: correlated noise, which draws member noise with a tunable temporal correlation instead of independent draws per lead, and per-channel band weighting, which re-weights the squared error across coefficient channels so the objective stops being dominated by the bands that carry the most energy. The report’s finding — that the re-weighting, not the joint architecture, is what unlocks coherence — is only statable because the three can be varied independently.
Implementation
- Head
- A joint 14-lead conditional DDPM denoiser over correction trajectories, with spatial residual blocks, temporal convolution and temporal attention
- Mechanisms
- Correlated member noise with a tunable correlation, and per-channel band weighting of the objective, each independently switchable
- Sweeps
- A sampler-stochasticity sweep that re-samples one trained checkpoint without retraining, plus a sequence-evaluation pass per variant — validation only
- Metrics
- Trajectory-level metrics — temporal coherence and path behaviour — alongside the per-lead scores
- Sealed test
- A documented protocol amendment authorizing one further read of the sealed split for the adopted long-trained variant
Components
models/sequence_diffusion.py- The joint-sequence head: spatial residual blocks, temporal convolution, temporal attention, correlated noise and band-weighted loss
trajectory_metrics.py- Temporal coherence and trajectory-level scoring
training/losses.py- The per-channel band weights the sequence objective consumes
probabilistic_evaluation.py- The shared ensemble evaluation surface, reused unchanged
scripts/phase_g0_correlated_noise.py- The correlated-noise study preceding the joint head
scripts/phase_g1_eta_sweep.py- The sampler sweep: how much fresh noise the reverse chain re-injects per step, and optionally how correlated it is across leads
scripts/phase_g1_sequence_evaluation.py- Sequence evaluation for each variant
scripts/phase_g_correction_structure.py- Structure of the learned corrections, band by band
scripts/phase_g3_recenter_trajectories.py- Trajectory recentring before calibration
scripts/phase_g_adopted_gallery.py- The adopted-variant case gallery
Workflow
trainFit the joint-sequence head on correction trajectories from the frozen deterministic centre.sweepVary the band weighting, the noise correlation and the sampler stochasticity independently of the architecture.evaluate sequencesScore temporal coherence and trajectory behaviour, not only per-lead error.recentre and calibrateApply the projection-space calibration from the previous report to the adopted trajectories.amend the protocolFreeze the adopted stack in a written amendment before any test access.sealed testRead the sealed split once under that amendment.
Reproducibility and validation
- Joint architecture, correlated noise and band weighting are independently switchable, which is what lets the report attribute the gain to the objective rather than to the architecture.
- Trajectory metrics are separate from per-lead metrics, so improved coherence cannot be inferred from unchanged per-lead error.
- The adopted stack is frozen in a protocol amendment before the test read, and the amendment is a checked document rather than an intention.
- The report states plainly that the validation advantage only partly transferred to the sealed period, which is the outcome the single-read protocol exists to make reportable.
Availability
Not publicly released. The sequence head is trained against correction trajectories from a specific frozen checkpoint, so the code is not independently runnable without artifacts the repository does not distribute.