VACUA

"""
VACUA lane — analog false-vacuum decay (Jenkins et al., PRD 109, 023506, 2024).

Two-component 1D BEC. The relative phase phi ~ (phi1 - phi2) obeys a relativistic
Klein-Gordon equation in a false-vacuum potential (paper Eq. 8-10). We evolve the
full two-field GPE with a Fourier pseudospectral, symplectic split-step integrator
(paper App. B), reproducing their bubble-nucleation physics.

Code units: hbar = m = gn = 1  =>  healing length xi = 1, sound speed c = 1.
We set n = nbar, g = 1/nbar so that gn = 1 exactly (paper Sec. V).

Model (paper Eq. 2,4,5):
  i d_t psi = [ -1/2 d_xx  - nu(t) sigma_x - mu_lin ] psi          (linear)
             + [ g |psi_i|^2 - mu_nlin ] psi_i                     (nonlinear, per component)
  nu(t) = eps*gn + lam*omega*sqrt(eps/2)*cos(omega t)             (Eq. 5, hbar=1)

The chemical potentials mu_lin, mu_nlin are arbitrary global-phase choices that do
NOT affect the relative phase or the Noether charges; we pick them to keep the
homogeneous field slowly varying.

The linear step is EXACT: kinetic (diagonal in k, scalar in component space) commutes
with the coupling (k-independent sigma_x), and the time-ordered coupling exponential is
exp(i R sigma_x) with R = ∫ nu dt' evaluated analytically over the sub-interval.
Split error lives only between linear and nonlinear pieces, controlled by the symmetric
composition order (2=Strang, 4/6/8 via Suzuki recursive triple-jump).
"""

import numpy as np

FFT = np.fft.fft
IFFT = np.fft.ifft


class Params:
    def __init__(self, nbar, L=500.0, N=2048, eps=2.5e-3, lam=np.sqrt(2.0),
                 omega=680.0, steps_per_period=16):
        self.nbar = float(nbar)
        self.L = float(L)
        self.N = int(N)
        self.eps = float(eps)
        self.lam = float(lam)
        self.omega = float(omega)
        self.g = 1.0 / self.nbar          # so that g*n = 1
        self.n = self.nbar
        self.dx = self.L / self.N
        self.dt = (1.0 / steps_per_period) * (2.0 * np.pi / self.omega)
        # grid
        self.x = np.arange(self.N) * self.dx
        self.k = 2.0 * np.pi * np.fft.fftfreq(self.N, d=self.dx)
        self.k2 = self.k ** 2
        # chemical potentials (global-phase bookkeeping only)
        self.mu_lin = self.eps          # cancels DC coupling on the antiphase state
        self.mu_nlin = self.g * self.n  # = 1, cancels mean nonlinear phase
        # derived physics
        self.m_fv = np.sqrt(4.0 * self.eps * (self.lam ** 2 - 1.0))  # false-vacuum mass, Eq. 11
        self.phi0 = np.sqrt(self.n / 2.0)                            # Eq. 7 (hbar=m=1)
        self.k_nyq = np.pi / self.dx

    def summary(self):
        return (f"nbar={self.nbar:g} L={self.L:g} N={self.N} dx={self.dx:.4f} "
                f"dt={self.dt:.3e} eps={self.eps:g} lam={self.lam:.4f} omega={self.omega:g} "
                f"m_fv={self.m_fv:.4f} k_nyq={self.k_nyq:.3f}")


def R_integral(p, t0, t1):
    """Exact integral of nu(t') over [t0, t1] (Eq. B4 rotation angle)."""
    return p.eps * (t1 - t0) + p.lam * np.sqrt(p.eps / 2.0) * (
        np.sin(p.omega * t1) - np.sin(p.omega * t0))


def linear_step(psi1, psi2, t0, dtau, p, kin_cache):
    """Exact linear propagation over [t0, t0+dtau]. kin_cache keyed by dtau."""
    if dtau in kin_cache:
        kin = kin_cache[dtau]
    else:
        kin = np.exp(-1j * dtau * (0.5 * p.k2 - p.mu_lin))
        kin_cache[dtau] = kin
    a = FFT(psi1) * kin
    b = FFT(psi2) * kin
    R = R_integral(p, t0, t0 + dtau)
    cR, sR = np.cos(R), np.sin(R)
    a2 = cR * a + 1j * sR * b
    b2 = 1j * sR * a + cR * b
    return IFFT(a2), IFFT(b2)


