visuals added

app.py

@@ -1,6 +1,12 @@
 import streamlit as st
 import pandas as pd
 import io
+from app_utils.model_utils import load_model_components
+from app_utils.viz_utils import run
+from app_utils.examples import EXAMPLES
+import numpy as np
+from streamlit_plotly_events import plotly_events
+
 
 st.set_page_config(page_title="Spectra Tool Demo", layout="wide")
 
@@ -8,68 +14,70 @@ st.title("Spectra Visualization Tool")
 
 st.markdown("Provide inputs below or load one of the example datasets.")
 
-# ------------------------
-# Example presets
-# ------------------------
-examples = {
-    "Example 1": {
-        "spectra": """mz,intensity
-100,10
-150,50
-200,80
-250,40
-300,20
-""",
-        "smiles": "CCO",   # ethanol
-        "formula": "C2H6O",
-        "adduct": "[M+H]+",
-    },
-    "Example 2": {
-        "spectra": """mz,intensity
-120,15
-180,60
-240,30
-300,70
-360,25
-""",
-        "smiles": "C6H6",  # benzene
-        "formula": "C6H6",
-        "adduct": "[M+Na]+",
-    },
-}
+FIELDS = ['mzs', 'intensities', 'smiles', 'formula', 'adduct', 'precursor_mz']
+def reset_fields():
+    for field in FIELDS:
+        st.session_state[field] = ""
 
 # ------------------------
 # Session state defaults
 # ------------------------
-if "spectra" not in st.session_state:
-    st.session_state.spectra = ""
-if "smiles" not in st.session_state:
-    st.session_state.smiles = ""
-if "formula" not in st.session_state:
-    st.session_state.formula = ""
-if "adduct" not in st.session_state:
-    st.session_state.adduct = ""
+if "run_clicked" not in st.session_state:
+    st.session_state.run_clicked = False
+if "selected_spectrum_idx" not in st.session_state:
+    st.session_state.selected_spectrum_idx = None
+if "selected_node_idx" not in st.session_state:
+    st.session_state.selected_node_idx = None
+for f in FIELDS:
+    if f not in st.session_state:
+        st.session_state[f] = ""
+
+
+if "model" not in st.session_state:
+    spec_featurizer, mol_featurizer, model = load_model_components()
+    st.session_state.spec_featurizer = spec_featurizer
+    st.session_state.mol_featurizer = mol_featurizer
+    st.session_state.model = model
 
 # ------------------------
-# Example loader buttons
+# Example loader dropdown
 # ------------------------
-cols = st.columns(len(examples))
-for i, (ex_name, ex_data) in enumerate(examples.items()):
-    if cols[i].button(f"Load {ex_name}"):
-        st.session_state.spectra = ex_data["spectra"]
-        st.session_state.smiles = ex_data["smiles"]
-        st.session_state.formula = ex_data["formula"]
-        st.session_state.adduct = ex_data["adduct"]
+example_names = list(EXAMPLES.keys())
+
+# Dropdown menu for selecting example
+selected_example = st.selectbox("Choose an example:", ["-- Select --"] + example_names)
+
+# Load button
+if st.button("Load Example") and selected_example != "-- Select --":
+
+    reset_fields()
+    ex_data = EXAMPLES[selected_example]
+    st.session_state.mzs = ex_data["mzs"]
+    st.session_state.intensities = ex_data['intensities']
+    st.session_state.smiles = ex_data["smiles"]
+    st.session_state.formula = ex_data["formula"]
+    st.session_state.adduct = ex_data["adduct"]
+    st.session_state.precursor_mz = ex_data["precursor_mz"]
+
+    # reset graph
+    st.session_state.run_clicked = False
+    st.session_state.selected_spectrum_idx = None
+    st.session_state.selected_node_idx = None
 
 # ------------------------
 # Inputs
 # ------------------------
