app.py

"""Gradio app (Hugging Face Spaces): shear-wave velocity prediction from a
phase-velocity dispersion curve with the phase-only DispFormer trained on
OpenSWI-shallow.

Gradio SDK so the Space also works on ZeroGPU hardware (ZeroGPU is
Gradio-only). The model is 2.44 M parameters — inference runs on CPU in
milliseconds, so no GPU decorator is needed.

Run locally with:
    python app.py
"""
import os
import tempfile

try:
    import spaces  # ZeroGPU: must be imported before torch
except ImportError:  # local run / CPU Space without the spaces package
    spaces = None

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import torch
import gradio as gr

from model import DispersionTransformerAblate

APP_DIR = os.path.dirname(os.path.abspath(__file__))

# proposed model of paper3 (v4): physically-coded encoder-decoder, 2.44 M params
CKPT = os.path.join(APP_DIR, "checkpoints", "best_model.pth")
PERIOD = np.load(os.path.join(APP_DIR, "assets/period_grid.npy"))
DEPTH = np.load(os.path.join(APP_DIR, "assets/depth_grid.npy"))
C1, C3 = 1 / 3, 1 / 2  # wavelength heuristic coefficients (Xia et al., 1999)

MODEL = DispersionTransformerAblate(
    model_dim=128, num_heads=8, num_layers=3, output_dim=72,
    scale_factor=4.5, local=False, decoder="depthq",
    depth_values=DEPTH, decoder_layers=1)
MODEL.load_state_dict(torch.load(CKPT, map_location="cpu"))
MODEL.train()  # evaluation protocol of the training pipeline

_EX = np.load(os.path.join(APP_DIR, "assets/examples.npz"))
EX_CURVES, EX_PROFILES = _EX["curves"], _EX["profiles"]
EX_LABELS = [f"{n} #{i}" for i, n in enumerate(_EX["names"])]

DEFAULT_PASTE = "\n".join(f"{t:.3f}, {c:.3f}" for t, c in
                          zip(PERIOD[::12], EX_CURVES[0][::12]) if c > 0)

K_PRESETS = ["Native (k = 1): 0.2–10 s, 0–2.76 km",
             "Engineering ≈ 30 m (k = 100): 2 ms–0.1 s, 0–27.6 m",
             "Custom k"]


def resolve_k(preset, k_custom):
    if preset.startswith("Custom"):
        return float(np.clip(k_custom or 100.0, 1.0, 10000.0))
    return 100.0 if preset.startswith("Engineering") else 1.0


def snap_to_grid(periods, velocities):
    """Place picks on the fixed 100-period grid (nearest period; picks that
    share a grid node are averaged). Returns (curve on grid, n_used, n_out)."""
    grid = np.full(len(PERIOD), -1.0, dtype=np.float32)
    counts = np.zeros(len(PERIOD))
    sums = np.zeros(len(PERIOD))
    n_out = 0
    for T, c in zip(periods, velocities):
        if not (PERIOD.min() <= T <= PERIOD.max()) or c <= 0:
            n_out += 1
            continue
        j = int(np.abs(PERIOD - T).argmin())
        sums[j] += c
        counts[j] += 1
    used = counts > 0
    grid[used] = (sums[used] / counts[used]).astype(np.float32)
    return grid, int(used.sum()), n_out


def usable_depth_range(curve):
    """Constrained depth interval from the wavelength heuristic."""
    valid = curve > 0
    if not valid.any():
        return 0, len(DEPTH)
    p, v = PERIOD[valid], curve[valid]
    dmin = C1 * p.min() * v[p.argmin()]
    dmax = C3 * p.max() * v[p.argmax()]
    lo = max(0, int(np.abs(DEPTH - dmin).argmin()) - 1)
    hi = min(len(DEPTH), int(np.abs(DEPTH - dmax).argmin()) + 1)
    return lo, hi


