| """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.") |
|
|
| |
| |
| |
| 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, |
| ) |
|
|
| |
| |
| |
| def _init_mpnn(self, emb_dim: int) -> None: |
| """Initialize MPNN graph encoder (placeholder for pre-trained weights).""" |
| |
| self.fingerprint_dim = 2048 |
| 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.") |
|
|
| |
| |
| |
| 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) |
| |
| 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) |
| 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 |
|
|
| |
| |
| |
| 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 |
|
|
| |
| |
| |
| @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()) |
|
|
|
|
| |
| |
| |
| 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()} |
|
|
