Ко всем библиотекам · scripts

scripts/make_figures.py

Build the visualisation set in reports/figures from runs/viz, runs/scan_pilot, runs/resource_probe.json.

235 строк · 16.3 KB · pbhgr @ 9e8e13d · как текст

#!/usr/bin/env python3
"""Build the visualisation set in reports/figures from runs/viz, runs/scan_pilot, runs/resource_probe.json."""
import json, pickle
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm

FIG = Path("reports/figures"); FIG.mkdir(parents=True, exist_ok=True)
plt.rcParams.update({"font.size": 9, "axes.grid": True, "grid.alpha": 0.3, "figure.dpi": 130})
OUTCOME_COLOR = {"BH_approach": "#d62728", "bounce": "#1f77b4", "dispersion": "#2ca02c", "unresolved": "#7f7f7f"}
CASE_TITLE = {"bounce_warm": "тёплое облако, ν=0.15, L_rms=0.5 → отскок",
              "bh_approach_flat": "плоский профиль, ν=0.30 → приближение к горизонту",
              "decay_warm": "тёплое облако + распад Γt_dyn=0.3",
              "cold_caustic": "холодная пыль, L=0 → каустика в центре"}

def load(name):
    return pickle.load(open(f"runs/viz/{name}.pkl", "rb"))

def rho_profile(d, snap):
    r_c, dr = d["r_c"], d["dr"]
    m_c = 0.5 * (snap["m_edge"][1:] + snap["m_edge"][:-1])
    eml = np.sqrt(np.maximum(1 - 2 * m_c / r_c, 1e-14))
    return eml * snap["SE"] / (4 * np.pi * r_c**2 * dr)

# ---------------------------------------------------------------- 1. outcome map
def fig_outcome_map():
    rows = json.load(open("runs/scan_pilot/scan_summary.json"))
    flats = sorted({r["flatness"] for r in rows}); Gs = sorted({r["Gamma_tdyn"] for r in rows})
    fig, axes = plt.subplots(len(Gs), len(flats), figsize=(9, 6.5), sharex=True, sharey=True)
    for i, G in enumerate(Gs):
        for j, fl in enumerate(flats):
            ax = axes[i, j]
            for r in rows:
                if r["flatness"] != fl or r["Gamma_tdyn"] != G: continue
                ok = r["M_total_drift"] < 1e-2
                ax.scatter(r["nu"], r["L_rms"], s=420, c=OUTCOME_COLOR[r["outcome"]], marker="o" if ok else "X",
                           edgecolor="k", linewidth=0.8, alpha=0.9, zorder=3)
                ax.annotate(f"{r['max_compactness']:.2f}", (r["nu"], r["L_rms"]), ha="center", va="center",
                            fontsize=7, color="w", fontweight="bold", zorder=4)
            ax.set_title(f"flatness={fl:.0f}, Γ·t_dyn={G:.1f}", fontsize=9)
            ax.set_xlim(0, 0.36); ax.set_ylim(-0.15, 0.65)
    for ax in axes[-1]: ax.set_xlabel("ν = 2M/R (начальная компактность)")
    for ax in axes[:, 0]: ax.set_ylabel("L_rms / sqrt(M/R)")
    from matplotlib.lines import Line2D
    h = [Line2D([], [], marker="o", ls="", ms=10, color=c, label=k) for k, c in OUTCOME_COLOR.items()]
    h.append(Line2D([], [], marker="X", ls="", ms=10, color="gray", label="ADM-дрейф > 1e-2 (недостоверно)"))
    fig.legend(handles=h, loc="lower center", ncol=5, fontsize=8, frameon=False)
    fig.suptitle("Пилотная карта исходов (число в кружке = max 2m/r за прогон)", fontsize=11)
    fig.tight_layout(rect=(0, 0.06, 1, 0.96)); fig.savefig(FIG / "fig1_outcome_map.png"); plt.close(fig)

