Skip to content

Posterior predictive checks for trained NPE posteriors #169

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
maho3 wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Posterior predictive checks for trained NPE posteriors #169

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b1c5749
adding PPC scripts
maho3 Aug 18, 2026
42acfce
adding rho_to_halo merge conflict
maho3 Aug 18, 2026
902c8a8
adding self-consistent PPC test
maho3 Aug 18, 2026
abab8eb
Merge branch 'main' into ppc
maho3 Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions cmass/bias/rho_to_halo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,8 +230,16 @@ def apply_charm_old(rho, fvel, charm_cfg, L, cosmo):
return hposs, hmasss, hvels, meta


def apply_charm_new(rho, fvel, L, cosmo):
"""Apply CHARM (gobig branch), accounting for the pre-trained resolution."""
CHARM_YAML = '/u/maho3/git/CHARM/run_configs/TRAIN_CHARM_JOINT_v2vel_finetune2.yaml'
CHARM_CKPT = '/work/hdd/bdne/maho3/cmass-ili/scratch/charm_joint_step001030_ph0_finetune5.pth'


def apply_charm_new(rho, fvel, L, cosmo, charm_yaml=None, charm_ckpt=None):
"""Apply CHARM (gobig branch), accounting for the pre-trained resolution.

`charm_yaml`/`charm_ckpt` default to the current model. Override them via
bias.halo.charm_yaml / bias.halo.charm_ckpt to reproduce an older suite.
"""

import torch
from charm.config_loader import load_config
Expand All @@ -241,7 +249,7 @@ def apply_charm_new(rho, fvel, L, cosmo):
build_dm_velocity_interpolators, reconstruct_catalog,
)

charm_yaml_path = '/u/maho3/git/CHARM/run_configs/TRAIN_CHARM_JOINT_v2vel_finetune2.yaml'
charm_yaml_path = charm_yaml or CHARM_YAML
cfg_charm = load_config(charm_yaml_path)

sc = cfg_charm['sim_settings']
Expand All @@ -260,7 +268,7 @@ def apply_charm_new(rho, fvel, L, cosmo):
Npix = nb * nax # pre-trained per-chunk resolution (128)
Lcharm = 1000.0 # CHARM trained box size [Mpc/h]

