"""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()