# ---------------------------------------------------------------- 2. time series
def fig_timeseries():
    names = list(CASE_TITLE)
    fig, axes = plt.subplots(2, 2, figsize=(10, 6.5))
    for n in names:
        d = load(n); dg = d["diag"]; t = np.array([x["t"] for x in dg]) / d["t_dyn"]
        lab = CASE_TITLE[n]
        axes[0, 0].plot(t, [x["max_2m_over_r"] for x in dg], label=lab)
        axes[0, 1].semilogy(t, [max(x["min_lapse"], 1e-6) for x in dg], label=lab)
        axes[1, 0].plot(t, [x["M_in"] for x in dg], label=lab)
        axes[1, 1].plot(t, [x["n_particles"] for x in dg], label=lab)
    axes[0, 0].axhline(0.99, ls="--", c="k", lw=0.8); axes[0, 0].text(0.05, 0.95, "порог 0.99: срыв полярно-ареальной калибровки", fontsize=7, transform=axes[0, 0].transAxes)
    axes[0, 0].set_ylabel("max 2m/r (компактность)"); axes[0, 1].set_ylabel("min лапс e^μ")
    axes[1, 0].set_ylabel("масса внутри r_out (M_in)"); axes[1, 1].set_ylabel("число макрочастиц (X + дочерние)")
    for ax in axes.flat: ax.set_xlabel("t / t_dyn")
    axes[0, 0].legend(fontsize=7, loc="center right")
    fig.suptitle("Временные ряды четырёх показательных прогонов", fontsize=11)
    fig.tight_layout(); fig.savefig(FIG / "fig2_timeseries.png"); plt.close(fig)

# ---------------------------------------------------------------- 3. spacetime diagrams
def fig_spacetime():
    names = ["bounce_warm", "bh_approach_flat", "cold_caustic"]
    fig, axes = plt.subplots(1, 3, figsize=(13, 4.6))
    for ax, n in zip(axes, names):
        d = load(n); pr = d["profiles"]; R = d["R"]; td = d["t_dyn"]
        t = np.array([s["t"] for s in pr]) / td
        comp = np.array([2 * s["m_edge"][1:] / d["r_edge"][1:] for s in pr])
        rr = d["r_edge"][1:] / R
        pc = ax.pcolormesh(rr, t, comp, cmap="magma", vmin=0, vmax=1, shading="nearest")
        # sample worldlines
        idx = np.linspace(0, len(pr[0]["particles"]["r"]) - 1, 25).astype(int)
        for i in idx:
            wl = [s["particles"]["r"][i] / R if i < len(s["particles"]["r"]) else np.nan for s in pr]
            ax.plot(wl, t, c="cyan", lw=0.5, alpha=0.7)
        ax.set_xlim(0, 2.5); ax.set_xlabel("r / R"); ax.set_title(CASE_TITLE[n], fontsize=8)
    axes[0].set_ylabel("t / t_dyn")
    fig.colorbar(pc, ax=axes, label="2m(r)/r", shrink=0.9)
    fig.suptitle("Пространственно-временные диаграммы: цвет = 2m/r, голубые линии = мировые линии макрочастиц", fontsize=10)
    fig.savefig(FIG / "fig3_spacetime.png", bbox_inches="tight"); plt.close(fig)

# ---------------------------------------------------------------- 4. phase space
def fig_phase_space():
    fig, axes = plt.subplots(2, 5, figsize=(14, 5.6))
    for row, n in enumerate(["cold_caustic", "bounce_warm"]):
        d = load(n); pr = d["profiles"]; R = d["R"]; td = d["t_dyn"]
        sel = np.linspace(0, len(pr) - 1, 5).astype(int)
        for ax, k in zip(axes[row], sel):
            p = pr[k]["particles"]; X = p["m"] > 0
            ax.scatter(p["r"][X] / R, p["w"][X], s=1.5, c=np.sqrt(p["L"][X]) / R, cmap="viridis", alpha=0.6)
            ax.set_title(f"t = {pr[k]['t']/td:.2f} t_dyn", fontsize=8); ax.set_xlim(0, 2.2)
            if row == 1: ax.set_xlabel("r / R")
        axes[row, 0].set_ylabel(CASE_TITLE[n].split("→")[0] + "\nw = p^r̂ (радиальный импульс)", fontsize=8)
    fig.suptitle("Фазовое пространство (r, w) частиц X; цвет = sqrt(L)/R. Многозначность w при одном r = пересечение потоков", fontsize=10)
    fig.tight_layout(); fig.savefig(FIG / "fig4_phase_space.png"); plt.close(fig)

