app.py

"""
Chromatographic Retention Time Predictor
Laboratory-conditioned retention-time prediction using archived fingerprint
neural-network fold models trained on 3,776 structure--laboratory observations
from 23 represented laboratories.

Reference: Accompanying retention-time modelling study.
"""

import os
import json
import hashlib
import warnings
warnings.filterwarnings("ignore")

import numpy as np
import torch
import torch.nn as nn
import gradio as gr
from rdkit import Chem
from rdkit.Chem import Draw, Descriptors, Crippen, rdFingerprintGenerator, DataStructs
from PIL import Image

# ─────────────────────────────────────────────────────────────────────────────
# Model architecture (FingerprintNN) — must match training exactly
# ─────────────────────────────────────────────────────────────────────────────

class FingerprintNN(nn.Module):
    """Feed-forward network for Morgan fingerprints with laboratory embedding."""

    def __init__(
        self,
        input_dim: int = 2048,
        hidden_dims=(768, 384, 192, 96),
        dropout: float = 0.15,
        use_batch_norm: bool = True,
        num_labs: int = 23,
        lab_embed_dim: int = 64,
        input_dropout: float = 0.1,
    ):
        super().__init__()
        self.lab_embedding = nn.Embedding(num_labs, lab_embed_dim)
        initial_dim = input_dim + lab_embed_dim
        self.input_dropout = nn.Dropout(input_dropout)

        layers = []
        prev_dim = initial_dim
        for idx, hidden_dim in enumerate(hidden_dims):
            layers.append(nn.Linear(prev_dim, hidden_dim))
            if use_batch_norm:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            layers.append(nn.Dropout(dropout if idx < len(hidden_dims) - 1 else dropout * 0.5))
            prev_dim = hidden_dim
        layers.append(nn.Linear(prev_dim, 1))
        self.network = nn.Sequential(*layers)

        self.target_mean: float = 0.0
        self.target_std: float = 1.0

    def forward(self, x: torch.Tensor, lab_indices: torch.Tensor) -> torch.Tensor:
        x = self.input_dropout(x)
        if lab_indices.dim() > 1:
            lab_indices = lab_indices.squeeze(-1)
        lab_emb = self.lab_embedding(lab_indices.long())
        x = torch.cat([x, lab_emb], dim=-1)
        return self.network(x).squeeze(-1)


# ─────────────────────────────────────────────────────────────────────────────
# Laboratory mapping (alphabetical LabelEncoder order, matching training)
# ─────────────────────────────────────────────────────────────────────────────

LAB_NAMES = [
    "Aarhus",
    "Academy of Forensic Science",
    "Adelaide",
    "Australian Racing Forensic Laboratory",
    "CFSRE",
    "ChemCentre",
    "Copenhagen",
    "Estonian Forensic Science Institute",
    "Finnish Customs Laboratory",
    "Ghent University",
    "IUPA, UJI I (E)",
    "King's College Hospital",
    "LADR",
    "Labor Krone",
    "Mainz",
    "Odense",
    "San Francisco OCME",
    "The University of Queensland",
    "Trondheim",
    "University Hospital of Northern Norway",
    "University of Athens",
    "Victorian Institute of Forensic Medicine",
    "Zurich Institute of Forensic Medicine",
]
LAB_TO_IDX = {name: idx for idx, name in enumerate(LAB_NAMES)}

# ─────────────────────────────────────────────────────────────────────────────
# Model loading (lazy, on first prediction)
# ─────────────────────────────────────────────────────────────────────────────

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = os.path.join(BASE_DIR, "models")
DATA_DIR = os.path.join(BASE_DIR, "data")

_models = []
_train_fps = None


