Earthquakes that cause earthquakes: Hawkes processes, ETAS and a 7-day aftershock forecast for Ridgecrest 2019

def G_omori(x, c, p):
    """int_0^x (s + c)^-p ds (Omori-Utsu), written to stay accurate near p = 1."""
    L = np.log1p(x / c)
    z = (1 - p) * L
    zs = np.where(np.abs(z) < 1e-8, 1.0, z)
    return c ** (1 - p) * L * np.where(np.abs(z) < 1e-8, 1 + z / 2, np.expm1(zs) / zs)


def Ginv_omori(y, c, p):
    """Inverse of G_omori: the delay x with G(x) = y."""
    if abs(1 - p) < 1e-8:
        return c * np.expm1(y / c)
    return c * np.expm1(np.log1p(y * (1 - p) * c ** (p - 1)) / (1 - p))


def G_exp(x, c):
    """int_0^x exp(-s / c) ds."""
    return -c * np.expm1(-x / c)


MMAX = 8.0  # Gutenberg-Richter truncation used throughout


def trunc_gr(rng, beta, lo, hi, size):
    """Gutenberg-Richter (exponential) magnitudes on [lo, hi]."""
    u = rng.uniform(size=size)
    return lo - np.log1p(-u * -np.expm1(-beta * (hi - lo))) / beta


def simulate_etas(par, hist_t, hist_m, T0, T1, Mc, rng, kernel="omori", max_events=200_000):
    """Events in (T0, T1]: background + offspring of the history + their cascades.
    par holds mu, K0, alpha, beta, c (and p for the Omori kernel); the kernel is K0 e^{alpha (m-Mc)} g."""
    mu, K0, al, beta, c = par["mu"], par["K0"], par["alpha"], par["beta"], par["c"]
    if kernel == "omori":
        Gf, Gi = (lambda x: G_omori(x, c, par["p"])), (lambda y: Ginv_omori(y, c, par["p"]))
    else:
        Gf, Gi = (lambda x: G_exp(x, c)), (lambda y: -c * np.log1p(-y / c))
    nb = rng.poisson(mu * (T1 - T0))
    gen_t, gen_m = rng.uniform(T0, T1, nb), trunc_gr(rng, beta, Mc, MMAX, nb)
    ht, hm = np.asarray(hist_t, float), np.asarray(hist_m, float)
    if ht.size:  # children of the history that fall inside (T0, T1]
        a, b = Gf(np.clip(T0 - ht, 0, None)), Gf(T1 - ht)
        k = rng.poisson(K0 * np.exp(al * (hm - Mc)) * (b - a))
        a, b, tp = np.repeat(a, k), np.repeat(b, k), np.repeat(ht, k)
        ch = tp + Gi(a + rng.uniform(size=tp.size) * (b - a))
        gen_t = np.concatenate([gen_t, ch])
        gen_m = np.concatenate([gen_m, trunc_gr(rng, beta, Mc, MMAX, ch.size)])
    out_t, out_m, exploded = [], [], False
    while gen_t.size:
        out_t.append(gen_t); out_m.append(gen_m)
        if sum(x.size for x in out_t) > max_events:
            exploded = True
            break
        b = Gf(T1 - gen_t)
        k = rng.poisson(K0 * np.exp(al * (gen_m - Mc)) * b)
        b, tp = np.repeat(b, k), np.repeat(gen_t, k)
        gen_t = tp + Gi(rng.uniform(size=tp.size) * b)
        gen_m = trunc_gr(rng, beta, Mc, MMAX, gen_t.size)
    if not out_t:
        return np.array([]), np.array([]), exploded
    tt, mm = np.concatenate(out_t), np.concatenate(out_m)
    o = np.argsort(tt)
    return tt[o], mm[o], exploded


# Poisson vs Hawkes with the same long-run rate (1 per unit time), exponential kernel
n_br, scale, T_demo = 0.8, 0.3, 60.0
hawkes_par = dict(mu=1.0 * (1 - n_br), K0=n_br / scale, alpha=0.0, beta=1.0, c=scale)
th, _, _ = simulate_etas(hawkes_par, [], [], 0.0, T_demo, 0.0, np.random.default_rng(1), kernel="exp")
tp_ = np.sort(np.random.default_rng(2).uniform(0, T_demo, np.random.default_rng(3).poisson(T_demo)))
grid = np.linspace(0, T_demo, 3000)
lam_h = hawkes_par["mu"] + hawkes_par["K0"] * np.where(
    grid[:, None] > th[None, :], np.exp(-(grid[:, None] - th[None, :]) / scale), 0).sum(1)

