#!/usr/bin/env python3
"""Run representative cases with full snapshot storage for visualisation -> runs/viz/*.pkl"""
import pickle, sys, time
from multiprocessing import Pool
from pathlib import Path
import numpy as np
from pbhgr.spherical_ev import Grid, Run, RunConfig, solve_metric
from pbhgr.initial_data import shell_family
CASES = {
"bounce_warm": dict(nu=0.15, flatness=1.0, L_rms=0.5, Gamma_tdyn=0.0),
"bh_approach_flat": dict(nu=0.30, flatness=3.0, L_rms=0.5, Gamma_tdyn=0.0),
"decay_warm": dict(nu=0.15, flatness=1.0, L_rms=0.5, Gamma_tdyn=0.3),
"cold_caustic": dict(nu=0.15, flatness=1.0, L_rms=0.0, Gamma_tdyn=0.0),
}
def run_case(name, n_part=20000, K=300, seed=11):
c = CASES[name]
M = 1.0; R = 2 * M / c["nu"]; t_dyn = np.sqrt(R**3 / M)
rng = np.random.default_rng(seed)
g = Grid(r_out=6.0 * R, K=K)
P = shell_family(M, R, n_part, rng, flatness=c["flatness"], L_rms=c["L_rms"])
for _ in range(4):
P.N *= M / solve_metric(g, P).M
cfg = RunConfig(Gamma=c["Gamma_tdyn"] / t_dyn, t_end=5.0 * t_dyn, cfl=0.4, diag_every=25, seed=seed,
emit_fraction=0.25, store_profiles=True, store_particles=3000)
run = Run(g, P, cfg)
t0 = time.time(); outcome = run.run()
out = dict(name=name, params=c, M=M, R=R, t_dyn=t_dyn, r_edge=g.r_edge, r_c=g.r_c, dr=g.dr,
outcome=outcome, diag=run.diag, profiles=run.profiles, wall=time.time() - t0)
Path("runs/viz").mkdir(parents=True, exist_ok=True)
with open(f"runs/viz/{name}.pkl", "wb") as f:
pickle.dump(out, f)
return name, outcome, out["wall"], len(run.profiles)
if __name__ == "__main__":
names = sys.argv[1:] or list(CASES)
with Pool(2) as pool:
for r in pool.imap_unordered(run_case, names):
print(r, flush=True)