def _sha256(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def _load_models():
    global _models, _train_fps
    if _models:
        return

    bundle_path = os.path.join(BASE_DIR, "model_bundle.json")
    if not os.path.isfile(bundle_path):
        raise FileNotFoundError("The versioned FPNN model bundle is not installed.")
    with open(bundle_path, "r", encoding="utf-8") as stream:
        bundle = json.load(stream)
    if bundle.get("component") != "FPNN only":
        raise ValueError("The installed bundle is not the scoped FPNN component.")
    fold_files = sorted(f for f in os.listdir(MODEL_DIR) if f.endswith(".pt"))
    if len(fold_files) != int(bundle["fold_models"]):
        raise FileNotFoundError(
            f"Expected {bundle['fold_models']} FPNN fold models; found {len(fold_files)}."
        )
    for item in bundle["files"]:
        path = os.path.join(BASE_DIR, *item["relative_path"].split("/"))
        if not os.path.isfile(path) or _sha256(path) != item["sha256"]:
            raise ValueError(f"Missing or hash-mismatched bundle file: {item['relative_path']}")
    for fname in fold_files:
        ckpt = torch.load(
            os.path.join(MODEL_DIR, fname), map_location="cpu", weights_only=False
        )
        model = FingerprintNN()
        model.load_state_dict(ckpt["model_state"])
        model.target_mean = float(ckpt["target_mean"])
        model.target_std = float(ckpt["target_std"])
        model.eval()
        _models.append(model)

    fps_path = os.path.join(DATA_DIR, "training_fps.npz")
    if not os.path.exists(fps_path):
        raise FileNotFoundError("Development-set fingerprints are missing from the bundle.")
    with np.load(fps_path) as fingerprint_bundle:
        _train_fps = fingerprint_bundle["fps"].astype(np.float32)


# ─────────────────────────────────────────────────────────────────────────────
# Molecular feature helpers
# ─────────────────────────────────────────────────────────────────────────────

def _smiles_to_fp(smiles: str):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    gen = rdFingerprintGenerator.GetMorganGenerator(
        radius=2, fpSize=2048, includeChirality=True
    )
    fp = gen.GetFingerprint(mol)
    arr = np.zeros(2048, dtype=np.float32)
    DataStructs.ConvertToNumpyArray(fp, arr)
    return arr


def _mol_descriptors(smiles: str) -> dict:
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return {}
    return {
        "MW": Descriptors.MolWt(mol),
        "LogP": Crippen.MolLogP(mol),
        "TPSA": Descriptors.TPSA(mol),
        "HBD": Descriptors.NumHDonors(mol),
        "HBA": Descriptors.NumHAcceptors(mol),
        "RotBonds": Descriptors.NumRotatableBonds(mol),
        "AromaticRings": Descriptors.NumAromaticRings(mol),
        "HeavyAtoms": mol.GetNumHeavyAtoms(),
    }


def _smiles_to_image(smiles: str):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None
    return Draw.MolToImage(mol, size=(300, 220))


# ─────────────────────────────────────────────────────────────────────────────
# Prospective structural-similarity context
# ─────────────────────────────────────────────────────────────────────────────

def _check_similarity_context(fp: np.ndarray, threshold: float = 0.4):
    if _train_fps is None:
        return float("nan"), False, "Development-set fingerprints are unavailable."
    # Binary logical reductions compute the exact Morgan-fingerprint Tanimoto
    # counts without dispatching to a BLAS matrix multiplication.  Avoiding
    # BLAS here also prevents a Windows OpenMP-runtime conflict between the
    # packaged RDKit and PyTorch wheels during first prediction.
    development = np.asarray(_train_fps) > 0
    query = np.asarray(fp) > 0
    intersections = np.logical_and(development, query).sum(axis=1, dtype=np.int32)
    unions = np.logical_or(development, query).sum(axis=1, dtype=np.int32)
    similarities = np.divide(
        intersections,
        unions,
        out=np.zeros(intersections.shape, dtype=np.float32),
        where=unions > 0,
    )
    maximum_similarity = float(np.max(similarities))
    represented = maximum_similarity >= threshold
    status = (
        f"Maximum development-set Tanimoto similarity: {maximum_similarity:.3f} "
        f"(reference threshold {threshold:.2f})."
    )
    return maximum_similarity, represented, status


def _format_prediction_summary(prediction_mean: float, prediction_sd: float) -> str:
    return (
        f"### Predicted RT: {prediction_mean:.2f} min\n\n"
        f"Uncalibrated fold-model disagreement (SD): {prediction_sd:.2f} min"
    )


# ─────────────────────────────────────────────────────────────────────────────
# Main prediction function
# ─────────────────────────────────────────────────────────────────────────────

def predict(smiles: str, lab_name: str):
    _load_models()

    smiles = smiles.strip()
    if not smiles:
        return None, "⚠️ Please enter a SMILES string.", "", "", ""

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return None, "❌ Invalid SMILES. Please check the input.", "", "", ""

    if lab_name not in LAB_TO_IDX:
        return None, f"Unknown laboratory: {lab_name}", "", "", ""

    fp = _smiles_to_fp(smiles)
    lab_idx = LAB_TO_IDX[lab_name]

    fp_tensor = torch.tensor(fp, dtype=torch.float32).unsqueeze(0)
    lab_tensor = torch.tensor([lab_idx], dtype=torch.long)

    fold_preds = []
    with torch.no_grad():
        for model in _models:
            raw = model(fp_tensor, lab_tensor).item()
            fold_preds.append(max(0.0, raw * model.target_std + model.target_mean))

    pred_mean = float(np.mean(fold_preds))
    pred_std = float(np.std(fold_preds))

    _, represented, similarity_status = _check_similarity_context(fp)
    desc = _mol_descriptors(smiles)
    mol_img = _smiles_to_image(smiles)

    rt_text = _format_prediction_summary(pred_mean, pred_std)

    ad_icon = "✅" if represented else "⚠️"
    ad_text = (
        f"{ad_icon} {similarity_status}\n\n"
        "This prospective similarity score provides structural context only; "
        "it is not a reliability guarantee or a predictive interval."
    )

    desc_text = (
        f"**MW:** {desc.get('MW', 0):.1f} Da  |  "
        f"**LogP:** {desc.get('LogP', 0):.2f}  |  "
        f"**TPSA:** {desc.get('TPSA', 0):.1f} Ų  |  "
        f"**HBD/HBA:** {int(desc.get('HBD', 0))}/{int(desc.get('HBA', 0))}  |  "
        f"**RotBonds:** {int(desc.get('RotBonds', 0))}  |  "
        f"**HeavyAtoms:** {int(desc.get('HeavyAtoms', 0))}"
    )

    fold_text = "Fold-model predictions: " + "  |  ".join(f"{p:.2f}" for p in fold_preds)

    return mol_img, rt_text, ad_text, desc_text, fold_text


# ─────────────────────────────────────────────────────────────────────────────
# Gradio interface (compatible with gradio 5.x)
# ─────────────────────────────────────────────────────────────────────────────

DESCRIPTION = """
# Chromatographic Retention Time Predictor

This interface exposes the archived **fingerprint neural-network (FPNN) fold models**, not the
GAT/GCN/ExtraTrees stack. The model requires one of the 23 laboratory labels represented during
training. The modelling table contains 3,776 structure--laboratory observations corresponding to
1,357 InChIKey connectivity groups.

> The fold-model standard deviation is uncalibrated model disagreement, not a confidence interval.
> The Tanimoto value is a prospective structural-similarity diagnostic, not a reliability guarantee.
"""

PERF_TABLE = """
The study evaluation uses molecular-identity-grouped and scaffold-aware outer holdouts over
three prespecified split repetitions. Results for the FPNN component and the full stack are reported separately
in the accompanying manuscript because they estimate performance under different validation tasks.
"""

EXAMPLES = [
    ["c1ccc2c(c1)cc1ccc3cccc4ccc2c1c34", "Aarhus"],
    ["CC(=O)Oc1ccccc1C(=O)O", "Copenhagen"],
    ["CN1CCC23c4c(ccc(O)c4OC2(CCN(C)CC3=O)C1)O", "Ghent University"],
    ["c1ccc(cc1)C(c1ccccc1)N1CCCC1", "Mainz"],
    ["CC12CCC3C(C1CCC2O)CCC4=CC(=O)CCC34C", "CFSRE"],
]

with gr.Blocks(
    title="RT Predictor — Multi-Lab Chromatography",
    theme=gr.themes.Soft(primary_hue="blue"),
) as demo:

    gr.Markdown(DESCRIPTION)
    gr.Markdown("---")

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Input")
            smiles_input = gr.Textbox(
                label="SMILES String",
                value="CC(=O)Oc1ccccc1C(=O)O",
                placeholder="e.g., c1ccc2c(c1)cc1ccc3cccc4ccc2c1c34",
                lines=2,
            )
            lab_input = gr.Dropdown(
                choices=LAB_NAMES,
                value="Aarhus",
                label="Target Laboratory",
            )
            predict_btn = gr.Button("Predict Retention Time", variant="primary")

            gr.Examples(
                examples=EXAMPLES,
                inputs=[smiles_input, lab_input],
                label="Example compounds (click to load)",
            )

        with gr.Column(scale=1):
            gr.Markdown("### Molecular Structure")
            mol_image = gr.Image(label="2D Structure", height=240)

    gr.Markdown("---")
    gr.Markdown("### Results")

    with gr.Row():
        with gr.Column(scale=1):
            rt_output = gr.Markdown()
        with gr.Column(scale=2):
            ad_output = gr.Markdown()

    desc_output = gr.Markdown()
    fold_output = gr.Markdown()

    predict_btn.click(
        fn=predict,
        inputs=[smiles_input, lab_input],
        outputs=[mol_image, rt_output, ad_output, desc_output, fold_output],
    )

    demo.load(
        fn=predict,
        inputs=[smiles_input, lab_input],
        outputs=[mol_image, rt_output, ad_output, desc_output, fold_output],
    )

    gr.Markdown("---")
    gr.Markdown("### Evaluation Scope")
    gr.Markdown(PERF_TABLE)

    gr.Markdown("""
---
**Dataset:** HighResNPS-derived forensic toxicology data; 23 represented laboratory labels
**Model repo:** [AI4deeperScience/chromatography-rt-prediction](https://huggingface.co/AI4deeperScience/chromatography-rt-prediction)
**Citation:** Manuscript under review
""")

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