Source code: Mogi model inversion in Python: magma depth and volume from InSAR at Svartsengi, Iceland

"""Volcano deformation with InSAR, part 2: Mogi and Okada source models fitted to Svartsengi InSAR.

Loads the ascending and descending Sentinel-1 LOS maps of Aug-Oct 2025 made in part 1 (or rebuilds
them from COMET LiCSAR if part 1 has not been run), checks a NumPy implementation of Okada (1985)
against the published test values, downsamples the maps with a quadtree, fits a Mogi point source
and an Okada horizontal sill by weighted non-linear least squares (scipy.optimize.least_squares),
estimates the noise covariance from the residuals, and propagates it with a Monte Carlo of
correlated noise to get uncertainties for depth, location and volume change.
Run with:  python mogi_okada_inversion.py
"""
import platform
import re
import warnings
from importlib.metadata import version
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.patheffects as pe
import matplotlib.pyplot as plt
import numpy as np
import rasterio
import scipy
from cmcrameri import cm
from matplotlib.collections import PatchCollection
from matplotlib.patches import Ellipse, Polygon, Rectangle
from rasterio.windows import from_bounds
from scipy.optimize import curve_fit, least_squares

HERE = Path(__file__).parent
OUT, DATA = HERE / "output", HERE / "data"
for d in (OUT, DATA):
    d.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9,
                     "figure.facecolor": "white", "savefig.facecolor": "white"})
warnings.filterwarnings("ignore", category=RuntimeWarning)  # all-NaN blocks over the sea

print("=== versions ===")
print(f"python {platform.python_version()} | numpy {np.__version__} | scipy {scipy.__version__} | "
      f"rasterio {rasterio.__version__} | matplotlib {matplotlib.__version__} | cmcrameri {version('cmcrameri')}")


def section(title: str) -> None:
    print(f"\n=== {title} ===")


NU = 0.25                                  # Poisson's ratio (lambda = mu)
SENG = (-22.4338, 63.8822)                 # GNSS station, used as the origin of the local x/y grid
PARKS_MOGI = (-22.454, 63.869, 4.0)        # Parks et al. (2025): Mogi source lon, lat, depth (km)
DAYS = 48                                  # length of the part 1 time window (15 Aug - 2 Oct 2025)
N_MC = 200                                 # Monte Carlo realisations of correlated noise
RNG = np.random.default_rng(42)
PART1 = HERE.parent / "insar-volcano-deformation-python" / "data" / "svartsengi_los_2025.npz"

# --- Step 0: inputs from part 1 (or rebuild them) ----------------------------------------------
LICSAR = "https://gws-access.jasmin.ac.uk/public/nceo_geohazards/LiCSAR_products.public"
TRACKS = {"ascending": ("016A_02504_162118", ["20250815_20250908", "20250908_20251002"]),
          "descending": ("155D_02484_191814", ["20250813_20250906", "20250906_20250930"])}
BBOX = (-22.80, 63.72, -22.05, 64.05)
RAD2MM = -0.0554658 / (4 * np.pi) * 1000


