import sys, pathlib, time, warnings
sys.path.insert(0, str(pathlib.Path.cwd().parent))
import numpy as np
import matplotlib.pyplot as plt
import bayeskit as bk
from bayeskit import problems as P, samplers as SM, diagnostics as D, plotting as BP, style as S
from bayeskit import afterslip as A
import jax, jax.numpy as jnp
import diffrax as dfx
from scipy.interpolate import RBFInterpolator
from scipy.linalg import cho_factor, cho_solve

T0 = time.time()
SIG = 2.0                                                    # GPS noise (mm)
C_FULL, C_GAL, C_DEIM, C_RBF = S.INK, S.BLUE, S.VIOLET, S.ORANGE

la_all, slip_all, surf_all, setup = A.training_set_fault(1500, seed=0)
zc, free, K, G40 = setup["zc"], setup["free"], setup["K"], setup["G"]
Kff, dtf, Gf = K[np.ix_(free, free)], setup["dtau"][free], G40[:, free]
sf = slip_all[:, :, free]                                    # (run, epoch, free patch), metres
T = A.T_OBS
TR, TE = np.arange(1200), np.arange(1200, 1500)
nf = int(free.sum())
print(f"{nf} free patches, {len(T)} epochs; snapshot matrix from the training runs: {len(TR) * len(T):,} profiles of {nf} numbers")


def surf_rms(s_pred, s_true):
    """rms surface error (mm) for free-patch slip arrays (..., nf)."""
    return np.sqrt(np.mean((1000.0 * (s_pred - s_true) @ Gf.T) ** 2))

31 free patches, 40 epochs; snapshot matrix from the training runs: 48,000 profiles of 31 numbers

Smat = sf[TR].reshape(-1, nf)
_, sv, Vt = np.linalg.svd(Smat, full_matrices=False)
Vt *= np.sign(Vt.sum(axis=1))[:, None]                       # a mode's sign is arbitrary; make each sum positive
energy_left = 1 - np.cumsum(sv ** 2) / np.sum(sv ** 2)
RS = np.arange(1, 17)
proj_err = [surf_rms(sf[TE] @ Vt[:r].T @ Vt[:r], sf[TE]) / SIG for r in RS]   # state (slip) projection


def surface_best(r):
    """The r-mode slip whose surface displacement is closest to the true one: Psi b*, b* = pinv(G Psi) G s."""
    Psi = Vt[:r].T
    M = Psi @ np.linalg.pinv(Gf @ Psi, rcond=1e-10) @ Gf       # (nf, nf): slip -> best r-mode slip for the surface
    # rcond drops directions whose surface signature is below 1e-10 of the largest; without it round-off in a nearly
    # singular G Psi makes the error grow again with r, although the best fit over nested sets of modes cannot get worse
    return sf[TE] @ M.T


surf_best_err = [surf_rms(surface_best(r), sf[TE]) / SIG for r in RS]
for r in [1, 2, 4, 8, 12]:
    print(f"r = {r:2d}: energy left {energy_left[r - 1]:.1e}; held-out rms surface error of the state projection "
          f"{proj_err[r - 1]:.3f} σ, of the best surface fit with the same modes {surf_best_err[r - 1]:.1e} σ")

r =  1: energy left 9.6e-02; held-out rms surface error of the state projection 17.663 σ, of the best surface fit with the same modes 4.9e+00 σ
r =  2: energy left 2.3e-02; held-out rms surface error of the state projection 4.919 σ, of the best surface fit with the same modes 9.8e-01 σ
r =  4: energy left 3.5e-03; held-out rms surface error of the state projection 0.590 σ, of the best surface fit with the same modes 1.4e-02 σ
r =  8: energy left 2.1e-04; held-out rms surface error of the state projection 0.027 σ, of the best surface fit with the same modes 5.0e-06 σ
r = 12: energy left 1.5e-05; held-out rms surface error of the state projection 0.001 σ, of the best surface fit with the same modes 7.3e-08 σ

fig, axs = plt.subplots(1, 3, figsize=(14.5, 3.9))
axs[0].semilogy(RS, energy_left[:len(RS)], "o-", color=S.INK, ms=4)
axs[0].set_xlabel("number of modes r"); axs[0].set_ylabel("share of energy not captured")
axs[0].set_title("Singular values: energy left out")
for k, c in enumerate([S.BLUE, S.ORANGE, S.AQUA, S.VIOLET]):
    full = np.full(len(zc), np.nan)
    full[free] = Vt[k]                                        # NaN on pinned patches: the line breaks across 2-11 km
    axs[1].plot(full, zc, color=c, lw=1.6, marker="o", ms=2.5, label=f"mode {k + 1}")