def nonlinear_step(psi1, psi2, dtau, p):
    """Local density-dependent phase rotation (Eq. B3); preserves |psi_i| pointwise."""
    psi1 = psi1 * np.exp(-1j * dtau * (p.g * np.abs(psi1) ** 2 - p.mu_nlin))
    psi2 = psi2 * np.exp(-1j * dtau * (p.g * np.abs(psi2) ** 2 - p.mu_nlin))
    return psi1, psi2


def strang(psi1, psi2, t, dtau, p, kin_cache):
    """2nd-order symmetric base step. Advances time by dtau."""
    psi1, psi2 = nonlinear_step(psi1, psi2, 0.5 * dtau, p)
    psi1, psi2 = linear_step(psi1, psi2, t, dtau, p, kin_cache)
    psi1, psi2 = nonlinear_step(psi1, psi2, 0.5 * dtau, p)
    return psi1, psi2, t + dtau


def composed_step(psi1, psi2, t, dtau, p, order, kin_cache):
    """Symmetric composition of Strang to arbitrary even order (Suzuki triple-jump)."""
    if order == 2:
        return strang(psi1, psi2, t, dtau, p, kin_cache)
    z = 1.0 / (2.0 - 2.0 ** (1.0 / (order - 1)))
    psi1, psi2, t = composed_step(psi1, psi2, t, z * dtau, p, order - 2, kin_cache)
    psi1, psi2, t = composed_step(psi1, psi2, t, (1.0 - 2.0 * z) * dtau, p, order - 2, kin_cache)
    psi1, psi2, t = composed_step(psi1, psi2, t, z * dtau, p, order - 2, kin_cache)
    return psi1, psi2, t


# ---------- diagnostics (Noether charges Eq. B6; relative-phase order parameter) ----------

def charges(psi1, psi2, p):
    """Total atom number N and momentum P (paper Eq. B6)."""
    N = (np.sum(np.abs(psi1) ** 2) + np.sum(np.abs(psi2) ** 2)) * p.dx
    d1 = IFFT(1j * p.k * FFT(psi1))
    d2 = IFFT(1j * p.k * FFT(psi2))
    P = (np.sum(np.imag(np.conj(psi1) * d1)) +
         np.sum(np.imag(np.conj(psi2) * d2))) * p.dx
    return N, P


def cos_phi(psi1, psi2):
    """Volume-averaged cos(phi/phi0) = <cos(arg psi1 - arg psi2)>. FV -> -1, TV -> +1."""
    rel = psi1 * np.conj(psi2)
    return np.mean(np.cos(np.angle(rel)))


# ---------- driver ----------

