A conditional Denoising Diffusion Probabilistic Model (DDPM) that emulates neutral-hydrogen (HI) 2D maps from the CAMELS Latin-Hypercube (LH) simulation suite, conditioned on the full 6 CAMELS LH parameters (Ωm, σ8, ASN1, AAGN1, ASN2, AAGN2). Sampling supports both full DDPM and accelerated DDIM.

This is the best-validation checkpoint from the training run under ddpm_hi_lh6/outputs_conditional_6param_20260413_132226/.

Files in this repo

Top level

File	Purpose
`model.pt`	PyTorch checkpoint (state-dict for `ConditionalDiffusionModel`)
`args.json` / `args.txt`	Training hyper-parameters and U-Net configuration
`config.json`	Architecture summary (for Hub discoverability)
`inference_example.py`	Runnable example: downloads weights and generates a sample

src/ — per-model Python

File	Purpose
`train_conditional.py`	Training entry point (`label_dim=6`, mixed-precision)
`evaluate_conditional.py`	Held-out evaluation: samples + metrics
`eval_model.py`	Lightweight evaluation helper used by the figure scripts
`posterior_inference.py`	Full posterior-inference pipeline (likelihood / sampling)
`figure9_posterior.py`	Paper figure 9 (posterior triangle for the 6-param model)
`plot_r2_cosmology_lhs.py`	Latin-hypercube R² map (μ, σ vs cosmology)
`unet_conditional.py`	`ConditionalUNet` module
`diffusion_conditional.py`	`GaussianDiffusion` (DDPM + DDIM) and the wrapping `ConditionalDiffusionModel`
`dataset_conditional.py`	CAMELS LH dataset loader + label normalisation

scripts/shell/ — SLURM launchers

File	Purpose
`train_conditional_lh6.sh`	Submit a training job (`label_dim=6`)
`evaluate_conditional_lh6.sh`	Submit evaluation against the held-out test split
`plot_r2_cosmology_lhs.sh`	Generate the R² cosmology figure

cross_model/ — posterior + comparison scripts that use BOTH models

File	Purpose
`compare_posterior_inference.py` (+ `run_compare_posterior.sh`)	End-to-end posterior comparison between 2-param and 6-param emulators
`ddpm_posterior_corrected.py` (+ `scripts/run_ddpm_posterior_corrected.sh`)	Corrected DDPM posterior inference
`poster.py` / `check_poster_env.py` (+ `scripts/run_poster.sh`)	Posterior orchestration and environment check
`submit_vlb_1000grid.py` / `run_vlb_inference_*.sh`	Variational-lower-bound grid inference (200 / 1000 grid)
`scripts/compare_ddpm_models.py` (+ `run_ddpm_comparison.sh`)	DDPM-2 vs DDPM-6 comparison figures
`scripts/ddpm_posterior_six_anchors.py` (+ `run_ddpm_posterior_six_anchors.sh`)	Six-anchor posterior visualisation
`scripts/ddpm_figure6_integration.py`, `figure6_2409_style.py`, `run_ddpm_figure6_suite.py` (+ `run_ddpm_figure6.sh`)	Figure 6 generation pipeline
`scripts/ddpm_triangle_integration.py`, `triangle_plot_posterior.py` (+ `run_triangle_ddpm_both.sh`)	Triangle-plot posterior figures
`scripts/sigma_contour_utils.py`	Confidence-contour helper used by the figure scripts
`scripts/compare_ddpm_training_curves.py`	Parses SLURM logs for combined train/val loss plots
`cross_model/README.md`	How to point these scripts at locally-downloaded weights/data

These cross-model scripts default to the original cluster paths (e.g. <CAMELS_LH_DATA_DIR>/params_6). After downloading this repo, supply --bundle-2param, --bundle-6param, --data-2param, --data-6param to override.

Architecture

Conditional U-Net + Gaussian diffusion process. Hyper-parameters (taken from args.json):

Field	Value
`label_dim`	6
`base_channels`	64
`channel_multipliers`	[1, 2, 4, 8]
`attention_levels`	[2, 3]
`dropout`	0.1
`timesteps`	1500 (linear β schedule: 1e-4 → 0.02)
EMA decay	0.9999
Mixed precision	Yes (`use_amp = true` during training)
Sampler	DDIM, 50 steps (DDPM also supported)
Image size	256 × 256, single channel
Image range	[-1, 1] (training data is rescaled by `x * 2 - 1`)

Labels are z-scored using the training-split mean / std. The inference_example.py shows how to recover this normalisation from the CAMELS LH params_6 dataset, or you can pass already-normalised conditioning values directly.

Quick start

python

from huggingface_hub import hf_hub_download
import sys, torch, json
from pathlib import Path

# 1) Download all needed files
repo = "collins909/DDPM-6param"
ckpt_path  = hf_hub_download(repo, "model.pt")
args_path  = hf_hub_download(repo, "args.json")
for name in ("unet_conditional.py", "diffusion_conditional.py", "__init__.py"):
    hf_hub_download(repo, f"src/{name}")
sys.path.insert(0, str(Path(ckpt_path).parent / "src"))

from unet_conditional import ConditionalUNet
from diffusion_conditional import GaussianDiffusion, ConditionalDiffusionModel

# 2) Rebuild the model from args.json
args = json.loads(Path(args_path).read_text())
unet = ConditionalUNet(
    in_channels=1, out_channels=1,
    label_dim=args["label_dim"],
    base_channels=args["base_channels"],
    channel_multipliers=tuple(args["channel_multipliers"]),
    attention_levels=tuple(args["attention_levels"]),
    dropout=args["dropout"],
)
diffusion = GaussianDiffusion(
    timesteps=args["timesteps"],
    beta_start=args["beta_start"],
    beta_end=args["beta_end"],
    schedule_type=args["schedule_type"],
)
model = ConditionalDiffusionModel(unet, diffusion)

# 3) Load the checkpoint and sample
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

# 6-parameter conditioning vector (order: Ωm, σ8, ASN1, AAGN1, ASN2, AAGN2),
# z-scored with training-split stats. See inference_example.py for the helper.
labels = torch.zeros((1, 6))
sample = model.sample(labels, channels=1, height=256, width=256,
                      device="cpu", use_ddim=True, ddim_steps=50)
# sample is in [-1, 1]; rescale to physical HI units as needed.

For an end-to-end runnable example (including label normalisation, GPU usage, and image saving), see inference_example.py in this repo.

Training data

Trained on CAMELS LH HI maps with full 6-parameter conditioning. The data layout consumed by src/dataset_conditional.py is:

<data_dir>/
  train_LH_6.npy, val_LH_6.npy, test_LH_6.npy
  train_labels_LH.npy, val_labels_LH.npy, test_labels_LH.npy

Images are rescaled to [-1, 1]; labels are z-scored using train-split statistics. Point your training/eval scripts at the local directory that contains those files (e.g. via --data_dir <CAMELS_LH_DATA_DIR>/params_6).

Intended use & limitations

Intended for research on diffusion emulators for cosmological fields, posterior inference, and sensitivity studies across cosmology / astrophysics nuisance parameters.
The companion 2-parameter model (collins909/DDPM-2param) is available for the simpler 2-label setup.
Outputs are 256 × 256 single-channel maps in the model's normalised range. Apply the inverse of any data-pipeline preprocessing before physical interpretation.

Citation

If you use this checkpoint, please cite the CAMELS project and the upstream DDPM HI emulation work. (Citation block to be filled in once the accompanying paper is published.)