Upload 2-parameter conditional DDPM (HI emulation, CAMELS LH params_2, epoch 200)

README.md

@@ -0,0 +1,143 @@
+---
+license: mit
+library_name: pytorch
+tags:
+  - diffusion
+  - ddpm
+  - ddim
+  - cosmology
+  - astrophysics
+  - camels
+  - emulator
+  - conditional-generation
+pipeline_tag: unconditional-image-generation
+---
+
+# DDPM HI Emulator — 2 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 **two cosmological parameters**
+(e.g. Ωm, σ8). Sampling supports both full DDPM and accelerated DDIM.
+
+This checkpoint is **epoch 200** of the training run carried out under
+`DDPM_HI_Emulation_improved/outputs_conditional_2label_20260408_125646/`.
+
+## Files in this repo
+
+| 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) |
+| `src/unet_conditional.py` | `ConditionalUNet` module |
+| `src/diffusion_conditional.py` | `GaussianDiffusion` (DDPM + DDIM) and the wrapping `ConditionalDiffusionModel` |
+| `src/dataset_conditional.py` | Helper for loading CAMELS LH data + label normalisation stats |
+| `src/evaluate_conditional.py` | Reference evaluation pipeline (samples + metrics) |
+| `inference_example.py` | Runnable example: downloads weights and generates a sample |
+
+## Architecture
+
+Conditional U-Net + Gaussian diffusion process. Hyper-parameters (taken from
+`args.json`):
+
+| Field | Value |
+|-------|-------|
+| `label_dim` | 2 |
+| `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 |
+| 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_2` 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 = "collinsmaripane/ddpm-hi-2param"
+ckpt_path  = hf_hub_download(repo, "model.pt")
+args_path  = hf_hub_download(repo, "args.json")
+# Pull the bundled source files so we can import the model classes.
+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()
+
+# Conditioning vector must be z-scored using training-split label statistics.
+labels = torch.tensor([[0.0, 0.0]])  # placeholder; see inference_example.py
+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 2-label conditioning. The exact data
+layout used by `src/dataset_conditional.py` is:
+
+```
+<data_dir>/
+  train_LH_2.npy, val_LH_2.npy, test_LH_2.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. The original training paths on the cluster were
+`/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2`.
+
+## Intended use & limitations
+
+- Intended for **research** on diffusion emulators for cosmological fields.
+- The 2-label setup is a simplified subset of the full CAMELS LH parameter
+  space; see the companion **6-parameter** model
+  (`collinsmaripane/ddpm-hi-6param`) for the full conditioning.
+- 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.)

args.json

@@ -0,0 +1,38 @@
+{
+  "label_dim": 2,
+  "base_channels": 64,
+  "channel_multipliers": [
+    1,
+    2,
+    4,
+    8
+  ],
+  "attention_levels": [
+    2,
+    3
+  ],
+  "dropout": 0.1,
+  "timesteps": 1500,
+  "beta_start": 0.0001,
+  "beta_end": 0.02,
+  "schedule_type": "linear",
+  "epochs": 210,
+  "batch_size": 8,
+  "lr": 0.0002,
+  "ema_decay": 0.9999,
+  "num_workers": 4,
+  "early_stop_patience": 100,
+  "use_amp": false,
+  "data_dir": "/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2",
+  "normalize_labels": true,
+  "output_dir": "outputs_conditional_2label",
+  "resume": "outputs_conditional_2label_20260401_164603/checkpoints/checkpoint_epoch_160.pt",
+  "resume_refresh_scheduler": true,
+  "sample_every": 10,
+  "use_ddim": true,
+  "ddim_steps": 50,
+  "use_wandb": false,
+  "wandb_project": "ddpm_cosmology",
+  "wandb_entity": "",
+  "wandb_run_name": ""
+}

args.txt

@@ -0,0 +1,28 @@
+label_dim: 2
+base_channels: 64
+channel_multipliers: [1, 2, 4, 8]
+attention_levels: [2, 3]
+dropout: 0.1
+timesteps: 1500
+beta_start: 0.0001
+beta_end: 0.02
+schedule_type: linear
+epochs: 210
+batch_size: 8
+lr: 0.0002
+ema_decay: 0.9999
+num_workers: 4
+early_stop_patience: 100
+use_amp: False
+data_dir: /scratch/mrpcol001/Diffusion_job/data/LH_data/params_2
+normalize_labels: True
+output_dir: outputs_conditional_2label
+resume: outputs_conditional_2label_20260401_164603/checkpoints/checkpoint_epoch_160.pt
+resume_refresh_scheduler: True
+sample_every: 10
+use_ddim: True
+ddim_steps: 50
+use_wandb: False
+wandb_project: ddpm_cosmology
+wandb_entity: 
+wandb_run_name: 

config.json

@@ -0,0 +1,24 @@
+{
+  "architecture": "ConditionalUNet + GaussianDiffusion",
+  "task": "conditional-image-generation",
+  "image_size": 256,
+  "in_channels": 1,
+  "out_channels": 1,
+  "label_dim": 2,
+  "base_channels": 64,
+  "channel_multipliers": [
+    1,
+    2,
+    4,
+    8
+  ],
+  "attention_levels": [
+    2,
+    3
+  ],
+  "dropout": 0.1,
+  "timesteps": 1500,
+  "beta_start": 0.0001,
+  "beta_end": 0.02,
+  "schedule_type": "linear"
+}

inference_example.py

@@ -0,0 +1,290 @@
+#!/usr/bin/env python
+# ---------------------------------------------------------------------------
+# inference_example.py
+#
+# Self-contained example that downloads a conditional-DDPM checkpoint from
+# the Hugging Face Hub and generates one HI map.
+#
+# Works for **both** uploaded models -- the script picks which one to load
+# from a CLI argument:
+#
+#     python inference_example.py --model 2param
+#     python inference_example.py --model 6param
+#     python inference_example.py --model 2param --repo myuser/my-fork
+#     python inference_example.py --model 6param --device cuda --ddim-steps 50
+#
+# The script:
+#   1. Downloads `model.pt`, `args.json`, and the bundled src/*.py files.
+#   2. Imports `ConditionalUNet` and `GaussianDiffusion` from the downloaded
+#      code (no need for a separate pip-installed package).
+#   3. Rebuilds the model from `args.json` so weights and architecture
+#      cannot drift apart.
+#   4. Samples one image with DDIM (or DDPM, with `--no-ddim`).
+#   5. Saves a `.npy` of the raw [-1, 1] output and a PNG visualisation.
+#
+# This file is bundled inside each HF repo so users can grab a single script
+# and immediately do inference.
+# ---------------------------------------------------------------------------
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+# huggingface_hub is the only "extra" dependency; everything else (torch,
+# numpy) is already required to run the model.
+from huggingface_hub import hf_hub_download
+
+
+# --------------------------------------------------------------------------
+# Defaults -- adjust here or override via CLI flags
+# --------------------------------------------------------------------------
+DEFAULT_REPOS = {
+    "2param": "collinsmaripane/ddpm-hi-2param",
+    "6param": "collinsmaripane/ddpm-hi-6param",
+}
+
+# All files we expect to find in every uploaded repo. We download each one
+# explicitly (rather than `snapshot_download`) so we can give a clear error
+# message if anything is missing.
+REQUIRED_FILES = [
+    "model.pt",
+    "args.json",
+    "src/__init__.py",
+    "src/unet_conditional.py",
+    "src/diffusion_conditional.py",
+]
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Sample one HI map from the HF-hosted DDPM.")
+    p.add_argument(
+        "--model",
+        choices=sorted(DEFAULT_REPOS.keys()),
+        required=True,
+        help="Which model to download. Picks the matching default HF repo.",
+    )
+    p.add_argument(
+        "--repo",
+        default=None,
+        help="Override the HF repo id (default: see DEFAULT_REPOS in this file).",
+    )
+    p.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Torch device for sampling. Defaults to cuda if available else cpu.",
+    )
+    p.add_argument(
+        "--ddim-steps",
+        type=int,
+        default=50,
+        help="Number of DDIM steps (ignored when --no-ddim).",
+    )
+    p.add_argument(
+        "--no-ddim",
+        action="store_true",
+        help="Use the full DDPM sampler (slow, all 1500 steps) instead of DDIM.",
+    )
+    p.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="RNG seed for reproducible sampling.",
+    )
+    p.add_argument(
+        "--labels",
+        type=float,
+        nargs="+",
+        default=None,
+        help=(
+            "Conditioning vector (already z-scored). Length must match label_dim "
+            "(2 or 6). If omitted, an all-zeros vector is used (i.e. the training-set mean)."
+        ),
+    )
+    p.add_argument(
+        "--output-dir",
+        type=Path,
+        default=Path("inference_outputs"),
+        help="Where to write the generated sample (.npy + .png).",
+    )
+    return p.parse_args()
+
+
+def download_repo(repo_id: str) -> Path:
+    """Download every required file from `repo_id`, return the local cache dir.
+
+    We rely on `hf_hub_download` to manage caching -- it stores files under
+    `~/.cache/huggingface/hub/` and returns the local path. We assume all the
+    required files end up in the same directory (which they do, modulo the
+    `src/` subfolder).
+    """
+    print(f"[inference] Downloading {len(REQUIRED_FILES)} files from {repo_id}")
+    local_paths = [Path(hf_hub_download(repo_id, f)) for f in REQUIRED_FILES]
+    # The repo root in the local cache is the parent of `model.pt`.
+    repo_root = local_paths[0].parent
+    print(f"[inference] Cached at: {repo_root}")
+    return repo_root
+
+
+def build_model(args_json: dict):
+    """Re-create `ConditionalDiffusionModel` from the training args dict.
+
+    Importing the model classes from the just-downloaded `src/` package is
+    the safest way to avoid drift between weights and architecture: if the
+    repo ships a particular version of the U-Net code, that's the version
+    we use.
+    """
+    from unet_conditional import ConditionalUNet
+    from diffusion_conditional import ConditionalDiffusionModel, GaussianDiffusion
+
+    unet = ConditionalUNet(
+        in_channels=1,
+        out_channels=1,
+        label_dim=args_json["label_dim"],
+        base_channels=args_json["base_channels"],
+        channel_multipliers=tuple(args_json["channel_multipliers"]),
+        attention_levels=tuple(args_json["attention_levels"]),
+        dropout=args_json["dropout"],
+    )
+    diffusion = GaussianDiffusion(
+        timesteps=args_json["timesteps"],
+        beta_start=args_json["beta_start"],
+        beta_end=args_json["beta_end"],
+        schedule_type=args_json["schedule_type"],
+    )
+    return ConditionalDiffusionModel(unet, diffusion)
+
+
+def load_weights(model: torch.nn.Module, ckpt_path: Path, device: str) -> None:
+    """Load the state-dict produced by `train_conditional.py`.
+
+    The checkpoint is a dict with keys:
+        model_state_dict, optimizer_state_dict, ema_shadow, epoch, loss, ...
+    We only need `model_state_dict` for inference.
+    """
+    # weights_only=False because the checkpoint also serialises optimizer
+    # state, EMA shadows, scheduler, etc. Safe here because we trust the
+    # source (the file came from our own training run on the cluster).
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    if "model_state_dict" not in ckpt:
+        raise KeyError(
+            f"{ckpt_path} doesn't contain 'model_state_dict' -- got keys: {list(ckpt)}"
+        )
+    model.load_state_dict(ckpt["model_state_dict"])
+    epoch = ckpt.get("epoch", "?")
+    loss = ckpt.get("loss", "?")
+    print(f"[inference] Loaded weights (epoch={epoch}, loss={loss})")
+
+
+def save_outputs(sample: torch.Tensor, output_dir: Path, model_name: str) -> None:
+    """Write the generated map to disk both as raw .npy and as a PNG preview."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # `sample` is shape (1, 1, 256, 256) in [-1, 1]; squeeze and bring to CPU.
+    arr = sample.squeeze().detach().cpu().numpy()
+    npy_path = output_dir / f"sample_{model_name}.npy"
+    np.save(npy_path, arr)
+    print(f"[inference] Wrote {npy_path}  shape={arr.shape}  range=[{arr.min():.3f}, {arr.max():.3f}]")
+
+    # Optional PNG -- only if matplotlib is around. Keeps the hard dependency
+    # list short (matplotlib isn't strictly needed for the science workflow).
+    try:
+        import matplotlib.pyplot as plt
+    except ImportError:
+        print("[inference] matplotlib not installed -- skipping PNG preview.")
+        return
+    png_path = output_dir / f"sample_{model_name}.png"
+    plt.figure(figsize=(5, 5))
+    plt.imshow(arr, cmap="inferno", origin="lower")
+    plt.axis("off")
+    plt.title(f"DDPM {model_name} sample")
+    plt.tight_layout()
+    plt.savefig(png_path, dpi=120, bbox_inches="tight")
+    plt.close()
+    print(f"[inference] Wrote {png_path}")
+
+
+def main() -> None:
+    args = parse_args()
+    repo_id = args.repo or DEFAULT_REPOS[args.model]
+
+    # ----------------------------------------------------------------------
+    # 1. Pull files from the Hub and make src/ importable
+    # ----------------------------------------------------------------------
+    repo_root = download_repo(repo_id)
+    sys.path.insert(0, str(repo_root / "src"))
+
+    # ----------------------------------------------------------------------
+    # 2. Rebuild the model from args.json
+    # ----------------------------------------------------------------------
+    with open(repo_root / "args.json") as f:
+        train_args = json.load(f)
+    expected_dim = train_args["label_dim"]
+    if expected_dim != (2 if args.model == "2param" else 6):
+        raise ValueError(
+            f"args.json says label_dim={expected_dim} but --model={args.model}; "
+            "did you point --repo at the wrong checkpoint?"
+        )
+
+    model = build_model(train_args).to(args.device)
+    load_weights(model, repo_root / "model.pt", args.device)
+    model.eval()
+
+    # ----------------------------------------------------------------------
+    # 3. Build the conditioning vector
+    # ----------------------------------------------------------------------
+    # By default we feed zeros, i.e. the training-set mean in the normalised
+    # space. To condition on physical (Ωm, σ8, ...) values, z-score them
+    # using the train-split statistics produced by `dataset_conditional.py`
+    # and pass the result via --labels.
+    if args.labels is None:
+        labels = torch.zeros((1, expected_dim), device=args.device)
+        print(f"[inference] Using zero (training-mean) conditioning, label_dim={expected_dim}")
+    else:
+        if len(args.labels) != expected_dim:
+            raise ValueError(
+                f"--labels has {len(args.labels)} entries but model expects {expected_dim}"
+            )
+        labels = torch.tensor([args.labels], dtype=torch.float32, device=args.device)
+        print(f"[inference] Using user-supplied labels: {args.labels}")
+
+    # ----------------------------------------------------------------------
+    # 4. Sample
+    # ----------------------------------------------------------------------
+    # Fix the RNG seed for reproducibility -- diffusion sampling is very
+    # sensitive to the initial Gaussian noise.
+    torch.manual_seed(args.seed)
+    if args.device.startswith("cuda"):
+        torch.cuda.manual_seed_all(args.seed)
+
+    use_ddim = not args.no_ddim
+    print(
+        f"[inference] Sampling 1 image with "
+        f"{'DDIM ' + str(args.ddim_steps) + ' steps' if use_ddim else 'DDPM ' + str(train_args['timesteps']) + ' steps'} "
+        f"on {args.device} ..."
+    )
+    with torch.no_grad():
+        sample = model.sample(
+            labels=labels,
+            channels=1,
+            height=256,
+            width=256,
+            device=args.device,
+            progress=True,
+            use_ddim=use_ddim,
+            ddim_steps=args.ddim_steps,
+            eta=0.0,
+        )
+
+    # ----------------------------------------------------------------------
+    # 5. Save outputs
+    # ----------------------------------------------------------------------
+    save_outputs(sample, args.output_dir, args.model)
+    print("[inference] Done.")
+
+
+if __name__ == "__main__":
+    main()

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:545b39a5a4bfb58a230c9abfaf73006069a26f9b50cd2cc9e01b3693764ca1ce
+size 1070314125

src/dataset_conditional.py

@@ -0,0 +1,86 @@
+"""
+Conditional Dataset loader with labels
+"""
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+import os
+
+
+class ConditionalImageDataset(Dataset):
+    def __init__(self, data_path, label_path, transform=None, label_stats=None):
+        self.data = np.load(data_path)
+        self.labels = np.load(label_path)
+        self.transform = transform
+        self.label_stats = label_stats
+
+        assert len(self.data) == len(self.labels), f"Data and labels length mismatch! {len(self.data)} vs {len(self.labels)}"
+
+        print(f"Loaded {len(self.data)} images | Image shape: {self.data.shape[1:]} | Label shape: {self.labels.shape[1:]}")
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        img = torch.from_numpy(self.data[idx]).float()
+        label = torch.from_numpy(self.labels[idx]).float()
+
+        # Normalize image to [-1, 1]
+        img = img * 2.0 - 1.0
+
+        # Normalize labels
+        if self.label_stats is not None:
+            label = (label - self.label_stats['mean']) / self.label_stats['std']
+
+        if img.dim() == 2:
+            img = img.unsqueeze(0)
+
+        return img, label
+
+
+def get_conditional_dataloaders(
+    data_dir='./data/params_2',
+    batch_size=8,
+    num_workers=4,
+    pin_memory=True,
+    normalize_labels=True
+):
+    is_6param = 'params_6' in data_dir
+
+    if is_6param:
+        train_data = os.path.join(data_dir, 'train_LH_6.npy')
+        val_data = os.path.join(data_dir, 'val_LH_6.npy')
+        test_data = os.path.join(data_dir, 'test_LH_6.npy')
+        train_labels = os.path.join(data_dir, 'train_labels_LH.npy')
+        val_labels = os.path.join(data_dir, 'val_labels_LH.npy')
+        test_labels = os.path.join(data_dir, 'test_labels_LH.npy')
+    else:
+        train_data = os.path.join(data_dir, 'train_LH.npy')
+        val_data = os.path.join(data_dir, 'val_LH.npy')
+        test_data = os.path.join(data_dir, 'test_LH.npy')
+        train_labels = os.path.join(data_dir, 'train_labels_LH_2.npy')
+        val_labels = os.path.join(data_dir, 'val_labels_LH_2.npy')
+        test_labels = os.path.join(data_dir, 'test_labels_LH_2.npy')
+
+    print(f"Loading dataset from {data_dir} ({'6-param' if is_6param else '2-param'})")
+
+    # Label normalization stats
+    label_stats = None
+    if normalize_labels:
+        train_labels_array = np.load(train_labels)
+        label_mean = train_labels_array.mean(axis=0)
+        label_std = train_labels_array.std(axis=0)
+        label_std = np.where(label_std == 0, 1.0, label_std)  # guard against zero-variance labels
+        label_stats = {'mean': torch.from_numpy(label_mean).float(), 'std': torch.from_numpy(label_std).float()}
+        print(f"Label normalization -> mean={label_mean}, std={label_std}")
+
+    train_dataset = ConditionalImageDataset(train_data, train_labels, label_stats=label_stats)
+    val_dataset   = ConditionalImageDataset(val_data,   val_labels,   label_stats=label_stats)
+    test_dataset  = ConditionalImageDataset(test_data,  test_labels,  label_stats=label_stats)
+
+    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,  num_workers=num_workers, pin_memory=pin_memory, drop_last=True)
+    val_loader   = DataLoader(val_dataset,   batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory, drop_last=False)
+    test_loader  = DataLoader(test_dataset,  batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=pin_memory, drop_last=False)
+
+    return train_loader, val_loader, test_loader

src/diffusion_conditional.py

@@ -0,0 +1,172 @@
+"""
+Diffusion Process Implementation (DDPM + DDIM)
+
+Changes from original:
+- Fixed q_posterior_mean_variance return value count (was 3, caller expected 4)
+- Fixed DDIM sigma/dir_xt formula inconsistency
+- Made GaussianDiffusion an nn.Module with registered buffers for proper device handling
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class GaussianDiffusion(nn.Module):
+    def __init__(self, timesteps=1500, beta_start=1e-4, beta_end=0.02, schedule_type="linear"):
+        super().__init__()
+        self.timesteps = timesteps
+
+        if schedule_type == "linear":
+            betas = torch.linspace(beta_start, beta_end, timesteps)
+        elif schedule_type == "cosine":
+            betas = self._cosine_beta_schedule(timesteps)
+        else:
+            raise ValueError(f"Unknown schedule: {schedule_type}")
+
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, dim=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+
+        # Register all schedule tensors as buffers so they move with .to(device)
+        self.register_buffer('betas', betas)
+        self.register_buffer('alphas', alphas)
+        self.register_buffer('alphas_cumprod', alphas_cumprod)
+        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1.0 - alphas_cumprod))
+
+        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        self.register_buffer('posterior_variance', posterior_variance)
+        self.register_buffer('posterior_log_variance_clipped', torch.log(torch.clamp(posterior_variance, min=1e-20)))
+        self.register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod))
+        self.register_buffer('posterior_mean_coef2', (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod))
+
+        # Precompute reciprocals used in _predict_xstart_from_noise (avoids recomputation per step)
+        self.register_buffer('recip_sqrt_alphas_cumprod', 1.0 / torch.sqrt(alphas_cumprod))
+        self.register_buffer('sqrt_recip_minus_one', torch.sqrt(1.0 / alphas_cumprod - 1.0))
+
+    def _cosine_beta_schedule(self, timesteps, s=0.008):
+        steps = timesteps + 1
+        x = torch.linspace(0, timesteps, steps)
+        alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
+        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+        return torch.clip(betas, 0.0001, 0.9999)
+
+    def q_sample(self, x_start, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+        sqrt_alphas_cumprod_t = self._extract(self.sqrt_alphas_cumprod, t, x_start.shape)
+        sqrt_one_minus_alphas_cumprod_t = self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
+
+    def p_mean_variance(self, model, x_t, t, labels, clip_denoised=True):
+        pred_noise = model(x_t, t, labels)
+        x_start = self._predict_xstart_from_noise(x_t, t, pred_noise)
+        if clip_denoised:
+            x_start = torch.clamp(x_start, -1.0, 1.0)
+        # FIX: q_posterior_mean_variance returns 3 values, not 4
+        model_mean, posterior_variance, posterior_log_variance = self.q_posterior_mean_variance(x_start, x_t, t)
+        return model_mean, posterior_variance, posterior_log_variance, x_start
+
+    def _predict_xstart_from_noise(self, x_t, t, noise):
+        return (
+            self._extract(self.recip_sqrt_alphas_cumprod, t, x_t.shape) * x_t -
+            self._extract(self.sqrt_recip_minus_one, t, x_t.shape) * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        posterior_mean = (
+            self._extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
+            self._extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = self._extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance = self._extract(self.posterior_log_variance_clipped, t, x_t.shape)
+        return posterior_mean, posterior_variance, posterior_log_variance
+
+    def p_sample(self, model, x_t, t, labels):
+        model_mean, _, model_log_variance, _ = self.p_mean_variance(model, x_t, t, labels)
+        noise = torch.randn_like(x_t)
+        nonzero_mask = ((t != 0).float().view(-1, *([1] * (len(x_t.shape) - 1))))
+        return model_mean + nonzero_mask * torch.exp(0.5 * model_log_variance) * noise
+
+    def ddim_sample_step(self, model, x_t, t, t_next, labels, eta=0.0):
+        pred_noise = model(x_t, t, labels)
+        alpha_t = self._extract(self.alphas_cumprod, t, x_t.shape)
+        alpha_t_next = self._extract(self.alphas_cumprod, t_next, x_t.shape) if t_next[0] >= 0 else torch.ones_like(alpha_t)
+
+        x0_pred = (x_t - torch.sqrt(1 - alpha_t) * pred_noise) / torch.sqrt(alpha_t)
+        x0_pred = torch.clamp(x0_pred, -1.0, 1.0)
+
+        # FIX: Consistent DDIM sigma computation
+        # sigma^2 = eta^2 * (1 - alpha_{t-1}) / (1 - alpha_t) * (1 - alpha_t / alpha_{t-1})
+        sigma_sq = eta**2 * (1 - alpha_t_next) / (1 - alpha_t) * (1 - alpha_t / alpha_t_next) if eta > 0 else 0
+        sigma_t = torch.sqrt(torch.clamp(sigma_sq, min=0)) if eta > 0 else 0
+
+        # dir_xt uses the same sigma^2
+        dir_xt = torch.sqrt(torch.clamp(1 - alpha_t_next - sigma_sq, min=0)) * pred_noise
+
+        noise = torch.randn_like(x_t) if eta > 0 else 0
+        return torch.sqrt(alpha_t_next) * x0_pred + dir_xt + sigma_t * noise
+
+    def sample(self, model, labels, channels, height, width, device, progress=False, use_ddim=True, ddim_steps=50, eta=0.0):
+        batch_size = labels.shape[0]
+        img = torch.randn((batch_size, channels, height, width), device=device)
+
+        if use_ddim:
+            skip = self.timesteps // ddim_steps
+            seq = list(range(0, self.timesteps, skip))
+            seq_next = [-1] + seq[:-1]
+            seq_iter = reversed(list(zip(seq, seq_next)))
+            if progress:
+                from tqdm import tqdm
+                seq_iter = tqdm(seq_iter, desc=f'DDIM Sampling ({ddim_steps} steps)', total=len(seq))
+            for i, j in seq_iter:
+                t = torch.full((batch_size,), i, device=device, dtype=torch.long)
+                t_next = torch.full((batch_size,), j, device=device, dtype=torch.long)
+                img = self.ddim_sample_step(model, img, t, t_next, labels, eta)
+        else:
+            if progress:
+                from tqdm import tqdm
+                timesteps_iter = tqdm(reversed(range(self.timesteps)), total=self.timesteps)
+            else:
+                timesteps_iter = reversed(range(self.timesteps))
+            for i in timesteps_iter:
+                t = torch.full((batch_size,), i, device=device, dtype=torch.long)
+                img = self.p_sample(model, img, t, labels)
+        return img
+
+    def training_losses(self, model, x_start, labels, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise)
+        pred_noise = model(x_t, t, labels)
+        return F.mse_loss(pred_noise, noise, reduction='none').mean(dim=list(range(1, len(pred_noise.shape))))
+
+    def _extract(self, a, t, x_shape):
+        batch_size = t.shape[0]
+        out = a.gather(0, t)
+        return out.reshape(batch_size, *((1,) * (len(x_shape) - 1)))
+
+
+class ConditionalDiffusionModel(nn.Module):
+    def __init__(self, unet, diffusion_process):
+        super().__init__()
+        self.unet = unet
+        self.diffusion = diffusion_process
+
+    def forward(self, x, t, labels):
+        return self.unet(x, t, labels)
+
+    def get_loss(self, x, labels, noise=None):
+        batch_size = x.shape[0]
+        device = x.device
+        t = torch.randint(0, self.diffusion.timesteps, (batch_size,), device=device).long()
+        return self.diffusion.training_losses(self, x, labels, t, noise=noise).mean()
+
+    def sample(self, labels, channels, height, width, device, progress=False, use_ddim=True, ddim_steps=50, eta=0.0):
+        self.eval()
+        with torch.no_grad():
+            return self.diffusion.sample(self, labels, channels, height, width, device, progress, use_ddim, ddim_steps, eta)

src/evaluate_conditional.py

@@ -0,0 +1,417 @@
+"""
+Evaluate Conditional Diffusion Model (2-label: Omega_m, sigma_8)
+
+Usage:
+    python evaluate_conditional.py --checkpoint outputs_conditional_YYYYMMDD_HHMMSS/checkpoints/best_model.pt
+
+Changes from original:
+- Loads args.json (saved by improved training script) for robust config parsing
+- Falls back to args.txt parsing if JSON not available
+- Vectorized power spectrum calculation (~100x speedup)
+- Added weights_only parameter to torch.load
+"""
+
+import argparse
+import ast
+import json
+import os
+from pathlib import Path
+from typing import Dict, Tuple
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+from diffusion_conditional import GaussianDiffusion, ConditionalDiffusionModel
+from unet_conditional import ConditionalUNet
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Evaluate conditional 2-label diffusion model")
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        required=True,
+        help="Path to trained checkpoint (e.g. outputs_conditional_*/checkpoints/best_model.pt)",
+    )
+    parser.add_argument(
+        "--training_args",
+        type=str,
+        default=None,
+        help="Path to args.json or args.txt from training (auto-detected if not provided)",
+    )
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default="./data/params_2",
+        help="Directory containing the CAMELS LH dataset (default matches repo structure)",
+    )
+    parser.add_argument(
+        "--split",
+        type=str,
+        default="test",
+        choices=["train", "val", "test"],
+        help="Which split to use for real images",
+    )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        default=8,
+        help="Number of examples to show in the comparison grid",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed for reproducibility",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="evaluation_outputs",
+        help="Where to save plots and results",
+    )
+    parser.add_argument(
+        "--ddim_steps",
+        type=int,
+        default=50,
+        help="Number of DDIM sampling steps",
+    )
+    return parser.parse_args()
+
+
+def load_training_config(path: str) -> Dict:
+    """Load training configuration. Prefers JSON, falls back to txt parsing."""
+    # Try JSON first (written by improved training script)
+    json_path = path.replace('.txt', '.json') if path.endswith('.txt') else path
+    if json_path.endswith('.json') and os.path.isfile(json_path):
+        with open(json_path, 'r') as f:
+            return json.load(f)
+
+    # Fall back to txt parsing
+    if not os.path.isfile(path):
+        raise FileNotFoundError(f"Training args file not found: {path}")
+
+    config = {}
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line or ":" not in line:
+                continue
+            key, value = line.split(":", 1)
+            key = key.strip()
+            value = value.strip()
+
+            if value.startswith("[") and value.endswith("]"):
+                try:
+                    config[key] = ast.literal_eval(value)
+                except (ValueError, SyntaxError):
+                    config[key] = value
+            elif value.isdigit():
+                config[key] = int(value)
+            elif value.replace(".", "", 1).replace("e-", "", 1).replace("e", "", 1).isdigit():
+                config[key] = float(value)
+            else:
+                config[key] = value
+
+    return config
+
+
+def _detect_label_suffix(data_dir: Path) -> str:
+    """Detect whether this is a 2-param or 6-param dataset."""
+    if (data_dir / "train_labels_LH_2.npy").exists():
+        return "_2"
+    elif (data_dir / "train_labels_LH.npy").exists():
+        return ""
+    else:
+        raise FileNotFoundError(f"No label files found in {data_dir}")
+
+
+def _detect_image_suffix(data_dir: Path) -> str:
+    """Detect whether images use _6 suffix (6-param) or not."""
+    if (data_dir / "train_LH.npy").exists():
+        return ""
+    elif (data_dir / "train_LH_6.npy").exists():
+        return "_6"
+    else:
+        raise FileNotFoundError(f"No image files found in {data_dir}")
+
+
+def load_label_stats(data_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
+    """Load mean and std from training labels (used for normalization)."""
+    suffix = _detect_label_suffix(data_dir)
+    labels_path = data_dir / f"train_labels_LH{suffix}.npy"
+    labels = np.load(labels_path)
+    mean, std = labels.mean(axis=0), labels.std(axis=0)
+    std = np.where(std == 0, 1.0, std)  # guard against zero-variance labels
+    return mean, std
+
+
+def load_split(data_dir: Path, split: str) -> Tuple[np.ndarray, np.ndarray]:
+    """Load images and labels for a given split."""
+    img_suffix = _detect_image_suffix(data_dir)
+    label_suffix = _detect_label_suffix(data_dir)
+
+    image_path = data_dir / f"{split}_LH{img_suffix}.npy"
+    label_path = data_dir / f"{split}_labels_LH{label_suffix}.npy"
+
+    if not image_path.exists():
+        raise FileNotFoundError(f"Image file not found: {image_path}")
+    if not label_path.exists():
+        raise FileNotFoundError(f"Label file not found: {label_path}")
+
+    images = np.load(image_path).astype(np.float32)
+    labels = np.load(label_path).astype(np.float32)
+    return images, labels
+
+
+def build_model(config: Dict, device: torch.device) -> ConditionalDiffusionModel:
+    """Rebuild the exact same model architecture used during training."""
+    unet = ConditionalUNet(
+        in_channels=1,
+        out_channels=1,
+        label_dim=int(config.get("label_dim", 2)),
+        base_channels=int(config.get("base_channels", 64)),
+        channel_multipliers=config.get("channel_multipliers", [1, 2, 4, 8]),
+        attention_levels=config.get("attention_levels", [2, 3]),
+        dropout=float(config.get("dropout", 0.1)),
+    )
+
+    diffusion = GaussianDiffusion(
+        timesteps=int(config.get("timesteps", 1500)),
+        beta_start=float(config.get("beta_start", 1e-4)),
+        beta_end=float(config.get("beta_end", 0.02)),
+        schedule_type=config.get("schedule_type", "linear"),
+    )
+
+    return ConditionalDiffusionModel(unet, diffusion).to(device)
+
+
+def load_checkpoint(model: ConditionalDiffusionModel, checkpoint_path: str, device: torch.device):
+    """Load model weights from checkpoint."""
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    state_dict = checkpoint["model_state_dict"] if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint else checkpoint
+
+    # If EMA weights are available, use them (they are the better weights)
+    if isinstance(checkpoint, dict) and "ema_shadow" in checkpoint:
+        print("Loading EMA shadow weights from checkpoint")
+        ema_shadow = checkpoint["ema_shadow"]
+        current_state = model.state_dict()
+        for name, param in ema_shadow.items():
+            if name in current_state:
+                current_state[name] = param
+        model.load_state_dict(current_state)
+    else:
+        model.load_state_dict(state_dict)
+
+    model.eval()
+    print(f"Loaded checkpoint: {checkpoint_path}")
+
+
+def PowerSpectrum(box: np.ndarray, N: int, dl: float) -> Tuple[np.ndarray, np.ndarray]:
+    """Vectorized 2D power spectrum computation."""
+    FT_box = np.fft.fftn(box, norm="ortho")
+    k = 2 * np.pi * np.fft.fftfreq(N, dl)
+    dk_val = 2 * np.pi / (N * dl)
+
+    # Vectorized: compute k magnitudes and bin indices for all pixels at once
+    ki, kj = np.meshgrid(k, k, indexing='ij')
+    kbar = np.sqrt(ki**2 + kj**2)
+    n_bins = N // 2  # only bins up to Nyquist frequency
+    t_idx = np.round(kbar / dk_val).astype(int)
+
+    # Mask out modes beyond Nyquist to avoid bin contamination
+    valid = t_idx < n_bins
+    power = (FT_box * np.conj(FT_box)).real
+
+    pk = np.zeros(n_bins)
+    count = np.zeros(n_bins)
+    np.add.at(pk, t_idx[valid], power[valid])
+    np.add.at(count, t_idx[valid], 1)
+
+    pk /= np.where(count == 0, 1, count)
+    pk *= dl**2
+    dk = np.arange(n_bins) * dk_val
+    return dk, pk
+
+
+def calculate_pdf_batch(images: np.ndarray, log_nhi_min=14.0, log_nhi_max=22.0, n_bins=100):
+    images_01 = np.clip(images, 0.0, 1.0)
+    log_nhi_bins = np.linspace(log_nhi_min, log_nhi_max, n_bins)
+    bin_centers = 0.5 * (log_nhi_bins[:-1] + log_nhi_bins[1:])
+
+    pdfs = []
+    for img in images_01:
+        log_nhi_values = log_nhi_min + (log_nhi_max - log_nhi_min) * img.reshape(-1)
+        hist, _ = np.histogram(log_nhi_values, bins=log_nhi_bins, density=True)
+        pdfs.append(hist)
+
+    pdf_array = np.stack(pdfs)
+    return bin_centers, pdf_array.mean(axis=0), pdf_array.std(axis=0)
+
+
+def calculate_power_spectrum_batch(images: np.ndarray, box_size: float = 25.0):
+    N = images.shape[-1]
+    dl = box_size / N
+
+    # Compute k-values once, then reuse for all images
+    dk, _ = PowerSpectrum(images[0], N=N, dl=dl)
+    power_spectra = [PowerSpectrum(img, N=N, dl=dl)[1] for img in images]
+    power_array = np.stack(power_spectra)
+    return dk, power_array.mean(axis=0), power_array.std(axis=0)
+
+
+def prepare_labels_for_model(labels: np.ndarray, mean: np.ndarray, std: np.ndarray) -> torch.Tensor:
+    normalized = (labels - mean) / std
+    return torch.from_numpy(normalized).float()
+
+
+def from_model_output(samples: torch.Tensor) -> np.ndarray:
+    arrays = samples.cpu().numpy()
+    return np.clip((arrays + 1.0) / 2.0, 0.0, 1.0)[:, 0, :, :]
+
+
+def plot_image_grid(generated, real, labels, output_path: Path, num_samples=8):
+    num = min(num_samples, generated.shape[0])
+    fig, axes = plt.subplots(num, 2, figsize=(6, 3 * num))
+    if num == 1:
+        axes = np.expand_dims(axes, axis=0)
+
+    for i in range(num):
+        label_str = ", ".join(f"{v:.3f}" for v in labels[i])
+        axes[i, 0].imshow(generated[i], cmap="magma", origin="lower")
+        axes[i, 0].set_title(f"Generated\n{label_str}")
+        axes[i, 0].axis("off")
+
+        axes[i, 1].imshow(real[i], cmap="magma", origin="lower")
+        axes[i, 1].set_title("Real")
+        axes[i, 1].axis("off")
+
+    plt.tight_layout()
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def plot_mean_std(x, mean_real, std_real, mean_gen, std_gen, xlabel, ylabel, title, output_path: Path, yscale="linear"):
+    fig, ax = plt.subplots(figsize=(10, 6))
+    ax.plot(x, mean_real, label="Real mean", color="tab:blue", linewidth=2)
+    ax.plot(x, mean_gen, label="Generated mean", color="tab:orange", linewidth=2)
+
+    ax.fill_between(x, mean_real - std_real, mean_real + std_real, color="tab:blue", alpha=0.15, label="Real +/-1s")
+    ax.fill_between(x, mean_real - 3*std_real, mean_real + 3*std_real, color="tab:blue", alpha=0.05)
+
+    ax.fill_between(x, mean_gen - std_gen, mean_gen + std_gen, color="tab:orange", alpha=0.15, label="Generated +/-1s")
+    ax.fill_between(x, mean_gen - 3*std_gen, mean_gen + 3*std_gen, color="tab:orange", alpha=0.05)
+
+    ax.set_xlabel(xlabel)
+    ax.set_ylabel(ylabel)
+    ax.set_title(title)
+    ax.set_yscale(yscale)
+    ax.legend()
+    ax.grid(alpha=0.3)
+    fig.tight_layout()
+    fig.savefig(output_path, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+
+
+def main():
+    args = parse_args()
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load training config
+    if args.training_args is None:
+        # Try JSON first, then txt
+        possible_json = list(Path(".").glob("outputs_conditional_*/args.json"))
+        possible_txt = list(Path(".").glob("outputs_conditional_*/args.txt"))
+        possible = possible_json + possible_txt
+        if possible:
+            args.training_args = str(max(possible, key=os.path.getctime))
+            print(f"Auto-detected training args: {args.training_args}")
+        else:
+            raise FileNotFoundError("Please provide --training_args path to your training args.json or args.txt")
+
+    config = load_training_config(args.training_args)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = build_model(config, device)
+    load_checkpoint(model, args.checkpoint, device)
+
+    # Load data
+    data_dir = Path(args.data_dir)
+    images_split, labels_split = load_split(data_dir, args.split)
+    label_mean, label_std = load_label_stats(data_dir)
+
+    # Select random samples
+    num_select = min(100, len(images_split))
+    indices = np.random.choice(len(images_split), num_select, replace=False)
+
+    real_images = images_split[indices]
+    original_labels = labels_split[indices]
+
+    # Generate samples in batches
+    batch_size = min(8, num_select)
+    generated_list = []
+
+    print(f"Generating {num_select} samples (batch size = {batch_size})...")
+    for i in range(0, num_select, batch_size):
+        batch_labels = original_labels[i:i+batch_size]
+        batch_labels_tensor = prepare_labels_for_model(batch_labels, label_mean, label_std).to(device)
+
+        with torch.no_grad():
+            batch_gen = model.sample(
+                labels=batch_labels_tensor,
+                channels=1,
+                height=real_images.shape[-2],
+                width=real_images.shape[-1],
+                device=device,
+                progress=False,
+                use_ddim=True,
+                ddim_steps=args.ddim_steps,
+            )
+        generated_list.append(from_model_output(batch_gen))
+        print(f"  Batch {i//batch_size + 1}/{(num_select+batch_size-1)//batch_size} done")
+
+    generated_images = np.concatenate(generated_list, axis=0)
+
+    # Plots
+    plot_image_grid(generated_images, real_images, original_labels,
+                    output_dir / "real_vs_generated.png", num_samples=args.num_samples)
+
+    # PDF
+    bin_centers, mean_pdf_real, std_pdf_real = calculate_pdf_batch(real_images)
+    _, mean_pdf_gen, std_pdf_gen = calculate_pdf_batch(generated_images)
+    plot_mean_std(bin_centers, mean_pdf_real, std_pdf_real, mean_pdf_gen, std_pdf_gen,
+                  "log N_HI [cm^-2]", "PDF", "Column Density PDF", output_dir / "pdf_mean_std.png")
+
+    # Power Spectrum (skip k=0 DC component for log-scale plotting)
+    dk, mean_pk_real, std_pk_real = calculate_power_spectrum_batch(real_images)
+    _, mean_pk_gen, std_pk_gen = calculate_power_spectrum_batch(generated_images)
+    plot_mean_std(dk[1:], mean_pk_real[1:], std_pk_real[1:], mean_pk_gen[1:], std_pk_gen[1:],
+                  "k [h/Mpc]", "P(k)", "Power Spectrum", output_dir / "power_spectrum_mean_std.png", yscale="log")
+
+    # Save numerical results
+    np.savez(
+        output_dir / "evaluation_data.npz",
+        indices=indices,
+        labels_original=original_labels,
+        bin_centers=bin_centers,
+        mean_pdf_real=mean_pdf_real, std_pdf_real=std_pdf_real,
+        mean_pdf_gen=mean_pdf_gen, std_pdf_gen=std_pdf_gen,
+        dk=dk,
+        mean_pk_real=mean_pk_real, std_pk_real=std_pk_real,
+        mean_pk_gen=mean_pk_gen, std_pk_gen=std_pk_gen,
+    )
+
+    print(f"\nEvaluation complete!")
+    print(f"   Plots saved to: {output_dir}")
+    print(f"   Numerical data saved to: {output_dir}/evaluation_data.npz")
+
+
+if __name__ == "__main__":
+    main()

src/unet_conditional.py

@@ -0,0 +1,179 @@
+"""
+Conditional U-Net Architecture for Diffusion Model
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+
+class TimeEmbedding(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, time):
+        device = time.device
+        half_dim = self.dim // 2
+        embeddings = math.log(10000) / (half_dim - 1)
+        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
+        embeddings = time[:, None] * embeddings[None, :]
+        return torch.cat([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
+
+
+class LabelEmbedding(nn.Module):
+    def __init__(self, label_dim, emb_dim):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(label_dim, emb_dim),
+            nn.SiLU(),
+            nn.Linear(emb_dim, emb_dim)
+        )
+
+    def forward(self, labels):
+        return self.mlp(labels)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, time_emb_dim, dropout=0.1):
+        super().__init__()
+        self.conv1 = nn.Sequential(
+            nn.GroupNorm(8, in_channels),
+            nn.SiLU(),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        )
+        self.time_emb = nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, out_channels))
+        self.conv2 = nn.Sequential(
+            nn.GroupNorm(8, out_channels),
+            nn.SiLU(),
+            nn.Dropout(dropout),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+        )
+        self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1) if in_channels != out_channels else nn.Identity()
+
+    def forward(self, x, emb):
+        h = self.conv1(x)
+        h = h + self.time_emb(emb)[:, :, None, None]
+        h = self.conv2(h)
+        return h + self.shortcut(x)
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, channels, num_heads=4):
+        super().__init__()
+        self.channels = channels
+        self.num_heads = num_heads
+        self.norm = nn.GroupNorm(8, channels)
+        self.qkv = nn.Conv2d(channels, channels * 3, kernel_size=1)
+        self.proj = nn.Conv2d(channels, channels, kernel_size=1)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        h = self.norm(x)
+        q, k, v = self.qkv(h).chunk(3, dim=1)
+        head_dim = C // self.num_heads
+        q = q.view(B, self.num_heads, head_dim, H*W).transpose(2, 3)
+        k = k.view(B, self.num_heads, head_dim, H*W).transpose(2, 3)
+        v = v.view(B, self.num_heads, head_dim, H*W).transpose(2, 3)
+
+        h = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
+        h = h.transpose(2, 3).reshape(B, C, H, W)
+        return x + self.proj(h)
+
+
+class ConditionalUNet(nn.Module):
+    def __init__(self, in_channels=1, out_channels=1, label_dim=2,
+                 base_channels=64, channel_multipliers=(1,2,4,8),
+                 attention_levels=(2,3), dropout=0.1, time_emb_dim=256, label_emb_dim=256):
+        super().__init__()
+        self.label_dim = label_dim
+
+        self.time_embedding = TimeEmbedding(time_emb_dim)
+        self.time_mlp = nn.Sequential(nn.Linear(time_emb_dim, time_emb_dim*4), nn.SiLU(), nn.Linear(time_emb_dim*4, time_emb_dim))
+
+        self.label_embedding = LabelEmbedding(label_dim, label_emb_dim)
+        self.combined_emb_dim = time_emb_dim + label_emb_dim
+        self.combined_mlp = nn.Sequential(nn.Linear(self.combined_emb_dim, time_emb_dim*4), nn.SiLU(), nn.Linear(time_emb_dim*4, time_emb_dim))
+
+        self.conv_in = nn.Conv2d(in_channels, base_channels, kernel_size=3, padding=1)
+
+        self.down_blocks = nn.ModuleList()
+        channels = [base_channels]
+        now_channels = base_channels
+
+        for i, mult in enumerate(channel_multipliers):
+            out_ch = base_channels * mult
+            for _ in range(2):
+                self.down_blocks.append(ResidualBlock(now_channels, out_ch, time_emb_dim, dropout))
+                if i in attention_levels:
+                    self.down_blocks.append(AttentionBlock(out_ch))
+                now_channels = out_ch
+                channels.append(now_channels)
+            if i != len(channel_multipliers) - 1:
+                self.down_blocks.append(nn.Conv2d(now_channels, now_channels, kernel_size=3, stride=2, padding=1))
+                channels.append(now_channels)
+
+        self.middle = nn.ModuleList([
+            ResidualBlock(now_channels, now_channels, time_emb_dim, dropout),
+            AttentionBlock(now_channels),
+            ResidualBlock(now_channels, now_channels, time_emb_dim, dropout)
+        ])
+
+        self.up_blocks = nn.ModuleList()
+        for i, mult in reversed(list(enumerate(channel_multipliers))):
+            out_ch = base_channels * mult
+            for _ in range(3):
+                self.up_blocks.append(ResidualBlock(now_channels + channels.pop(), out_ch, time_emb_dim, dropout))
+                if i in attention_levels:
+                    self.up_blocks.append(AttentionBlock(out_ch))
+                now_channels = out_ch
+            if i != 0:
+                self.up_blocks.append(nn.ConvTranspose2d(now_channels, now_channels, kernel_size=4, stride=2, padding=1))
+
+        self.conv_out = nn.Sequential(
+            nn.GroupNorm(8, now_channels),
+            nn.SiLU(),
+            nn.Conv2d(now_channels, out_channels, kernel_size=3, padding=1)
+        )
+
+    def forward(self, x, t, labels=None):
+        t_emb = self.time_embedding(t)
+        t_emb = self.time_mlp(t_emb)
+
+        if labels is not None:
+            label_emb = self.label_embedding(labels)
+            combined = torch.cat([t_emb, label_emb], dim=-1)
+            emb = self.combined_mlp(combined)
+        else:
+            emb = t_emb
+
+        h = self.conv_in(x)
+        skips = [h]
+
+        for module in self.down_blocks:
+            if isinstance(module, ResidualBlock):
+                h = module(h, emb)
+                skips.append(h)
+            elif isinstance(module, AttentionBlock):
+                h = module(h)
+            else:
+                h = module(h)
+                skips.append(h)
+
+        for module in self.middle:
+            if isinstance(module, ResidualBlock):
+                h = module(h, emb)
+            else:
+                h = module(h)
+
+        for module in self.up_blocks:
+            if isinstance(module, ResidualBlock):
+                h = torch.cat([h, skips.pop()], dim=1)
+                h = module(h, emb)
+            elif isinstance(module, AttentionBlock):
+                h = module(h)
+            else:
+                h = module(h)
+
+        return self.conv_out(h)