fig, axes = plt.subplots(3, 1, figsize=(12, 6.5), sharex=True, height_ratios=[1, 1, 2])
axes[0].eventplot(tp_, colors=GREY, lineoffsets=0, linelengths=0.8)
axes[0].set(yticks=[], title=f"Poisson, rate 1: {tp_.size} events")
axes[1].eventplot(th, colors=ORANGE, lineoffsets=0, linelengths=0.8)
axes[1].set(yticks=[], title=f"Hawkes, background 0.2 + branching ratio 0.8 (same long-run rate): {th.size} events")
axes[2].plot(grid, lam_h, color=ORANGE, lw=1)
axes[2].axhline(hawkes_par["mu"], color=INK, ls=":", lw=1)
axes[2].text(0.5, hawkes_par["mu"] + 0.2, "background $\\mu$", fontsize=9)
axes[2].set(xlabel="time", ylabel="Hawkes intensity $\\lambda(t)$");

def aki_b(m, Mc, dm=0.01):
    m = np.asarray(m)
    m = m[m >= Mc - 1e-9]
    return np.log10(np.e) / (m.mean() - (Mc - dm / 2)), m.size


def mc_steps(S, T, Mc, mains, a=4.5, b=0.75, dm=0.1):
    """Piecewise-constant completeness Mc(t) = max(Mc, M - a - b log10(t - t_M)) on [S, T], rounded UP
    to dm. Returns (breaks, levels): levels[k] holds on [breaks[k], breaks[k+1])."""
    br = {S, T}
    for tm, M in mains:
        for lev in np.arange(Mc, M + 1e-9, dm):
            tt = tm + 10 ** ((M - a - lev) / b)
            if S < tt < T:
                br.add(tt)
        if S < tm < T:
            br.add(tm)
    br = np.array(sorted(br))
    mid = 0.5 * (br[:-1] + br[1:])
    lev = np.full(mid.shape, Mc)
    for tm, M in mains:
        after = mid > tm
        f = M - a - b * np.log10(np.where(after, mid - tm, 1.0))
        f = np.ceil(np.round(f / dm, 6)) * dm
        lev = np.where(after, np.maximum(lev, np.minimum(f, M)), lev)
    return br, lev


def mc_at(t, steps):
    br, lev = steps
    return lev[np.clip(np.searchsorted(br, t, side="left") - 1, 0, len(lev) - 1)]


MC = 3.0
MAINS = [(t64, 6.4), (0.0, 7.1)]
steps_demo = mc_steps(S0, 30.0, MC, MAINS)

bins_t = [(-1400, t64, "2016 to the M6.4"), (0, 0.01, "M7.1 + 0-15 min"), (0.01, 0.1, "15 min - 2.4 h"),
          (0.1, 1, "2.4 h - 1 day"), (1, 8, "days 1-8")]
rows = []
for lo, hi, lab in bins_t:
    mm = cat.mag[(cat.t >= lo) & (cat.t < hi)].to_numpy()
    h, e = np.histogram(mm, bins=np.arange(1.95, 7.25, 0.1))
    b3, n3 = aki_b(mm, 3.0)
    rows.append(dict(window=lab, events=mm.size, Mc_maxcurv=round(e[np.argmax(h)] + 0.05, 1),
                     n_M3=n3, b_above_3=round(b3, 2)))
display(pd.DataFrame(rows).set_index("window"))

fig, axes = plt.subplots(1, 2, figsize=(13, 4.6))
ax = axes[0]
aft = cat[(cat.t > 1e-4) & (cat.t < 30)]
ax.scatter(aft.t, aft.mag, s=2, color=INK, alpha=0.35)
tt_ = np.logspace(-4, np.log10(30), 400)
ax.step(tt_, mc_at(tt_, steps_demo), where="post", color=RED, lw=2, label="$M_c(t)$ used (Helmstetter et al. 2006, steps of 0.1)")
ax.axhline(3.0, color=BLUE, ls="--", lw=1.5, label="constant $M_c$ = 3.0")
ax.set(xscale="log", xlabel="days after the M7.1 (log scale)", ylabel="magnitude",
       title="small events are missing in the first hours", ylim=(1.9, 7.3))