# ---------------------------------------------------------------- 5. radial profiles
def fig_profiles():
    d = load("bh_approach_flat"); pr = d["profiles"]; R = d["R"]; td = d["t_dyn"]
    sel = np.linspace(0, len(pr) - 1, 6).astype(int)
    fig, axes = plt.subplots(2, 2, figsize=(10, 6.5))
    cm = plt.cm.plasma(np.linspace(0, 0.9, len(sel)))
    for c, k in zip(cm, sel):
        s = pr[k]; lab = f"t={s['t']/td:.2f} t_dyn"
        axes[0, 0].semilogy(d["r_c"] / R, np.maximum(rho_profile(d, s), 1e-12), c=c, label=lab)
        axes[0, 1].plot(d["r_edge"] / R, s["m_edge"], c=c)
        axes[1, 0].plot(d["r_edge"] / R, 2 * s["m_edge"] / np.maximum(d["r_edge"], 1e-9), c=c)
        axes[1, 1].plot(d["r_edge"] / R, np.exp(s["mu_edge"]), c=c)
    axes[0, 0].set_ylabel("ρ(r) энергия/собств. объём"); axes[0, 0].legend(fontsize=7)
    axes[0, 1].set_ylabel("m(r) масса Мизнера–Шарпа"); axes[1, 0].set_ylabel("2m/r"); axes[1, 0].axhline(1, ls="--", c="k", lw=0.8)
    axes[1, 1].set_ylabel("лапс e^μ (темп собственного времени)")
    for ax in axes.flat: ax.set_xlim(0, 3); ax.set_xlabel("r / R")
    axes[0, 0].set_ylim(1e-6, None)
    fig.suptitle(f"Радиальные профили: {CASE_TITLE['bh_approach_flat']}", fontsize=11)
    fig.tight_layout(); fig.savefig(FIG / "fig5_profiles.png"); plt.close(fig)

