Software companion
Building and Validating a Complex-Wavelet Earthformer Forecasting System
Overview
One repository backs all eleven documents in this series. This report is the part that builds the substrate and proves it correct, so its deliverable is not forecast skill but a pipeline in which every transformation can be checked: a wavelet representation that inverts exactly, an adaptation of an external pretrained backbone whose every modified tensor is accounted for, and a data boundary that cannot be crossed by accident.
The representation is the centre of it. A two-dimensional dual-tree complex wavelet transform is applied to each monthly map, and the packer moves the orientation and complex axes into channels — [B, C, 6, H, W, 2] → [B, C·6·2, H, W] — while retaining enough metadata for the inverse to be exact rather than approximate: per-level geometry, orientation order, complex encoding, and the original unpadded map shape. Real and imaginary channels stay paired through normalization, loss, unpacking and reconstruction, and the packing metadata travels inside the checkpoint, so a model cannot be reloaded against a layout it was not trained on.
The other half is access control. Climatology and coefficient-normalization statistics are fitted on training timestamps only, supervised loaders fail closed before 2016, and this milestone neither materializes sealed-period chunks nor evaluates them. That machinery is what makes the single authorized test read, several reports later, worth anything.
Implementation
- Core libraries
- torch, xarray, netCDF4, zarr, dask, pytorch-wavelets and PyWavelets, cdsapi, Matplotlib
- Representation
- A per-map 2D DTCWT packed into 37 coefficient channels on a 24×48 tropical-Pacific grid
- Forecast geometry
- Twelve context months in, fourteen forecast months out
- Data boundary
- ERA5 monthly SST: 1981–2010 training, 2011–2015 validation, 2016–2020 sealed test
- Runtime
- Official training and evaluation profiles require CUDA and fail rather than silently falling back to CPU
- Entry points
- Thirteen console scripts, from
cwe-build-real-datasetthroughcwe-validate-coeff-nativeandcwe-overfit-one-batch - Tests
- 62 pytest modules across the repository, roughly 14,000 lines
Components
wavelets/dtcwt.py- The transform wrapper, asserting the library’s tensor conventions at the boundary rather than assuming them
wavelets/packing.py- CoefficientPacker: lossless pack and unpack plus the versioned layout metadata that travels with every checkpoint
wavelets/normalization.py- CoefficientNormalizer: streaming mean and standard deviation at global, per-group or per-channel scope, fitted on the training split only
models/earthformer_adapter.py- The coefficient-native adaptation of the external Earthformer backbone
models/checkpointing.py- Checkpoint surgery, trainability audits, and the embedded provenance that every later phase reloads
data/cds.py, download.py, regrid.py- ERA5 acquisition and regridding to the analysis grid
data/anomalies.py, preprocess.py- Climatology construction and the anomaly reduction, fitted on training timestamps only
data/splits.py- The chronological split and the fail-closed sealed-period rule
validation.py, coefficient_native_validation.py, checkpoint_validation.py- The three standing validation entry points for the backbone, the coefficient path, and a saved projection checkpoint
training/overfit.py- The one-batch overfit harness used as a plumbing check
inference/forecast.py, viz/- Reconstruction to SST-anomaly maps and the map, Hovmöller and band-Hovmöller renders
Method pipeline
build-real-datasetAcquire ERA5 monthly SST, regrid to 24×48, and split chronologically.anomaliesFit climatology and normalization statistics on training timestamps only.DTCWTTransform each monthly map and pack the pyramid into 37 channels.trainFit the coefficient-native forecaster from 12 context months to 14 forecast months.inverse DTCWTReconstruct predicted coefficient pyramids to SST-anomaly maps.evaluateScore fields and regional indices and render maps and Hovmöller diagnostics.
Reproducibility and validation
- Round-trip and packing tests assert exact inverse reconstruction directly, rather than inferring correctness from a downstream metric that looks reasonable.
- Checkpoint surgery is audited rather than trusted: the official checkpoint hash, the surgery record and the trainability audit are embedded in the saved state and re-checked on load.
- Supervised loaders fail closed before 2016; this milestone neither materializes sealed-period chunks nor evaluates them.
- A one-batch overfit harness is a standing check — if the model cannot memorize a single batch, no skill claim is examined at all.
- Loading normalization metadata with a different coefficient layout, channel order, encoding, level count or geometry raises rather than rescaling silently.
Availability
Not publicly released. ERA5 acquisition needs Copernicus credentials and an accepted licence, and raw data, processed stores, checkpoints and generated products are all kept outside version control — the working copy reaches them through directory junctions to a separate volume. The synthetic and unit-test paths are network-free, which is what makes the pipeline testable without any of that.