ax.legend(fontsize=8)
ax = axes[1]
for (lo, hi, lab), col in zip(bins_t, [GREY, RED, ORANGE, PURPLE, BLUE]):
    mm = np.sort(cat.mag[(cat.t >= lo) & (cat.t < hi)].to_numpy())
    ax.plot(mm[::-1], np.arange(1, mm.size + 1) / mm.size, color=col, lw=2, label=lab)
xx = np.linspace(3, 6, 10)
ax.plot(xx, 0.35 * 10 ** (-(xx - 3.0)), color=INK, ls=":", lw=1.5, label="slope b = 1")
ax.set(yscale="log", xlabel="magnitude M", ylabel="fraction of events with magnitude >= M",
       title="frequency-magnitude distributions", xlim=(2, 7.2))
ax.legend(fontsize=8);

PRIORS = dict(mu=(np.log(0.02), 1.0), log_c=(np.log(0.01), 1.5), p=(1.1, 0.15),
              alpha=(1.8, 0.5), b=(1.0, 0.3), k1=(np.log(0.1), 1.0))


def branching_ratio(k1, c, p, alpha, beta, Mc=MC):
    """Mean number of direct children of an average event (inf when p <= 1)."""
    K0 = k1 / G_omori(365.0, c, p)
    G_inf = np.where(p > 1, c ** (1 - p) / np.where(p > 1, p - 1, 1), np.inf)
    d = beta - alpha
    span = MMAX - Mc
    Em = np.where(np.abs(d) > 1e-9, beta / np.where(np.abs(d) > 1e-9, d, 1) * -np.expm1(-d * span), beta * span)
    Em = Em / -np.expm1(-beta * span)
    return K0 * G_inf * Em


prng = np.random.default_rng(RANDOM_SEED)
npr = 4000
pr = dict(mu=np.exp(prng.normal(*PRIORS["mu"], npr)), c=np.exp(prng.normal(*PRIORS["log_c"], npr)),
          p=prng.normal(*PRIORS["p"], npr), alpha=prng.normal(*PRIORS["alpha"], npr),
          b=stats.truncnorm.rvs((0.3 - 1) / 0.3, np.inf, 1, 0.3, size=npr, random_state=prng),
          k1=np.exp(prng.normal(*PRIORS["k1"], npr)))
pr["beta"] = pr["b"] * np.log(10)
pr["K0"] = pr["k1"] / G_omori(365.0, pr["c"], pr["p"])
direct_day1 = pr["K0"] * np.exp(pr["alpha"] * (7.1 - MC)) * G_omori(1.0, pr["c"], pr["p"])
n_prior = branching_ratio(pr["k1"], pr["c"], pr["p"], pr["alpha"], pr["beta"])
obs_day1 = int(((cat.t > 0) & (cat.t < 1) & (cat.mag >= MC)).sum())

fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
axes[0].hist(np.log10(direct_day1), bins=60, color=BLUE, alpha=0.7)
axes[0].axvline(np.log10(obs_day1), color=RED, lw=2)
axes[0].text(np.log10(obs_day1) + 0.1, axes[0].get_ylim()[1] * 0.85,
             f"{obs_day1} M3+ recorded\n(all generations,\nincomplete)", color=RED, fontsize=8)
axes[0].set(xlabel="log10 direct M3+ aftershocks, day 1", yticks=[],
            title="prior: M7.1 productivity")
nn = np.where(np.isfinite(n_prior), n_prior, 50)
axes[1].hist(np.log10(np.clip(nn, 1e-3, 50)), bins=60, color=PURPLE, alpha=0.7)
axes[1].axvline(0, color=INK, ls="--")
axes[1].set(xlabel="log10 n (p <= 1: right edge)", yticks=[],
            title="prior: branching ratio")
tt_ = np.logspace(-4, 2, 200)
for i in range(60):
    axes[2].plot(tt_, pr["K0"][i] * np.exp(pr["alpha"][i] * 4.1) * (tt_ + pr["c"][i]) ** -pr["p"][i],
                 color=BLUE, alpha=0.25, lw=0.8)
