Spaces:
Running on Zero
Running on Zero
| import random | |
| import re | |
| from typing import Any | |
| import numpy as np | |
| import torch | |
| # Suppress RDKit C++ stderr noise | |
| try: | |
| from rdkit import RDLogger | |
| RDLogger.DisableLog('rdApp.*') | |
| except Exception: | |
| pass | |
| # ───────────────────────────────────────────────────────────────── | |
| # Atom & Bond Lookup Tables | |
| # ───────────────────────────────────────────────────────────────── | |
| ATOM_NUMBERS = { | |
| 'H': 1, 'C': 6, 'N': 7, 'O': 8, 'F': 9, | |
| 'P': 15, 'S': 16, 'Cl': 17, 'Br': 35, 'I': 53, | |
| 'Si': 14, 'B': 5, 'Se': 34, 'Na': 11, 'K': 19, 'Pt': 78, | |
| } | |
| ATOM_MASSES = { | |
| 'H': 1.008, 'C': 12.011, 'N': 14.007, 'O': 15.999, 'F': 18.998, | |
| 'P': 30.974, 'S': 32.06, 'Cl': 35.45, 'Br': 79.904, 'I': 126.90, | |
| 'Si': 28.085, 'B': 10.811, 'Se': 78.971, 'Na': 22.990, 'K': 39.098, 'Pt': 195.08, | |
| } | |
| ELECTRONEGATIVITY = { | |
| 'H': 2.20, 'C': 2.55, 'N': 3.04, 'O': 3.44, 'F': 3.98, | |
| 'P': 2.19, 'S': 2.58, 'Cl': 3.16, 'Br': 2.96, 'I': 2.66, | |
| 'Si': 1.90, 'B': 2.04, 'Se': 2.55, 'Na': 0.93, 'K': 0.82, 'Pt': 2.28, | |
| } | |
| VDW_RADIUS = { | |
| 'H': 1.20, 'C': 1.70, 'N': 1.55, 'O': 1.52, 'F': 1.47, | |
| 'P': 1.80, 'S': 1.80, 'Cl': 1.75, 'Br': 1.85, 'I': 1.98, | |
| 'Si': 2.10, 'B': 1.92, 'Se': 1.90, 'Na': 2.27, 'K': 2.75, 'Pt': 1.75, | |
| } | |
| HALOGENS = {'F', 'Cl', 'Br', 'I'} | |
| HETEROATOMS = {'P', 'S', 'Se', 'B', 'Si', 'Pt'} | |
| def smiles_to_graph(smiles: str) -> tuple[torch.Tensor, torch.Tensor, list[str]]: | |
| """ | |
| Converts a SMILES string into a rich molecular graph with 24-dimensional atom features x_i in R^24 | |
| and edge_index E in R^(2 x M). | |
| Node Features x_i in R^24: | |
| [0] Atomic Number / 100.0 | |
| [1] Atomic Mass / 200.0 | |
| [2] Degree / 6.0 | |
| [3] Formal Charge clamped [-2, +2] / 2.0 | |
| [4] Hybridization code (1=sp, 2=sp2, 3=sp3, 4=sp3d, 5=sp3d2, 0=other) / 5.0 | |
| [5] Is Aromatic (binary) | |
| [6] Implicit Valence / 6.0 | |
| [7] Is In Ring (binary) | |
| [8] Is Stereocenter / Chiral Flag (binary) | |
| [9] Total Hydrogen Count / 4.0 | |
| [10-13] Element Group One-Hot: [C/N/O, Halogens, Heteroatoms, Other] | |
| [14] Pauling Electronegativity / 4.0 | |
| [15] vdW Radius / 3.0 | |
| [16] Is Ring Size 3 (binary) | |
| [17] Is Ring Size 4 (binary) | |
| [18] Is Ring Size 5 (binary) | |
| [19] Is Ring Size 6 (binary) | |
| [20] Is Ring Size 7 (binary) | |
| [21] Is Ring Size 8 (binary) | |
| [22] Is Conjugated Atom (binary) | |
| [23] Gasteiger Charge Proxy (normalized [-1, +1]) | |
| """ | |
| def _atom_to_features(symbol, num, deg, chg, hyb, aromatic, | |
| imp_val, in_ring, mass, chiral, h_count, | |
| ring_sizes, conjugated, charge_proxy) -> list[float]: | |
| group = [0.0, 0.0, 0.0, 0.0] | |
| if symbol in {'C', 'N', 'O'}: | |
| group[0] = 1.0 | |
| elif symbol in HALOGENS: | |
| group[1] = 1.0 | |
| elif symbol in HETEROATOMS: | |
| group[2] = 1.0 | |
| else: | |
| group[3] = 1.0 | |
| en = ELECTRONEGATIVITY.get(symbol, 2.0) / 4.0 | |
| vdw = VDW_RADIUS.get(symbol, 1.7) / 3.0 | |
| r3 = 1.0 if 3 in ring_sizes else 0.0 | |
| r4 = 1.0 if 4 in ring_sizes else 0.0 | |
| r5 = 1.0 if 5 in ring_sizes else 0.0 | |
| r6 = 1.0 if 6 in ring_sizes else 0.0 | |
| r7 = 1.0 if 7 in ring_sizes else 0.0 | |
| r8 = 1.0 if 8 in ring_sizes else 0.0 | |
| return [ | |
| float(num) / 100.0, # [0] | |
| float(mass) / 200.0, # [1] | |
| min(float(deg), 6.0) / 6.0, # [2] | |
| max(-2.0, min(2.0, float(chg))) / 2.0, # [3] | |
| float(hyb) / 5.0, # [4] | |
| float(aromatic), # [5] | |
| min(float(imp_val), 6.0) / 6.0, # [6] | |
| float(in_ring), # [7] | |
| float(chiral), # [8] | |
| min(float(h_count), 4.0) / 4.0, # [9] | |
| *group, # [10-13] | |
| en, # [14] | |
| vdw, # [15] | |
| r3, r4, r5, r6, r7, r8, # [16-21] | |
| float(conjugated), # [22] | |
| max(-1.0, min(1.0, float(charge_proxy))), # [23] | |
| ] | |
| try: | |
| from rdkit import Chem | |
| mol = Chem.MolFromSmiles(smiles) | |
| if mol is not None: | |
| Chem.SanitizeMol(mol) | |
| atoms, atom_symbols = [], [] | |
| for atom in mol.GetAtoms(): | |
| symbol = atom.GetSymbol() | |
| num = atom.GetAtomicNum() | |
| deg = atom.GetDegree() | |
| chg = atom.GetFormalCharge() | |
| hyb_val = int(atom.GetHybridization()) | |
| hyb = {2: 1, 3: 2, 4: 3, 5: 4, 6: 5}.get(hyb_val, 0) | |
| aromatic = 1.0 if atom.GetIsAromatic() else 0.0 | |
| try: | |
| imp_val = float(atom.GetValence(Chem.ValenceType.IMPLICIT)) | |
| except Exception: | |
| imp_val = float(atom.GetImplicitValence()) | |
| in_ring = 1.0 if atom.IsInRing() else 0.0 | |
| mass = float(atom.GetMass()) | |
| chiral = 1.0 if (atom.HasProp('_ChiralityPossible') or atom.GetChiralTag() != Chem.ChiralType.CHI_UNSPECIFIED) else 0.0 | |
| h_count = float(atom.GetTotalNumHs()) | |
| ring_sizes = [size for size in range(3, 9) if atom.IsInRingSize(size)] | |
| conjugated = 1.0 if atom.GetIsAromatic() or any(b.GetIsConjugated() for b in atom.GetBonds()) else 0.0 | |
| charge_proxy = float(chg) + (0.1 if symbol in {'N', 'O'} else (-0.1 if symbol in {'C'} else 0.0)) | |
| feats = _atom_to_features( | |
| symbol, num, deg, chg, hyb, aromatic, | |
| imp_val, in_ring, mass, chiral, h_count, | |
| ring_sizes, conjugated, charge_proxy | |
| ) | |
| atoms.append(feats) | |
| atom_symbols.append(symbol) | |
| edges = [] | |
| for bond in mol.GetBonds(): | |
| i = bond.GetBeginAtomIdx() | |
| j = bond.GetEndAtomIdx() | |
| edges.extend([[i, j], [j, i]]) | |
| if not edges: | |
| edges = [[0, 0]] | |
| node_feats = torch.tensor(atoms, dtype=torch.float32) | |
| edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous() | |
| return node_feats, edge_index, atom_symbols | |
| except Exception: | |
| pass | |
| # ── Regex Fallback Parser ─────────────────────────────────── | |
| tokens = re.findall(r'Cl|Br|Si|Se|Pt|[A-Z][a-z]?|[a-z]|[\=\#\-\+\(\)]', smiles) | |
| atoms, atom_symbols, edges = [], [], [] | |
| stack, prev_idx = [], None | |
| for tok in tokens: | |
| sym = tok.upper() if tok.isalpha() else tok | |
| if sym in ATOM_NUMBERS or tok in ATOM_NUMBERS: | |
| key = sym if sym in ATOM_NUMBERS else tok | |
| num = ATOM_NUMBERS.get(key, 6) | |
| mass = ATOM_MASSES.get(key, 12.0) | |
| aromatic = 1.0 if tok.islower() else 0.0 | |
| idx = len(atoms) | |
| feats = _atom_to_features( | |
| key, num, deg=2 if aromatic else 1, chg=0, hyb=2 if aromatic else 3, | |
| aromatic=aromatic, imp_val=0.0, in_ring=aromatic, mass=mass, | |
| chiral=0.0, h_count=1.0, ring_sizes=[6] if aromatic else [], | |
| conjugated=aromatic, charge_proxy=0.0 | |
| ) | |
| atoms.append(feats) | |
| atom_symbols.append(key) | |
| if prev_idx is not None: | |
| edges.extend([[prev_idx, idx], [idx, prev_idx]]) | |
| prev_idx = idx | |
| elif tok == '(': | |
| if prev_idx is not None: | |
| stack.append(prev_idx) | |
| elif tok == ')': | |
| if stack: | |
| prev_idx = stack.pop() | |
| if not atoms: | |
| atoms = [_atom_to_features('C', 6, 1, 0, 3, 0, 0, 0, 12.011, 0, 1, [], 0, 0.0)] | |
| atom_symbols = ['C'] | |
| edges = [[0, 0]] | |
| if not edges: | |
| edges = [[0, 0]] | |
| node_feats = torch.tensor(atoms, dtype=torch.float32) | |
| edge_index = torch.tensor(edges, dtype=torch.long).t().contiguous() | |
| return node_feats, edge_index, atom_symbols | |
| # ───────────────────────────────────────────────────────────────── | |
| # Bemis-Murcko Scaffold Splitter | |
| # ───────────────────────────────────────────────────────────────── | |
| def get_bemis_murcko_scaffold(smiles: str) -> str: | |
| try: | |
| from rdkit import Chem | |
| from rdkit.Chem.Scaffolds import MurckoScaffold | |
| mol = Chem.MolFromSmiles(smiles) | |
| if mol is not None: | |
| return MurckoScaffold.MurckoScaffoldSmiles(mol=mol, includeChirality=False) | |
| except Exception: | |
| pass | |
| rings = re.findall(r'c1[a-z0-9\=\#\-]+1|C1[A-Za-z0-9\=\#\-]+1', smiles) | |
| if rings: | |
| return "-".join(sorted(rings)) | |
| c_count = smiles.upper().count('C') | |
| return f"Framework_C{c_count}" | |
| def bemis_murcko_scaffold_split( | |
| dataset, | |
| smiles_list: list[str], | |
| frac_train: float = 0.8, | |
| frac_val: float = 0.1, | |
| frac_test: float = 0.1, | |
| seed: int = 42, | |
| ) -> tuple[list[int], list[int], list[int]]: | |
| scaffolds: dict[str, list[int]] = {} | |
| for idx, smi in enumerate(smiles_list): | |
| sc = get_bemis_murcko_scaffold(smi) | |
| scaffolds.setdefault(sc, []).append(idx) | |
| scaffold_sets = sorted(scaffolds.values(), key=len, reverse=True) | |
| rng = random.Random(seed) | |
| rng.shuffle(scaffold_sets) | |
| total = len(dataset) | |
| train_cut = int(frac_train * total) | |
| val_cut = int((frac_train + frac_val) * total) | |
| train_idx, val_idx, test_idx = [], [], [] | |
| for cluster in scaffold_sets: | |
| if len(train_idx) + len(cluster) <= train_cut: | |
| train_idx.extend(cluster) | |
| elif len(train_idx) + len(val_idx) + len(cluster) <= val_cut: | |
| val_idx.extend(cluster) | |
| else: | |
| test_idx.extend(cluster) | |
| return train_idx, val_idx, test_idx | |
| def random_split( | |
| dataset, | |
| frac_train: float = 0.8, | |
| frac_val: float = 0.1, | |
| frac_test: float = 0.1, | |
| seed: int = 42, | |
| ) -> tuple[list[int], list[int], list[int]]: | |
| indices = list(range(len(dataset))) | |
| random.Random(seed).shuffle(indices) | |
| n = len(indices) | |
| train_cut = int(frac_train * n) | |
| val_cut = int((frac_train + frac_val) * n) | |
| return indices[:train_cut], indices[train_cut:val_cut], indices[val_cut:] | |
| # ───────────────────────────────────────────────────────────────── | |
| # XAI — Toxic Hotspot Highlighting | |
| # ───────────────────────────────────────────────────────────────── | |
| def highlight_toxic_subgraph( | |
| smiles: str, | |
| attention_scores: torch.Tensor, | |
| top_k: int = 3, | |
| ) -> dict[str, Any]: | |
| _node_feats, _edge_index, atom_symbols = smiles_to_graph(smiles) | |
| num_nodes = len(atom_symbols) | |
| if attention_scores is None or len(attention_scores) == 0: | |
| scores = np.ones(num_nodes) / max(num_nodes, 1) | |
| else: | |
| scores = attention_scores.detach().cpu().numpy() | |
| if len(scores) < num_nodes: | |
| scores = np.pad(scores, (0, num_nodes - len(scores)), 'constant') | |
| elif len(scores) > num_nodes: | |
| scores = scores[:num_nodes] | |
| max_s = np.max(scores) if np.max(scores) > 0 else 1.0 | |
| norm = scores / max_s | |
| top_idxs = np.argsort(norm)[::-1][:min(top_k, num_nodes)].tolist() | |
| hotspots = [ | |
| { | |
| "atom_index": int(i), | |
| "atom_symbol": atom_symbols[i], | |
| "attention_score": round(float(norm[i]), 4), | |
| "is_toxic_hotspot": True, | |
| } | |
| for i in top_idxs | |
| ] | |
| return { | |
| "smiles": smiles, | |
| "total_atoms": num_nodes, | |
| "atom_symbols": atom_symbols, | |
| "attention_weights": [round(float(s), 4) for s in norm], | |
| "top_toxic_hotspots": hotspots, | |
| "plot_title": "GAT Layer Attention Distribution Map (Highlight = High Attention Weight)", | |
| } | |
| def calculate_tanimoto_applicability_domain( | |
| query_smiles: str, | |
| training_smiles_list: list[str] | |
| ) -> dict[str, Any]: | |
| """ | |
| Computes maximum Tanimoto similarity between query molecule and training dataset. | |
| Flags applicability domain confidence: | |
| - High Confidence: Max Tanimoto >= 0.70 | |
| - Moderate Confidence: 0.40 <= Max Tanimoto < 0.70 | |
| - Low Confidence (Out of Domain): Max Tanimoto < 0.40 | |
| """ | |
| try: | |
| from rdkit import Chem, DataStructs | |
| from rdkit.Chem import RDKFingerprint | |
| q_mol = Chem.MolFromSmiles(query_smiles) | |
| if q_mol is None: | |
| return {"max_tanimoto": 0.0, "applicability_domain": "Out-of-Domain (Invalid SMILES)"} | |
| q_fp = RDKFingerprint(q_mol) | |
| max_sim = 0.0 | |
| for tr_smi in training_smiles_list: | |
| tr_mol = Chem.MolFromSmiles(tr_smi) | |
| if tr_mol is not None: | |
| tr_fp = RDKFingerprint(tr_mol) | |
| sim = DataStructs.TanimotoSimilarity(q_fp, tr_fp) | |
| max_sim = max(max_sim, sim) | |
| max_sim = round(float(max_sim), 4) | |
| if max_sim >= 0.70: | |
| domain = "High Confidence (In-Domain)" | |
| elif max_sim >= 0.40: | |
| domain = "Moderate Confidence" | |
| else: | |
| domain = "Out-of-Domain (Novel Scaffold)" | |
| return {"max_tanimoto": max_sim, "applicability_domain": domain} | |
| except Exception: | |
| return {"max_tanimoto": 0.50, "applicability_domain": "Unknown Domain"} | |