def rebuild_part1(path: Path) -> None:
    """Same steps as part 1: add two 24-day pairs per track, mask coherence < 0.3, remove a far-field plane."""
    rasterio.Env(GDAL_DISABLE_READDIR_ON_OPEN="EMPTY_DIR", GDAL_HTTP_MAX_RETRY=4).__enter__()

    def rd(url):
        with rasterio.open(url) as src:
            win = from_bounds(*BBOX, src.transform).round_offsets().round_lengths()
            return src.read(1, window=win).astype("float32"), src.window_transform(win)

    out = {}
    for track, (frame, pairs) in TRACKS.items():
        b = f"{LICSAR}/{int(frame[:3])}/{frame}"
        stack = [(rd(f"{b}/interferograms/{p}/{p}.geo.unw.tif"), rd(f"{b}/interferograms/{p}/{p}.geo.cc.tif")[0])
                 for p in pairs]
        t = stack[0][0][1]
        H, W = stack[0][1].shape
        lon = t.c + t.a * (np.arange(W) + 0.5)
        lat = t.f + t.e * (np.arange(H) + 0.5)
        LON, LAT = np.meshgrid(lon, lat)
        far = np.hypot((LON - SENG[0]) * 49.0, (LAT - SENG[1]) * 110.6) > 15
        total, good = np.zeros((H, W)), np.ones((H, W), bool)
        for (u, _), c in stack:
            ok = (u != 0) & (c / 255 >= 0.3)
            good &= ok
            use = ok & far
            A = np.column_stack([np.ones(use.sum()), LON[use] - LON.mean(), LAT[use] - LAT.mean()])
            k = np.linalg.lstsq(A, u[use], rcond=None)[0]
            total += np.where(ok, u - k[0] - k[1] * (LON - LON.mean()) - k[2] * (LAT - LAT.mean()), 0) * RAD2MM
        out[track[:3] if track == "ascending" else "desc"] = np.where(good, total, np.nan)
        out["vec_" + ("asc" if track == "ascending" else "desc")] = np.stack(
            [rd(f"{b}/metadata/{frame}.geo.{k}.tif")[0] for k in "ENU"])
    np.savez_compressed(path, lon=lon, lat=lat, **out)


src = PART1 if PART1.exists() else DATA / "svartsengi_los_2025.npz"
if not src.exists():
    print("part 1 output not found: rebuilding it from LiCSAR (about 2 minutes)")
    rebuild_part1(src)
z = np.load(src)
lon, lat = z["lon"], z["lat"]
X = (lon - SENG[0]) * 111.32 * np.cos(np.radians(SENG[1]))    # km east of SENG
Y = (lat - SENG[1]) * 110.6                                     # km north of SENG
XX, YY = np.meshgrid(X, Y)
DATASETS = {"ascending": (z["asc"], z["vec_asc"]), "descending": (z["desc"], z["vec_desc"])}
section("inputs")
print(f"read {src.relative_to(HERE.parent)}: grid {len(lat)} x {len(lon)}, "
      f"x {X.min():.1f}..{X.max():.1f} km, y {Y.min():.1f}..{Y.max():.1f} km from SENG")
for k, (d, _) in DATASETS.items():
    print(f"{k:10s} LOS: {np.isfinite(d).sum():,} pixels, {np.nanmin(d):+.1f} to {np.nanmax(d):+.1f} mm")


# --- Step 1: forward models --------------------------------------------------------------------
def mogi(x, y, x0, y0, depth, dV, nu=NU):
    """Mogi (1958) point pressure source in an elastic half-space.
    x, y, x0, y0, depth in metres; dV = volume change of the source in m3. Returns (ue, un, uz) in metres."""
    dx, dy = x - x0, y - y0
    R3 = (dx ** 2 + dy ** 2 + depth ** 2) ** 1.5
    C = (1 - nu) * dV / np.pi
    return C * dx / R3, C * dy / R3, C * depth / R3