axes[2].set(xscale="log", yscale="log", xlabel="days after an M7.1", ylabel="M3+ aftershocks per day",
            title="prior: direct aftershock rate", ylim=(1e-3, 1e7))
print(f"prior: median direct day-1 aftershocks of an M7.1 = {np.median(direct_day1):.0f} "
      f"(90%: {np.quantile(direct_day1, 0.05):.0f} - {np.quantile(direct_day1, 0.95):.0f}); "
      f"P(n >= 1 or p <= 1) = {np.mean(~(n_prior < 1)):.2f}");

U_NODES = np.arange(-25.0, 14.0 + 1e-9, 0.5)  # log s nodes: 1/s from ~1e-6 to 7e10 days


def pt_G(x, c, p):
    L = pt.log1p(x / c)
    z = (1 - p) * L
    zs = pt.switch(pt.abs(z) < 1e-8, 1.0, z)
    return c ** (1 - p) * L * pt.switch(pt.abs(z) < 1e-8, 1 + z / 2, pt.expm1(zs) / zs)


def etas_model(t, m, Mc, S, T, steps=None, kernel="omori", method="soe", priors=PRIORS):
    """Temporal ETAS (kernel="omori") or exponential Hawkes (kernel="exp") with GR magnitudes.
    t: event times (days, sorted, all <= T), m: magnitudes >= Mc. Every event is a parent; the
    likelihood covers the events in [S, T] with m >= Mc(t), where steps = (breaks, levels) is a
    piecewise-constant completeness threshold (None: constant Mc).
    method: "soe" (mixture-of-exponentials recursion) or "exact" (all pairs)."""
    t, m = np.asarray(t, float), np.asarray(m, float)
    steps = steps if steps is not None else (np.array([S, T]), np.array([Mc]))
    br, lev = steps
    lo, hi = br[:-1], br[1:]
    tgt = np.flatnonzero((t >= S) & (t <= T) & (m >= mc_at(t, steps) - 1e-9))
    dm = m - Mc
    with pm.Model() as model:
        mu = pm.LogNormal("mu", *priors["mu"])
        log_c = pm.Normal("log_c", *priors["log_c"])
        c = pm.Deterministic("c", pt.exp(log_c))
        alpha = pm.Normal("alpha", *priors["alpha"])
        b = pm.TruncatedNormal("b", *priors["b"], lower=0.3)
        beta = b * np.log(10)
        k1 = pm.LogNormal("k1", *priors["k1"])
        ea = pt.exp(alpha * dm)                                    # productivity of every event
        if kernel == "omori":
            p = pm.Normal("p", *priors["p"])
            Gf = lambda x: pt_G(x, c, p)
        else:
            Gf = lambda x: -c * pt.expm1(-x / c)
        K0 = pm.Deterministic("K0", k1 / Gf(365.0))
        # ---- sum over parents for every event: lambda(t_i) = mu + K0 * trig_i
        if method == "exact":
            I, J = np.tril_indices(len(t), -1)                     # (child, parent) pairs
            keep = np.isin(I, tgt)
            I, J = I[keep], J[keep]
            x = t[I] - t[J]
            g = (x + c) ** (-p) if kernel == "omori" else pt.exp(-x / c)
            trig = pt.zeros(len(t))[I].inc(ea[J] * g)[tgt]
        elif kernel == "omori":
            sk = np.exp(U_NODES)
            decay = np.exp(-np.outer(np.diff(t), sk))              # fixed: (n - 1) x K
            wk = pt.exp(p * U_NODES - sk * c - pt.gammaln(p) + np.log(U_NODES[1] - U_NODES[0]))
            A = pytensor.scan(lambda dec, e, a: dec * (a + e), sequences=[pt.as_tensor(decay), ea[:-1]],
                              outputs_info=[pt.zeros(len(sk))], return_updates=False)
            A = pt.concatenate([pt.zeros((1, len(sk))), A], axis=0)  # states just before each event
            trig = (A @ wk)[tgt]
        else:
            A = pytensor.scan(lambda dec, e, a: dec * (a + e), sequences=[pt.exp(-np.diff(t) / c), ea[:-1]],
                              outputs_info=[pt.zeros(())], return_updates=False)
            trig = pt.concatenate([pt.zeros(1), A])[tgt]
        lam = mu + K0 * trig
        # ---- compensator of the DETECTABLE events: int q(t) lambda(t) dt, q = P(M >= Mc(t) | M >= Mc)
        Z = -pt.expm1(-beta * (MMAX - Mc))
        q = (pt.exp(-beta * (lev - Mc)) - pt.exp(-beta * (MMAX - Mc))) / Z
        comp = mu * (q * (hi - lo)).sum() + K0 * (ea * (Gf(np.clip(T - t, 0, None))
                                                        - Gf(np.clip(S - t, 0, None)))).sum()
        inc = np.flatnonzero((lev > Mc + 1e-9) & (hi > lo))      # incomplete steps: subtract (1 - q)
        if inc.size and kernel == "omori" and method == "soe":
            # parents born before the step: integrate their exponential mixture from the scan state
            first = np.searchsorted(t, lo[inc], side="left")
            last = np.maximum(first - 1, 0)
            dec_lo = np.exp(-np.outer(lo[inc] - t[last], sk)) * (first > 0)[:, None]
            integ = -np.expm1(-np.outer(hi[inc] - lo[inc], sk)) / sk
            old = ((A[last] + ea[last][:, None]) * dec_lo * integ) @ wk
            # parents born inside the step: exact
            kin = np.clip(np.searchsorted(br, t, side="right") - 1, 0, len(lev) - 1)
            jin = np.flatnonzero(np.isin(kin, inc))
            pos = np.searchsorted(inc, kin[jin])
            new = pt.zeros(inc.size)[pos].inc(ea[jin] * Gf(hi[kin[jin]] - t[jin]))
            comp = comp - K0 * ((1 - q[inc]) * (old + new)).sum()
        elif inc.size:
            ks, js = np.nonzero((lev[:, None] > Mc + 1e-9) & (t[None, :] < hi[:, None]) & (hi > lo)[:, None])
            Gkj = Gf(np.clip(hi[ks] - t[js], 0, None)) - Gf(np.clip(lo[ks] - t[js], 0, None))
            comp = comp - K0 * ((1 - q[ks]) * ea[js] * Gkj).sum()
        pm.Potential("time_loglik", pt.log(lam).sum() - comp)
        pm.Potential("magnitude_loglik", (pt.log(beta) - beta * dm[tgt] - pt.log(Z)).sum())
    return model, tgt

