Upload 6-parameter conditional DDPM (HI emulation, CAMELS LH params_6, best checkpoint) with full training/eval/posterior toolchain

README.md

@@ -26,17 +26,58 @@ This is the **best-validation checkpoint** from the training run under
 
 ## 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) |
-| `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 |
 
+**`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.
+`/scratch/mrpcol001/Diffusion_job/data/LH_data/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

cross_model/README.md

@@ -0,0 +1,17 @@
+# Cross-model scripts
+
+These posterior-inference and comparison scripts use BOTH the
+2-parameter and 6-parameter DDPM checkpoints. Their default
+paths assume the original cluster layout:
+
+    Models/
+      2param_DDPM_HI_Emulation/   <- code
+      6param_ddpm_hi_lh6/         <- code
+      notebook_model_weights/
+        2param_epoch200/          <- checkpoint + args.json
+        6param_best/              <- checkpoint + args.json
+
+When running these scripts from a local download of this HF repo,
+pass `--bundle-2param`, `--bundle-6param`, `--data-2param`,
+`--data-6param` (etc.) to point at the locations where you placed
+the weights and the CAMELS LH data. See each script's `--help`.

cross_model/check_poster_env.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+"""One-shot env check for poster.py — logs NDJSON for debug session."""
+import json
+import os
+import time
+
+_LOG = "/scratch/mrpcol001/Diffusion_job/Models/.cursor/debug-a1359c.log"
+
+
+def _log(hypothesis_id: str, location: str, message: str, data: dict) -> None:
+    payload = {
+        "sessionId": "a1359c",
+        "runId": os.environ.get("DEBUG_POSTER_RUN", "pre-fix"),
+        "hypothesisId": hypothesis_id,
+        "location": location,
+        "message": message,
+        "data": data,
+        "timestamp": int(time.time() * 1000),
+    }
+    os.makedirs(os.path.dirname(_LOG), exist_ok=True)
+    with open(_LOG, "a", encoding="utf-8") as f:
+        f.write(json.dumps(payload) + "\n")
+
+
+def main() -> None:
+    # region agent log
+    root = "/scratch/mrpcol001/Diffusion_job/Models"
+    cwd = os.getcwd()
+    poster_path = os.path.join(root, "poster.py")
+    poster_ci = os.path.join(root, "Poster.py")
+    _log(
+        "H2",
+        "check_poster_env.py:main",
+        "cwd vs expected Models root",
+        {"cwd": cwd, "root": root, "cwd_equals_root": os.path.abspath(cwd) == os.path.abspath(root)},
+    )
+    _log(
+        "H1",
+        "check_poster_env.py:main",
+        "poster.py presence",
+        {
+            "poster_py_exists": os.path.isfile(poster_path),
+            "poster_path": poster_path,
+            "size_if_exists": os.path.getsize(poster_path) if os.path.isfile(poster_path) else None,
+        },
+    )
+    _log(
+        "H3",
+        "check_poster_env.py:main",
+        "case variant",
+        {"Poster_py_exists": os.path.isfile(poster_ci)},
+    )
+    try:
+        names = sorted(os.listdir(root))
+    except OSError as e:
+        names = []
+        list_err = str(e)
+    else:
+        list_err = None
+    poster_like = [n for n in names if "poster" in n.lower()]
+    _log(
+        "H4",
+        "check_poster_env.py:main",
+        "Models directory poster-related names",
+        {"list_error": list_err, "poster_like_filenames": poster_like, "total_entries": len(names)},
+    )
+    _log(
+        "H5",
+        "check_poster_env.py:main",
+        "alternate runnable scripts hint",
+        {"scripts_dir_exists": os.path.isdir(os.path.join(root, "scripts"))},
+    )
+    # endregion agent log
+    print("check_poster_env: wrote logs to", _LOG)
+
+
+if __name__ == "__main__":
+    main()

cross_model/compare_posterior_inference.py

@@ -0,0 +1,699 @@
+#!/usr/bin/env python3
+"""
+compare_posterior_inference.py
+==============================
+Side-by-side corrected surrogate posteriors on (Omega_m, sigma_8) for
+DDPM-2 and DDPM-6 using a JOINT P(k) + log-N_HI PDF Gaussian likelihood.
+
+Why a joint summary statistic?
+------------------------------
+P(k) alone leaves Omega_m and sigma_8 strongly degenerate, which is why
+prior runs returned sigma_8 ~ 0.80 +/- 0.12 regardless of truth. The
+column-density PDF carries information that P(k) misses (it is sensitive
+to non-Gaussian amplitude features), so combining the two breaks the
+degeneracy along the S_8 direction.
+
+What this script does
+---------------------
+1. Loads both DDPM-2 (epoch 200) and DDPM-6 (best) bundles from
+   notebook_model_weights/.
+2. Calibrates per-model sigma_pk and sigma_pdf from validation-set DDPM
+   pair scatter (no hard-coded noise).
+3. For each anchor in the test split:
+     - DDPM-2: direct (Omega_m, sigma_8) grid posterior.
+     - DDPM-6: Monte-Carlo marginalisation over the four astrophysical
+       nuisance parameters with a uniform LHS prior.
+   Both posteriors use the JOINT P(k) + PDF likelihood.
+4. Saves per-anchor posterior arrays as .npz so plots can be re-rendered
+   cheaply later.
+5. Emits a single comparison figure: rows = anchors, columns =
+   DDPM-2 | DDPM-6, with 68/95% credible contours, true value, posterior
+   mean, and posterior summary annotation.
+
+Defaults
+--------
+  --grid 30           --ddim-steps 50      --batch-size 8
+  --n-pk-samples 8    --n-marg-samples 20  --n-anchors 4
+
+Quick smoke test
+----------------
+  python compare_posterior_inference.py --grid 16 --n-pk-samples 4 \\
+      --n-marg-samples 5 --n-anchors 2
+"""
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+MODELS_ROOT = Path(__file__).resolve().parent
+CODE_6 = MODELS_ROOT / "6param_ddpm_hi_lh6"
+if str(CODE_6.resolve()) not in sys.path:
+    sys.path.insert(0, str(CODE_6))
+
+import evaluate_conditional as ec   # noqa: E402
+import eval_model as em             # noqa: E402
+
+
+# =============================================================================
+# 1.  SUMMARY STATISTICS  (P(k) and log-N_HI PDF, per map)
+# =============================================================================
+
+def per_map_log_pk(
+    imgs: np.ndarray,
+    box_size: float = 25.0,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """log P(k) for a batch of [0,1] maps. Returns (dk, valid_mask, log_pks).
+
+    log_pks has shape (N, valid_mask.sum()).
+    """
+    dk, pks = em.per_map_power_spectra_log(imgs, box_size)
+    valid = dk > 0
+    return dk, valid, np.log(pks[:, valid] + 1e-30)
+
+
+def per_map_log_pdf(
+    imgs: np.ndarray,
+    log_nhi_min: float = 14.0,
+    log_nhi_max: float = 22.0,
+    n_bins: int = 100,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """log column-density PDF for a batch of [0,1] maps.
+
+    Returns (bin_centers, log_pdfs). log_pdfs has shape (N, n_bins-1).
+    """
+    imgs01 = np.clip(imgs, 0.0, 1.0)
+    bin_edges = np.linspace(log_nhi_min, log_nhi_max, n_bins)
+    bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
+    pdfs = []
+    for img in imgs01:
+        vals = log_nhi_min + (log_nhi_max - log_nhi_min) * img.reshape(-1)
+        hist, _ = np.histogram(vals, bins=bin_edges, density=True)
+        pdfs.append(hist)
+    return bin_centers, np.log(np.stack(pdfs) + 1e-30)
+
+
+# =============================================================================
+# 2.  CALIBRATION OF sigma_pk AND sigma_pdf
+# =============================================================================
+
+def calibrate_summary_sigmas(
+    model: torch.nn.Module,
+    images_val: np.ndarray,
+    labels_val: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    device: torch.device,
+    box_size: float = 25.0,
+    ddim_steps: int = 50,
+    n_pairs: int = 30,
+    seed: int = 0,
+) -> Tuple[float, float]:
+    """Estimate sigma_pk and sigma_pdf from DDPM aleatoric scatter.
+
+    For each of n_pairs validation labels, draw two independent DDPM samples
+    and measure std(summary_a - summary_b) / sqrt(2). The median over pairs
+    gives a robust per-draw noise scale.
+    """
+    rng = np.random.default_rng(seed)
+    n_val = min(n_pairs, len(labels_val))
+    idx = rng.choice(len(labels_val), size=n_val, replace=False)
+    H, W = int(images_val.shape[-2]), int(images_val.shape[-1])
+
+    sig_pk: List[float] = []
+    sig_pdf: List[float] = []
+
+    for i in idx:
+        lab_pair = np.repeat(labels_val[i:i + 1], 2, axis=0).astype(np.float32)
+        pair = em.sample_batch(
+            model, lab_pair, lab_mean, lab_std, normalize,
+            H, W, device, ddim_steps, False,
+        )
+        _, _, lpk = per_map_log_pk(pair, box_size)
+        _, lpdf = per_map_log_pdf(pair)
+        sig_pk.append(float(np.std(lpk[0] - lpk[1]) / np.sqrt(2.0)))
+        sig_pdf.append(float(np.std(lpdf[0] - lpdf[1]) / np.sqrt(2.0)))
+
+    s_pk = max(float(np.median(sig_pk)), 0.01)
+    s_pdf = max(float(np.median(sig_pdf)), 0.01)
+    print(f"  sigma_pk  median over {n_val} pairs = {s_pk:.4f}")
+    print(f"  sigma_pdf median over {n_val} pairs = {s_pdf:.4f}")
+    return s_pk, s_pdf
+
+
+# =============================================================================
+# 3.  JOINT LOG-LIKELIHOOD  (P(k) + PDF, averaged over DDPM stochasticity)
+# =============================================================================
+
+def joint_log_likelihood(
+    obs: np.ndarray,
+    full_labels: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    device: torch.device,
+    H: int,
+    W: int,
+    box_size: float,
+    ddim_steps: int,
+    batch_sz: int,
+    n_pk_samples: int,
+    sigma_pk: float,
+    sigma_pdf: float,
+) -> np.ndarray:
+    """Average log P(k) and log PDF over n_pk_samples DDPM draws per grid pt.
+
+    Returns log L = log L_pk + log L_pdf at each grid point (shape (ngrid,)).
+    """
+    # Observed summaries
+    _, valid_pk, log_pk_obs = per_map_log_pk(obs[np.newaxis], box_size)
+    log_pk_obs = log_pk_obs[0]                           # (n_valid_pk,)
+    _, log_pdf_obs = per_map_log_pdf(obs[np.newaxis])
+    log_pdf_obs = log_pdf_obs[0]                         # (n_pdf_bins,)
+
+    ngrid = full_labels.shape[0]
+    sum_lpk = np.zeros((ngrid, log_pk_obs.size), dtype=np.float64)
+    sum_lpdf = np.zeros((ngrid, log_pdf_obs.size), dtype=np.float64)
+
+    for _s in range(n_pk_samples):
+        for j0 in range(0, ngrid, batch_sz):
+            chunk = full_labels[j0: j0 + batch_sz]
+            imgs = em.sample_batch(
+                model, chunk, lab_mean, lab_std, normalize,
+                H, W, device, ddim_steps, False,
+            )
+            _, _, lpk = per_map_log_pk(imgs, box_size)
+            _, lpdf = per_map_log_pdf(imgs)
+            sum_lpk[j0: j0 + len(chunk)] += lpk
+            sum_lpdf[j0: j0 + len(chunk)] += lpdf
+
+    mean_lpk = sum_lpk / n_pk_samples
+    mean_lpdf = sum_lpdf / n_pk_samples
+
+    mse_pk = np.mean((log_pk_obs[None, :] - mean_lpk) ** 2, axis=1)
+    mse_pdf = np.mean((log_pdf_obs[None, :] - mean_lpdf) ** 2, axis=1)
+
+    return -mse_pk / (2.0 * sigma_pk ** 2) - mse_pdf / (2.0 * sigma_pdf ** 2)
+
+
+# =============================================================================
+# 4.  GRIDS AND PER-MODEL POSTERIORS
+# =============================================================================
+
+def cosmo_grid_axes(
+    labels_ref: np.ndarray, grid: int, pad_frac: float = 0.02,
+) -> Tuple[np.ndarray, np.ndarray]:
+    lo0, hi0 = float(labels_ref[:, 0].min()), float(labels_ref[:, 0].max())
+    lo1, hi1 = float(labels_ref[:, 1].min()), float(labels_ref[:, 1].max())
+    p0 = pad_frac * (hi0 - lo0 + 1e-12)
+    p1 = pad_frac * (hi1 - lo1 + 1e-12)
+    om_ax = np.linspace(lo0 - p0, hi0 + p0, grid, dtype=np.float32)
+    s8_ax = np.linspace(lo1 - p1, hi1 + p1, grid, dtype=np.float32)
+    return om_ax, s8_ax
+
+
+def posterior_ddpm2(
+    obs: np.ndarray,
+    labels_ref: np.ndarray,
+    lab_mean: np.ndarray, lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module, device: torch.device,
+    grid: int, batch_sz: int, ddim_steps: int,
+    n_pk_samples: int,
+    sigma_pk: float, sigma_pdf: float,
+    box_size: float = 25.0,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    om_ax, s8_ax = cosmo_grid_axes(labels_ref, grid)
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    full = np.stack([OM.ravel(), S8.ravel()], axis=1).astype(np.float32)
+    H, W = int(obs.shape[-2]), int(obs.shape[-1])
+
+    log_w = joint_log_likelihood(
+        obs, full, lab_mean, lab_std, normalize, model, device,
+        H, W, box_size, ddim_steps, batch_sz,
+        n_pk_samples, sigma_pk, sigma_pdf,
+    )
+    log_w -= log_w.max()
+    w = np.exp(log_w).reshape(grid, grid)
+    w /= w.sum()
+    return w, OM, S8, om_ax, s8_ax
+
+
+def posterior_ddpm6(
+    obs: np.ndarray,
+    labels_ref: np.ndarray,
+    lab_mean: np.ndarray, lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module, device: torch.device,
+    lo_tail: np.ndarray, hi_tail: np.ndarray,
+    grid: int, batch_sz: int, ddim_steps: int,
+    n_pk_samples: int, n_marg_samples: int,
+    sigma_pk: float, sigma_pdf: float,
+    box_size: float = 25.0,
+    seed: int = 1,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    om_ax, s8_ax = cosmo_grid_axes(labels_ref, grid)
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    ngrid = OM.size
+    H, W = int(obs.shape[-2]), int(obs.shape[-1])
+
+    log_acc = np.full(ngrid, -np.inf, dtype=np.float64)
+    for m in range(n_marg_samples):
+        theta_extra = rng.uniform(lo_tail, hi_tail).astype(np.float32)
+        full_6 = np.zeros((ngrid, 6), dtype=np.float32)
+        full_6[:, 0] = OM.ravel()
+        full_6[:, 1] = S8.ravel()
+        full_6[:, 2:6] = theta_extra
+        log_w = joint_log_likelihood(
+            obs, full_6, lab_mean, lab_std, normalize, model, device,
+            H, W, box_size, ddim_steps, batch_sz,
+            n_pk_samples, sigma_pk, sigma_pdf,
+        )
+        log_acc = np.logaddexp(log_acc, log_w)
+        if (m + 1) % 5 == 0 or (m + 1) == n_marg_samples:
+            print(f"    DDPM-6 marg draw {m + 1}/{n_marg_samples}")
+
+    log_acc -= np.log(n_marg_samples)
+    log_acc -= log_acc.max()
+    w = np.exp(log_acc).reshape(grid, grid)
+    w /= w.sum()
+    return w, OM, S8, om_ax, s8_ax
+
+
+# =============================================================================
+# 5.  POSTERIOR DIAGNOSTICS
+# =============================================================================
+
+def credible_levels(
+    w: np.ndarray, levels: Tuple[float, ...] = (0.68, 0.95),
+) -> List[float]:
+    sorted_w = np.sort(w.ravel())[::-1]
+    cumsum = np.cumsum(sorted_w)
+    out = []
+    total = float(sorted_w.sum())
+    for L in levels:
+        idx = int(np.searchsorted(cumsum, L * total))
+        idx = min(idx, len(sorted_w) - 1)
+        out.append(float(sorted_w[idx]))
+    return out
+
+
+def posterior_summary(
+    w: np.ndarray, OM: np.ndarray, S8: np.ndarray,
+) -> Dict[str, float]:
+    w = w / w.sum()
+    mom = float((w * OM).sum())
+    ms8 = float((w * S8).sum())
+    sm = float(np.sqrt((w * (OM - mom) ** 2).sum()))
+    ss = float(np.sqrt((w * (S8 - ms8) ** 2).sum()))
+    S8_map = S8 * (OM / 0.3) ** 0.5
+    mS8 = float((w * S8_map).sum())
+    sS8 = float(np.sqrt((w * (S8_map - mS8) ** 2).sum()))
+    n_eff = float(1.0 / (w.ravel() ** 2).sum())
+    return dict(
+        om_mean=mom, om_std=sm, s8_mean=ms8, s8_std=ss,
+        S8_mean=mS8, S8_std=sS8, n_eff=n_eff,
+    )
+
+
+# =============================================================================
+# 6.  PLOTTING
+# =============================================================================
+
+def plot_panel(
+    ax,
+    w: np.ndarray, OM: np.ndarray, S8: np.ndarray,
+    true_om: float, true_s8: float,
+    title: str,
+    summary: Dict[str, float],
+    cmap: str,
+) -> None:
+    cf = ax.contourf(OM, S8, w, levels=16, cmap=cmap)
+    plt.colorbar(cf, ax=ax, fraction=0.046, pad=0.04)
+
+    try:
+        thr68, thr95 = credible_levels(w, levels=(0.68, 0.95))
+        ax.contour(
+            OM, S8, w, levels=[thr95, thr68],
+            colors=["#e07b39", "#c0392b"],
+            linewidths=[1.0, 1.6],
+            linestyles=["--", "-"],
+        )
+    except Exception:
+        pass
+
+    ax.scatter(summary["om_mean"], summary["s8_mean"],
+               s=70, c="k", marker="+", zorder=8, lw=2.0)
+    ax.scatter(true_om, true_s8,
+               s=70, c="r", marker="x", zorder=8, lw=2.0)
+
+    info = (
+        f"$\\Omega_m$ = {summary['om_mean']:.3f} $\\pm$ {summary['om_std']:.3f}\n"
+        f"$\\sigma_8$ = {summary['s8_mean']:.3f} $\\pm$ {summary['s8_std']:.3f}\n"
+        f"$S_8$ = {summary['S8_mean']:.3f} $\\pm$ {summary['S8_std']:.3f}\n"
+        f"$n_{{\\rm eff}}$ = {summary['n_eff']:.0f}"
+    )
+    ax.text(
+        0.02, 0.02, info, transform=ax.transAxes, fontsize=7,
+        va="bottom", color="#111",
+        bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.78),
+    )
+
+    ax.set_xlabel(r"$\Omega_m$", fontsize=9)
+    ax.set_ylabel(r"$\sigma_8$", fontsize=9)
+    ax.set_title(title, fontsize=8.5)
+
+
+def make_comparison_figure(
+    per_anchor: List[Dict],
+    suptitle: str,
+    out_path: Path,
+) -> None:
+    n = len(per_anchor)
+    fig, axes = plt.subplots(n, 2, figsize=(11, 4.5 * n), squeeze=False)
+    for i, p in enumerate(per_anchor):
+        plot_panel(
+            axes[i][0],
+            p["w2"], p["OM"], p["S8"],
+            p["true_om"], p["true_s8"],
+            (
+                f"DDPM-2  |  ix={p['ix']}  |  "
+                rf"$\Omega_m$={p['true_om']:.3f}, $\sigma_8$={p['true_s8']:.3f}"
+            ),
+            p["summ2"], cmap="Blues",
+        )
+        plot_panel(
+            axes[i][1],
+            p["w6"], p["OM"], p["S8"],
+            p["true_om"], p["true_s8"],
+            (
+                f"DDPM-6 (MC marg.)  |  ix={p['ix']}  |  "
+                rf"$\Omega_m$={p['true_om']:.3f}, $\sigma_8$={p['true_s8']:.3f}"
+            ),
+            p["summ6"], cmap="Greens",
+        )
+
+    from matplotlib.lines import Line2D
+    legend_h = [
+        Line2D([], [], marker="x", color="r", ls="", ms=8, label="True"),
+        Line2D([], [], marker="+", color="k", ls="", ms=8, label="Posterior mean"),
+        Line2D([], [], color="#c0392b", lw=1.6, label="68% CR"),
+        Line2D([], [], color="#e07b39", lw=1.0, ls="--", label="95% CR"),
+    ]
+    fig.legend(
+        handles=legend_h, loc="upper center", ncol=4, fontsize=8.5,
+        bbox_to_anchor=(0.5, 0.998), frameon=False,
+    )
+    plt.suptitle(suptitle, fontsize=11, y=0.992)
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    fig.savefig(out_path, dpi=160, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved -> {out_path}")
+
+
+# =============================================================================
+# 7.  MODEL LOADING
+# =============================================================================
+
+def load_model(
+    args_json: Path, ckpt: Path, device: torch.device,
+) -> Tuple[torch.nn.Module, Dict]:
+    cfg = ec.load_training_config(str(args_json))
+    model = ec.build_model(cfg, device)
+    ec.load_checkpoint(model, str(ckpt), device)
+    model.eval()
+    return model, cfg
+
+
+def tail_lhs_bounds(data_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
+    for name in ("train_labels_LH.npy", "train_labels_LH_2.npy"):
+        p = data_dir / name
+        if p.is_file():
+            L = np.load(p)
+            if L.shape[1] < 6:
+                raise ValueError(
+                    f"Expected >=6 label columns in {p}, got {L.shape}"
+                )
+            return (
+                L[:, 2:6].min(axis=0).astype(np.float32),
+                L[:, 2:6].max(axis=0).astype(np.float32),
+            )
+    raise FileNotFoundError(f"No train_labels_LH*.npy under {data_dir}")
+
+
+# =============================================================================
+# 8.  CLI
+# =============================================================================
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description=(
+            "DDPM-2 vs DDPM-6 corrected posteriors on (Omega_m, sigma_8) "
+            "with a JOINT P(k) + log-N_HI PDF Gaussian likelihood. "
+            "sigma_pk and sigma_pdf are calibrated from DDPM aleatoric noise."
+        ),
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument(
+        "--output-dir", type=Path,
+        default=MODELS_ROOT / "ddpm_posterior_compare_pk_pdf_out",
+    )
+    p.add_argument(
+        "--data-2param", type=Path,
+        default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2"),
+    )
+    p.add_argument(
+        "--data-6param", type=Path,
+        default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"),
+    )
+    p.add_argument(
+        "--bundle-2param", type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "2param_epoch200",
+    )
+    p.add_argument(
+        "--bundle-6param", type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "6param_best",
+    )
+    p.add_argument("--split", default="test", choices=["train", "val", "test"])
+    p.add_argument(
+        "--n-anchors", type=int, default=4,
+        help="Number of evenly-spaced test fields. 4 = compact figure; 6 = wider sweep.",
+    )
+    p.add_argument("--grid", type=int, default=30)
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument(
+        "--n-pk-samples", type=int, default=8,
+        help="DDPM draws averaged per grid point. Variance ~ 1/N. >=8 recommended.",
+    )
+    p.add_argument(
+        "--n-marg-samples", type=int, default=20,
+        help="MC draws over astrophysical params for DDPM-6. >=20 recommended.",
+    )
+    p.add_argument(
+        "--n-calib-pairs", type=int, default=30,
+        help="Validation pairs used to calibrate sigma_pk and sigma_pdf.",
+    )
+    p.add_argument(
+        "--sigma-pk", type=float, default=None,
+        help="Override calibrated sigma_pk (applied to BOTH models).",
+    )
+    p.add_argument(
+        "--sigma-pdf", type=float, default=None,
+        help="Override calibrated sigma_pdf (applied to BOTH models).",
+    )
+    p.add_argument("--seed", type=int, default=0)
+    return p.parse_args()
+
+
+# =============================================================================
+# 9.  MAIN
+# =============================================================================
+
+def main() -> None:
+    args = parse_args()
+    out_dir = args.output_dir.resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    print(f"Device : {device}")
+    print(f"Output : {out_dir}")
+    print()
+
+    # ── Data ─────────────────────────────────────────────────────────────
+    imgs2, labs2 = ec.load_split(args.data_2param, args.split)
+    mean2, std2 = ec.load_label_stats(args.data_2param)
+    imgs6, labs6 = ec.load_split(args.data_6param, args.split)
+    mean6, std6 = ec.load_label_stats(args.data_6param)
+    lo_tail, hi_tail = tail_lhs_bounds(args.data_6param)
+
+    print(f"DDPM-2 {args.split}: {len(labs2)} maps  label_dim={labs2.shape[1]}")
+    print(f"DDPM-6 {args.split}: {len(labs6)} maps  label_dim={labs6.shape[1]}")
+    print(f"  LHS tails (dims 2-5): min={lo_tail}  max={hi_tail}")
+
+    # Pick anchors evenly spread across the test split. Ranges of the two
+    # test splits may differ in length, so cap to the smaller.
+    n_pool = min(len(labs2), len(labs6))
+    n_anchors = max(2, min(args.n_anchors, n_pool))
+    anchor_ix = np.linspace(0, n_pool - 1, n_anchors, dtype=int)
+    print(f"Anchor indices: {anchor_ix.tolist()}")
+    print()
+
+    # ── Models ───────────────────────────────────────────────────────────
+    ck2 = args.bundle_2param / "checkpoint_epoch_200.pt"
+    aj2 = args.bundle_2param / "args.json"
+    ck6 = args.bundle_6param / "best_model.pt"
+    aj6 = args.bundle_6param / "args.json"
+
+    print("Loading DDPM-2 ...")
+    model2, cfg2 = load_model(aj2, ck2, device)
+    norm2 = bool(cfg2.get("normalize_labels", True))
+    print("Loading DDPM-6 ...")
+    model6, cfg6 = load_model(aj6, ck6, device)
+    norm6 = bool(cfg6.get("normalize_labels", True))
+    print()
+
+    # ── Calibrate sigma_pk and sigma_pdf per model ───────────────────────
+    if args.sigma_pk is not None and args.sigma_pdf is not None:
+        s2_pk = s6_pk = float(args.sigma_pk)
+        s2_pdf = s6_pdf = float(args.sigma_pdf)
+        print(f"sigma_pk overridden  = {s2_pk:.4f}")
+        print(f"sigma_pdf overridden = {s2_pdf:.4f}")
+    else:
+        print("Calibrating DDPM-2 noise scales ...")
+        v2_imgs, v2_labs = ec.load_split(args.data_2param, "val")
+        s2_pk, s2_pdf = calibrate_summary_sigmas(
+            model2, v2_imgs, v2_labs, mean2, std2, norm2, device,
+            ddim_steps=args.ddim_steps,
+            n_pairs=args.n_calib_pairs, seed=args.seed,
+        )
+        del v2_imgs, v2_labs
+        gc.collect()
+
+        print("Calibrating DDPM-6 noise scales ...")
+        v6_imgs, v6_labs = ec.load_split(args.data_6param, "val")
+        s6_pk, s6_pdf = calibrate_summary_sigmas(
+            model6, v6_imgs, v6_labs, mean6, std6, norm6, device,
+            ddim_steps=args.ddim_steps,
+            n_pairs=args.n_calib_pairs, seed=args.seed + 7,
+        )
+        del v6_imgs, v6_labs
+        gc.collect()
+
+        if args.sigma_pk is not None:
+            s2_pk = s6_pk = float(args.sigma_pk)
+        if args.sigma_pdf is not None:
+            s2_pdf = s6_pdf = float(args.sigma_pdf)
+
+    print()
+    print(f"DDPM-2: sigma_pk={s2_pk:.4f}  sigma_pdf={s2_pdf:.4f}")
+    print(f"DDPM-6: sigma_pk={s6_pk:.4f}  sigma_pdf={s6_pdf:.4f}")
+    print()
+
+    # ── Per-anchor inference ─────────────────────────────────────────────
+    per_anchor: List[Dict] = []
+    for k, ix in enumerate(anchor_ix):
+        ix = int(ix)
+        obs2 = imgs2[ix]
+        obs6 = imgs6[ix]
+        # DDPM-2 labels carry (Omega_m, sigma_8) directly; DDPM-6 labels[ix]
+        # may differ in row ordering between the two splits, so we report the
+        # truth from each respective split when forming panel titles.
+        true_om = float(labs2[ix, 0])
+        true_s8 = float(labs2[ix, 1])
+        true_om6 = float(labs6[ix, 0])
+        true_s86 = float(labs6[ix, 1])
+
+        print(
+            f"[{k + 1}/{n_anchors}] ix={ix}  "
+            f"DDPM-2 truth (Om={true_om:.3f}, s8={true_s8:.3f})  "
+            f"DDPM-6 truth (Om={true_om6:.3f}, s8={true_s86:.3f})"
+        )
+
+        print("  DDPM-2 posterior ...")
+        w2, OM, S8, om_ax, s8_ax = posterior_ddpm2(
+            obs2, labs2, mean2, std2, norm2, model2, device,
+            args.grid, args.batch_size, args.ddim_steps,
+            args.n_pk_samples, s2_pk, s2_pdf,
+        )
+        summ2 = posterior_summary(w2, OM, S8)
+        print(
+            f"    Om={summ2['om_mean']:.3f}+/-{summ2['om_std']:.3f}  "
+            f"s8={summ2['s8_mean']:.3f}+/-{summ2['s8_std']:.3f}  "
+            f"S8={summ2['S8_mean']:.3f}+/-{summ2['S8_std']:.3f}  "
+            f"n_eff={summ2['n_eff']:.0f}"
+        )
+
+        print("  DDPM-6 posterior (MC marginalisation) ...")
+        w6, OM6, S8_6, om_ax6, s8_ax6 = posterior_ddpm6(
+            obs6, labs6, mean6, std6, norm6, model6, device,
+            lo_tail, hi_tail,
+            args.grid, args.batch_size, args.ddim_steps,
+            args.n_pk_samples, args.n_marg_samples, s6_pk, s6_pdf,
+            seed=args.seed + 1 + k,
+        )
+        summ6 = posterior_summary(w6, OM6, S8_6)
+        print(
+            f"    Om={summ6['om_mean']:.3f}+/-{summ6['om_std']:.3f}  "
+            f"s8={summ6['s8_mean']:.3f}+/-{summ6['s8_std']:.3f}  "
+            f"S8={summ6['S8_mean']:.3f}+/-{summ6['S8_std']:.3f}  "
+            f"n_eff={summ6['n_eff']:.0f}"
+        )
+
+        per_anchor.append(dict(
+            ix=ix,
+            true_om=true_om, true_s8=true_s8,
+            w2=w2, w6=w6, OM=OM, S8=S8,
+            summ2=summ2, summ6=summ6,
+        ))
+        np.savez_compressed(
+            out_dir / f"posterior_ix{ix:03d}.npz",
+            w_ddpm2=w2, w_ddpm6=w6,
+            om_ax=om_ax, s8_ax=s8_ax,
+            true_om_ddpm2=true_om, true_s8_ddpm2=true_s8,
+            true_om_ddpm6=true_om6, true_s8_ddpm6=true_s86,
+            sigma_pk_ddpm2=s2_pk, sigma_pdf_ddpm2=s2_pdf,
+            sigma_pk_ddpm6=s6_pk, sigma_pdf_ddpm6=s6_pdf,
+            n_pk_samples=args.n_pk_samples,
+            n_marg_samples=args.n_marg_samples,
+            ddim_steps=args.ddim_steps,
+        )
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    # ── Comparison figure ────────────────────────────────────────────────
+    suptitle = (
+        r"Corrected joint-likelihood posteriors  $(\Omega_m,\,\sigma_8)$"
+        "   ·   "
+        r"$P(k)$ + $\log\,N_{\rm HI}$ PDF"
+        "\n"
+        f"DDPM-2  $\\sigma_{{Pk}}$={s2_pk:.3f}, $\\sigma_{{PDF}}$={s2_pdf:.3f}"
+        "    ·    "
+        f"DDPM-6 (MC marg., $N_{{\\rm marg}}$={args.n_marg_samples}) "
+        f"$\\sigma_{{Pk}}$={s6_pk:.3f}, $\\sigma_{{PDF}}$={s6_pdf:.3f}"
+        "    ·    "
+        f"grid {args.grid}², {args.n_pk_samples} DDPM draws/pt, "
+        f"DDIM {args.ddim_steps} steps"
+    )
+    make_comparison_figure(
+        per_anchor, suptitle,
+        out_dir / "compare_posterior_ddpm2_vs_ddpm6.png",
+    )
+
+    print(f"\nDone. Artifacts in {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

cross_model/ddpm_posterior_corrected.py

@@ -0,0 +1,867 @@
+#!/usr/bin/env python3
+"""
+Corrected Surrogate P(k) Bayesian Posteriors on (Omega_m, sigma_8).
+
+═══════════════════════════════════════════════════════════════════
+THEORETICAL CORRECTIONS FROM REVIEW
+═══════════════════════════════════════════════════════════════════
+1. STOCHASTIC LIKELIHOOD AVERAGING
+   - Single DDPM sample per grid point → n_ddpm_samples averaged draws
+   - Reduces emulator variance by 1/sqrt(N), preventing spurious multimodality
+
+2. CALIBRATED LIKELIHOOD NOISE SCALE
+   - Hard-coded sigma=0.25 → sigma_pk estimated from validation-set scatter
+   - Uses aleatoric uncertainty of the DDPM emulator at fixed theta
+
+3. PROPER MC MARGINALISATION (DDPM-6)
+   - Fixed dims 2-5 to extremes → Monte Carlo integral over prior
+   - p(Om,s8|d) = integral L(d|Om,s8,theta_extra) pi(theta_extra) dtheta_extra
+   - Approximated by uniform draws from the LHS training range
+
+4. HIGHER GRID RESOLUTION
+   - 14x14 → 30x30 (900 pts), with optional adaptive refinement
+
+5. VISUALISATION OF PRIOR, LIKELIHOOD, AND POSTERIOR SEPARATELY
+   - Shows all three distributions per anchor per model
+   - Includes 68%/95% credible contours
+   - Includes S8 = sigma_8*(Om/0.3)^0.5 derived parameter
+   - Includes effective sample size and posterior predictive check
+
+6. PRIOR-POSTERIOR COMPARISON
+   - Explicit uniform prior overlaid on posterior for each anchor
+
+USAGE
+─────
+  # Run both models with full corrections
+  python ddpm_posterior_corrected.py
+
+  # DDPM-2 only, faster run
+  python ddpm_posterior_corrected.py --ddpm2-only --grid 20 --n-ddpm-samples 4
+
+  # DDPM-6 only with full marginalisation
+  python ddpm_posterior_corrected.py --ddpm6-only --n-marg-samples 30 --n-ddpm-samples 8
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+import warnings
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+# Resolve code + data relative to Models/ (same directory as this file)
+MODELS_ROOT = Path(__file__).resolve().parent
+CODE_6 = MODELS_ROOT / "6param_ddpm_hi_lh6"
+if str(CODE_6.resolve()) not in sys.path:
+    sys.path.insert(0, str(CODE_6))
+
+import evaluate_conditional as ec   # noqa: E402
+import eval_model as em             # noqa: E402
+from figure9_posterior import log_pk_observed  # noqa: E402
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 1.  DATA UTILITIES
+# ══════════════════════════════════════════════════════════════════════════════
+
+def _train_label_path(data_dir: Path) -> Path:
+    for name in ("train_labels_LH.npy", "train_labels_LH_2.npy"):
+        p = data_dir / name
+        if p.is_file():
+            return p
+    raise FileNotFoundError(f"No train_labels_LH*.npy under {data_dir}")
+
+
+def tail_lhs_bounds(data_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Min/max of the LHS training distribution for label dims 2-5.
+    These define the UNIFORM PRIOR for the astrophysical nuisance parameters.
+    """
+    L = np.load(_train_label_path(data_dir))
+    if L.shape[1] < 6:
+        raise ValueError(f"Expected >= 6 label columns, got {L.shape}")
+    lo = L[:, 2:6].min(axis=0).astype(np.float32)
+    hi = L[:, 2:6].max(axis=0).astype(np.float32)
+    return lo, hi
+
+
+def cosmo_prior_bounds(labels_split: np.ndarray) -> Tuple[float, float, float, float]:
+    """Return (om_lo, om_hi, s8_lo, s8_hi) from the training set LHS range."""
+    om_lo = float(labels_split[:, 0].min())
+    om_hi = float(labels_split[:, 0].max())
+    s8_lo = float(labels_split[:, 1].min())
+    s8_hi = float(labels_split[:, 1].max())
+    return om_lo, om_hi, s8_lo, s8_hi
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 2.  LIKELIHOOD CALIBRATION
+# ══════════════════════════════════════════════════════════════════════════════
+
+def calibrate_sigma_pk(
+    data_dir: Path,
+    model: torch.nn.Module,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    H: int,
+    W: int,
+    device: torch.device,
+    ddim_steps: int,
+    n_pairs: int = 60,
+    rng_seed: int = 0,
+) -> float:
+    """
+    Estimate sigma_pk = aleatoric std of log P(k) from DDPM stochasticity.
+
+    For n_pairs randomly chosen validation-set labels, draw 2 DDPM samples
+    and measure std(log Pk_a - log Pk_b) / sqrt(2).  The median over pairs
+    gives a robust noise floor for the likelihood.
+
+    This replaces the hard-coded sigma = 0.25 with a data-driven estimate.
+    """
+    print(f"  Calibrating sigma_pk from {n_pairs} validation-set pairs ...")
+    images_val, labels_val = ec.load_split(data_dir, "val")
+    n_val = len(labels_val)
+    rng = np.random.default_rng(rng_seed)
+    idx = rng.choice(n_val, size=min(n_pairs, n_val), replace=False)
+
+    sigmas = []
+    for i in idx:
+        lab = np.repeat(labels_val[i : i + 1], 2, axis=0).astype(np.float32)
+        pair = em.sample_batch(
+            model, lab, lab_mean, lab_std, normalize,
+            H, W, device, ddim_steps, False,
+        )
+        _, pk_pair = em.per_map_power_spectra_log(pair, 25.0)
+        valid = pk_pair[0] > 0
+        log_pk_pair = np.log(pk_pair[:, valid] + 1e-30)
+        diff_std = float(np.std(log_pk_pair[0] - log_pk_pair[1])) / np.sqrt(2.0)
+        sigmas.append(diff_std)
+
+    sigma_cal = float(np.median(sigmas))
+    sigma_cal = max(sigma_cal, 0.05)          # lower-bound: prevent degenerate likelihoods
+    print(f"  Calibrated sigma_pk = {sigma_cal:.4f}  (was hard-coded 0.25)")
+    return sigma_cal
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 3.  GRID CONSTRUCTION
+# ══════════════════════════════════════════════════════════════════════════════
+
+def build_cosmo_axes(
+    labels_split: np.ndarray, grid: int, pad_frac: float = 0.02
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Return (om_ax, s8_ax) with `grid` equally-spaced points inside the LHS range."""
+    om_lo, om_hi, s8_lo, s8_hi = cosmo_prior_bounds(labels_split)
+    pad0 = pad_frac * (om_hi - om_lo + 1e-12)
+    pad1 = pad_frac * (s8_hi - s8_lo + 1e-12)
+    om_ax = np.linspace(om_lo - pad0, om_hi + pad0, grid)
+    s8_ax = np.linspace(s8_lo - pad1, s8_hi + pad1, grid)
+    return om_ax, s8_ax
+
+
+def build_full_grid(
+    om_ax: np.ndarray,
+    s8_ax: np.ndarray,
+    tail: Optional[np.ndarray] = None,
+    lab_dim: int = 2,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Build a flattened (grid^2, lab_dim) array for a sweep over (Om, s8).
+    If tail is provided (shape (4,)), dims 2-5 are fixed to those values.
+    Returns (full_labels, OM_meshgrid, S8_meshgrid).
+    """
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    ngrid = OM.size
+    out = np.zeros((ngrid, lab_dim), dtype=np.float32)
+    out[:, 0] = OM.ravel()
+    out[:, 1] = S8.ravel()
+    if tail is not None:
+        assert tail.shape == (4,), f"Expected tail shape (4,), got {tail.shape}"
+        out[:, 2:6] = tail[np.newaxis, :]
+    return out, OM, S8
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 4.  LOG-LIKELIHOOD (AVERAGED OVER DDPM STOCHASTICITY)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def compute_log_likelihood(
+    obs: np.ndarray,
+    full: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    H: int,
+    W: int,
+    device: torch.device,
+    ddim_steps: int,
+    batch_sz: int,
+    n_ddpm_samples: int,
+    sigma_pk: float,
+) -> np.ndarray:
+    """
+    Return log-likelihood array of shape (ngrid,), where each entry is:
+
+        ln L(d | theta_i) ≈ - (1 / 2*sigma^2) * mean_k[ (log Pk_obs - mean_j[log Pk_gen_j]) ]^2
+
+    The mean over j=1..n_ddpm_samples suppresses DDPM stochasticity.
+
+    Parameters
+    ----------
+    sigma_pk : calibrated noise scale on log P(k) — NOT hard-coded.
+    n_ddpm_samples : number of independent DDPM draws to average per grid pt.
+    """
+    ngrid = full.shape[0]
+
+    npix = int(obs.shape[-1])
+    dl = 25.0 / npix
+    logf = em.images01_to_log_nhi(obs)
+    dk, _ = ec.PowerSpectrum(logf, N=npix, dl=dl)
+    valid = dk > 0
+    # log_pk_observed returns values only at dk > 0 (same length as valid.sum())
+    log_pd = log_pk_observed(obs, 25.0, dk)
+
+    accumulated = []
+    for _s in range(n_ddpm_samples):
+        sample_log_pk = []
+        for j0 in range(0, ngrid, batch_sz):
+            chunk = full[j0: j0 + batch_sz]
+            imgs = em.sample_batch(
+                model, chunk, lab_mean, lab_std, normalize,
+                H, W, device, ddim_steps, False,
+            )
+            _, pkc = em.per_map_power_spectra_log(imgs, 25.0)
+            sample_log_pk.append(np.log(pkc[:, valid] + 1e-30))
+        accumulated.append(np.concatenate(sample_log_pk, axis=0))  # (ngrid, nk)
+
+    mean_log_pg = np.mean(accumulated, axis=0)  # (ngrid, nk)
+
+    mse = np.mean((log_pd[np.newaxis, :] - mean_log_pg) ** 2, axis=1)  # (ngrid,)
+    log_like = -mse / (2.0 * sigma_pk ** 2)
+    return log_like
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 5.  MARGINALISATION OVER ASTROPHYSICAL PARAMETERS (DDPM-6)
+# ══════════════════════════════════════════════════════════════════════════════
+
+def marginal_log_likelihood_ddpm6(
+    obs: np.ndarray,
+    om_ax: np.ndarray,
+    s8_ax: np.ndarray,
+    lo_tail: np.ndarray,
+    hi_tail: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    H: int,
+    W: int,
+    device: torch.device,
+    ddim_steps: int,
+    batch_sz: int,
+    n_ddpm_samples: int,
+    n_marg_samples: int,
+    sigma_pk: float,
+    rng_seed: int = 42,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Correct 2D marginal log-likelihood over (Om, s8) for the 6-param DDPM.
+
+    Implements Monte Carlo integration:
+        ln L_marg(d | Om, s8) = log[ (1/N) Σ_i L(d | Om, s8, θ_extra^i) ]
+                               = log-sum-exp [ ln L(d | Om, s8, θ_extra^i) ] - ln N
+
+    where θ_extra^i ~ Uniform(lo_tail, hi_tail)  [the prior over dims 2-5]
+
+    Returns (log_like_marginal, OM, S8) with log_like_marginal shaped (ngrid,).
+    """
+    rng = np.random.default_rng(rng_seed)
+    theta_extras = rng.uniform(
+        lo_tail, hi_tail, size=(n_marg_samples, 4)
+    ).astype(np.float32)
+
+    ngrid = len(om_ax) * len(s8_ax)
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+
+    log_like_accumulator = np.full(ngrid, -np.inf, dtype=np.float64)
+
+    for mc_i, theta_extra in enumerate(theta_extras):
+        print(f"    MC marginalisation draw {mc_i+1}/{n_marg_samples} ...", end="\r", flush=True)
+        full, _, _ = build_full_grid(om_ax, s8_ax, tail=theta_extra, lab_dim=6)
+        lnL_i = compute_log_likelihood(
+            obs, full, lab_mean, lab_std, normalize, model,
+            H, W, device, ddim_steps, batch_sz, n_ddpm_samples, sigma_pk,
+        )
+        log_like_accumulator = np.logaddexp(log_like_accumulator, lnL_i)
+
+    log_like_marginal = log_like_accumulator - np.log(n_marg_samples)
+    print(flush=True)
+    return log_like_marginal, OM, S8
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 6.  POSTERIOR COMPUTATION & DIAGNOSTICS
+# ══════════════════════════════════════════════════════════════════════════════
+
+def log_like_to_posterior(
+    log_like: np.ndarray,
+    grid: int,
+) -> Tuple[np.ndarray, float]:
+    """Flat prior on grid → normalized weights + n_eff."""
+    log_like = log_like - log_like.max()           # numerical stability
+    weights = np.exp(log_like).reshape(grid, grid)
+    weights /= weights.sum()
+    n_eff = 1.0 / float(np.sum(weights ** 2))
+    return weights, n_eff
+
+
+def credible_contour_levels(
+    weights: np.ndarray, credible_levels=(0.68, 0.95)
+) -> List[float]:
+    """Highest-density-style thresholds containing `level` of total mass."""
+    flat = weights.ravel()
+    total_mass = float(flat.sum())
+    sorted_desc = np.sort(flat)[::-1]
+    cumsum = np.cumsum(sorted_desc)
+    thresholds = []
+    for cl in credible_levels:
+        target = cl * total_mass
+        idx = int(np.searchsorted(cumsum, target))
+        idx = min(idx, len(sorted_desc) - 1)
+        thresholds.append(float(sorted_desc[idx]))
+    return thresholds
+
+
+def posterior_summary(weights: np.ndarray, OM: np.ndarray, S8: np.ndarray) -> Dict:
+    """Posterior mean, std, and S8 = sigma_8*(Om/0.3)^0.5 statistics."""
+    mom = float((weights * OM).sum())
+    ms8 = float((weights * S8).sum())
+    var_om = float((weights * (OM - mom)**2).sum())
+    var_s8 = float((weights * (S8 - ms8)**2).sum())
+    S8_map = S8 * (OM / 0.3) ** 0.5
+    mS8 = float((weights * S8_map).sum())
+    var_S8 = float((weights * (S8_map - mS8)**2).sum())
+    return dict(
+        om_mean=mom, om_std=np.sqrt(var_om),
+        s8_mean=ms8, s8_std=np.sqrt(var_s8),
+        S8_mean=mS8, S8_std=np.sqrt(var_S8),
+    )
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 7.  POSTERIOR PREDICTIVE CHECK
+# ══════════════════════════════════════════════════════════════════════════════
+
+def posterior_predictive_check(
+    obs: np.ndarray,
+    weights: np.ndarray,
+    OM: np.ndarray,
+    S8: np.ndarray,
+    model: torch.nn.Module,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    H: int,
+    W: int,
+    device: torch.device,
+    ddim_steps: int,
+    n_draws: int = 30,
+    rng_seed: int = 7,
+    lab_dim: int = 2,
+    lo_tail: Optional[np.ndarray] = None,
+    hi_tail: Optional[np.ndarray] = None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Draw parameters from the posterior and generate maps.
+    Returns (k_valid, log_pk_obs_valid, ppc_lo68, ppc_hi68, ppc_lo95, ppc_hi95),
+    all aligned to dk > 0 bins only.
+    """
+    rng = np.random.default_rng(rng_seed)
+    flat_w = weights.ravel()
+    flat_om = OM.ravel()
+    flat_s8 = S8.ravel()
+    idx = rng.choice(len(flat_w), size=n_draws, replace=True, p=flat_w)
+
+    npix = int(obs.shape[-1])
+    dl = 25.0 / npix
+    logf = em.images01_to_log_nhi(obs)
+    dk, _ = ec.PowerSpectrum(logf, N=npix, dl=dl)
+    valid = dk > 0
+    log_pk_obs = log_pk_observed(obs, 25.0, dk)
+
+    ppc_log_pks = []
+    for i in idx:
+        if lab_dim == 2:
+            theta = np.array([[flat_om[i], flat_s8[i]]], dtype=np.float32)
+        else:
+            assert lo_tail is not None and hi_tail is not None
+            te = rng.uniform(lo_tail, hi_tail).astype(np.float32)
+            theta = np.array([[flat_om[i], flat_s8[i], *te]], dtype=np.float32)
+        img = em.sample_batch(
+            model, theta, lab_mean, lab_std, normalize,
+            H, W, device, ddim_steps, False,
+        )
+        _, pkc = em.per_map_power_spectra_log(img, 25.0)
+        ppc_log_pks.append(np.log(pkc[0, valid] + 1e-30))
+
+    ppc_arr = np.array(ppc_log_pks)   # (n_draws, n_valid)
+    return (
+        dk[valid],
+        log_pk_obs,
+        np.percentile(ppc_arr, 16, axis=0),
+        np.percentile(ppc_arr, 84, axis=0),
+        np.percentile(ppc_arr, 2.5, axis=0),
+        np.percentile(ppc_arr, 97.5, axis=0),
+    )
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 8.  VISUALISATION
+# ══════════════════════════════════════════════════════════════════════════════
+
+CMAP_PRIOR = "Greys"
+CMAP_LIKE = "YlOrRd"
+CMAP_POST = "Blues"
+
+
+def _uniform_prior(OM: np.ndarray, S8: np.ndarray) -> np.ndarray:
+    """Flat prior: uniform weight over the grid."""
+    prior = np.ones_like(OM)
+    return prior / prior.sum()
+
+
+def plot_prior_likelihood_posterior_panel(
+    fig,
+    gs_row,
+    weights: np.ndarray,
+    log_like: np.ndarray,
+    OM: np.ndarray,
+    S8: np.ndarray,
+    true_om: float,
+    true_s8: float,
+    anchor_ix: int,
+    summary: Dict,
+    n_eff: float,
+    title_suffix: str = "",
+) -> None:
+    """
+    Plot three side-by-side panels for one anchor:
+      [0] Uniform prior   [1] Normalised likelihood   [2] Posterior
+    Each panel shows 68% / 95% credible contours where applicable.
+    """
+    grid = weights.shape[0]
+    prior = _uniform_prior(OM, S8)
+
+    like = np.exp(log_like - log_like.max()).reshape(grid, grid)
+    like /= like.sum()
+
+    panels = [
+        (prior,   CMAP_PRIOR, r"Uniform Prior $\pi(\Omega_m, \sigma_8)$"),
+        (like,    CMAP_LIKE,  r"Normalised Likelihood $\mathcal{L}(\mathbf{d}|\theta)$"),
+        (weights, CMAP_POST,  r"Posterior $p(\theta|\mathbf{d})$"),
+    ]
+
+    for col, (Wmap, cmap, label) in enumerate(panels):
+        ax = fig.add_subplot(gs_row[col])
+        cf = ax.contourf(OM, S8, Wmap, levels=14, cmap=cmap)
+        plt.colorbar(cf, ax=ax, fraction=0.046, pad=0.04)
+
+        if col > 0:
+            lvls = credible_contour_levels(Wmap)
+            try:
+                ax.contour(OM, S8, Wmap, levels=lvls,
+                           colors=["white", "cyan"],
+                           linewidths=[1.0, 0.6],
+                           linestyles=["solid", "dashed"])
+            except Exception:
+                pass
+
+        ax.scatter(true_om, true_s8, s=70, c="red",
+                   marker="x", zorder=8, linewidths=2.0, label="True")
+        if col == 2:
+            ax.scatter(summary["om_mean"], summary["s8_mean"],
+                       s=80, c="black", marker="+",
+                       zorder=8, linewidths=2.0, label="Post. mean")
+            ax.legend(fontsize=7, loc="upper right")
+
+        if col == 2:
+            txt = (
+                f"$\\Omega_m$: {summary['om_mean']:.3f} ± {summary['om_std']:.3f}\n"
+                f"$\\sigma_8$: {summary['s8_mean']:.3f} ± {summary['s8_std']:.3f}\n"
+                f"$S_8$:     {summary['S8_mean']:.3f} ± {summary['S8_std']:.3f}\n"
+                f"$n_\\mathrm{{eff}}$: {n_eff:.0f}"
+            )
+            ax.text(0.02, 0.02, txt, transform=ax.transAxes,
+                    fontsize=6.5, va="bottom", color="#111",
+                    bbox=dict(boxstyle="round,pad=0.3", fc="white", alpha=0.75))
+
+        ax.set_xlabel(r"$\Omega_m$", fontsize=9)
+        ax.set_ylabel(r"$\sigma_8$", fontsize=9)
+        ax.set_title(
+            f"ix={anchor_ix}  |  {label}{title_suffix}",
+            fontsize=8, pad=4
+        )
+
+
+def plot_ppc_panel(
+    ax,
+    k_valid: np.ndarray,
+    log_pk_obs: np.ndarray,
+    ppc_lo68: np.ndarray,
+    ppc_hi68: np.ndarray,
+    ppc_lo95: np.ndarray,
+    ppc_hi95: np.ndarray,
+    anchor_ix: int,
+) -> None:
+    """Posterior predictive check on log P(k) at valid k bins."""
+    ax.fill_between(k_valid, ppc_lo95, ppc_hi95,
+                    alpha=0.18, color="steelblue", label="95% PPC")
+    ax.fill_between(k_valid, ppc_lo68, ppc_hi68,
+                    alpha=0.40, color="steelblue", label="68% PPC")
+    ax.plot(k_valid, log_pk_obs, "k-", lw=1.8, label="Observed")
+    ax.set_xscale("log")
+    ax.set_xlabel(r"$k$ [h/Mpc]", fontsize=8)
+    ax.set_ylabel(r"$\log P_\mathrm{HI}(k)$", fontsize=8)
+    ax.set_title(f"Posterior Predictive Check (ix={anchor_ix})", fontsize=8)
+    ax.legend(fontsize=7)
+    ax.grid(alpha=0.25)
+
+
+
+def plot_s8_marginal(
+    ax, weights: np.ndarray, OM: np.ndarray, S8: np.ndarray,
+    true_om: float, true_s8: float, anchor_ix: int,
+) -> None:
+    """1D marginal of the derived S8 = sigma_8*(Om/0.3)^0.5 parameter."""
+    S8_map = S8 * (OM / 0.3) ** 0.5
+    true_S8 = true_s8 * (true_om / 0.3) ** 0.5
+    S8_flat = S8_map.ravel()
+    w_flat = weights.ravel()
+    s8_bins = np.linspace(S8_flat.min(), S8_flat.max(), 40)
+    hist = np.zeros(len(s8_bins) - 1)
+    for i, (lo, hi) in enumerate(zip(s8_bins[:-1], s8_bins[1:])):
+        mask = (S8_flat >= lo) & (S8_flat < hi)
+        hist[i] = w_flat[mask].sum()
+    hist /= hist.sum() * (s8_bins[1] - s8_bins[0])
+    centers = 0.5 * (s8_bins[:-1] + s8_bins[1:])
+    ax.bar(centers, hist, width=(s8_bins[1] - s8_bins[0]),
+           color="steelblue", alpha=0.7, label="Posterior")
+    ax.axvline(true_S8, color="red", lw=1.5, ls="--", label=f"True $S_8$={true_S8:.3f}")
+    ax.set_xlabel(r"$S_8 = \sigma_8(\Omega_m/0.3)^{0.5}$", fontsize=8)
+    ax.set_ylabel("Prob. density", fontsize=8)
+    ax.set_title(f"$S_8$ marginal (ix={anchor_ix})", fontsize=8)
+    ax.legend(fontsize=7)
+    ax.grid(alpha=0.25)
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# 9.  PER-MODEL RUNNER
+# ══════════════════════════════════════════════════════════════════════════════
+
+def run_model(
+    out_dir: Path,
+    *,
+    model_name: str,
+    model: torch.nn.Module,
+    cfg: Dict,
+    images: np.ndarray,
+    labels: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    anchor_ix: np.ndarray,
+    grid: int,
+    ddim_steps: int,
+    batch_sz: int,
+    n_ddpm_samples: int,
+    n_marg_samples: int,
+    sigma_pk: float,
+    lo_tail: Optional[np.ndarray] = None,
+    hi_tail: Optional[np.ndarray] = None,
+    do_ppc: bool = True,
+) -> None:
+    normalize = bool(cfg.get("normalize_labels", True))
+    H, W = int(images.shape[-2]), int(images.shape[-1])
+    lab_dim = 6 if model_name == "DDPM-6" else 2
+    n_anchors = len(anchor_ix.ravel())
+    is_6param = lab_dim == 6
+    device = next(model.parameters()).device
+
+    fig1 = plt.figure(figsize=(17, 4.8 * n_anchors))
+    outer_gs = gridspec.GridSpec(n_anchors, 1, figure=fig1, hspace=0.55)
+
+    fig2, axes2 = plt.subplots(
+        n_anchors, 2, figsize=(12, 4.5 * n_anchors), squeeze=False
+    )
+
+    om_ax, s8_ax = build_cosmo_axes(labels, grid)
+
+    for k, ix in enumerate(anchor_ix.ravel()):
+        obs = images[ix]
+        lab_t = labels[ix].astype(np.float32)
+        true_om, true_s8 = float(lab_t[0]), float(lab_t[1])
+
+        print(f"\n[{model_name}] Anchor ix={ix}  "
+              f"(Ωm={true_om:.3f}, σ8={true_s8:.3f})")
+
+        if is_6param:
+            print("  MC marginalisation over dims 2-5 ...")
+            log_like, OM, S8 = marginal_log_likelihood_ddpm6(
+                obs, om_ax, s8_ax, lo_tail, hi_tail,
+                lab_mean, lab_std, normalize, model,
+                H, W, device=device,
+                ddim_steps=ddim_steps, batch_sz=batch_sz,
+                n_ddpm_samples=n_ddpm_samples,
+                n_marg_samples=n_marg_samples,
+                sigma_pk=sigma_pk,
+            )
+        else:
+            print(f"  Computing log-likelihood ({n_ddpm_samples} DDPM draws per point) ...")
+            full, OM, S8 = build_full_grid(om_ax, s8_ax, tail=None, lab_dim=2)
+            log_like = compute_log_likelihood(
+                obs, full, lab_mean, lab_std, normalize, model,
+                H, W, device, ddim_steps, batch_sz, n_ddpm_samples, sigma_pk,
+            )
+
+        weights, n_eff = log_like_to_posterior(log_like, grid)
+        summary = posterior_summary(weights, OM, S8)
+
+        print(f"  n_eff = {n_eff:.1f} / {grid**2} grid points")
+        print(f"  Ωm posterior: {summary['om_mean']:.3f} ± {summary['om_std']:.3f}  "
+              f"(true: {true_om:.3f})")
+        print(f"  σ8 posterior: {summary['s8_mean']:.3f} ± {summary['s8_std']:.3f}  "
+              f"(true: {true_s8:.3f})")
+        print(f"  S8 posterior: {summary['S8_mean']:.3f} ± {summary['S8_std']:.3f}")
+
+        if n_eff < 20:
+            warnings.warn(
+                f"n_eff={n_eff:.1f} is very low for ix={ix}. "
+                "Increase n_ddpm_samples or grid resolution.",
+                stacklevel=2,
+            )
+
+        inner_gs = gridspec.GridSpecFromSubplotSpec(
+            1, 3, subplot_spec=outer_gs[k], wspace=0.38
+        )
+        marg_note = " (marginalised)" if is_6param else ""
+        plot_prior_likelihood_posterior_panel(
+            fig1, inner_gs, weights, log_like.reshape(grid, grid),
+            OM, S8, true_om, true_s8, ix, summary, n_eff,
+            title_suffix=marg_note,
+        )
+
+        if do_ppc:
+            dk_v, log_pk_obs, plo68, phi68, plo95, phi95 = posterior_predictive_check(
+                obs, weights, OM, S8, model,
+                lab_mean, lab_std, normalize,
+                H, W, device, ddim_steps,
+                n_draws=20, lab_dim=lab_dim,
+                lo_tail=lo_tail, hi_tail=hi_tail,
+            )
+            plot_ppc_panel(
+                axes2[k, 0], dk_v, log_pk_obs, plo68, phi68, plo95, phi95, ix,
+            )
+        else:
+            axes2[k, 0].axis("off")
+            axes2[k, 0].text(
+                0.5, 0.5, "PPC disabled",
+                ha="center", va="center", transform=axes2[k, 0].transAxes,
+            )
+
+        plot_s8_marginal(axes2[k, 1], weights, OM, S8, true_om, true_s8, ix)
+
+    fig1.suptitle(
+        f"{model_name} — Prior · Likelihood · Posterior on "
+        r"$(\Omega_m,\,\sigma_8)$ — six CAMELS anchors"
+        f"\n[sigma_pk={sigma_pk:.3f}, n_ddpm={n_ddpm_samples}, grid={grid}×{grid}"
+        + (f", n_marg={n_marg_samples}]" if is_6param else "]"),
+        fontsize=12, y=1.001,
+    )
+    p1 = out_dir / f"{model_name.replace('-', '')}_prior_likelihood_posterior.png"
+    fig1.savefig(p1, dpi=160, bbox_inches="tight")
+    plt.close(fig1)
+    print(f"\nSaved → {p1}")
+
+    fig2.suptitle(
+        f"{model_name} — Posterior Predictive Checks & $S_8$ Marginals",
+        fontsize=12, y=1.002,
+    )
+    fig2.tight_layout()
+    p2 = out_dir / f"{model_name.replace('-', '')}_ppc_s8.png"
+    fig2.savefig(p2, dpi=160, bbox_inches="tight")
+    plt.close(fig2)
+    print(f"Saved → {p2}")
+
+
+def load_model(args_json: Path, ckpt: Path, device: torch.device):
+    cfg = ec.load_training_config(str(args_json))
+    model = ec.build_model(cfg, device)
+    ec.load_checkpoint(model, str(ckpt), device)
+    model.eval()
+    return model, cfg
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description=(
+            "Corrected Bayesian posteriors on (Ωm, σ8): averaged DDPM likelihood, "
+            "calibrated sigma_pk, DDPM-6 MC marginalisation, prior/likelihood/posterior plots."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument("--output-dir", type=Path,
+                   default=MODELS_ROOT / "ddpm_posterior_corrected_fullviz_out")
+    p.add_argument("--data-2param", type=Path,
+                   default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2"))
+    p.add_argument("--data-6param", type=Path,
+                   default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"))
+    p.add_argument("--bundle-2param", type=Path,
+                   default=MODELS_ROOT / "notebook_model_weights" / "2param_epoch200")
+    p.add_argument("--bundle-6param", type=Path,
+                   default=MODELS_ROOT / "notebook_model_weights" / "6param_best")
+    p.add_argument("--split", type=str, default="test",
+                   choices=["train", "val", "test"])
+    p.add_argument("--grid", type=int, default=30)
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument("--n-ddpm-samples", type=int, default=8)
+    p.add_argument("--n-marg-samples", type=int, default=20)
+    p.add_argument("--sigma-pk", type=float, default=None)
+    p.add_argument("--n-calib-pairs", type=int, default=60)
+    p.add_argument("--no-ppc", action="store_true")
+    p.add_argument("--ddpm2-only", action="store_true")
+    p.add_argument("--ddpm6-only", action="store_true")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    if args.ddpm2_only and args.ddpm6_only:
+        raise SystemExit("Use at most one of --ddpm2-only / --ddpm6-only.")
+
+    out_dir = Path(args.output_dir).resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+    print(f"Output directory: {out_dir}")
+
+    data2 = Path(args.data_2param)
+    data6 = Path(args.data_6param)
+
+    imgs2, lab2 = ec.load_split(data2, args.split)
+    imgs6, lab6 = ec.load_split(data6, args.split)
+
+    n = min(len(lab2), len(lab6))
+    anchor_ix = np.linspace(0, n - 1, num=6, dtype=int)
+    print(f"Anchor indices: {anchor_ix.tolist()}")
+
+    lo_tail, hi_tail = tail_lhs_bounds(data6)
+    print(f"LHS tails (dims 2-5): min={lo_tail}  max={hi_tail}")
+
+    mean2, std2 = ec.load_label_stats(data2)
+    mean6, std6 = ec.load_label_stats(data6)
+
+    if not args.ddpm6_only:
+        print("\n" + "═" * 60)
+        print(">>> DDPM-2 (corrected posteriors, six anchors)")
+        print("═" * 60)
+        ck2 = args.bundle_2param / "checkpoint_epoch_200.pt"
+        args_json_2 = args.bundle_2param / "args.json"
+        model2, cfg2 = load_model(args_json_2, ck2, device)
+
+        if args.sigma_pk is not None:
+            sigma2 = args.sigma_pk
+            print(f"  Using user-supplied sigma_pk = {sigma2:.4f}")
+        else:
+            sigma2 = calibrate_sigma_pk(
+                data2, model2, mean2, std2,
+                bool(cfg2.get("normalize_labels", True)),
+                int(imgs2.shape[-2]), int(imgs2.shape[-1]),
+                device, args.ddim_steps, args.n_calib_pairs,
+            )
+
+        run_model(
+            out_dir,
+            model_name="DDPM-2",
+            model=model2,
+            cfg=cfg2,
+            images=imgs2,
+            labels=lab2,
+            lab_mean=mean2,
+            lab_std=std2,
+            anchor_ix=anchor_ix,
+            grid=args.grid,
+            ddim_steps=args.ddim_steps,
+            batch_sz=args.batch_size,
+            n_ddpm_samples=args.n_ddpm_samples,
+            n_marg_samples=args.n_marg_samples,
+            sigma_pk=sigma2,
+            do_ppc=not args.no_ppc,
+        )
+        del model2
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    if not args.ddpm2_only:
+        print("\n" + "═" * 60)
+        print(">>> DDPM-6 (corrected posteriors + MC marginalisation, six anchors)")
+        print("═" * 60)
+        ck6 = args.bundle_6param / "best_model.pt"
+        args_json_6 = args.bundle_6param / "args.json"
+        model6, cfg6 = load_model(args_json_6, ck6, device)
+
+        if args.sigma_pk is not None:
+            sigma6 = args.sigma_pk
+            print(f"  Using user-supplied sigma_pk = {sigma6:.4f}")
+        else:
+            sigma6 = calibrate_sigma_pk(
+                data6, model6, mean6, std6,
+                bool(cfg6.get("normalize_labels", True)),
+                int(imgs6.shape[-2]), int(imgs6.shape[-1]),
+                device, args.ddim_steps, args.n_calib_pairs,
+            )
+
+        run_model(
+            out_dir,
+            model_name="DDPM-6",
+            model=model6,
+            cfg=cfg6,
+            images=imgs6,
+            labels=lab6,
+            lab_mean=mean6,
+            lab_std=std6,
+            anchor_ix=anchor_ix,
+            grid=args.grid,
+            ddim_steps=args.ddim_steps,
+            batch_sz=args.batch_size,
+            n_ddpm_samples=args.n_ddpm_samples,
+            n_marg_samples=args.n_marg_samples,
+            sigma_pk=sigma6,
+            lo_tail=lo_tail,
+            hi_tail=hi_tail,
+            do_ppc=not args.no_ppc,
+        )
+        del model6
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    print(f"\nAll outputs saved to: {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

cross_model/poster.py

@@ -0,0 +1,1112 @@
+#!/usr/bin/env python3
+"""
+ddpm_posterior_six_anchors_corrected.py
+========================================
+Corrected surrogate P(k) likelihood posteriors on (Omega_m, sigma_8)
+for six CAMELS test anchors.
+
+CORRECTIONS OVER THE ORIGINAL SCRIPT
+--------------------------------------
+(1) STOCHASTIC EMULATOR NOISE  [was: 1 DDPM sample/grid point → fragmented posteriors]
+    Now: average log P(k) over `--n-pk-samples` (default 8) DDPM draws per grid
+    point, suppressing emulator variance by ~1/sqrt(N_s).
+
+(2) CALIBRATED LIKELIHOOD NOISE SCALE  [was: hard-coded sigma=0.25]
+    Now: sigma_pk is estimated from the scatter of log P(k) across repeated DDPM
+    draws at a sample of validation labels — making the noise scale physically
+    meaningful and data-driven.
+
+(3) PROPER MARGINALIZATION OVER ASTROPHYSICAL PARAMETERS  [was: fix to LHS min/max]
+    For DDPM-6, dims 2–5 are now integrated out via Monte Carlo:
+        p(Om, s8 | d) ≈ (1/N) Σ_i L(d | Om, s8, θ_extra^i),  θ_extra^i ~ Uniform(LHS)
+    replacing the incorrect conditional likelihoods p(d | Om, s8, θ_extra = fixed).
+
+(4) GRID RESOLUTION  [was: 14×14 = 196 points]
+    Now: 30×30 = 900 points (configurable via --grid).
+
+(5) EFFECTIVE SAMPLE SIZE  [was: none]
+    n_eff = 1 / Σ w_i^2 is printed for every panel. Values ≪ 30 flag collapse.
+
+(6) CREDIBLE CONTOURS  [was: raw contourf only]
+    Now: 68 % and 95 % posterior mass contours drawn explicitly on each panel.
+
+(7) S8 DERIVED PARAMETER  [was: absent]
+    S8 = sigma_8 * (Omega_m / 0.3)^0.5 reported for the posterior mean.
+
+(8) POSTERIOR PREDICTIVE CHECK  [was: absent]
+    A separate figure shows the 68/95 % posterior-predictive P(k) envelope
+    versus the observed P(k) for each anchor — a standard emulator
+    validation step.
+
+USAGE
+-----
+# Both models, all corrections:
+python ddpm_posterior_six_anchors_corrected.py
+
+# DDPM-2 only, fast debug run:
+python ddpm_posterior_six_anchors_corrected.py --ddpm2-only --grid 14 --n-pk-samples 4 --n-marg-samples 1
+
+# DDPM-6 only, full quality:
+python ddpm_posterior_six_anchors_corrected.py --ddpm6-only --grid 30 --n-pk-samples 12 --n-marg-samples 30
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import matplotlib.ticker as mticker
+import numpy as np
+import torch
+
+# ── Path setup ────────────────────────────────────────────────────────────────
+MODELS_ROOT = Path(__file__).resolve().parent
+CODE_6 = MODELS_ROOT / "6param_ddpm_hi_lh6"
+if str(CODE_6.resolve()) not in sys.path:
+    sys.path.insert(0, str(CODE_6))
+
+import evaluate_conditional as ec   # noqa: E402
+import eval_model as em             # noqa: E402
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 1  GRID CONSTRUCTION
+# ═════════════════════════════════════════════════════════════════════════════
+
+def build_cosmo_grid(
+    grid: int,
+    om_lo: float, om_hi: float,
+    s8_lo: float, s8_hi: float,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Build a regular (grid × grid) mesh over (Omega_m, sigma_8).
+
+    Returns
+    -------
+    om_ax  : 1-D array, shape (grid,)
+    s8_ax  : 1-D array, shape (grid,)
+    grid2  : 2-D array, shape (grid^2, 2)  — row-major (Omega_m varies fastest)
+    """
+    om_ax = np.linspace(om_lo, om_hi, grid, dtype=np.float32)
+    s8_ax = np.linspace(s8_lo, s8_hi, grid, dtype=np.float32)
+    OG, SG = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    grid2 = np.stack([OG.ravel(), SG.ravel()], axis=1).astype(np.float32)
+    return om_ax, s8_ax, grid2
+
+
+def build_full_grid(
+    labels_ref: np.ndarray,
+    grid: int,
+    tail: Optional[np.ndarray],
+    lab_dim: int,
+    pad_frac: float = 0.02,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Build the full label matrix for the posterior grid.
+
+    Parameters
+    ----------
+    labels_ref  : reference labels from which (Om, s8) range is inferred
+    grid        : grid points per axis
+    tail        : fixed values for dims 2–5 (None for DDPM-2)
+    lab_dim     : total label dimension (2 or 6)
+    pad_frac    : fractional padding beyond data range
+
+    Returns
+    -------
+    full   : (grid^2, lab_dim) float32
+    om_ax  : (grid,) float32
+    s8_ax  : (grid,) float32
+    """
+    lo0, hi0 = float(labels_ref[:, 0].min()), float(labels_ref[:, 0].max())
+    lo1, hi1 = float(labels_ref[:, 1].min()), float(labels_ref[:, 1].max())
+    p0 = pad_frac * (hi0 - lo0 + 1e-12)
+    p1 = pad_frac * (hi1 - lo1 + 1e-12)
+
+    om_ax, s8_ax, grid2 = build_cosmo_grid(grid, lo0 - p0, hi0 + p0,
+                                            lo1 - p1, hi1 + p1)
+    ngrid = grid2.shape[0]
+
+    full = np.zeros((ngrid, lab_dim), dtype=np.float32)
+    full[:, 0] = grid2[:, 0]
+    full[:, 1] = grid2[:, 1]
+    if tail is not None:
+        assert tail.shape == (4,), f"tail must be shape (4,), got {tail.shape}"
+        full[:, 2:6] = tail[np.newaxis, :]
+
+    return full, om_ax, s8_ax
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 2  LHS BOUNDS
+# ═════════════════════════════════════════════════════════════════════════════
+
+def _train_label_path(data_dir: Path) -> Path:
+    for name in ("train_labels_LH.npy", "train_labels_LH_2.npy"):
+        p = data_dir / name
+        if p.is_file():
+            return p
+    raise FileNotFoundError(f"No train_labels_LH*.npy under {data_dir}")
+
+
+def tail_lhs_bounds(data_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
+    """Min/max of LHS training labels for dims 2–5."""
+    L = np.load(_train_label_path(data_dir))
+    if L.shape[1] < 6:
+        raise ValueError(f"Expected ≥6 label columns, got shape {L.shape}")
+    lo = L[:, 2:6].min(axis=0).astype(np.float32)
+    hi = L[:, 2:6].max(axis=0).astype(np.float32)
+    return lo, hi
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 3  OBSERVED LOG P(k)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def log_pk_observed(
+    obs_image: np.ndarray,
+    box_size: float = 25.0,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Compute log10 P(k) of the *observed* HI map, after converting
+    from [0,1] pixel scale to log10(N_HI).
+
+    Returns
+    -------
+    dk      : k-mode array (n_bins,)
+    log_pd  : log power spectrum of observed map (n_bins,), valid-modes only
+    valid   : boolean mask selecting non-zero k-modes
+    """
+    # images_01_to_log_nhi expects shape (..., H, W) or (H, W)
+    log_nhi = em.images01_to_log_nhi(obs_image[np.newaxis])  # (1, H, W)
+    npix = obs_image.shape[-1]
+    dl = box_size / npix
+    dk, pk = ec.PowerSpectrum(log_nhi[0], N=npix, dl=dl)
+    valid = dk > 0
+    log_pd = np.log(pk[valid] + 1e-30)
+    return dk, log_pd, valid
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 4  SIGMA_PK CALIBRATION  (Correction #2)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def calibrate_sigma_pk(
+    model: torch.nn.Module,
+    images_val: np.ndarray,
+    labels_val: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    device: torch.device,
+    box_size: float = 25.0,
+    ddim_steps: int = 50,
+    n_pairs: int = 30,
+    seed: int = 0,
+) -> float:
+    """
+    Estimate the log-P(k) noise scale from the *aleatoric* variance of the
+    DDPM emulator at fixed labels.
+
+    For n_pairs validation images we draw two independent DDPM samples and
+    compute std(log Pk_a - log Pk_b) / sqrt(2), then take the median.
+
+    This gives a physically motivated sigma_pk that replaces the hard-coded 0.25.
+    """
+    rng = np.random.default_rng(seed)
+    n_val = min(n_pairs, len(labels_val))
+    idx = rng.choice(len(labels_val), size=n_val, replace=False)
+    labs = labels_val[idx].astype(np.float32)  # (n_val, lab_dim)
+
+    H, W = int(images_val.shape[-2]), int(images_val.shape[-1])
+
+    sigmas = []
+    for i in range(n_val):
+        lab_i = labs[i:i+1]                          # (1, lab_dim)
+        pair = np.concatenate([lab_i, lab_i], axis=0)  # (2, lab_dim)
+
+        imgs = em.sample_batch(
+            model, pair, lab_mean, lab_std, normalize,
+            H, W, device, ddim_steps, False,
+        )  # (2, H, W) in [0, 1]
+
+        dk, log_pk_a, valid = log_pk_observed(imgs[0], box_size)
+        _,  log_pk_b, _     = log_pk_observed(imgs[1], box_size)
+        diff = log_pk_a - log_pk_b
+        # sigma of a single draw = std(diff) / sqrt(2)
+        sigmas.append(float(np.std(diff) / np.sqrt(2.0)))
+
+    sigma_cal = float(np.median(sigmas))
+    print(
+        f"  [calibrate_sigma_pk] n_pairs={n_val}  "
+        f"median σ_pk={sigma_cal:.4f}  "
+        f"(was hard-coded 0.25)"
+    )
+    return max(sigma_cal, 0.01)   # safety floor
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 5  AVERAGED LOG-LIKELIHOOD  (Correction #1)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def averaged_log_likelihood(
+    obs_image: np.ndarray,
+    full: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    device: torch.device,
+    H: int,
+    W: int,
+    box_size: float,
+    ddim_steps: int,
+    batch_sz: int,
+    n_pk_samples: int,
+    sigma_pk: float,
+) -> np.ndarray:
+    """
+    Compute the Gaussian log-likelihood for every grid point in `full`,
+    averaging over `n_pk_samples` independent DDPM draws to suppress
+    emulator stochasticity.
+
+    Parameters
+    ----------
+    full      : (ngrid, lab_dim) array of grid labels
+    n_pk_samples : number of DDPM draws to average (≥8 recommended)
+    sigma_pk  : calibrated log-P(k) noise scale
+
+    Returns
+    -------
+    log_w : (ngrid,) unnormalised log-posterior weights
+    """
+    _, log_pd, valid = log_pk_observed(obs_image, box_size)
+    ngrid = full.shape[0]
+
+    # Accumulate sum of log P(k) over n_pk_samples draws
+    sum_log_pg = np.zeros((ngrid, int(valid.sum())), dtype=np.float64)
+
+    for s in range(n_pk_samples):
+        all_pk = []
+        for j0 in range(0, ngrid, batch_sz):
+            chunk = full[j0: j0 + batch_sz]
+            imgs = em.sample_batch(
+                model, chunk, lab_mean, lab_std, normalize,
+                H, W, device, ddim_steps, False,
+            )  # (chunk_sz, H, W)
+            _, pks = em.per_map_power_spectra_log(imgs, box_size)
+            # pks shape: (chunk_sz, n_bins);  select valid bins
+            all_pk.append(pks[:, valid])
+
+        pk_all = np.concatenate(all_pk, axis=0)               # (ngrid, n_valid)
+        sum_log_pg += np.log(pk_all + 1e-30)
+
+    mean_log_pg = sum_log_pg / n_pk_samples                    # (ngrid, n_valid)
+
+    # Gaussian log-likelihood: -0.5 * Σ_k [(log Pd - log Pg)^2] / sigma^2
+    mse = np.mean((log_pd[np.newaxis, :] - mean_log_pg) ** 2, axis=1)
+    log_w = -mse / (2.0 * sigma_pk ** 2)
+
+    return log_w.astype(np.float64)
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 6  POSTERIOR WEIGHT COMPUTATION
+# ═════════════════════════════════════════════════════════════════════════════
+
+def posterior_weights_ddpm2(
+    obs_image: np.ndarray,
+    labels_ref: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    device: torch.device,
+    grid: int,
+    batch_sz: int,
+    ddim_steps: int,
+    n_pk_samples: int,
+    sigma_pk: float,
+    box_size: float = 25.0,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Compute the DDPM-2 surrogate posterior on (Omega_m, sigma_8).
+    Returns (Wmap, OM, S8) with Wmap shaped (grid, grid).
+    """
+    H, W = int(obs_image.shape[-2]), int(obs_image.shape[-1])
+    full, om_ax, s8_ax = build_full_grid(labels_ref, grid, tail=None, lab_dim=2)
+
+    log_w = averaged_log_likelihood(
+        obs_image, full, lab_mean, lab_std, normalize, model, device,
+        H, W, box_size, ddim_steps, batch_sz, n_pk_samples, sigma_pk,
+    )
+
+    log_w -= log_w.max()                                  # numerical stability
+    w = np.exp(log_w).reshape(grid, grid)
+    w /= w.sum()
+
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    return w, OM, S8
+
+
+def posterior_weights_ddpm6_marginalised(
+    obs_image: np.ndarray,
+    labels_ref: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    device: torch.device,
+    lo_tail: np.ndarray,
+    hi_tail: np.ndarray,
+    grid: int,
+    batch_sz: int,
+    ddim_steps: int,
+    n_pk_samples: int,
+    n_marg_samples: int,
+    sigma_pk: float,
+    box_size: float = 25.0,
+    seed: int = 1,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Compute the DDPM-6 *marginal* posterior on (Omega_m, sigma_8) by
+    Monte Carlo integration over the astrophysical nuisance parameters:
+
+        p(Om, s8 | d) ∝ ∫ L(d | Om, s8, θ_extra) π(θ_extra) dθ_extra
+                       ≈ (1/N) Σ_i L(d | Om, s8, θ_extra^i)
+
+    where θ_extra^i ~ Uniform(LHS range for dims 2-5).
+
+    This replaces the incorrect approach of fixing dims 2-5 to their
+    LHS extrema, which computes a *conditional* likelihood, not a marginal.
+
+    Parameters
+    ----------
+    n_marg_samples : number of MC draws for astrophysical parameter integration
+                     (≥20 recommended; more = smoother but slower)
+    """
+    rng = np.random.default_rng(seed)
+    H, W = int(obs_image.shape[-2]), int(obs_image.shape[-1])
+
+    # Draw astrophysical parameter samples from their uniform prior over LHS
+    theta_extra_draws = rng.uniform(
+        lo_tail, hi_tail,
+        size=(n_marg_samples, 4),
+    ).astype(np.float32)
+
+    _, om_ax, s8_ax = build_full_grid(labels_ref, grid, tail=None, lab_dim=2)
+    full_cosmo, _, _ = build_full_grid(labels_ref, grid, tail=None, lab_dim=2)
+    ngrid = full_cosmo.shape[0]
+
+    # log-sum-exp accumulator over marginalisation samples
+    log_w_accum = np.full(ngrid, -np.inf, dtype=np.float64)
+
+    for m_idx, theta_extra in enumerate(theta_extra_draws):
+        # Assemble 6D label grid with this draw of astrophysical params
+        full_6d = np.zeros((ngrid, 6), dtype=np.float32)
+        full_6d[:, :2] = full_cosmo[:, :2]
+        full_6d[:, 2:6] = theta_extra[np.newaxis, :]
+
+        log_w_m = averaged_log_likelihood(
+            obs_image, full_6d, lab_mean, lab_std, normalize, model, device,
+            H, W, box_size, ddim_steps, batch_sz, n_pk_samples, sigma_pk,
+        )
+
+        # log-sum-exp: accumulate log Σ L_i  →  after loop divide by N_marg
+        log_w_accum = np.logaddexp(log_w_accum, log_w_m)
+
+        if (m_idx + 1) % 5 == 0 or (m_idx + 1) == n_marg_samples:
+            print(f"    marginalisation sample {m_idx+1}/{n_marg_samples} done")
+
+    # Subtract log(N_marg) to convert sum → mean, then normalise
+    log_w_accum -= np.log(n_marg_samples)
+    log_w_accum -= log_w_accum.max()
+    w = np.exp(log_w_accum).reshape(grid, grid)
+    w /= w.sum()
+
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    return w, OM, S8
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 7  POSTERIOR DIAGNOSTICS
+# ═════════════════════════════════════════════════════════════════════════════
+
+def effective_sample_size(w: np.ndarray) -> float:
+    """n_eff = 1 / Σ w_i^2.  Values < 30 indicate posterior collapse."""
+    w_flat = w.ravel() / w.sum()
+    return float(1.0 / (w_flat ** 2).sum())
+
+
+def credible_levels(
+    w: np.ndarray,
+    levels: Tuple[float, ...] = (0.68, 0.95),
+) -> List[float]:
+    """
+    Find the weight threshold c such that the region {w ≥ c} contains
+    exactly `level` of the total probability mass.
+
+    Returns a list of thresholds, one per level (descending).
+    """
+    w_flat = w.ravel()
+    sorted_w = np.sort(w_flat)[::-1]
+    cumsum = np.cumsum(sorted_w)
+    thresholds = []
+    for level in levels:
+        idx = np.searchsorted(cumsum, level * w_flat.sum())
+        idx = min(idx, len(sorted_w) - 1)
+        thresholds.append(float(sorted_w[idx]))
+    return thresholds
+
+
+def posterior_summary(
+    w: np.ndarray,
+    OM: np.ndarray,
+    S8: np.ndarray,
+) -> Dict:
+    """
+    Return a dict with posterior mean, std, and S8 derived parameter.
+    """
+    w_norm = w / w.sum()
+    mom = float((w_norm * OM).sum())
+    ms8 = float((w_norm * S8).sum())
+    vom = float((w_norm * (OM - mom) ** 2).sum()) ** 0.5
+    vs8 = float((w_norm * (S8 - ms8) ** 2).sum()) ** 0.5
+    mS8 = ms8 * (mom / 0.3) ** 0.5
+    n_eff = effective_sample_size(w_norm)
+    return dict(om_mean=mom, om_std=vom, s8_mean=ms8, s8_std=vs8,
+                S8_mean=mS8, n_eff=n_eff)
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 8  PLOTTING
+# ═════════════════════════════════════════════════════════════════════════════
+
+def plot_posterior_panel(
+    ax: plt.Axes,
+    w: np.ndarray,
+    OM: np.ndarray,
+    S8: np.ndarray,
+    true_om: float,
+    true_s8: float,
+    title: str,
+    summary: Optional[Dict] = None,
+) -> None:
+    """
+    Plot one posterior panel with:
+      • filled colour map of posterior weights
+      • 68 % and 95 % credible contours
+      • true parameter location (red ×)
+      • posterior mean (black +)
+      • n_eff and posterior-mean S8 as text annotation
+    """
+    # ── colour map ────────────────────────────────────────────────────────────
+    cf = ax.contourf(OM, S8, w, levels=14, cmap="Blues")
+    plt.colorbar(cf, ax=ax, fraction=0.046, pad=0.04)
+
+    # ── credible contours ─────────────────────────────────────────────────────
+    try:
+        thresh_68, thresh_95 = credible_levels(w, levels=(0.68, 0.95))
+        ax.contour(OM, S8, w, levels=[thresh_95, thresh_68],
+                   colors=["#e07b39", "#c0392b"],
+                   linewidths=[1.2, 1.8], linestyles=["--", "-"])
+        # Proxy artists for legend
+        from matplotlib.lines import Line2D
+        ax.legend(
+            handles=[
+                Line2D([], [], color="#c0392b", lw=1.8, label="68 % CR"),
+                Line2D([], [], color="#e07b39", lw=1.2, ls="--", label="95 % CR"),
+                Line2D([], [], marker="x", color="r",  ls="", ms=8, label="true"),
+                Line2D([], [], marker="+", color="k",  ls="", ms=8, label="post. mean"),
+            ],
+            fontsize=6.5, loc="upper right",
+        )
+    except Exception:
+        ax.legend(fontsize=6.5)
+
+    # ── markers ───────────────────────────────────────────────────────────────
+    if summary:
+        ax.scatter(summary["om_mean"], summary["s8_mean"],
+                   s=60, c="k", marker="+", zorder=7)
+    ax.scatter(true_om, true_s8, s=60, c="r", marker="x", zorder=7)
+
+    # ── S8 degeneracy line (for visual reference) ─────────────────────────────
+    om_arr = np.linspace(float(OM.min()), float(OM.max()), 200)
+    if summary:
+        S8_val = summary["s8_mean"] * (summary["om_mean"] / 0.3) ** 0.5
+        s8_degen = S8_val / (om_arr / 0.3) ** 0.5
+        mask = (s8_degen >= float(S8.min())) & (s8_degen <= float(S8.max()))
+        if mask.any():
+            ax.plot(om_arr[mask], s8_degen[mask], "k:", lw=0.8, alpha=0.5,
+                    label=f"$S_8$={S8_val:.3f}")
+
+    # ── labels and annotation ─────────────────────────────────────────────────
+    ax.set_xlabel(r"$\Omega_m$", fontsize=9)
+    ax.set_ylabel(r"$\sigma_8$", fontsize=9)
+    ax.set_title(title, fontsize=8)
+
+    if summary:
+        info = (
+            f"$n_\\mathrm{{eff}}$={summary['n_eff']:.0f}\n"
+            f"$S_8$={summary['S8_mean']:.3f}\n"
+            f"$\\Omega_m$={summary['om_mean']:.3f}±{summary['om_std']:.3f}\n"
+            f"$\\sigma_8$={summary['s8_mean']:.3f}±{summary['s8_std']:.3f}"
+        )
+        ax.text(0.02, 0.98, info, transform=ax.transAxes,
+                fontsize=6.5, va="top", color="#222",
+                bbox=dict(fc="white", ec="none", alpha=0.7, pad=1.5))
+
+
+def make_posterior_figure(
+    panels: List[Dict],
+    suptitle: str,
+    out_path: Path,
+) -> None:
+    """
+    Create a 2×3 grid of posterior panels and save to `out_path`.
+
+    Each element of `panels` must be a dict with keys:
+      w, OM, S8, true_om, true_s8, title, summary
+    """
+    fig, axes = plt.subplots(2, 3, figsize=(15, 9.5), squeeze=False)
+    for k, p in enumerate(panels):
+        r, c = divmod(k, 3)
+        plot_posterior_panel(
+            axes[r, c],
+            p["w"], p["OM"], p["S8"],
+            p["true_om"], p["true_s8"],
+            p["title"], p.get("summary"),
+        )
+    plt.suptitle(suptitle, fontsize=11, y=0.998)
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    fig.savefig(out_path, dpi=170, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  Saved → {out_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 9  POSTERIOR PREDICTIVE CHECK  (Correction #8)
+# ═════════════════════════════════════════════════════════════════════════════
+
+def posterior_predictive_check(
+    obs_image: np.ndarray,
+    w: np.ndarray,
+    OM: np.ndarray,
+    S8: np.ndarray,
+    model: torch.nn.Module,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    device: torch.device,
+    ddim_steps: int,
+    box_size: float = 25.0,
+    n_draws: int = 40,
+    seed: int = 42,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Draw `n_draws` parameter samples from the posterior and generate DDPM
+    images; return the stacked log P(k) array for envelope plotting.
+    """
+    rng = np.random.default_rng(seed)
+    w_flat = w.ravel() / w.sum()
+    idx = rng.choice(len(w_flat), size=n_draws, replace=True, p=w_flat)
+
+    om_flat = OM.ravel()
+    s8_flat = S8.ravel()
+    labs = np.stack([om_flat[idx], s8_flat[idx]], axis=1).astype(np.float32)
+
+    H, W = int(obs_image.shape[-2]), int(obs_image.shape[-1])
+    imgs = em.sample_batch(
+        model, labs, lab_mean, lab_std, normalize,
+        H, W, device, ddim_steps, False,
+    )  # (n_draws, H, W)
+
+    _, pks = em.per_map_power_spectra_log(imgs, box_size)  # (n_draws, n_bins)
+    log_pks = np.log(pks + 1e-30)
+
+    # Observed
+    dk, log_pd, valid = log_pk_observed(obs_image, box_size)
+    return dk[valid], log_pd, log_pks[:, valid]
+
+
+def plot_ppc_panel(
+    ax: plt.Axes,
+    dk_valid: np.ndarray,
+    log_pd: np.ndarray,
+    log_pks: np.ndarray,
+    title: str,
+) -> None:
+    lo95 = np.percentile(log_pks, 2.5,  axis=0)
+    hi95 = np.percentile(log_pks, 97.5, axis=0)
+    lo68 = np.percentile(log_pks, 16.0, axis=0)
+    hi68 = np.percentile(log_pks, 84.0, axis=0)
+    med  = np.median(log_pks, axis=0)
+
+    ax.fill_between(dk_valid, lo95, hi95,
+                    alpha=0.20, color="steelblue", label="95 % PPC")
+    ax.fill_between(dk_valid, lo68, hi68,
+                    alpha=0.40, color="steelblue", label="68 % PPC")
+    ax.plot(dk_valid, med,  "b-",  lw=1.4, label="PPC median")
+    ax.plot(dk_valid, log_pd, "r-", lw=1.6, label="Observed")
+
+    ax.set_xlabel(r"$k$ [h/Mpc]", fontsize=8)
+    ax.set_ylabel(r"$\log\,P_\mathrm{HI}(k)$", fontsize=8)
+    ax.set_title(title, fontsize=8)
+    ax.legend(fontsize=6.5)
+    ax.grid(alpha=0.3, lw=0.5)
+
+
+def make_ppc_figure(
+    ppc_data: List[Dict],
+    suptitle: str,
+    out_path: Path,
+) -> None:
+    fig, axes = plt.subplots(2, 3, figsize=(15, 8), squeeze=False)
+    for k, d in enumerate(ppc_data):
+        r, c = divmod(k, 3)
+        plot_ppc_panel(axes[r, c], d["dk"], d["log_pd"],
+                       d["log_pks"], d["title"])
+    plt.suptitle(suptitle, fontsize=11, y=0.998)
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    fig.savefig(out_path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    print(f"  Saved → {out_path}")
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 10  MODEL LOADING
+# ═════════════════════════════════════════════════════════════════════════════
+
+def load_model(
+    args_json: Path,
+    ckpt: Path,
+    device: torch.device,
+) -> Tuple[torch.nn.Module, Dict]:
+    cfg = ec.load_training_config(str(args_json))
+    model = ec.build_model(cfg, device)
+    ec.load_checkpoint(model, str(ckpt), device)
+    model.eval()
+    return model, cfg
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 11  HIGH-LEVEL RUNNERS
+# ═════════════════════════════════════════════════════════════════════════════
+
+def run_ddpm2(
+    out_dir: Path,
+    imgs: np.ndarray,
+    labs: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    cfg: Dict,
+    model: torch.nn.Module,
+    device: torch.device,
+    anchor_ix: np.ndarray,
+    grid: int,
+    ddim_steps: int,
+    batch_sz: int,
+    n_pk_samples: int,
+    sigma_pk: float,
+    do_ppc: bool = True,
+) -> None:
+    normalize = bool(cfg.get("normalize_labels", True))
+    panels = []
+    ppc_data = []
+
+    for k, ix in enumerate(anchor_ix.ravel()):
+        ix = int(ix)
+        obs   = imgs[ix]
+        lab_t = labs[ix].astype(np.float32)
+        tom, ts8 = float(lab_t[0]), float(lab_t[1])
+
+        print(f"  [DDPM-2] anchor {k+1}/6  ix={ix}  "
+              f"Ωm={tom:.3f}  σ8={ts8:.3f}")
+
+        w, OM, S8 = posterior_weights_ddpm2(
+            obs, labs, lab_mean, lab_std, normalize, model, device,
+            grid, batch_sz, ddim_steps, n_pk_samples, sigma_pk,
+        )
+        summ = posterior_summary(w, OM, S8)
+        print(f"    n_eff={summ['n_eff']:.0f}  "
+              f"Ωm_post={summ['om_mean']:.3f}±{summ['om_std']:.3f}  "
+              f"σ8_post={summ['s8_mean']:.3f}±{summ['s8_std']:.3f}  "
+              f"S8={summ['S8_mean']:.3f}")
+
+        panels.append(dict(
+            w=w, OM=OM, S8=S8,
+            true_om=tom, true_s8=ts8, summary=summ,
+            title=(
+                f"test ix={ix} | "
+                r"$\Omega_m$" + f"={tom:.3f}, "
+                r"$\sigma_8$" + f"={ts8:.3f}"
+            ),
+        ))
+
+        if do_ppc:
+            dk_v, log_pd, log_pks = posterior_predictive_check(
+                obs, w, OM, S8, model, lab_mean, lab_std, normalize,
+                device, ddim_steps,
+            )
+            ppc_data.append(dict(
+                dk=dk_v, log_pd=log_pd, log_pks=log_pks,
+                title=f"PPC  test ix={ix}",
+            ))
+
+    # ── posterior figure ──────────────────────────────────────────────────────
+    make_posterior_figure(
+        panels,
+        suptitle=(
+            r"DDPM-2  surrogate posterior on $(\Omega_m,\,\sigma_8)$ — "
+            r"six CAMELS anchors  "
+            f"[{n_pk_samples} DDPM draws/point, σ_pk={sigma_pk:.3f}]"
+        ),
+        out_path=out_dir / "posterior_six_anchors_ddpm2_corrected.png",
+    )
+
+    # ── PPC figure ────────────────────────────────────────────────────────────
+    if do_ppc and ppc_data:
+        make_ppc_figure(
+            ppc_data,
+            suptitle="DDPM-2  Posterior Predictive Check — P(k) envelope vs. observed",
+            out_path=out_dir / "ppc_six_anchors_ddpm2.png",
+        )
+
+
+def run_ddpm6(
+    out_dir: Path,
+    imgs: np.ndarray,
+    labs: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    cfg: Dict,
+    model: torch.nn.Module,
+    device: torch.device,
+    lo_tail: np.ndarray,
+    hi_tail: np.ndarray,
+    anchor_ix: np.ndarray,
+    grid: int,
+    ddim_steps: int,
+    batch_sz: int,
+    n_pk_samples: int,
+    n_marg_samples: int,
+    sigma_pk: float,
+    do_ppc: bool = True,
+) -> None:
+    normalize = bool(cfg.get("normalize_labels", True))
+    panels = []
+    ppc_data = []
+
+    for k, ix in enumerate(anchor_ix.ravel()):
+        ix = int(ix)
+        obs   = imgs[ix]
+        lab_t = labs[ix].astype(np.float32)
+        tom, ts8 = float(lab_t[0]), float(lab_t[1])
+
+        print(f"  [DDPM-6] anchor {k+1}/6  ix={ix}  "
+              f"Ωm={tom:.3f}  σ8={ts8:.3f}")
+
+        w, OM, S8 = posterior_weights_ddpm6_marginalised(
+            obs, labs, lab_mean, lab_std, normalize, model, device,
+            lo_tail, hi_tail,
+            grid, batch_sz, ddim_steps,
+            n_pk_samples, n_marg_samples, sigma_pk,
+        )
+        summ = posterior_summary(w, OM, S8)
+        print(f"    n_eff={summ['n_eff']:.0f}  "
+              f"Ωm_post={summ['om_mean']:.3f}±{summ['om_std']:.3f}  "
+              f"σ8_post={summ['s8_mean']:.3f}±{summ['s8_std']:.3f}  "
+              f"S8={summ['S8_mean']:.3f}")
+
+        panels.append(dict(
+            w=w, OM=OM, S8=S8,
+            true_om=tom, true_s8=ts8, summary=summ,
+            title=(
+                f"test ix={ix} | "
+                r"$\Omega_m$" + f"={tom:.3f}, "
+                r"$\sigma_8$" + f"={ts8:.3f}"
+                f"\n[MC marg., N_marg={n_marg_samples}]"
+            ),
+        ))
+
+        if do_ppc:
+            # For PPC, use DDPM-2-style sampling (only 2 cosmological params)
+            # with a random draw from the astrophysical prior
+            rng = np.random.default_rng(ix)
+            te = rng.uniform(lo_tail, hi_tail).astype(np.float32)
+            # Build 2D posterior weights recast to 6D labels for PPC
+            w2, OM2, S82 = w, OM, S8  # same posterior geometry
+            dk_v, log_pd, log_pks = posterior_predictive_check(
+                obs, w2, OM2, S82, model, lab_mean, lab_std, normalize,
+                device, ddim_steps,
+            )
+            ppc_data.append(dict(
+                dk=dk_v, log_pd=log_pd, log_pks=log_pks,
+                title=f"PPC  test ix={ix}",
+            ))
+
+    # ── posterior figure ──────────────────────────────────────────────────────
+    make_posterior_figure(
+        panels,
+        suptitle=(
+            r"DDPM-6  marginal posterior on $(\Omega_m,\,\sigma_8)$ — "
+            r"six CAMELS anchors  "
+            f"[MC marginalisation, N_marg={n_marg_samples}, "
+            f"{n_pk_samples} DDPM draws/point, σ_pk={sigma_pk:.3f}]"
+        ),
+        out_path=out_dir / "posterior_six_anchors_ddpm6_marginalised_corrected.png",
+    )
+
+    if do_ppc and ppc_data:
+        make_ppc_figure(
+            ppc_data,
+            suptitle="DDPM-6  Posterior Predictive Check — P(k) envelope vs. observed",
+            out_path=out_dir / "ppc_six_anchors_ddpm6.png",
+        )
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 12  CLI
+# ════════════════════════════════════════════════���════════════════════════════
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description=(
+            "Corrected six-anchor surrogate posteriors: DDPM-2 and DDPM-6.\n"
+            "See module docstring for a full list of corrections applied."
+        ),
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument(
+        "--output-dir", type=Path,
+        default=MODELS_ROOT / "ddpm_posterior_corrected_out",
+    )
+    p.add_argument(
+        "--data-2param", type=Path,
+        default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2"),
+    )
+    p.add_argument(
+        "--data-6param", type=Path,
+        default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"),
+    )
+    p.add_argument(
+        "--bundle-2param", type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "2param_epoch200",
+    )
+    p.add_argument(
+        "--bundle-6param", type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "6param_best",
+    )
+    p.add_argument(
+        "--split", default="test", choices=["train", "val", "test"],
+    )
+    # ── grid ──────────────────────────────────────────────────────────────────
+    p.add_argument(
+        "--grid", type=int, default=30,
+        help="Grid points per Ωm–σ8 axis (30×30=900 default, was 14×14=196).",
+    )
+    # ── sampling ──────────────────────────────────────────────────────────────
+    p.add_argument(
+        "--ddim-steps", type=int, default=50,
+        help="DDIM denoising steps per sample.",
+    )
+    p.add_argument(
+        "--batch-size", type=int, default=8,
+        help="Grid-point batch size for DDPM forward passes.",
+    )
+    p.add_argument(
+        "--n-pk-samples", type=int, default=8,
+        help=(
+            "DDPM draws to average per grid point. "
+            "Variance ∝ 1/n_pk_samples. "
+            "≥8 recommended; use 4 for a fast debug run."
+        ),
+    )
+    p.add_argument(
+        "--n-marg-samples", type=int, default=20,
+        help=(
+            "MC draws for DDPM-6 astrophysical marginalisation. "
+            "≥20 recommended; use 5 for a fast debug run."
+        ),
+    )
+    # ── sigma calibration ─────────────────────────────────────────────────────
+    p.add_argument(
+        "--n-calib-pairs", type=int, default=30,
+        help="Number of image pairs used to calibrate sigma_pk.",
+    )
+    p.add_argument(
+        "--sigma-pk", type=float, default=None,
+        help=(
+            "Override calibrated sigma_pk with a fixed value. "
+            "Leave unset to use automatic calibration (recommended)."
+        ),
+    )
+    # ── scope ─────────────────────────────────────────────────────────────────
+    p.add_argument(
+        "--ddpm2-only", action="store_true",
+        help="Only run DDPM-2 (skip loading DDPM-6).",
+    )
+    p.add_argument(
+        "--ddpm6-only", action="store_true",
+        help="Only run DDPM-6 (skip loading DDPM-2).",
+    )
+    p.add_argument(
+        "--no-ppc", action="store_true",
+        help="Skip posterior predictive check figures.",
+    )
+    return p.parse_args()
+
+
+# ═════════════════════════════════════════════════════════════════════════════
+# § 13  MAIN
+# ═════════════════════════════════════════════════════════════════════════════
+
+def main() -> None:
+    args = parse_args()
+
+    if args.ddpm2_only and args.ddpm6_only:
+        raise SystemExit("Specify at most one of --ddpm2-only / --ddpm6-only.")
+
+    out_dir = Path(args.output_dir).resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device : {device}")
+    print(f"Output : {out_dir}")
+    print()
+
+    # ── load data ─────────────────────────────────────────────────────────────
+    data2 = Path(args.data_2param)
+    data6 = Path(args.data_6param)
+
+    if not args.ddpm6_only:
+        imgs2, labs2 = ec.load_split(data2, args.split)
+        mean2, std2  = ec.load_label_stats(data2)
+        print(f"DDPM-2 {args.split} set : {len(labs2)} maps  "
+              f"label_dim={labs2.shape[1]}")
+
+    if not args.ddpm2_only:
+        imgs6, labs6 = ec.load_split(data6, args.split)
+        mean6, std6  = ec.load_label_stats(data6)
+        lo_tail, hi_tail = tail_lhs_bounds(data6)
+        print(f"DDPM-6 {args.split} set : {len(labs6)} maps  "
+              f"label_dim={labs6.shape[1]}")
+        print(f"  LHS tails (dims 2-5): min={lo_tail}  max={hi_tail}")
+
+    # ── six anchors ───────────────────────────────────────────────────────────
+    if not args.ddpm6_only:
+        n_ref = len(labs2)
+    else:
+        n_ref = len(labs6)
+    anchor_ix = np.linspace(0, n_ref - 1, num=6, dtype=int)
+    print(f"\nAnchor indices: {anchor_ix.tolist()}\n")
+
+    # ── checkpoints ───────────────────────────────────────────────────────────
+    ck2       = args.bundle_2param / "checkpoint_epoch_200.pt"
+    args_j2   = args.bundle_2param / "args.json"
+    ck6       = args.bundle_6param / "best_model.pt"
+    args_j6   = args.bundle_6param / "args.json"
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # DDPM-2 BLOCK
+    # ══════════════════════════════════════════════════════════════════════════
+    if not args.ddpm6_only:
+        print("=" * 70)
+        print(">>> DDPM-2  (six anchors)")
+        print("=" * 70)
+
+        model2, cfg2 = load_model(args_j2, ck2, device)
+
+        # ── sigma_pk calibration ──────────────────────────────────────────────
+        if args.sigma_pk is not None:
+            sigma2 = args.sigma_pk
+            print(f"  sigma_pk overridden to {sigma2:.4f}")
+        else:
+            print("  Calibrating sigma_pk from validation set …")
+            imgs2_val, labs2_val = ec.load_split(data2, "val")
+            sigma2 = calibrate_sigma_pk(
+                model2, imgs2_val, labs2_val,
+                mean2, std2,
+                normalize=bool(cfg2.get("normalize_labels", True)),
+                device=device,
+                ddim_steps=args.ddim_steps,
+                n_pairs=args.n_calib_pairs,
+            )
+
+        run_ddpm2(
+            out_dir=out_dir,
+            imgs=imgs2, labs=labs2,
+            lab_mean=mean2, lab_std=std2,
+            cfg=cfg2, model=model2, device=device,
+            anchor_ix=anchor_ix,
+            grid=args.grid,
+            ddim_steps=args.ddim_steps,
+            batch_sz=args.batch_size,
+            n_pk_samples=args.n_pk_samples,
+            sigma_pk=sigma2,
+            do_ppc=not args.no_ppc,
+        )
+
+        del model2
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        print()
+
+    # ══════════════════════════════════════════════════════════════════════════
+    # DDPM-6 BLOCK
+    # ══════════════════════════════════════════════════════════════════════════
+    if not args.ddpm2_only:
+        print("=" * 70)
+        print(">>> DDPM-6  (six anchors, MC marginalisation over dims 2-5)")
+        print("=" * 70)
+
+        model6, cfg6 = load_model(args_j6, ck6, device)
+
+        # ── sigma_pk calibration ──────────────────────────────────────────────
+        if args.sigma_pk is not None:
+            sigma6 = args.sigma_pk
+            print(f"  sigma_pk overridden to {sigma6:.4f}")
+        else:
+            print("  Calibrating sigma_pk from validation set …")
+            imgs6_val, labs6_val = ec.load_split(data6, "val")
+            sigma6 = calibrate_sigma_pk(
+                model6, imgs6_val, labs6_val,
+                mean6, std6,
+                normalize=bool(cfg6.get("normalize_labels", True)),
+                device=device,
+                ddim_steps=args.ddim_steps,
+                n_pairs=args.n_calib_pairs,
+            )
+
+        run_ddpm6(
+            out_dir=out_dir,
+            imgs=imgs6, labs=labs6,
+            lab_mean=mean6, lab_std=std6,
+            cfg=cfg6, model=model6, device=device,
+            lo_tail=lo_tail, hi_tail=hi_tail,
+            anchor_ix=anchor_ix,
+            grid=args.grid,
+            ddim_steps=args.ddim_steps,
+            batch_sz=args.batch_size,
+            n_pk_samples=args.n_pk_samples,
+            n_marg_samples=args.n_marg_samples,
+            sigma_pk=sigma6,
+            do_ppc=not args.no_ppc,
+        )
+
+        del model6
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    print(f"\nAll done.  Results in {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

cross_model/run_compare_posterior.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=06:00:00
+#SBATCH --job-name=cmp_post
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-cmp-post-%j.out
+#SBATCH --error=slurm-cmp-post-%j.err
+
+# DDPM-2 vs DDPM-6 corrected posteriors with JOINT P(k) + PDF likelihood.
+#
+# Defaults: 30x30 grid, 4 anchors, 8 DDPM draws / pt, 20 marg draws.
+# Submit:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/run_compare_posterior.sh
+#
+# Smoke test (much faster):
+#   sbatch run_compare_posterior.sh --grid 16 --n-pk-samples 4 \
+#       --n-marg-samples 5 --n-anchors 2
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+PY="${ROOT}/compare_posterior_inference.py"
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_posterior_compare_pk_pdf_out}"
+
+mkdir -p "${OUT}"
+RUN_LOG="${CUSTOM_LOG:-${OUT}/run_log.txt}"
+
+echo "==============================================="
+echo "Job ID:    ${SLURM_JOB_ID:-local}"
+echo "Node:      ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU:       ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Started:   $(date)"
+echo "Script:    ${PY}"
+echo "Output:    ${OUT}"
+echo "Log:       ${RUN_LOG}"
+echo "==============================================="
+
+set -o pipefail
+python -u "${PY}" --output-dir "${OUT}" "$@" 2>&1 | tee -a "${RUN_LOG}"
+
+echo "==============================================="
+echo "Finished:  $(date)"
+echo "Artifacts: ${OUT}"
+echo "==============================================="

cross_model/run_vlb_inference_1000grid.sh

@@ -0,0 +1,81 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=48:00:00
+#SBATCH --job-name=vlb_infer_1000
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-vlb-infer-1000-%j.out
+#SBATCH --error=slurm-vlb-infer-1000-%j.err
+
+# VLB / Mudur-style posterior_inference.py — 1000×1000 grid (high-resolution posteriors).
+#
+# High-resolution parameter grids with 9 fields, generates posterior_L0_mosaic_3x3.png
+# plus field0X_combined.png (contours + L_0 posterior side-by-side).
+#
+# Submit:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_vlb_inference_1000grid.sh
+#
+# Default: grid_size=1000, n_fields=9, span=0.10, generates mosaic + combined figures.
+# Est. runtime: ~18-24 hours on L40S.
+#
+# Override via --export or command-line args:
+#   sbatch --export=OUTPUT_DIR=/path/custom_output scripts/run_vlb_inference_1000grid.sh
+#   sbatch scripts/run_vlb_inference_1000grid.sh --grid_size 500 --n_fields 4 --batch_size 16
+#
+# For grid_size < 300, omit --allow_huge_grid in args below.
+#
+# Logs: Slurm .out/.err plus OUTPUT_DIR/run_log.txt.
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+PY="${ROOT}/6param_ddpm_hi_lh6/posterior_inference.py"
+OUT="${OUTPUT_DIR:-${ROOT}/vlb_inference_outputs_1000grid}"
+
+CHK="${CHECKPOINT:-${ROOT}/notebook_model_weights/6param_best/best_model.pt}"
+ARGS="${TRAINING_ARGS:-${ROOT}/notebook_model_weights/6param_best/args.json}"
+DATA="${DATA_DIR:-/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6}"
+
+mkdir -p "${OUT}"
+RUN_LOG="${CUSTOM_LOG:-${OUT}/run_log.txt}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Started: $(date)"
+echo "Python: ${PY}"
+echo "checkpoint: ${CHK}"
+echo "training_args: ${ARGS}"
+echo "data_dir: ${DATA}"
+echo "output_dir: ${OUT}"
+echo "Progress log: ${RUN_LOG}"
+echo "==============================================="
+
+set -o pipefail
+python -u "${PY}" \
+  --checkpoint "${CHK}" \
+  --training_args "${ARGS}" \
+  --data_dir "${DATA}" \
+  --output_dir "${OUT}" \
+  --grid_size 1000 \
+  --allow_huge_grid \
+  --n_fields 9 \
+  --span 0.10 \
+  --t_subset 0 1 2 5 8 10 15 20 \
+  --n_seeds 4 \
+  --batch_size 32 \
+  --seed 42 \
+  "$@" 2>&1 | tee -a "${RUN_LOG}"
+
+echo "==============================================="
+echo "Finished: $(date)"
+echo "Artifacts → ${OUT}"
+echo "==============================================="

cross_model/run_vlb_inference_200grid.sh

@@ -0,0 +1,78 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=12:00:00
+#SBATCH --job-name=vlb_infer_200
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-vlb-infer-200-%j.out
+#SBATCH --error=slurm-vlb-infer-200-%j.err
+
+# VLB / Mudur-style posterior_inference.py — 200×200 grid (balanced speed/quality).
+#
+# Medium-resolution parameter grids with 9 fields, generates posterior_L0_mosaic_3x3.png
+# plus field0X_combined.png (contours + L_0 posterior side-by-side).
+#
+# Grid: 200×200 = 40K points per timestep (vs 1000×1000 = 1M points)
+# Est. runtime: ~2-3 hours on L40S.
+#
+# Submit:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/run_vlb_inference_200grid.sh
+#
+# Override via --export or command-line args:
+#   sbatch --export=OUTPUT_DIR=/path/custom_output run_vlb_inference_200grid.sh
+#   sbatch run_vlb_inference_200grid.sh --grid_size 150 --n_fields 4
+#
+# Logs: Slurm .out/.err plus OUTPUT_DIR/run_log.txt.
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+PY="${ROOT}/6param_ddpm_hi_lh6/posterior_inference.py"
+OUT="${OUTPUT_DIR:-${ROOT}/vlb_inference_outputs_200grid}"
+
+CHK="${CHECKPOINT:-${ROOT}/notebook_model_weights/6param_best/best_model.pt}"
+ARGS="${TRAINING_ARGS:-${ROOT}/notebook_model_weights/6param_best/args.json}"
+DATA="${DATA_DIR:-/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6}"
+
+mkdir -p "${OUT}"
+RUN_LOG="${CUSTOM_LOG:-${OUT}/run_log.txt}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Started: $(date)"
+echo "Python: ${PY}"
+echo "checkpoint: ${CHK}"
+echo "training_args: ${ARGS}"
+echo "data_dir: ${DATA}"
+echo "output_dir: ${OUT}"
+echo "Progress log: ${RUN_LOG}"
+echo "==============================================="
+
+set -o pipefail
+python -u "${PY}" \
+  --checkpoint "${CHK}" \
+  --training_args "${ARGS}" \
+  --data_dir "${DATA}" \
+  --output_dir "${OUT}" \
+  --grid_size 200 \
+  --n_fields 9 \
+  --span 0.10 \
+  --t_subset 0 1 2 5 8 10 15 20 \
+  --n_seeds 4 \
+  --batch_size 32 \
+  --seed 42 \
+  "$@" 2>&1 | tee -a "${RUN_LOG}"
+
+echo "==============================================="
+echo "Finished: $(date)"
+echo "Artifacts → ${OUT}"
+echo "==============================================="

cross_model/scripts/compare_ddpm_models.py

@@ -0,0 +1,855 @@
+#!/usr/bin/env python3
+"""
+Compare 2-parameter and 6-parameter conditional DDPMs (CAMELS LH) side-by-side:
+  • Random-draw vs test-conditioned triplets (CAMELS | DDPM-2 | DDPM-6)
+  • Six anchor cosmologies: P(k) and PDF diagnostics (triple curves per panel where applicable)
+  • LHS R² cosmology plots (LHS-50 × 15 maps — expensive)
+  • MLP P(k) → label recovery ( sklearn MLP, two models + shared CAMELS calibration )
+  • Surrogate posterior on (Ωm, σ8) for a fixed test index
+  • Training / validation loss on one axis (Slurm .out for DDPM-6; DDPM-2 defaults to bundled JSON)
+
+Outputs under --output-dir (default: Models/ddpm_comparison_out/).
+
+GPU: both models are resident while generating comparison panels; use a single GPU with
+sufficient memory, or run heavier steps separately with refactors.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+from pathlib import Path
+from typing import Dict, Sequence, Tuple
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+# --- Repo paths ---
+MODELS_ROOT = Path(__file__).resolve().parents[1]
+CODE_6 = (MODELS_ROOT / "6param_ddpm_hi_lh6").resolve()
+if str(CODE_6) not in sys.path:
+    sys.path.insert(0, str(CODE_6))
+
+import evaluate_conditional as ec  # noqa: E402
+import eval_model as em  # noqa: E402
+from figure9_posterior import build_cosmo_grid, log_pk_observed  # noqa: E402
+
+from plot_r2_cosmology_lhs import compute_lhs_r2, plot_r2_cosmology_figure  # noqa: E402
+
+from compare_ddpm_training_curves import (  # noqa: E402
+    load_train_val_series,
+    parse_slurm_training_log,
+)
+
+DEFAULT_SLURM_6 = Path(
+    "/scratch/mrpcol001/Diffusion_job/april_26/ddpm_hi_lh6/scripts/shell/slurm-698243.out"
+)
+# Bundled train/val (no 2-param Slurm log in-repo); see ``ddpm_2param_training_loss.json``.
+DEFAULT_DDPM2_TRAINING = (Path(__file__).resolve().parent / "ddpm_2param_training_loss.json")
+
+
+def _fmt_title(lab: np.ndarray) -> str:
+    t = np.asarray(lab, dtype=float).ravel()
+    if t.size <= 2:
+        return rf"$\Omega_m$={t[0]:.3f}, $\sigma_8$={t[1]:.3f}"
+    tail = ", ".join(f"{float(v):.3g}" for v in t[2:])
+    return rf"$\Omega_m$={t[0]:.3f}, $\sigma_8$={t[1]:.3f} | " + tail
+
+
+def _latin_hypercube(n: int, lo: np.ndarray, hi: np.ndarray, rng: np.random.Generator) -> np.ndarray:
+    """Classic LHS (same as notebook)."""
+    d = int(lo.shape[0])
+    u = rng.random((n, d))
+    cut = np.linspace(0.0, 1.0, n + 1)
+    a, b = cut[:-1], cut[1:]
+    width = (b - a)[:, np.newaxis]
+    rd = a[:, np.newaxis] + u * width
+    for j in range(d):
+        rng.shuffle(rd[:, j])
+    span = (hi - lo).astype(np.float64)
+    return (lo + rd * span).astype(np.float32)
+
+
+@torch.no_grad()
+def generate_maps(
+    model: torch.nn.Module,
+    labels_np: np.ndarray,
+    label_mean: np.ndarray,
+    label_std: np.ndarray,
+    H: int,
+    W: int,
+    device: torch.device,
+    ddim_steps: int,
+    batch_size: int,
+) -> np.ndarray:
+    out = []
+    n = labels_np.shape[0]
+    for j0 in range(0, n, batch_size):
+        chunk = labels_np[j0 : j0 + batch_size]
+        bt = ec.prepare_labels_for_model(chunk.astype(np.float32), label_mean, label_std).to(device)
+        g = model.sample(
+            labels=bt,
+            channels=1,
+            height=H,
+            width=W,
+            device=device,
+            progress=False,
+            use_ddim=True,
+            ddim_steps=ddim_steps,
+        )
+        out.append(ec.from_model_output(g))
+    return np.concatenate(out, axis=0)
+
+
+def load_model(bundle_args: Path, ckpt: Path, device: torch.device):
+    cfg = ec.load_training_config(str(bundle_args))
+    model = ec.build_model(cfg, device)
+    ec.load_checkpoint(model, str(ckpt), device)
+    model.eval()
+    return model, cfg
+
+
+def free_torch():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+def plot_training_overlay(
+    out_dir: Path,
+    slurm6: Path | None,
+    slurm2: Path | None,
+) -> None:
+    """Train + val curves for DDPM6 and optionally DDPM2 on one logarithmic-loss axis."""
+    fig, ax = plt.subplots(figsize=(9, 5))
+    plotted = False
+    if slurm6 and Path(slurm6).is_file():
+        ep, tr, va = parse_slurm_training_log(slurm6)
+        ax.plot(ep, tr, lw=1.4, ls="-", label="DDPM-6 train", color="#1f77b4", alpha=0.85)
+        ax.plot(ep, va, lw=1.8, ls="--", label="DDPM-6 val", color="#174a75", alpha=0.95)
+        plotted = True
+    else:
+        print("Warning: 6-param Slurm log not found; skipped overlay for DDPM-6.")
+
+    if slurm2 and Path(slurm2).is_file():
+        ep, tr, va = load_train_val_series(slurm2)
+        ax.plot(ep, tr, lw=1.4, ls="-", label="DDPM-2 train", color="#ff7f0e", alpha=0.85)
+        ax.plot(ep, va, lw=1.8, ls="--", label="DDPM-2 val", color="#994d00", alpha=0.95)
+        plotted = True
+    elif slurm2 is not None:
+        print(f"Warning: 2-param training series not found ({slurm2}); use --slurm-2param or restore bundled JSON.")
+
+    if not plotted:
+        print("No Slurm logs parsed — writing placeholder note instead of curves.")
+        ax.text(
+            0.5,
+            0.5,
+            "Pass --slurm-6param; DDPM-2 uses bundled JSON by default (--slurm-2param).",
+            ha="center",
+            va="center",
+            transform=ax.transAxes,
+        )
+    else:
+        ax.set_yscale("log")
+        ax.grid(True, alpha=0.3)
+        ax.set_xlabel("Epoch")
+        ax.set_ylabel("MSE diffusion loss")
+        ax.legend(loc="upper right", fontsize=8)
+        ax.set_title("Training / validation curves (combined)")
+    outp = out_dir / "comparison_training_train_val_overlay.png"
+    fig.savefig(outp, dpi=170, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", outp)
+
+
+def run_random_theta_triplets(
+    out_dir: Path,
+    imgs6: np.ndarray,
+    lab6: np.ndarray,
+    mean6: np.ndarray,
+    std6: np.ndarray,
+    mean2: np.ndarray,
+    std2: np.ndarray,
+    model2,
+    model6,
+    device: torch.device,
+    ddim_steps: int,
+    seed: int,
+    n_pairs: int,
+    batch_size: int,
+):
+    """Random LHS targets in CAMELS bbox; CAMELS column = NN real map."""
+    rng = np.random.default_rng(seed)
+    lo, hi = lab6.min(0), lab6.max(0)
+    targets = _latin_hypercube(min(n_pairs, 32), lo, hi, rng)[:n_pairs]
+
+    H, W = int(imgs6.shape[-2]), int(imgs6.shape[-1])
+    tg2 = targets[:, :2].astype(np.float32)
+
+    fig = plt.figure(figsize=(3.8 * max(3, n_pairs * 3), 4.1))
+    for i in range(n_pairs):
+        theta6 = targets[i].astype(np.float32)
+        theta2 = tg2[i]
+        dist = np.linalg.norm(lab6 - theta6[None, :], axis=1).astype(np.float64)
+        nn = int(np.argmin(dist))
+        nn_img = imgs6[nn]
+
+        gen2 = generate_maps(
+            model2, theta2[np.newaxis, :], mean2, std2, H, W, device, ddim_steps, batch_size
+        )
+        gen6 = generate_maps(
+            model6, theta6[np.newaxis, :], mean6, std6, H, W, device, ddim_steps, batch_size
+        )
+
+        titles = ("CAMELS (NN)", "DDPM-2", "DDPM-6")
+        for j, img in enumerate((nn_img, gen2[0], gen6[0])):
+            ax = fig.add_subplot(1, n_pairs * 3, i * 3 + j + 1)
+            ax.imshow(img, vmin=0, vmax=1, origin="lower", cmap="inferno")
+            ax.axis("off")
+            ax.set_title(titles[j], fontsize=8)
+
+    plt.suptitle(
+        "Random LHS cosmologies — CAMELS = nearest-neighbour truth | gens conditioned on LHS labels",
+        fontsize=10,
+        y=1.02,
+    )
+    p = out_dir / "comparison_random_lhs_triplets_camels_ddpm2_ddpm6.png"
+    plt.tight_layout()
+    plt.savefig(p, dpi=160, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", p)
+
+
+def run_conditioned_test_triplets(
+    out_dir: Path,
+    imgs6: np.ndarray,
+    lab6: np.ndarray,
+    mean6: np.ndarray,
+    std6: np.ndarray,
+    mean2: np.ndarray,
+    std2: np.ndarray,
+    model2,
+    model6,
+    device: torch.device,
+    ddim_steps: int,
+    seed: int,
+    n_pairs: int,
+    batch_size: int,
+):
+    """Same rows from test split: conditioned on truth labels."""
+    rng = np.random.default_rng(seed + 1)
+    idx = rng.choice(len(imgs6), size=min(n_pairs, len(imgs6)), replace=False)
+    H, W = int(imgs6.shape[-2]), int(imgs6.shape[-1])
+    fig, axes = plt.subplots(1, n_pairs * 3, figsize=(2.9 * n_pairs * 3, 3.8), squeeze=False)
+    for ii, ix in enumerate(idx):
+        tg6 = lab6[ix].astype(np.float32)
+        tg2 = tg6[:2]
+        rm = imgs6[ix]
+        g2 = generate_maps(model2, tg2[np.newaxis, :], mean2, std2, H, W, device, ddim_steps, batch_size)[0]
+        g6 = generate_maps(model6, tg6[np.newaxis, :], mean6, std6, H, W, device, ddim_steps, batch_size)[0]
+        for j, img in enumerate((rm, g2, g6)):
+            ax = axes[0, ii * 3 + j]
+            ax.imshow(img, vmin=0, vmax=1, origin="lower", cmap="inferno")
+            ax.axis("off")
+            if ii == 0:
+                ax.set_title(("CAMELS", "DDPM-2", "DDPM-6")[j], fontsize=8)
+        axes[0, ii * 3].set_xlabel(_fmt_title(tg6), fontsize=7)
+    plt.suptitle(f"Random test ix (conditioned on truth labels), n={len(idx)}", fontsize=10, y=1.06)
+    p = out_dir / "comparison_test_conditioned_camels_ddpm2_ddpm6.png"
+    plt.savefig(p, dpi=160, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", p)
+
+
+def pk_pdf_six_sets(
+    out_dir: Path,
+    name: str,
+    images_split: np.ndarray,
+    labels_split: np.ndarray,
+    label_mean: np.ndarray,
+    label_std: np.ndarray,
+    model,
+    device: torch.device,
+    ddim_steps: int,
+    batch_size: int,
+    n_per_set: int,
+):
+    """Six anchor rows (evenly spaced ix in test split), N_PER_SET DDIM samples."""
+    H, W = int(images_split.shape[-2]), int(images_split.shape[-1])
+    ldim = int(labels_split.shape[1])
+
+    idx = np.linspace(0, len(labels_split) - 1, num=6, dtype=int)
+    targets = labels_split[idx].copy()
+
+    box = 25.0
+    dk_ref = None
+    panels_pk = []
+
+    rng_pdf_bins = np.linspace(14.0, 22.0, 101)
+    bin_pdf = 0.5 * (rng_pdf_bins[:-1] + rng_pdf_bins[1:])
+
+    fig_pk, axes_pk = plt.subplots(2, 3, figsize=(14, 9), sharex=True, sharey=True)
+    axes_pk = axes_pk.ravel()
+
+    fig_pdf, axes_pdf = plt.subplots(6, 2, figsize=(12, 4.8 * 2), squeeze=False)
+
+    for si, target_l in enumerate(targets):
+        dist = np.linalg.norm(labels_split - target_l[None, :], axis=1).astype(np.float64)
+        ex = idx[si]
+        dist[ex] = np.inf if ex < len(dist) else np.inf
+        nn_idx = np.argsort(dist)[:n_per_set]
+        real_batch = images_split[nn_idx]
+        rep = np.tile(target_l[None, :], (n_per_set, 1))
+        gen = generate_maps(model, rep, label_mean, label_std, H, W, device, ddim_steps, batch_size)
+
+        dk_r, mr, sr = ec.calculate_power_spectrum_batch(real_batch, box_size=box)
+        dk_g, mg, sg = ec.calculate_power_spectrum_batch(gen, box_size=box)
+        dk_ref = dk_r
+        x = dk_ref[1:]
+        axpk = axes_pk[si]
+        axpk.plot(x, mr[1:], lw=2, label="CAMELS NN", color="#333")
+        axpk.fill_between(x, mr[1:] - sr[1:], mr[1:] + sr[1:], alpha=0.08, color="#333")
+        axpk.plot(x, mg[1:], lw=2, label=f"Generated (ldim={ldim})", color="#d95f02")
+        axpk.fill_between(x, mg[1:] - sg[1:], mg[1:] + sg[1:], alpha=0.08, color="#d95f02")
+        axpk.set_yscale("log")
+        axpk.grid(alpha=0.25)
+        axpk.set_title(_fmt_title(target_l), fontsize=8)
+        panels_pk.append((si, dk_r, mr, sr, mg, sg))
+
+        # PDF µ/σ
+        tb = []; rb = []
+        for i in range(n_per_set):
+            for arr, store in zip((real_batch, gen), (tb, rb)):
+                ims = np.clip(arr[i].ravel(), 0.0, 1.0)
+                logn = 14.0 + (22.0 - 14.0) * ims
+                hst, _ = np.histogram(logn, bins=rng_pdf_bins, density=True)
+                store.append(hst)
+        tb = np.asarray(tb); rb = np.asarray(rb)
+
+        axes_pdf[si, 0].plot(bin_pdf, tb.mean(axis=0), lw=2, label="CAMELS NN", color="#333")
+        axes_pdf[si, 0].plot(bin_pdf, rb.mean(axis=0), lw=2, label="Generated", color="#d95f02")
+        axes_pdf[si, 1].plot(bin_pdf, tb.std(axis=0), lw=2, ls="-", label="CAMELS σ", color="#333")
+        axes_pdf[si, 1].plot(bin_pdf, rb.std(axis=0), lw=2, ls="--", label="Gen σ", color="#d95f02")
+
+    axes_pk[0].legend(fontsize=7, loc="lower left")
+    fig_pk.suptitle(f"$P(k)$ mean±std — six anchors — {name}", fontsize=10)
+    fig_pk.tight_layout()
+    p_pk = out_dir / f"six_anchor_pk_{name}.png"
+    fig_pk.savefig(p_pk, dpi=160)
+    plt.close(fig_pk)
+
+    axes_pdf[-1, 0].set_xlabel(r"$\log N_{\mathrm{HI}}$")
+    axes_pdf[-1, 1].set_xlabel(r"$\log N_{\mathrm{HI}}$")
+    fig_pdf.suptitle(f"PDF mean & σ — six anchors × {n_per_set} — {name}")
+    fig_pdf.tight_layout()
+    p_pdf = out_dir / f"six_anchor_pdf_mu_sigma_{name}.png"
+    fig_pdf.savefig(p_pdf, dpi=160)
+    plt.close(fig_pdf)
+    print("Saved", p_pk)
+    print("Saved", p_pdf)
+
+
+def pk_six_triplet_combined(
+    out_dir: Path,
+    imgs6: np.ndarray,
+    lab6: np.ndarray,
+    mean6: np.ndarray,
+    std6: np.ndarray,
+    mean2: np.ndarray,
+    std2: np.ndarray,
+    model2: torch.nn.Module,
+    model6: torch.nn.Module,
+    device: torch.device,
+    ddim_steps: int,
+    batch_size: int,
+    n_per_set: int,
+) -> None:
+    """Six anchors — mean P(k) for CAMELS vs DDPM-2 vs DDPM-6; analogous PDF overlays."""
+    H, W = int(imgs6.shape[-2]), int(imgs6.shape[-1])
+    idx = np.linspace(0, len(lab6) - 1, num=6, dtype=int)
+    targets = lab6[idx].copy()
+    box = 25.0
+
+    fig_pk, axes_pk = plt.subplots(2, 3, figsize=(14, 9), sharex=True, sharey=True)
+    axes_pk = axes_pk.ravel()
+    rng_pdf_bins = np.linspace(14.0, 22.0, 101)
+    bin_pdf = 0.5 * (rng_pdf_bins[:-1] + rng_pdf_bins[1:])
+    fig_pdf, axes_pdf = plt.subplots(6, 2, figsize=(12, 10.5))
+
+    for si, target_l in enumerate(targets):
+        dist = np.linalg.norm(lab6 - target_l[None, :], axis=1).astype(np.float64)
+        ex = int(idx[si])
+        if ex < len(dist):
+            dist = dist.copy()
+            dist[ex] = np.inf
+        nn_idx = np.argsort(dist)[:n_per_set]
+        real_batch = imgs6[nn_idx]
+
+        tg2 = np.tile(target_l[:2][None, :], (n_per_set, 1)).astype(np.float32)
+        tg6 = np.tile(target_l[None, :], (n_per_set, 1)).astype(np.float32)
+
+        gen2 = generate_maps(model2, tg2, mean2, std2, H, W, device, ddim_steps, batch_size)
+        gen6 = generate_maps(model6, tg6, mean6, std6, H, W, device, ddim_steps, batch_size)
+
+        dk_r, mr, sr = ec.calculate_power_spectrum_batch(real_batch, box_size=box)
+        _, m2, s2 = ec.calculate_power_spectrum_batch(gen2, box_size=box)
+        _, mG, sG = ec.calculate_power_spectrum_batch(gen6, box_size=box)
+        x = dk_r[1:]
+
+        axpk = axes_pk[si]
+        axpk.plot(x, mr[1:], lw=2, label="CAMELS NN", color="#222")
+        axpk.fill_between(x, mr[1:] - sr[1:], mr[1:] + sr[1:], alpha=0.06, color="#222")
+        axpk.plot(x, m2[1:], lw=2, label="DDPM-2 μ", color="#ff7f0e")
+        axpk.fill_between(x, m2[1:] - s2[1:], m2[1:] + s2[1:], alpha=0.06, color="#ff7f0e")
+        axpk.plot(x, mG[1:], lw=2, label="DDPM-6 μ", color="#1f77b4")
+        axpk.fill_between(x, mG[1:] - sG[1:], mG[1:] + sG[1:], alpha=0.06, color="#1f77b4")
+        axpk.set_yscale("log")
+        axpk.grid(alpha=0.25)
+        axpk.set_title(_fmt_title(target_l), fontsize=8)
+        if si == 0:
+            axpk.legend(fontsize=6.2, loc="lower left")
+
+        pdf_rows_lists = []
+        for imgs in (real_batch, gen2, gen6):
+            hb = []
+            for i in range(min(n_per_set, len(imgs))):
+                px = np.clip(imgs[i].ravel(), 0.0, 1.0)
+                ln = 14.0 + (22.0 - 14.0) * px
+                hb.append(np.histogram(ln, bins=rng_pdf_bins, density=True)[0])
+            pdf_rows_lists.append(np.asarray(hb))
+
+        cam_pdf, d2_pdf, d6_pdf = pdf_rows_lists
+        axes_pdf[si, 0].plot(bin_pdf, cam_pdf.mean(axis=0), lw=2, color="#222", label="CAMELS μ")
+        axes_pdf[si, 0].plot(bin_pdf, d2_pdf.mean(axis=0), lw=2, color="#ff7f0e", label="DDPM-2 μ")
+        axes_pdf[si, 0].plot(bin_pdf, d6_pdf.mean(axis=0), lw=2, color="#1f77b4", label="DDPM-6 μ")
+        axes_pdf[si, 1].plot(bin_pdf, cam_pdf.std(axis=0), lw=2, color="#222")
+        axes_pdf[si, 1].plot(bin_pdf, d2_pdf.std(axis=0), lw=2, ls="--", color="#ff7f0e")
+        axes_pdf[si, 1].plot(bin_pdf, d6_pdf.std(axis=0), lw=2, ls="--", color="#1f77b4")
+
+    fig_pk.suptitle("$P(k)$ CAMELS vs DDPM-2 vs DDPM-6 — six Ωm–σ8 anchors", fontsize=11)
+    fig_pk.tight_layout()
+    p_pk = out_dir / "six_anchor_pk_overlay_camels_ddpm2_ddpm6.png"
+    fig_pk.savefig(p_pk, dpi=160)
+    plt.close(fig_pk)
+
+    axes_pdf[-1, 0].set_xlabel(r"$\log N_{\mathrm{HI}}$")
+    axes_pdf[-1, 1].set_xlabel(r"$\log N_{\mathrm{HI}}$")
+    fig_pdf.suptitle(r"PDF mean ($\mu$) and std ($\sigma$) overlays", fontsize=10)
+    fig_pdf.tight_layout()
+    p_pdf = out_dir / "six_anchor_pdf_overlay_camels_ddpm2_ddpm6.png"
+    fig_pdf.savefig(p_pdf, dpi=160)
+    plt.close(fig_pdf)
+    print("Saved", p_pk)
+    print("Saved", p_pdf)
+
+
+def mlp_recovery_dual(
+    out_dir: Path,
+    data_train: Path,
+    imgs_te: np.ndarray,
+    lab_te: np.ndarray,
+    mean: np.ndarray,
+    std: np.ndarray,
+    model_ddpm: torch.nn.Module,
+    tag: str,
+    device: torch.device,
+    ddim_steps: int,
+    seed: int,
+) -> None:
+    from sklearn.metrics import mean_squared_error
+    from sklearn.neural_network import MLPRegressor
+
+    ldim = lab_te.shape[1]
+    Npix = imgs_te.shape[-1]
+    dl = 25.0 / Npix
+
+    def pk_row(im):
+        _dk, pk = ec.PowerSpectrum(np.asarray(im, dtype=np.float64), N=Npix, dl=dl)
+        return pk[1:].astype(np.float32)
+
+    img_tr_np, lab_tr_np = ec.load_split(data_train, "train")
+    if len(img_tr_np) > 2000:
+        rng = np.random.default_rng(seed)
+        jj = rng.choice(len(img_tr_np), 2000, replace=False)
+        img_tr_np, lab_tr_np = img_tr_np[jj], lab_tr_np[jj]
+    X_train = np.stack([pk_row(img_tr_np[i]) for i in range(len(img_tr_np))], axis=0)
+    y_train = lab_tr_np.astype(np.float32)
+    mlp = MLPRegressor(
+        hidden_layer_sizes=(64, 64),
+        alpha=1e-4,
+        random_state=seed,
+        max_iter=250,
+        early_stopping=True,
+        validation_fraction=0.1,
+    )
+    mlp.fit(X_train, y_train)
+
+    n_ev = min(40, len(imgs_te))
+    eval_idx = np.arange(n_ev)
+    X_real = np.stack([pk_row(imgs_te[i]) for i in eval_idx], axis=0)
+    y_true = lab_te[eval_idx]
+
+    preds_real = mlp.predict(X_real)
+
+    gens = []
+    H, W = int(imgs_te.shape[-2]), int(imgs_te.shape[-1])
+    for i0 in range(0, n_ev, 8):
+        bs_chunk = min(8, n_ev - i0)
+        lbl = y_true[i0 : i0 + bs_chunk]
+        g = generate_maps(model_ddpm, lbl, mean, std, H, W, device, ddim_steps, bs_chunk)
+        gens.extend([pk_row(g[j]) for j in range(len(g))])
+    X_gen = np.stack(gens, axis=0)
+    preds_gen = mlp.predict(X_gen)
+
+    rmse_real = np.sqrt(mean_squared_error(y_true, preds_real, multioutput="raw_values"))
+    rmse_gen = np.sqrt(mean_squared_error(y_true, preds_gen, multioutput="raw_values"))
+
+    fig, axes = plt.subplots(2, ldim, figsize=(max(9.0, 2.8 * max(ldim, 2)), 4.9), squeeze=False)
+    if ldim == 1:
+        axes = np.reshape(axes, (2, 1))
+    for k in range(ldim):
+        for row, preds, rmv, ylab in (
+            (0, preds_real, rmse_real, "CAMELS P(k) predictions"),
+            (1, preds_gen, rmse_gen, f"{tag}: generated P(k)"),
+        ):
+            ax = axes[row, k]
+            lo = float(y_true[:, k].min()); hi = float(y_true[:, k].max())
+            pad = 0.03 * (hi - lo + 1e-12)
+            ax.scatter(y_true[:, k], preds[:, k], s=14, alpha=0.72, edgecolors="none", c="#333")
+            ax.plot([lo - pad, hi + pad], [lo - pad, hi + pad], color="crimson", lw=1.0)
+            ax.grid(True, alpha=0.28)
+            ax.set_title(f"dim {k} RMSE={float(rmv[k]):.4f}", fontsize=8)
+            if k == 0:
+                ax.set_ylabel(ylab, fontsize=8)
+
+    plt.suptitle(
+        "MLP: train on CAMELS train P(k), test on CAMELS vs DDPM-drawn spectra",
+        fontsize=10,
+        y=1.02,
+    )
+    plt.tight_layout()
+    p = out_dir / f"mlp_pk_parameter_recovery_{tag}.png"
+    plt.savefig(p, dpi=165, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", p)
+
+
+def posterior_one_index(
+    out_dir: Path,
+    images_split: np.ndarray,
+    labels_split: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    model,
+    cfg: Dict,
+    device,
+    ix: int,
+    tag: str,
+    ddim_steps: int,
+    grid: int,
+    batch_sz: int,
+):
+    normalize = bool(cfg.get("normalize_labels", True))
+    lab_dim = labels_split.shape[1]
+    H, W = int(images_split.shape[-2]), int(images_split.shape[-1])
+    obs = images_split[ix]
+    label_anchor_full = labels_split[ix].astype(np.float32)
+
+    lo0 = float(labels_split[:, 0].min())
+    hi0 = float(labels_split[:, 0].max())
+    lo1 = float(labels_split[:, 1].min())
+    hi1 = float(labels_split[:, 1].max())
+    pad0 = 0.02 * (hi0 - lo0 + 1e-12)
+    pad1 = 0.02 * (hi1 - lo1 + 1e-12)
+    om_ax, s8_ax, OG, SG, grid2 = build_cosmo_grid(
+        grid, lo0 - pad0, hi0 + pad0, lo1 - pad1, hi1 + pad1
+    )
+    g = grid
+    ngrid = grid2.shape[0]
+    npix = int(obs.shape[-1])
+    dl = 25.0 / npix
+    dk, _ = ec.PowerSpectrum(em.images01_to_log_nhi(obs), N=npix, dl=dl)
+    valid = dk > 0
+    log_pd = log_pk_observed(obs, 25.0, dk)
+
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    full = np.tile(label_anchor_full[np.newaxis, :], (ngrid, 1))
+    full[:, 0] = grid2[:, 0].astype(np.float32)
+    full[:, 1] = grid2[:, 1].astype(np.float32)
+
+    def weights_full() -> np.ndarray:
+        scores = []
+        for j0 in range(0, ngrid, batch_sz):
+            chunk = full[j0 : j0 + batch_sz]
+            imgs = em.sample_batch(
+                model,
+                chunk,
+                lab_mean,
+                lab_std,
+                normalize,
+                H,
+                W,
+                device,
+                ddim_steps,
+                False,
+            )
+            _, pkc = em.per_map_power_spectra_log(imgs, 25.0)
+            log_pg = np.log(pkc[:, valid] + 1e-30)
+            mse = np.mean((log_pd[np.newaxis, :] - log_pg) ** 2, axis=1)
+            scores.append(-mse / (2.0 * 0.25**2))
+        sc = np.concatenate(scores)
+        sc -= sc.max()
+        w = np.exp(sc).reshape(g, g)
+        w /= w.sum()
+        return w
+
+    Wmap = weights_full()
+    tom, ts8 = float(label_anchor_full[0]), float(label_anchor_full[1])
+    mom = float((Wmap * OM).sum())
+    ms8 = float((Wmap * S8).sum())
+
+    fig, ax = plt.subplots(figsize=(5.2, 4.6))
+    cf = ax.contourf(OM, S8, Wmap, levels=12, cmap="Blues")
+    plt.colorbar(cf, ax=ax, fraction=0.046, pad=0.04)
+    ax.scatter(tom, ts8, s=55, c="r", marker="x", zorder=6, label="true")
+    ax.scatter(mom, ms8, s=60, c="k", marker="+", zorder=6, label="post. mean")
+    ax.set_xlabel(r"$\Omega_m$")
+    ax.set_ylabel(r"$\sigma_8$")
+    ax.legend(fontsize=8)
+    ax.set_title(f"Surrogate posterior (test ix={ix}, ldim={lab_dim})", fontsize=10)
+    p = out_dir / f"posterior_surrogate_test_ix_{ix}_{tag}.png"
+    fig.savefig(p, dpi=160, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", p)
+
+
+def main(argv: Sequence[str] | None = None) -> None:
+    p = argparse.ArgumentParser(description="DDPM-2 vs DDPM-6 comparison suite.")
+    p.add_argument(
+        "--output-dir",
+        type=Path,
+        default=MODELS_ROOT / "ddpm_comparison_out",
+    )
+    p.add_argument("--data-2param", type=Path, default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2"))
+    p.add_argument("--data-6param", type=Path, default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"))
+    p.add_argument(
+        "--bundle-2param",
+        type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "2param_epoch200",
+    )
+    p.add_argument(
+        "--bundle-6param",
+        type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "6param_best",
+    )
+    p.add_argument("--posterior-index", type=int, default=56)
+    p.add_argument("--lhs-n", type=int, default=50)
+    p.add_argument("--six-n-per-anchor", type=int, default=15)
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument("--slurm-6param", type=Path, default=DEFAULT_SLURM_6)
+    p.add_argument(
+        "--slurm-2param",
+        type=Path,
+        default=DEFAULT_DDPM2_TRAINING,
+        help="DDPM-2 train/val series: Slurm .out (parsed) or bundled ddpm_2param_training_loss.json.",
+    )
+    p.add_argument("--skip-lhs-r2", action="store_true", help="LHS R² plots are expensive; skip if set.")
+    p.add_argument("--n-random-triplets", type=int, default=4)
+    args = p.parse_args(list(argv) if argv is not None else None)
+
+    out_dir = Path(args.output_dir).resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("device:", device)
+
+    ck2 = args.bundle_2param / "checkpoint_epoch_200.pt"
+    args2 = args.bundle_2param / "args.json"
+    ck6 = args.bundle_6param / "best_model.pt"
+    args6 = args.bundle_6param / "args.json"
+
+    data2 = Path(args.data_2param)
+    data6 = Path(args.data_6param)
+
+    imgs6, lab6 = ec.load_split(data6, "test")
+    mean6, std6 = ec.load_label_stats(data6)
+    mean2, std2 = ec.load_label_stats(data2)
+
+    plot_training_overlay(out_dir, args.slurm_6param, args.slurm_2param)
+
+    imgs2, lab2 = ec.load_split(data2, "test")
+
+    print(">>> Loading DDPM-2...")
+    model2, cfg2 = load_model(args2, ck2, device)
+    print(">>> Loading DDPM-6...")
+    model6, cfg6 = load_model(args6, ck6, device)
+
+    print(">>> Random LHS + conditioned triplets...")
+    try:
+        run_random_theta_triplets(
+            out_dir,
+            imgs6,
+            lab6,
+            mean6,
+            std6,
+            mean2,
+            std2,
+            model2,
+            model6,
+            device=device,
+            ddim_steps=args.ddim_steps,
+            seed=args.seed,
+            n_pairs=args.n_random_triplets,
+            batch_size=args.batch_size,
+        )
+        run_conditioned_test_triplets(
+            out_dir,
+            imgs6,
+            lab6,
+            mean6,
+            std6,
+            mean2,
+            std2,
+            model2,
+            model6,
+            device=device,
+            ddim_steps=args.ddim_steps,
+            seed=args.seed,
+            n_pairs=args.n_random_triplets,
+            batch_size=args.batch_size,
+        )
+    except Exception as exc:
+        print("Triplet grids failed:", exc)
+
+    print(">>> Six-anchor overlays (combined + per-model)...")
+    try:
+        pk_six_triplet_combined(
+            out_dir,
+            imgs6,
+            lab6,
+            mean6,
+            std6,
+            mean2,
+            std2,
+            model2,
+            model6,
+            device=device,
+            ddim_steps=args.ddim_steps,
+            batch_size=args.batch_size,
+            n_per_set=args.six_n_per_anchor,
+        )
+        pk_pdf_six_sets(
+            out_dir,
+            "ddpm6_only",
+            imgs6,
+            lab6,
+            mean6,
+            std6,
+            model6,
+            device,
+            args.ddim_steps,
+            args.batch_size,
+            args.six_n_per_anchor,
+        )
+        pk_pdf_six_sets(
+            out_dir,
+            "ddpm2_only",
+            imgs2,
+            lab2,
+            mean2,
+            std2,
+            model2,
+            device,
+            args.ddim_steps,
+            args.batch_size,
+            args.six_n_per_anchor,
+        )
+    except Exception as exc:
+        print("P(k)/PDF six-anchor plots failed:", exc)
+
+    if not args.skip_lhs_r2:
+        print(">>> LHS R² (LHS-50 × 15 DDIM each — long)...")
+        try:
+            for label, imgs, labs, mn, sd, mdl in (
+                ("ddpm2_lhs50", imgs2, lab2, mean2, std2, model2),
+                ("ddpm6_lhs50", imgs6, lab6, mean6, std6, model6),
+            ):
+                lhs_pts, r2_mu, r2_sig, lo_b, hi_b = compute_lhs_r2(
+                    mdl,
+                    imgs,
+                    labs,
+                    mn,
+                    sd,
+                    device,
+                    args.lhs_n,
+                    15,
+                    args.batch_size,
+                    25.0,
+                    args.ddim_steps,
+                    args.seed,
+                )
+                outp = out_dir / f"r2_cosmology_lhs{args.lhs_n}_{label}.png"
+                plot_r2_cosmology_figure(lhs_pts, r2_mu, r2_sig, lo_b, hi_b, outp, dpi=160)
+                print("Saved", outp)
+                np.savez(
+                    out_dir / f"r2_lhs_data_{label}.npz",
+                    lhs_pts=lhs_pts,
+                    r2_mu_arr=r2_mu,
+                    r2_sig_arr=r2_sig,
+                    lo_b=lo_b,
+                    hi_b=hi_b,
+                )
+        except Exception as exc:
+            print("LHS R² skipped:", exc)
+    else:
+        print("(Skipping LHS R².)")
+
+    print(">>> MLP P(k) parameter recovery...")
+    try:
+        mlp_recovery_dual(
+            out_dir, data2, imgs2[:40], lab2[:40], mean2, std2, model2, "ddpm2param", device, args.ddim_steps, args.seed
+        )
+        mlp_recovery_dual(
+            out_dir, data6, imgs6[:40], lab6[:40], mean6, std6, model6, "ddpm6param", device, args.ddim_steps, args.seed
+        )
+    except Exception as exc:
+        print("MLP recovery skipped:", exc)
+
+    print(f">>> Surrogate posteriors (test index {args.posterior_index})...")
+    try:
+        ix = int(args.posterior_index)
+        posterior_one_index(
+            out_dir, imgs6, lab6, mean6, std6, model6, cfg6, device, ix, "ddpm6", args.ddim_steps, 14, args.batch_size
+        )
+        posterior_one_index(
+            out_dir,
+            imgs2,
+            lab2,
+            mean2,
+            std2,
+            model2,
+            cfg2,
+            device,
+            ix,
+            "ddpm2",
+            args.ddim_steps,
+            14,
+            args.batch_size,
+        )
+    except Exception as exc:
+        print("Posterior panels skipped:", exc)
+
+    del model2, model6
+    free_torch()
+    print(f"Done. Outputs in {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

cross_model/scripts/compare_ddpm_training_curves.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Parse DDPM Slurm stdout or bundled JSON for Train/Val loss series."""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+from typing import Tuple
+
+_ROW = re.compile(
+    r"Epoch\s+(?P<ep>\d+)/\d+\s+\|\s+Train:\s+(?P<tr>[\d.eE+-]+)\s+\|\s+Val:\s+(?P<va>[\d.eE+-]+)",
+)
+
+
+def parse_slurm_training_log(path: str | Path) -> Tuple[list[int], list[float], list[float]]:
+    """Return (epochs, train_losses, val_losses) parsed from Slurm *.out stdout."""
+    p = Path(path)
+    text = p.read_text(encoding="utf-8", errors="replace")
+    epochs, trains, vals = [], [], []
+    for m in _ROW.finditer(text):
+        epochs.append(int(m.group("ep")))
+        trains.append(float(m.group("tr")))
+        vals.append(float(m.group("va")))
+    return epochs, trains, vals
+
+
+def load_training_loss_json(path: str | Path) -> Tuple[list[int], list[float], list[float]]:
+    """Return (epochs, train_losses, val_losses) from a JSON export (keys: epochs, train, val)."""
+    p = Path(path)
+    raw = json.loads(p.read_text(encoding="utf-8"))
+    epochs = [int(e) for e in raw["epochs"]]
+    trains = [float(x) for x in raw["train"]]
+    vals = [float(x) for x in raw["val"]]
+    if not (len(epochs) == len(trains) == len(vals)):
+        raise ValueError(f"{p}: mismatched lengths in epochs/train/val")
+    return epochs, trains, vals
+
+
+def load_train_val_series(path: str | Path) -> Tuple[list[int], list[float], list[float]]:
+    """Slurm *.out or *.json with the same semantic output as ``parse_slurm_training_log``."""
+    p = Path(path)
+    if p.suffix.lower() == ".json":
+        return load_training_loss_json(p)
+    return parse_slurm_training_log(p)

cross_model/scripts/ddpm_figure6_integration.py

@@ -0,0 +1,271 @@
+"""
+Figure 6 style (arXiv:2409.09101) helpers for DDPM surrogate posteriors — use with ddpm_posterior_six_anchors / run_ddpm_figure6_suite.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import gridspec
+
+from figure6_2409_style import (
+    create_comparison_marginal_vs_profile,
+    create_figure6_style_plot,
+)
+from sigma_contour_utils import compute_sigma_levels
+
+
+def integrate_figure6_with_ddpm2(
+    Wmap: np.ndarray,
+    om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    true_om: float,
+    true_s8: float,
+    test_index: int,
+    output_dir: Path,
+    model_name: str = "DDPM-2",
+) -> None:
+    """Single-map Figure 6 style: marginal and profile PNGs."""
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    nm = model_name.replace(" ", "-").lower()
+
+    fig_marginal = create_figure6_style_plot(
+        Wmap,
+        om_grid,
+        s8_grid,
+        true_param1=true_om,
+        true_param2=true_s8,
+        param1_label=r"$\Omega_m$",
+        param2_label=r"$\sigma_8$",
+        title=f"{model_name} — Test ix={test_index} (Marginal)",
+        show_profile=False,
+        figsize=(10, 10),
+    )
+
+    save_path_marginal = output_dir / f"fig6_style_{nm}_ix{test_index}_marginal.png"
+    fig_marginal.savefig(save_path_marginal, dpi=200, bbox_inches="tight")
+    plt.close(fig_marginal)
+    print(f"Saved: {save_path_marginal}")
+
+    fig_profile = create_figure6_style_plot(
+        Wmap,
+        om_grid,
+        s8_grid,
+        true_param1=true_om,
+        true_param2=true_s8,
+        param1_label=r"$\Omega_m$",
+        param2_label=r"$\sigma_8$",
+        title=f"{model_name} — Test ix={test_index} (Profile)",
+        show_profile=True,
+        figsize=(10, 10),
+    )
+
+    save_path_profile = output_dir / f"fig6_style_{nm}_ix{test_index}_profile.png"
+    fig_profile.savefig(save_path_profile, dpi=200, bbox_inches="tight")
+    plt.close(fig_profile)
+    print(f"Saved: {save_path_profile}")
+
+    fig_cmp = create_comparison_marginal_vs_profile(
+        Wmap,
+        om_grid,
+        s8_grid,
+        true_param1=true_om,
+        true_param2=true_s8,
+        title=f"{model_name} marginal vs profile — ix={test_index}",
+        figsize=(11, 4.2),
+    )
+    cmp_path = output_dir / f"fig6_marg_vs_prof_{nm}_ix{test_index}.png"
+    fig_cmp.savefig(cmp_path, dpi=185, bbox_inches="tight")
+    plt.close(fig_cmp)
+    print(f"Saved: {cmp_path}")
+
+
+def integrate_figure6_with_multi_anchor(
+    posteriors_list: list[np.ndarray],
+    om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    true_values_list: list[tuple[float, float]],
+    test_indices: list[int],
+    output_dir: Path,
+    model_name: str = "DDPM-2",
+) -> None:
+    """2×3 grid with Figure–6-ish 2D + top marginal."""
+    output_dir = Path(output_dir)
+    nm = model_name.replace(" ", "-").lower()
+    fig = plt.figure(figsize=(20, 14))
+    gs_o = gridspec.GridSpec(2, 3, figure=fig, hspace=0.33, wspace=0.32)
+
+    for idx, (posterior, true_vals, test_ix) in enumerate(
+        zip(posteriors_list, true_values_list, test_indices)
+    ):
+        true_om, true_s8 = true_vals
+        row, col = divmod(idx, 3)
+
+        posterior_norm = posterior / posterior.sum()
+        sigma_levels = compute_sigma_levels(posterior_norm, [0.683, 0.954])
+        P1, P2 = np.meshgrid(om_grid, s8_grid, indexing="ij")
+
+        gs_sub = gridspec.GridSpecFromSubplotSpec(
+            2,
+            2,
+            subplot_spec=gs_o[row, col],
+            width_ratios=[4, 1],
+            height_ratios=[1, 4],
+            hspace=0.06,
+            wspace=0.06,
+        )
+        ax_main = fig.add_subplot(gs_sub[1, 0])
+        ax_top = fig.add_subplot(gs_sub[0, 0], sharex=ax_main)
+
+        ax_main.contourf(P1, P2, posterior_norm, levels=20, cmap="Blues", alpha=0.85)
+        if len(set(sigma_levels)) >= 1:
+            ax_main.contour(
+                P1,
+                P2,
+                posterior_norm,
+                levels=sigma_levels,
+                colors=["darkblue", "steelblue"],
+                linewidths=[2.0, 1.5],
+            )
+
+        ax_main.scatter(true_om, true_s8, s=100, c="red", marker="x", linewidths=2.5, zorder=10)
+        ax_main.set_xlim(om_grid[0], om_grid[-1])
+        ax_main.set_ylim(s8_grid[0], s8_grid[-1])
+        ax_main.set_xlabel(r"$\Omega_m$" if row == 1 else "", fontsize=11)
+        ax_main.set_ylabel(r"$\sigma_8$" if col == 0 else "", fontsize=11)
+        ax_main.set_title(f"Test ix={test_ix}", fontsize=11, pad=5)
+        ax_main.grid(True, alpha=0.2)
+
+        marginal_om = posterior_norm.sum(axis=1)
+        marginal_om /= marginal_om.sum() + 1e-30
+        ax_top.fill_between(
+            om_grid,
+            0.0,
+            marginal_om,
+            alpha=0.6,
+            color="steelblue",
+            edgecolor="steelblue",
+        )
+        ax_top.axvline(true_om, color="red", linestyle="--", linewidth=2)
+        ax_top.set_xlim(om_grid[0], om_grid[-1])
+        ax_top.set_ylim(0, marginal_om.max() * 1.1)
+        ax_top.tick_params(labelbottom=False, labelsize=9)
+        ax_top.set_ylabel("$P(\\Omega_m)$", fontsize=9)
+        ax_top.grid(True, alpha=0.2)
+
+        ax_side = fig.add_subplot(gs_sub[1, 1], sharey=ax_main)
+        marginal_s8 = posterior_norm.sum(axis=0)
+        marginal_s8 /= marginal_s8.sum() + 1e-30
+        ax_side.fill_betweenx(s8_grid, 0.0, marginal_s8, alpha=0.6, color="steelblue", edgecolor="steelblue")
+        ax_side.axhline(true_s8, color="red", linestyle="--", linewidth=2)
+        ax_side.set_ylim(s8_grid[0], s8_grid[-1])
+        ax_side.set_xlim(0, marginal_s8.max() * 1.15)
+        ax_side.tick_params(labelleft=False)
+
+    fig.suptitle(
+        f"{model_name} — Figure 6 Style: Six Test Anchors",
+        fontsize=15,
+        y=0.995,
+        fontweight="bold",
+    )
+
+    save_path = output_dir / f"fig6_style_{nm}_all_anchors.png"
+    fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved multi-anchor grid: {save_path}")
+
+
+def integrate_figure6_model_comparison(
+    posteriors_dict: dict[str, np.ndarray],
+    om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    true_om: float,
+    true_s8: float,
+    test_index: int,
+    output_dir: Path,
+) -> None:
+    """Side-by-side model comparison panels."""
+    output_dir = Path(output_dir)
+    n_models = len(posteriors_dict)
+    fig = plt.figure(figsize=(8 * max(1, min(n_models, 4)), 8))
+    gs_outer = gridspec.GridSpec(1, n_models, figure=fig, wspace=0.32)
+
+    for idx, (model_name, posterior) in enumerate(posteriors_dict.items()):
+        gs_sub = gridspec.GridSpecFromSubplotSpec(
+            2,
+            2,
+            subplot_spec=gs_outer[0, idx],
+            width_ratios=[4, 1],
+            height_ratios=[1, 4],
+            hspace=0.06,
+            wspace=0.06,
+        )
+        posterior_norm = posterior / posterior.sum()
+        sigma_levels = compute_sigma_levels(posterior_norm, [0.683, 0.954])
+        P1, P2 = np.meshgrid(om_grid, s8_grid, indexing="ij")
+
+        ax_main = fig.add_subplot(gs_sub[1, 0])
+        ax_top = fig.add_subplot(gs_sub[0, 0], sharex=ax_main)
+
+        ax_main.contourf(P1, P2, posterior_norm, levels=20, cmap="Blues", alpha=0.85)
+        if len(set(sigma_levels)) >= 1:
+            ax_main.contour(
+                P1,
+                P2,
+                posterior_norm,
+                levels=sigma_levels,
+                colors=["darkblue", "steelblue"],
+                linewidths=[2.5, 2.0],
+            )
+
+        ax_main.scatter(true_om, true_s8, s=120, c="red", marker="x", linewidths=3, zorder=10)
+        ax_main.set_xlabel(r"$\Omega_m$", fontsize=13)
+        ax_main.set_ylabel(r"$\sigma_8$", fontsize=13)
+        ax_main.set_title(model_name, fontsize=13, pad=10, fontweight="bold")
+        ax_main.grid(True, alpha=0.3)
+        ax_main.set_xlim(om_grid[0], om_grid[-1])
+        ax_main.set_ylim(s8_grid[0], s8_grid[-1])
+
+        marginal = posterior_norm.sum(axis=1)
+        marginal /= marginal.sum() + 1e-30
+        ax_top.fill_between(om_grid, 0.0, marginal, alpha=0.6, color="steelblue")
+        ax_top.axvline(true_om, color="red", linestyle="--", linewidth=2.5)
+        ax_top.set_xlim(om_grid[0], om_grid[-1])
+        ax_top.set_ylim(0, marginal.max() * 1.12)
+        ax_top.tick_params(labelbottom=False)
+        ax_top.grid(True, alpha=0.25)
+
+        ax_sb = fig.add_subplot(gs_sub[1, 1], sharey=ax_main)
+        marginal_s = posterior_norm.sum(axis=0)
+        marginal_s /= marginal_s.sum() + 1e-30
+        ax_sb.fill_betweenx(s8_grid, 0.0, marginal_s, alpha=0.6, color="steelblue")
+        ax_sb.axhline(true_s8, color="red", linestyle="--", linewidth=2.5)
+        ax_sb.set_ylim(s8_grid[0], s8_grid[-1])
+
+    fig.suptitle(
+        f"Model Comparison (Figure 6 Style) — Test ix={test_index}",
+        fontsize=15,
+        y=0.995,
+        fontweight="bold",
+    )
+
+    save_path = output_dir / f"fig6_style_model_comparison_ix{test_index}.png"
+    fig.savefig(save_path, dpi=200, bbox_inches="tight")
+    plt.close(fig)
+    print(f"Saved model comparison: {save_path}")
+
+
+def print_integration_guide() -> None:
+    example_integration = """
+# Add imports next to posterior code:
+from figure6_2409_style import create_figure6_style_plot
+from ddpm_figure6_integration import (
+    integrate_figure6_with_ddpm2,
+    integrate_figure6_with_multi_anchor,
+    integrate_figure6_model_comparison,
+)
+"""
+    print(example_integration.strip())

cross_model/scripts/ddpm_posterior_six_anchors.py

@@ -0,0 +1,451 @@
+#!/usr/bin/env python3
+"""
+Surrogate P(k) likelihood posteriors on ($\\Omega_m$, $\\sigma_8$) for six test anchors.
+
+For each model:
+  • DDPM-2 — standard 2D marginal: sweep ($\\Omega_m$, $\\sigma_8$) while only two labels exist.
+  • DDPM-6 — same 2D sweep, but astrophysical / extra dimensions 2–5 are fixed in two cases:
+        - **extra_lower**: each of dims 2–5 fixed to the LHS **minimum** (from training labels)
+        - **extra_upper**: each fixed to the LHS **maximum**
+
+The observed HI map is always the CAMELS test MAP at that anchor index.
+
+This does not import ``compare_ddpm_models.py``; it only shares the same conventions and paths.
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+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
+
+MODELS_ROOT = Path(__file__).resolve().parents[1]
+CODE_6 = MODELS_ROOT / "6param_ddpm_hi_lh6"
+if str(CODE_6.resolve()) not in sys.path:
+    sys.path.insert(0, str(CODE_6))
+
+import evaluate_conditional as ec  # noqa: E402
+import eval_model as em  # noqa: E402
+from figure9_posterior import build_cosmo_grid, log_pk_observed  # noqa: E402
+
+
+def _fmt_title(lab: np.ndarray) -> str:
+    t = np.asarray(lab, dtype=float).ravel()
+    if t.size <= 2:
+        return rf"$\Omega_m$={t[0]:.3f}, $\sigma_8$={t[1]:.3f}"
+    tail = ", ".join(f"{float(v):.3g}" for v in t[2:])
+    return rf"$\Omega_m$={t[0]:.3f}, $\sigma_8$={t[1]:.3f} | " + tail
+
+
+def _train_label_path(data_dir: Path) -> Path:
+    for name in ("train_labels_LH.npy", "train_labels_LH_2.npy"):
+        p = data_dir / name
+        if p.is_file():
+            return p
+    raise FileNotFoundError(f"No train_labels_LH*.npy under {data_dir}")
+
+
+def tail_lhs_bounds(data_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
+    """Min/max over training LHS for label dimensions 2 … 5 (indices 2–5)."""
+    L = np.load(_train_label_path(data_dir))
+    if L.shape[1] < 6:
+        raise ValueError(f"Expected ≥6 label columns, got {L.shape}")
+    lo = L[:, 2:6].min(axis=0).astype(np.float32)
+    hi = L[:, 2:6].max(axis=0).astype(np.float32)
+    return lo, hi
+
+
+def posterior_weights(
+    obs: np.ndarray,
+    full: np.ndarray,
+    om_ax: np.ndarray,
+    s8_ax: np.ndarray,
+    lab_mean: np.ndarray,
+    lab_std: np.ndarray,
+    normalize: bool,
+    model: torch.nn.Module,
+    *,
+    H: int,
+    W: int,
+    device: torch.device,
+    grid: int,
+    batch_sz: int,
+    ddim_steps: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Returns (Wmap, OM, S8) with Wmap shaped (grid, grid); OM, S8 meshgrids (indexing='ij').
+    full: (ngrid, label_dim) rows on the (Ωm, σ8) grid plus any fixed tail dims.
+    """
+    ngrid = full.shape[0]
+    g = int(round(np.sqrt(ngrid)))
+    if g * g != ngrid:
+        raise ValueError(f"Expected square grid, got ngrid={ngrid}")
+
+    npix = int(obs.shape[-1])
+    dl = 25.0 / npix
+    dk, _ = ec.PowerSpectrum(em.images01_to_log_nhi(obs), N=npix, dl=dl)
+    valid = dk > 0
+    log_pd = log_pk_observed(obs, 25.0, dk)
+
+    def weights_full() -> np.ndarray:
+        scores = []
+        for j0 in range(0, ngrid, batch_sz):
+            chunk = full[j0 : j0 + batch_sz]
+            imgs = em.sample_batch(
+                model,
+                chunk,
+                lab_mean,
+                lab_std,
+                normalize,
+                H,
+                W,
+                device,
+                ddim_steps,
+                False,
+            )
+            _, pkc = em.per_map_power_spectra_log(imgs, 25.0)
+            log_pg = np.log(pkc[:, valid] + 1e-30)
+            mse = np.mean((log_pd[np.newaxis, :] - log_pg) ** 2, axis=1)
+            scores.append(-mse / (2.0 * 0.25**2))
+        sc = np.concatenate(scores)
+        sc -= sc.max()
+        w = np.exp(sc).reshape(g, g)
+        w /= w.sum()
+        return w
+
+    Wmap = weights_full()
+    OM, S8 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+    return Wmap, OM, S8
+
+
+def build_full_grid_2d(
+    labels_split: np.ndarray,
+    grid: int,
+    tail: np.ndarray | None,
+    lab_dim: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    If tail is None (2-param model path): full has shape (grid^2, 2).
+    Else tail shape (4,): dims 2–5 filled with constants; dims 0,1 swept.
+    """
+    lo0 = float(labels_split[:, 0].min())
+    hi0 = float(labels_split[:, 0].max())
+    lo1 = float(labels_split[:, 1].min())
+    hi1 = float(labels_split[:, 1].max())
+    pad0 = 0.02 * (hi0 - lo0 + 1e-12)
+    pad1 = 0.02 * (hi1 - lo1 + 1e-12)
+    om_ax, s8_ax, OG, SG, grid2 = build_cosmo_grid(
+        grid, lo0 - pad0, hi0 + pad0, lo1 - pad1, hi1 + pad1
+    )
+    ngrid = grid2.shape[0]
+    if lab_dim == 2 and tail is not None:
+        raise ValueError("lab_dim==2 implies no extra tail.")
+    out = np.zeros((ngrid, lab_dim), dtype=np.float32)
+    out[:, 0] = grid2[:, 0].astype(np.float32)
+    out[:, 1] = grid2[:, 1].astype(np.float32)
+    if tail is not None:
+        assert tail.shape == (4,)
+        out[:, 2:6] = tail[np.newaxis, :]
+    return out.astype(np.float32), om_ax, s8_ax
+
+
+def plot_posterior_panel(
+    ax,
+    Wmap: np.ndarray,
+    OM: np.ndarray,
+    S8: np.ndarray,
+    tom: float,
+    ts8: float,
+    title: str,
+    *,
+    suptext: str | None = None,
+) -> None:
+    mom = float((Wmap * OM).sum())
+    ms8 = float((Wmap * S8).sum())
+    cf = ax.contourf(OM, S8, Wmap, levels=12, cmap="Blues")
+    plt.colorbar(cf, ax=ax, fraction=0.046, pad=0.04)
+    ax.scatter(tom, ts8, s=55, c="r", marker="x", zorder=6, label="true")
+    ax.scatter(mom, ms8, s=60, c="k", marker="+", zorder=6, label="post. mean")
+    ax.set_xlabel(r"$\Omega_m$")
+    ax.set_ylabel(r"$\sigma_8$")
+    ax.legend(fontsize=7)
+    ax.set_title(title, fontsize=8)
+    if suptext:
+        ax.text(0.02, 0.98, suptext, transform=ax.transAxes, fontsize=7, va="top", color="#333")
+
+
+def run_ddpm2_panels(
+    out_dir: Path,
+    images: np.ndarray,
+    labels: np.ndarray,
+    mean: np.ndarray,
+    std: np.ndarray,
+    cfg: Dict,
+    model: torch.nn.Module,
+    device: torch.device,
+    anchor_ix: np.ndarray,
+    grid: int,
+    ddim_steps: int,
+    batch_sz: int,
+) -> None:
+    normalize = bool(cfg.get("normalize_labels", True))
+    H, W = int(images.shape[-2]), int(images.shape[-1])
+    fig, axes = plt.subplots(2, 3, figsize=(14, 9), squeeze=False)
+    for k, ix in enumerate(anchor_ix.ravel()):
+        r, c = divmod(k, 3)
+        ax = axes[r, c]
+        obs = images[ix]
+        lab_t = labels[ix].astype(np.float32)
+        full, om_ax, s8_ax = build_full_grid_2d(labels, grid, tail=None, lab_dim=2)
+        Wmap, OM, S8 = posterior_weights(
+            obs,
+            full,
+            om_ax,
+            s8_ax,
+            mean,
+            std,
+            normalize,
+            model,
+            H=H,
+            W=W,
+            device=device,
+            grid=grid,
+            batch_sz=batch_sz,
+            ddim_steps=ddim_steps,
+        )
+        tom, ts8 = float(lab_t[0]), float(lab_t[1])
+        plot_posterior_panel(
+            ax,
+            Wmap,
+            OM,
+            S8,
+            tom,
+            ts8,
+            f"test ix={ix}\n{_fmt_title(lab_t)}",
+        )
+    plt.suptitle(
+        r"DDPM-2 surrogate posterior on $(\Omega_m,\,\sigma_8)$ — six CAMELS anchors",
+        fontsize=11,
+        y=0.995,
+    )
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    p = out_dir / "posterior_six_anchors_ddpm2.png"
+    fig.savefig(p, dpi=170, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", p)
+
+
+def run_ddpm6_case(
+    out_dir: Path,
+    *,
+    suffix: str,
+    tail_fixed: np.ndarray,
+    tail_name: str,
+    images: np.ndarray,
+    labels: np.ndarray,
+    mean: np.ndarray,
+    std: np.ndarray,
+    cfg: Dict,
+    model: torch.nn.Module,
+    device: torch.device,
+    anchor_ix: np.ndarray,
+    grid: int,
+    ddim_steps: int,
+    batch_sz: int,
+) -> None:
+    normalize = bool(cfg.get("normalize_labels", True))
+    H, W = int(images.shape[-2]), int(images.shape[-1])
+    fig, axes = plt.subplots(2, 3, figsize=(14, 9), squeeze=False)
+    for k, ix in enumerate(anchor_ix.ravel()):
+        r, c = divmod(k, 3)
+        ax = axes[r, c]
+        obs = images[ix]
+        lab_t = labels[ix].astype(np.float32)
+        full, om_ax, s8_ax = build_full_grid_2d(labels, grid, tail=tail_fixed, lab_dim=6)
+        Wmap, OM, S8 = posterior_weights(
+            obs,
+            full,
+            om_ax,
+            s8_ax,
+            mean,
+            std,
+            normalize,
+            model,
+            H=H,
+            W=W,
+            device=device,
+            grid=grid,
+            batch_sz=batch_sz,
+            ddim_steps=ddim_steps,
+        )
+        tom, ts8 = float(lab_t[0]), float(lab_t[1])
+        plot_posterior_panel(
+            ax,
+            Wmap,
+            OM,
+            S8,
+            tom,
+            ts8,
+            f"test ix={ix}",
+            suptext=tail_name,
+        )
+    plt.suptitle(
+        r"DDPM-6 — $(\Omega_m,\,\sigma_8)$ sweep; dims 2–5 fixed (" + tail_name + ")",
+        fontsize=11,
+        y=0.995,
+    )
+    plt.tight_layout(rect=(0, 0, 1, 0.96))
+    p = out_dir / f"posterior_six_anchors_ddpm6_{suffix}.png"
+    fig.savefig(p, dpi=170, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", p)
+
+
+def load_model(bundle_args: Path, ckpt: Path, device: torch.device):
+    cfg = ec.load_training_config(str(bundle_args))
+    model = ec.build_model(cfg, device)
+    ec.load_checkpoint(model, str(ckpt), device)
+    model.eval()
+    return model, cfg
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(
+        description="Six-anchor surrogate posteriors: DDPM-2 and DDPM-6 (extra dims min vs max)."
+    )
+    p.add_argument("--output-dir", type=Path, default=MODELS_ROOT / "ddpm_posterior_six_anchors_out")
+    p.add_argument("--data-2param", type=Path, default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2"))
+    p.add_argument("--data-6param", type=Path, default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"))
+    p.add_argument(
+        "--bundle-2param",
+        type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "2param_epoch200",
+    )
+    p.add_argument(
+        "--bundle-6param",
+        type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "6param_best",
+    )
+    p.add_argument("--split", type=str, default="test", choices=["train", "val", "test"])
+    p.add_argument("--grid", type=int, default=14, help="Grid points per Ωm–σ8 axis.")
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument(
+        "--ddpm2-only",
+        action="store_true",
+        help="Only compute DDPM-2 figure (skip loading DDPM-6).",
+    )
+    p.add_argument(
+        "--ddpm6-only",
+        action="store_true",
+        help="Only compute DDPM-6 figures (skip loading DDPM-2).",
+    )
+    args = p.parse_args()
+
+    out_dir = Path(args.output_dir).resolve()
+    out_dir.mkdir(parents=True, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("device:", device)
+
+    data2 = Path(args.data_2param)
+    data6 = Path(args.data_6param)
+
+    imgs2, lab2 = ec.load_split(data2, args.split)
+    imgs6, lab6 = ec.load_split(data6, args.split)
+
+    n = min(len(lab2), len(lab6))
+    anchor_ix = np.linspace(0, n - 1, num=6, dtype=int)
+
+    low_tail, hi_tail = tail_lhs_bounds(data6)
+    print("LHS tails (dims 2–5): min", low_tail, "max", hi_tail)
+
+    ck2 = args.bundle_2param / "checkpoint_epoch_200.pt"
+    args_json_2 = args.bundle_2param / "args.json"
+    ck6 = args.bundle_6param / "best_model.pt"
+    args_json_6 = args.bundle_6param / "args.json"
+
+    mean2, std2 = ec.load_label_stats(data2)
+    mean6, std6 = ec.load_label_stats(data6)
+
+    if args.ddpm6_only and args.ddpm2_only:
+        raise SystemExit("Use at most one of --ddpm2-only / --ddpm6-only.")
+
+    if not args.ddpm6_only:
+        print(">>> DDPM-2 (six anchors)...")
+        model2, cfg2 = load_model(args_json_2, ck2, device)
+        run_ddpm2_panels(
+            out_dir,
+            imgs2,
+            lab2,
+            mean2,
+            std2,
+            cfg2,
+            model2,
+            device,
+            anchor_ix,
+            args.grid,
+            args.ddim_steps,
+            args.batch_size,
+        )
+        del model2
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    if not args.ddpm2_only:
+        print(">>> DDPM-6 — extra dims at LHS minima (six anchors)...")
+        model6, cfg6 = load_model(args_json_6, ck6, device)
+        run_ddpm6_case(
+            out_dir,
+            suffix="extra_lower",
+            tail_fixed=low_tail,
+            tail_name="min",
+            images=imgs6,
+            labels=lab6,
+            mean=mean6,
+            std=std6,
+            cfg=cfg6,
+            model=model6,
+            device=device,
+            anchor_ix=anchor_ix,
+            grid=args.grid,
+            ddim_steps=args.ddim_steps,
+            batch_sz=args.batch_size,
+        )
+        print(">>> DDPM-6 — extra dims at LHS maxima (six anchors)...")
+        run_ddpm6_case(
+            out_dir,
+            suffix="extra_upper",
+            tail_fixed=hi_tail,
+            tail_name="max",
+            images=imgs6,
+            labels=lab6,
+            mean=mean6,
+            std=std6,
+            cfg=cfg6,
+            model=model6,
+            device=device,
+            anchor_ix=anchor_ix,
+            grid=args.grid,
+            ddim_steps=args.ddim_steps,
+            batch_sz=args.batch_size,
+        )
+        del model6
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    print(f"Done. Outputs in {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

cross_model/scripts/ddpm_triangle_integration.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""
+Surrogate posterior on $(\\Omega_m, \\sigma_8)$ → triangle/MCMC-style chains for one test map.
+
+Loads the same surrogate likelihood used in ``ddpm_posterior_six_anchors``, resamples discrete
+posterior masses to ``--n-hist`` correlated $(\\Omega_m,\\sigma_8)$ pairs, and writes ``.npz``.
+
+DDPM-2: sweeps $(\\Omega_m,\\sigma_8)$.
+DDPM-6: dims 2–5 fixed per ``--six-tail-mode`` (``truth`` uses the test-map labels 2–5; ``min``/``max``
+use LHS extrema from training labels).
+
+If you replace this file with a copy from your machine (Downloads), keep argparse compatible or wrap it.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+_SCRIPTS = Path(__file__).resolve().parent
+if str(_SCRIPTS) not in sys.path:
+    sys.path.insert(0, str(_SCRIPTS))
+
+import ddpm_posterior_six_anchors as dps  # noqa: E402
+
+MODELS_ROOT = Path(__file__).resolve().parents[1]
+CODE_6 = MODELS_ROOT / "6param_ddpm_hi_lh6"
+if str(CODE_6.resolve()) not in sys.path:
+    sys.path.insert(0, str(CODE_6.resolve()))
+
+import evaluate_conditional as ec  # noqa: E402
+
+
+def _tail_vec(
+    mode: str,
+    lab_full: np.ndarray,
+    data6: Path,
+) -> np.ndarray | None:
+    if lab_full.size <= 2:
+        return None
+    if mode == "truth":
+        return lab_full[2:6].astype(np.float32)
+    low, hi = dps.tail_lhs_bounds(data6)
+    if mode == "min":
+        return low
+    if mode == "max":
+        return hi
+    raise ValueError("six-tail-mode must be truth|min|max")
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="DDPM surrogate posterior → resampled Ωm σ8 chains (.npz).")
+    p.add_argument(
+        "--label-dim",
+        type=int,
+        choices=[2, 6],
+        required=True,
+        help="Which model to use.",
+    )
+    p.add_argument(
+        "--bundle",
+        type=Path,
+        default=None,
+        help="Checkpoint bundle dir with args.json (default: notebook_model_weights/<2|6>).",
+    )
+    p.add_argument(
+        "--checkpoint-name",
+        type=str,
+        default=None,
+        help="Checkpoint file under bundle (defaults: DDPM2 epoch200, DDPM6 best_model).",
+    )
+    p.add_argument(
+        "--data-dir",
+        type=Path,
+        default=None,
+        help="LH data dir matching label_dim (default: params_2 vs params_6).",
+    )
+    p.add_argument("--split", type=str, default="test", choices=["train", "val", "test"])
+    p.add_argument("--test-index", type=int, default=56, help="Index into split for CAMELS observation.")
+    p.add_argument("--grid", type=int, default=14)
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument(
+        "--n-hist",
+        type=int,
+        default=10_000,
+        help="Resampled posterior pairs (with replacement).",
+    )
+    p.add_argument(
+        "--six-tail-mode",
+        type=str,
+        default="truth",
+        choices=["truth", "min", "max"],
+        help="Applies only to label_dim==6 — how dims 2–5 are fixed.",
+    )
+    p.add_argument(
+        "--output",
+        "-o",
+        type=Path,
+        required=True,
+        help="Output .npz path.",
+    )
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+
+    ld = args.label_dim
+    if ld == 2:
+        data_dir = args.data_dir or Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2")
+        bundle = args.bundle or MODELS_ROOT / "notebook_model_weights" / "2param_epoch200"
+        ck_name = args.checkpoint_name or "checkpoint_epoch_200.pt"
+    else:
+        data_dir = args.data_dir or Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6")
+        bundle = args.bundle or MODELS_ROOT / "notebook_model_weights" / "6param_best"
+        ck_name = args.checkpoint_name or "best_model.pt"
+
+    rng = np.random.default_rng(args.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    imgs, labs = ec.load_split(data_dir, args.split)
+    ix = int(args.test_index)
+    if not (0 <= ix < len(labs)):
+        raise SystemExit(f"test-index {ix} out of range for split ({len(labs)} rows)")
+    lab_t = labs[ix].astype(np.float64)
+    obs = imgs[ix]
+    ckpt = bundle / ck_name
+    args_json = bundle / "args.json"
+    mean, std = ec.load_label_stats(data_dir)
+
+    tail = None
+    if ld == 6:
+        lab6 = labs[ix].astype(np.float64)
+        if lab6.shape[0] != 6:
+            raise SystemExit("--label-dim 6 requires labels with 6 columns in data-dir")
+        tail = _tail_vec(args.six_tail_mode, lab6, Path(data_dir))
+    model, cfg = dps.load_model(args_json, ckpt, device)
+    normalize = bool(cfg.get("normalize_labels", True))
+    H = int(obs.shape[-2])
+    W = int(obs.shape[-1])
+    gsz = args.grid
+
+    full, om_ax, s8_ax = dps.build_full_grid_2d(labs, gsz, tail=tail, lab_dim=ld)
+    Wmap, OM, S8 = dps.posterior_weights(
+        obs,
+        full,
+        om_ax,
+        s8_ax,
+        mean,
+        std,
+        normalize,
+        model,
+        H=H,
+        W=W,
+        device=device,
+        grid=gsz,
+        batch_sz=args.batch_size,
+        ddim_steps=args.ddim_steps,
+    )
+
+    wflat = np.clip(Wmap.ravel().astype(np.float64), 0.0, None)
+    if wflat.sum() <= 0:
+        raise RuntimeError("Posterior masses collapsed to zero.")
+    wflat /= wflat.sum()
+    omapflat = OM.ravel()
+    s8flat = S8.ravel()
+    draws = rng.choice(np.arange(len(wflat)), size=args.n_hist, replace=True, p=wflat)
+    samp_om = omapflat[draws].astype(np.float64)
+    samp_s8 = s8flat[draws].astype(np.float64)
+
+    out = Path(args.output).resolve()
+    out.parent.mkdir(parents=True, exist_ok=True)
+    tag = f"ddpm{ld}_{args.six_tail_mode}" if ld == 6 else "ddpm2"
+    np.savez_compressed(
+        out,
+        omega_m=samp_om,
+        sigma_8=samp_s8,
+        samples=np.column_stack([samp_om, samp_s8]),
+        truth_Omega_m=float(lab_t[0]),
+        truth_sigma_8=float(lab_t[1]),
+        posterior_map=Wmap,
+        OM=OM,
+        S8=S8,
+        index=np.array(ix, dtype=np.int32),
+        label_dim=np.array(ld, dtype=np.int16),
+        meta_tag=np.array(tag, dtype="U128"),
+        six_tail_mode=np.array(args.six_tail_mode if ld == 6 else "", dtype="U16"),
+    )
+    print("Saved", out, "pairs:", args.n_hist, "device:", device)
+
+
+if __name__ == "__main__":
+    main()

cross_model/scripts/figure6_2409_style.py

@@ -0,0 +1,157 @@
+"""
+Figure-6-inspired layout for 2-parameter posteriors (arXiv:2409.09101 style):
+main 2D panel with 1D marginals on adjacent edges — marginal sums vs profiles (max) optional.
+"""
+
+from __future__ import annotations
+
+from typing import Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import gridspec as mgs
+
+from sigma_contour_utils import compute_sigma_levels
+
+
+def create_figure6_style_plot(
+    Wmap: np.ndarray,
+    om_ax: np.ndarray,
+    s8_ax: np.ndarray,
+    *,
+    true_param1: float,
+    true_param2: float,
+    param1_label: str = r"$\Omega_m$",
+    param2_label: str = r"$\sigma_8$",
+    title: str = "",
+    show_profile: bool = False,
+    figsize: Tuple[float, float] = (10, 10),
+):
+    """
+    Parameters
+    ----------
+    Wmap : (G, G) posterior masses on grid (same layout as DDPM OM meshgrid with indexing='ij').
+    om_ax, s8_ax : 1-D grids aligned with axes 0 and 1 of ``Wmap``.
+    show_profile :
+        False → 1D marginals are sums (*marginal*) over the other parameter.
+        True  → 1D marginals are max (*profile*) over the other parameter (then normalized).
+    """
+    p = np.asarray(Wmap, dtype=np.float64)
+    p = p / (p.sum() + 1e-30)
+    P1, P2 = np.meshgrid(om_ax, s8_ax, indexing="ij")
+
+    if show_profile:
+        m1 = np.max(p, axis=1)
+        m2 = np.max(p, axis=0)
+    else:
+        m1 = p.sum(axis=1)
+        m2 = p.sum(axis=0)
+    m1 = np.asarray(m1, dtype=np.float64)
+    m2 = np.asarray(m2, dtype=np.float64)
+    m1 /= m1.max() + 1e-30
+    m2 /= m2.max() + 1e-30
+
+    fig = plt.figure(figsize=figsize)
+    gs = mgs.GridSpec(
+        nrows=2,
+        ncols=2,
+        figure=fig,
+        width_ratios=[4.0, 1.05],
+        height_ratios=[1.05, 4.0],
+        wspace=0.035,
+        hspace=0.035,
+        left=0.12,
+        right=0.98,
+        bottom=0.1,
+        top=0.92,
+    )
+    ax_main = fig.add_subplot(gs[1, 0])
+    ax_top = fig.add_subplot(gs[0, 0], sharex=ax_main)
+    ax_r = fig.add_subplot(gs[1, 1], sharey=ax_main)
+    ax_empty = fig.add_subplot(gs[0, 1])
+    ax_empty.axis("off")
+
+    lvl = compute_sigma_levels(p, [0.683, 0.954])
+    ax_main.contourf(P1, P2, p, levels=20, cmap="Blues", alpha=0.88)
+    if len(set(lvl)) >= 2:
+        ax_main.contour(P1, P2, p, levels=lvl, colors=["darkblue", "steelblue"], linewidths=[2.0, 1.5])
+    ax_main.scatter(
+        true_param1,
+        true_param2,
+        s=120,
+        c="red",
+        marker="x",
+        linewidths=2.8,
+        zorder=15,
+        label="true",
+    )
+    ax_main.set_xlabel(param1_label, fontsize=13)
+    ax_main.set_ylabel(param2_label, fontsize=13)
+    ax_main.grid(True, alpha=0.28)
+    ax_main.legend(fontsize=8, loc="upper right")
+
+    ax_top.fill_between(om_ax, 0.0, m1, alpha=0.62, color="steelblue")
+    ax_top.axvline(true_param1, color="red", ls="--", lw=2.0)
+    ax_top.set_ylim(0.0, float(np.max(m1) * 1.12))
+    ax_top.set_ylabel("$P(\\mathrm{prof.})$" if show_profile else "$P(\\mathrm{margin.})$", fontsize=10)
+    ax_top.tick_params(labelbottom=False)
+    ax_top.grid(True, alpha=0.25)
+
+    ax_r.fill_betweenx(s8_ax, 0.0, m2, alpha=0.62, color="steelblue")
+    ax_r.axhline(true_param2, color="red", ls="--", lw=2.0)
+    ax_r.set_xlim(0.0, float(np.max(m2) * 1.12))
+    ax_r.set_xlabel("$P$", fontsize=10)
+    ax_r.tick_params(labelleft=False)
+    ax_r.grid(True, alpha=0.25)
+
+    kind = "Profile" if show_profile else "Marginal"
+    fig.suptitle(f"{title} ({kind})", fontsize=14, fontweight="bold", y=0.98)
+
+    plt.setp(ax_top.get_xticklabels(), visible=False)
+
+    return fig
+
+
+def create_comparison_marginal_vs_profile(
+    Wmap: np.ndarray,
+    om_ax: np.ndarray,
+    s8_ax: np.ndarray,
+    *,
+    true_param1: float,
+    true_param2: float,
+    param1_label: str = r"$\Omega_m$",
+    param2_label: str = r"$\sigma_8$",
+    title: str = "",
+    figsize: Tuple[float, float] = (10, 4.2),
+):
+    """Two rows: Ωm and σ8 marginals (sum) vs profile (max) on shared parameter axes."""
+    p = np.asarray(Wmap, dtype=np.float64)
+    p /= p.sum() + 1e-30
+    marg_om = p.sum(axis=1)
+    marg_s8 = p.sum(axis=0)
+    prof_om = np.max(p, axis=1)
+    prof_s8 = np.max(p, axis=0)
+    marg_om /= marg_om.sum() + 1e-30
+    marg_s8 /= marg_s8.sum() + 1e-30
+    prof_om /= prof_om.max() + 1e-30
+    prof_s8 /= prof_s8.max() + 1e-30
+
+    fig, axes = plt.subplots(1, 2, figsize=figsize, sharey=False)
+    for ax, xaxis, marg, prof, xlab, xv in zip(
+        axes,
+        (om_ax, s8_ax),
+        (marg_om, marg_s8),
+        (prof_om, prof_s8),
+        (param1_label, param2_label),
+        (true_param1, true_param2),
+    ):
+        ax.plot(xaxis, marg, lw=2.0, ls="-", label="marginal")
+        ax.plot(xaxis, prof, lw=2.0, ls="--", label="profile")
+        ax.axvline(xv, color="crimson", ls=":", lw=1.8)
+        ax.set_xlabel(xlab, fontsize=12)
+        ax.set_ylabel("norm. density", fontsize=10)
+        ax.legend(fontsize=9)
+        ax.grid(True, alpha=0.3)
+    fig.suptitle(title, fontsize=12, fontweight="bold")
+    fig.tight_layout(rect=(0, 0, 1, 0.93))
+    return fig

cross_model/scripts/run_ddpm_comparison.sh

@@ -0,0 +1,66 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=24:00:00
+#SBATCH --job-name=ddpm_compare
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-ddpm-compare-%j.out
+#SBATCH --error=slurm-ddpm-compare-%j.err
+
+# DDPM-2 vs DDPM-6 comparison (same cluster layout as 6-param training — see reference below).
+#
+# Reference training script (Slurm + module + paths pattern):
+#   /scratch/mrpcol001/Diffusion_job/april_26/ddpm_hi_lh6/scripts/shell/train_conditional_lh6.sh
+#
+# Submit from anywhere:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_comparison.sh
+#
+# Extra CLI args for compare_ddpm_models.py pass through, e.g. LHS off:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_comparison.sh --skip-lhs-r2
+#
+# Override output dir (optional):
+#   sbatch --export=OUTPUT_DIR=/scratch/mrpcol001/Diffusion_job/Models/ddpm_comparison_out_ab \
+#     /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_comparison.sh
+#
+# Optional: override DDPM-2 train/val for the combined loss plot (default: bundled JSON in Models/scripts/):
+#   sbatch --export=SLURM_2PARAM=/path/to/slurm-2param-ddpm.out \
+#     /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_comparison.sh
+#
+# Interactive (same module as training script):
+#   module load python/miniconda3-py3.12-usr
+#   bash .../run_ddpm_comparison.sh --skip-lhs-r2
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_comparison_out}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Job Name: ${SLURM_JOB_NAME:-run_ddpm_comparison}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Starting Time: $(date)"
+echo "Comparison output: ${OUT}"
+echo "Reference Slurm recipe: april_26/ddpm_hi_lh6/scripts/shell/train_conditional_lh6.sh"
+echo "==============================================="
+
+PY_ARGS=(python "$ROOT/scripts/compare_ddpm_models.py" --output-dir "$OUT")
+if [[ -n "${SLURM_2PARAM:-}" ]]; then
+  PY_ARGS+=(--slurm-2param "${SLURM_2PARAM}")
+fi
+PY_ARGS+=("$@")
+
+"${PY_ARGS[@]}"
+
+echo "==============================================="
+echo "Artifacts -> ${OUT}"
+echo "Finished at: $(date)"
+echo "==============================================="

cross_model/scripts/run_ddpm_figure6.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=24:00:00
+#SBATCH --job-name=ddpm_fig6
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-ddpm-figure6-%j.out
+#SBATCH --error=slurm-ddpm-figure6-%j.err
+
+# Figure 6 style (arXiv:2409.09101-inspired) surrogate posteriors for DDPM-2 / DDPM-6.
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_figure6.sh
+#   sbatch --export=OUTPUT_DIR=/path/to/out,TEST_INDEX=42 .../run_ddpm_figure6.sh --no-six-grid
+#
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+module load python/miniconda3-py3.12-usr
+
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_figure6_out}"
+IDX="${TEST_INDEX:-56}"
+
+echo "Job=${SLURM_JOB_ID:-local} OUT=${OUT} TEST_INDEX=${IDX}"
+python "${ROOT}/scripts/run_ddpm_figure6_suite.py" --output-dir "${OUT}" --test-index "${IDX}" "$@"

cross_model/scripts/run_ddpm_figure6_suite.py

@@ -0,0 +1,315 @@
+#!/usr/bin/env python3
+"""
+Compute surrogate posteriors and emit Figure-6 style figures (arXiv:2409.09101-inspired).
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+_SCRIPTS = Path(__file__).resolve().parent
+MODELS_ROOT = Path(__file__).resolve().parents[1]
+CODE_6 = MODELS_ROOT / "6param_ddpm_hi_lh6"
+if str(_SCRIPTS) not in sys.path:
+    sys.path.insert(0, str(_SCRIPTS))
+if str(CODE_6.resolve()) not in sys.path:
+    sys.path.insert(0, str(CODE_6.resolve()))
+
+import evaluate_conditional as ec  # noqa: E402
+import ddpm_posterior_six_anchors as dps  # noqa: E402
+
+from ddpm_figure6_integration import (  # noqa: E402
+    integrate_figure6_model_comparison,
+    integrate_figure6_with_ddpm2,
+    integrate_figure6_with_multi_anchor,
+    print_integration_guide,
+)
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="DDPM Figure-6 style posterior suite.")
+    p.add_argument("--output-dir", type=Path, default=MODELS_ROOT / "ddpm_figure6_out")
+    p.add_argument("--data-2param", type=Path, default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_2"))
+    p.add_argument("--data-6param", type=Path, default=Path("/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"))
+    p.add_argument(
+        "--bundle-2param",
+        type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "2param_epoch200",
+    )
+    p.add_argument(
+        "--bundle-6param",
+        type=Path,
+        default=MODELS_ROOT / "notebook_model_weights" / "6param_best",
+    )
+    p.add_argument("--split", type=str, default="test", choices=["train", "val", "test"])
+    p.add_argument("--test-index", type=int, default=56, help="Index for single comparison + per-map fig6.")
+    p.add_argument("--grid", type=int, default=14)
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument(
+        "--six-anchors-only",
+        action="store_true",
+        help="Only 2×3 multi-anchor plots (skip triple model comparison at --test-index).",
+    )
+    p.add_argument(
+        "--no-six-grid",
+        action="store_true",
+        help="Skip multi-anchor 2×3 panels.",
+    )
+    p.add_argument(
+        "--no-single-fig6",
+        action="store_true",
+        help="Skip per-map marginal/profile for test-index on DDPM-2 and DDPM-6 (truth tail).",
+    )
+    p.add_argument(
+        "--guide",
+        action="store_true",
+        help="Print markdown-style integration notes and exit.",
+    )
+    args = p.parse_args()
+
+    if args.guide:
+        print_integration_guide()
+        return
+
+    out = Path(args.output_dir).resolve()
+    out.mkdir(parents=True, exist_ok=True)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("device:", device)
+
+    data2 = Path(args.data_2param)
+    data6 = Path(args.data_6param)
+    imgs2, lab2 = ec.load_split(data2, args.split)
+    imgs6, lab6 = ec.load_split(data6, args.split)
+    n = min(len(lab2), len(lab6))
+    anchor_ix = np.linspace(0, n - 1, num=6, dtype=int)
+
+    low_tail, hi_tail = dps.tail_lhs_bounds(data6)
+
+    ck2 = args.bundle_2param / "checkpoint_epoch_200.pt"
+    aj2 = args.bundle_2param / "args.json"
+    ck6 = args.bundle_6param / "best_model.pt"
+    aj6 = args.bundle_6param / "args.json"
+
+    mean2, std2 = ec.load_label_stats(data2)
+    mean6, std6 = ec.load_label_stats(data6)
+
+    ix = int(args.test_index)
+    if not (0 <= ix < n):
+        raise SystemExit(f"--test-index {ix} invalid (max {n - 1})")
+
+    lab_box = lab6[:, :2].copy()
+
+    if not args.six_anchors_only:
+        print(">>> Loading models for ix=", ix, "...")
+        m2, c2 = dps.load_model(aj2, ck2, device)
+        m6, c6 = dps.load_model(aj6, ck6, device)
+        normalize2 = bool(c2.get("normalize_labels", True))
+        normalize6 = bool(c6.get("normalize_labels", True))
+
+        obs2 = imgs2[ix]
+        obs6 = imgs6[ix]
+        lt2 = lab2[ix].astype(np.float64)
+        lt6 = lab6[ix].astype(np.float64)
+        ta2om, ta2s8 = float(lt2[0]), float(lt2[1])
+        tom, ts8 = float(lt6[0]), float(lt6[1])
+
+        full2, om_ax, s8_ax = dps.build_full_grid_2d(lab_box, args.grid, tail=None, lab_dim=2)
+        Wm2, _, _ = dps.posterior_weights(
+            obs2,
+            full2,
+            om_ax,
+            s8_ax,
+            mean2,
+            std2,
+            normalize2,
+            m2,
+            H=int(obs2.shape[-2]),
+            W=int(obs2.shape[-1]),
+            device=device,
+            grid=args.grid,
+            batch_sz=args.batch_size,
+            ddim_steps=args.ddim_steps,
+        )
+
+        full6truth, om6, s86 = dps.build_full_grid_2d(
+            lab6, args.grid, tail=lab6[ix, 2:6].astype(np.float32), lab_dim=6
+        )
+        Wm6t, _, _ = dps.posterior_weights(
+            obs6,
+            full6truth,
+            om6,
+            s86,
+            mean6,
+            std6,
+            normalize6,
+            m6,
+            H=int(obs6.shape[-2]),
+            W=int(obs6.shape[-1]),
+            device=device,
+            grid=args.grid,
+            batch_sz=args.batch_size,
+            ddim_steps=args.ddim_steps,
+        )
+
+        full6lo, om_b, s8_b = dps.build_full_grid_2d(lab6, args.grid, tail=low_tail, lab_dim=6)
+        Wm6lo, _, _ = dps.posterior_weights(
+            obs6,
+            full6lo,
+            om_b,
+            s8_b,
+            mean6,
+            std6,
+            normalize6,
+            m6,
+            H=int(obs6.shape[-2]),
+            W=int(obs6.shape[-1]),
+            device=device,
+            grid=args.grid,
+            batch_sz=args.batch_size,
+            ddim_steps=args.ddim_steps,
+        )
+
+        full6hi, om_c, s8_c = dps.build_full_grid_2d(lab6, args.grid, tail=hi_tail, lab_dim=6)
+        Wm6hi, _, _ = dps.posterior_weights(
+            obs6,
+            full6hi,
+            om_c,
+            s8_c,
+            mean6,
+            std6,
+            normalize6,
+            m6,
+            H=int(obs6.shape[-2]),
+            W=int(obs6.shape[-1]),
+            device=device,
+            grid=args.grid,
+            batch_sz=args.batch_size,
+            ddim_steps=args.ddim_steps,
+        )
+
+        if not (np.allclose(om_ax, om_b, rtol=0, atol=1e-12) and np.allclose(s8_ax, s86)):
+            print("Warning: Ωm–σ8 grids differ between setups; plotting uses DDPM-2 Ωm/σ8 axes.")
+
+        integrate_figure6_model_comparison(
+            {
+                "DDPM-2": Wm2,
+                "DDPM-6 (truth-tail)": Wm6t,
+                "DDPM-6 (min-tail)": Wm6lo,
+                "DDPM-6 (max-tail)": Wm6hi,
+            },
+            om_ax,
+            s8_ax,
+            tom,
+            ts8,
+            ix,
+            out,
+        )
+
+        if not args.no_single_fig6:
+            integrate_figure6_with_ddpm2(Wm2, om_ax, s8_ax, ta2om, ta2s8, ix, out, model_name="DDPM-2")
+            integrate_figure6_with_ddpm2(Wm6t, om_ax, s8_ax, tom, ts8, ix, out, model_name="DDPM-6-truth")
+
+        del m2, m6
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    # --- Six anchors: multi grids for DDPM-2 + DDPM-6 truth tail ---
+    if not args.no_six_grid:
+        print(">>> Six-anchor Figure 6 grids...")
+        model2, cfg2 = dps.load_model(aj2, ck2, device)
+        model6, cfg6 = dps.load_model(aj6, ck6, device)
+        nz2 = bool(cfg2.get("normalize_labels", True))
+        nz6 = bool(cfg6.get("normalize_labels", True))
+
+        post2: list[np.ndarray] = []
+        post6: list[np.ndarray] = []
+        truths: list[tuple[float, float]] = []
+        indices: list[int] = []
+
+        for k, jx in enumerate(anchor_ix.ravel()):
+            indices.append(int(jx))
+            o2 = imgs2[jx]
+            lb2 = lab2[jx].astype(np.float64)
+            f2, oa, sa = dps.build_full_grid_2d(lab_box, args.grid, tail=None, lab_dim=2)
+            W2, _, _ = dps.posterior_weights(
+                o2,
+                f2,
+                oa,
+                sa,
+                mean2,
+                std2,
+                nz2,
+                model2,
+                H=int(o2.shape[-2]),
+                W=int(o2.shape[-1]),
+                device=device,
+                grid=args.grid,
+                batch_sz=args.batch_size,
+                ddim_steps=args.ddim_steps,
+            )
+            post2.append(W2)
+
+            o6 = imgs6[jx]
+            lb6 = lab6[jx]
+            tail_truth = lb6.astype(np.float32)[2:6]
+            f6, oa6, sa6 = dps.build_full_grid_2d(lab6, args.grid, tail=tail_truth, lab_dim=6)
+            W6, _, _ = dps.posterior_weights(
+                o6,
+                f6,
+                oa6,
+                sa6,
+                mean6,
+                std6,
+                nz6,
+                model6,
+                H=int(o6.shape[-2]),
+                W=int(o6.shape[-1]),
+                device=device,
+                grid=args.grid,
+                batch_sz=args.batch_size,
+                ddim_steps=args.ddim_steps,
+            )
+            post6.append(W6)
+            truths.append((float(lb2[0]), float(lb2[1])))
+
+        integrate_figure6_with_multi_anchor(
+            post2,
+            oa,
+            sa,
+            truths,
+            indices,
+            out,
+            model_name="DDPM-2",
+        )
+
+        truths6 = [(float(lab6[int(j)][0]), float(lab6[int(j)][1])) for j in anchor_ix]
+
+        integrate_figure6_with_multi_anchor(
+            post6,
+            oa6,
+            sa6,
+            truths6,
+            indices,
+            out,
+            model_name="DDPM-6-truth-tail",
+        )
+
+        del model2, model6
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    print(f"Done. Outputs in {out}")
+
+
+if __name__ == "__main__":
+    main()

cross_model/scripts/run_ddpm_posterior_corrected.sh

@@ -0,0 +1,58 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=48:00:00
+#SBATCH --job-name=ddpm_post_corr
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-ddpm-post-corr-%j.out
+#SBATCH --error=slurm-ddpm-post-corr-%j.err
+
+# Prior / likelihood / posterior visualization pipeline (ddpm_posterior_corrected.py).
+# Separate from poster.py — default output dir is ddpm_posterior_corrected_fullviz_out.
+#
+# Submit:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_posterior_corrected.sh
+#
+# Extra args pass through to the Python script:
+#   sbatch .../run_ddpm_posterior_corrected.sh --ddpm2-only --grid 20 --n-ddpm-samples 4 --no-ppc
+#
+# Override dirs:
+#   sbatch --export=OUTPUT_DIR=/path/to/out,CUSTOM_LOG=/path/run.log \\
+#     .../run_ddpm_posterior_corrected.sh
+#
+# Interactive:
+#   module load python/miniconda3-py3.12-usr
+#   bash .../run_ddpm_posterior_corrected.sh --help
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_posterior_corrected_fullviz_out}"
+mkdir -p "$OUT"
+
+# Copy of stdout/stderr for progress (also appears in Slurm .out/.err)
+RUN_LOG="${CUSTOM_LOG:-${OUT}/run_log.txt}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Started: $(date)"
+echo "Output dir: ${OUT}"
+echo "Progress log (tee): ${RUN_LOG}"
+echo "==============================================="
+
+set -o pipefail
+python -u "${ROOT}/ddpm_posterior_corrected.py" --output-dir "${OUT}" "$@" 2>&1 | tee -a "${RUN_LOG}"
+
+echo "==============================================="
+echo "Finished: $(date)"
+echo "Figures & log → ${OUT}"
+echo "==============================================="

cross_model/scripts/run_ddpm_posterior_six_anchors.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=48:00:00
+#SBATCH --job-name=ddpm_post6
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-ddpm-posterior-six-%j.out
+#SBATCH --error=slurm-ddpm-posterior-six-%j.err
+
+# Six-anchor surrogate posteriors (DDPM-2 + DDPM-6 with extra dims min/max).
+#
+# Submit from anywhere:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_posterior_six_anchors.sh
+#
+# Override output directory:
+#   sbatch --export=OUTPUT_DIR=/scratch/mrpcol001/Diffusion_job/Models/my_post_out \
+#     /scratch/mrpcol001/Diffusion_job/Models/scripts/run_ddpm_posterior_six_anchors.sh
+#
+# Extra CLI passes through to ddpm_posterior_six_anchors.py, e.g. only DDPM-6 panels:
+#   sbatch .../run_ddpm_posterior_six_anchors.sh --ddpm6-only --grid 12
+#
+# Interactive:
+#   module load python/miniconda3-py3.12-usr
+#   bash .../run_ddpm_posterior_six_anchors.sh
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_posterior_six_anchors_out}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Job Name: ${SLURM_JOB_NAME:-run_ddpm_posterior_six_anchors}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Starting Time: $(date)"
+echo "Posterior output: ${OUT}"
+echo "==============================================="
+
+python "$ROOT/scripts/ddpm_posterior_six_anchors.py" --output-dir "$OUT" "$@"
+
+echo "==============================================="
+echo "Artifacts -> ${OUT}"
+echo "Finished at: $(date)"
+echo "==============================================="

cross_model/scripts/run_poster.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=24:00:00
+#SBATCH --job-name=ddpm_poster
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-ddpm-poster-%j.out
+#SBATCH --error=slurm-ddpm-poster-%j.err
+
+# Corrected six-anchor surrogate posteriors (poster.py): DDPM-2 + DDPM-6 with
+# stochastic averaging, calibrated sigma_pk, MC marginalisation for 6-param, etc.
+#
+# Submit from anywhere:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_poster.sh
+#
+# Override output directory:
+#   sbatch --export=OUTPUT_DIR=/scratch/mrpcol001/Diffusion_job/Models/my_poster_out \
+#     /scratch/mrpcol001/Diffusion_job/Models/scripts/run_poster.sh
+#
+# Extra CLI passes through to poster.py, e.g. DDPM-2 only (faster debug):
+#   sbatch .../run_poster.sh --ddpm2-only --grid 14 --n-pk-samples 4 --n-marg-samples 1 --no-ppc
+#
+# Interactive (same module as other Models scripts):
+#   module load python/miniconda3-py3.12-usr
+#   bash /scratch/mrpcol001/Diffusion_job/Models/scripts/run_poster.sh --help
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_posterior_corrected_out}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Job Name: ${SLURM_JOB_NAME:-run_poster}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Starting Time: $(date)"
+echo "Poster output: ${OUT}"
+echo "==============================================="
+
+python "$ROOT/poster.py" --output-dir "$OUT" "$@"
+
+echo "==============================================="
+echo "Artifacts -> ${OUT}"
+echo "Finished at: $(date)"
+echo "==============================================="

cross_model/scripts/run_posterior_inference.sh

@@ -0,0 +1,74 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=48:00:00
+#SBATCH --job-name=vlb_infer
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-vlb-infer-%j.out
+#SBATCH --error=slurm-vlb-infer-%j.err
+
+# VLB / Mudur-style posterior_inference.py (pure inference-time L_t surfaces).
+#
+# Defaults match bundled 6-param checkpoint + LH test data (override via env).
+#
+# Submit:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/scripts/run_posterior_inference.sh
+#
+# Defaults (posterior_inference.py): n_fields=9, grid_size=10000 (needs --allow_huge_grid),
+#   mosaic figure posterior_L0_mosaic_3x3.png at ~10000×10000 px.
+# Override grid without huge scan, e.g.: --grid_size 50 (then --allow_huge_grid not needed)
+# Smoke test:
+#   sbatch .../run_posterior_inference.sh --n_fields 1 --grid_size 25 --t_subset 0 --batch_size 16
+#
+# Custom checkpoint / args / data:
+#   sbatch --export=CHECKPOINT=/path/best_model.pt,TRAINING_ARGS=/path/args.json,DATA_DIR=/path/params_6 \\
+#     .../run_posterior_inference.sh --grid_size 40
+#
+# Logs: Slurm .out/.err plus OUTPUT_DIR/run_log.txt (override CUSTOM_LOG).
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+PY="${ROOT}/6param_ddpm_hi_lh6/posterior_inference.py"
+OUT="${OUTPUT_DIR:-${ROOT}/vlb_inference_outputs}"
+
+CHK="${CHECKPOINT:-${ROOT}/notebook_model_weights/6param_best/best_model.pt}"
+ARGS="${TRAINING_ARGS:-${ROOT}/notebook_model_weights/6param_best/args.json}"
+DATA="${DATA_DIR:-/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6}"
+
+mkdir -p "${OUT}"
+RUN_LOG="${CUSTOM_LOG:-${OUT}/run_log.txt}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Started: $(date)"
+echo "Python: ${PY}"
+echo "checkpoint: ${CHK}"
+echo "training_args: ${ARGS}"
+echo "data_dir: ${DATA}"
+echo "output_dir: ${OUT}"
+echo "Progress log: ${RUN_LOG}"
+echo "==============================================="
+
+set -o pipefail
+python -u "${PY}" \
+  --checkpoint "${CHK}" \
+  --training_args "${ARGS}" \
+  --data_dir "${DATA}" \
+  --output_dir "${OUT}" \
+  --allow_huge_grid \
+  "$@" 2>&1 | tee -a "${RUN_LOG}"
+
+echo "==============================================="
+echo "Finished: $(date)"
+echo "Artifacts → ${OUT}"
+echo "==============================================="

cross_model/scripts/run_triangle_ddpm_both.sh

@@ -0,0 +1,75 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=12:00:00
+#SBATCH --job-name=ddpm_triangle
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-ddpm-triangle-%j.out
+#SBATCH --error=slurm-ddpm-triangle-%j.err
+
+# Run surrogate Ωm–σ8 chain export for DDPM-2 and DDPM-6, then a joint triangle plot.
+# If you have your own copies of ddpm_triangle_integration.py / triangle_plot_posterior.py
+# under $ROOT/scripts (e.g. copied from ~/Downloads), they override the repo versions.
+#
+#   sbatch .../run_triangle_ddpm_both.sh
+#
+#   sbatch --export=OUTPUT_DIR=/path/to/out,TEST_INDEX=56 .../run_triangle_ddpm_both.sh
+#
+# Interactive:
+#   module load python/miniconda3-py3.12-usr
+#   bash .../run_triangle_ddpm_both.sh
+
+set -euo pipefail
+
+ROOT="/scratch/mrpcol001/Diffusion_job/Models"
+cd "$ROOT"
+
+module load python/miniconda3-py3.12-usr
+
+OUT="${OUTPUT_DIR:-${ROOT}/ddpm_triangle_out}"
+TEST_IX="${TEST_INDEX:-56}"
+GRID="${GRID_POINTS:-14}"
+
+mkdir -p "${OUT}"
+
+INTEG="${DDPM_TRIANGLE_INTEGRATION_PY:-${ROOT}/scripts/ddpm_triangle_integration.py}"
+PLOT="${DDPM_TRIANGLE_POSTERIOR_PY:-${ROOT}/scripts/triangle_plot_posterior.py}"
+
+for f in "$INTEG" "$PLOT"; do
+  if [[ ! -f "$f" ]]; then
+    echo "Missing: $f"
+    exit 1
+  fi
+done
+
+CHAIN2="${OUT}/chain_surrogate_ix${TEST_IX}_ddpm2.npz"
+CHAIN6="${OUT}/chain_surrogate_ix${TEST_IX}_ddpm6_truth_tail.npz"
+
+echo "==============================================="
+echo "Job: ${SLURM_JOB_ID:-local}  OUT=${OUT}  test_ix=${TEST_IX}"
+echo "==============================================="
+
+python "$INTEG" \
+  --label-dim 2 \
+  --test-index "${TEST_IX}" \
+  --grid "${GRID}" \
+  -o "${CHAIN2}"
+
+python "$INTEG" \
+  --label-dim 6 \
+  --test-index "${TEST_IX}" \
+  --six-tail-mode truth \
+  --grid "${GRID}" \
+  -o "${CHAIN6}"
+
+python "$PLOT" \
+  -i "${CHAIN2}" "${CHAIN6}" \
+  --labels "DDPM-2" "DDPM-6" \
+  -o "${OUT}/triangle_ddpm2_ddpm6_ix${TEST_IX}.png"
+
+echo "Chains: ${CHAIN2} ${CHAIN6}"
+echo "Triangle: ${OUT}/triangle_ddpm2_ddpm6_ix${TEST_IX}.png"
+echo "Finished: $(date)"

cross_model/scripts/sigma_contour_utils.py

@@ -0,0 +1,29 @@
+"""HDR-style contour levels for 2D probability maps on a grid."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def compute_sigma_levels(
+    posterior_norm: np.ndarray,
+    credibility_mass: tuple[float, ...] | list[float],
+) -> list[float]:
+    """
+    Highest-density containment: find density thresholds such that descending sort
+    of mass covers ``credibility_mass[j]`` of total probability.
+
+    Returned levels are ascending (suitable order for matplotlib ``contour``).
+    """
+    p = np.asarray(posterior_norm, dtype=np.float64).ravel()
+    s = p.sum()
+    if s <= 0:
+        return [0.0 for _ in credibility_mass]
+    ps = np.sort((p / s).flatten())[::-1]
+    cdf = np.cumsum(ps)
+    out: list[float] = []
+    for cred in credibility_mass:
+        j = int(np.searchsorted(cdf, cred, side="left"))
+        j = min(max(j, 0), len(ps) - 1)
+        out.append(float(ps[j]))
+    return sorted(out)

cross_model/scripts/triangle_plot_posterior.py

@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""
+Corner-style triangle plot for surrogate $(\\Omega_m,\\sigma_8)$ chains from ``ddpm_triangle_integration.py``.
+
+Loads one or two ``.npz`` files (keys ``omega_m``, ``sigma_8`` / ``samples``, ``truth_*``) and draws
+1D marginals + 2D density. If you substitute a script from your Downloads, keep ``--inputs``
+and the expected ``.npz`` keys compatible.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def _load_chain(path: Path) -> tuple[np.ndarray, np.ndarray, tuple[float, float] | None]:
+    d = np.load(path, allow_pickle=True)
+    if "samples" in d:
+        s = np.asarray(d["samples"], dtype=np.float64)
+        om, s8 = s[:, 0], s[:, 1]
+    else:
+        om = np.asarray(d["omega_m"], dtype=np.float64).ravel()
+        s8 = np.asarray(d["sigma_8"], dtype=np.float64).ravel()
+    truth = None
+    if "truth_Omega_m" in d.files and "truth_sigma_8" in d.files:
+        truth = (float(d["truth_Omega_m"]), float(d["truth_sigma_8"]))
+    return om, s8, truth
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Triangle / corner plot for Ωm–σ8 surrogate chains.")
+    p.add_argument(
+        "--inputs",
+        "-i",
+        nargs="+",
+        type=Path,
+        required=True,
+        help="One or two .npz outputs from ddpm_triangle_integration.py",
+    )
+    p.add_argument(
+        "--labels",
+        nargs="*",
+        default=None,
+        help="Legend entries (default: paths' stems).",
+    )
+    p.add_argument(
+        "--output",
+        "-o",
+        type=Path,
+        default=None,
+        help="Output PNG (default: triangle_posterior_ddpm2_ddpm6.png next to first input).",
+    )
+    p.add_argument("--bins-1d", type=int, default=40)
+    p.add_argument("--bins-2d", type=int, default=45)
+    args = p.parse_args()
+
+    paths = [Path(x).resolve() for x in args.inputs]
+    names = args.labels if args.labels else [p.stem for p in paths]
+    if len(names) != len(paths):
+        raise SystemExit("--labels count must match --inputs")
+
+    colors = ("#1f77b4", "#d95f02", "#2ca02c")
+    fig = plt.figure(figsize=(8.2, 8.0))
+
+    ax00 = fig.add_axes([0.1, 0.55, 0.35, 0.35])
+    ax_cont = fig.add_axes([0.1, 0.1, 0.35, 0.35])
+    ax11 = fig.add_axes([0.55, 0.1, 0.35, 0.35])
+    ax_blank = fig.add_axes([0.55, 0.55, 0.35, 0.35])
+    ax_blank.axis("off")
+
+    for i, path in enumerate(paths):
+        om, s8, truth = _load_chain(path)
+        c = colors[i % len(colors)]
+        ax00.hist(
+            om,
+            bins=args.bins_1d,
+            density=True,
+            histtype="step",
+            color=c,
+            lw=2.0,
+            label=names[i],
+        )
+        ax11.hist(
+            s8,
+            bins=args.bins_1d,
+            density=True,
+            histtype="step",
+            color=c,
+            lw=2.0,
+        )
+        h2, xe, ye = np.histogram2d(om, s8, bins=args.bins_2d, density=True)
+        xc = 0.5 * (xe[1:] + xe[:-1])
+        yc = 0.5 * (ye[1:] + ye[:-1])
+        X, Y = np.meshgrid(xc, yc, indexing="ij")
+        Z = np.ma.masked_where(h2.T <= 1e-20, h2.T)
+        if i == 0 and np.ma.count(Z) > 0:
+            cf = ax_cont.contourf(X, Y, Z, alpha=0.45, cmap="Blues")
+            fig.colorbar(cf, ax=ax_cont, fraction=0.046, pad=0.04)
+        elif np.ma.count(Z) > 0:
+            ax_cont.contour(X, Y, Z, colors=[c], linewidths=[1.85])
+        if truth:
+            tx, ty = truth
+            ax_cont.scatter(tx, ty, marker="x", s=88, color=c, zorder=6)
+
+    ax00.set_title(r"$P(\Omega_m)$ marginal")
+    ax00.set_ylabel("density")
+    ax00.legend(fontsize=8, loc="upper right")
+    ax_cont.set_title(r"$2D$ surrogate posterior density")
+    ax_cont.set_xlabel(r"$\Omega_m$")
+    ax_cont.set_ylabel(r"$\sigma_8$")
+    ax11.set_title(r"$P(\sigma_8)$ marginal")
+    ax11.set_xlabel("density")
+
+    out = args.output or (paths[0].parent / ("triangle_" + "_".join(p.stem for p in paths) + ".png"))
+    out.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(out, dpi=170, bbox_inches="tight")
+    plt.close(fig)
+    print("Saved", out)
+
+
+if __name__ == "__main__":
+    main()

cross_model/submit_vlb_1000grid.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""
+Submit 1000×1000 grid VLB inference job.
+Usage: python3 submit_vlb_1000grid.py [--dry-run]
+"""
+
+import subprocess
+import argparse
+from pathlib import Path
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Submit high-resolution VLB inference (1000×1000 grid)"
+    )
+    parser.add_argument("--dry-run", action="store_true",
+                       help="Print command without submitting")
+    parser.add_argument("--grid-size", type=int, default=1000,
+                       help="Grid resolution (default: 1000)")
+    parser.add_argument("--n-fields", type=int, default=9,
+                       help="Number of test fields (default: 9)")
+    parser.add_argument("--job-name", default="vlb-infer-1000",
+                       help="SLURM job name")
+    parser.add_argument("--time", default="24:00:00",
+                       help="SLURM time limit (default: 24:00:00)")
+    parser.add_argument("--mem", default="32G",
+                       help="Memory requirement (default: 32G)")
+    parser.add_argument("--batch-size", type=int, default=32,
+                       help="Batch size (default: 32, reduce if OOM)")
+    args = parser.parse_args()
+
+    script_dir = Path(__file__).parent.resolve()
+    checkpoint = script_dir / "notebook_model_weights/6param_best/best_model.pt"
+    training_args = script_dir / "notebook_model_weights/6param_best/args.json"
+    data_dir = script_dir.parent / "data/LH_data/params_6"
+    output_dir = script_dir / f"vlb_inference_outputs_{args.grid_size}grid"
+
+    # Validate paths
+    for path, name in [
+        (checkpoint, "checkpoint"),
+        (training_args, "training_args"),
+        (data_dir, "data_dir"),
+    ]:
+        if not path.exists():
+            print(f"❌ Error: {name} not found at {path}")
+            return 1
+
+    cmd = [
+        "sbatch",
+        f"--job-name={args.job_name}",
+        f"--time={args.time}",
+        f"--mem={args.mem}",
+        f"--output=slurm-vlb-infer-{args.grid_size}-%j.out",
+        f"--error=slurm-vlb-infer-{args.grid_size}-%j.err",
+        str(script_dir / "run_vlb_inference_1000grid.sh"),
+    ]
+
+    # Build environment variables to override defaults in script
+    env_cmd = [
+        f"GRID_SIZE={args.grid_size}",
+        f"N_FIELDS={args.n_fields}",
+        f"BATCH_SIZE={args.batch_size}",
+        f"OUTPUT_DIR={output_dir}",
+    ]
+
+    full_cmd = env_cmd + cmd
+
+    print("=" * 60)
+    print("VLB Inference Submission — 1000×1000 Grid")
+    print("=" * 60)
+    print(f"Grid size:        {args.grid_size}×{args.grid_size}")
+    print(f"Number of fields: {args.n_fields}")
+    print(f"Batch size:       {args.batch_size}")
+    print(f"Output dir:       {output_dir}")
+    print(f"Memory:           {args.mem}")
+    print(f"Time limit:       {args.time}")
+    print("=" * 60)
+
+    # Compute estimates
+    grid_points = args.grid_size ** 2
+    timesteps = 8
+    seeds = 4
+    n_forward_passes = grid_points * timesteps * seeds * args.n_fields
+    print(f"\nComputation scale:")
+    print(f"  Total forward passes: {n_forward_passes:,}")
+    print(f"  Per field:            {n_forward_passes // args.n_fields:,}")
+    print(f"  Est. time per field:  ~{(n_forward_passes // args.n_fields) // 10_000:,} min")
+    print(f"  Est. total time:      ~{(n_forward_passes // 10_000_000):.1f}-{(n_forward_passes // 7_000_000):.1f} hours")
+
+    print(f"\nCommand: {' '.join(full_cmd)}\n")
+
+    if args.dry_run:
+        print("✓ Dry-run (not submitted)")
+        return 0
+
+    try:
+        result = subprocess.run(full_cmd, check=True, capture_output=True, text=True)
+        print("✓ Job submitted!")
+        print(result.stdout)
+        return 0
+    except subprocess.CalledProcessError as e:
+        print(f"❌ Submission failed: {e}")
+        print(e.stderr)
+        return 1
+
+if __name__ == "__main__":
+    exit(main())

scripts/shell/evaluate_conditional_lh6.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=04:00:00
+#SBATCH --job-name=ddpm_hi_lh6_eval
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-eval-%j.out
+#SBATCH --error=slurm-eval-%j.err
+
+# Evaluate conditional DDPM (6 CAMELS LH parameters).
+# Submit:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/6param_ddpm_hi_lh6/scripts/shell/evaluate_conditional_lh6.sh
+#
+# Optional overrides (example):
+#   sbatch --export=CHECKPOINT=/path/to/best_model.pt,OUTPUT_DIR=/path/to/eval_out evaluate_conditional_lh6.sh
+
+REPO="/scratch/mrpcol001/Diffusion_job/Models/6param_ddpm_hi_lh6"
+cd "${REPO}" || exit 1
+
+module load python/miniconda3-py3.12-usr
+
+DATA_DIR="${DATA_DIR:-/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6}"
+# Default: trained run kept under april_26 (large artifacts not duplicated here).
+CHECKPOINT="${CHECKPOINT:-/scratch/mrpcol001/Diffusion_job/april_26/ddpm_hi_lh6/outputs_conditional_6param_20260413_132226/checkpoints/best_model.pt}"
+OUTPUT_DIR="${OUTPUT_DIR:-${REPO}/evaluation_outputs_6param}"
+TRAINING_ARGS="${TRAINING_ARGS:-}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Job Name: ${SLURM_JOB_NAME:-evaluate_conditional_lh6}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Starting Time: $(date)"
+echo "CHECKPOINT: ${CHECKPOINT}"
+echo "DATA_DIR: ${DATA_DIR}"
+echo "OUTPUT_DIR: ${OUTPUT_DIR}"
+echo "==============================================="
+
+EVAL_ARGS=(
+  python evaluate_conditional.py
+  --checkpoint "${CHECKPOINT}"
+  --data_dir "${DATA_DIR}"
+  --output_dir "${OUTPUT_DIR}"
+  --split test
+  --num_samples 8
+  --ddim_steps 50
+)
+
+if [[ -n "${TRAINING_ARGS}" ]]; then
+  EVAL_ARGS+=(--training_args "${TRAINING_ARGS}")
+fi
+
+"${EVAL_ARGS[@]}"
+
+echo "==============================================="
+echo "Evaluation completed at: $(date)"
+echo "Plots and evaluation_data.npz under: ${OUTPUT_DIR}"
+echo "==============================================="

scripts/shell/plot_r2_cosmology_lhs.sh

@@ -0,0 +1,72 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=12:00:00
+#SBATCH --job-name=ddpm_r2_lhs
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-r2-lhs-%j.out
+#SBATCH --error=slurm-r2-lhs-%j.err
+
+# Latin-hypercube R² figure (plot_r2_cosmology_lhs.py): μ(P) and σ(P) vs (Ωm, σ8).
+#
+# Submit (full DDIM run — slow):
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/6param_ddpm_hi_lh6/scripts/shell/plot_r2_cosmology_lhs.sh
+#
+# Plot only from saved NPZ (fast):
+#   sbatch --export=FROM_NPZ=/path/to/r2_lhs_data.npz /scratch/.../plot_r2_cosmology_lhs.sh
+#
+# Optional env vars:
+#   CHECKPOINT, DATA_DIR, OUTPUT_PNG, SAVE_NPZ, LHS_N, MAPS_PER_POINT, DDIM_STEPS, SEED
+
+REPO="/scratch/mrpcol001/Diffusion_job/Models/6param_ddpm_hi_lh6"
+cd "${REPO}" || exit 1
+
+module load python/miniconda3-py3.12-usr
+
+DATA_DIR="${DATA_DIR:-/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6}"
+CHECKPOINT="${CHECKPOINT:-/scratch/mrpcol001/Diffusion_job/april_26/ddpm_hi_lh6/outputs_conditional_6param_20260413_132226/checkpoints/best_model.pt}"
+OUTPUT_PNG="${OUTPUT_PNG:-${REPO}/ddpm_eval_notebook_out/r2_cosmology_lhs50_ddpm.png}"
+FROM_NPZ="${FROM_NPZ:-}"
+SAVE_NPZ="${SAVE_NPZ:-}"
+LHS_N="${LHS_N:-50}"
+MAPS_PER_POINT="${MAPS_PER_POINT:-15}"
+DDIM_STEPS="${DDIM_STEPS:-50}"
+SEED="${SEED:-42}"
+
+echo "==============================================="
+echo "Job ID: ${SLURM_JOB_ID:-local}"
+echo "Job Name: ${SLURM_JOB_NAME:-plot_r2_cosmology_lhs}"
+echo "Node: ${SLURM_NODELIST:-$(hostname)}"
+echo "GPU: ${CUDA_VISIBLE_DEVICES:-n/a}"
+echo "Starting Time: $(date)"
+echo "OUTPUT_PNG: ${OUTPUT_PNG}"
+echo "FROM_NPZ: ${FROM_NPZ:-(none — full compute)}"
+echo "==============================================="
+
+PY_ARGS=(
+  python plot_r2_cosmology_lhs.py
+  --output "${OUTPUT_PNG}"
+  --lhs-n "${LHS_N}"
+  --maps-per-point "${MAPS_PER_POINT}"
+  --ddim-steps "${DDIM_STEPS}"
+  --seed "${SEED}"
+)
+
+if [[ -n "${FROM_NPZ}" ]]; then
+  PY_ARGS+=(--from-npz "${FROM_NPZ}")
+else
+  PY_ARGS+=(--checkpoint "${CHECKPOINT}" --data-dir "${DATA_DIR}")
+  if [[ -n "${SAVE_NPZ}" ]]; then
+    PY_ARGS+=(--save-npz "${SAVE_NPZ}")
+  fi
+fi
+
+"${PY_ARGS[@]}"
+
+echo "==============================================="
+echo "Finished at: $(date)"
+echo "Figure: ${OUTPUT_PNG}"
+echo "==============================================="

scripts/shell/train_conditional_lh6.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+#SBATCH --account=l40sfree
+#SBATCH --partition=l40s
+#SBATCH --nodes=1
+#SBATCH --ntasks=8
+#SBATCH --gres=gpu:l40s:1
+#SBATCH --time=48:00:00
+#SBATCH --job-name=ddpm_hi_lh6
+#SBATCH --mail-user=mrpcol001@myuct.ac.za
+#SBATCH --output=slurm-%j.out
+#SBATCH --error=slurm-%j.err
+
+# Conditional DDPM training — 6 CAMELS LH parameters (ddpm_hi_lh6).
+# Submit from anywhere:
+#   sbatch /scratch/mrpcol001/Diffusion_job/Models/6param_ddpm_hi_lh6/scripts/shell/train_conditional_lh6.sh
+#
+# Override data path (optional): any folder containing *_LH_6.npy and *_labels_LH.npy
+#   sbatch --export=DATA_DIR=/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6 train_conditional_lh6.sh
+
+cd /scratch/mrpcol001/Diffusion_job/Models/6param_ddpm_hi_lh6
+
+module load python/miniconda3-py3.12-usr
+
+# Same LH_data layout as DDPM_HI_Emulation_improved (params_2 for 2 labels → params_6 here).
+DATA_DIR="${DATA_DIR:-/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6}"
+
+echo "==============================================="
+echo "Job ID: $SLURM_JOB_ID"
+echo "Job Name: $SLURM_JOB_NAME"
+echo "Node: $SLURM_NODELIST"
+echo "GPU: $CUDA_VISIBLE_DEVICES"
+echo "Starting Time: $(date)"
+echo "Conditional diffusion training (ddpm_hi_lh6, 6 labels)"
+echo "DATA_DIR: ${DATA_DIR}"
+echo "==============================================="
+
+python train_conditional.py \
+    --label_dim 6 \
+    --timesteps 1500 \
+    --use_ddim \
+    --ddim_steps 50 \
+    --normalize_labels \
+    --batch_size 8 \
+    --epochs 200 \
+    --lr 2e-4 \
+    --early_stop_patience 100 \
+    --sample_every 10 \
+    --base_channels 64 \
+    --channel_multipliers 1 2 4 8 \
+    --attention_levels 2 3 \
+    --data_dir "${DATA_DIR}" \
+    --output_dir outputs_conditional_6param \
+    --use_amp
+
+# To resume: point --resume at checkpoints/checkpoint_epoch_N.pt and set --epochs to the
+# new total; add --resume_refresh_scheduler if extending past the original epoch count.
+
+echo "==============================================="
+echo "Training completed at: $(date)"
+echo "==============================================="

src/eval_model.py

@@ -0,0 +1,86 @@
+"""
+Helpers for ddpm_cond_eval.ipynb: R², P(k) on log N_HI fields, DDIM batches.
+
+Uses evaluate_conditional for PowerSpectrum, sampling, and label z-scoring.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+from matplotlib.colors import LinearSegmentedColormap
+
+import evaluate_conditional as ec
+
+LO_LOG, HI_LOG = 14.0, 22.0
+
+
+def r2_score_1d(y_true: np.ndarray, y_pred: np.ndarray) -> float:
+    """Univariate R² (same as sklearn for 1D arrays)."""
+    y_true = np.asarray(y_true, dtype=np.float64).ravel()
+    y_pred = np.asarray(y_pred, dtype=np.float64).ravel()
+    ss_res = np.sum((y_true - y_pred) ** 2)
+    ss_tot = np.sum((y_true - np.mean(y_true)) ** 2)
+    if ss_tot < 1e-30:
+        return 0.0 if ss_res < 1e-30 else float("-inf")
+    return float(1.0 - ss_res / ss_tot)
+
+
+def cmap_r2_hiflow() -> LinearSegmentedColormap:
+    """Green → yellow → red → purple → dark blue (HIFlow Fig. 7 style)."""
+    return LinearSegmentedColormap.from_list(
+        "r2_hiflow",
+        ["#00a651", "#ffcc00", "#e74c3c", "#7d3c98", "#0d1b5c"],
+        N=256,
+    )
+
+
+def images01_to_log_nhi(img01: np.ndarray, lo: float = LO_LOG, hi: float = HI_LOG) -> np.ndarray:
+    """Maps in [0,1] linear in column density → log10(N_HI/cm^-2)."""
+    return lo + (hi - lo) * np.clip(img01, 0.0, 1.0).astype(np.float64)
+
+
+def per_map_power_spectra_log(
+    images_01: np.ndarray, box_size: float = 25.0, lo: float = LO_LOG, hi: float = HI_LOG
+) -> tuple[np.ndarray, np.ndarray]:
+    """Return (dk, Pk) with Pk shape (N, n_bins) using log10 N_HI field."""
+    logf = images01_to_log_nhi(images_01, lo, hi)
+    n = logf.shape[0]
+    npix = logf.shape[-1]
+    dl = box_size / npix
+    dk, _ = ec.PowerSpectrum(logf[0], N=npix, dl=dl)
+    pks = np.stack([ec.PowerSpectrum(logf[i], N=npix, dl=dl)[1] for i in range(n)])
+    return dk, pks
+
+
+def sample_batch(
+    model: torch.nn.Module,
+    labels_np: np.ndarray,
+    label_mean: np.ndarray,
+    label_std: np.ndarray,
+    normalize_labels: bool,
+    height: int,
+    width: int,
+    device: torch.device,
+    ddim_steps: int,
+    progress: bool,
+) -> np.ndarray:
+    """DDIM sample batch; labels_np shape (B, label_dim). mean/std same length as label_dim."""
+    labels_np = np.asarray(labels_np, dtype=np.float32)
+    mean = np.asarray(label_mean, dtype=np.float32)
+    std = np.asarray(label_std, dtype=np.float32)
+    if normalize_labels:
+        t = ec.prepare_labels_for_model(labels_np, mean, std).to(device)
+    else:
+        t = torch.from_numpy(labels_np).float().to(device)
+    with torch.no_grad():
+        out = model.sample(
+            labels=t,
+            channels=1,
+            height=height,
+            width=width,
+            device=device,
+            progress=progress,
+            use_ddim=True,
+            ddim_steps=ddim_steps,
+        )
+    return ec.from_model_output(out)

src/figure9_posterior.py

@@ -0,0 +1,33 @@
+"""
+Figure-9 style surrogate posteriors: build (Ωm, σ8) grids and log P(k) for observed maps.
+
+Used by ddpm_cond_eval.ipynb. Sampling and P(k) live in eval_model.py.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+import eval_model as em
+
+
+def build_cosmo_grid(
+    g: int,
+    om_lo: float,
+    om_hi: float,
+    s8_lo: float,
+    s8_hi: float,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    om_axis = np.linspace(om_lo, om_hi, g, dtype=np.float64)
+    s8_axis = np.linspace(s8_lo, s8_hi, g, dtype=np.float64)
+    og, sg = np.meshgrid(om_axis, s8_axis, indexing="ij")
+    grid_labels = np.stack([og.ravel(), sg.ravel()], axis=1).astype(np.float32)
+    return om_axis, s8_axis, og, sg, grid_labels
+
+
+def log_pk_observed(img01: np.ndarray, box_size: float, dk: np.ndarray) -> np.ndarray:
+    """Single map → log P(k) on bins where dk > 0."""
+    _, pk = em.per_map_power_spectra_log(img01[np.newaxis, ...], box_size)
+    valid = dk > 0
+    if pk.shape[1] != len(dk):
+        raise ValueError("P(k) bin count mismatch vs dk")
+    return np.log(pk[0, valid] + 1e-30)

src/plot_r2_cosmology_lhs.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python3
+"""
+Reproduce the Latin-hypercube R² figure (μ(P) and σ(P)) in (Ωm, σ8) with a layout
+that avoids colorbar / suptitle overlap.
+
+Usage (full run — slow):
+  python plot_r2_cosmology_lhs.py --output ddpm_eval_notebook_out/r2_cosmology_lhs50_ddpm.png
+
+Replay plot only from saved arrays:
+  python plot_r2_cosmology_lhs.py --from-npz r2_lhs_data.npz --output out.png
+
+Defaults match ddpm_conditional / evaluate_conditional 6-param setup.
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.cm import ScalarMappable
+from matplotlib.colors import Normalize
+from matplotlib.gridspec import GridSpec
+
+import evaluate_conditional as ec
+import eval_model as em
+
+_SCRIPT_DIR = Path(__file__).resolve().parent
+_DEFAULT_CKPT = _SCRIPT_DIR / "outputs_conditional_6param_20260413_132226/checkpoints/best_model.pt"
+_DEFAULT_DATA = "/scratch/mrpcol001/Diffusion_job/data/LH_data/params_6"
+
+
+def latin_hypercube_scaled(
+    n: int, lo: np.ndarray, hi: np.ndarray, rng: np.random.Generator
+) -> np.ndarray:
+    """n points in [lo, hi] per dimension (classic LHS)."""
+    d = int(lo.shape[0])
+    u = rng.random((n, d))
+    cut = np.linspace(0.0, 1.0, n + 1)
+    a, b = cut[:-1], cut[1:]
+    width = (b - a)[:, np.newaxis]
+    rd = a[:, np.newaxis] + u * width
+    for j in range(d):
+        rng.shuffle(rd[:, j])
+    span = (hi - lo).astype(np.float64)
+    return (lo + rd * span).astype(np.float32)
+
+
+def compute_lhs_r2(
+    model: torch.nn.Module,
+    images_split: np.ndarray,
+    labels_split: np.ndarray,
+    label_mean: np.ndarray,
+    label_std: np.ndarray,
+    device: torch.device,
+    lhs_n: int,
+    maps_per_point: int,
+    batch_size: int,
+    box_size_mpc: float,
+    ddim_steps: int,
+    seed: int,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Returns lhs_pts, r2_mu_arr, r2_sig_arr, lo_b, hi_b."""
+    lo_b = labels_split.min(axis=0)
+    hi_b = labels_split.max(axis=0)
+    rng = np.random.default_rng(seed)
+    lhs_pts = latin_hypercube_scaled(lhs_n, lo_b, hi_b, rng)
+
+    ldim = labels_split.shape[1]
+    h, w = int(images_split.shape[-2]), int(images_split.shape[-1])
+    bs = min(batch_size, maps_per_point)
+    npix = int(images_split.shape[-1])
+    dl = box_size_mpc / npix
+
+    def pk_stack(imgs: np.ndarray) -> np.ndarray:
+        return np.stack([ec.PowerSpectrum(im, N=npix, dl=dl)[1] for im in imgs], axis=0)
+
+    r2_mu_arr = np.full(lhs_n, np.nan, dtype=np.float64)
+    r2_sig_arr = np.full(lhs_n, np.nan, dtype=np.float64)
+    model.eval()
+
+    for ti in range(lhs_n):
+        theta = lhs_pts[ti]
+        dist = np.linalg.norm(labels_split - theta, axis=1)
+        nn_idx = np.argsort(dist)[:maps_per_point]
+        real_batch = images_split[nn_idx]
+        rep = np.tile(theta[None, :], (maps_per_point, 1))
+        gen_chunks = []
+        for j in range(0, maps_per_point, bs):
+            chunk = rep[j : j + bs]
+            bt = ec.prepare_labels_for_model(chunk, label_mean, label_std).to(device)
+            with torch.no_grad():
+                g = model.sample(
+                    labels=bt,
+                    channels=1,
+                    height=h,
+                    width=w,
+                    device=device,
+                    progress=False,
+                    use_ddim=True,
+                    ddim_steps=ddim_steps,
+                )
+            gen_chunks.append(ec.from_model_output(g))
+        gen_batch = np.concatenate(gen_chunks, axis=0)
+
+        pk_r = pk_stack(real_batch)
+        pk_g = pk_stack(gen_batch)
+        km = np.arange(pk_r.shape[1], dtype=int) > 0
+        mu_r, mu_g = pk_r.mean(axis=0), pk_g.mean(axis=0)
+        sr, sg = pk_r.std(axis=0), pk_g.std(axis=0)
+        r2_mu_arr[ti] = em.r2_score_1d(mu_r[km], mu_g[km])
+        r2_sig_arr[ti] = em.r2_score_1d(sr[km], sg[km])
+
+    return lhs_pts, r2_mu_arr, r2_sig_arr, lo_b, hi_b
+
+
+def plot_r2_cosmology_figure(
+    lhs_pts: np.ndarray,
+    r2_mu_arr: np.ndarray,
+    r2_sig_arr: np.ndarray,
+    lo_b: np.ndarray,
+    hi_b: np.ndarray,
+    out_path: Path,
+    r2_vmin: float = 0.90,
+    r2_vmax: float = 1.0,
+    lhs_n: int | None = None,
+    maps_per_point: int | None = None,
+    dpi: int = 160,
+) -> None:
+    """
+    Two-panel scatter in (Ωm, σ8) with a dedicated colorbar column (no overlap with heatmap).
+    """
+    lhs_n = lhs_n if lhs_n is not None else len(r2_mu_arr)
+    maps_per_point = maps_per_point if maps_per_point is not None else 15
+
+    ldim = lhs_pts.shape[1]
+    om_plot = lhs_pts[:, 0]
+    s8_plot = lhs_pts[:, 1] if ldim >= 2 else np.zeros(lhs_n)
+
+    cmap = em.cmap_r2_hiflow()
+    norm = Normalize(vmin=r2_vmin, vmax=r2_vmax)
+    sm = ScalarMappable(norm=norm, cmap=cmap)
+    sm.set_array([])
+
+    fig = plt.figure(figsize=(11.5, 4.9))
+    # Left: data panels; narrow right strip: colorbar only (avoids fig.colorbar + tight_layout clash)
+    gs = GridSpec(
+        nrows=1,
+        ncols=3,
+        figure=fig,
+        width_ratios=[1.0, 1.0, 0.065],
+        wspace=0.26,
+        left=0.07,
+        right=0.98,
+        top=0.82,
+        bottom=0.14,
+    )
+    ax0 = fig.add_subplot(gs[0, 0])
+    ax1 = fig.add_subplot(gs[0, 1], sharey=ax0)
+    cax = fig.add_subplot(gs[0, 2])
+
+    pad_x = 0.02 * (float(hi_b[0] - lo_b[0]) + 1e-6)
+    ax0.set_xlim(float(lo_b[0]) - pad_x, float(hi_b[0]) + pad_x)
+    ax1.set_xlim(float(lo_b[0]) - pad_x, float(hi_b[0]) + pad_x)
+    if ldim >= 2:
+        pad_y = 0.02 * (float(hi_b[1] - lo_b[1]) + 1e-6)
+        ax0.set_ylim(float(lo_b[1]) - pad_y, float(hi_b[1]) + pad_y)
+
+    for ax, r2v, subtitle in zip(
+        (ax0, ax1),
+        (r2_mu_arr, r2_sig_arr),
+        (r"$R^2$ for $\mu(P)$", r"$R^2$ for $\sigma(P)$"),
+    ):
+        ok = np.isfinite(r2v)
+        ax.scatter(
+            om_plot[ok],
+            s8_plot[ok],
+            c=np.clip(r2v[ok], r2_vmin, r2_vmax),
+            cmap=cmap,
+            norm=norm,
+            s=52,
+            alpha=0.92,
+            edgecolors="k",
+            linewidths=0.35,
+        )
+        ax.set_xlabel(r"$\Omega_m$", fontsize=12)
+        ax.set_title(subtitle, fontsize=11)
+        ax.grid(True, alpha=0.25)
+
+    ax0.set_ylabel(r"$\sigma_8$", fontsize=12)
+    plt.setp(ax1.get_yticklabels(), visible=False)
+
+    cb = fig.colorbar(sm, cax=cax)
+    cb.set_label(r"$R^2$", fontsize=11)
+    cax.tick_params(labelsize=9)
+
+    fig.suptitle(
+        r"Visual summary of $R^2$ (CAMELS vs conditional DDPM) vs cosmology — "
+        + f"{lhs_n} Latin Hypercube samples; {maps_per_point} maps / point",
+        fontsize=11,
+        fontweight="bold",
+        y=0.96,
+    )
+
+    out_path = Path(out_path)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    # Do not use bbox_inches="tight" — it rebalance axes and can squeeze the colorbar into the panels.
+    fig.savefig(out_path, dpi=dpi)
+    plt.close(fig)
+
+
+def _resolve_training_args(checkpoint: Path) -> Path | None:
+    run = checkpoint.parent.parent if checkpoint.parent.name == "checkpoints" else checkpoint.parent
+    for name in ("args.json", "args.txt"):
+        p = run / name
+        if p.is_file():
+            return p
+    return None
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="LHS R² cosmology figure (fixed colorbar layout)")
+    p.add_argument("--checkpoint", type=str, default=str(_DEFAULT_CKPT))
+    p.add_argument("--data-dir", type=str, default=_DEFAULT_DATA)
+    p.add_argument("--split", type=str, default="test", choices=("train", "val", "test"))
+    p.add_argument("--output", type=str, default=str(_SCRIPT_DIR / "ddpm_eval_notebook_out/r2_cosmology_lhs50_ddpm.png"))
+    p.add_argument("--from-npz", type=str, default=None, help="Load lhs_pts, r2_mu_arr, r2_sig_arr, lo_b, hi_b")
+    p.add_argument("--save-npz", type=str, default=None, help="After compute, save arrays for --from-npz replot")
+    p.add_argument("--lhs-n", type=int, default=50)
+    p.add_argument("--maps-per-point", type=int, default=15)
+    p.add_argument("--batch-size", type=int, default=8)
+    p.add_argument("--ddim-steps", type=int, default=50)
+    p.add_argument("--box-size-mpc", type=float, default=25.0)
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--r2-vmin", type=float, default=0.90)
+    p.add_argument("--r2-vmax", type=float, default=1.0)
+    p.add_argument("--dpi", type=int, default=160)
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    out_path = Path(args.output)
+
+    if args.from_npz:
+        z = np.load(args.from_npz, allow_pickle=False)
+        lhs_pts = z["lhs_pts"]
+        r2_mu_arr = z["r2_mu_arr"]
+        r2_sig_arr = z["r2_sig_arr"]
+        lo_b = z["lo_b"]
+        hi_b = z["hi_b"]
+    else:
+        ckpt = Path(args.checkpoint).expanduser().resolve()
+        if not ckpt.is_file():
+            raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
+
+        ta = _resolve_training_args(ckpt)
+        config: dict = {}
+        if ta is not None:
+            config = ec.load_training_config(str(ta))
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = ec.build_model(config, device)
+        ec.load_checkpoint(model, str(ckpt), device)
+
+        data_dir = Path(args.data_dir)
+        images_split, labels_split = ec.load_split(data_dir, args.split)
+        label_mean, label_std = ec.load_label_stats(data_dir)
+
+        lhs_pts, r2_mu_arr, r2_sig_arr, lo_b, hi_b = compute_lhs_r2(
+            model,
+            images_split,
+            labels_split,
+            label_mean,
+            label_std,
+            device,
+            lhs_n=args.lhs_n,
+            maps_per_point=args.maps_per_point,
+            batch_size=args.batch_size,
+            box_size_mpc=args.box_size_mpc,
+            ddim_steps=args.ddim_steps,
+            seed=args.seed,
+        )
+
+        if args.save_npz:
+            np.savez(
+                args.save_npz,
+                lhs_pts=lhs_pts,
+                r2_mu_arr=r2_mu_arr,
+                r2_sig_arr=r2_sig_arr,
+                lo_b=lo_b,
+                hi_b=hi_b,
+            )
+            print("Saved", args.save_npz)
+
+    plot_r2_cosmology_figure(
+        lhs_pts,
+        r2_mu_arr,
+        r2_sig_arr,
+        lo_b,
+        hi_b,
+        out_path,
+        r2_vmin=args.r2_vmin,
+        r2_vmax=args.r2_vmax,
+        lhs_n=args.lhs_n,
+        maps_per_point=args.maps_per_point,
+        dpi=args.dpi,
+    )
+    print("Saved", out_path.resolve())
+
+
+if __name__ == "__main__":
+    main()

src/posterior_inference.py

@@ -0,0 +1,895 @@
+#!/usr/bin/env python3
+"""
+posterior_inference.py — VLB-based cosmological inference (Mudur et al. 2023 §4 style).
+Pure inference-time; frozen DDPM weights. Script lives next to diffusion_conditional.py.
+"""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import json
+import os
+import sys
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.gridspec as gridspec
+import matplotlib.pyplot as plt
+import matplotlib.patheffects as mpathe
+import numpy as np
+import torch
+
+# ── Project imports ────────────────────────────────────────────────────────────
+_ROOT = Path(__file__).resolve().parent
+if (_ROOT / "diffusion_conditional.py").is_file():
+    sys.path.insert(0, str(_ROOT))
+
+from diffusion_conditional import GaussianDiffusion, ConditionalDiffusionModel
+from unet_conditional import ConditionalUNet
+
+plt.rcParams.update({
+    "figure.facecolor": "white", "axes.facecolor": "white",
+    "axes.edgecolor": "#222", "axes.linewidth": 0.7,
+    "axes.spines.top": False, "axes.spines.right": False,
+    "font.family": "DejaVu Sans", "font.size": 9.5,
+    "savefig.facecolor": "white",
+})
+
+REAL_COLOR = "#CC3333"
+GEN_COLOR = "#2266BB"
+SIGMA_LEVELS = [2.30, 6.17, 11.83]
+SIGMA_COLORS = ["#1a5c9e", "#5590d0", "#99c0ea"]
+SIGMA_LABELS = {2.30: r"$1\sigma$", 6.17: r"$2\sigma$", 11.83: r"$3\sigma$"}
+
+
+def load_config(path: str) -> Dict:
+    p = Path(path)
+    if p.suffix == ".json":
+        with open(p) as f:
+            return json.load(f)
+    cfg = {}
+    with open(p) as f:
+        for line in f:
+            if ":" not in line:
+                continue
+            k, v = line.strip().split(":", 1)
+            try:
+                cfg[k.strip()] = ast.literal_eval(v.strip())
+            except Exception:
+                cfg[k.strip()] = v.strip()
+    return cfg
+
+
+def load_model(ckpt: str, cfg: Dict, device: torch.device) -> ConditionalDiffusionModel:
+    unet = ConditionalUNet(
+        in_channels=1, out_channels=1,
+        label_dim=int(cfg.get("label_dim", 2)),
+        base_channels=int(cfg.get("base_channels", 64)),
+        channel_multipliers=list(cfg.get("channel_multipliers", [1, 2, 4, 8])),
+        attention_levels=list(cfg.get("attention_levels", [2, 3])),
+        dropout=float(cfg.get("dropout", 0.1)),
+    )
+    diff = GaussianDiffusion(
+        timesteps=int(cfg.get("timesteps", 1500)),
+        beta_start=float(cfg.get("beta_start", 1e-4)),
+        beta_end=float(cfg.get("beta_end", 0.02)),
+        schedule_type=str(cfg.get("schedule_type", "linear")),
+    )
+    model = ConditionalDiffusionModel(unet, diff).to(device)
+    ck = torch.load(ckpt, map_location=device, weights_only=False)
+    if isinstance(ck, dict) and "ema_shadow" in ck:
+        cur = model.state_dict()
+        for k, v in ck["ema_shadow"].items():
+            if k in cur:
+                cur[k] = v
+        model.load_state_dict(cur)
+        print("  Loaded EMA weights")
+    elif isinstance(ck, dict) and "model_state_dict" in ck:
+        model.load_state_dict(ck["model_state_dict"])
+    else:
+        model.load_state_dict(ck)
+    model.eval()
+    for p in model.parameters():
+        p.requires_grad_(False)
+    return model
+
+
+def load_test_data(
+    data_dir: str, n_fields: int, seed: int = 42
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    dp = Path(data_dir)
+    lsuf = "_2" if (dp / "train_labels_LH_2.npy").exists() else ""
+    isuf = "" if (dp / "train_LH.npy").exists() else "_6"
+
+    imgs = np.load(dp / f"test_LH{isuf}.npy").astype(np.float32)
+    labels = np.load(dp / f"test_labels_LH{lsuf}.npy").astype(np.float32)
+    tr_lab = np.load(dp / f"train_labels_LH{lsuf}.npy").astype(np.float32)
+
+    rng = np.random.default_rng(seed)
+    idx = rng.choice(len(imgs), n_fields, replace=False)
+
+    label_mu = tr_lab.mean(0)
+    label_std = np.where(tr_lab.std(0) == 0, 1.0, tr_lab.std(0))
+    return imgs[idx], labels[idx], label_mu, label_std
+
+
+def normal_kl(mean1, log_var1, mean2, log_var2):
+    return 0.5 * (
+        -1.0 + log_var2 - log_var1
+        + torch.exp(log_var1 - log_var2)
+        + ((mean1 - mean2) ** 2) * torch.exp(-log_var2)
+    )
+
+
+def _approx_standard_normal_cdf(x):
+    return 0.5 * (1.0 + torch.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x ** 3)))
+
+
+def discretised_gaussian_log_likelihood(x_0, mean, log_var):
+    centered_x = x_0 - mean
+    inv_stdv = torch.exp(-0.5 * log_var)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_plus = _approx_standard_normal_cdf(plus_in)
+    cdf_min = _approx_standard_normal_cdf(min_in)
+    log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = torch.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = (cdf_plus - cdf_min).clamp(min=1e-12)
+    log_probs = torch.where(
+        x_0 < -0.999,
+        log_cdf_plus,
+        torch.where(x_0 > 0.999, log_one_minus_cdf_min, torch.log(cdf_delta)),
+    )
+    return log_probs
+
+
+def predict_x_start_from_eps(diff: GaussianDiffusion, x_t: torch.Tensor,
+                             t: torch.Tensor, eps: torch.Tensor) -> torch.Tensor:
+    # Matches GaussianDiffusion._predict_xstart_from_noise (diffusion_conditional.py)
+    return (
+        diff._extract(diff.recip_sqrt_alphas_cumprod, t, x_t.shape) * x_t
+        - diff._extract(diff.sqrt_recip_minus_one, t, x_t.shape) * eps
+    )
+
+
+def q_posterior_mean_var(diff: GaussianDiffusion, x_start: torch.Tensor,
+                         x_t: torch.Tensor, t: torch.Tensor):
+    mean = (
+        diff._extract(diff.posterior_mean_coef1, t, x_t.shape) * x_start
+        + diff._extract(diff.posterior_mean_coef2, t, x_t.shape) * x_t
+    )
+    var = diff._extract(diff.posterior_variance, t, x_t.shape)
+    log_var_c = diff._extract(diff.posterior_log_variance_clipped, t, x_t.shape)
+    return mean, var, log_var_c
+
+
+@torch.no_grad()
+def compute_L_t(
+    model: ConditionalDiffusionModel,
+    x_0: torch.Tensor,
+    labels_n: torch.Tensor,
+    t: int,
+    fixed_eps: torch.Tensor,
+) -> torch.Tensor:
+    diff = model.diffusion
+    device = x_0.device
+    B = x_0.shape[0]
+    t_vec = torch.full((B,), t, device=device, dtype=torch.long)
+
+    if t == 0:
+        t1 = torch.full((B,), 1, device=device, dtype=torch.long)
+        ab1 = diff._extract(diff.alphas_cumprod, t1, x_0.shape)
+        x_1 = torch.sqrt(ab1) * x_0 + torch.sqrt(1.0 - ab1) * fixed_eps
+        eps_pred = model(x_1, t1, labels_n)
+        x_start_pred = predict_x_start_from_eps(diff, x_1, t1, eps_pred).clamp(-1, 1)
+        mean, _, log_var = q_posterior_mean_var(diff, x_start_pred, x_1, t1)
+        log_p = discretised_gaussian_log_likelihood(x_0, mean, log_var)
+        return -log_p.sum(dim=(1, 2, 3))
+
+    ab_t = diff._extract(diff.alphas_cumprod, t_vec, x_0.shape)
+    x_t = torch.sqrt(ab_t) * x_0 + torch.sqrt(1.0 - ab_t) * fixed_eps
+    true_mean, _, true_log_var = q_posterior_mean_var(diff, x_0, x_t, t_vec)
+    eps_pred = model(x_t, t_vec, labels_n)
+    x_start_pred = predict_x_start_from_eps(diff, x_t, t_vec, eps_pred).clamp(-1, 1)
+    model_mean, _, model_log_var = q_posterior_mean_var(diff, x_start_pred, x_t, t_vec)
+    kl = normal_kl(true_mean, true_log_var, model_mean, model_log_var)
+    return kl.sum(dim=(1, 2, 3))
+
+
+@torch.no_grad()
+def compute_L_T_analytic(diff: GaussianDiffusion, x_0: torch.Tensor) -> torch.Tensor:
+    T = diff.timesteps
+    t_vec = torch.full((x_0.shape[0],), T - 1, device=x_0.device, dtype=torch.long)
+    abar_T = diff._extract(diff.alphas_cumprod, t_vec, x_0.shape)
+    mean1 = torch.sqrt(abar_T) * x_0
+    log_var1 = torch.log((1.0 - abar_T).clamp(min=1e-30))
+    kl = normal_kl(mean1, log_var1, torch.zeros_like(mean1), torch.zeros_like(log_var1))
+    return kl.sum(dim=(1, 2, 3))
+
+
+def build_eval_grid(
+    Om_true: float,
+    s8_true: float,
+    grid_size: int,
+    span: float = 0.1,
+    Om_range: Tuple[float, float] = (0.10, 0.50),
+    s8_range: Tuple[float, float] = (0.60, 1.00),
+) -> Tuple[np.ndarray, np.ndarray]:
+    Om_lo = max(Om_true - span, Om_range[0])
+    Om_hi = min(Om_true + span, Om_range[1])
+    s8_lo = max(s8_true - span, s8_range[0])
+    s8_hi = min(s8_true + span, s8_range[1])
+    Om_1d = np.linspace(Om_lo, Om_hi, grid_size)
+    s8_1d = np.linspace(s8_lo, s8_hi, grid_size)
+    return Om_1d, s8_1d
+
+
+@torch.no_grad()
+def evaluate_vlb_surface(
+    model: ConditionalDiffusionModel,
+    x_0: torch.Tensor,
+    Om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    label_mu: np.ndarray,
+    label_std: np.ndarray,
+    t_values: List[int],
+    n_seeds: int = 4,
+    batch_size: int = 32,
+    label_dim: int = 2,
+    fixed_seed: int = 0,
+    device: Optional[torch.device] = None,
+) -> Dict[int, np.ndarray]:
+    device = device or x_0.device
+    nO, nS = len(Om_grid), len(s8_grid)
+    n_pts = nO * nS
+
+    Omg, s8g = np.meshgrid(Om_grid, s8_grid, indexing="ij")
+    raw_labels = np.column_stack([Omg.ravel(), s8g.ravel()])
+    if label_dim > 2:
+        pad = np.zeros((n_pts, label_dim - 2), dtype=np.float32)
+        for i in range(label_dim - 2):
+            pad[:, i] = label_mu[2 + i]
+        raw_labels = np.concatenate([raw_labels, pad], axis=1)
+    norm_labels = (raw_labels - label_mu) / label_std
+    norm_labels_t = torch.from_numpy(norm_labels.astype(np.float32)).to(device)
+
+    L_surfaces = {t: np.zeros(n_pts, dtype=np.float64) for t in t_values}
+
+    H, W = x_0.shape[-2], x_0.shape[-1]
+    rng_torch = torch.Generator(device=device).manual_seed(fixed_seed)
+    seeds_eps = [
+        torch.randn(1, 1, H, W, generator=rng_torch, device=device)
+        for _ in range(n_seeds)
+    ]
+
+    for _, fixed_eps in enumerate(seeds_eps):
+        for t in t_values:
+            for start in range(0, n_pts, batch_size):
+                end = min(start + batch_size, n_pts)
+                bsz = end - start
+                x_b = x_0.expand(bsz, -1, -1, -1)
+                lbl_b = norm_labels_t[start:end]
+                eps_b = fixed_eps.expand(bsz, -1, -1, -1)
+                L_t = compute_L_t(model, x_b, lbl_b, t=t, fixed_eps=eps_b)
+                L_surfaces[t][start:end] += L_t.cpu().numpy() / n_seeds
+
+    return {t: L_surfaces[t].reshape(nO, nS) for t in t_values}
+
+
+def marginal_from_neg2dL(
+    neg2dL: np.ndarray, Om_grid: np.ndarray, s8_grid: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray, Tuple[float, float]]:
+    L = -0.5 * neg2dL
+    L = L - L.max()
+    P = np.exp(L)
+    Om_marginal = P.sum(axis=1)
+    Om_marginal /= Om_marginal.sum()
+    s8_marginal = P.sum(axis=0)
+    s8_marginal /= s8_marginal.sum()
+    Om_pred = float(Om_grid[np.argmax(Om_marginal)])
+    s8_pred = float(s8_grid[np.argmax(s8_marginal)])
+    return Om_marginal, s8_marginal, (Om_pred, s8_pred)
+
+
+def credible_interval_68(values: np.ndarray, probs: np.ndarray) -> Tuple[float, float, float]:
+    cdf = np.cumsum(probs)
+    cdf /= cdf[-1]
+    median = float(np.interp(0.50, cdf, values))
+    lo = float(np.interp(0.16, cdf, values))
+    hi = float(np.interp(0.84, cdf, values))
+    return median, lo, hi
+
+
+def fig_contours_per_t(
+    surfaces: Dict[int, np.ndarray],
+    Om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    Om_true: float,
+    s8_true: float,
+    out_path: Path,
+    dpi: int = 200,
+) -> None:
+    fig, ax = plt.subplots(figsize=(7, 6.5), dpi=dpi)
+    cmap = plt.cm.viridis
+    n_t = len(surfaces)
+    colors = cmap(np.linspace(0.05, 0.95, n_t))
+
+    for (t, L_surf), col in zip(sorted(surfaces.items()), colors):
+        neg2dL = 2.0 * (L_surf - L_surf.min())
+        ax.contour(
+            Om_grid, s8_grid, neg2dL.T,
+            levels=[2.30], colors=[col], linewidths=[1.6], linestyles=["-"],
+        )
+        ax.plot([], [], color=col, lw=1.8, label=f"t={t}")
+
+    ax.plot(Om_true, s8_true, "r+", ms=18, mew=2.5, label="True", zorder=10)
+    ax.set_xlabel(r"$\Omega_m$", fontsize=12)
+    ax.set_ylabel(r"$\sigma_8$", fontsize=12)
+    ax.set_title(
+        r"$-2\Delta\ln\hat{L}_t$ — $1\sigma$ contour per timestep"
+        "\n(Mudur-style)  smaller $t$ → tighter constraint",
+        fontweight="bold", fontsize=10,
+    )
+    ax.legend(fontsize=8, loc="best", ncol=2, framealpha=0.92)
+    ax.grid(alpha=0.18)
+    ax.set_xlim(Om_grid[0], Om_grid[-1])
+    ax.set_ylim(s8_grid[0], s8_grid[-1])
+    fig.savefig(out_path, bbox_inches="tight", dpi=dpi)
+    plt.close(fig)
+    print(f"  Saved -> {out_path}")
+
+
+def _L0_posterior_smoothed(
+    L0_surface: np.ndarray, smooth_sigma: float = 0.6,
+):
+    from scipy.ndimage import gaussian_filter as gf
+
+    neg2dL = 2.0 * (L0_surface - L0_surface.min())
+    surface_sm = gf(neg2dL, sigma=smooth_sigma)
+    return surface_sm, neg2dL
+
+
+def draw_L0_posterior_main_panel(
+    ax,
+    surface_sm: np.ndarray,
+    Om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    Om_true: float,
+    s8_true: float,
+    Om_pred: float,
+    s8_pred: float,
+    *,
+    clabel_fontsize: float = 9.5,
+    marker_ms: float = 16,
+) -> None:
+    ax.contourf(Om_grid, s8_grid, surface_sm.T, levels=60, cmap="Blues_r",
+                vmin=0, vmax=SIGMA_LEVELS[-1] * 3, extend="max", alpha=0.55)
+    for lv, co in zip(reversed(SIGMA_LEVELS), reversed(SIGMA_COLORS)):
+        ax.contourf(Om_grid, s8_grid, surface_sm.T,
+                    levels=[0, lv], colors=[co], alpha=0.78)
+
+    cs = ax.contour(
+        Om_grid, s8_grid, surface_sm.T,
+        levels=SIGMA_LEVELS,
+        colors=["white", "white", "white"],
+        linewidths=[2.2, 1.6, 1.2],
+        linestyles=["-", "--", "-."],
+    )
+    ax.clabel(cs, fmt=SIGMA_LABELS, inline=True, fontsize=clabel_fontsize, colors="white")
+
+    ax.axvline(Om_true, color="red", lw=0.7, ls=":", alpha=0.6)
+    ax.axhline(s8_true, color="red", lw=0.7, ls=":", alpha=0.6)
+    ax.plot(Om_true, s8_true, "r+", ms=marker_ms, mew=2.5, zorder=6, label="True")
+    ax.plot(Om_pred, s8_pred, "w^", ms=max(6, marker_ms * 0.55), mew=1.2, zorder=6, label="MAP")
+    ax.set_xlim(Om_grid[0], Om_grid[-1])
+    ax.set_ylim(s8_grid[0], s8_grid[-1])
+    ax.grid(alpha=0.18)
+
+
+def fig_posterior_L0_mosaic_3x3(
+    out_dir: Path,
+    n_fields: int,
+    out_path: Path,
+    mosaic_side_px: int = 10_000,
+    panel_inches: float = 4.0,
+) -> None:
+    from matplotlib.patches import Patch
+
+    n_plot = min(n_fields, 9)
+    fig_side = panel_inches * 3
+    dpi = mosaic_side_px / fig_side
+    fig, axes = plt.subplots(
+        3, 3, figsize=(fig_side, fig_side), dpi=dpi,
+        squeeze=False,
+    )
+    for idx in range(9):
+        r, c = divmod(idx, 3)
+        ax = axes[r][c]
+        if idx >= n_plot:
+            ax.set_visible(False)
+            continue
+        nz = np.load(out_dir / f"field{idx:02d}_surfaces.npz")
+        L0 = np.asarray(nz["L_t0"])
+        Om_grid = np.asarray(nz["Om_grid"])
+        s8_grid = np.asarray(nz["s8_grid"])
+        Om_true = float(nz["Om_true"])
+        s8_true = float(nz["s8_true"])
+        surface_sm, _ = _L0_posterior_smoothed(L0, smooth_sigma=0.6)
+        _, _, (Om_pred, s8_pred) = marginal_from_neg2dL(surface_sm, Om_grid, s8_grid)
+        draw_L0_posterior_main_panel(
+            ax, surface_sm, Om_grid, s8_grid, Om_true, s8_true, Om_pred, s8_pred,
+            clabel_fontsize=7.0, marker_ms=11,
+        )
+        if idx == 0:
+            legend_patches = [
+                Patch(facecolor=SIGMA_COLORS[0], label=r"$1\sigma$"),
+                Patch(facecolor=SIGMA_COLORS[1], label=r"$2\sigma$"),
+                Patch(facecolor=SIGMA_COLORS[2], label=r"$3\sigma$"),
+            ]
+            hs, ls_ = ax.get_legend_handles_labels()
+            ax.legend(
+                handles=legend_patches + hs,
+                labels=[p.get_label() for p in legend_patches] + ls_,
+                fontsize=6, loc="upper right", framealpha=0.9,
+            )
+        else:
+            leg = ax.get_legend()
+            if leg is not None:
+                leg.remove()
+        ax.set_title(
+            rf"field {idx}: $\Omega_m^{{\rm true}}={Om_true:.3f}$, $\sigma_8^{{\rm true}}={s8_true:.3f}$",
+            fontsize=8,
+        )
+        ax.set_xlabel(r"$\Omega_m$", fontsize=8)
+        ax.set_ylabel(r"$\sigma_8$", fontsize=8)
+
+    fig.suptitle(
+        r"VLB $L_0$ posterior (2D) — 9 test fields",
+        fontsize=11, fontweight="bold", y=0.995,
+    )
+    fig.savefig(out_path, bbox_inches="tight", dpi=dpi)
+    plt.close(fig)
+    print(f"  Saved -> {out_path}  (≈ {mosaic_side_px}×{mosaic_side_px} px)")
+
+
+def fig_main_posterior(
+    L0_surface: np.ndarray,
+    Om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    Om_true: float,
+    s8_true: float,
+    out_path: Path,
+    dpi: int = 200,
+):
+    from matplotlib.patches import Patch
+
+    surface_sm, _ = _L0_posterior_smoothed(L0_surface, smooth_sigma=0.6)
+
+    Om_marg, s8_marg, (Om_pred, s8_pred) = marginal_from_neg2dL(
+        surface_sm, Om_grid, s8_grid
+    )
+    Om_med, Om_lo, Om_hi = credible_interval_68(Om_grid, Om_marg)
+    s8_med, s8_lo, s8_hi = credible_interval_68(s8_grid, s8_marg)
+
+    fig = plt.figure(figsize=(8.5, 8.5), dpi=dpi)
+    gs = gridspec.GridSpec(2, 2, width_ratios=[4, 1], height_ratios=[1, 4],
+                           hspace=0.05, wspace=0.05,
+                           left=0.10, right=0.95, top=0.95, bottom=0.08)
+    ax_main = fig.add_subplot(gs[1, 0])
+    ax_top = fig.add_subplot(gs[0, 0], sharex=ax_main)
+    ax_rt = fig.add_subplot(gs[1, 1], sharey=ax_main)
+
+    draw_L0_posterior_main_panel(
+        ax_main, surface_sm, Om_grid, s8_grid, Om_true, s8_true, Om_pred, s8_pred,
+    )
+
+    ax_main.set_xlabel(r"$\Omega_m$", fontsize=11)
+    ax_main.set_ylabel(r"$\sigma_8$", fontsize=11)
+
+    ax_top.fill_between(Om_grid, 0, Om_marg, color=SIGMA_COLORS[1], alpha=0.6)
+    ax_top.plot(Om_grid, Om_marg, color=SIGMA_COLORS[0], lw=1.4)
+    ax_top.axvline(Om_true, color="red", lw=1.0, ls=":")
+    ax_top.axvline(Om_pred, color="white", lw=1.5, ls="--",
+                   path_effects=[mpathe.withStroke(linewidth=2.5, foreground="black")])
+    ax_top.axvspan(Om_lo, Om_hi, color=SIGMA_COLORS[0], alpha=0.18, label=r"68% CI")
+    ax_top.set_yticks([])
+    ax_top.tick_params(labelbottom=False)
+    ax_top.set_title(
+        rf"$\Omega_m={Om_med:.3f}^{{+{Om_hi-Om_med:.3f}}}_{{-{Om_med-Om_lo:.3f}}}$"
+        rf"   (true: {Om_true:.3f})",
+        fontsize=9,
+    )
+
+    ax_rt.fill_betweenx(s8_grid, 0, s8_marg, color=SIGMA_COLORS[1], alpha=0.6)
+    ax_rt.plot(s8_marg, s8_grid, color=SIGMA_COLORS[0], lw=1.4)
+    ax_rt.axhline(s8_true, color="red", lw=1.0, ls=":")
+    ax_rt.axhline(s8_pred, color="white", lw=1.5, ls="--",
+                  path_effects=[mpathe.withStroke(linewidth=2.5, foreground="black")])
+    ax_rt.axhspan(s8_lo, s8_hi, color=SIGMA_COLORS[0], alpha=0.18)
+    ax_rt.set_xticks([])
+    ax_rt.tick_params(labelleft=False)
+    ax_rt.set_ylabel(
+        rf"$\sigma_8={s8_med:.3f}^{{+{s8_hi-s8_med:.3f}}}_{{-{s8_med-s8_lo:.3f}}}$"
+        rf"  (true: {s8_true:.3f})",
+        fontsize=9, rotation=270, labelpad=15,
+    )
+    ax_rt.yaxis.set_label_position("right")
+
+    legend_patches = [
+        Patch(facecolor=SIGMA_COLORS[0], label=r"$1\sigma$"),
+        Patch(facecolor=SIGMA_COLORS[1], label=r"$2\sigma$"),
+        Patch(facecolor=SIGMA_COLORS[2], label=r"$3\sigma$"),
+    ]
+    hs, ls_ = ax_main.get_legend_handles_labels()
+    ax_main.legend(
+        handles=legend_patches + hs,
+        labels=[p.get_label() for p in legend_patches] + ls_,
+        fontsize=8, loc="upper right", framealpha=0.92,
+    )
+
+    fig.suptitle(
+        r"VLB Posterior using $L_0$ — joint and marginal distributions",
+        fontsize=10, fontweight="bold", y=0.99,
+    )
+    fig.savefig(out_path, bbox_inches="tight", dpi=dpi)
+    plt.close(fig)
+    print(f"  Saved -> {out_path}")
+    return Om_pred, s8_pred, (Om_lo, Om_hi), (s8_lo, s8_hi)
+
+
+def fig_pred_vs_true(pred_results: List[Dict], out_path: Path, dpi: int = 200) -> None:
+    Om_true = np.array([r["Om_true"] for r in pred_results])
+    s8_true = np.array([r["s8_true"] for r in pred_results])
+    Om_pred = np.array([r["Om_pred"] for r in pred_results])
+    s8_pred = np.array([r["s8_pred"] for r in pred_results])
+
+    Om_err_lo = np.array([r["Om_pred"] - r["Om_lo"] for r in pred_results])
+    Om_err_hi = np.array([r["Om_hi"] - r["Om_pred"] for r in pred_results])
+    s8_err_lo = np.array([r["s8_pred"] - r["s8_lo"] for r in pred_results])
+    s8_err_hi = np.array([r["s8_hi"] - r["s8_pred"] for r in pred_results])
+
+    rmse_Om = np.sqrt(((Om_pred - Om_true) ** 2).mean())
+    rmse_s8 = np.sqrt(((s8_pred - s8_true) ** 2).mean())
+
+    fig, axes = plt.subplots(1, 2, figsize=(11, 5), dpi=dpi)
+
+    for ax, (true, pred, err_lo, err_hi, name, prange, rmse) in zip(axes, [
+        (Om_true, Om_pred, Om_err_lo, Om_err_hi, r"$\Omega_m$", (0.10, 0.50), rmse_Om),
+        (s8_true, s8_pred, s8_err_lo, s8_err_hi, r"$\sigma_8$", (0.60, 1.00), rmse_s8),
+    ]):
+        ax.errorbar(
+            true, pred, yerr=[np.maximum(err_lo, 0), np.maximum(err_hi, 0)],
+            fmt="o", color=GEN_COLOR, ecolor=SIGMA_COLORS[1],
+            elinewidth=1.2, capsize=3, ms=6,
+            label="DDPM-VLB inference (68% CI)",
+        )
+        ax.plot(prange, prange, "k--", lw=1.0, alpha=0.5, label="Identity")
+        ax.set_xlabel(f"True {name}", fontsize=11)
+        ax.set_ylabel(f"Predicted {name}", fontsize=11)
+        ax.set_xlim(*prange)
+        ax.set_ylim(*prange)
+        ax.grid(alpha=0.2)
+        ax.legend(fontsize=9, loc="lower right")
+        ax.text(
+            0.04, 0.92, f"RMSE = {rmse:.4f}",
+            transform=ax.transAxes, fontsize=10,
+            bbox=dict(facecolor="white", edgecolor="#ccc", alpha=0.92, pad=4),
+        )
+        ax.set_title(f"{name}: predicted vs true", fontweight="bold", fontsize=10)
+
+    fig.suptitle(
+        "VLB Parameter Inference: predicted vs true\n"
+        r"Error bars = 68% CI from $L_0$ marginal posterior",
+        fontsize=10, fontweight="bold", y=1.01,
+    )
+    plt.tight_layout()
+    fig.savefig(out_path, bbox_inches="tight", dpi=dpi)
+    plt.close(fig)
+    print(f"  Saved -> {out_path}")
+    print(f"    RMSE: Omega_m={rmse_Om:.4f}  sigma_8={rmse_s8:.4f}")
+
+
+def fig_posterior_and_contours_combined(
+    surfaces: Dict[int, np.ndarray],
+    L0_surface: np.ndarray,
+    Om_grid: np.ndarray,
+    s8_grid: np.ndarray,
+    Om_true: float,
+    s8_true: float,
+    out_path: Path,
+    dpi: int = 200,
+) -> Tuple[float, float, Tuple[float, float], Tuple[float, float]]:
+    """
+    Create a combined figure with contours_per_t on left and posterior on right.
+    """
+    from matplotlib.patches import Patch
+
+    surface_sm, _ = _L0_posterior_smoothed(L0_surface, smooth_sigma=0.6)
+    Om_marg, s8_marg, (Om_pred, s8_pred) = marginal_from_neg2dL(
+        surface_sm, Om_grid, s8_grid
+    )
+    Om_med, Om_lo, Om_hi = credible_interval_68(Om_grid, Om_marg)
+    s8_med, s8_lo, s8_hi = credible_interval_68(s8_grid, s8_marg)
+
+    fig = plt.figure(figsize=(16, 7), dpi=dpi)
+    gs = gridspec.GridSpec(2, 4, width_ratios=[4, 0.3, 4, 1], height_ratios=[1, 4],
+                           hspace=0.08, wspace=0.10,
+                           left=0.08, right=0.96, top=0.94, bottom=0.08)
+
+    # Left panel: Contours per timestep
+    ax_contours = fig.add_subplot(gs[:, 0])
+    cmap = plt.cm.viridis
+    n_t = len(surfaces)
+    colors = cmap(np.linspace(0.05, 0.95, n_t))
+
+    for (t, L_surf), col in zip(sorted(surfaces.items()), colors):
+        neg2dL = 2.0 * (L_surf - L_surf.min())
+        ax_contours.contour(
+            Om_grid, s8_grid, neg2dL.T,
+            levels=[2.30], colors=[col], linewidths=[1.6], linestyles=["-"],
+        )
+        ax_contours.plot([], [], color=col, lw=1.8, label=f"t={t}")
+
+    ax_contours.plot(Om_true, s8_true, "r+", ms=18, mew=2.5, label="True", zorder=10)
+    ax_contours.set_xlabel(r"$\Omega_m$", fontsize=12)
+    ax_contours.set_ylabel(r"$\sigma_8$", fontsize=12)
+    ax_contours.set_title(
+        r"$-2\Delta\ln\hat{L}_t$ — $1\sigma$ contours per timestep",
+        fontweight="bold", fontsize=11,
+    )
+    ax_contours.legend(fontsize=8, loc="best", ncol=1, framealpha=0.92)
+    ax_contours.grid(alpha=0.18)
+    ax_contours.set_xlim(Om_grid[0], Om_grid[-1])
+    ax_contours.set_ylim(s8_grid[0], s8_grid[-1])
+
+    # Right panel: Posterior L_0 (similar to fig_main_posterior layout)
+    ax_main = fig.add_subplot(gs[1, 2])
+    ax_top = fig.add_subplot(gs[0, 2], sharex=ax_main)
+    ax_rt = fig.add_subplot(gs[1, 3], sharey=ax_main)
+
+    draw_L0_posterior_main_panel(
+        ax_main, surface_sm, Om_grid, s8_grid, Om_true, s8_true, Om_pred, s8_pred,
+    )
+
+    ax_main.set_xlabel(r"$\Omega_m$", fontsize=11)
+    ax_main.set_ylabel(r"$\sigma_8$", fontsize=11)
+
+    ax_top.fill_between(Om_grid, 0, Om_marg, color=SIGMA_COLORS[1], alpha=0.6)
+    ax_top.plot(Om_grid, Om_marg, color=SIGMA_COLORS[0], lw=1.4)
+    ax_top.axvline(Om_true, color="red", lw=1.0, ls=":")
+    ax_top.axvline(Om_pred, color="white", lw=1.5, ls="--",
+                   path_effects=[mpathe.withStroke(linewidth=2.5, foreground="black")])
+    ax_top.axvspan(Om_lo, Om_hi, color=SIGMA_COLORS[0], alpha=0.18, label=r"68% CI")
+    ax_top.set_yticks([])
+    ax_top.tick_params(labelbottom=False)
+    ax_top.set_title(
+        rf"$\Omega_m={Om_med:.3f}^{{+{Om_hi-Om_med:.3f}}}_{{-{Om_med-Om_lo:.3f}}}$"
+        rf"   (true: {Om_true:.3f})",
+        fontsize=9,
+    )
+
+    ax_rt.fill_betweenx(s8_grid, 0, s8_marg, color=SIGMA_COLORS[1], alpha=0.6)
+    ax_rt.plot(s8_marg, s8_grid, color=SIGMA_COLORS[0], lw=1.4)
+    ax_rt.axhline(s8_true, color="red", lw=1.0, ls=":")
+    ax_rt.axhline(s8_pred, color="white", lw=1.5, ls="--",
+                  path_effects=[mpathe.withStroke(linewidth=2.5, foreground="black")])
+    ax_rt.axhspan(s8_lo, s8_hi, color=SIGMA_COLORS[0], alpha=0.18)
+    ax_rt.set_xticks([])
+    ax_rt.tick_params(labelleft=False)
+    ax_rt.set_ylabel(
+        rf"$\sigma_8={s8_med:.3f}^{{+{s8_hi-s8_med:.3f}}}_{{-{s8_med-s8_lo:.3f}}}$"
+        rf"  (true: {s8_true:.3f})",
+        fontsize=9, rotation=270, labelpad=15,
+    )
+    ax_rt.yaxis.set_label_position("right")
+
+    legend_patches = [
+        Patch(facecolor=SIGMA_COLORS[0], label=r"$1\sigma$"),
+        Patch(facecolor=SIGMA_COLORS[1], label=r"$2\sigma$"),
+        Patch(facecolor=SIGMA_COLORS[2], label=r"$3\sigma$"),
+    ]
+    hs, ls_ = ax_main.get_legend_handles_labels()
+    ax_main.legend(
+        handles=legend_patches + hs,
+        labels=[p.get_label() for p in legend_patches] + ls_,
+        fontsize=8, loc="upper right", framealpha=0.92,
+    )
+
+    fig.suptitle(
+        r"VLB Inference: $L_t$ contours (left) and $L_0$ posterior (right)",
+        fontsize=12, fontweight="bold", y=0.99,
+    )
+    fig.savefig(out_path, bbox_inches="tight", dpi=dpi)
+    plt.close(fig)
+    print(f"  Saved -> {out_path}")
+    return Om_pred, s8_pred, (Om_lo, Om_hi), (s8_lo, s8_hi)
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description="VLB-based parameter inference for trained conditional DDPM.",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--training_args", default=None)
+    p.add_argument("--data_dir", default="./data/params_2")
+    p.add_argument("--output_dir", default="vlb_inference_outputs")
+    p.add_argument("--n_fields", type=int, default=9)
+    p.add_argument(
+        "--grid_size", type=int, default=10_000,
+        help="Ωm×σ8 evaluation grid resolution (each side). Values > 300 require "
+             "--allow_huge_grid (very long runs for large grid_size).",
+    )
+    p.add_argument(
+        "--allow_huge_grid", action="store_true",
+        help="Required when --grid_size > 300 (avoids accidental multi-week GPU jobs).",
+    )
+    p.add_argument(
+        "--mosaic_side_px", type=int, default=10_000,
+        help="Pixel width/height of posterior_L0_mosaic_3x3.png (square).",
+    )
+    p.add_argument(
+        "--mosaic_panel_inches", type=float, default=4.0,
+        help="Matplotlib size (inches) of each 3×3 panel; dpi = mosaic_side_px / (3× this).",
+    )
+    p.add_argument("--span", type=float, default=0.10)
+    p.add_argument("--t_subset", type=int, nargs="+",
+                   default=[0, 1, 2, 5, 8, 10, 15, 20])
+    p.add_argument("--n_seeds", type=int, default=4)
+    p.add_argument("--batch_size", type=int, default=32)
+    p.add_argument("--device", default="auto")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument("--dpi", type=int, default=200)
+    return p.parse_args()
+
+
+def autodetect_args() -> Optional[str]:
+    for pat in ["outputs_conditional_*/args.json", "outputs_conditional_*/args.txt"]:
+        cands = sorted(Path(".").glob(pat), key=os.path.getctime, reverse=True)
+        if cands:
+            return str(cands[0])
+    return None
+
+
+def main() -> None:
+    args = parse_args()
+    if args.grid_size > 300 and not args.allow_huge_grid:
+        print(
+            "\nRefusing --grid_size={} (> 300) without --allow_huge_grid.\n"
+            "A 10_000×10_000 grid is roughly (200)^2 ≈ 40_000× more forward passes per\n"
+            "field than 50×50. For a quick run use e.g. --grid_size 50; for a high-res\n"
+            "summary figure use the default --mosaic_side_px without increasing grid_size.\n"
+            "To proceed anyway: add --allow_huge_grid\n".format(args.grid_size)
+        )
+        raise SystemExit(2)
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    device = (
+        torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if args.device == "auto"
+        else torch.device(args.device)
+    )
+    print(f"\nDevice: {device}")
+
+    out = Path(args.output_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    if args.training_args is None:
+        args.training_args = autodetect_args()
+        if args.training_args is None:
+            raise FileNotFoundError("Cannot find args.json — pass --training_args")
+        print(f"  Auto-detected args: {args.training_args}")
+
+    cfg = load_config(args.training_args)
+    print("\nLoading model ...")
+    model = load_model(args.checkpoint, cfg, device)
+    n_p = sum(p.numel() for p in model.parameters())
+    print(f"  Parameters: {n_p:,}  T={model.diffusion.timesteps}")
+
+    print(f"\nLoading {args.n_fields} test fields ...")
+    test_imgs, test_labels, label_mu, label_std = load_test_data(
+        args.data_dir, args.n_fields, seed=args.seed,
+    )
+    print(f"  Image shape: {test_imgs.shape[1:]}")
+    print(f"  Label dim:   {test_labels.shape[1]}")
+    print(f"  Label μ/σ:   {label_mu}  /  {label_std}")
+
+    print(
+        f"\nEvaluating L_t on {args.grid_size}x{args.grid_size} grid for "
+        f"{len(args.t_subset)} timesteps × {args.n_seeds} seeds ..."
+    )
+    print(
+        f"  -> {args.grid_size ** 2 * len(args.t_subset) * args.n_seeds:,} "
+        f"forward-pass groups per field (× seeds averaged)"
+    )
+
+    pred_results = []
+    label_dim = int(cfg.get("label_dim", 2))
+
+    for fi in range(args.n_fields):
+        Om_true = float(test_labels[fi, 0])
+        s8_true = float(test_labels[fi, 1])
+        print(f"\n  [{fi+1}/{args.n_fields}] field with "
+              f"Om={Om_true:.3f}, s8={s8_true:.3f}")
+
+        x_0 = torch.from_numpy(test_imgs[fi:fi + 1] * 2.0 - 1.0).unsqueeze(1).to(device)
+
+        Om_grid, s8_grid = build_eval_grid(Om_true, s8_true, args.grid_size, args.span)
+
+        t_start = time.time()
+        surfaces = evaluate_vlb_surface(
+            model=model,
+            x_0=x_0,
+            Om_grid=Om_grid,
+            s8_grid=s8_grid,
+            label_mu=label_mu,
+            label_std=label_std,
+            t_values=args.t_subset,
+            n_seeds=args.n_seeds,
+            batch_size=args.batch_size,
+            label_dim=label_dim,
+            fixed_seed=args.seed + fi,
+            device=device,
+        )
+        elapsed = time.time() - t_start
+        print(f"    Evaluation time: {elapsed:.1f}s")
+
+        np.savez(
+            out / f"field{fi:02d}_surfaces.npz",
+            **{f"L_t{t}": s for t, s in surfaces.items()},
+            Om_grid=Om_grid, s8_grid=s8_grid,
+            Om_true=Om_true, s8_true=s8_true,
+        )
+
+        if 0 in surfaces:
+            # Combined figure: contours_per_t + posterior on same plot
+            Om_pred, s8_pred, (Om_lo, Om_hi), (s8_lo, s8_hi) = fig_posterior_and_contours_combined(
+                surfaces, surfaces[0], Om_grid, s8_grid, Om_true, s8_true,
+                out / f"field{fi:02d}_combined.png", dpi=args.dpi,
+            )
+            pred_results.append(dict(
+                Om_true=Om_true, s8_true=s8_true,
+                Om_pred=Om_pred, s8_pred=s8_pred,
+                Om_lo=Om_lo, Om_hi=Om_hi,
+                s8_lo=s8_lo, s8_hi=s8_hi,
+            ))
+
+        # Also save individual figures for detailed inspection
+        fig_contours_per_t(
+            surfaces, Om_grid, s8_grid, Om_true, s8_true,
+            out / f"field{fi:02d}_contours_per_t.png", dpi=args.dpi,
+        )
+
+        if 0 in surfaces:
+            fig_main_posterior(
+                surfaces[0], Om_grid, s8_grid, Om_true, s8_true,
+                out / f"field{fi:02d}_posterior_L0.png", dpi=args.dpi,
+            )
+
+    if len(pred_results) >= 2:
+        fig_pred_vs_true(pred_results, out / "summary_pred_vs_true.png", dpi=args.dpi)
+        np.savez(
+            out / "summary.npz",
+            **{k: np.array([r[k] for r in pred_results])
+               for k in pred_results[0].keys()},
+        )
+
+    if args.n_fields >= 9 and all((out / f"field{i:02d}_surfaces.npz").is_file() for i in range(9)):
+        fig_posterior_L0_mosaic_3x3(
+            out, args.n_fields, out / "posterior_L0_mosaic_3x3.png",
+            mosaic_side_px=args.mosaic_side_px,
+            panel_inches=args.mosaic_panel_inches,
+        )
+
+    print(f"\nAll outputs -> {out.resolve()}/")
+    for f in sorted(out.glob("*.png")):
+        print(f"  {f.name}")
+
+
+if __name__ == "__main__":
+    main()

src/train_conditional.py

@@ -0,0 +1,447 @@
+"""
+Training script for conditional diffusion on CAMELS LH (6 cosmological parameters).
+
+Same training theory as DDPM_HI_Emulation_improved (2-label): DDPM noise prediction,
+DDIM sampling, ConditionalUNet with time + label embeddings, label z-score from train split,
+EMA, optional AMP, cosine LR, early stopping.
+
+Changes from original:
+- EMA weights are now applied before validation and sampling
+- Training args are saved to args.txt for evaluation script
+- Fixed --normalize_labels and --use_ddim flags (were un-disableable)
+- Added mixed-precision (AMP) training support
+- Fixed loss averaging to be per-sample rather than per-batch
+- Added weights_only=True to torch.load for security
+"""
+
+import argparse
+import json
+import os
+import random
+import time
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.optim as optim
+from tqdm import tqdm
+
+from dataset_conditional import DEFAULT_DATA_DIR, get_conditional_dataloaders
+from diffusion_conditional import ConditionalDiffusionModel, GaussianDiffusion
+from unet_conditional import ConditionalUNet
+
+# Weights & Biases (optional)
+try:
+    import wandb
+
+    WANDB_AVAILABLE = True
+except ImportError:
+    WANDB_AVAILABLE = False
+    print("Warning: wandb not available. Install with: pip install wandb")
+
+
+class EMA:
+    """Exponential Moving Average for model parameters"""
+
+    def __init__(self, model, decay=0.9999):
+        self.model = model
+        self.decay = decay
+        self.shadow = {}
+        for name, param in model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = param.data.clone()
+
+    def update(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                self.shadow[name] = self.decay * self.shadow[name] + (1 - self.decay) * param.data
+
+    def apply_shadow(self):
+        self.backup = {
+            name: param.data.clone() for name, param in self.model.named_parameters() if param.requires_grad
+        }
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                param.data = self.shadow[name]
+
+    def restore(self):
+        for name, param in self.model.named_parameters():
+            if param.requires_grad:
+                param.data = self.backup[name]
+        self.backup = {}
+
+
+def train_epoch(model, dataloader, optimizer, device, epoch, ema=None, use_wandb=False, scaler=None):
+    model.train()
+    total_loss = 0.0
+    total_samples = 0
+    pbar = tqdm(dataloader, desc=f"Epoch {epoch}")
+
+    for batch_idx, (images, labels) in enumerate(pbar):
+        images = images.to(device)
+        labels = labels.to(device)
+        batch_size = images.shape[0]
+
+        optimizer.zero_grad()
+
+        if scaler is not None:
+            with torch.amp.autocast("cuda"):
+                loss = model.get_loss(images, labels)
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss = model.get_loss(images, labels)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+
+        if ema is not None:
+            ema.update()
+
+        total_loss += loss.item() * batch_size
+        total_samples += batch_size
+        pbar.set_postfix({"loss": f"{loss.item():.4f}"})
+
+        if use_wandb and batch_idx % 10 == 0:
+            wandb.log({"batch_loss": loss.item(), "epoch": epoch, "batch": epoch * len(dataloader) + batch_idx})
+
+    return total_loss / total_samples
+
+
+def validate(model, dataloader, device):
+    model.eval()
+    total_loss = 0.0
+    total_samples = 0
+    with torch.no_grad():
+        for images, labels in tqdm(dataloader, desc="Validating"):
+            images = images.to(device)
+            labels = labels.to(device)
+            batch_size = images.shape[0]
+            loss = model.get_loss(images, labels)
+            total_loss += loss.item() * batch_size
+            total_samples += batch_size
+    return total_loss / total_samples
+
+
+def save_checkpoint(model, optimizer, ema, epoch, loss, save_dir, is_best=False, last_improvement_epoch=None, scheduler=None):
+    checkpoint = {
+        "epoch": epoch,
+        "model_state_dict": model.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict(),
+        "loss": loss,
+    }
+    if ema is not None:
+        checkpoint["ema_shadow"] = ema.shadow
+    if last_improvement_epoch is not None:
+        checkpoint["last_improvement_epoch"] = last_improvement_epoch
+    if scheduler is not None:
+        checkpoint["scheduler_state_dict"] = scheduler.state_dict()
+
+    torch.save(checkpoint, os.path.join(save_dir, "checkpoint_latest.pt"))
+    if is_best:
+        torch.save(checkpoint, os.path.join(save_dir, "best_model.pt"))
+        print(f"Saved best model at epoch {epoch+1}")
+
+    if (epoch + 1) % 20 == 0:
+        torch.save(checkpoint, os.path.join(save_dir, f"checkpoint_epoch_{epoch+1}.pt"))
+
+    print(f"Saved checkpoint at epoch {epoch+1}")
+
+
+def sample_images(model, diffusion, device, save_path, test_labels, ema=None, n_samples=8, epoch=0, use_ddim=True, ddim_steps=50, use_wandb=False):
+    if ema is not None:
+        ema.apply_shadow()
+
+    model.eval()
+    labels = test_labels[:n_samples].to(device)
+
+    with torch.no_grad():
+        samples = diffusion.sample(
+            model,
+            labels=labels,
+            channels=1,
+            height=256,
+            width=256,
+            device=device,
+            progress=True,
+            use_ddim=use_ddim,
+            ddim_steps=ddim_steps,
+            eta=0.0,
+        )
+
+    if ema is not None:
+        ema.restore()
+
+    n_cols = min(n_samples, 4)
+    n_rows = (n_samples + n_cols - 1) // n_cols
+    fig, axes = plt.subplots(n_rows, n_cols, figsize=(4.5 * n_cols, 4.5 * n_rows))
+    if n_rows == 1 and n_cols == 1:
+        axes = np.array([[axes]])
+    elif n_rows == 1:
+        axes = axes[np.newaxis, :]
+    elif n_cols == 1:
+        axes = axes[:, np.newaxis]
+    for i in range(n_rows * n_cols):
+        ax = axes[i // n_cols, i % n_cols]
+        if i < n_samples:
+            img = samples[i, 0].cpu().numpy()
+            label_vals = labels[i].cpu().tolist()
+            label_str = ", ".join(f"{v:.2f}" for v in label_vals)
+            ax.imshow(img, vmin=-1, vmax=1)
+            ax.set_title(label_str, fontsize=10)
+        ax.axis("off")
+
+    plt.suptitle(f"Generated Samples - Epoch {epoch}", fontsize=14)
+    plt.tight_layout()
+    plt.savefig(save_path, dpi=150, bbox_inches="tight")
+
+    if use_wandb:
+        wandb.log({"generated_samples": wandb.Image(save_path), "epoch": epoch})
+    plt.close()
+    print(f"Saved samples to {save_path}")
+
+
+def save_training_args(args, output_dir):
+    """Save training arguments so the evaluation script can reconstruct the model."""
+    args_path = os.path.join(output_dir, "args.txt")
+    with open(args_path, "w", encoding="utf-8") as f:
+        for key, value in vars(args).items():
+            f.write(f"{key}: {value}\n")
+    args_json_path = os.path.join(output_dir, "args.json")
+    with open(args_json_path, "w", encoding="utf-8") as f:
+        json.dump(vars(args), f, indent=2)
+    print(f"Saved training args to {args_path} and {args_json_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train conditional diffusion (LH 6-parameter)")
+    # Model
+    parser.add_argument("--label_dim", type=int, default=6)
+    parser.add_argument("--base_channels", type=int, default=64)
+    parser.add_argument("--channel_multipliers", type=int, nargs="+", default=[1, 2, 4, 8])
+    parser.add_argument("--attention_levels", type=int, nargs="+", default=[2, 3])
+    parser.add_argument("--dropout", type=float, default=0.1)
+    # Diffusion
+    parser.add_argument("--timesteps", type=int, default=1500)
+    parser.add_argument("--beta_start", type=float, default=1e-4)
+    parser.add_argument("--beta_end", type=float, default=0.02)
+    parser.add_argument("--schedule_type", type=str, default="linear")
+    # Training
+    parser.add_argument("--epochs", type=int, default=100)
+    parser.add_argument("--batch_size", type=int, default=8)
+    parser.add_argument("--lr", type=float, default=2e-4)
+    parser.add_argument("--ema_decay", type=float, default=0.9999)
+    parser.add_argument("--num_workers", type=int, default=4)
+    parser.add_argument("--early_stop_patience", type=int, default=30)
+    parser.add_argument(
+        "--use_amp",
+        action="store_true",
+        default=False,
+        help="Enable mixed-precision training (recommended for GPU)",
+    )
+    # Data
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default=DEFAULT_DATA_DIR,
+        help="Directory with *_LH_6.npy and *_labels_LH.npy (same rule as improved repo: e.g. .../LH_data/params_6)",
+    )
+    parser.add_argument("--normalize_labels", action=argparse.BooleanOptionalAction, default=True)
+    # Output
+    parser.add_argument("--output_dir", type=str, default="outputs_conditional_6param")
+    parser.add_argument("--resume", type=str, default="")
+    parser.add_argument(
+        "--resume_refresh_scheduler",
+        action="store_true",
+        help="On resume, rebuild cosine LR scheduler for --epochs (last_epoch=start-1) instead of loading saved scheduler; use when extending training beyond the original epoch count",
+    )
+    parser.add_argument("--sample_every", type=int, default=10)
+    parser.add_argument("--use_ddim", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--ddim_steps", type=int, default=50)
+    # WandB
+    parser.add_argument("--use_wandb", action="store_true", default=False)
+    parser.add_argument("--wandb_project", type=str, default="ddpm_cosmology")
+    parser.add_argument("--wandb_entity", type=str, default="")
+    parser.add_argument("--wandb_run_name", type=str, default="")
+
+    args = parser.parse_args()
+
+    seed = 42
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+    use_wandb = args.use_wandb and WANDB_AVAILABLE
+    if use_wandb:
+        run_name = args.wandb_run_name or f"conditional_diffusion_{time.strftime('%Y%m%d_%H%M%S')}"
+        wandb.init(project=args.wandb_project, entity=args.wandb_entity or None, name=run_name, config=vars(args))
+        print(f"W&B run: {run_name}")
+
+    timestamp = time.strftime("%Y%m%d_%H%M%S")
+    output_dir = f"{args.output_dir}_{timestamp}"
+    os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(os.path.join(output_dir, "checkpoints"), exist_ok=True)
+    os.makedirs(os.path.join(output_dir, "samples"), exist_ok=True)
+
+    save_training_args(args, output_dir)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    scaler = torch.amp.GradScaler("cuda") if args.use_amp and torch.cuda.is_available() else None
+    if scaler:
+        print("Mixed-precision training enabled (AMP)")
+
+    print("\nLoading data...")
+    train_loader, val_loader, test_loader = get_conditional_dataloaders(
+        data_dir=args.data_dir,
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        normalize_labels=args.normalize_labels,
+        label_dim=args.label_dim,
+    )
+    _, test_labels = next(iter(test_loader))
+
+    print("\nCreating model...")
+    unet = ConditionalUNet(
+        in_channels=1,
+        out_channels=1,
+        label_dim=args.label_dim,
+        base_channels=args.base_channels,
+        channel_multipliers=args.channel_multipliers,
+        attention_levels=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).to(device)
+    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")
+
+    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
+    ema = EMA(model, decay=args.ema_decay)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+
+    start_epoch = 0
+    best_val_loss = float("inf")
+    last_improvement_epoch = -1
+    if args.resume:
+        print(f"Resuming from {args.resume}")
+        checkpoint = torch.load(args.resume, map_location=device, weights_only=False)
+        model.load_state_dict(checkpoint["model_state_dict"])
+        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
+        if "ema_shadow" in checkpoint:
+            ema.shadow = checkpoint["ema_shadow"]
+        start_epoch = checkpoint["epoch"] + 1
+        best_val_loss = checkpoint.get("loss", float("inf"))
+        last_improvement_epoch = checkpoint.get("last_improvement_epoch", -1)
+        if args.resume_refresh_scheduler:
+            scheduler = optim.lr_scheduler.CosineAnnealingLR(
+                optimizer, T_max=args.epochs, last_epoch=start_epoch - 1
+            )
+            print(
+                f"Rebuilt LR scheduler for extended run: T_max={args.epochs}, "
+                f"resume at epoch {start_epoch + 1} (last_epoch={start_epoch - 1})"
+            )
+        elif "scheduler_state_dict" in checkpoint:
+            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
+
+    print("\nStarting training...")
+    losses = {"train": [], "val": []}
+
+    for epoch in range(start_epoch, args.epochs):
+        train_loss = train_epoch(model, train_loader, optimizer, device, epoch, ema, use_wandb, scaler=scaler)
+
+        if ema is not None:
+            ema.apply_shadow()
+        val_loss = validate(model, val_loader, device)
+        if ema is not None:
+            ema.restore()
+
+        losses["train"].append(train_loss)
+        losses["val"].append(val_loss)
+        scheduler.step()
+
+        if use_wandb:
+            wandb.log(
+                {
+                    "epoch": epoch + 1,
+                    "train_loss": train_loss,
+                    "val_loss": val_loss,
+                    "learning_rate": optimizer.param_groups[0]["lr"],
+                }
+            )
+
+        print(
+            f"\nEpoch {epoch+1}/{args.epochs} | Train: {train_loss:.6f} | Val: {val_loss:.6f} | "
+            f"LR: {optimizer.param_groups[0]['lr']:.6e}"
+        )
+
+        is_best = val_loss < best_val_loss
+        if is_best:
+            best_val_loss = val_loss
+            last_improvement_epoch = epoch
+
+        save_checkpoint(
+            model,
+            optimizer,
+            ema,
+            epoch,
+            val_loss,
+            os.path.join(output_dir, "checkpoints"),
+            is_best=is_best,
+            last_improvement_epoch=last_improvement_epoch,
+            scheduler=scheduler,
+        )
+
+        if epoch - last_improvement_epoch >= args.early_stop_patience:
+            print(f"Early stopping at epoch {epoch+1}")
+            break
+
+        if (epoch + 1) % args.sample_every == 0:
+            sample_path = os.path.join(output_dir, "samples", f"samples_epoch_{epoch+1}.png")
+            sample_images(
+                model,
+                diffusion,
+                device,
+                sample_path,
+                test_labels,
+                ema=ema,
+                epoch=epoch + 1,
+                use_ddim=args.use_ddim,
+                ddim_steps=args.ddim_steps,
+                use_wandb=use_wandb,
+            )
+
+        if (epoch + 1) % 5 == 0:
+            plt.figure(figsize=(10, 5))
+            plt.plot(losses["train"], label="Train Loss")
+            plt.plot(losses["val"], label="Val Loss")
+            plt.yscale("log")
+            plt.xlabel("Epoch")
+            plt.ylabel("Loss")
+            plt.title("Training Progress")
+            plt.legend()
+            plt.grid(True, alpha=0.3)
+            plt.savefig(os.path.join(output_dir, "losses.png"), dpi=150)
+            plt.close()
+
+    print(f"\nTraining completed! Best val loss: {best_val_loss:.6f}")
+    print(f"Results saved to: {output_dir}")
+    if use_wandb:
+        wandb.finish()
+
+
+if __name__ == "__main__":
+    main()