def _gpu(fn):
    """ZeroGPU hardware refuses to start without a @spaces.GPU function.
    The model itself runs on CPU in milliseconds, so the short duration just
    satisfies the check while keeping queue priority high."""
    return spaces.GPU(duration=10)(fn) if spaces is not None else fn


@_gpu
def predict(curve):
    x = np.full((1, 3, len(PERIOD)), -1.0, dtype=np.float32)
    x[0, 0] = PERIOD
    x[0, 1] = curve
    x = torch.from_numpy(x)
    mask = (x[:, 1] == -1) & (x[:, 2] == -1)
    with torch.no_grad():
        out = MODEL(x, mask)
    return out[0, :len(DEPTH)].numpy()


def render(curve, k, true_vs=None, note=None):
    """Run the inversion and build (input fig, profile fig, summary, CSV)."""
    if curve is None or (curve > 0).sum() == 0:
        raise gr.Error("No valid picks inside the accepted period band — "
                       "check the values, units, and scale factor k.")

    lines = [] if note is None else [note]
    if (curve > 0).sum() < 5:
        lines.append("⚠️ Very few valid picks — the prediction will be "
                     "poorly constrained.")
    vmin_meas = float(curve[curve > 0].min())
    if vmin_meas < 0.3:
        lines.append(
            f"⚠️ Lowest measured phase velocity is {vmin_meas*1000:.0f} m/s — "
            "below the training support (≈ 0.3–4.5 km/s, unchanged by the "
            "scale factor). Typical of soft-soil sites; predictions there are "
            "extrapolation and the recommended path is fine-tuning on "
            "engineering-scale synthetics (paper3 §6.2).")

    vs = predict(curve)
    lo, hi = usable_depth_range(curve)
    valid = curve > 0

    # display-unit helpers: show everything in the user's field units
    depth_disp = DEPTH / k
    period_disp = PERIOD / k
    depth_in_m = depth_disp.max() < 0.2       # show meters for shallow scales
    dunit = "m" if depth_in_m else "km"
    dsc = 1000.0 if depth_in_m else 1.0
    D = depth_disp * dsc

    fig1, ax = plt.subplots(figsize=(5.5, 4), constrained_layout=True)
    ax.plot(period_disp[valid], curve[valid], "o-", ms=3.5, lw=1.2,
            color="#1f77b4")
    ax.set_xscale("log")
    ax.set_xlabel("period (s)" + (f"  [measured; k = {k:g}]" if k != 1 else ""))
    ax.set_ylabel("phase velocity (km/s)")
    ax.set_title(f"Input curve ({int(valid.sum())} of {len(PERIOD)} grid periods)")
    ax.grid(alpha=0.3)

    fig2, ax = plt.subplots(figsize=(5.5, 4), constrained_layout=True)
    if true_vs is not None:
        ax.step(true_vs, D, where="mid", color="k", lw=1.6, label="true Vs")
    ax.step(vs, D, where="mid", color="#d62728", lw=1.6, label="predicted Vs")
    if lo > 0:
        ax.axhspan(0, D[lo], color="gray", alpha=0.15)
    if hi < len(DEPTH):
        ax.axhspan(D[hi - 1], D[-1], color="gray", alpha=0.15)
    ax.invert_yaxis()
    ax.set_xlabel("Vs (km/s)")
    ax.set_ylabel(f"depth ({dunit})")
    ax.set_title("Predicted 1-D Vs profile")
    ax.legend(fontsize=8)
    ax.grid(alpha=0.3)

    if true_vs is not None:
        err = vs - true_vs
        lines.append(
            f"**RMSE vs truth:** {np.sqrt((err ** 2).mean()):.3f} km/s · "
            f"**MAE:** {np.abs(err).mean():.3f} km/s · "
            f"**MAPE:** {(np.abs(err) / true_vs).mean() * 100:.1f} %")

    lines.append(
        f"Gray bands mark depths outside the range the input band physically "
        f"constrains (wavelength heuristic: ≈ ⅓·λ_min to ½·λ_max → "
        f"{D[lo]:.2f}–{D[min(hi, len(DEPTH)) - 1]:.2f} {dunit} here); treat "
        f"the profile there as extrapolation."
        + (f" Scale factor k = {k:g}: periods ×{k:g} into the model, depths "
           f"÷{k:g} on output; velocities unchanged." if k != 1 else ""))

    out_df = pd.DataFrame({f"depth_{dunit}": D, "vs_km_s": vs,
                           "constrained": [(lo <= i < hi)
                                           for i in range(len(DEPTH))]})
    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False,
                                      prefix="predicted_vs_profile_")
    out_df.to_csv(tmp.name, index=False)
    tmp.close()

    return fig1, fig2, "\n\n".join(lines), tmp.name