def params_list(idata, n, seed, kernel="omori"):
    post = idata.posterior
    names = ["mu", "c", "alpha", "b", "K0"] + (["p"] if kernel == "omori" else [])
    arr = {k: post[k].values.ravel() for k in names}
    arr["beta"] = arr["b"] * np.log(10)
    idx = np.random.default_rng(seed).choice(arr["mu"].size, n, replace=False)
    return [{k: v[i] for k, v in arr.items()} for i in idx]


def Gk(x, par, kernel):
    x = np.clip(x, 0, None)
    return G_omori(x, par["c"], par["p"]) if kernel == "omori" else G_exp(x, par["c"])


def g_kernel(x, par, kernel):
    xs = np.where(x > 0, x, 1.0)
    g = (xs + par["c"]) ** -par["p"] if kernel == "omori" else np.exp(-xs / par["c"])
    return np.where(x > 0, g, 0.0)


def intensity(tau, t, m, par, kernel="omori", Mc=MC):
    """lambda(tau) for M >= Mc events; parents are events strictly before tau."""
    ea = np.exp(par["alpha"] * (m - Mc))
    return par["mu"] + par["K0"] * (g_kernel(tau[:, None] - t[None, :], par, kernel) * ea).sum(1)


def compensator(tau, t, m, S, steps, par, kernel="omori", Mc=MC):
    """int_S^tau q(s) lambda(s) ds at sorted times tau (recorded events only)."""
    br, lev = steps
    be = par["beta"]
    q = (np.exp(-be * (lev - Mc)) - np.exp(-be * (MMAX - Mc))) / -np.expm1(-be * (MMAX - Mc))
    pts = np.union1d(tau, br[(br > S) & (br <= tau.max())])
    pts = np.concatenate([[S], pts[pts > S]])
    ea = np.exp(par["alpha"] * (m - Mc))
    L = par["mu"] * (pts - S) + par["K0"] * (Gk(pts[:, None] - t[None, :], par, kernel) * ea).sum(1)
    k = np.clip(np.searchsorted(br, 0.5 * (pts[1:] + pts[:-1]), side="right") - 1, 0, len(lev) - 1)
    cum = np.concatenate([[0.0], np.cumsum(q[k] * np.diff(L))])
    return cum[np.searchsorted(pts, tau)]