ckpt_path = '/work/hdd/bdne/maho3/cmass-ili/scratch/charm_joint_step001030_ph0_finetune5.pth'
ckpt_path = charm_ckpt or CHARM_CKPT

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Loading CHARM gobig model on {device}...')
Expand Down Expand Up @@ -397,7 +405,9 @@ def run_snapshot(rho, fvel, a, cfg, ppos=None, pvel=None):
hpos, hmass, hvel, meta = apply_charm_new(
rho,
fvel*a, # physical velocities in km/s
cfg.nbody.L, cfg.nbody.cosmo
cfg.nbody.L, cfg.nbody.cosmo,
charm_yaml=cfg.bias.halo.get('charm_yaml', None),
charm_ckpt=cfg.bias.halo.get('charm_ckpt', None),
)
elif cfg.bias.halo.model == "LIMD":
logging.info('Using LIMD model...')
Expand Down
100 changes: 100 additions & 0 deletions params/ppc_obs01880_cosmo.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
0.27675214 0.06087762 0.80103314 1.00657928 0.72702891
0.28667587 0.06438865 0.80647767 1.01571190 0.83903348
0.29826128 0.06003127 0.73945498 0.99380291 0.76001418
0.27670872 0.06884170 0.80064249 1.03419960 0.70753324
0.28721008 0.04821378 0.63572979 1.02552176 0.74674994
0.29134357 0.04856814 0.62440348 1.06178784 0.74835837
0.29340559 0.06035429 0.72910798 1.00190520 0.77136254
0.28541660 0.04503233 0.62893587 1.04063475 0.77609593
0.28925928 0.05602067 0.68206257 1.02603304 0.76984280
0.29725754 0.05667123 0.66964722 1.00942218 0.69314826
0.27924520 0.05084693 0.61939484 1.08430755 0.77861488
0.30607384 0.06382248 0.81837249 0.93221819 0.74947715
0.27714795 0.04609121 0.56634867 1.11298239 0.68213844
0.32800832 0.06105687 0.81210923 0.84159958 0.68183368
0.29820734 0.04901345 0.60427833 1.05812085 0.80438656
0.28288579 0.05392755 0.62231672 1.09999216 0.81865859
0.29045823 0.05255372 0.71905369 0.96149701 0.72077650
0.27582949 0.04948111 0.63225406 1.09955144 0.79645330
0.28844270 0.05503631 0.73080999 0.99438578 0.79953384
0.30072021 0.06764337 0.84002960 0.91865528 0.71764046
0.29475787 0.05297923 0.68405610 1.02001631 0.77376461
0.28258589 0.05274764 0.68282771 1.00580001 0.71466959
0.29786751 0.06589658 0.78269434 0.97580510 0.74690568
0.28313857 0.04244839 0.58259207 1.07300889 0.76418418
0.28918588 0.05535994 0.62562090 1.07586133 0.78003472
0.28838465 0.06339586 0.75300241 1.03224754 0.79105306
0.30275989 0.06415081 0.77840072 0.93716550 0.69463342
0.28605530 0.06809600 0.79685855 1.03275418 0.81581807
0.28825495 0.04370315 0.58118188 1.06157482 0.72303849
0.29548335 0.05483468 0.68853897 1.00313342 0.81419557
0.29879156 0.06906509 0.84834468 0.93496519 0.76236057
0.28290546 0.05727346 0.76486152 1.01250446 0.82166719
0.28310955 0.05552996 0.70961297 1.02813888 0.73823893
0.29039380 0.04520540 0.64977044 1.00837350 0.72608209
0.29648787 0.06338184 0.74505049 0.99652517 0.77938336
0.29569846 0.05488655 0.72086227 0.97769028 0.74694556
0.28069583 0.04116175 0.57193732 1.05546808 0.74381709
0.30704558 0.04844869 0.68689334 0.94681287 0.70039684
0.30149856 0.05932021 0.78891444 0.91895735 0.82752401
0.27757999 0.05189757 0.59850007 1.12709475 0.76871961
0.26873428 0.05622272 0.64998221 1.16653943 0.94728917
0.29949802 0.06542403 0.81011629 0.92104411 0.66341853
0.28955072 0.06114673 0.74219400 1.01285183 0.81135195
0.30062020 0.05750507 0.74509799 0.96183556 0.72427827
0.24748853 0.03146946 0.50747555 1.19165456 0.77386993
0.28869647 0.06042621 0.74825853 0.98873878 0.73344809
0.27584162 0.05749136 0.77644783 0.99020988 0.75203031
0.33354425 0.04766504 0.63616556 0.91012001 0.66836387
0.31617472 0.04793444 0.64662355 0.96314049 0.80233508
0.29516935 0.05648018 0.78127939 0.94875550 0.71517122
0.27652317 0.05534474 0.72896403 1.02857602 0.73511392
0.30099788 0.06177065 0.74604970 0.98777109 0.74322391
0.30696875 0.06382020 0.86002600 0.83638233 0.75878316
0.30317223 0.05864483 0.78095168 0.95426023 0.79168987
0.30752850 0.06951352 0.89254779 0.91021413 0.80555636
0.30458423 0.05617470 0.78876704 0.89003670 0.75729597
0.29128227 0.06091065 0.88039815 0.88415766 0.73424876
0.31483412 0.06050878 0.78254247 0.92211461 0.80731302
0.29852515 0.05670508 0.68944848 1.00186372 0.71491253
0.34141114 0.05439284 0.70967162 0.85940450 0.67420745
0.31392139 0.04527319 0.58977026 1.01493025 0.79114038
0.29308742 0.04875853 0.62011588 1.04119074 0.81534028
0.31901792 0.06037047 0.75446552 0.94039643 0.74751318
0.30037466 0.04919076 0.68921214 0.94481736 0.78726685
0.28957531 0.04846058 0.58735967 1.10388291 0.78802145
0.29291627 0.04251602 0.59600556 1.03297997 0.74633551
0.31184685 0.06941345 0.85101867 0.91323280 0.77994984
0.33531442 0.05747644 0.74993724 0.86816931 0.77976024
0.30026445 0.05271457 0.69845265 0.95851743 0.80991149
0.32163137 0.05145400 0.70991814 0.89062345 0.77373445
0.31171402 0.05992209 0.73863924 0.94755101 0.76910633
0.28806958 0.04680296 0.62136936 1.01532328 0.70441049
0.29675168 0.04234955 0.55441815 1.08226335 0.76740152
0.27865875 0.03933500 0.58668697 1.09077597 0.71357185
0.33350423 0.05824307 0.69760138 0.90555459 0.76556224
0.29363468 0.03763252 0.52290273 1.09948599 0.77208495
0.31026453 0.05382432 0.67019498 0.93170184 0.77486038
0.31514233 0.06841953 0.84632653 0.91844165 0.80549991
0.30788982 0.04306333 0.63031584 0.88958091 0.70046300
0.32178277 0.04219004 0.71729285 0.87701881 0.73028654
0.27391002 0.06486482 0.60764349 1.19865191 0.79235429
0.30513713 0.05288957 0.78737527 0.84576368 0.67651677
0.31101727 0.03801847 0.59431005 0.93368876 0.70473278
0.28402326 0.06143761 0.62183326 1.04084063 0.78252625
0.33154067 0.03472157 0.56859517 0.91364682 0.70128369
0.28940696 0.04844623 0.70668995 1.02459896 0.78569061
0.28802955 0.05615541 0.70035845 0.98479700 0.76232362
0.29747766 0.04381624 0.70468503 0.95476228 0.72484672
0.33944589 0.06264764 0.65049404 1.00571513 0.85388607
0.32395485 0.05086449 0.63585520 0.96044940 0.71982354
0.28857893 0.04511701 0.62866497 1.02781737 0.78945786
0.28749228 0.06740821 0.87731558 0.95294112 0.82308495
0.27053094 0.03798579 0.58970886 1.11352706 0.79746419
0.31810576 0.03501521 0.73844635 0.82003546 0.69481272
0.30361909 0.06347793 0.69019061 1.02587664 0.78648412
0.31513464 0.05938600 0.78327918 0.89155459 0.79948014
0.32905859 0.04863289 0.85454845 0.80014139 0.73486245
0.31408709 0.05646950 0.73071271 0.90498686 0.73989260
0.33134058 0.06265547 0.62378323 1.03212404 0.77948666
0.33078876 0.04393242 0.52024376 1.07683766 0.77517629
79 changes: 79 additions & 0 deletions ppc/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
# Posterior predictive checks