def invert_example(label, lo_p, hi_p):
    idx = EX_LABELS.index(label)
    curve = EX_CURVES[idx].copy()
    curve[(PERIOD < lo_p) | (PERIOD > hi_p)] = -1.0
    return render(curve, 1.0, true_vs=EX_PROFILES[idx])


def invert_csv(file, freq_input, vel_unit, preset, k_custom):
    if file is None:
        raise gr.Error("Upload a CSV file first.")
    k = resolve_k(preset, k_custom)
    path = file if isinstance(file, str) else file.name
    df = pd.read_csv(path)
    if df.shape[1] < 2:
        raise gr.Error("The CSV needs at least two columns: period (s) or "
                       "frequency (Hz), then phase velocity.")
    p = pd.to_numeric(df.iloc[:, 0], errors="coerce").to_numpy(dtype=float)
    v = pd.to_numeric(df.iloc[:, 1], errors="coerce").to_numpy(dtype=float)
    ok = np.isfinite(p) & np.isfinite(v)
    p, v = p[ok], v[ok]
    if freq_input:
        p = np.where(p > 0, 1.0 / p, np.nan)
        v = v[np.isfinite(p)]
        p = p[np.isfinite(p)]
    if vel_unit == "m/s":
        v = v / 1000.0
    curve, n_used, n_out = snap_to_grid(p * k, v)
    note = (f"{n_used} grid periods filled"
            + (f" · {n_out} picks outside the accepted band dropped"
               if n_out else ""))
    return render(curve, k, note=note)


def invert_paste(text, preset, k_custom):
    k = resolve_k(preset, k_custom)
    try:
        rows = [list(map(float, ln.replace(",", " ").split()))
                for ln in text.strip().splitlines() if ln.strip()]
        arr = np.array([r[:2] for r in rows if len(r) >= 2])
        curve, n_used, n_out = snap_to_grid(arr[:, 0] * k, arr[:, 1])
    except gr.Error:
        raise
    except Exception as e:
        raise gr.Error(f"Could not parse input: {e}")
    return render(curve, k, note=f"{n_used} grid periods filled")


INTRO = (
    "# Shear-wave velocity from a phase-velocity dispersion curve\n\n"
    "Inverts a fundamental-mode Rayleigh **phase-velocity** curve (the "
    "SASW/MASW observable) for a 70-layer 1-D Vs profile in one forward pass "
    "— no initial model. Proposed physically-coded transformer "
    "encoder–decoder (2.44 M params: period tokens in, depth-query tokens "
    "out) trained on **OpenSWI-shallow** (22 M curve/profile pairs).")

SCALE_NOTE = (
    "**Site scale:** the physics is scale-invariant — measured periods are "
    "multiplied by k to enter the model's 0.2–10 s band, and output depths "
    "are divided by k. Velocities are never rescaled (support "
    "≈ 0.3–4.5 km/s at any scale). Examples are native-scale; k applies to "
    "uploaded/pasted curves only. Picks are snapped to the nearest of the "
    f"{len(PERIOD)} grid periods; gaps and band-limited curves are handled "
    "natively.")

