""" Hybrid Enterprise Chemical Search Engine Includes: FAISS Retrieval, Multi-Fingerprint Fusion, Chemical-Aware Reranking, Z-Score Calibration, and MMR Diversity Control. ============================================================================ """ import warnings warnings.filterwarnings("ignore", category=DeprecationWarning) import os import pickle import numpy as np import faiss from rdkit import Chem from rdkit.Chem import AllChem, MACCSkeys, rdMolDescriptors from rdkit import DataStructs class ChemicalSearchEngine: """ Advanced Hybrid Chemical Similarity Search Engine (Drug-Discovery Grade). Pipeline Steps: 1. FAISS Binary Retrieval: Rapid screen over Morgan fingerprints (Top 200). 2. Multi-Fingerprint Fusion: Combines Morgan, MACCS, Atom Pairs, and Torsion. 3. Chemical-Aware Reranking: Applies strict domain constraints (Aromaticity, Rings, Charge, Fragments). 4. Similarity Calibration: Z-score normalization to probabilities. 5. Diversity Control (MMR): Maximizes scaffold diversity. """ def __init__(self, bit_size=2048): self.bit_size = bit_size self.index = None self.index_built = False self.total_compounds = 0 self.metadata = [] self.morgan_fps = [] self.maccs_fps = [] self.atom_pair_fps = [] self.torsion_fps = [] self.faiss_fingerprints = [] def _mol_to_all_fingerprints(self, mol): """Generates all 4 distinct tactical chemical fingerprints.""" if mol is None: return None, None, None, None try: morgan = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=self.bit_size) maccs = MACCSkeys.GenMACCSKeys(mol) atom_pairs = rdMolDescriptors.GetHashedAtomPairFingerprintAsBitVect(mol, nBits=self.bit_size) torsions = rdMolDescriptors.GetHashedTopologicalTorsionFingerprintAsBitVect(mol, nBits=self.bit_size) return morgan, maccs, atom_pairs, torsions except Exception: return None, None, None, None def _extract_chemical_features(self, mol): """ Extracts structural semantics for the Chemical-Aware Reranking Layer. Pre-computed during ingestion for zero overhead during search. """ if mol is None: return 0, 0, 0, 1 try: arom = sum(1 for atom in mol.GetAtoms() if atom.GetIsAromatic()) rings = mol.GetRingInfo().NumRings() charge = sum(atom.GetFormalCharge() for atom in mol.GetAtoms()) frags = len(Chem.GetMolFrags(mol)) return arom, rings, charge, frags except Exception: return 0, 0, 0, 1 def _bitvect_to_numpy(self, bv): """Converts an RDKit bit vector to a NumPy uint8 array.""" arr = np.zeros((len(bv),), dtype=np.uint8) DataStructs.ConvertToNumpyArray(bv, arr) return arr def add_compounds(self, smiles_list, metadata_list=None): """Processes compounds, pre-computes chemical semantics, and builds engine.""" print(f"[START] Processing {len(smiles_list)} compounds for Hybrid Engine...") for i, smiles in enumerate(smiles_list): mol = Chem.MolFromSmiles(smiles) if mol is None: continue morgan, maccs, atom_pairs, torsions = self._mol_to_all_fingerprints(mol) if morgan is None: continue # Extract and cache chemical features (Aromaticity, Rings, Charge, Fragments) arom, rings, charge, frags = self._extract_chemical_features(mol) meta = metadata_list[i] if metadata_list else {"smiles": smiles} meta["smiles"] = smiles # Cache semantic features directly in metadata meta["chem_features"] = {"arom": arom, "rings": rings, "charge": charge, "frags": frags} self.metadata.append(meta) self.morgan_fps.append(morgan) self.maccs_fps.append(maccs) self.atom_pair_fps.append(atom_pairs) self.torsion_fps.append(torsions) self.faiss_fingerprints.append(self._bitvect_to_numpy(morgan)) self.total_compounds = len(self.metadata) print(f"[SUCCESS] Feature extraction complete: {self.total_compounds} compounds") self._build_faiss_index() def _build_faiss_index(self): """Builds FAISS index with proper byte packing.""" if len(self.faiss_fingerprints) == 0: return self.index = faiss.IndexBinaryFlat(self.bit_size) np_fps = np.array(self.faiss_fingerprints, dtype=np.uint8) packed_fps = np.packbits(np_fps, axis=1) self.index.add(packed_fps) self.index_built = True print(f"[SUCCESS] FAISS Binary index successfully built.") def search(self, query_smiles, k=3, lambda_param=0.6): """Executes the complete Drug-Discovery Grade Search Pipeline.""" query_mol = Chem.MolFromSmiles(query_smiles) if query_mol is None or not self.index_built: return [] # 1. Query fingerprints & Semantic features q_morgan, q_maccs, q_atom_pairs, q_torsions = self._mol_to_all_fingerprints(query_mol) q_arom, q_rings, q_charge, q_frags = self._extract_chemical_features(query_mol) if q_morgan is None: return [] # 2. FAISS Fast Retrieval k_search = min(max(k * 20, 200), self.total_compounds) q_packed = np.packbits(self._bitvect_to_numpy(q_morgan).reshape(1, -1), axis=1) _, indices = self.index.search(q_packed, k_search) candidate_indices = [int(idx) for idx in indices[0] if idx >= 0] if not candidate_indices: return [] # 3. Multi-FP & Chemical-Aware Scoring candidate_pool = [] for idx in candidate_indices: # 3a. Multi-Fingerprint Fusion (Base Structural Score) score_morgan = DataStructs.FingerprintSimilarity(q_morgan, self.morgan_fps[idx]) score_maccs = DataStructs.FingerprintSimilarity(q_maccs, self.maccs_fps[idx]) score_atom_pairs = DataStructs.FingerprintSimilarity(q_atom_pairs, self.atom_pair_fps[idx]) score_torsions = DataStructs.FingerprintSimilarity(q_torsions, self.torsion_fps[idx]) base_score = (0.50 * score_morgan + 0.20 * score_maccs + 0.20 * score_atom_pairs + 0.10 * score_torsions) # 3b. Chemical-Aware Penalties & Bonuses (O(1) lookup, ultra-fast) c_feats = self.metadata[idx]["chem_features"] c_arom, c_rings, c_charge, c_frags = c_feats["arom"], c_feats["rings"], c_feats["charge"], c_feats["frags"] # Aromaticity Bonus max_arom = max(q_arom, c_arom) arom_score = 1.0 if max_arom == 0 else 1.0 - (abs(q_arom - c_arom) / max_arom) # Ring System Bonus max_rings = max(q_rings, c_rings) ring_score = 1.0 if max_rings == 0 else 1.0 - (abs(q_rings - c_rings) / max_rings) # Strict Domain Penalties charge_pen = min(abs(c_charge) / 3.0, 1.0) frag_pen = 0.0 if c_frags <= 1 else min((c_frags - 1) * 0.3, 1.0) # 3c. Final Hybrid Score chem_aware_score = ( 0.70 * base_score + 0.15 * arom_score + 0.10 * ring_score - 0.15 * charge_pen - 0.10 * frag_pen ) candidate_pool.append({ "index": idx, "smiles": self.metadata[idx]["smiles"], "metadata": self.metadata[idx], "chem_aware_score": chem_aware_score, "base_structural_score": base_score }) # 4. Calibration Layer (Z-Score) scores = [c["chem_aware_score"] for c in candidate_pool] mean_score = np.mean(scores) std_score = np.std(scores) if np.std(scores) > 0 else 1.0 for c in candidate_pool: z = (c["chem_aware_score"] - mean_score) / std_score c["calibrated_score"] = 1.0 / (1.0 + np.exp(-z)) # 5. Diversity Control (MMR) selected_results = [] remaining_candidates = list(candidate_pool) remaining_candidates.sort(key=lambda x: x["calibrated_score"], reverse=True) selected_results.append(remaining_candidates.pop(0)) while len(selected_results) < k and remaining_candidates: best_mmr = -float('inf') best_idx = -1 for idx, cand in enumerate(remaining_candidates): max_sim = -float('inf') cand_fp = self.morgan_fps[cand["index"]] for sel in selected_results: sim = DataStructs.FingerprintSimilarity(cand_fp, self.morgan_fps[sel["index"]]) if sim > max_sim: max_sim = sim mmr_val = (lambda_param * cand["calibrated_score"]) - ((1.0 - lambda_param) * max_sim) if mmr_val > best_mmr: best_mmr = mmr_val best_idx = idx if best_idx != -1: selected_results.append(remaining_candidates.pop(best_idx)) else: break # Output formatting return [{ "smiles": res["smiles"], "similarity_score": round(float(res["chem_aware_score"]), 4), "calibrated_score": round(float(res["calibrated_score"]), 4), "metadata": res["metadata"], "index": res["index"] } for res in selected_results] def save_index(self, filepath): """Serializes the multi-fingerprint infrastructure and metadata safely.""" if not self.index_built: return try: faiss_file = filepath.replace(".pkl", ".faiss") faiss.write_index_binary(self.index, faiss_file) data = { "metadata": self.metadata, "morgan_fps": self.morgan_fps, "maccs_fps": self.maccs_fps, "atom_pair_fps": self.atom_pair_fps, "torsion_fps": self.torsion_fps, "faiss_fingerprints": self.faiss_fingerprints, "bit_size": self.bit_size, "total_compounds": self.total_compounds } with open(filepath, "wb") as f: pickle.dump(data, f) print(f"[SAVE] Engine saved at: {filepath}") except Exception as e: print(f"[ERROR] Serialization failed: {e}") def load_index(self, filepath): """Loads and provisions the full hybrid matrix and FAISS architecture.""" if not os.path.exists(filepath): return False try: with open(filepath, "rb") as f: data = pickle.load(f) self.metadata = data["metadata"] self.morgan_fps = data["morgan_fps"] self.maccs_fps = data["maccs_fps"] self.atom_pair_fps = data["atom_pair_fps"] self.torsion_fps = data["torsion_fps"] self.faiss_fingerprints = data["faiss_fingerprints"] self.bit_size = data.get("bit_size", 2048) self.total_compounds = data.get("total_compounds", len(self.metadata)) faiss_file = filepath.replace(".pkl", ".faiss") self.index = faiss.read_index_binary(faiss_file) self.index_built = True print(f"[LOAD] Engine loaded completely: {self.total_compounds} compounds.") return True except Exception as e: print(f"[ERROR] Loading error: {e}") return False