app_streamlit.py

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

Self-contained: model code, checkpoint, and assets live in this folder.
Run locally with:
    streamlit run app.py
"""
import io
import os
import sys

import streamlit as st
from streamlit import runtime

if not runtime.exists():
    # Launched as `python app.py` (e.g. a Space whose SDK is not `streamlit`):
    # replace this process with the Streamlit server on the port HF expects.
    port = os.environ.get("PORT", "7860")
    os.execvp(sys.executable, [
        sys.executable, "-m", "streamlit", "run",
        os.path.abspath(__file__),
        "--server.port", port,
        "--server.address", "0.0.0.0",
        "--server.headless", "true",
        "--server.enableCORS", "false",
        "--server.enableXsrfProtection", "false"])

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch

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)

st.set_page_config(page_title="Vs from dispersion curve", page_icon="🌍",
                   layout="wide")


@st.cache_resource
def load_model():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    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).to(device)
    model.load_state_dict(torch.load(CKPT, map_location=device))
    model.train()  # evaluation protocol of the training pipeline
    return model, device


@st.cache_data
def load_examples():
    d = np.load(os.path.join(APP_DIR, "assets/examples.npz"))
    return d["curves"], d["profiles"], list(d["names"])


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 predict(curve):
    model, device = load_model()
    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).to(device)
    mask = (x[:, 1] == -1) & (x[:, 2] == -1)
    with torch.no_grad():
        out = model(x, mask)
    return out[0, :len(DEPTH)].cpu().numpy()


def parse_table(df, period_col, vel_col, freq_input, vel_unit):
    p = pd.to_numeric(df[period_col], errors="coerce").to_numpy(dtype=float)
    v = pd.to_numeric(df[vel_col], 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
    return p, v


# ----------------------------------------------------------------- interface
st.title("Shear-wave velocity from a phase-velocity dispersion curve")
st.markdown(
    "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).")

with st.sidebar:
    st.header("Input")
    source = st.radio("Curve source",
                      ["Example from test data", "Upload CSV", "Paste values"])
    st.divider()
    st.header("Site scale")
    st.caption("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.")
    preset = st.selectbox(
        "Scale factor k",
        ["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"])
    if preset.startswith("Custom"):
        k = float(st.number_input("k (period multiplier)", min_value=1.0,
                                  max_value=10000.0, value=100.0, step=1.0))
    elif preset.startswith("Engineering"):
        k = 100.0
    else:
        k = 1.0
    if source == "Example from test data" and k != 1.0:
        st.info("Examples are native-scale; k is applied to uploaded/pasted "
                "curves only.")
        k = 1.0
    st.caption(
        f"Accepted measured band at k = {k:g}: "
        f"{PERIOD.min()/k:.4g}–{PERIOD.max()/k:.4g} s "
        f"({k/PERIOD.max():.3g}–{k/PERIOD.min():.3g} Hz). "
        f"Output depth range: {DEPTH.max()/k*1000:.3g} m."
        if k > 1 else
        f"Model grid: {len(PERIOD)} periods, {PERIOD.min():.1f}–{PERIOD.max():.0f} s. "
        f"Output: {len(DEPTH)} layers, 0–{DEPTH.max():.2f} km (40 m spacing).")
    st.divider()
    st.caption("Picks are snapped to the nearest grid period; the model accepts "
               "gaps and band-limited curves natively. Velocity support "
               "≈ 0.3–4.5 km/s at any scale (velocities are not rescaled).")

# 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

curve = None
true_vs = None

if source == "Example from test data":
    curves, profiles, names = load_examples()
    labels = [f"{n} #{i}" for i, n in enumerate(names)]
    pick = st.sidebar.selectbox("Example", labels)
    idx = labels.index(pick)
    curve = curves[idx].copy()
    true_vs = profiles[idx]
    lo_p, hi_p = st.sidebar.slider(
        "Restrict period band (s) — simulates a band-limited survey",
        float(PERIOD.min()), float(PERIOD.max()),
        (float(PERIOD.min()), float(PERIOD.max())))
    curve[(PERIOD < lo_p) | (PERIOD > hi_p)] = -1.0

elif source == "Upload CSV":
    st.sidebar.markdown("CSV with one column of period (s) **or** frequency (Hz), "
                        "and one of phase velocity.")
    up = st.sidebar.file_uploader("CSV file", type=["csv", "txt"])
    freq_input = st.sidebar.checkbox("First column is frequency (Hz)", False)
    vel_unit = st.sidebar.radio("Velocity unit", ["km/s", "m/s"], horizontal=True)
    if up is not None:
        df = pd.read_csv(up)
        cols = list(df.columns)
        c1_, c2_ = st.sidebar.selectbox("Period/frequency column", cols, index=0), \
                   st.sidebar.selectbox("Velocity column", cols,
                                        index=min(1, len(cols) - 1))
        p, v = parse_table(df, c1_, c2_, freq_input, vel_unit)
        curve, n_used, n_out = snap_to_grid(p * k, v)
        st.sidebar.success(
            f"{n_used} grid periods filled"
            + (f" · {n_out} picks outside the accepted band dropped" if n_out else ""))

else:  # paste
    st.sidebar.markdown("One `period_s, velocity_km_s` pair per line:")
    default = "\n".join(f"{t:.3f}, {c:.3f}" for t, c in
                        zip(PERIOD[::12], load_examples()[0][0][::12])
                        if c > 0)
    txt = st.sidebar.text_area("Values", default, height=200)
    try:
        rows = [list(map(float, ln.replace(",", " ").split()))
                for ln in txt.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])
        st.sidebar.success(f"{n_used} grid periods filled")
    except Exception as e:
        st.sidebar.error(f"Could not parse input: {e}")

# ----------------------------------------------------------------- results
if curve is None or (curve > 0).sum() == 0:
    st.info("Provide a dispersion curve in the sidebar to run the inversion.")
    st.stop()

if (curve > 0).sum() < 5:
    st.warning("Very few valid picks — the prediction will be poorly constrained.")

vmin_meas = float(curve[curve > 0].min())
if vmin_meas < 0.3:
    st.warning(
        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)

col1, col2 = st.columns(2)
valid = curve > 0

with col1:
    fig, ax = plt.subplots(figsize=(5.5, 4))
    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)
    st.pyplot(fig, use_container_width=True)
    plt.close(fig)

with col2:
    fig, ax = plt.subplots(figsize=(5.5, 4))
    D = DEPTH_DISP * DSC
    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)
    st.pyplot(fig, use_container_width=True)
    plt.close(fig)

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

if true_vs is not None:
    err = vs - true_vs
    m1, m2, m3 = st.columns(3)
    m1.metric("RMSE vs truth", f"{np.sqrt((err ** 2).mean()):.3f} km/s")
    m2.metric("MAE vs truth", f"{np.abs(err).mean():.3f} km/s")
    m3.metric("MAPE vs truth", f"{(np.abs(err) / true_vs).mean() * 100:.1f} %")

out_df = pd.DataFrame({f"depth_{DUNIT}": DEPTH_DISP * DSC, "vs_km_s": vs,
                       "constrained": [(lo <= i < hi) for i in range(len(DEPTH))]})
buf = io.StringIO()
out_df.to_csv(buf, index=False)
st.download_button("Download predicted profile (CSV)", buf.getvalue(),
                   file_name="predicted_vs_profile.csv", mime="text/csv")

with st.expander("Model & protocol details"):
    st.markdown(
        "- **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.")