model_specs = {"naive ETAS (constant Mc)": ((np.array([S0, T_B]), np.array([MC])), "omori"),
               "ETAS with Mc(t)": (steps_B, "omori"),
               "exponential Hawkes, Mc(t)": (steps_B, "exp")}
tB, mB = evB.t.values, evB.mag.values
resid = {}
for label, (st, kern) in model_specs.items():
    tgt = np.flatnonzero(mB >= mc_at(tB, st) - 1e-9)
    taus, us = [], []
    for par in params_list(fits[label], 200, 1, kern):
        tau = compensator(tB[tgt], tB, mB, S0, st, par, kern)
        taus.append(tau)
        us.append(np.sort(-np.expm1(-np.diff(np.concatenate([[0.0], tau])))))
    resid[label] = dict(tgt=tgt, tau=np.array(taus), u=np.array(us))
    D = np.array([stats.kstest(u, "uniform").statistic for u in resid[label]["u"]])
    print(f"{label:28s} n = {len(tgt)}: KS distance median {np.median(D):.3f} "
          f"(5% critical value {1.36 / np.sqrt(len(tgt)):.3f}); share of draws rejected at 5%: "
          f"{np.mean(D > 1.36 / np.sqrt(len(tgt))):.2f}")

fig, axes = plt.subplots(2, 2, figsize=(13, 8.5))
cols = {"naive ETAS (constant Mc)": ORANGE, "ETAS with Mc(t)": BLUE, "exponential Hawkes, Mc(t)": AQUA}
for label, r in resid.items():
    n_ = len(r["tgt"])
    tau_med = np.median(r["tau"], 0)
    res_count = np.arange(1, n_ + 1) - tau_med                 # observed - expected count
    axes[0, 0].plot(tau_med / n_, res_count / np.sqrt(n_), color=cols[label], lw=1.5, label=label)
    grid_u = np.linspace(0, 1, 200)
    ec = np.array([np.searchsorted(u, grid_u) / n_ for u in r["u"]])
    lo_, hi_ = np.quantile((ec - grid_u) * np.sqrt(n_), [0.05, 0.95], axis=0)
    axes[0, 1].fill_between(grid_u, lo_, hi_, color=cols[label], alpha=0.35, lw=0, label=label)
    # where in real time? residual count restarted at each mainshock
    tt_ = tB[r["tgt"]]
    for ax, (a, b_) in zip(axes[1], [(t64, 0.0), (0.0, T_B)]):
        sel = (tt_ >= a) & (tt_ < b_)
        base = res_count[np.searchsorted(tt_, a) - 1]
        ax.plot((tt_[sel] - a) * 24, res_count[sel] - base, color=cols[label], lw=1.5)
axes[0, 0].axhspan(-2, 2, color=GREY, alpha=0.15)
axes[0, 0].axhline(0, color=INK, lw=0.8)
axes[0, 0].set(xlabel="transformed time / number of events", ylabel="(observed - expected) / sqrt(n)",
               title="cumulative residual in transformed time")
axes[0, 0].legend(fontsize=8)
axes[0, 1].axhline(0, color=INK, lw=0.8)
for y in (-1.36, 1.36):
    axes[0, 1].axhline(y, color=INK, ls=":", lw=1)
axes[0, 1].set(xlabel="u", ylabel="sqrt(n) (empirical CDF - u)",
               title="KS check of rescaled gaps (90% of draws; dotted: 5% test)")
for ax, what in zip(axes[1], ["M6.4", "M7.1"]):
    ax.axhline(0, color=INK, lw=0.8)
    ax.set(xscale="log", xlabel=f"hours after the {what} (log scale)", ylabel="observed - expected count",
           title=f"where the misfit happens: after the {what}")