axs[1].axhspan(zc[~free].min() - 0.5, zc[~free].max() + 0.5, color=S.GRID, alpha=0.6, lw=0)
axs[1].axvline(0, color=S.AXIS, lw=0.8)
axs[1].invert_yaxis(); axs[1].set_xlabel("mode value"); axs[1].set_ylabel("depth (km)")
axs[1].set_title("The first four POD modes"); axs[1].legend(loc="lower right")
axs[2].semilogy(RS, proj_err, "o-", color=C_GAL, ms=4, label="POD projection of the slip")
axs[2].semilogy(RS, surf_best_err, "s--", color=S.ORANGE, ms=3.5, label="best surface fit, same r modes")
axs[2].axhline(1.0, color=S.INK2, lw=0.8, ls=":")
axs[2].text(RS[-1], 0.9, "error = GPS noise", ha="right", va="top", fontsize=8, color=S.INK2)
axs[2].set_xlabel("number of modes r"); axs[2].set_ylabel("rms surface error / σ (held-out runs)")
axs[2].set_ylim(1e-4, None)                                   # below this: smaller than the solver's own accuracy
axs[2].set_title("Surface error of r-mode slip"); axs[2].legend(fontsize=8, loc="upper right")
fig.tight_layout()
bk.savefig(fig, "3-06_pod")
plt.show()

saved figures/3-06_pod.png

i0 = TE[0]
s1 = sf[i0, -1]
for r in [2, 4, 8]:
    b = Vt[:r] @ s1
    print(f"r = {r}: coefficients {np.round(b, 3)}; reconstruction rms slip error {1000 * np.sqrt(np.mean((Vt[:r].T @ b - s1) ** 2)):.1f} mm "
          f"(slip rms {1000 * np.sqrt(np.mean(s1 ** 2)):.0f} mm)")

r = 2: coefficients [1.947 0.148]; reconstruction rms slip error 19.9 mm (slip rms 351 mm)
r = 4: coefficients [ 1.947  0.148  0.005 -0.094]; reconstruction rms slip error 10.4 mm (slip rms 351 mm)
r = 8: coefficients [ 1.947  0.148  0.005 -0.094  0.01   0.04   0.029  0.014]; reconstruction rms slip error 4.4 mm (slip rms 351 mm)

Yflat = sf.reshape(len(sf), -1)
xm = la_all[TR].mean(0)
_, sx, Vx = np.linalg.svd(la_all[TR] - xm, full_matrices=False)
R_OUT = 16
KXS = [2, 4, 8, 12, 16, 24]


def rbf_rom(idx_tr, kx, r=R_OUT):
    ym = Yflat[idx_tr].mean(0)
    _, _, Vy = np.linalg.svd(Yflat[idx_tr] - ym, full_matrices=False)
    Zx = (la_all - xm) @ Vx[:kx].T
    rb = RBFInterpolator(Zx[idx_tr], (Yflat[idx_tr] - ym) @ Vy[:r].T, kernel="thin_plate_spline", degree=1)
    return lambda la_: (ym + rb(np.atleast_2d((la_ - xm) @ Vx[:kx].T)) @ Vy[:r]).reshape(-1, len(T), nf)


rbf_err = {}
for n in [300, 1200]:
    rbf_err[n] = [surf_rms(rbf_rom(TR[:n], kx)(la_all[TE]), sf[TE]) / SIG for kx in KXS]
    print(f"{n:5d} training runs: rms surface error / σ for k_x = {KXS}: {np.round(rbf_err[n], 2)}")
left = 1 - np.cumsum(sx ** 2) / np.sum(sx ** 2)
print("share of input variance left out by k_x components: " + ", ".join(f"k_x = {k}: {left[k - 1]:.3f}" for k in KXS))
rom_rbf = rbf_rom(TR, 16)
t1 = time.time()
for i in TE[:50]:
    rom_rbf(la_all[i])