DETAILS = (
    "- **Model:** physically-coded encoder–decoder (paper3): each dispersion "
    "pick is a token carrying its physical period; each output depth is a query "
    "token carrying its physical depth, reading the period tokens by masked "
    "cross-attention — band-limited and gappy curves are handled natively "
    "(no interpolation), and its learned attention reproduces the classical "
    "λ/3 sensitivity rule (paper3, Fig. 7).\n"
    "- **Checkpoint:** v4 finalist (valid masked MSE 0.0216 (km/s)²); test "
    "accuracy 0.151 km/s full-profile RMSE / 6.0 % MAPE / R² 0.949 on 50 k "
    "held-out samples; Long Beach field data 38 m/s MAE vs the tomographic "
    "reference (5,297 real curves).\n"
    "- **Scope:** trained on 0.2–10 s periods, 0–2.76 km depth, Vs ≈ 0.3–4.5 km/s "
    "(OpenSWI-shallow). Curves outside this envelope — e.g. soft-soil sites with "
    "Vs < 0.3 km/s — are out of distribution.\n"
    "- **Site scale factor k:** the elastodynamic problem is scale-invariant, so "
    "a high-frequency engineering curve is inverted by stretching its periods "
    "×k into the training band and shrinking the output depths ÷k (e.g. k = 100 "
    "→ 70 layers over 0.4–27.6 m at 0.4 m spacing). Velocities are never "
    "rescaled — the ≈ 0.3–4.5 km/s support applies at every scale; see paper3 §6.2.")


with gr.Blocks(title="Vs from dispersion curve") as demo:
    gr.Markdown(INTRO)
    with gr.Row():
        with gr.Column(scale=1):
            preset = gr.Dropdown(K_PRESETS, value=K_PRESETS[0],
                                 label="Scale factor k")
            k_custom = gr.Number(value=100.0, minimum=1.0, maximum=10000.0,
                                 label="Custom k (used when preset is "
                                       "'Custom k')")
            gr.Markdown(SCALE_NOTE)
            with gr.Tab("Example from test data"):
                ex = gr.Dropdown(EX_LABELS, value=EX_LABELS[0],
                                 label="Example (native scale, k = 1)")
                lo_p = gr.Slider(float(PERIOD.min()), float(PERIOD.max()),
                                 value=float(PERIOD.min()),
                                 label="Min period (s) — simulates a "
                                       "band-limited survey")
                hi_p = gr.Slider(float(PERIOD.min()), float(PERIOD.max()),
                                 value=float(PERIOD.max()),
                                 label="Max period (s)")
                btn_ex = gr.Button("Invert example", variant="primary")
            with gr.Tab("Upload CSV"):
                gr.Markdown("First column: period (s) **or** frequency (Hz); "
                            "second column: phase velocity. A header row is "
                            "expected.")
                up = gr.File(file_types=[".csv", ".txt"], label="CSV file")
                freq_in = gr.Checkbox(False,
                                      label="First column is frequency (Hz)")
                unit = gr.Radio(["km/s", "m/s"], value="km/s",
                                label="Velocity unit")
                btn_csv = gr.Button("Invert CSV", variant="primary")
            with gr.Tab("Paste values"):
                txt = gr.Textbox(DEFAULT_PASTE, lines=12,
                                 label="One 'period_s, velocity_km_s' pair "
                                       "per line")
                btn_txt = gr.Button("Invert pasted curve", variant="primary")
        with gr.Column(scale=2):
            with gr.Row():
                plot_in = gr.Plot(label="Input curve")
                plot_out = gr.Plot(label="Predicted 1-D Vs profile")
            summary = gr.Markdown()
            dl = gr.File(label="Predicted profile (CSV)")
    with gr.Accordion("Model & protocol details", open=False):
        gr.Markdown(DETAILS)

    outputs = [plot_in, plot_out, summary, dl]
    btn_ex.click(invert_example, [ex, lo_p, hi_p], outputs)
    btn_csv.click(invert_csv, [up, freq_in, unit, preset, k_custom], outputs)
    btn_txt.click(invert_paste, [txt, preset, k_custom], outputs)


if __name__ == "__main__":
    demo.launch()