diff --git a/cmass/bias/rho_to_halo.py b/cmass/bias/rho_to_halo.py index 5774e6d9..ec37eae4 100644 --- a/cmass/bias/rho_to_halo.py +++ b/cmass/bias/rho_to_halo.py @@ -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 @@ -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'] @@ -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}...') @@ -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...') diff --git a/params/ppc_obs01880_cosmo.txt b/params/ppc_obs01880_cosmo.txt new file mode 100644 index 00000000..34deaca0 --- /dev/null +++ b/params/ppc_obs01880_cosmo.txt @@ -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 diff --git a/ppc/README.md b/ppc/README.md new file mode 100644 index 00000000..755fc4d2 --- /dev/null +++ b/ppc/README.md @@ -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 `/ppc/_/_//`: + +| | | +|---|---| +| `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. diff --git a/ppc/collect.py b/ppc/collect.py new file mode 100644 index 00000000..0deb9917 --- /dev/null +++ b/ppc/collect.py @@ -0,0 +1,499 @@ +""" +Step 4 of the posterior predictive check (PPC) campaign. + +Collects the per-draw diagnostics into the deliverable arrays and plots the +posterior predictive bands against the observed data vector. + + /theta_ppc.npy (Ndraw, 17) theta actually simulated + /x_ppc.npy (Ndraw, 117) the INFERENCE blocks only, in training + x ordering -- the deliverable + /x_ppc_all.npz every plotted summary block, inference or not + /plots/ppc_bands.png, ppc_corner.png, ppc_logprob.png + +Rows align: x_ppc[i] is the data vector simulated at theta_ppc[i]. Draws whose +diagnostics are missing or whose recorded parameters disagree with the draw are +dropped from BOTH arrays together and marked in manifest.tsv. + +The plot covers ALL available z-space summaries, not just the ones the posterior +was conditioned on. Held-out summaries (zBk0, zEqBk0, zSqBk0, ...) carry no +weight in the inference, so disagreement there is informative: it says the +forward model reproduces the constrained statistics but not the unconstrained +ones. Panel titles mark which is which. + +For summaries outside the inference set there is no x_obs in the trained +experiment, so the observed vector is recomputed from the observed lhid's own +diagnostics file -- located by matching its recorded HOD parameters to +theta_obs, not by hardcoding a filename. + +k-grid caveat: the training summaries predate ltu-cmass ba2334f (2026-07-14, +pylians -> pypower for periodic-box P(k)), so training P(k) sits on a uniform +0.01 grid and these PPC runs sit on pypower's effective-k grid. Both give 39 +bins in k <= 0.4. Bk is unaffected. Per instruction this is left as-is and +everything is plotted against the PPC (pypower) k, with residuals taken +element-wise -- matching how x_ppc and x_obs are used downstream. + +No Mahalanobis distances, no p-values -- bands only. +""" + +import argparse +import os +import re +import shutil +from os.path import join, exists +import numpy as np +import h5py +from omegaconf import OmegaConf +import matplotlib +matplotlib.use('Agg') +from matplotlib import pyplot as plt # noqa: E402 + +from cmass.infer.loaders import ( # noqa: E402 + preprocess_Pk, preprocess_Bk, load_Pk, load_Bk, + _is_in_kminmax, _get_Bk_mask) +from cmass.infer.resim import ( # noqa: E402 + load_pool, plot_logprob, plot_corner, batched_log_prob) +from cmass.infer.tools import resolve_kmax # noqa: E402 + +WDIR = '/work/hdd/bdne/maho3/cmass-ili' +PPC = join(WDIR, 'ppc/abacuslike_fastpm_charm6_comphod', + 'zPk0+zPk2+zPk4_kmin-0.0_kmax-0.4/obs01880') +AF = 0.666660000066666 # analysis snapshot (a), key '0.666660' +N_COSMO, N_NOISE = 5, 2 +TAGS = ('Eq', 'Sq', 'Ss', 'Is', '') + +# Entity -> colour, fixed. Observed is ink, the PPC ensemble is one hue, the +# training pool is recessive fill. Separated by lightness as well as hue, so it +# survives CVD and greyscale print. +C_OBS, C_PPC, C_POOL = 'k', 'C0', '0.85' + + +def build_argparser(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--ppc_dir', default=PPC) + p.add_argument('--sim_sub', default='fastpm/L2000-N256') + p.add_argument('--hod_seed', type=int, default=1, + help='which diag/galaxies/hodNNNNN.h5 to read per draw') + p.add_argument('--atol', type=float, default=1e-5, + help='tolerance when checking recorded vs drawn params') + p.add_argument('--summaries', default=None, + help='comma-separated blocks to plot. Default: the ' + 'inference blocks plus every other available z-space ' + 'summary') + p.add_argument('--bk_kmax', type=float, default=None, + help='kmax for bispectrum blocks (default: same as the ' + 'experiment k-cut)') + p.add_argument('--obs_dir', default=None, + help='per-lhid dir of the observed sim. Default: derived ' + 'from exp_path and id_obs') + p.add_argument('--no_plot', action='store_true') + p.add_argument('--n_post', type=int, default=10000, + help='direct posterior samples for the theta-space plots') + p.add_argument('--batch_size', type=int, default=2048) + p.add_argument('--device', default='cpu') + p.add_argument('--no_theta_plots', action='store_true') + return p + + +def actual_nnets(exp_path, nnets): + """How many nets load_ensemble will really use. + + cfg.infer.Nnets is what was *requested*; load_ensemble silently skips top + trials whose posterior.pkl is missing. Reporting the requested count in a + figure title would overstate the ensemble, so resolve it the same way + load_ensemble does -- without paying to unpickle anything. + """ + import optuna + from cmass.infer.tools import select_top_trials, study_name_from_path + db = join(exp_path, 'optuna_study.db') + if not exists(db): + return None + try: + study = optuna.load_study(storage=f'sqlite:///{db}', + study_name=study_name_from_path(exp_path)) + top = select_top_trials(study, nnets) + except Exception: + return None + return sum(exists(join(exp_path, 'nets', f'net-{t.number}', + 'posterior.pkl')) for t in top) + + +def kcut_from_path(exp_path): + m = re.match(r'kmin-([\d.]+)_kmax-([\d.]+)', os.path.basename(exp_path)) + if not m: + raise ValueError(f'Cannot parse k-cut from {exp_path}') + return float(m.group(1)), float(m.group(2)) + + +def split_tag(summ): + """'zEqBk0' -> ('zBk0', 'Eq'), replicating run_preprocessing's loop.""" + for tag in TAGS: + if tag in summ: + return summ.replace(tag, ''), tag + return summ, '' + + +def load_summ(diagfile): + s = {} + s.update(load_Pk(diagfile, AF)) + s.update(load_Bk(diagfile, AF)) + return s + + +def preprocess_block(summ, data, cfg, kmin, kmax, bk_kmax): + """One summary block, preprocessed exactly as run_preprocessing does.""" + base, tag = split_tag(summ) + norm_key = base[:-1] + '0' + is_bk = ('Bk' in base) or ('Qk' in base) + skmax = (bk_kmax if (is_bk and bk_kmax is not None) + else resolve_kmax(kmax, summ)) + norm = None if '0' in base else data[norm_key] + if is_bk: + x = preprocess_Bk(data[base], kmin=kmin, kmax=skmax, norm=norm, + mode=tag, correct_shot=cfg.infer.correct_shot) + else: + x = preprocess_Pk(data[base], kmin=kmin, kmax=skmax, norm=norm, + correct_shot=cfg.infer.correct_shot, + loglinear_start_idx=cfg.infer.loglinear_start_idx) + return x + + +def block_axis(summ, kdata, kmin, kmax, bk_kmax): + """x values and axis label for a block. + + P(k) blocks and equilateral bispectra have a single k per feature, so they + get a real k axis. General triangle configurations do not -- three k's per + point -- so those fall back to a triangle index. + """ + base, tag = split_tag(summ) + is_bk = ('Bk' in base) or ('Qk' in base) + skmax = (bk_kmax if (is_bk and bk_kmax is not None) + else resolve_kmax(kmax, summ)) + if not is_bk: + k = np.asarray(kdata) + return k[_is_in_kminmax(k, kmin, skmax)], r'$k$ [$h$/Mpc]', skmax + k123 = np.asarray(kdata) + mask = _get_Bk_mask(k123, kmin, skmax, equilateral=(tag == 'Eq'), + squeezed=(tag == 'Sq'), subsampled=(tag == 'Ss'), + isoceles=(tag == 'Is')) + if tag == 'Eq': # k1 == k2 == k3, so a k axis is meaningful + return k123[0][mask], r'$k$ [$h$/Mpc]', skmax + return np.arange(int(mask.sum())), 'triangle index', skmax + + +def ylabel_for(summ): + base, _ = split_tag(summ) + if '0' in base: + stat = 'B' if 'Bk' in base else ('Q' if 'Qk' in base else 'P') + return rf'signed $\log_{{10}} {stat}_0$' + ell = base[-1] + stat = 'B' if 'Bk' in base else ('Q' if 'Qk' in base else 'P') + return rf'${stat}_{ell}/{stat}_0$' + + +def find_obs_diag(obs_dir, theta_obs, names, atol): + """The observed lhid has several HOD realizations; pick the one whose + recorded HOD parameters are theta_obs. Identifying it by content rather + than filename keeps this correct if the pool ordering ever changes.""" + diagdir = join(obs_dir, 'diag', 'galaxies') + want = theta_obs[N_COSMO:-N_NOISE] + hodnames = list(names[N_COSMO:-N_NOISE]) + for fn in sorted(os.listdir(diagdir)): + path = join(diagdir, fn) + try: + with h5py.File(path, 'r') as f: + if [str(s) for s in f.attrs['HOD_names']] != hodnames: + continue + if not np.allclose(np.asarray(f.attrs['HOD_params'], float), + want, atol=atol): + continue + noise = np.array([f.attrs['noise_radial'], + f.attrs['noise_transverse']], float) + if not np.allclose(noise, theta_obs[-N_NOISE:], atol=atol): + continue + except (OSError, KeyError): + continue + return path + raise SystemExit( + f'No diagnostics file in {diagdir} matches theta_obs. Cannot ' + 'reconstruct the observed vector for held-out summaries.') + + +def check_draw(diagfile, theta, names, atol): + """Confirm the sim actually used the parameters we drew (TODO.md ยง7.6).""" + problems = [] + with h5py.File(diagfile, 'r') as f: + cosmo = np.asarray(f.attrs['cosmo_params'], dtype=float) + hod = np.asarray(f.attrs['HOD_params'], dtype=float) + hodnames = [str(s) for s in f.attrs['HOD_names']] + noise = np.array([f.attrs['noise_radial'], + f.attrs['noise_transverse']], dtype=float) + if not np.allclose(cosmo, theta[:N_COSMO], atol=atol): + problems.append(f'cosmo {cosmo} != {theta[:N_COSMO]}') + if hodnames != list(names[N_COSMO:-N_NOISE]): + problems.append(f'HOD name order {hodnames}') + elif not np.allclose(hod, theta[N_COSMO:-N_NOISE], atol=atol): + problems.append(f'HOD {hod} != {theta[N_COSMO:-N_NOISE]}') + if not np.allclose(noise, theta[-N_NOISE:], atol=atol): + problems.append(f'noise {noise} != {theta[-N_NOISE:]}') + return problems + + +def plot_bands(blocks, title, out_path, ncols=4): + """Value + residual panels per summary block. + + Two rows per group of columns: the vector itself with 68/95 bands, and the + residual against the observed. Held-out blocks are flagged in the panel + title (text, never colour alone). + """ + n = len(blocks) + ncols = min(ncols, n) + ngroup = int(np.ceil(n / ncols)) + fig, axs = plt.subplots(2 * ngroup, ncols, squeeze=False, + figsize=(4.3 * ncols, 3.1 * 2 * ngroup)) + q = [2.5, 16, 50, 84, 97.5] + for idx, b in enumerate(blocks): + g, c = divmod(idx, ncols) + ax, axr = axs[2 * g, c], axs[2 * g + 1, c] + x, obs, X = b['x'], b['obs'], b['x_ppc'] + qs = np.percentile(X, q, axis=0) + + if b['pool'] is not None: + lo, hi = np.percentile(b['pool'], [2.5, 97.5], axis=0) + ax.fill_between(x, lo, hi, color=C_POOL, label='training pool 95%') + ax.fill_between(x, qs[0], qs[4], color=C_PPC, alpha=0.25, + label='PPC 95%') + ax.fill_between(x, qs[1], qs[3], color=C_PPC, alpha=0.45, + label='PPC 68%') + ax.plot(x, qs[2], color=C_PPC, lw=2, label='PPC median') + ax.plot(x, obs, color=C_OBS, lw=2, label='observed') + used = 'inference' if b['used'] else 'held out' + ax.set_title(f"{b['name']} ({used}, {b['n']} bins)", fontsize=10) + ax.set_ylabel(b['ylabel'], fontsize=9) + + axr.fill_between(x, qs[0] - obs, qs[4] - obs, color=C_PPC, alpha=0.25) + axr.fill_between(x, qs[1] - obs, qs[3] - obs, color=C_PPC, alpha=0.45) + axr.plot(x, qs[2] - obs, color=C_PPC, lw=2) + axr.axhline(0, color=C_OBS, lw=2) + axr.set_xlabel(b['xlabel'], fontsize=9) + axr.set_ylabel('PPC $-$ observed', fontsize=9) + + for a in (ax, axr): + a.tick_params(labelsize=8) + a.grid(alpha=0.25, lw=0.5) + a.set_axisbelow(True) + for side in ('top', 'right'): + a.spines[side].set_visible(False) + if idx == 0: + ax.legend(fontsize=7.5, loc='best', framealpha=0.9) + + for idx in range(n, ngroup * ncols): # blank unused panels + g, c = divmod(idx, ncols) + axs[2 * g, c].axis('off') + axs[2 * g + 1, c].axis('off') + + fig.suptitle(title, fontsize=11) + fig.tight_layout(rect=[0, 0, 1, 0.985]) + fig.savefig(out_path, dpi=150, bbox_inches='tight') + plt.close(fig) + + +def main(): + args = build_argparser().parse_args() + out = args.ppc_dir + + draws = np.load(join(out, 'posterior_draws.npz'), allow_pickle=True) + theta_draws = draws['theta_draws'] + names = [str(s) for s in draws['param_names']] + inf_labels = [str(s) for s in draws['labels']] + startidx_ref = list(draws['startidx']) + exp_path = str(draws['exp_path']) + x_obs = draws['x_obs'] + theta_obs = np.asarray(draws['theta_obs']) + id_obs = str(draws['id_obs']) + nnets_req = int(draws['nnets']) if 'nnets' in draws else None + + cfg = OmegaConf.load(join(exp_path, 'config.yaml')) + if cfg.infer.pca_features or exists(join(exp_path, 'pca.pkl')): + raise SystemExit('Experiment uses PCA; must apply it, never refit.') + kmin, kmax = kcut_from_path(exp_path) + print(f'exp_path = {exp_path}') + print(f'k-cut = {kmin} <= k <= {kmax}') + print(f'inference = {inf_labels}, startidx {startidx_ref}') + print(f'correct_shot={cfg.infer.correct_shot}, ' + f'loglinear_start_idx={cfg.infer.loglinear_start_idx}') + + # --- the observed sim's own diagnostics (for held-out summaries) -------- + obs_dir = args.obs_dir + if obs_dir is None: + suite_root = exp_path.split(os.sep + 'models' + os.sep)[0] + obs_dir = join(suite_root, f'L{cfg.nbody.L}-N{cfg.nbody.N}', id_obs) + obs_diag = find_obs_diag(obs_dir, theta_obs, names, args.atol) + print(f'x_obs from = {obs_diag}') + obs_data = load_summ(obs_diag) + + # --- which blocks to plot ---------------------------------------------- + if args.summaries: + plot_labels = [s.strip() for s in args.summaries.split(',')] + else: + # every z-space summary present, plus the equilateral/squeezed slices + # of the bispectrum monopole + avail = sorted(k for k in obs_data if k.startswith('z')) + plot_labels = ( + inf_labels + + [s for s in avail if s not in inf_labels] + + [f'z{t}Bk0' for t in ('Eq', 'Sq') if 'zBk0' in avail]) + print(f'plotting = {plot_labels}') + + # --- gather per-draw summaries ----------------------------------------- + draw_summs = [] + needed_bases = {split_tag(lab)[0] for lab in plot_labels} + needed_bases |= {split_tag(lab)[0][:-1] + '0' for lab in plot_labels} + kept, status = [], {} + for i, theta in enumerate(theta_draws): + simdir = join(out, args.sim_sub, str(i)) + diagfile = join(simdir, 'diag', 'galaxies', + f'hod{args.hod_seed:05d}.h5') + if not exists(diagfile): + status[i] = 'missing_diag' + continue + s = load_summ(diagfile) + if any(b not in s for b in needed_bases): + status[i] = 'incomplete_summ' + continue + problems = check_draw(diagfile, theta, names, args.atol) + if problems: + status[i] = 'param_mismatch' + print(f' draw {i}: PARAM MISMATCH -- ' + '; '.join(problems)) + continue + draw_summs.append(s) + kept.append(i) + status[i] = 'ok' + + n_ok = len(kept) + print(f'\n{n_ok}/{len(theta_draws)} draws usable') + for st in sorted(set(status.values()) - {'ok'}): + bad = [i for i, v in status.items() if v == st] + print(f' {st}: {len(bad)} -> {bad}') + if n_ok == 0: + raise SystemExit('No usable draws.') + + # per-draw dicts, reindexed by base summary for preprocess_* + draw_data = {b: [s[b] for s in draw_summs] for b in needed_bases} + + # the training pool, for the grey context band on inference blocks + POOL, _, _, _ = load_pool(exp_path, ('train', 'val', 'test')) + + # --- preprocess every plotted block ------------------------------------ + blocks, xs_inf = [], [] + for lab in plot_labels: + base, _ = split_tag(lab) + X = preprocess_block(lab, draw_data, cfg, kmin, kmax, args.bk_kmax) + o = preprocess_block(lab, {b: [obs_data[b]] for b in needed_bases}, + cfg, kmin, kmax, args.bk_kmax)[0] + xax, xlabel, skmax = block_axis( + lab, obs_data[base]['k'], kmin, kmax, args.bk_kmax) + if len(xax) != X.shape[1]: + raise SystemExit( + f'{lab}: axis has {len(xax)} points but block has ' + f'{X.shape[1]} features.') + used = lab in inf_labels + pool = None + if used: # pool band only where a trained x exists + j = inf_labels.index(lab) + pool = POOL[:, startidx_ref[j]:startidx_ref[j + 1]] + o = x_obs[startidx_ref[j]:startidx_ref[j + 1]] + xs_inf.append(X) + blocks.append(dict(name=lab, x=xax, xlabel=xlabel, + ylabel=ylabel_for(lab), obs=o, x_ppc=X, + pool=pool, used=used, n=X.shape[1], kmax=skmax)) + print(f' {lab:9s} {X.shape[1]:3d} bins ' + f'{"inference" if used else "held out"} kmax={skmax}') + + # --- deliverables: inference blocks only, unchanged --------------------- + x_ppc = np.concatenate(xs_inf, axis=-1) + startidx = list(np.cumsum([0] + [b.shape[1] for b in xs_inf])) + if startidx != list(startidx_ref): + raise SystemExit( + f'Block layout {startidx} != training {list(startidx_ref)}.') + theta_ppc = theta_draws[kept] + np.save(join(out, 'x_ppc.npy'), x_ppc) + np.save(join(out, 'theta_ppc.npy'), theta_ppc) + np.savez(join(out, 'x_ppc_all.npz'), + **{b['name']: b['x_ppc'] for b in blocks}, + **{b['name'] + '_obs': b['obs'] for b in blocks}, + **{b['name'] + '_k': b['x'] for b in blocks}) + print(f'\nWrote x_ppc.npy {x_ppc.shape}, theta_ppc.npy {theta_ppc.shape}, ' + f'x_ppc_all.npz ({len(blocks)} blocks)') + + # --- manifest ----------------------------------------------------------- + mpath = join(out, 'manifest.tsv') + rows = open(mpath).read().splitlines() + header, body = rows[0], rows[1:] + row_of = {int(r.split('\t')[0]): r.split('\t') for r in body} + with open(mpath, 'w') as f: + f.write(header + '\n') + for i in range(len(theta_draws)): + r = row_of[i] + r[1] = status.get(i, 'pending') + f.write('\t'.join(r) + '\n') + + if args.no_plot: + return + + # --- bands -------------------------------------------------------------- + parts = exp_path.rstrip('/').split(os.sep) + suite, sim, tracer = parts[-6], parts[-5], parts[-3] + n_nets = actual_nnets(exp_path, nnets_req or cfg.infer.Nnets) + nets = f'{n_nets}-net ' if n_nets else '' + if n_nets and nnets_req and n_nets != nnets_req: + print(f'NOTE: ensemble is {n_nets} nets, not the {nnets_req} ' + f'requested (missing posterior.pkl for some top trials)') + title = ( + f'Posterior predictive check | {suite}/{sim}, tracer={tracer}\n' + f'conditioned on {"+".join(inf_labels)} at {kmin} $\\leq k \\leq$ ' + f'{kmax} | {nets}{cfg.infer.backend}/{cfg.infer.engine} ensemble, ' + f'correct_shot={cfg.infer.correct_shot}\n' + f'$x_{{\\rm obs}}$ = lhid {id_obs} ({os.path.basename(obs_diag)}), ' + f'{n_ok} joint draws from $q(\\theta|x_{{\\rm obs}})$') + plotdir = join(out, 'plots') + os.makedirs(plotdir, exist_ok=True) + plot_bands(blocks, title, join(plotdir, 'ppc_bands.png')) + print(f'Wrote {join(plotdir, "ppc_bands.png")}') + + if args.no_theta_plots: + return + + # --- theta-space plots -------------------------------------------------- + # resim.py's importance-sampling diagnostics, reused. For a PPC these are + # bookkeeping checks, not tests of the model: theta_ppc are direct + # posterior draws, so both SHOULD agree with the direct sample. What they + # catch is drift between what we drew and what got simulated. The ensemble + # is unweighted, so ESS = N and max weight = 1/N trivially. + import torch + from cmass.infer.validate import load_ensemble + w = np.full(n_ok, 1. / n_ok) + ensemble = load_ensemble(exp_path, cfg.infer.Nnets, plot=False, + clean=False).to(args.device) + with torch.no_grad(): + theta_post = ensemble.sample( + (args.n_post,), torch.Tensor(x_obs).to(args.device), + show_progress_bars=False).cpu().numpy() + logq_post = batched_log_prob(ensemble, theta_post, x_obs, + args.batch_size, args.device) + logq_ppc = batched_log_prob(ensemble, theta_ppc, x_obs, + args.batch_size, args.device) + print(f'log q(theta|x_obs): direct median {np.median(logq_post):.2f}, ' + f'simulated median {np.median(logq_ppc):.2f}') + + plot_logprob(logq_post, logq_ppc, float(n_ok), 1.0, 1. / n_ok, plotdir) + shutil.move(join(plotdir, 'plot_resim_logprob.png'), + join(plotdir, 'ppc_logprob.png')) + plot_corner(theta_post, theta_ppc, w, theta_obs, names, plotdir) + shutil.move(join(plotdir, 'plot_resim_corner.png'), + join(plotdir, 'ppc_corner.png')) + np.save(join(out, 'logq_ppc.npy'), logq_ppc) + print('Wrote ppc_logprob.png, ppc_corner.png, logq_ppc.npy') + + +if __name__ == '__main__': + main() diff --git a/ppc/draw.py b/ppc/draw.py new file mode 100644 index 00000000..4d7f1e2a --- /dev/null +++ b/ppc/draw.py @@ -0,0 +1,222 @@ +""" +Step 1 of the posterior predictive check (PPC) campaign. + +Draws theta from q(theta | x_obs) at a single fiducial test point, and writes +everything the SLURM stages need to simulate those draws: + + params/ppc__cosmo.txt cosmology rows, indexed by draw id (= lhid) + /posterior_draws.npz theta_draws, x_obs, theta_obs, id_obs, ... + /manifest.tsv one row per draw + /overrides/.txt hydra overrides for the HOD + noise stage + +Draws are joint: one theta per simulation, straight from ensemble.sample(). +Nothing here is expensive; it runs on the login node. + +Appending (campaign step 5): pass --start N to extend an existing campaign. +The full stream of start+ndraw draws is regenerated from the same seed and the +first N are checked against the existing npz, so earlier draws are never +silently altered. +""" + +import argparse +import os +from os.path import join, exists +import numpy as np +import torch +from omegaconf import OmegaConf + +from cmass.infer.resim import load_pool, load_labels, param_names, \ + select_test_point, empirical_quantiles +from cmass.infer.validate import load_ensemble + +WDIR = '/work/hdd/bdne/maho3/cmass-ili' +EXP_PATH = join( + WDIR, 'abacuslike/fastpm_charm6_comphod/models/galaxy', + 'zPk0+zPk2+zPk4/kmin-0.0_kmax-0.4') + +# theta layout: 5 cosmology, then HOD (alphabetical, from hodprior.csv), +# then noise_radial, noise_transverse. +N_COSMO = 5 +N_NOISE = 2 + + +def build_argparser(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument('--exp_path', default=EXP_PATH) + p.add_argument('--wdir', default=WDIR) + p.add_argument('--ndraw', type=int, default=10, + help='number of new draws to generate') + p.add_argument('--start', type=int, default=0, + help='draw id to start at; earlier draws are reproduced ' + 'from the same seed and verified, not rewritten') + p.add_argument('--seed', type=int, default=0) + p.add_argument('--nnets', type=int, default=None, + help='ensemble size (default: infer.Nnets from exp config)') + p.add_argument('--device', default='cpu') + p.add_argument('--expect_lhid', type=int, default=1880, + help='abort if select_test_point disagrees; -1 to disable') + p.add_argument('--tag', default='obs01880') + p.add_argument('--outroot', default=None, + help='default: /ppc///') + return p + + +def default_outroot(wdir, exp_path, tag): + # ...///models/// + parts = exp_path.rstrip('/').split(os.sep) + kcut, summ = parts[-1], parts[-2] + sim, suite = parts[-5], parts[-6] + return join(wdir, 'ppc', f'{suite}_{sim}', f'{summ}_{kcut}', tag) + + +def draw_theta(ensemble, x_obs, n, seed, device): + """n joint draws from q(theta|x_obs), rejecting any outside prior support. + + Deterministic given seed: the accept/reject stream is fixed. + """ + torch.manual_seed(seed) + xt = torch.Tensor(x_obs).to(device) + kept, n_rejected = [], 0 + with torch.no_grad(): + while sum(len(k) for k in kept) < n: + want = n - sum(len(k) for k in kept) + t = ensemble.sample((max(want, 64),), xt, show_progress_bars=False) + lp = ensemble.prior.log_prob(t) + good = torch.isfinite(lp) + n_rejected += int((~good).sum()) + kept.append(t[good].cpu().numpy()[:want]) + return np.concatenate(kept)[:n], n_rejected + + +def write_cosmofile(path, cosmo, start): + """Space-delimited, one row per draw id, matching latin_hypercube_params.""" + if start > 0: + if not exists(path): + raise FileNotFoundError( + f'--start={start} but no existing cosmofile at {path}') + old = np.loadtxt(path, ndmin=2) + if len(old) < start: + raise ValueError( + f'{path} has {len(old)} rows, need at least {start}') + if not np.allclose(old[:start], cosmo[:start], atol=1e-8): + raise ValueError( + f'Regenerated draws disagree with existing rows in {path}. ' + 'Refusing to overwrite. Check --seed.') + np.savetxt(path, cosmo, fmt='%.8f', delimiter=' ') + + +def override_string(names, theta): + """Hydra overrides injecting one draw's HOD and noise. No spaces, so the + job script can splice it in unquoted.""" + hod = dict(zip(names[N_COSMO:-N_NOISE], theta[N_COSMO:-N_NOISE])) + hodstr = '{' + ','.join(f'{k}:{float(v)!r}' for k, v in hod.items()) + '}' + return (f'bias.hod.theta={hodstr} ' + f'noise.params.radial={float(theta[-2])!r} ' + f'noise.params.transverse={float(theta[-1])!r}') + + +def main(): + args = build_argparser().parse_args() + out = args.outroot or default_outroot(args.wdir, args.exp_path, args.tag) + os.makedirs(join(out, 'overrides'), exist_ok=True) + + cfg = OmegaConf.load(join(args.exp_path, 'config.yaml')) + nnets = args.nnets if args.nnets is not None else cfg.infer.Nnets + if cfg.infer.pca_features: + raise RuntimeError( + 'Experiment uses PCA; ppc_collect must load pca.pkl. Aborting.') + + # --- the observed point ------------------------------------------------- + x, theta, ids, tags = load_pool(args.exp_path, ('train', 'val', 'test')) + labels, startidx = load_labels(args.exp_path) + names = param_names(args.exp_path) + assert len(names) == theta.shape[1], (names, theta.shape) + + iobs = select_test_point(theta, tags, theta) + x_obs, theta_obs, id_obs = x[iobs], theta[iobs], ids[iobs] + if args.expect_lhid >= 0 and int(id_obs) != args.expect_lhid: + raise SystemExit( + f'select_test_point returned lhid {id_obs}, expected ' + f'{args.expect_lhid}. Stopping (see TODO.md Phase 0.1).') + q_obs = empirical_quantiles(theta, theta_obs[None])[0] + print(f'Pool: {x.shape[0]} vectors, {theta.shape[1]} params, ' + f'x is {x.shape[1]}-dim') + print(f'x_obs: index {iobs}, lhid {id_obs}, split {tags[iobs]}') + for n_, v_, q_ in zip(names, theta_obs, q_obs): + print(f' {n_:46s} {v_:12.6g} (q={q_:.2f})') + + # --- draw --------------------------------------------------------------- + ensemble = load_ensemble(args.exp_path, nnets, plot=False, + clean=False).to(args.device) + npz_path = join(out, 'posterior_draws.npz') + n_total = args.start + args.ndraw + + # Append-only: earlier draws are carried over verbatim from the npz, never + # regenerated. New draws come from their own RNG stream (seed + start), so + # correctness doesn't depend on torch reproducing a prefix of a longer + # sample. They are still i.i.d. from the same q(theta|x_obs). + if args.start > 0: + old = np.load(npz_path, allow_pickle=True) + if len(old['theta_draws']) < args.start: + raise SystemExit( + f'{npz_path} holds {len(old["theta_draws"])} draws, ' + f'need at least --start={args.start}') + if not np.allclose(old['x_obs'], x_obs): + raise SystemExit( + 'x_obs differs from the existing posterior_draws.npz. ' + 'Refusing to append.') + prev = old['theta_draws'][:args.start] + prev_seeds = (old['seed_blocks'].tolist() + if 'seed_blocks' in old else [int(old['seed'])]) + else: + prev = np.empty((0, theta.shape[1])) + prev_seeds = [] + + block_seed = args.seed + args.start + new_draws, n_rejected = draw_theta( + ensemble, x_obs, args.ndraw, block_seed, args.device) + theta_draws = np.concatenate([prev, new_draws]) + seed_blocks = prev_seeds + [block_seed] + print(f'\nDrew {args.ndraw} new theta ({n_rejected} rejected outside ' + f'prior support), seed={block_seed}; {len(theta_draws)} total') + + # --- cosmofile (round-trip so theta_ppc matches what the sims read) ----- + cosmofile = join('params', f'ppc_{args.tag}_cosmo.txt') + write_cosmofile(cosmofile, theta_draws[:, :N_COSMO], args.start) + theta_draws[:, :N_COSMO] = np.loadtxt(cosmofile, ndmin=2) + print(f'Wrote {cosmofile} ({n_total} rows)') + + np.savez( + npz_path, + theta_draws=theta_draws, x_obs=x_obs, theta_obs=theta_obs, + id_obs=id_obs, index_obs=iobs, split_obs=tags[iobs], + seed=args.seed, seed_blocks=np.array(seed_blocks), + n_rejected=n_rejected, exp_path=args.exp_path, + param_names=np.array(names), labels=np.array(labels), + startidx=np.array(startidx), nnets=nnets, + ) + + # --- per-draw overrides + manifest -------------------------------------- + for i in range(args.start, n_total): + with open(join(out, 'overrides', f'{i}.txt'), 'w') as f: + f.write(override_string(names, theta_draws[i]) + '\n') + + with open(join(out, 'manifest.tsv'), 'w') as f: + f.write('\t'.join(['draw_id', 'status', 'wall_s'] + names + + ['sim_dir', 'diag_file']) + '\n') + for i in range(n_total): + simdir = join(out, 'fastpm', 'L2000-N256', str(i)) + f.write('\t'.join( + [str(i), 'pending', ''] + + [f'{v:.10g}' for v in theta_draws[i]] + + [simdir, join(simdir, 'diag', 'galaxies', 'hod00000.h5')] + ) + '\n') + + print(f'Wrote {npz_path}, manifest.tsv, and {n_total - args.start} ' + f'override files under {out}') + print(f'\nSuite for the job scripts:\n ' + f'{os.path.relpath(out, args.wdir)}') + + +if __name__ == '__main__': + main() diff --git a/ppc/run_charm.sh b/ppc/run_charm.sh new file mode 100644 index 00000000..1fa81bd6 --- /dev/null +++ b/ppc/run_charm.sh @@ -0,0 +1,66 @@ +#!/bin/bash +# Stage B of the PPC campaign: CHARM (cmass.bias.rho_to_halo). +# +# NOT a SLURM script -- run this by hand on a GPU machine, after +# ppc/slurm_nbody.sh has produced nbody.h5 for every draw. +# +# usage: bash ppc/run_charm.sh [first_draw] [last_draw] +# e.g. bash ppc/run_charm.sh 0 99 +# +# CHARM VERSION: the training suite (abacuslike/fastpm_charm6) was built with +# charm6 = charm_joint_best_val_ft15.pth (ltu-cmass commit 23ff0a8, 2026-06-10). +# HEAD defaults to charm_joint_v19.pth (charm7), which is a DIFFERENT halo bias +# model -- using it would make the PPC test a forward model the posterior was +# never trained on. CHARM_CKPT below pins the right one. +# +# rho_to_halo reads nbody.h5 and writes halos.h5 into the SAME directory, so +# nothing needs symlinking. + +set -u + +first=${1:-0} +last=${2:-99} + +# --- adjust these three for the machine you are running on ------------------- +WDIR=/work/hdd/bdne/maho3/cmass-ili +CHARM_CKPT=$WDIR/scratch/charm_joint_best_val_ft15.pth +CHARM_YAML=/u/maho3/git/CHARM/run_configs/TRAIN_CHARM_JOINT_v2vel_finetune2.yaml +# ----------------------------------------------------------------------------- + +cd /u/maho3/git/ltu-cmass + +tag=obs01880 +ppcdir=ppc/abacuslike_fastpm_charm6_comphod/zPk0+zPk2+zPk4_kmin-0.0_kmax-0.4/$tag +L=2000 +N=256 + +for f in "$CHARM_CKPT" "$CHARM_YAML" "./params/ppc_${tag}_cosmo.txt"; do + if [ ! -f "$f" ]; then echo "MISSING: $f"; exit 1; fi +done + +outdir=$WDIR/$ppcdir/fastpm/L$L-N$N +echo "outdir=$outdir" +echo "ckpt=$CHARM_CKPT" + +for lhid in $(seq "$first" "$last"); do + if [ ! -f "$outdir/$lhid/nbody.h5" ]; then + echo "draw=$lhid: no nbody.h5 yet, skipping" + continue + fi + if [ -f "$outdir/$lhid/halos.h5" ]; then + echo "draw=$lhid: halos.h5 exists, skipping" + continue + fi + + start=$(date +%s) + python -m cmass.bias.rho_to_halo \ + nbody=abacuslike sim=fastpm bias=zheng_composite \ + nbody.suite="'$ppcdir'" nbody.lhid=$lhid \ + multisnapshot=False nbody.zf=0.500015 nbody.matchIC=0 \ + meta.wdir=$WDIR \ + meta.cosmofile=./params/ppc_${tag}_cosmo.txt \ + +bias.halo.charm_ckpt=$CHARM_CKPT \ + +bias.halo.charm_yaml=$CHARM_YAML \ + hydra/job_logging=disabled + echo "draw=$lhid stage=charm status=$? wall_s=$(( $(date +%s) - start ))" +done diff --git a/ppc/slurm_collect.sh b/ppc/slurm_collect.sh new file mode 100644 index 00000000..0c52eebb --- /dev/null +++ b/ppc/slurm_collect.sh @@ -0,0 +1,28 @@ +#!/bin/bash +#SBATCH --job-name=ppc_collect # Job name +#SBATCH --nodes=1 # Number of nodes +#SBATCH --ntasks=16 # Number of tasks +#SBATCH --mem=64G # Amount of memory +#SBATCH --time=2:00:00 # Time limit +#SBATCH --partition=cpu # Partition name +#SBATCH --account=bdne-delta-cpu # Account name +#SBATCH --output=/work/hdd/bdne/maho3/jobout/%x_%A.out # Output file +#SBATCH --error=/work/hdd/bdne/maho3/jobout/%x_%A.out # Error file + +# Step 4 of the PPC campaign: collect x_ppc/theta_ppc and make the plots. +# +# The bands stage is cheap, but the theta-space plots (corner + logprob) load +# the 8-net ensemble and draw n_post samples from it, which blew a 40-minute +# wall on the login node. Hence a job. + +source ~/.bashrc +conda activate cmass + +cd /u/maho3/git/ltu-cmass + +export OMP_NUM_THREADS=16 +export MKL_NUM_THREADS=16 +export TQDM_DISABLE=0 + +PYTHONPATH=. python -u ppc/collect.py --n_post 5000 --device cpu +echo "ppc_collect status=$?" diff --git a/ppc/slurm_hod.sh b/ppc/slurm_hod.sh new file mode 100644 index 00000000..077f751f --- /dev/null +++ b/ppc/slurm_hod.sh @@ -0,0 +1,97 @@ +#!/bin/bash +#SBATCH --job-name=ppc_hod # Job name +#SBATCH --array=0-99%50 # One task per posterior draw (array idx == draw id == lhid) +#SBATCH --nodes=1 # Number of nodes +#SBATCH --ntasks=16 # Number of tasks +#SBATCH --mem=64G # Amount of memory +#SBATCH --time=2:00:00 # Time limit +#SBATCH --partition=cpu # Partition name +#SBATCH --account=bdne-delta-cpu # Account name +#SBATCH --output=/work/hdd/bdne/maho3/jobout/%x_%A_%a.out # Output file for each array task +#SBATCH --error=/work/hdd/bdne/maho3/jobout/%x_%A_%a.out # Error file for each array task + +# Stage C of the PPC campaign: apply_hod + diagnostics. +# Run after stage A (ppc/slurm_nbody.sh) and stage B (ppc/run_charm.sh). +# +# Mirrors jobs/slurm_abacuslike_bias.sh, which produced the training summaries: +# bias=zheng_composite, noise_uniform=False, multisnapshot=False, +# nbody.zf=0.500015, diag.from_scratch=True, rm_galaxies=True. +# +# Two deliberate differences from that job, both required by the PPC: +# * HOD parameters and noise are INJECTED per draw from overrides/.txt +# rather than sampled. bias.hod.seed=1 still resolves to lhid*1e4+1 inside +# parse_hod (so galaxy-placement RNG varies per draw, as in training), but +# the theta dict is applied afterwards and overwrites all 10 sampled +# params. Verified: resolved theta == the drawn theta, exactly. +# * noise=fixed instead of noise=reciprocal, so noise_radial/transverse take +# the drawn values instead of being redrawn from the prior. +# +# diag.summaries is pinned to [Pk,Bk] because HEAD's diag/default.yaml now also +# lists nz, which the training runs did not have. + +echo "SLURM_ARRAY_TASK_ID=$SLURM_ARRAY_TASK_ID" + +source ~/.bashrc +conda activate cmass + +lhid=$SLURM_ARRAY_TASK_ID + +cd /u/maho3/git/ltu-cmass + +tag=obs01880 +ppcdir=ppc/abacuslike_fastpm_charm6_comphod/zPk0+zPk2+zPk4_kmin-0.0_kmax-0.4/$tag +ppcbase=/work/hdd/bdne/maho3/cmass-ili/$ppcdir + +nbody=abacuslike +sim=fastpm +L=2000 +N=256 +multisnapshot=False +rm_galaxies=True + +outdir=$ppcbase/$sim/L$L-N$N +ovrfile=$ppcbase/overrides/$lhid.txt + +if [ ! -f "$outdir/$lhid/halos.h5" ]; then + echo "draw=$lhid: no halos.h5 (CHARM not run yet). Skipping." + exit 0 +fi +if [ ! -f "$ovrfile" ]; then + echo "draw=$lhid: no override file at $ovrfile. Aborting." + exit 1 +fi + +# bias.hod.theta={...} and the two noise values. No spaces inside the braces, +# so this splices in as three words. +draw_extras=$(cat "$ovrfile") +echo "draw=$lhid overrides: $draw_extras" + +postfix="nbody=$nbody sim=$sim nbody.suite='$ppcdir' nbody.lhid=$lhid" +postfix="$postfix multisnapshot=$multisnapshot nbody.zf=0.500015 nbody.matchIC=0" +postfix="$postfix bias=zheng_composite noise=fixed" +postfix="$postfix bias.hod.seed=1 bias.hod.from_samples=False" +postfix="$postfix bias.hod.noise_uniform=False" +postfix="$postfix meta.cosmofile=./params/ppc_${tag}_cosmo.txt" +postfix="$postfix $draw_extras hydra/job_logging=disabled" + +export TQDM_DISABLE=0 + +diag_file=$outdir/$lhid/diag/galaxies/hod00001.h5 +if [ -f "$diag_file" ]; then + echo "draw=$lhid: $diag_file exists. Skipping." +else + start=$(date +%s) + python -m cmass.bias.apply_hod $postfix + st_hod=$? + python -m cmass.diagnostics.summ $postfix \ + diag.galaxy=True diag.from_scratch=True 'diag.summaries=[Pk,Bk]' + st_summ=$? + echo "draw=$lhid stage=hod status=$st_hod summ_status=$st_summ wall_s=$(( $(date +%s) - start ))" +fi + +# Trash collection. Only removes galaxies/, which this job created; nbody.h5 +# and halos.h5 are left alone. +if [ "$rm_galaxies" = "True" ]; then + echo "Removing galaxies/ for draw=$lhid" + rm -rf "$outdir/$lhid/galaxies" +fi diff --git a/ppc/slurm_nbody.sh b/ppc/slurm_nbody.sh new file mode 100644 index 00000000..3a48a94d --- /dev/null +++ b/ppc/slurm_nbody.sh @@ -0,0 +1,84 @@ +#!/bin/bash +#SBATCH --job-name=ppc_nbody # Job name +#SBATCH --array=0-99%50 # One task per posterior draw (array idx == draw id == lhid) +#SBATCH --nodes=1 # Number of nodes +#SBATCH --ntasks=128 # Number of tasks +#SBATCH --mem=240G # Amount of memory +#SBATCH --time=3:00:00 # Time limit +#SBATCH --partition=cpu # Partition name +#SBATCH --account=bdne-delta-cpu # Account name +#SBATCH --output=/work/hdd/bdne/maho3/jobout/%x_%A_%a.out # Output file for each array task +#SBATCH --error=/work/hdd/bdne/maho3/jobout/%x_%A_%a.out # Error file for each array task + +# Stage A of the PPC campaign: FastPM only. +# Stage B = CHARM (cmass.bias.rho_to_halo), run by hand on a GPU machine. +# Stage C = ppc/slurm_hod.sh (apply_hod + diagnostics). +# +# The forward chain replicates jobs/slurm_abacuslike_bias.sh, which produced the +# training suite: nbody=abacuslike, bias=zheng_composite, multisnapshot=False, +# nbody.zf=0.500015 (a=0.666660, the analysis snapshot). +# +# Cosmology per draw comes from params/ppc__cosmo.txt, written by +# ppc/draw.py, indexed by nbody.lhid. matchIC=0 so each draw gets its +# own IC phase (gen_white_noise is seeded by lhid). +# +# Storage strategy copied from jobs/slurm_fastpm_2000.sh: the transient FastPM +# particle snapshots are staged on node-local /tmp via meta.scratchdir so they +# never touch the quota'd /work/hdd (which is at 17.6T of a 19.5T soft quota). + +echo "SLURM_ARRAY_TASK_ID=$SLURM_ARRAY_TASK_ID" + +module load cray-mpich/8.1.32 gsl +export LD_LIBRARY_PATH=/sw/rh9.4/spack/v1.0.0/sw/linux-x86_64_v2/gsl-2.8-zty4u3k/lib:$LD_LIBRARY_PATH + +source ~/.bashrc +conda activate cmass + +lhid=$SLURM_ARRAY_TASK_ID + +cd /u/maho3/git/ltu-cmass + +tag=obs01880 +ppcdir=ppc/abacuslike_fastpm_charm6_comphod/zPk0+zPk2+zPk4_kmin-0.0_kmax-0.4/$tag + +nbody=abacuslike +sim=fastpm +L=2000 +N=256 +multisnapshot=False + +# Node-local scratch base for transient particle snapshots (per job+node). +# get_source_path appends suite/sim/L-N/lhid underneath this. +scratchbase=/tmp/$USER/cmass_${SLURM_ARRAY_JOB_ID}_${SLURM_ARRAY_TASK_ID} +mkdir -p "$scratchbase" + +extras="bias=zheng_composite nbody.zf=0.500015 nbody.matchIC=0" +extras="$extras meta.cosmofile=./params/ppc_${tag}_cosmo.txt" +extras="$extras meta.scratchdir=$scratchbase" +extras="$extras hydra/job_logging=disabled" + +export TQDM_DISABLE=0 + +outdir=/work/hdd/bdne/maho3/cmass-ili/$ppcdir/$sim/L$L-N$N +echo "outdir=$outdir" +echo "scratchbase=$scratchbase" + +postfix="nbody=$nbody sim=$sim nbody.suite='$ppcdir' nbody.lhid=$lhid" +postfix="$postfix multisnapshot=$multisnapshot $extras" + +# Resume guard, same convention as slurm_fastpm_2000.sh: config.yaml is written +# last, so a crashed mid-run leaves an incomplete dir and re-runs. +cfgfile=$outdir/$lhid/config.yaml +h5file=$outdir/$lhid/nbody.h5 +if [ -f "$cfgfile" ] || [ -f "$h5file" ]; then + echo "$outdir/$lhid already done. Skipping." +else + echo "$outdir/$lhid not done. Running FastPM for draw $lhid." + start=$(date +%s) + python -m cmass.nbody.fastpm $postfix + status=$? + echo "draw=$lhid stage=nbody status=$status wall_s=$(( $(date +%s) - start ))" +fi + +# Safety net: clear this job's node-local scratch (also auto-purged at job end). +rm -rf "$scratchbase"