t_rbf = (time.time() - t1) / 50
print(f"one RBF-ROM evaluation (k_x = 16): {1000 * t_rbf:.2f} ms")

  300 training runs: rms surface error / σ for k_x = [2, 4, 8, 12, 16, 24]: [24.13 13.2   5.02  2.85  2.47  2.4 ]

 1200 training runs: rms surface error / σ for k_x = [2, 4, 8, 12, 16, 24]: [24.02 13.22  4.57  2.03  1.68  1.59]
share of input variance left out by k_x components: k_x = 2: 0.340, k_x = 4: 0.109, k_x = 8: 0.019, k_x = 12: 0.006, k_x = 16: 0.002, k_x = 24: 0.001
one RBF-ROM evaluation (k_x = 16): 0.03 ms

fig, axs = plt.subplots(1, 2, figsize=(12.0, 3.9))
for n, c, mk in [(300, S.MUTED, "s"), (1200, C_RBF, "o")]:
    axs[0].semilogy(KXS, rbf_err[n], marker=mk, color=c, lw=1.6, label=f"{n} training runs")
axs[0].axhline(1.0, color=S.INK2, lw=0.8, ls=":")
axs[0].set_xlabel("input dimensions k_x (principal components of the profile)")
axs[0].set_ylabel("rms surface error / σ (held-out runs)"); axs[0].set_title("Data-driven ROM: error against input dimension")
axs[0].legend()
j_st = 8
s_rb = rom_rbf(la_all[i0])[0]
axs[1].plot(T * 365.25, 1000 * sf[i0] @ Gf[j_st], color=C_FULL, lw=2.6, alpha=0.35, label="solver")
axs[1].plot(T * 365.25, 1000 * s_rb @ Gf[j_st], color=C_RBF, lw=1.4, ls="--", label="RBF ROM, k_x = 16")
axs[1].set_xscale("log"); axs[1].set_xlabel("days after the earthquake")
axs[1].set_ylabel(f"displacement at {setup['x_obs'][j_st]:.1f} km (mm)"); axs[1].set_title("One held-out run at one station")
axs[1].legend()
fig.tight_layout()
bk.savefig(fig, "3-06_data_driven_rom")
plt.show()

saved figures/3-06_data_driven_rom.png

v_snap = 2 * 0.03 * np.sinh((dtf[None, None, :] + sf[TR] @ Kff.T) / np.exp(la_all[TR][:, None, free]))
Uv = np.linalg.svd(v_snap.reshape(-1, nf), full_matrices=False)[2].T
print(f"slip-rate snapshots span {v_snap.min():.1e} to {v_snap.max():.1e} m/yr")


def deim_points(U):
    """Greedy DEIM selection: each new point is where the current interpolant of the next basis vector is worst."""
    p = [int(np.argmax(np.abs(U[:, 0])))]
    for j in range(1, U.shape[1]):
        c = np.linalg.solve(U[p, :j], U[p, j])
        p.append(int(np.argmax(np.abs(U[:, j] - U[:, :j] @ c))))
    return np.array(p)


def make_rom(r, m=None):
    Psi = jnp.asarray(Vt[:r].T)
    KP, d0 = jnp.asarray(Kff) @ Psi, jnp.asarray(dtf)
    if m is None:
        rhs_v = lambda b, a: Psi.T @ (2 * 0.03 * jnp.sinh((d0 + KP @ b) / a))
        sel = jnp.arange(nf)
    else:
        p = deim_points(Uv[:, :m])
        M = jnp.asarray(Vt[:r] @ Uv[:, :m] @ np.linalg.inv(Uv[p, :m]))
        KPp, d0p, sel = KP[p], d0[p], jnp.asarray(p)
        rhs_v = lambda b, a: M @ (2 * 0.03 * jnp.sinh((d0p + KPp @ b) / a))

    @jax.jit
    def solve(la_f):
        a = jnp.exp(la_f)[sel]
        sol = dfx.diffeqsolve(dfx.ODETerm(lambda t, b, args: rhs_v(b, a)), dfx.Kvaerno5(), t0=0.0, t1=T[-1], dt0=1e-8,
                              y0=jnp.zeros(r), saveat=dfx.SaveAt(ts=jnp.asarray(T)),
                              stepsize_controller=dfx.PIDController(rtol=1e-5, atol=1e-7), max_steps=400_000)
        return sol.ys @ Psi.T
    return solve


def run_timed(f, inputs):
    jax.block_until_ready(f(inputs[0]))                      # compile first, then time
    t1 = time.time()
    out = np.stack([np.asarray(f(x)) for x in inputs])
    return out, (time.time() - t1) / len(inputs)


