README.md

---
license: mit
library_name: pytorch
tags:
  - diffusion
  - ddpm
  - ddim
  - cosmology
  - astrophysics
  - camels
  - emulator
  - conditional-generation
pipeline_tag: unconditional-image-generation
---

# DDPM HI Emulator — 6 Parameter (CAMELS LH)

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.)