"""Drug molecular encoder for structure-based drug representation. Provides two encoding strategies: 1. Morgan fingerprints (ECFP4) → MLP projection [MVI, no training needed] 2. MPNN graph encoder (pre-trained, frozen) [Phase 2] The drug embedding bridges the gap between molecular structure and gene-level intervention effects, enabling the model to differentiate drugs that target the same genes but have different downstream transcriptional signatures. Usage ----- encoder = DrugEncoder(encoding="morgan", emb_dim=128) smiles_list = ["CC(=O)Oc1ccccc1C(=O)O", "CN1C=NC2=C1C(=O)N(C(=O)N2C)C"] drug_emb = encoder(smiles_list) # [n_drugs, emb_dim] References ---------- - Morgan fingerprints: Rogers & Hahn (2010) "Extended-Connectivity Fingerprints" - MPNN: Gilmer et al. (2018) "Neural Message Passing for Quantum Chemistry" """ from __future__ import annotations import logging from typing import Optional import numpy as np import torch import torch.nn as nn import torch.nn.functional as F logger = logging.getLogger(__name__) class DrugEncoder(nn.Module): """Encode drug SMILES strings into continuous embeddings. Parameters ---------- encoding : str "morgan" for ECFP4 fingerprints (MVI, no training needed). "mpnn" for graph neural network (requires RDKit + pre-trained weights). emb_dim : int Output embedding dimension. Default 128. morgan_radius : int Radius for Morgan fingerprint computation. Default 2. morgan_nbits : int Number of bits in Morgan fingerprint. Default 2048. freeze : bool If True, freeze all encoder parameters (recommended for pre-trained MPNN). """ def __init__( self, encoding: str = "morgan", emb_dim: int = 128, morgan_radius: int = 2, morgan_nbits: int = 2048, freeze: bool = True, ) -> None: super().__init__() self.encoding = encoding self.emb_dim = emb_dim self.freeze = freeze if encoding == "morgan": self._init_morgan(morgan_nbits, morgan_radius, emb_dim) elif encoding == "mpnn": self._init_mpnn(emb_dim) else: raise ValueError(f"Unknown encoding: {encoding}. Choose 'morgan' or 'mpnn'.") if freeze: for param in self.parameters(): param.requires_grad = False logger.info("DrugEncoder: all parameters frozen.") # ------------------------------------------------------------------ # Morgan fingerprint backend # ------------------------------------------------------------------ def _init_morgan(self, nbits: int, radius: int, emb_dim: int) -> None: """Initialize Morgan fingerprint encoder with MLP projection.""" self.fingerprint_dim = nbits self.projection = nn.Sequential( nn.Linear(nbits, max(nbits // 2, emb_dim)), nn.LayerNorm(max(nbits // 2, emb_dim)), nn.GELU(), nn.Linear(max(nbits // 2, emb_dim), emb_dim), ) logger.info( "DrugEncoder: Morgan fingerprint (radius=%d, nbits=%d) → MLP(%d→%d→%d)", radius, nbits, nbits, max(nbits // 2, emb_dim), emb_dim, ) # ------------------------------------------------------------------ # MPNN backend (placeholder for Phase 2) # ------------------------------------------------------------------ def _init_mpnn(self, emb_dim: int) -> None: """Initialize MPNN graph encoder (placeholder for pre-trained weights).""" # Placeholder — will be replaced with actual MPNN in Phase 2 self.fingerprint_dim = 2048 # fallback self.projection = nn.Sequential( nn.Linear(2048, 512), nn.GELU(), nn.Linear(512, emb_dim), ) logger.warning("DrugEncoder: MPNN not yet implemented, using fallback MLP on 2048-dim input.") # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ def forward(self, smiles_list: list[str]) -> torch.Tensor: """Encode a list of SMILES strings to drug embeddings. Parameters ---------- smiles_list : list of str SMILES strings for each drug molecule. Returns ------- emb : [n_drugs, emb_dim] float tensor Drug embeddings on the same device as the module's parameters. """ if len(smiles_list) == 0: device = next(self.parameters()).device return torch.empty(0, self.emb_dim, device=device) if self.encoding == "morgan": fps = self._compute_morgan_fps(smiles_list) fps_tensor = torch.as_tensor(fps, dtype=torch.float32) # Move to same device as model parameters device = next(self.parameters()).device fps_tensor = fps_tensor.to(device) emb = self.projection(fps_tensor) elif self.encoding == "mpnn": fps = self._compute_morgan_fps(smiles_list) # fallback fps_tensor = torch.as_tensor(fps, dtype=torch.float32) device = next(self.parameters()).device fps_tensor = fps_tensor.to(device) emb = self.projection(fps_tensor) return emb # ------------------------------------------------------------------ # Morgan fingerprint computation # ------------------------------------------------------------------ def _compute_morgan_fps(self, smiles_list: list[str]) -> np.ndarray: """Compute ECFP4 Morgan fingerprints for a list of SMILES. Parameters ---------- smiles_list : list of str Returns ------- fps : [n_drugs, morgan_nbits] binary array """ try: from rdkit import Chem from rdkit.Chem import AllChem except ImportError: raise ImportError( "RDKit is required for Morgan fingerprint computation. " "Install with: conda install -c conda-forge rdkit" ) nbits = self.fingerprint_dim fps = np.zeros((len(smiles_list), nbits), dtype=np.float32) for i, smiles in enumerate(smiles_list): mol = Chem.MolFromSmiles(smiles) if mol is None: logger.warning("Failed to parse SMILES: %s", smiles) continue fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=nbits) fps[i] = np.array(list(fp), dtype=np.float32) return fps # ------------------------------------------------------------------ # Utility: compute similarity between two drug embeddings # ------------------------------------------------------------------ @torch.no_grad() def similarity(self, emb_a: torch.Tensor, emb_b: torch.Tensor) -> torch.Tensor: """Compute cosine similarity between drug embeddings. Parameters ---------- emb_a : [n, emb_dim] emb_b : [m, emb_dim] Returns ------- sim : [n, m] cosine similarity matrix """ a_norm = F.normalize(emb_a, p=2, dim=-1) b_norm = F.normalize(emb_b, p=2, dim=-1) return torch.mm(a_norm, b_norm.t()) # --------------------------------------------------------------------------- # Cache for computed embeddings (avoid recomputation) # --------------------------------------------------------------------------- class DrugEmbeddingCache: """Cache drug embeddings by SMILES to avoid recomputation. Parameters ---------- encoder : DrugEncoder cache_path : str, optional Path to pickle file for persistent caching across sessions. """ def __init__(self, encoder: DrugEncoder, cache_path: Optional[str] = None): self.encoder = encoder self.cache_path = cache_path self._cache: dict[str, torch.Tensor] = {} def get(self, smiles: str) -> torch.Tensor: """Get embedding for a single SMILES string (cached).""" if smiles not in self._cache: emb = self.encoder([smiles]).squeeze(0).detach().cpu() self._cache[smiles] = emb return self._cache[smiles] def get_batch(self, smiles_list: list[str]) -> torch.Tensor: """Get embeddings for a batch of SMILES strings.""" uncached = [s for s in smiles_list if s not in self._cache] if uncached: embs = self.encoder(uncached).detach().cpu() for s, e in zip(uncached, embs): self._cache[s] = e return torch.stack([self._cache[s] for s in smiles_list]) def save(self) -> None: """Save cache to disk.""" if self.cache_path: import pickle with open(self.cache_path, "wb") as f: pickle.dump({k: v.numpy() for k, v in self._cache.items()}, f) def load(self) -> None: """Load cache from disk.""" if self.cache_path: import os import pickle if os.path.exists(self.cache_path): with open(self.cache_path, "rb") as f: raw = pickle.load(f) device = next(self.encoder.parameters()).device self._cache = {k: torch.as_tensor(v, device=device) for k, v in raw.items()}