-st.subheader("Spectra (two-column CSV: mz, intensity)")
-spectra_text = st.text_area(
-    "Paste spectra data here:",
-    value=st.session_state.spectra,
-    height=150,
-    placeholder="mz,intensity\n100,10\n150,50\n..."
+st.subheader("Spectra")
+mz_input = st.text_input(
+    "m/z values (comma-separated):",
+    value=st.session_state.mzs,
+    placeholder="100,150,200,250,300"
+)
+
+intensity_input = st.text_input(
+    "Intensities (comma-separated):",
+    value=st.session_state.intensities,
+    placeholder="10,50,80,40,20"
 )
 
 st.subheader("SMILES")
@@ -81,22 +89,101 @@ formula_input = st.text_input("Enter molecular formula:", value=st.session_state
 st.subheader("Adduct")
 adduct_input = st.text_input("Enter adduct:", value=st.session_state.adduct)
 
-# ------------------------
-# Run button
-# ------------------------
+st.subheader("Precursor mz")
+precursor_input = st.text_input("Enter precursor mz:", value=st.session_state.precursor_mz)
+
+# --- Run button toggles flag ---
 if st.button("Run"):
-    st.subheader("Results")
 
-    # Try parsing spectra
-    try:
-        spectra_df = pd.read_csv(io.StringIO(spectra_text))
-        st.write("Spectra Preview:")
-        st.dataframe(spectra_df.head())
-    except Exception as e:
-        st.error(f"Could not parse spectra: {e}")
+    for f in FIELDS:
+        if not st.session_state[f]:
+            st.error(f"Field {f} is empty.")
+            reset_fields()
+            st.stop()
+
+    st.session_state.mzs = mz_input
+    st.session_state.intensities = intensity_input
+    st.session_state.smiles = smiles_input
+    st.session_state.formula = formula_input
+    st.session_state.adduct = adduct_input
+    st.session_state.precursor_mz = precursor_input
+
+    mz_input = [float(x) for x in st.session_state.mzs.split(",") if x.strip()]
+    intensity_input = [float(x) for x in st.session_state.intensities.split(",") if x.strip()]
+
+    if len(mz_input) != len(intensity_input):
+        st.error("Number of m/z values must match the number of intensty values")
+        reset_fields()
+        st.stop()
+
+    ms = np.array(list(zip(mz_input, intensity_input)))
+
+    st.session_state.fig, st.session_state.sim_norm = run(
+            ms,
+            st.session_state.smiles,
+            st.session_state.formula,
+            st.session_state.precursor_mz,
+            st.session_state.adduct,
+            st.session_state.spec_featurizer,
+            st.session_state.mol_featurizer,
+            st.session_state.model,
+            mass_diff_thresh=20,
+            precursor_intensity=1.1
+        )
+    st.session_state.selected_spectrum_idx = None
+    st.session_state.selected_node_idx = None
+    st.session_state.run_clicked = True
+
+# --- Main results ---
+if st.session_state.run_clicked:
+    st.subheader("Peak-to-Node Similarity")
+    st.text("Double click on a peak or node to visualize similarity scores")
+
+    fig = st.session_state.fig  # get the figure
+    # Apply any coloring before rendering
+    if st.session_state.selected_spectrum_idx is not None:
+        idx = st.session_state.selected_spectrum_idx
+        scores = st.session_state.sim_norm[idx, :]
+        st.session_state.fig.data[2].marker.color = scores
+        st.session_state.fig.data[0].marker.color = [
+            "red" if i == idx else "lightgray" for i in range(st.session_state.sim_norm.shape[0])
+        ]
+    elif st.session_state.selected_node_idx is not None:
+        idx = st.session_state.selected_node_idx
+        scores = st.session_state.sim_norm[:, idx]
+        st.session_state.fig.data[0].marker.color = scores
+        st.session_state.fig.data[2].marker.color = [
+            "red" if i == idx else "lightgray" for i in range(st.session_state.sim_norm.shape[1])
+        ]
+
+    # Render figure once with plotly_events
+    selected = plotly_events(
+        st.session_state.fig,
+        click_event=True,
+        hover_event=False,
+        key="events"
+    )
+
+    # Handle click and update figure immediately
+    if selected:
+        point = selected[0]
+        curve, idx = point["curveNumber"], point["pointIndex"]
+
+        if curve == 0:  # Spectrum clicked
+            st.session_state.selected_spectrum_idx = idx
+            st.session_state.selected_node_idx = None
+            scores = st.session_state.sim_norm[idx, :]
+            st.session_state.fig.data[2].marker.color = scores
+            st.session_state.fig.data[0].marker.color = [
+                "red" if i == idx else "lightgray" for i in range(st.session_state.sim_norm.shape[0])
+            ]
 
-    st.write("**SMILES:**", smiles_input)
-    st.write("**Formula:**", formula_input)
-    st.write("**Adduct:**", adduct_input)
+        elif curve == 2:  # Node clicked
+            st.session_state.selected_node_idx = idx
+            st.session_state.selected_spectrum_idx = None
+            scores = st.session_state.sim_norm[:, idx]
+            st.session_state.fig.data[0].marker.color = scores
+            st.session_state.fig.data[2].marker.color = [
+                "red" if i == idx else "lightgray" for i in range(st.session_state.sim_norm.shape[1])
+            ]
 
-    st.info("🔬 Interactive visualization will be displayed here.")

app_utils/examples.py

@@ -0,0 +1,6 @@
+import pandas as pd
+
+data = pd.read_csv("/data/yzhouc01/FILIP-MS/data/sample/data.tsv", sep='\t')
+data = data[['identifier', 'mzs', 'intensities', 'smiles', 'formula', 'precursor_mz', 'adduct']]
+
+EXAMPLES = data.set_index('identifier').to_dict('index')

app_utils/model_utils.py

@@ -0,0 +1,31 @@
+import sys
+sys.path.insert(0, "/data/yzhouc01/MassSpecGym")
+sys.path.insert(0, "/data/yzhouc01/FILIP-MS")
+
+from rdkit import RDLogger
+from mvp.utils.data import get_spec_featurizer, get_mol_featurizer, get_ms_dataset
+from mvp.utils.models import get_model
+
+import yaml
+
+# Suppress RDKit warnings and errors
+lg = RDLogger.logger()
+lg.setLevel(RDLogger.CRITICAL)
+
+# Load model and data
+
+def load_model_components():
+    param_pth = '/data/yzhouc01/FILIP-MS/experiments/20250913_optimized_filip-model/lightning_logs/version_0/hparams.yaml'
+    with open(param_pth) as f:
+        params = yaml.load(f, Loader=yaml.FullLoader)
+
+    spec_featurizer = get_spec_featurizer(params['spectra_view'], params)
+    mol_featurizer = get_mol_featurizer(params['molecule_view'], params)
+
+    # load model
+
+    checkpoint_pth = "/data/yzhouc01/FILIP-MS/experiments/20250913_optimized_filip-model/epoch=1993-train_loss=0.10.ckpt"
+    params['checkpoint_pth'] = checkpoint_pth
+    model = get_model(params['model'], params)
+
+    return spec_featurizer, mol_featurizer, model

app_utils/viz_utils.py

@@ -0,0 +1,186 @@
+
+import torch.nn.functional as F
+import torch
+import numpy as np
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+from rdkit import Chem
+from rdkit.Chem import rdDepictor
+import pandas as pd
+
+def mol_to_graph_coords(mol):
+    """Return atom coordinates and bond list for a molecule."""
+    rdDepictor.Compute2DCoords(mol)
+    conf = mol.GetConformer()
+    coords = {i: conf.GetAtomPosition(i) for i in range(mol.GetNumAtoms())}
+    bonds = [(b.GetBeginAtomIdx(), b.GetEndAtomIdx()) for b in mol.GetBonds()]
+    return coords, bonds
+
+import torch
+import torch.nn.functional as F
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+
+def interactive_attention_visualization(
+    spectral_embeds,
+    graph_embeds,
+    peak_mzs,
+    peak_intensities,
+    peak_formulas,
+    mol
+):
+    """
+    Build base Plotly figure + similarity matrix for Streamlit interactivity.
+    - Streamlit will handle clicks & recoloring using sim_norm
+    """
+
+    # --- Similarity matrix ---
+    spectral_embeds = F.normalize(spectral_embeds, p=2, dim=-1)
+    graph_embeds = F.normalize(graph_embeds, p=2, dim=-1)
+
+    similarity = torch.matmul(spectral_embeds, graph_embeds.T).detach().cpu().numpy()
+    sim_norm = (similarity - similarity.min()) / (similarity.max() - similarity.min() + 1e-8)
+
+    num_peaks, num_nodes = similarity.shape
+
+    # --- Molecule graph ---
+    coords, bonds = mol_to_graph_coords(mol)
+    atom_labels = [a.GetSymbol() for a in mol.GetAtoms()]
+    atom_x = [coords[i].x for i in range(num_nodes)]
+    atom_y = [coords[i].y for i in range(num_nodes)]
+
+    # --- Spectrum trace ---
+    spectrum_trace = go.Scatter(
+        x=peak_mzs,
+        y=peak_intensities,
+        mode='markers',  # crucial for clickable peaks
+        name="peak",
+        marker=dict(
+            size=12,
+            color="lightgray",
+            colorscale="Viridis",
+            cmin=0,
+            cmax=1,
+            colorbar=dict(title="Similarity", len=0.8, y=0.5),
+        ),
+        hovertext=[f"{f} \n ({m:,.2f}, {i:.2})" for f, m, i in zip(peak_formulas, peak_mzs, peak_intensities)],
+        hoverinfo='text',
+        customdata=list(range(num_peaks)),  # actual peak indices
+    )
+
+    # --- Graph nodes ---
+    graph_nodes = go.Scatter(
+        x=atom_x,
+        y=atom_y,
+        mode="markers+text",
+        name="node",
+        text=atom_labels,
+        textposition="middle center",
+        marker=dict(
+            size=20,
+            color="lightgray",
+            colorscale="Viridis",
+            cmin=0,
+            cmax=1,
+            colorbar=dict(title="Similarity", len=0.8, y=0.5),
+        ),
+        customdata=list(range((num_nodes+1))),
+    )
+
+    # --- Graph bonds ---
+    edge_x, edge_y = [], []
+    for i, j in bonds:
+        edge_x += [coords[i].x, coords[j].x, None]
+        edge_y += [coords[i].y, coords[j].y, None]
+    graph_edges = go.Scatter(
+        x=edge_x,
+        y=edge_y,
+        mode="lines",
+        line=dict(color="gray", width=2),
+        hoverinfo="none",
+        showlegend=False,
+    )
+
+    # --- Subplots ---
+    fig = make_subplots(
+        rows=1,
+        cols=2,
+        subplot_titles=("Spectrum", "Molecule"),
+        column_widths=[0.6, 0.4],
+    )
+
+    fig.add_trace(spectrum_trace, row=1, col=1)
+    fig.add_trace(graph_edges, row=1, col=2)
+    fig.add_trace(graph_nodes, row=1, col=2)
+
+    fig.update_xaxes(title="m/z", row=1, col=1)
+    fig.update_yaxes(title="Intensity", row=1, col=1)
+    fig.update_xaxes(visible=False, row=1, col=2)
+    fig.update_yaxes(visible=False, row=1, col=2)
+
+    fig.update_layout(showlegend=False)
+
+    return fig, sim_norm
+
+
+# ------------------------
+# Model set up
+# ------------------------
+
+from mvp.subformula_assign.utils.spectra_utils import assign_subforms
+import matchms
+
+def run(ms, smiles, formula, precursor_mz, adduct, spec_featurizer, mol_featurizer,model, mass_diff_thresh=20, precursor_intensity=1.1):
+
+    # step 1 - label peaks with formula, setup matchms spectrum
+    x = assign_subforms(formula, np.array(ms), adduct, mass_diff_thresh=mass_diff_thresh)
+    if x['output_tbl'] is None:
+        return None, None
+    
+    formulas = np.array(x['output_tbl']['formula'])
+    mzs = x['output_tbl']['mz']
+    intensities = x['output_tbl']['ms2_inten']
+    mzs = np.array([float(m) for m in mzs])
+    intensities = np.array([float(i) for i in intensities])
+
+    # add precursor if not already present
+    if formula not in formulas:    
+        mzs = np.concatenate([mzs, [float(precursor_mz)]])
+        formulas = np.concatenate([formulas, [formula]])
+        intensities = np.concatenate([intensities, [float(precursor_intensity)]])
+    else:
+        i = np.where(formulas==formula)[0]
+        intensities[i] = precursor_intensity
+
+    sorted_idx = np.argsort(mzs)
+    mzs = mzs[sorted_idx]
+    intensities = intensities[sorted_idx]
+    formulas = formulas[sorted_idx]
+
+    spectrum = matchms.Spectrum(
+            mz = mzs,
+            intensities = intensities,
+            metadata = {'precursor_mz': precursor_mz, 'formulas': formulas}
+        )
+
+    # step 2 - featurize spectra
+    spectrum_encoding = spec_featurizer['SpecFormula'](spectrum)
+
+    # step 3 - featuraize molecule
+    molecule_encoding = mol_featurizer(smiles)
+    
+    # step 4 - Embed spectra & molecules
+    model_input = {'mol': molecule_encoding, 'SpecFormula': spectrum_encoding}
+    
+    model = model.to(torch.device('cpu'))
+    model.eval()
+    with torch.no_grad():
+        spec_embed, mol_embed = model.forward(model_input, stage='test')
+
+
+    # step 5 - visualization
+    mol = Chem.MolFromSmiles(smiles)
+    fig, sim_norm = interactive_attention_visualization(spec_embed, mol_embed, mzs, intensities, formulas, mol)
+    
+    return fig, sim_norm