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