def okada85(x, y, depth, dip, L, W, U1=0.0, U2=0.0, U3=0.0, nu=NU):
    """Okada (1985) surface displacement of a rectangular dislocation in an elastic half-space.
    Okada's own frame: x along strike, fault from x = 0..L and from the bottom edge (depth `depth`,
    at y = 0) W up-dip. U1 strike-slip, U2 dip-slip, U3 tensile opening. Returns (ux, uy, uz)."""
    d = np.radians(dip)
    cs, sn = np.cos(d), np.sin(d)
    if abs(cs) < 1e-12:
        cs = 0.0
    a = 1 - 2 * nu                                    # mu / (lambda + mu)
    p = y * cs + depth * sn
    q = y * sn - depth * cs

    def f(xi, eta):
        yt, dt = eta * cs + q * sn, eta * sn - q * cs
        R = np.sqrt(xi ** 2 + eta ** 2 + q ** 2)
        Xb = np.sqrt(xi ** 2 + q ** 2)
        with np.errstate(divide="ignore", invalid="ignore"):
            th = np.where(q == 0, 0.0, np.arctan(xi * eta / (q * R)))
            if cs != 0:
                I5 = np.where(xi == 0, 0.0, a * 2 / cs * np.arctan(
                    (eta * (Xb + q * cs) + Xb * (R + Xb) * sn) / (xi * (R + Xb) * cs)))
                I4 = a / cs * (np.log(R + dt) - sn * np.log(R + eta))
                I3 = a * (yt / (cs * (R + dt)) - np.log(R + eta)) + sn / cs * I4
                I1 = a * (-xi / (cs * (R + dt))) - sn / cs * I5
            else:
                I1 = -a / 2 * xi * q / (R + dt) ** 2
                I3 = a / 2 * (eta / (R + dt) + yt * q / (R + dt) ** 2 - np.log(R + eta))
                I4 = -a * q / (R + dt)
                I5 = -a * xi * sn / (R + dt)
            I2 = a * -np.log(R + eta) - I3
            ux = -U1 / (2 * np.pi) * (xi * q / (R * (R + eta)) + th + I1 * sn) \
                - U2 / (2 * np.pi) * (q / R - I3 * sn * cs) \
                + U3 / (2 * np.pi) * (q ** 2 / (R * (R + eta)) - I3 * sn ** 2)
            uy = -U1 / (2 * np.pi) * (yt * q / (R * (R + eta)) + q * cs / (R + eta) + I2 * sn) \
                - U2 / (2 * np.pi) * (yt * q / (R * (R + xi)) + cs * th - I1 * sn * cs) \
                + U3 / (2 * np.pi) * (-dt * q / (R * (R + xi)) - sn * (xi * q / (R * (R + eta)) - th) - I1 * sn ** 2)
            uz = -U1 / (2 * np.pi) * (dt * q / (R * (R + eta)) + q * sn / (R + eta) + I4 * sn) \
                - U2 / (2 * np.pi) * (dt * q / (R * (R + xi)) + sn * th - I5 * sn * cs) \
                + U3 / (2 * np.pi) * (yt * q / (R * (R + xi)) + cs * (xi * q / (R * (R + eta)) - th) - I5 * sn ** 2)
        return np.array([ux, uy, uz])

    return f(x, p) - f(x, p - W) - f(x - L, p) + f(x - L, p - W)   # Chinnery's notation


def okada_rect(x, y, xc, yc, depth, L, W, strike, dip, opening):
    """Tensile Okada source in map coordinates (metres). (xc, yc, depth) is the centre of the rectangle,
    strike is in degrees clockwise from north, dip 0 = horizontal sill, 90 = vertical dyke. Returns (ue, un, uz)."""
    s = np.radians(strike)
    se, sn_ = np.sin(s), np.cos(s)                      # unit vector along strike (east, north)
    dx, dy = x - xc, y - yc
    xs = dx * se + dy * sn_ + L / 2                     # Okada x: along strike, from the fault end
    ys = -dx * sn_ + dy * se + W / 2 * np.cos(np.radians(dip))  # Okada y: left of strike
    ux, uy, uz = okada85(xs, ys, depth + W / 2 * np.sin(np.radians(dip)), dip, L, W, U3=opening)
    return ux * se - uy * sn_, ux * sn_ + uy * se, uz


section("Okada (1985) check against Table 2 (x=2, y=3, d=4, dip=70, L=3, W=2, lambda=mu)")
okada_table2 = {"strike-slip": (-8.689e-3, -4.298e-3, -2.747e-3),
                "dip-slip": (-4.682e-3, -3.527e-2, -3.564e-2),
                "tensile": (-2.660e-4, 1.056e-2, 3.214e-3)}
for (name, ref), kw in zip(okada_table2.items(), [{"U1": 1}, {"U2": 1}, {"U3": 1}]):
    got = okada85(np.array(2.0), np.array(3.0), 4.0, 70.0, 3.0, 2.0, **kw)
    print(f"{name:12s} ux {got[0]:+.3e} uy {got[1]:+.3e} uz {got[2]:+.3e} | published "
          f"{ref[0]:+.3e} {ref[1]:+.3e} {ref[2]:+.3e} | max diff {np.max(np.abs(got - ref)):.1e}")
    assert np.allclose(got, ref, rtol=2e-3, atol=0), "Okada implementation does not match Okada (1985) Table 2"