TEST = TE[:20]
fom = jax.jit(lambda l: A.fault_afterslip(l, setup)[:, free])          # bayeskit form: stress and slip on all 40 patches
Kj_, dj_ = jnp.asarray(Kff), jnp.asarray(dtf)


@jax.jit
def fom31(la_f):
    """The same physics as an ODE for the 31 free slips only (the FOM that the ROMs below reduce)."""
    a = jnp.exp(la_f)
    sol = dfx.diffeqsolve(dfx.ODETerm(lambda t, s_, args: 2 * 0.03 * jnp.sinh((dj_ + Kj_ @ s_) / a)), dfx.Kvaerno5(),
                          t0=0.0, t1=T[-1], dt0=1e-8, y0=jnp.zeros(nf), saveat=dfx.SaveAt(ts=jnp.asarray(T)),
                          stepsize_controller=dfx.PIDController(rtol=1e-5, atol=1e-7), max_steps=400_000)
    return sol.ys


s_fom, t_fom = run_timed(fom, [jnp.asarray(la_all[i]) for i in TEST])
s_f31, t_f31 = run_timed(fom31, [jnp.asarray(la_all[i, free]) for i in TEST])
print(f"FOM as in bayeskit (80 unknowns: stress and slip on every patch): {1000 * t_fom:.0f} ms per run; agrees with the "
      f"cache to {1000 * np.abs(s_fom - sf[TEST]).max():.4f} mm")
print(f"FOM as 31 free slips only: {1000 * t_f31:.0f} ms per run; agrees with the cache to "
      f"{1000 * np.abs(s_f31 - sf[TEST]).max():.4f} mm")
gal, deim = {}, {}
for r in [2, 4, 6, 8, 12]:
    out, tt = run_timed(make_rom(r), [jnp.asarray(la_all[i, free]) for i in TEST])
    gal[r] = (surf_rms(out, sf[TEST]) / SIG, tt, surf_rms(sf[TEST] @ Vt[:r].T @ Vt[:r], sf[TEST]) / SIG, out)
    print(f"Galerkin r = {r:2d}: rms surface error {gal[r][0]:.3f} σ (POD projection of the slip {gal[r][2]:.3f} σ), {1000 * tt:.0f} ms per run")
for r, m in [(8, 8), (8, 12), (8, 16)]:
    out, tt = run_timed(make_rom(r, m), [jnp.asarray(la_all[i, free]) for i in TEST])
    deim[(r, m)] = (surf_rms(out, sf[TEST]) / SIG, tt)
    print(f"Galerkin r = {r} + DEIM m = {m:2d}: rms surface error {deim[(r, m)][0]:.3f} σ, {1000 * tt:.0f} ms per run")

slip-rate snapshots span 2.9e-03 to 3.7e+01 m/yr

FOM as in bayeskit (80 unknowns: stress and slip on every patch): 84 ms per run; agrees with the cache to 0.0002 mm
FOM as 31 free slips only: 14 ms per run; agrees with the cache to 0.0008 mm

Galerkin r =  2: rms surface error 14.209 σ (POD projection of the slip 4.740 σ), 10 ms per run

Galerkin r =  4: rms surface error 3.431 σ (POD projection of the slip 0.659 σ), 13 ms per run

Galerkin r =  6: rms surface error 0.705 σ (POD projection of the slip 0.108 σ), 15 ms per run

Galerkin r =  8: rms surface error 0.126 σ (POD projection of the slip 0.030 σ), 16 ms per run

Galerkin r = 12: rms surface error 0.011 σ (POD projection of the slip 0.002 σ), 18 ms per run

Galerkin r = 8 + DEIM m =  8: rms surface error 0.854 σ, 11 ms per run

Galerkin r = 8 + DEIM m = 12: rms surface error 0.193 σ, 14 ms per run

Galerkin r = 8 + DEIM m = 16: rms surface error 0.119 σ, 14 ms per run

fig, axs = plt.subplots(1, 3, figsize=(15.8, 3.9), gridspec_kw=dict(width_ratios=[1.05, 1.25, 1.05]))
rr = sorted(gal)
axs[0].semilogy(rr, [gal[r][2] for r in rr], "o:", color=S.MUTED, ms=4, label="POD projection of the slip")
axs[0].semilogy(rr, [gal[r][0] for r in rr], "o-", color=C_GAL, ms=4, label="Galerkin ROM")
axs[0].semilogy([8, 8, 8], [deim[(8, m)][0] for m in (8, 12, 16)], "D", color=C_DEIM, ms=5,
                label="Galerkin + DEIM, m = 8, 12, 16")