Draw parameters from `q(theta | x_obs)` at one test point, resimulate each draw
through the identical forward chain, and compare the resulting data vectors
against the observed one.

This is the replacement for importance-sampling `cmass/infer/resim.py`, which
reweights existing simulations instead of generating new ones and collapses in
high dimension (ESS = 3 of 19777 on a 17-parameter model). Nothing here modifies
`resim.py`; the scripts reuse its loaders and plotting helpers.

Validated on a 100-simulation self-consistent run against
`abacuslike/fastpm_charm6_comphod`, zPk0+zPk2+zPk4 at kmax=0.4. Write-up:
`ltu-gobig-notes/experiments/2026-08-13_ppc_abacuslike-fastpm_charm6_comphod/`.

## Stages

Run from the repo root. `PYTHONPATH=.` is needed for the two python entry
points, as in `power_tests/`; the sbatch scripts set it themselves. Stage B is
deliberately not a SLURM job: CHARM wants a GPU, which is often a different
machine from the one running FastPM.

```bash
PYTHONPATH=. python ppc/draw.py --ndraw 10 # draw theta, write cosmofile + overrides
sbatch ppc/slurm_nbody.sh # A: FastPM (~14 min/draw, 128 cores)
bash ppc/run_charm.sh 0 99 # B: CHARM (by hand, GPU machine)
sbatch ppc/slurm_hod.sh # C: HOD + summaries (~3 min/draw, 16 cores)
sbatch ppc/slurm_collect.sh # D: arrays + figures (~1 min)
```

Every stage is idempotent and resume-safe: A and C skip draws that already have
their output, B skips draws with no `nbody.h5` yet or an existing `halos.h5`, so
C and B can be re-run as the one before them drains.

To extend a campaign, `PYTHONPATH=. python ppc/draw.py --start 10 --ndraw 90`
appends. Earlier draws are carried over verbatim from the existing npz and never
regenerated, so extending cannot perturb draws already simulated.

## Adapting to a new experiment

`ppc/draw.py --exp_path` picks the trained model; the output root and the `tag`
in the job scripts derive from it. The job scripts hardcode the suite, box and
forward-chain config of the run being reproduced, which **must** match the
config that produced the training summaries, not the current defaults. Check
before reusing:

- `nbody=`, `bias=`, `multisnapshot`, `nbody.zf` in `slurm_nbody.sh` and
`slurm_hod.sh`
- `CHARM_CKPT` in `run_charm.sh`. The checkpoint path is a `bias.halo.charm_ckpt`
override; HEAD's default is the current CHARM, which is not necessarily the
one that built the suite you are reproducing
- `diag.summaries` in `slurm_hod.sh`, pinned to match the training run rather
than inheriting today's `diag/default.yaml`

A mismatch here does not error. It produces a PPC that tests a different forward
model than the posterior was trained on, and a failed check then tells you about
the config difference rather than the model.

## Outputs

Under `<wdir>/ppc/<suite>_<sim>/<summaries>_<kcut>/<tag>/`:

| | |
|---|---|
| `x_ppc.npy` | (Ndraw, Nfeat) inference blocks only, training x ordering |
| `theta_ppc.npy` | (Ndraw, Nparam), row-aligned with `x_ppc` |
| `x_ppc_all.npz` | every plotted block, its observed vector, its k axis |
| `posterior_draws.npz` | `x_obs`, `theta_obs`, `theta_draws`, `seed_blocks`, `param_names` |
| `logq_ppc.npy` | log q(theta \| x_obs) per simulated draw |
| `manifest.tsv` | per-draw status and all parameters |
| `plots/ppc_{bands,corner,logprob}.png` | |

`collect.py` verifies every draw's recorded cosmology, HOD and noise parameters
against the drawn values, drops mismatches from `x_ppc` and `theta_ppc`
together, and records the reason in the manifest. It aborts rather than write a
misaligned array if the block layout disagrees with the training data.

No Mahalanobis distance or p-value is computed. `x_ppc` / `theta_ppc` are
row-aligned and `x_obs` is in the npz so a distance statistic is trivial to add.
Loading