collins909/DDPM-HI-models

DDPM HI Models — CAMELS

Conditional Denoising Diffusion Probabilistic Models that emulate neutral hydrogen (HI) 2D maps from the CAMELS Latin-Hypercube (LH) suite, and infer cosmological parameters back out of an observed map. Two variants ship here, each self-contained in its own folder.

DDPM-HI-models/
├── 2param/     # conditioned on (Ωm, σ8)
└── 6param/     # conditioned on (Ωm, σ8, A_SN1, A_AGN1, A_SN2, A_AGN2)

Each folder holds the checkpoint (model.pt), its training configuration (args.json, config.json), the label normalisation (label_stats.json), the full source tree (src/), SLURM launchers (scripts/shell/), the cross-model scripts (cross_model/), and a model card (README.md) documenting that variant in detail. Start with the model card in the folder you care about.

Both directions

Generation — parameters → map. inference_example.py downloads a checkpoint and samples one 256×256 HI map, via DDIM (50 steps) or full DDPM (1500 steps).

python 2param/inference_example.py --model 2param --device cuda --ddim-steps 50

Inference — map → parameters. vlb_posterior_example.py runs the inverse problem. The DDPM training objective

L(θ) = E_{t,ε} ‖ ε − ε_θ(x_t, t, y(θ)) ‖²

is a variational bound on −log p(x | θ). Freeze an observed map and read it as a function of the conditioning label instead: it becomes a proxy negative log-likelihood over parameters (Mudur, Cuesta-Lazaro & Finkbeiner 2023, arXiv:2312.07534; posterior form as in Diffusion-HMC, arXiv:2405.05255). Scanning it over the prior box yields a point estimate argmin_θ L(θ), a −2Δlnℒ̂ χ² map, and a posterior p(θ | x) ∝ exp(−βL(θ)).

python 2param/vlb_posterior_example.py --model 2param --from-dataset --index 0
python 6param/vlb_posterior_example.py --model 6param --map my_map.npy --astro-preset low

Two variance-control tricks make the scan usable, both on by default: one common random number bank of (t, ε) reused at every grid point, and stratified timesteps spread evenly over [0, T).

β is not calibrated. It defaults to the pixel count (65536), the naive unit-variance-Gaussian guess. The argmin and the shape of the loss surface stand on their own, but the credible-interval widths do not until β is calibrated by coverage over many held-out fields — see src/calibrate_beta.py. Each run prints n_eff, the number of grid cells carrying posterior mass; n_eff ≈ 1 means β has collapsed the softmax, not that the constraint is tight.

Reported behaviour on held-out CAMELS fields: Ωm is recovered to ≈0.01 (2-param) / ≈0.02 (6-param), σ8 is the weak direction, and the feedback amplitudes are essentially unconstrained — consistent with HIFlow.

Architecture

Both variants share one architecture; only label_dim differs.

Field	Value
Backbone	ConditionalUNet + GaussianDiffusion (plain PyTorch, not diffusers)
base_channels	64
channel_multipliers	[1, 2, 4, 8]
attention_levels	[2, 3]
timesteps	1500 (linear β schedule: 1e-4 → 0.02)
EMA decay	0.9999
Image	256 × 256, single channel, [-1, 1]

These are plain PyTorch state dicts, so DDPMPipeline.from_pretrained will not load them. Rebuild the module from args.json and load model.pt["model_state_dict"] — every bundled example does exactly that, and the model cards show the code.

Conditioning vectors are z-scored with train-split statistics. Those statistics ship as label_stats.json in each folder (mean/std plus the min/max of the Latin-hypercube box), so you do not need the dataset to condition on physical parameter values.

Data

Both models are trained on, and evaluated against, collins909/DDPM-HI-CAMELS-LH — CAMELS LH HI maps, 256 × 256, stored in [0, 1] on disk and rescaled to [-1, 1] on the way into the model.

Split	Maps	2-param arrays	6-param arrays
train	13500	params_2/train_LH.npy, params_2/train_labels_LH_2.npy	params_6/train_LH_6.npy, params_6/train_labels_LH.npy
val	750	params_2/val_LH.npy, params_2/val_labels_LH_2.npy	params_6/val_LH_6.npy, params_6/val_labels_LH.npy
test	750	params_2/test_LH.npy, params_2/test_labels_LH_2.npy	params_6/test_LH_6.npy, params_6/test_labels_LH.npy

vlb_posterior_example.py --from-dataset wires the two together for you: it pulls a held-out map and its true parameters from the dataset repo, runs the scan, and reports the recovery error against ground truth. Nothing to download by hand.

python 2param/vlb_posterior_example.py --model 2param --from-dataset --split test --index 0

To go through the arrays yourself:

import numpy as np
from huggingface_hub import hf_hub_download

D = "collins909/DDPM-HI-CAMELS-LH"
images = np.load(hf_hub_download(D, "params_2/test_LH.npy", repo_type="dataset"), mmap_mode="r")
labels = np.load(hf_hub_download(D, "params_2/test_labels_LH_2.npy", repo_type="dataset"))

x = images[0] * 2.0 - 1.0          # [0,1] on disk -> [-1,1] for the model
theta = labels[0]                  # (Omega_m, sigma_8)

Note the asymmetric filenames: the 2-param images carry no suffix while their labels end in _2, and the 6-param set is the other way round. Training and evaluation scripts want the directory containing these files, e.g. --data_dir .../params_2.

Conditioning on physical parameter values needs the train-split z-scoring, which ships as label_stats.json in each model folder — so you do not need to download the 7.3 GB dataset just to generate maps.

Intended use & limitations

• Intended for research on diffusion emulators for cosmological fields and on field-level, simulation-based parameter inference.
• The models only ever saw 256 × 256 CAMELS LH maps; other resolutions or simulation suites are off-distribution.
• The VLB posterior is a bound-based likelihood proxy with an uncalibrated temperature. Treat point estimates as the trustworthy output and interval widths as provisional.

Citation

Please cite the CAMELS project and the upstream DDPM HI emulation work. (Citation block to be filled in once the accompanying paper is published.)

Contact

Collins Maripane — University of Cape Town, Department of Mathematics and Applied Mathematics.