axs[0].axhline(1.0, color=S.INK2, lw=0.8, ls=":")
axs[0].set_xlabel("number of modes r"); axs[0].set_ylabel("rms surface error / σ (20 held-out runs)")
axs[0].set_title("Intrusive ROM: error against modes"); axs[0].legend(fontsize=7.5)
labs = ["FOM, 80\nunknowns", "FOM, 31\nunknowns", "Galerkin\nr = 8", "Galerkin r = 8\nDEIM m = 16", "RBF ROM\nk_x = 16"]
tvals = [t_fom, t_f31, gal[8][1], deim[(8, 16)][1], t_rbf]
axs[1].bar(labs, [1000 * t for t in tvals], color=[S.MUTED, C_FULL, C_GAL, C_DEIM, C_RBF], width=0.6)
axs[1].set_yscale("log"); axs[1].set_ylabel("time per run (ms)"); axs[1].tick_params(axis="x", labelsize=7.5)
for k, t in enumerate(tvals):
    axs[1].text(k, 1000 * t, f"{1000 * t:.2g}", ha="center", va="bottom", fontsize=8)
axs[1].set_title("Cost of one forward run")
k0 = 0
axs[2].plot(T * 365.25, 1000 * s_fom[k0] @ Gf[j_st], color=C_FULL, lw=2.8, alpha=0.3, label="FOM")
for r, ls in [(2, ":"), (4, "--"), (8, "-")]:
    axs[2].plot(T * 365.25, 1000 * gal[r][3][k0] @ Gf[j_st], color=C_GAL, lw=1.2, ls=ls, label=f"Galerkin r = {r}")
axs[2].set_xscale("log"); axs[2].set_xlabel("days after the earthquake")
axs[2].set_ylabel(f"displacement at {setup['x_obs'][j_st]:.1f} km (mm)"); axs[2].set_title("Same held-out run, Galerkin ROMs")
axs[2].legend(fontsize=7.5)
fig.tight_layout()
bk.savefig(fig, "3-06_galerkin_rom")
plt.show()

saved figures/3-06_galerkin_rom.png

la_ood = A.sample_log_asig(np.random.default_rng(3), 20, zc, ell=3.0)
s_ood, _ = run_timed(fom, [jnp.asarray(l) for l in la_ood])          # truth from the bayeskit solver
rom8, rom12 = make_rom(8), make_rom(12)
ood = {
    "Galerkin r = 8": (gal[8][0], surf_rms(np.stack([np.asarray(rom8(jnp.asarray(l[free]))) for l in la_ood]), s_ood) / SIG),
    "Galerkin r = 12": (gal[12][0], surf_rms(np.stack([np.asarray(rom12(jnp.asarray(l[free]))) for l in la_ood]), s_ood) / SIG),
    "RBF ROM, k_x = 16": (surf_rms(rom_rbf(la_all[TEST]), sf[TEST]) / SIG, surf_rms(rom_rbf(la_ood), s_ood) / SIG),
}
for k, (a_, b_) in ood.items():
    print(f"{k:18s}: rms surface error / σ, ordinary held-out {a_:.3f}, rough profiles {b_:.3f} (ratio {b_ / a_:.1f})")

Galerkin r = 8    : rms surface error / σ, ordinary held-out 0.126, rough profiles 0.204 (ratio 1.6)
Galerkin r = 12   : rms surface error / σ, ordinary held-out 0.011, rough profiles 0.043 (ratio 4.1)
RBF ROM, k_x = 16 : rms surface error / σ, ordinary held-out 2.125, rough profiles 3.696 (ratio 1.7)

fig, ax = plt.subplots(figsize=(7.2, 3.6))
names = list(ood)
x = np.arange(len(names))
ax.bar(x - 0.18, [ood[k][0] for k in names], width=0.34, color=[C_GAL, C_GAL, C_RBF], label="ordinary held-out runs")
ax.bar(x + 0.18, [ood[k][1] for k in names], width=0.34, color=[C_GAL, C_GAL, C_RBF], alpha=0.45, hatch="//",
       edgecolor=S.SURFACE, label="rough profiles (ℓ = 3 km)")