# A small sill should behave like a point tensile crack: uz = 3 dV d^3 / (2 pi R^5) (Okada 1985, point source)
r = np.array([0.0, 2e3, 4e3, 8e3])
dV_test = 100.0 * 100.0 * 1.0
crack = 3 * dV_test * 4e3 ** 3 / (2 * np.pi * (r ** 2 + 4e3 ** 2) ** 2.5)
sill_small = okada_rect(r, 0 * r, 0, 0, 4e3, 100.0, 100.0, 0, 0, 1.0)[2]
mogi_same = mogi(r, 0 * r, 0, 0, 4e3, dV_test)[2]
print("100 m x 100 m sill, 1 m opening, 4 km deep; uz in micrometres at r = 0, 2, 4, 8 km:")
print(f"  Okada {np.round(sill_small * 1e6, 3)} | point crack formula {np.round(crack * 1e6, 3)} | "
      f"Mogi, same volume {np.round(mogi_same * 1e6, 3)}")
print(f"  above the source a sill lifts the surface {sill_small[0] / mogi_same[0]:.2f}x more than a Mogi "
      f"source of the same volume change")


# --- Step 2: quadtree downsampling -------------------------------------------------------------
def quadtree(data, max_std=2.5, min_px=4, start=32, min_valid=0.5):
    """Split blocks while the LOS inside them varies more than max_std (mm). Returns leaf rows/cols/values."""
    leaves = []

    def split(r0, c0, h, w):
        blk = data[r0:r0 + h, c0:c0 + w]
        valid = np.isfinite(blk)
        if valid.mean() < min_valid:
            if h > min_px:                              # coastline: try smaller blocks
                for dr, dc in [(0, 0), (0, w // 2), (h // 2, 0), (h // 2, w // 2)]:
                    split(r0 + dr, c0 + dc, h // 2 if dr == 0 else h - h // 2, w // 2 if dc == 0 else w - w // 2)
            return
        if np.nanstd(blk) > max_std and h > min_px:
            for dr, dc in [(0, 0), (0, w // 2), (h // 2, 0), (h // 2, w // 2)]:
                split(r0 + dr, c0 + dc, h // 2 if dr == 0 else h - h // 2, w // 2 if dc == 0 else w - w // 2)
            return
        rr, cc = np.nonzero(valid)
        leaves.append((r0 + rr.mean(), c0 + cc.mean(), np.nanmean(blk), h, w))

    H, Wd = data.shape
    bh, bw = start, start * 2                           # pixels are ~49 m x 111 m: 2:1 keeps blocks square
    for r0 in range(0, H, bh):
        for c0 in range(0, Wd, bw):
            split(r0, c0, min(bh, H - r0), min(bw, Wd - c0))
    return np.array(leaves)


section("quadtree downsampling")
obs = {}
for k, (d, vec) in DATASETS.items():
    q = quadtree(d)
    r, c = q[:, 0].round().astype(int), q[:, 1].round().astype(int)
    obs[k] = dict(x=np.interp(q[:, 1], np.arange(len(X)), X) * 1e3, y=np.interp(q[:, 0], np.arange(len(Y)), Y) * 1e3,
                  d=q[:, 2], los=vec[:, r, c], h=q[:, 3], w=q[:, 4])
    print(f"{k:10s} {np.isfinite(d).sum():,} pixels -> {len(q)} quadtree points "
          f"(block sizes {q[:, 3].min():.0f}-{q[:, 3].max():.0f} rows)")


def los_of(model, o):
    e, n, u = model
    return (o["los"][0] * e + o["los"][1] * n + o["los"][2] * u) * 1000    # metres -> mm


def nuisance(p3, o):
    return p3[0] + p3[1] * o["x"] / 1e3 + p3[2] * o["y"] / 1e3              # offset + ramp per track (mm)


def predict(kind, p, o, track):
    nus = p[-6:-3] if track == "ascending" else p[-3:]
    if kind == "mogi":
        x0, y0, dep, dV = p[:4]
        m = mogi(o["x"], o["y"], x0 * 1e3, y0 * 1e3, dep * 1e3, dV * 1e6)
    else:
        x0, y0, dep, L, W, strike, op = p[:7]
        m = okada_rect(o["x"], o["y"], x0 * 1e3, y0 * 1e3, dep * 1e3, L * 1e3, W * 1e3, strike, 0.0, op)
    return los_of(m, o) + nuisance(nus, o)


def residuals(p, kind, data, whiten):
    return np.concatenate([whiten[k] @ (data[k] - predict(kind, p, obs[k], k)) for k in obs])


START = {"mogi": ([0, -1, 4, 5] + [0] * 6, [[-10, -10, 0.5, 0] + [-np.inf] * 6, [10, 10, 15, 100] + [np.inf] * 6]),
         "sill": ([0, -1, 4, 5, 3, 45, 0.5] + [0] * 6,
                  [[-10, -10, 0.5, 0.2, 0.2, -90, 0] + [-np.inf] * 6, [10, 10, 15, 20, 20, 90, 10] + [np.inf] * 6])}
NAMES = {"mogi": ["x0 (km)", "y0 (km)", "depth (km)", "dV (1e6 m3)"],
         "sill": ["x0 (km)", "y0 (km)", "depth (km)", "length (km)", "width (km)", "strike (deg)", "opening (m)"]}


def fit(kind, data, whiten, p0=None):
    x0, bounds = START[kind]
    return least_squares(residuals, p0 if p0 is not None else x0, bounds=bounds, args=(kind, data, whiten),
                         x_scale="jac", method="trf")


def volume(kind, p):
    return p[3] * 1e6 if kind == "mogi" else p[3] * p[4] * p[6] * 1e6


# --- Step 3: first (unweighted) fit, then estimate the noise covariance from the residuals --------
section("noise covariance from the Mogi residuals")
data = {k: o["d"] for k, o in obs.items()}
eye = {k: np.eye(len(o["d"])) for k, o in obs.items()}
first = fit("mogi", data, eye)
cov_par = {}
for k, (d, vec) in DATASETS.items():
    sub = {"x": XX * 1e3, "y": YY * 1e3, "los": vec}
    model_full = predict("mogi", first.x, sub, k)
    res = d - model_full
    rr, cc = np.nonzero(np.isfinite(res))
    pick = RNG.choice(len(rr), 3000, replace=False)
    px, py, pv = XX[rr[pick], cc[pick]], YY[rr[pick], cc[pick]], res[rr[pick], cc[pick]]
    dist = np.hypot(px[:, None] - px[None], py[:, None] - py[None])
    prod = (pv[:, None] - pv.mean()) * (pv[None] - pv.mean())
    iu = np.triu_indices(len(pv), 1)
    du, pu = dist[iu], prod[iu]
    bins = np.arange(0, 16, 0.5)
    idx = np.digitize(du, bins)
    lag = np.array([du[idx == i].mean() for i in range(1, len(bins))])
    cv = np.array([pu[idx == i].mean() for i in range(1, len(bins))])
    (s2, a), _ = curve_fit(lambda h, s2, a: s2 * np.exp(-h / a), lag, cv, p0=(cv[0], 3.0),
                           bounds=([0, 0.1], [1e4, 50]))
    cov_par[k] = (s2, a, lag, cv, pv.var())
    print(f"{k:10s} residual std {np.nanstd(res):.1f} mm; covariance C(r) = {s2:.1f} mm2 * exp(-r / {a:.1f} km) "
          f"(sigma {np.sqrt(s2):.1f} mm, e-folding distance {a:.1f} km)")

whiten, chol = {}, {}
for k, o in obs.items():
    s2, a = cov_par[k][:2]
    dist = np.hypot(o["x"][:, None] - o["x"][None], o["y"][:, None] - o["y"][None]) / 1e3
    C = s2 * np.exp(-dist / a) + np.eye(len(dist)) * 0.1 * s2   # 10 % nugget keeps C well conditioned
    Lc = np.linalg.cholesky(C)
    chol[k], whiten[k] = Lc, np.linalg.inv(Lc)          # r_w = L^-1 r  makes  r_w^T r_w = r^T C^-1 r

# --- Step 4: weighted inversions -------------------------------------------------------------
section("weighted least-squares inversions")
best = {}
n_obs = sum(len(o["d"]) for o in obs.values())
for kind in ["mogi", "sill"]:
    sol = fit(kind, data, whiten)
    best[kind] = sol.x
    chi2 = 2 * sol.cost
    rms = np.sqrt(np.mean(np.concatenate([(data[k] - predict(kind, sol.x, obs[k], k)) ** 2 for k in obs])))
    lon0 = SENG[0] + sol.x[0] / (111.32 * np.cos(np.radians(SENG[1])))
    lat0 = SENG[1] + sol.x[1] / 110.6
    vals = ", ".join(f"{n} {v:.2f}" for n, v in zip(NAMES[kind], sol.x))
    print(f"{kind:5s}: {vals}")
    print(f"       centre {lat0:.4f} N, {abs(lon0):.4f} W; volume change {volume(kind, sol.x) / 1e6:.2f} million m3 "
          f"in {DAYS} days = {volume(kind, sol.x) / (DAYS * 86400):.2f} m3/s; RMS misfit {rms:.1f} mm at the quadtree points; "
          f"chi2 {chi2:.0f} for {n_obs} points and {len(sol.x)} parameters; {sol.nfev} evaluations")

# --- Step 5: Monte Carlo of correlated noise --------------------------------------------------
section(f"Monte Carlo uncertainty ({N_MC} realisations of correlated noise)")
mc = {}
for kind in ["mogi", "sill"]:
    n_mc = N_MC if kind == "mogi" else N_MC // 2
    draws = []
    for _ in range(n_mc):
        noisy = {k: data[k] + chol[k] @ RNG.standard_normal(len(data[k])) for k in obs}
        draws.append(fit(kind, noisy, whiten, p0=best[kind]).x)
    draws = np.array(draws)
    mc[kind] = draws
    vol = np.array([volume(kind, p) for p in draws]) / 1e6
    for i, n in enumerate(NAMES[kind]):
        lo, hi = np.percentile(draws[:, i], [2.5, 97.5])
        print(f"{kind:5s} {n:13s} {best[kind][i]:8.2f}  95% interval {lo:8.2f} .. {hi:8.2f}  (std {draws[:, i].std():.2f})")
    lo, hi = np.percentile(vol, [2.5, 97.5])
    print(f"{kind:5s} {'volume (1e6 m3)':13s} {volume(kind, best[kind]) / 1e6:8.2f}  95% interval {lo:8.2f} .. {hi:8.2f}"
          f"  -> {lo * 1e6 / (DAYS * 86400):.2f}-{hi * 1e6 / (DAYS * 86400):.2f} m3/s")

section("comparison with Parks et al. (2025)")
px = (PARKS_MOGI[0] - SENG[0]) * 111.32 * np.cos(np.radians(SENG[1]))
py = (PARKS_MOGI[1] - SENG[1]) * 110.6
for kind in ["mogi", "sill"]:
    off = np.hypot(best[kind][0] - px, best[kind][1] - py)
    print(f"{kind:5s}: our centre is {off:.2f} km from the published Mogi location "
          f"({PARKS_MOGI[1]} N, {abs(PARKS_MOGI[0])} W); depth {best[kind][2]:.2f} km vs 4.0 km (Mogi) "
          f"and 4.35-4.7 km (sill)")

# --- Step 6: figures -------------------------------------------------------------------------
asp = 1.0
ext = (X.min(), X.max(), Y.min(), Y.max())


def sill_outline(p):
    x0, y0, _, L, W, strike = p[:6]
    s = np.radians(strike)
    a, b = np.array([np.sin(s), np.cos(s)]), np.array([-np.cos(s), np.sin(s)])
    return np.array([[x0, y0] + sa * L / 2 * a + sb * W / 2 * b for sa, sb in [(-1, -1), (1, -1), (1, 1), (-1, 1)]])


# Figure 1: quadtree points
fig, axes = plt.subplots(1, 2, figsize=(11, 4.4), sharey=True)
for ax, (k, o) in zip(axes, obs.items()):
    ax.set_facecolor("#dfe7ee")
    hk, wk = o["h"] * abs(Y[1] - Y[0]), o["w"] * (X[1] - X[0])          # block size in km
    boxes = [Rectangle((xi / 1e3 - wi / 2, yi / 1e3 - hi / 2), wi, hi) for xi, yi, wi, hi in zip(o["x"], o["y"], wk, hk)]
    sc = PatchCollection(boxes, cmap=cm.vik, edgecolor="white", linewidth=0.2)
    sc.set_array(o["d"])
    sc.set_clim(-60, 60)
    ax.add_collection(sc)
    ax.set_title(f"{k}: {len(o['d'])} quadtree points")
    ax.set_xlabel("x, km east of SENG")
    ax.set_aspect(asp)
    ax.set_xlim(ext[:2])
    ax.set_ylim(ext[2:])
axes[0].set_ylabel("y, km north of SENG")
fig.colorbar(sc, ax=axes, shrink=0.85, label="LOS displacement (mm, + = towards satellite)")
fig.savefig(OUT / "01-quadtree-downsampling.png")
plt.close(fig)

# Figures 2 and 3: data, model and residual on the full grid for both models
for num, kind in [(2, "mogi"), (3, "sill")]:
    fig, axes = plt.subplots(2, 3, figsize=(12, 6.9), sharex=True, sharey=True, layout="constrained")
    for row, (k, (d, vec)) in enumerate(DATASETS.items()):
        sub = {"x": XX * 1e3, "y": YY * 1e3, "los": vec}
        m = np.where(np.isfinite(d), predict(kind, best[kind], sub, k), np.nan)
        res = d - m
        for col, (arr, title, lim) in enumerate([(d, "data", 60), (m, "model", 60), (res, "residual", 20)]):
            ax = axes[row, col]
            ax.set_facecolor("#dfe7ee")
            im = ax.imshow(arr, cmap=cm.vik, vmin=-lim, vmax=lim, extent=ext, origin="upper",
                           interpolation="nearest")
            if kind == "sill":
                ax.add_patch(Polygon(sill_outline(best[kind]), fill=False, ec="black", lw=1))
            ax.plot(best[kind][0], best[kind][1], "k+", ms=8)
            ax.set_aspect(asp)
            ax.set_title(f"{k} {title}" + (f" (RMS {np.nanstd(res):.1f} mm)" if col == 2 else ""))
            fig.colorbar(im, ax=ax, shrink=0.8, label="mm")
    for ax in axes[1]:
        ax.set_xlabel("x, km east of SENG")
    for ax in axes[:, 0]:
        ax.set_ylabel("y, km north of SENG")
    label = "Mogi point source" if kind == "mogi" else "Okada horizontal sill"
    fig.suptitle(f"{label}: data, model and residual (LOS, Aug\u2013Oct 2025)")
    fig.savefig(OUT / f"0{num}-{kind}-data-model-residual.png")
    plt.close(fig)

# Figure 4: empirical covariance
fig, ax = plt.subplots(figsize=(6.5, 3.6))
for (k, (s2, a, lag, cv, var)), col in zip(cov_par.items(), ["#ef5675", "#003f5c"]):
    ax.plot(lag, cv, "o", ms=4, color=col, label=f"{k} residuals")
    hh = np.linspace(0, 15, 200)
    ax.plot(hh, s2 * np.exp(-hh / a), "-", color=col, label=f"{s2:.0f} exp(-r/{a:.1f} km)")
ax.axhline(0, color="0.5", lw=0.6)
ax.set_xlabel("distance between pixels r (km)")
ax.set_ylabel("covariance (mm\u00b2)")
ax.set_title("Spatially correlated noise in the residuals (mostly atmosphere)")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
fig.savefig(OUT / "04-noise-covariance.png")
plt.close(fig)

# Figure 5: Monte Carlo trade-offs: depth vs volume (Mogi), width vs opening (sill)
fig, axes = plt.subplots(1, 2, figsize=(11, 4))
ax = axes[0]
vol = mc["mogi"][:, 3]
ax.scatter(mc["mogi"][:, 2], vol, s=8, color="#ef5675", alpha=0.6, label="Monte Carlo fits")
ax.plot(best["mogi"][2], best["mogi"][3], "k*", ms=12, label="best fit")
ax.axvline(4.0, color="0.3", ls="--", lw=1, label="Parks et al. (2025), 4.0 km")
ax.set_xlabel("depth (km)")
ax.set_ylabel("volume change in 48 days (million m\u00b3)")
ax.set_title("Mogi: deeper source needs more volume")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
ax = axes[1]
vol_s = np.array([volume("sill", p) for p in mc["sill"]]) / 1e6
sc = ax.scatter(mc["sill"][:, 4], mc["sill"][:, 6], c=vol_s, s=10, cmap=cm.batlow)
ax.plot(best["sill"][4], best["sill"][6], "k*", ms=12, label="best fit")
ww = np.linspace(0.2, mc["sill"][:, 4].max(), 200)
ax.plot(ww, volume("sill", best["sill"]) / 1e6 / (best["sill"][3] * ww), "k--", lw=1,
        label="same volume and length")
fig.colorbar(sc, ax=ax, label="sill volume change (million m\u00b3)")
ax.set_xlabel("sill width (km)")
ax.set_ylabel("opening (m)")
ax.set_title("Okada sill: width and opening trade off")
ax.legend(fontsize=8)
ax.grid(alpha=0.3)
fig.savefig(OUT / "05-monte-carlo-trade-offs.png")
plt.close(fig)

# Figure 6: map of the sources on the vertical displacement from part 1
up = z["up"] if "up" in z.files else None
fig, ax = plt.subplots(figsize=(7, 5.2))
ax.set_facecolor("#dfe7ee")
if up is not None:
    im = ax.imshow(up, cmap=cm.vik, vmin=-120, vmax=120, extent=ext, interpolation="nearest", alpha=0.9)
    fig.colorbar(im, ax=ax, shrink=0.8, label="vertical displacement, part 1 (mm)")
else:
    im = ax.imshow(DATASETS["ascending"][0], cmap=cm.vik, vmin=-120, vmax=120, extent=ext, interpolation="nearest")
    fig.colorbar(im, ax=ax, shrink=0.8, label="ascending LOS (mm)")
m = mc["mogi"]
cov2 = np.cov(m[:, 0], m[:, 1])
w, v = np.linalg.eigh(cov2)
ang = np.degrees(np.arctan2(v[1, 1], v[0, 1]))
ax.add_patch(Ellipse((best["mogi"][0], best["mogi"][1]), 2 * 2.45 * np.sqrt(w[1]), 2 * 2.45 * np.sqrt(w[0]),
                     angle=ang, fill=False, ec="black", lw=1.2))
ax.plot(m[:, 0], m[:, 1], ".", ms=2, color="0.25", alpha=0.5)
ax.plot(best["mogi"][0], best["mogi"][1], "k*", ms=9, label=f"Mogi (this fit), {best['mogi'][2]:.1f} km deep")
ax.add_patch(Polygon(sill_outline(best["sill"]), fill=False, ec="black", lw=1.5, ls="--",
                     label=f"Okada sill (this fit), {best['sill'][2]:.1f} km deep"))
ax.plot(px, py, "wo", mec="black", ms=8, label="Mogi, Parks et al. (2025), 4.0 km deep")
ax.plot(0, 0, "^", mfc="white", mec="black", ms=7)
ax.annotate("SENG", (0, 0), xytext=(5, 5), textcoords="offset points", fontsize=8,
            path_effects=[pe.withStroke(linewidth=2.5, foreground="white")])
ax.set_xlim(-12, 10)
ax.set_ylim(-8, 9)
ax.set_aspect(asp)
ax.set_xlabel("x, km east of SENG")
ax.set_ylabel("y, km north of SENG")
ax.set_title("Best-fitting sources (dots and ellipse: Monte Carlo, 95 %)")
ax.legend(fontsize=7.5, loc="upper left")
fig.savefig(OUT / "06-source-map.png")
plt.close(fig)
print("\nfigures written to output/")