# ---------------------------------------------------------------- 6. decay
def fig_decay():
    d = load("decay_warm"); dg = d["diag"]; td = d["t_dyn"]; R = d["R"]
    t = np.array([x["t"] for x in dg]); G = d["params"]["Gamma_tdyn"] / td
    fig, axes = plt.subplots(2, 2, figsize=(10, 6.5))
    ax = axes[0, 0]
    restX = np.array([x["rest_mass_X"] for x in dg])
    ax.plot(t / td, restX / restX[0], label="нераспавшаяся масса покоя X (внутри r_out)")
    ax.plot(t / td, np.exp(-G * t), "k--", label="exp(−Γ t): без замедления времени и без утечки X")
    ax.plot(t / td, [x["decayed_rest_mass"] for x in dg], label="распавшаяся масса X (M_dec, по событиям)")
    ax.set_ylabel("доля начальной массы"); ax.legend(fontsize=7)
    ax = axes[0, 1]
    ax.plot(t / td, [x["M_in"] for x in dg], label="M_in (внутри r_out)")
    ax.plot(t / td, [x["escaped_energy"] for x in dg], label="энергия Киллинга, ушедшая через r_out")
    ax.plot(t / td, [x["M_total"] for x in dg], "k", label="M_in + ушедшая = ADM (должна быть const)")
    ax.set_ylabel("масса / энергия"); ax.legend(fontsize=7)
    ax = axes[1, 0]
    ax.plot(t / td, [x["n_daughters"] for x in dg], label="дочерних макрочастиц"); ax.plot(t / td, [x["n_particles"] - x["n_daughters"] for x in dg], label="макрочастиц X")
    ax.set_ylabel("число макрочастиц"); ax.legend(fontsize=7)
    ax = axes[1, 1]
    pr = d["profiles"]; k = int(np.argmax([np.sum(s["particles"]["m"] == 0) for s in pr[: len(pr) // 2]]))
    p = pr[k]["particles"]; X = p["m"] > 0
    ax.scatter(p["r"][X] / R, p["w"][X], s=2, c="tab:blue", label="X (m=1)", alpha=0.6)
    ax.scatter(p["r"][~X] / R, p["w"][~X], s=2, c="tab:orange", label="дочерние (m=0)", alpha=0.6)
    ax.set_xlim(0, 6); ax.set_xlabel("r / R"); ax.set_ylabel("w"); ax.set_title(f"фазовое пространство при t={pr[k]['t']/td:.2f} t_dyn", fontsize=8); ax.legend(fontsize=7)
    for ax in axes.flat[:3]: ax.set_xlabel("t / t_dyn")
    fig.suptitle(f"Распад: {CASE_TITLE['decay_warm']}", fontsize=11)
    fig.tight_layout(); fig.savefig(FIG / "fig6_decay.png"); plt.close(fig)

# ---------------------------------------------------------------- 7. validation
def fig_validation():
    from pbhgr.spherical_ev import Grid, Particles, Run, RunConfig, solve_metric, metric_at
    from pbhgr.initial_data import dust_ball, normalise_mass
    from pbhgr.analytic import schwarzschild_radial_infall, eta_of_t
    from scipy.optimize import brentq
    fig, axes = plt.subplots(1, 3, figsize=(13, 4.2))
    # LTB shells
    M, R0 = 1.0, 10.0
    rng = np.random.default_rng(3); g = Grid(r_out=20.0, K=400)
    P = normalise_mass(g, dust_ball(M, R0, 20000, rng), M)
    m0 = metric_at(g, solve_metric(g, P), P.r)[0]; r0 = P.r.copy()
    run = Run(g, P, RunConfig(t_end=25.0, cfl=0.4, diag_every=50)); run.run()
    def r_of_tau(m, a, tau):
        if m <= 0: return a
        if tau >= np.sqrt(a**3 / (8 * m)) * np.pi: return 0.0     # shell already reached the centre
        e = brentq(lambda e: np.sqrt(a**3 / (8 * m)) * (e + np.sin(e)) - tau, 0, np.pi, xtol=1e-12)
        return 0.5 * a * (1 + np.cos(e))
    sel = np.linspace(0, len(r0) - 1, 400).astype(int)
    pred = np.array([r_of_tau(m0[i], r0[i], run.P.tau[i]) for i in sel])
    ax = axes[0]; ax.scatter(r0[sel] / R0, run.P.r[sel] / R0, s=6, label="решатель (t=25M)")
    ax.plot(r0[sel] / R0, pred / R0, "r-", lw=1, label="аналитика LTB: r(τ_i) с m(r0_i)")
    ax.set_xlabel("начальный радиус r0 / R0"); ax.set_ylabel("текущий радиус r / R0"); ax.set_title("Коллапс пыли: каждая оболочка — геодезика", fontsize=9); ax.legend(fontsize=7)
    # radial infall in Schwarzschild
    g = Grid(r_out=40.0, K=1200, r_in=2.5, m_in=1.0)
    P = Particles(np.array([15.0]), np.array([0.0]), np.array([0.0]), np.array([0.0]), np.array([1.0]))
    tr = []
    run = Run(g, P, RunConfig(t_end=70.0, cfl=0.4, diag_every=10**9))
    cb = lambda run_, d: None
    # manual loop to record trajectory
    from pbhgr.spherical_ev import solve_metric as sm
    met = sm(g, run.P); traj = [(0.0, 15.0)]
    while run.t < 70.0:
        run._step(min(run.dt(met), 70.0 - run.t)); traj.append((run.t, run.P.r[0]))
    traj = np.array(traj)
    eta = np.linspace(0, np.arccos(4 / 15 - 1) * 0.999, 400)
    r_ex, _, t_ex = schwarzschild_radial_infall(1.0, 15.0, eta)
    ax = axes[1]; ax.plot(traj[::20, 0], traj[::20, 1], "o", ms=3, label="решатель"); ax.plot(t_ex, r_ex, "r-", lw=1, label="точная геодезика")
    ax.axhline(2, ls="--", c="k", lw=0.8); ax.text(2, 2.2, "r = 2M", fontsize=7); ax.set_xlim(0, 70); ax.set_xlabel("t / M"); ax.set_ylabel("r / M")
    ax.set_title("Радиальное падение в Шварцшильде (t → ∞ у горизонта)", fontsize=9); ax.legend(fontsize=7)
    # convergence
    ax = axes[2]; Ks = [300, 600, 1200]; e_mu = []; e_orb = []
    from pbhgr.analytic import circular_orbit_L
    for K in Ks:
        g = Grid(r_out=40.0, K=K, r_in=2.5, m_in=1.0)
        P = Particles(np.array([10.0]), np.array([0.0]), np.array([0.0]), np.array([0.0]), np.array([1.0]))
        e_mu.append(np.max(np.abs(solve_metric(g, P).mu_edge - 0.5 * np.log(1 - 2 / g.r_edge))))
        P = Particles(np.array([8.0]), np.array([0.0]), np.array([circular_orbit_L(1.0, 8.0)]), np.array([0.0]), np.array([1.0]))
        run = Run(g, P, RunConfig(t_end=2 * np.pi * np.sqrt(512) * 3, cfl=0.4, diag_every=10**9)); run.run(); e_orb.append(abs(run.P.r[0] - 8.0) / 8.0)
    ax.loglog(Ks, e_mu, "o-", label="ошибка μ(r) вакуума"); ax.loglog(Ks, e_orb, "s-", label="дрейф радиуса круговой орбиты (3 оборота)")
    ax.loglog(Ks, e_mu[0] * (Ks[0] / np.array(Ks))**2, "k--", lw=0.8, label="∝ K⁻² (2-й порядок)")
    ax.set_xlabel("число ячеек K"); ax.set_ylabel("относительная ошибка"); ax.set_title("Сходимость по сетке", fontsize=9); ax.legend(fontsize=7)
    fig.suptitle("Валидация решателя на аналитических решениях", fontsize=11)
    fig.tight_layout(); fig.savefig(FIG / "fig7_validation.png"); plt.close(fig)

# ---------------------------------------------------------------- 8. resources
def fig_resources():
    rows = json.load(open("runs/resource_probe.json"))
    ok = [r for r in rows if "error" not in r]
    n = np.array([r["n_part"] for r in ok]); s = np.array([r["s_per_step"] for r in ok]); rss = np.array([r["peak_rss_MB"] for r in ok])
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    axes[0].loglog(n, s, "o-"); axes[0].set_xlabel("число частиц"); axes[0].set_ylabel("секунд на шаг RK4 (1 ядро)")
    axes[0].loglog(n, s[-1] * n / n[-1], "k--", lw=0.8, label="линейно"); axes[0].legend(fontsize=8)
    axes[0].set_title("Время шага: 3.8 мкс/частица (numpy), цель ×10 с C/OpenMP", fontsize=9)
    axes[1].loglog(n, rss, "o-", label="измерено"); axes[1].axhline(3700, c="r", ls="--", label="RAM этого VPS (3.7 GB)")
    bad = [r for r in rows if "error" in r]
    if bad: axes[1].scatter([bad[0]["n_part"]], [3030], marker="X", s=100, c="r", label="OOM-killed (2·10⁷ частиц)")
    for lab, gb in (("CCX33: 4 GB/ядро", 4096), ("AX102-1: 8 GB/ядро", 8192)): axes[1].axhline(gb, ls=":", c="gray"); axes[1].text(2.5e4, gb * 1.08, lab, fontsize=7)
    axes[1].set_xlabel("число частиц"); axes[1].set_ylabel("пик RSS, MB"); axes[1].legend(fontsize=8, loc="lower right")
    axes[1].set_title("Память: ≈280 B/частица, дочерние частицы дают до ×8", fontsize=9)
    fig.tight_layout(); fig.savefig(FIG / "fig8_resources.png"); plt.close(fig)

if __name__ == "__main__":
    import sys
    todo = sys.argv[1:] or ["outcome_map", "timeseries", "spacetime", "phase_space", "profiles", "decay", "validation", "resources"]
    for name in todo:
        globals()["fig_" + name](); print("done", name, flush=True)