ax.set_yscale("log"); ax.axhline(1.0, color=S.INK2, lw=0.8, ls=":")
ax.set_xticks(x); ax.set_xticklabels(names); ax.set_ylabel("rms surface error / σ")
ax.set_title("Errors inside and outside the training range"); ax.legend(fontsize=8)
fig.tight_layout()
bk.savefig(fig, "3-06_rough_inputs")
plt.show()

saved figures/3-06_rough_inputs.png

nt, npch = len(T), len(zc)
dz = np.diff(setup["z_edges"])
Cz = A.matern_cov(zc, 8.0, 1.5, 1.0)
lt = np.log(T)
Ct = 0.6 ** 2 * np.exp(-0.5 * (lt[:, None] - lt[None, :]) ** 2)
Cm = np.kron(Ct, Cz)                                          # (1600, 1600), epoch-major ordering
Gbig = np.kron(np.eye(nt), 1000.0 * G40)                      # mm per m, (1200, 1600)
Cd = SIG ** 2 * np.eye(Gbig.shape[0])
x_q = np.geomspace(0.3, 150.0, 60)
Gq = 1000.0 * A.surface_greens(x_q, setup["z_edges"])
Q = np.zeros((60 + nt, nt * npch))
Q[:60, (nt - 1) * npch:] = Gq                                 # displacement at year 1 (mm)
for k in range(nt):
    Q[60 + k, k * npch:(k + 1) * npch] = dz                    # potency at epoch k (m km)
print(f"unknowns {Cm.shape[0]}, data {Gbig.shape[0]}, quantities of interest {Q.shape[0]}")

t1 = time.time()
CmGt = Cm @ Gbig.T
Sfac = cho_factor(Gbig @ CmGt + Cd)
Aq = cho_solve(Sfac, (Q @ CmGt).T).T                          # (n_q, n_d)
Kgain = cho_solve(Sfac, CmGt.T).T                             # full gain (n_m, n_d), for the slip itself
Cq = Q @ Cm @ Q.T - Aq @ (Q @ CmGt).T
t_off = time.time() - t1
sd_q = np.sqrt(np.diag(Cq))
print(f"offline phase (once): {t_off:.2f} s")


def direct_posterior_q(d):
    """What one would do without the offline phase: build and solve the whole linear system for this data set."""
    CmGt_ = Cm @ Gbig.T
    fac = cho_factor(Gbig @ CmGt_ + Cd)
    return Q @ (CmGt_ @ cho_solve(fac, d))


ir = TE[5]
s_true = slip_all[ir]                                         # (40 epochs, 40 patches), from the friction model
rng = np.random.default_rng(0)
d_new = Gbig @ s_true.ravel() + SIG * rng.standard_normal(Gbig.shape[0])
reps = 2000
t1 = time.perf_counter()
for _ in range(reps):
    q_on = Aq @ d_new
t_on = (time.perf_counter() - t1) / reps
t1 = time.perf_counter()
for _ in range(200):
    m_on = Kgain @ d_new
t_on_m = (time.perf_counter() - t1) / 200
t1 = time.perf_counter()
for _ in range(3):
    q_dir = direct_posterior_q(d_new)
t_dir = (time.perf_counter() - t1) / 3
print(f"online, quantities of interest: {1e6 * t_on:.1f} µs per data set; online, all 1,600 slip values: {1e6 * t_on_m:.0f} µs")
print(f"direct solve for one data set: {1000 * t_dir:.0f} ms  ->  ratio direct / online (QoI) {t_dir / t_on:,.0f}")
print(f"online vs direct: largest difference {np.abs(q_on - q_dir).max():.2e} (largest |q| {np.abs(q_dir).max():.1f}); "
      f"QoI from the full slip mean {np.abs(Q @ m_on - q_on).max():.2e}")

q_true = Q @ s_true.ravel()
inside = []
for i in TE[:100]:
    d_i = Gbig @ slip_all[i].ravel() + SIG * np.random.default_rng(1000 + i).standard_normal(Gbig.shape[0])
    inside.append(np.abs(Aq @ d_i - Q @ slip_all[i].ravel()) <= 1.645 * sd_q)
inside = np.array(inside)
print(f"100 held-out runs: true QoI inside the 90 % interval, displacement {inside[:, :60].mean():.2f}, potency {inside[:, 60:].mean():.2f}")

