spark / app.py
bingyan user
Fix reconstruction unavailable: SIGALRM timeout only on main thread (Gradio runs callbacks in worker threads)
632bb58
Raw
History Blame Contribute Delete
22.8 kB
"""
SPARK (Simulation-based Posterior Amortization for Reaction Kinetics)
Amortized Bayesian inference for outer-sphere electron-transfer cyclic voltammetry.
Gradio web interface for the v4 continuous-manifold model: from one or more cyclic
voltammograms it returns a calibrated posterior over a continuous space of reaction
mechanisms, the alternative mechanisms the data admit, per-parameter posteriors, and a
posterior-predictive reconstruction. Mechanism identity is read post-hoc from which
elementary-step gates the posterior places above kinetic silence -- there is no classifier.
"""
import os
import sys
import json
from pathlib import Path
import numpy as np
import gradio as gr
# spark_app dir first (local inference.py/plotting.py), repo root appended for `manifold`
_HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(_HERE))
sys.path.append(str(_HERE.parent))
from inference import get_predictor, GATE_PRETTY
from preprocessing import parse_cv_csv, estimate_E0
from plotting import (
plot_mechanism_posterior, plot_alternatives, plot_presence,
plot_parameter_posteriors, plot_reconstruction, parameter_table,
)
from manifold import reaction_network as RN
DEMO_DIR = _HERE / "demo_data"
# lazy predictor (loaded on first request so the Space boots fast / shows errors cleanly)
_PRED = None
_PRED_ERR = None
def _predictor():
global _PRED, _PRED_ERR
if _PRED is None and _PRED_ERR is None:
try:
_PRED = get_predictor()
except Exception as e: # noqa: BLE001
_PRED_ERR = str(e)
return _PRED
def _mech_choice_label(m):
return f"{m['label']} ({m['prob']*100:.0f}%)"
def _error(msg):
# order: mech_plot, presence_plot, summary, dropdown, rec_plot, param_plot, table, state
return (None, None, f"### Error\n\n{msg}",
gr.update(choices=[], value=None), None, None, "", None)
# ---------------------------------------------------------------------------
# Core analysis: raw scans -> fast posterior outputs + a mechanism selector.
# Reconstruction + per-mechanism parameter posteriors are computed for the
# selected mechanism (top by default; user can switch via the dropdown).
# ---------------------------------------------------------------------------
def _analyze(scans, E0_V, T_K, A_cm2, C_mM, D_cm2s, n_electrons,
n_post, n_pred, n_scans):
pred = _predictor()
if pred is None:
return _error(f"Model unavailable: {_PRED_ERR}")
if not scans:
return _error("No usable voltammograms were parsed.")
C_molcm3 = float(C_mM) * 1e-6 if C_mM else 1e-6
n = int(n_electrons) if n_electrons else 1
T = float(T_K) if T_K else 298.15
A = float(A_cm2) if A_cm2 else 0.0707
D = float(D_cm2s) if D_cm2s else 1e-5
n_scans = int(n_scans); n_pred = int(n_pred)
e0 = float(E0_V) if E0_V else float(np.median(
[estimate_E0(s["E_V"], s["i_A"]) for s in scans]))
try:
exp = pred.build_exp(scans, E0_V=e0, T_K=T, A_cm2=A, C_A_molcm3=C_molcm3,
D_A_cm2s=D, n=n)
s = pred.sample_posterior(exp, n_scans=n_scans, n_post=int(n_post))
except Exception as e: # noqa: BLE001
return _error(f"Inference failed: {e}")
mechs, _ = pred.mechanism_posterior(s, top_k=8)
fig_mech = plot_mechanism_posterior(mechs) # fast: probabilities only
fig_pres = plot_presence(pred.presence(s), GATE_PRETTY)
summary = _summary_md(mechs)
choices = [_mech_choice_label(m) for m in mechs]
labelmap = {c: m["gates"] for c, m in zip(choices, mechs)}
state = {"s": s, "exp": exp, "labelmap": labelmap,
"n_scans": n_scans, "n_pred": n_pred}
# reconstruct + inspect the top mechanism for the initial view
top = mechs[0]
info = pred.inspect_mechanism(s, top["gates"], exp, n_scans=n_scans, n_pred=n_pred)
fig_rec = plot_reconstruction(info["panels"], top["label"], info["nrmse"])
fig_par = plot_parameter_posteriors(info["subset"], info["params"], RN.SLOT_IDX)
table = "### Parameter posteriors (selected mechanism)\n\n" + parameter_table(info["params"])
dd = gr.update(choices=choices, value=choices[0])
return fig_mech, fig_pres, summary, dd, fig_rec, fig_par, table, state
def on_mechanism_change(label, state):
"""Reconstruct + show parameter posteriors for the mechanism picked in the dropdown."""
if not state or not label:
return None, None, ""
pred = _predictor()
gates = state["labelmap"].get(label, [])
try:
info = pred.inspect_mechanism(state["s"], gates, state["exp"],
n_scans=state["n_scans"], n_pred=state["n_pred"])
except Exception as e: # noqa: BLE001
return None, None, f"Reconstruction failed: {e}"
name = label.split(" (")[0]
fig_rec = plot_reconstruction(info["panels"], name, info["nrmse"])
fig_par = plot_parameter_posteriors(info["subset"], info["params"], RN.SLOT_IDX)
table = "### Parameter posteriors (selected mechanism)\n\n" + parameter_table(info["params"])
return fig_rec, fig_par, table
def _summary_md(mechs):
lines = ["### Mechanism inference"]
top = mechs[0]
lines.append(f"Most probable mechanism: **{top['label']}** "
f"(posterior probability {top['prob']*100:.0f}%).")
strong = [m for m in mechs if m["prob"] >= 0.15]
if len(strong) >= 2:
alt = "; ".join(f"{m['label']} ({m['prob']*100:.0f}%)" for m in strong)
lines.append(
f"\n**This voltammogram is consistent with more than one mechanism:** {alt}. "
"Use the selector below to reconstruct each and inspect its parameter posterior. "
"A single point fit, given one assumed mechanism, would report only one of these.")
else:
lines.append("\nUse the selector below to reconstruct a mechanism and inspect its "
"parameter posterior.")
return "\n".join(lines)
# ---------------------------------------------------------------------------
# CSV entry point
# ---------------------------------------------------------------------------
def analyze_cv_csv(files, scan_rates_text, E0_V, T_K, A_cm2, C_mM, D_cm2s,
n_electrons, n_post, n_pred, n_scans):
if not files:
return _error("Please upload at least one CSV file (one per scan rate).")
rates_text = (scan_rates_text or "").strip()
user_rates = None
if rates_text:
try:
user_rates = [float(s.strip()) for s in rates_text.split(",")]
except ValueError:
return _error("Invalid scan rates. Enter comma-separated numbers in V/s.")
if len(user_rates) != len(files):
return _error(f"Number of files ({len(files)}) must match number of "
f"scan rates ({len(user_rates)}).")
scans = []
for idx, f in enumerate(files):
try:
parsed = parse_cv_csv(Path(f.name).read_text())
except Exception as e: # noqa: BLE001
return _error(f"Could not parse '{Path(f.name).name}': {e}")
if user_rates is not None:
v = user_rates[idx]
elif "scan_rate_Vs" in parsed:
v = parsed["scan_rate_Vs"]
else:
return _error(f"Cannot determine scan rate for '{Path(f.name).name}'. "
"Provide scan rates (V/s) or include a Time (s) column.")
scans.append({"E_V": parsed["E_V"], "i_A": parsed["i_A"], "v_Vs": v})
return _analyze(scans, E0_V, T_K, A_cm2, C_mM, D_cm2s, n_electrons,
n_post, n_pred, n_scans)
# ---------------------------------------------------------------------------
# Image entry point (digitize plot images)
# ---------------------------------------------------------------------------
def analyze_cv_image(files, scan_rates_text, E0_V, threshold, n_post, n_pred,
n_scans, x_min, x_max, y_min, y_max):
if not files:
return _error("Please upload at least one plot image.")
try:
from digitizer import digitize_plot, auto_detect_axis_bounds
from PIL import Image as PILImage
except ImportError:
return _error("Image digitization requires opencv-python-headless and Pillow.")
rates_text = (scan_rates_text or "").strip()
if not rates_text:
return _error("Please enter scan rate(s) (V/s), comma-separated.")
try:
rates = [float(s.strip()) for s in rates_text.split(",")]
except ValueError:
return _error("Invalid scan rates.")
if len(rates) == 1 and len(files) > 1:
rates = rates * len(files)
if len(rates) != len(files):
return _error(f"Number of scan rates ({len(rates)}) must match images ({len(files)}).")
has_bounds = all(v not in (None, 0) for v in [x_min, x_max, y_min, y_max])
scans = []
for idx, f in enumerate(files):
fpath = f.name if hasattr(f, "name") else str(f)
try:
img = np.array(PILImage.open(fpath).convert("RGB"))
except Exception as e: # noqa: BLE001
return _error(f"Could not read image {idx+1}: {e}")
if has_bounds:
b = {"x_min": float(x_min), "x_max": float(x_max),
"y_min": float(y_min), "y_max": float(y_max)}
else:
try:
b = auto_detect_axis_bounds(img)
except Exception: # noqa: BLE001 (e.g. easyocr unavailable)
b = None
if b is None:
return _error(f"Could not auto-detect axes for image {idx+1}. "
"Enter E min/max and I min/max under 'Axis overrides' and retry.")
try:
E_V, I_raw = digitize_plot(img, b["x_min"], b["x_max"], b["y_min"], b["y_max"],
threshold=int(threshold))
except Exception as e: # noqa: BLE001
return _error(f"Digitization failed for image {idx+1}: {e}")
i_max = np.max(np.abs(I_raw))
i_A = I_raw * (1e-6 if i_max > 100 else 1e-3 if i_max > 0.1 else 1.0)
scans.append({"E_V": E_V, "i_A": i_A, "v_Vs": rates[idx]})
return _analyze(scans, E0_V, 298.15, 0.0707, 1.0, 1e-5, 1, n_post, n_pred, n_scans)
# ---------------------------------------------------------------------------
# Demo examples
# ---------------------------------------------------------------------------
def _demo_examples():
"""Group bundled demo CSVs by mechanism -> (label, [csv paths], scan_rates, E0)."""
out = []
# os_* = real experimental case studies; sim_* = synthetic held-out test examples
metas = sorted(DEMO_DIR.glob("os_*_metadata.json")) + sorted(DEMO_DIR.glob("sim_*_metadata.json"))
for meta in metas:
try:
m = json.load(open(meta))
except Exception: # noqa: BLE001
continue
stem = meta.name.replace("_metadata.json", "")
csvs = sorted(DEMO_DIR.glob(f"{stem}_*mVs.csv"))
if not csvs:
continue
# files are sorted by zero-padded mV (ascending); pair with ascending rates
out.append({"label": m.get("system", stem), "csvs": [str(c) for c in csvs],
"rates": sorted(m.get("scan_rates_Vs", [])),
"E0": m.get("physical_params", {}).get("E0_V", 0.0)})
return out
# ---------------------------------------------------------------------------
# UI: theme, CSS, hero header, About
# ---------------------------------------------------------------------------
HERO_HTML = """
<div id="hero">
<div class="hero-title">SPARK</div>
<div class="hero-sub">Simulation-based Posterior Amortization for Reaction Kinetics</div>
<div class="hero-tag">
Calibrated Bayesian inference of electrochemical reaction mechanisms from cyclic
voltammetry in a single forward pass. SPARK returns a full posterior over a
<b>continuous space of mechanisms</b>, so it reports not just the single most likely
mechanism but the <b>alternatives your data genuinely support</b>, with honest uncertainty.
</div>
<div class="hero-flow">
upload voltammograms &nbsp;&rarr;&nbsp; posterior over mechanisms &nbsp;&rarr;&nbsp;
pick a mechanism to reconstruct &amp; inspect its parameters
</div>
</div>
"""
CSS = """
.gradio-container {max-width: 1200px !important; margin: auto !important;}
#hero {text-align: center; padding: 22px 18px 18px; margin-bottom: 8px;
border-radius: 16px;
background: linear-gradient(135deg, rgba(99,102,241,0.12), rgba(16,185,129,0.12));
border: 1px solid rgba(99,102,241,0.20);}
#hero .hero-title {font-size: 2.6rem; font-weight: 800; letter-spacing: 2px;
background: linear-gradient(90deg,#4f46e5,#059669); -webkit-background-clip: text;
-webkit-text-fill-color: transparent; background-clip: text;}
#hero .hero-sub {font-size: 1.05rem; font-weight: 600; color: #475569; margin-top: 2px;}
#hero .hero-tag {font-size: 0.97rem; color: #334155; max-width: 820px; margin: 12px auto 0;
line-height: 1.5;}
#hero .hero-flow {font-size: 0.9rem; color: #4f46e5; margin-top: 12px; font-weight: 600;}
.card {border: 1px solid rgba(100,116,139,0.18); border-radius: 14px; padding: 14px 16px;
box-shadow: 0 1px 3px rgba(15,23,42,0.06);}
.verdict {background: rgba(16,185,129,0.06); border-color: rgba(16,185,129,0.28);}
.gr-plot {min-height: 340px;}
"""
ABOUT_MD = """
### What SPARK does
SPARK turns one or more cyclic voltammograms into a **calibrated posterior over a continuous
space of reaction mechanisms**. A single amortized neural model returns a joint posterior over
a 42-dimensional reaction-network parameter vector in milliseconds; mechanism identity is read
post-hoc from which elementary-step "gates" the posterior places above kinetic silence. There
is no discrete classifier and no per-mechanism refitting.
### Why a posterior over mechanisms
A voltammogram often does not determine a unique mechanism. Because SPARK's posterior spans
the whole manifold, it surfaces the **alternative mechanisms the data genuinely admit** and how
strongly &mdash; something a single point fit (DigiElch / gradient fitting), handed one assumed
mechanism, cannot discover. Use the mechanism selector to reconstruct each candidate and compare
its fit and parameter posterior.
### Elementary steps covered
Built on a unified Marcus-Hush-Chidsey / Butler-Volmer / Nernst electron-transfer law: plain
electron transfer (E), following chemical reaction (EC), second electron transfer (ECE),
disproportionation (DISP), catalytic regeneration (EC'), preceding equilibrium (CE),
two-electron (EE), radical-radical dimerization, radical-substrate coupling, proton-coupled
electron transfer (PCET), surface adsorption (Laviron), plus the double-layer capacitance.
### How to use
Upload potentiostat CSVs (one per scan rate; columns: potential in V, current in A/mA/µA, and
optionally time in s to auto-detect the scan rate) or plot images. Enter scan rates and, if
known, the formal potential and cell parameters for accurate nondimensionalization. Then pick a
mechanism to reconstruct and inspect its parameter posterior.
### Model
Continuous-manifold simulation-based inference (calibrated RQ-NSF density estimator) trained on a
coupled reaction-diffusion-adsorption voltammetry simulator. Outer-sphere electron transfer.
"""
def _output_components():
"""Returns (click_outputs, mech_dropdown, state, per_mech_outputs).
click_outputs order matches _analyze's return:
[mech_plot, pres_plot, summary_md, mech_dropdown, rec_plot, par_plot, table_md, state]
per_mech_outputs (for the dropdown change): [rec_plot, par_plot, table_md]."""
state = gr.State(None)
# 1) verdict card
with gr.Group(elem_classes="card verdict"):
summary_md = gr.Markdown("Run an analysis to see the mechanism posterior.")
# 2) overview: posterior over mechanisms + elementary-step presence
with gr.Group(elem_classes="card"):
gr.Markdown("#### Posterior over mechanisms")
with gr.Row():
mech_plot = gr.Plot(show_label=False)
pres_plot = gr.Plot(show_label=False)
# 3) prominent mechanism selector
with gr.Group(elem_classes="card"):
gr.Markdown("#### Reconstruct & inspect a mechanism")
mech_dropdown = gr.Dropdown(
label="Mechanism (posterior probability shown)", choices=[], interactive=True)
with gr.Row():
rec_plot = gr.Plot(show_label=False)
par_plot = gr.Plot(show_label=False)
with gr.Accordion("Parameter table (selected mechanism)", open=False):
table_md = gr.Markdown()
click_outputs = [mech_plot, pres_plot, summary_md, mech_dropdown,
rec_plot, par_plot, table_md, state]
return click_outputs, mech_dropdown, state, [rec_plot, par_plot, table_md]
def _example_blocks(inputs):
"""Two labeled example blocks: experimental case studies + simulated test examples."""
ex = _demo_examples()
real = [e for e in ex if not e["label"].lower().startswith("simulated")]
sim = [e for e in ex if e["label"].lower().startswith("simulated")]
def rows(items):
return [[e["csvs"], ", ".join(f"{r:.3g}" for r in e["rates"]), e["E0"]] for e in items]
if real:
gr.Examples(examples=rows(real), inputs=inputs, label="Experimental case studies",
example_labels=[e["label"] for e in real])
if sim:
gr.Examples(examples=rows(sim), inputs=inputs, label="Simulated test examples (known ground truth)",
example_labels=[e["label"] for e in sim])
def build_app():
theme = gr.themes.Soft(primary_hue="indigo", secondary_hue="emerald", neutral_hue="slate")
with gr.Blocks(title="SPARK", theme=theme, css=CSS) as demo:
gr.HTML(HERO_HTML)
with gr.Tabs():
# ---- CSV tab ----
with gr.Tab("Analyze CV (CSV)"):
with gr.Row():
with gr.Column(scale=4, min_width=320):
with gr.Group(elem_classes="card"):
gr.Markdown("#### Input")
csv_files = gr.File(file_count="multiple",
label="CV CSV files (one per scan rate)")
csv_rates = gr.Textbox(label="Scan rates (V/s, comma-separated)",
placeholder="e.g. 0.05, 0.1, 0.2")
with gr.Accordion("Cell parameters (for nondimensionalization)", open=False):
csv_e0 = gr.Number(label="Formal potential E0 (V) - blank = auto", value=None)
csv_T = gr.Number(label="Temperature (K)", value=298.15)
csv_A = gr.Number(label="Electrode area (cm^2)", value=0.0707)
csv_C = gr.Number(label="Concentration (mM)", value=1.0)
csv_D = gr.Number(label="Diffusion coeff. (cm^2/s)", value=1e-5)
csv_n = gr.Number(label="Electrons n", value=1, precision=0)
with gr.Accordion("Inference settings", open=False):
csv_npost = gr.Slider(500, 4000, value=2000, step=500, label="Posterior samples")
csv_npred = gr.Slider(4, 24, value=8, step=2, label="Reconstruction draws")
csv_nscans = gr.Slider(1, 3, value=3, step=1, label="Scans used")
csv_btn = gr.Button("Analyze", variant="primary", size="lg")
_example_blocks([csv_files, csv_rates, csv_e0])
with gr.Column(scale=7, min_width=480):
csv_out, csv_dd, csv_state, csv_permech = _output_components()
csv_btn.click(
analyze_cv_csv,
inputs=[csv_files, csv_rates, csv_e0, csv_T, csv_A, csv_C, csv_D,
csv_n, csv_npost, csv_npred, csv_nscans],
outputs=csv_out)
csv_dd.change(on_mechanism_change, inputs=[csv_dd, csv_state],
outputs=csv_permech)
# ---- Image tab ----
with gr.Tab("Analyze CV (image)"):
with gr.Row():
with gr.Column(scale=4, min_width=320):
with gr.Group(elem_classes="card"):
gr.Markdown("#### Input")
img_files = gr.File(file_count="multiple",
label="CV plot images (one per scan rate)")
img_rates = gr.Textbox(label="Scan rates (V/s, comma-separated)",
placeholder="e.g. 0.05, 0.1, 0.2")
img_e0 = gr.Number(label="Formal potential E0 (V) - blank = auto", value=None)
img_thr = gr.Slider(20, 200, value=90, step=5, label="Digitization threshold")
with gr.Accordion("Axis overrides (if auto-detect fails)", open=False):
img_xmin = gr.Number(label="E min (V)", value=None)
img_xmax = gr.Number(label="E max (V)", value=None)
img_ymin = gr.Number(label="I min", value=None)
img_ymax = gr.Number(label="I max", value=None)
with gr.Accordion("Inference settings", open=False):
img_npost = gr.Slider(500, 4000, value=2000, step=500, label="Posterior samples")
img_npred = gr.Slider(4, 24, value=8, step=2, label="Reconstruction draws")
img_nscans = gr.Slider(1, 3, value=3, step=1, label="Scans used")
img_btn = gr.Button("Analyze", variant="primary", size="lg")
with gr.Column(scale=7, min_width=480):
img_out, img_dd, img_state, img_permech = _output_components()
img_btn.click(
analyze_cv_image,
inputs=[img_files, img_rates, img_e0, img_thr, img_npost, img_npred,
img_nscans, img_xmin, img_xmax, img_ymin, img_ymax],
outputs=img_out)
img_dd.change(on_mechanism_change, inputs=[img_dd, img_state],
outputs=img_permech)
with gr.Tab("About"):
gr.Markdown(ABOUT_MD)
return demo
if __name__ == "__main__":
build_app().launch(server_name="0.0.0.0", server_port=7860)