def evolve(psi1, psi2, p, t_end, order=8, record_every=None, t0=0.0):
    """Evolve to ct/xi = t_end. Returns final state + recorded diagnostics."""
    kin_cache = {}
    nsteps = int(round((t_end - t0) / p.dt))
    rec = {"t": [], "N": [], "P": [], "cos": []}
    if record_every is None:
        record_every = max(1, nsteps // 200)
    t = t0
    N0, _ = charges(psi1, psi2, p)
    for i in range(nsteps):
        if i % record_every == 0:
            N, P = charges(psi1, psi2, p)
            rec["t"].append(t)
            rec["N"].append(N)
            rec["P"].append(P)
            rec["cos"].append(cos_phi(psi1, psi2))
        psi1, psi2, t = composed_step(psi1, psi2, t, p.dt, p, order, kin_cache)
    N, P = charges(psi1, psi2, p)
    rec["t"].append(t); rec["N"].append(N); rec["P"].append(P)
    rec["cos"].append(cos_phi(psi1, psi2))
    for kk in rec:
        rec[kk] = np.array(rec[kk])
    rec["N0"] = N0
    return psi1, psi2, rec


def false_vacuum_meanfield(p):
    """Homogeneous antiphase false vacuum: psi1 = +sqrt(n), psi2 = -sqrt(n)."""
    amp = np.sqrt(p.n)
    psi1 = np.full(p.N, amp, dtype=complex)
    psi2 = np.full(p.N, -amp, dtype=complex)
    return psi1, psi2

"""
Bogoliubov vacuum initial conditions for the analog false vacuum (paper Sec. III, App. A).

The relative-phase field phi has the quantum-vacuum fluctuation power spectrum that
interpolates between Klein-Gordon (IR, xi k << 1) and white-noise (UV, xi k >> 1).
We sample it classically (truncated-Wigner) to seed the lattice simulations (Sec. V).

Code units hbar = m = gn = 1, xi = 1, c = 1.

Dispersion (Eq. 20):
  omega_k = 1/2 * sqrt(k^2 + 4 eps (lam^2 - 1)) * sqrt(k^2 + 4 - 4 eps (lam^2 + 1))
Bogoliubov coefficients (Eq. A2), with X = (1/(2 omega_k)) (k^2 + 2 - 4 eps):
  u_k^2 = 1/2 (X + 1),   v_k^2 = 1/2 (X - 1),   u^2 - v^2 = 1
Vacuum spectra (derived; check the KG and UV limits):
  P_phi(k) = 1/4 (u+v)^2   ->  1/(2 omega_k) in IR,  ->  1/4 in UV   (matches Eq. 13 / Eq. 21)
  P_pi(k)  = (u-v)^2       (conjugate momentum; P_phi * P_pi = 1/4, vacuum-saturated)

Crucially (Sec. V): the quasiparticle amplitudes a_k, a_{-k} are independent, but this
mixes phi_k and phi_{-k}; and BOTH the phase (phi) and momentum (pi/relative density)
quadratures must be seeded, else the state is excited and over-nucleates.
"""

import numpy as np

FFT = np.fft.fft
IFFT = np.fft.ifft

K_UV = 3.22  # UV truncation of the ICs (paper: xi k_UV ~= 3.22)


def dispersion(k, p):
    """omega_k, Eq. 20 (gn = hbar = xi = 1)."""
    a = k ** 2 + 4.0 * p.eps * (p.lam ** 2 - 1.0)
    b = k ** 2 + 4.0 - 4.0 * p.eps * (p.lam ** 2 + 1.0)
    return 0.5 * np.sqrt(a) * np.sqrt(b)


def bogoliubov_uv(k, p):
    """u_k, v_k (Eq. A2). Returns arrays; robust at k=0."""
    w = dispersion(k, p)
    X = (k ** 2 + 2.0 - 4.0 * p.eps) / (2.0 * w)
    u2 = 0.5 * (X + 1.0)
    v2 = 0.5 * (X - 1.0)
    v2 = np.clip(v2, 0.0, None)  # guard tiny negative from roundoff
    return np.sqrt(u2), np.sqrt(v2)


def spectra(k, p):
    """P_phi(k) = 1/4 (u+v)^2, P_pi(k) = (u-v)^2."""
    u, v = bogoliubov_uv(k, p)
    P_phi = 0.25 * (u + v) ** 2
    P_pi = (u - v) ** 2
    return P_phi, P_pi


def _grf(rng, M, target_k, cutoff_mask):
    """Vectorized real Gaussian random field(s) with discrete spectrum target_k.
    Recipe: white ~ N(0,1) real -> fft is hermitian with <|W_k|^2> = N;
    multiply by sqrt(target_k) (symmetric in k) -> ifft is real with
    <|fft(field)_k|^2> = N * target_k. Measuring |fft(field)_k|^2 / N recovers target_k
    exactly, with NO free normalization constant."""
    N = target_k.size
    white = rng.standard_normal((M, N))
    Wk = FFT(white, axis=1) * (np.sqrt(target_k) * cutoff_mask)[None, :]
    return IFFT(Wk, axis=1).real  # real field(s), shape (M, N)


def generate_fields(p, M=1, seed=0, cutoff=True, white_noise=False):
    """Draw M realizations of the relative phase field phi(x) and conjugate pi(x).
    white_noise=True uses the nonrelativistic UV spectrum everywhere (u=1, v=0):
    P_phi = 1/4 const -> the 'excited' state the paper contrasts against (Fig. 6)."""
    rng = np.random.default_rng(seed)
    k = p.k
    if white_noise:
        P_phi = 0.25 * np.ones_like(k)
        P_pi = np.ones_like(k)
    else:
        P_phi, P_pi = spectra(k, p)
    mask = (np.abs(k) <= K_UV).astype(float) if cutoff else np.ones_like(k)
    phi = _grf(rng, M, P_phi, mask)
    pi = _grf(rng, M, P_pi, mask)
    return phi, pi


def measure_spectrum(field, p):
    """Discrete power spectrum P(k) = <|fft(field)_k|^2> / N, ensemble-averaged over axis 0."""
    Fk = FFT(field, axis=1)
    return np.mean(np.abs(Fk) ** 2, axis=0) / p.N


def seed_state(p, phi, pi):
    """Map one (phi, pi) realization onto the two-component field around the antiphase
    false vacuum. Relative phase angle theta_rel = phi/phi0 (split +/- between species);
    relative density from the conjugate momentum. Amplitude ~ 1/sqrt(nbar) via phi0 (Eq. 23)."""
    amp = np.sqrt(p.n)
    theta = phi / p.phi0                 # relative phase angle fluctuation
    # conjugate momentum -> relative density perturbation (canonical: delta_n_rel = 2 phi0 * pi)
    dn_rel = 2.0 * p.phi0 * pi
    n1 = p.n + 0.5 * dn_rel
    n2 = p.n - 0.5 * dn_rel
    psi1 = np.sqrt(np.clip(n1, 0.0, None)) * np.exp(1j * 0.5 * theta)
    psi2 = -np.sqrt(np.clip(n2, 0.0, None)) * np.exp(-1j * 0.5 * theta)
    return psi1.astype(complex), psi2.astype(complex)

"""
Vectorized batch evolution: many truncated-Wigner realizations at once.
The evolution steps in solver.py already act on the last axis, so a batch of shape
(M, N) evolves with the same code; only the reductions (charges, <cos>) need an axis.
Used by G3 (find a nucleating trajectory) and G4 (survival probability -> decay rate).
"""
import numpy as np
from solver import Params, composed_step
from bogoliubov import generate_fields, seed_state

FFT = np.fft.fft


def cos_phi_batch(psi1, psi2):
    """<cos(phi/phi0)>_V per realization -> shape (M,)."""
    rel = psi1 * np.conj(psi2)
    return np.mean(np.cos(np.angle(rel)), axis=-1)


def make_batch(p, M, seed=0, white_noise=False):
    """M realizations of the Bogoliubov (or white-noise) vacuum around the false vacuum."""
    phi, pi = generate_fields(p, M=M, seed=seed, cutoff=True, white_noise=white_noise)
    return seed_state(p, phi, pi)  # (M, N) each


def evolve_batch(psi1, psi2, p, t_end, order=2, record_every=None, keep_field=False):
    """Evolve a batch to ct/xi = t_end.
    Returns t (nrec,), cos (nrec, M), and optionally the cos-field of realization 0
    at each recorded time (nrec, N) for the space-time diagram."""
    kin_cache = {}
    nsteps = int(round(t_end / p.dt))
    if record_every is None:
        record_every = max(1, nsteps // 300)
    t = 0.0
    ts, coss, fields = [], [], []
    for i in range(nsteps):
        if i % record_every == 0:
            ts.append(t)
            coss.append(cos_phi_batch(psi1, psi2))
            if keep_field:
                rel0 = psi1[0] * np.conj(psi2[0])
                fields.append(np.cos(np.angle(rel0)))
        psi1, psi2, t = composed_step(psi1, psi2, t, p.dt, p, order, kin_cache)
    ts.append(t); coss.append(cos_phi_batch(psi1, psi2))
    if keep_field:
        rel0 = psi1[0] * np.conj(psi2[0])
        fields.append(np.cos(np.angle(rel0)))
    out = (np.array(ts), np.array(coss))
    if keep_field:
        out = out + (np.array(fields),)
    return out

"""
JAX port of the VACUA split-step solver: jit + lax.scan over timesteps, batched over
realizations (M, N). Same physics as solver.py (validated at G0), but ~orders of magnitude
faster on the ensemble, and differentiable (for Phase 2 inference).

Order-2 (Strang) is used for ensembles -- G0 confirmed the composition orders; Strang at the
paper's dt is the production integrator for the decay statistics.
"""
import jax
jax.config.update("jax_enable_x64", True)  # physics needs float64
import jax.numpy as jnp
from functools import partial
from solver import Params  # reuse grid/constant computation


def make_evolver(p):
    """Build a jitted chunk-evolver + cos diagnostic for a given Params."""
    k2 = jnp.asarray(p.k2)
    dt = float(p.dt)
    kin = jnp.exp(-1j * dt * (0.5 * k2 - p.mu_lin))
    eps, lam, omega = float(p.eps), float(p.lam), float(p.omega)
    g, mu_nlin = float(p.g), float(p.mu_nlin)
    amp_ac = lam * (eps / 2.0) ** 0.5

    def body(carry, _):
        psi1, psi2, t = carry
        # half nonlinear
        psi1 = psi1 * jnp.exp(-1j * (0.5 * dt) * (g * jnp.abs(psi1) ** 2 - mu_nlin))
        psi2 = psi2 * jnp.exp(-1j * (0.5 * dt) * (g * jnp.abs(psi2) ** 2 - mu_nlin))
        # linear (exact): kinetic phase + 2x2 coupling rotation
        a = jnp.fft.fft(psi1, axis=-1) * kin
        b = jnp.fft.fft(psi2, axis=-1) * kin
        R = eps * dt + amp_ac * (jnp.sin(omega * (t + dt)) - jnp.sin(omega * t))
        cR, sR = jnp.cos(R), jnp.sin(R)
        a2 = cR * a + 1j * sR * b
        b2 = 1j * sR * a + cR * b
        psi1 = jnp.fft.ifft(a2, axis=-1)
        psi2 = jnp.fft.ifft(b2, axis=-1)
        # half nonlinear
        psi1 = psi1 * jnp.exp(-1j * (0.5 * dt) * (g * jnp.abs(psi1) ** 2 - mu_nlin))
        psi2 = psi2 * jnp.exp(-1j * (0.5 * dt) * (g * jnp.abs(psi2) ** 2 - mu_nlin))
        return (psi1, psi2, t + dt), None

    @partial(jax.jit, static_argnames=("stride",))
    def evolve_chunk(psi1, psi2, t, stride):
        (psi1, psi2, t), _ = jax.lax.scan(body, (psi1, psi2, t), None, length=stride)
        return psi1, psi2, t

    return evolve_chunk


def cos_phi(psi1, psi2):
    return jnp.mean(jnp.cos(jnp.angle(psi1 * jnp.conj(psi2))), axis=-1)


def run_ensemble(p, psi1, psi2, t_end, nrec=300):
    """Evolve a batch (M, N) to ct/xi=t_end, recording <cos> at nrec points.
    Returns ts (nrec+1,), coss (nrec+1, M) as numpy."""
    import numpy as np
    evolve_chunk = make_evolver(p)
    nsteps = int(round(t_end / p.dt))
    stride = max(1, nsteps // nrec)
    psi1 = jnp.asarray(psi1); psi2 = jnp.asarray(psi2)
    t = 0.0
    ts = [0.0]; coss = [np.asarray(cos_phi(psi1, psi2))]
    done = 0
    while done < nsteps:
        s = min(stride, nsteps - done)
        psi1, psi2, tt = evolve_chunk(psi1, psi2, t, s)
        t = float(tt); done += s
        ts.append(t); coss.append(np.asarray(cos_phi(psi1, psi2)))
    return np.array(ts), np.array(coss)

"""
Gate G4 (HEADLINE) -- survival probability and vacuum decay rate (paper Fig. 6).
For each nbar and IC type (Bogoliubov vacuum vs nonrelativistic white-noise), evolve an
ensemble, extract per-realization nucleation times, build P(survive; t) ~ exp(-Gamma t),
fit Gamma, and show:
  (i)  Gamma decreases (roughly exponentially) with nbar  [Eq. 22 / Fig. 6 right]
  (ii) white-noise ICs decay FASTER than the true vacuum ICs at matched nbar [Fig. 6]
       (white noise is an excited state -> over-nucleates; Sec. V).

Reduced-statistics laptop run by default (M per point, fast-decaying nbar band). Honest:
this samples the fast end of the paper's nbar=10-50 range; the TREND is the reproduction.
Scale M and NBARS up (or move to GPU/JAX) for publication statistics.
"""
import sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from solver import Params
from ensemble import make_batch, evolve_batch

THRESH = -0.3  # <cos> climbing through this = a bubble has nucleated & is growing


def nucleation_times(ts, coss):
    """First time each realization's <cos> crosses above THRESH; inf if it survived."""
    M = coss.shape[1]
    tnuc = np.full(M, np.inf)
    for j in range(M):
        idx = np.argmax(coss[:, j] > THRESH)
        if coss[idx, j] > THRESH:
            tnuc[j] = ts[idx]
    return tnuc


def survival_curve(tnuc, ts):
    return np.array([np.mean(tnuc > t) for t in ts])


def fit_gamma(ts, surv, lo=0.02, hi=0.5):
    """Fit P ~ exp(-Gamma t) over the survival window [lo, hi] (paper uses ~[1%,50%])."""
    m = (surv <= hi) & (surv >= lo)
    if m.sum() < 3:
        return np.nan
    A = np.polyfit(ts[m], np.log(surv[m]), 1)
    return -A[0]


def run_point(nbar, M, t_end, seed, white_noise):
    p = Params(nbar=nbar)
    psi1, psi2 = make_batch(p, M, seed=seed, white_noise=white_noise)
    ts, coss = evolve_batch(psi1, psi2, p, t_end, order=2)
    tnuc = nucleation_times(ts, coss)
    surv = survival_curve(tnuc, ts)
    gamma = fit_gamma(ts, surv)
    frac = np.mean(np.isfinite(tnuc))
    return ts, surv, gamma, frac


def main(M=96, nbars=(8, 10, 12, 15), tag="reduced"):
    tend = {8: 450, 10: 600, 12: 800, 15: 1100, 20: 1500}
    results = {"vacuum": {}, "white": {}}
    plt.figure(figsize=(6, 4))
    for kind, white in [("vacuum", False), ("white", True)]:
        for i, nbar in enumerate(nbars):
            te = tend.get(nbar, 60 * nbar)
            ts, surv, gamma, frac = run_point(nbar, M, te, seed=100 + nbar, white_noise=white)
            results[kind][nbar] = {"gamma": float(gamma), "decayed_frac": float(frac)}
            print(f"[{kind:6s}] nbar={nbar:2d} t_end={te:4d} decayed={frac*100:4.0f}%  "
                  f"Gamma={gamma:.3e}")
            ls = "-" if kind == "vacuum" else "--"
            plt.semilogy(ts, np.clip(surv, 1e-3, 1), ls, color=f"C{i}",
                         label=f"{kind} n={nbar}" if i < 2 else None)
            sys.stdout.flush()
    plt.ylim(1e-2, 1.2); plt.xlabel(r"$ct/\xi$"); plt.ylabel("P(survive)")
    plt.title("VAC G4 -- survival probability (cf. Fig. 6 left)")
    plt.legend(fontsize=7, ncol=2); plt.tight_layout()
    plt.savefig("figs/g4_survival.png", dpi=130); plt.close()

    # Gamma(nbar): vacuum vs white
    plt.figure(figsize=(6, 4))
    for kind, mk in [("vacuum", "o-"), ("white", "s--")]:
        ns = [n for n in nbars if np.isfinite(results[kind][n]["gamma"])]
        gs = [results[kind][n]["gamma"] for n in ns]
        plt.semilogy(ns, gs, mk, label=kind)
    plt.xlabel(r"$\bar n$"); plt.ylabel(r"$\Gamma$ (decay rate)")
    plt.title("VAC G4 -- decay rate vs amplitude (cf. Fig. 6 right)")
    plt.legend(); plt.tight_layout(); plt.savefig("figs/g4_gamma_vs_nbar.png", dpi=130); plt.close()

    np.savez(f"data/g4_{tag}.npz", results=np.array([results], dtype=object),
             nbars=np.array(nbars), M=M)
    # checks
    gv = [results["vacuum"][n]["gamma"] for n in nbars]
    gw = [results["white"][n]["gamma"] for n in nbars]
    mono = all(np.diff([g for g in gv if np.isfinite(g)]) < 0)  # Gamma falls with nbar
    faster = np.nanmean(gw) > np.nanmean(gv)                    # white decays faster
    print("-" * 60)
    print(f"G4 (i)  Gamma decreases with nbar : {'PASS' if mono else 'CHECK'}")
    print(f"G4 (ii) white-noise decays faster : {'PASS' if faster else 'CHECK'}")
    print("figures -> vacua/figs/g4_survival.png, g4_gamma_vs_nbar.png")


if __name__ == "__main__":
    M = int(sys.argv[1]) if len(sys.argv) > 1 else 96
    main(M=M)