unknowns 1600, data 1200, quantities of interest 100

offline phase (once): 0.06 s

online, quantities of interest: 4.7 µs per data set; online, all 1,600 slip values: 54 µs
direct solve for one data set: 34 ms  ->  ratio direct / online (QoI) 7,220
online vs direct: largest difference 1.11e-09 (largest |q| 537.6); QoI from the full slip mean 1.14e-09
100 held-out runs: true QoI inside the 90 % interval, displacement 0.86, potency 0.93

fig, axs = plt.subplots(1, 3, figsize=(14.8, 3.9), gridspec_kw=dict(width_ratios=[1, 1.15, 1.15]))
labs = ["offline\n(once)", "direct solve\n(per data set)", "online, all slip\n(per data set)", "online, QoI\n(per data set)"]
tv = [t_off, t_dir, t_on_m, t_on]
axs[0].bar(labs, tv, color=[S.MUTED, S.INK2, S.BLUE, S.BLUE], width=0.62)
axs[0].set_yscale("log"); axs[0].set_ylabel("seconds"); axs[0].tick_params(axis="x", labelsize=7.5)
for k, t in enumerate(tv):
    axs[0].text(k, t, f"{t:.1e} s", ha="center", va="bottom", fontsize=7.5)
axs[0].set_title("Where the time goes")
dq = q_on[:60] - q_true[:60]                                 # the displacement itself spans hundreds of mm, the band a few
axs[1].fill_between(x_q, dq - 1.645 * sd_q[:60], dq + 1.645 * sd_q[:60], color=S.BLUE, alpha=0.2, lw=0, label="posterior 90 %")
axs[1].plot(x_q, dq, color=S.BLUE, lw=1.5, label="posterior mean (online) − truth")
axs[1].axhline(0, color=S.INK, lw=1.0, ls="--", label="truth (friction model)")
lo_ = (dq - 1.645 * sd_q[:60]).min()
axs[1].plot(setup["x_obs"], np.full(30, lo_ * 1.15), "|", color=S.INK2, ms=10, label="GPS stations")
axs[1].set_xscale("log"); axs[1].set_xlabel("distance from the fault (km)"); axs[1].set_ylabel("displacement after 1 yr, minus truth (mm)")
axs[1].set_title("Goal 1: displacement between and beyond stations"); axs[1].legend(fontsize=7.5, loc="upper right")
qp = q_on[60:]
axs[2].fill_between(T * 365.25, qp - 1.645 * sd_q[60:], qp + 1.645 * sd_q[60:], color=S.BLUE, alpha=0.2, lw=0)
axs[2].plot(T * 365.25, qp, color=S.BLUE, lw=1.5, label="posterior mean ± 90 %")
axs[2].plot(T * 365.25, q_true[60:], color=S.INK, lw=1.2, ls="--", label="truth")
axs[2].set_xscale("log"); axs[2].set_xlabel("days after the earthquake"); axs[2].set_ylabel("potency (m km)")
axs[2].set_title("Goal 2: afterslip potency through time"); axs[2].legend(fontsize=7.5)
fig.tight_layout()
bk.savefig(fig, "3-06_realtime")
plt.show()

saved figures/3-06_realtime.png

print(f"notebook run time: {time.time() - T0:.0f} s")

notebook run time: 19 s

problem	unknowns	data or output	forward model	linear?
POD and ROMs (sections 1-4)	$\ln a\sigma(z)$ on 40 patches	afterslip history of the 31 patches that slip, 40 epochs (day 1 to year 1)	1-D antiplane fault, $v = 2v_0\sinh(\tau/a\sigma)$, $\dot\tau = \mathbf K v$ (`bayeskit.afterslip`)	nonlinear
real-time inference (sections 5-6)	afterslip $s(z, t)$ on 40 patches × 40 epochs = 1,600 numbers, no friction law (kinematic), Gaussian prior	30 GPS stations × 40 epochs = 1,200 numbers, σ = 2 mm	$\mathbf d_k = \mathbf G\,\mathbf s_k$ at every epoch	linear

3.06 · Reduced-order models and real-time inference¶

1. Snapshots and POD: how many shapes does afterslip need?¶

2. A data-driven ROM: interpolate the POD coefficients¶

3. An intrusive ROM: project the equations¶

4. Outside the training range: physics against interpolation¶

5. Real-time Bayesian inference: offline once, online in microseconds¶

6. Where the field is¶