Upload mulgit/perturb/encoder.py

mulgit/perturb/encoder.py

@@ -1,35 +1,22 @@
 """
 Perturbation Encoder for MuLGIT-Perturb.
 
-Encodes both drug (SMILES) and genetic (gene ID) perturbations into a
-unified embedding space. Uses pretrained transformers as frozen feature
-extractors, with a learned fusion gate that determines the relative
-contribution of drug vs. genetic perturbation.
-
-Architecture:
-  Drug: SMILES → ChemBERTa/MolFormer (frozen) → drug_embed (768-dim)
-  Genetic: Gene symbol → Geneformer (frozen) → gene_embed (256-dim)
-           + Perturbation type (KO/KD/OE) → 3-dim one-hot
-  Fusion: Learned convex combination z_pert = α·z_drug + (1-α)·z_gene
+Encodes drug (SMILES) and genetic perturbations into a unified embedding.
+Drug encoding uses ChemBERTa/MolFormer when available and falls back to
+RDKit Morgan fingerprints. The encoder always returns the configured
+`embed_dim`, so fallback fingerprints cannot break downstream projection
+layers.
 """
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing import Optional, List, Dict, Tuple
+from typing import Optional, List
 import warnings
 
 
 class DrugEncoder(nn.Module):
-    """
-    Encodes drug SMILES strings into molecular embeddings.
-
-    Two modes:
-      1. ChemBERTa/MolFormer (primary): pretrained transformer, frozen.
-         Provides 768-dim continuous embedding capturing molecular properties.
-      2. Morgan fingerprints (fallback): 2048-bit ECFP4 fingerprint via RDKit.
-         Used when SMILES parsing fails or transformer is unavailable.
-    """
+    """SMILES -> fixed-size molecular embedding."""
 
     def __init__(
         self,
@@ -43,39 +30,47 @@ class DrugEncoder(nn.Module):
         self.use_morgan_fallback = use_morgan_fallback
         self.max_smiles_len = max_smiles_len
         self.transformer_model = transformer_model
-        self._transformer = None  # lazy load
+        self._transformer = None
         self._tokenizer = None
-
-    @property
-    def transformer(self):
-        """Lazy-load the transformer to avoid loading if not needed."""
-        if self._transformer is None:
-            try:
-                from transformers import AutoModel, AutoTokenizer
-                self._tokenizer = AutoTokenizer.from_pretrained(self.transformer_model)
-                self._transformer = AutoModel.from_pretrained(self.transformer_model)
-                for p in self._transformer.parameters():
-                    p.requires_grad = False
-                self._transformer.eval()
-            except Exception as e:
-                warnings.warn(f"Failed to load {self.transformer_model}: {e}. "
-                              f"Falling back to Morgan fingerprints.")
-                self._transformer = None
-        return self._transformer
+        self._tried_loading = False
+
+    def _match_dim(self, x: torch.Tensor) -> torch.Tensor:
+        """Pad/truncate any encoder output to self.embed_dim."""
+        if x.shape[-1] == self.embed_dim:
+            return x
+        if x.shape[-1] < self.embed_dim:
+            return F.pad(x, (0, self.embed_dim - x.shape[-1]))
+        return x[:, : self.embed_dim]
+
+    def _load_transformer(self):
+        if self._tried_loading:
+            return
+        self._tried_loading = True
+        try:
+            from transformers import AutoModel, AutoTokenizer
+            self._tokenizer = AutoTokenizer.from_pretrained(self.transformer_model)
+            self._transformer = AutoModel.from_pretrained(self.transformer_model)
+            for p in self._transformer.parameters():
+                p.requires_grad = False
+            self._transformer.eval()
+        except Exception as e:
+            warnings.warn(
+                f"Failed to load {self.transformer_model}: {e}. Falling back to Morgan fingerprints."
+            )
+            self._transformer = None
+            self._tokenizer = None
 
     def encode_morgan(self, smiles_list: List[str], device: torch.device) -> torch.Tensor:
-        """Encode SMILES using Morgan (ECFP4) fingerprints via RDKit."""
         try:
             from rdkit import Chem
             from rdkit.Chem import AllChem
-        except ImportError:
-            raise ImportError("RDKit required for Morgan fingerprint fallback. "
-                              "Install with: pip install rdkit")
+        except ImportError as e:
+            raise ImportError("RDKit required for Morgan fingerprint fallback. Install with: pip install rdkit") from e
 
         fps = []
         for smi in smiles_list:
             try:
-                mol = Chem.MolFromSmiles(smi)
+                mol = Chem.MolFromSmiles(smi or "")
                 if mol is not None:
                     fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048)
                     fp_array = torch.tensor(list(fp), dtype=torch.float32)
@@ -84,18 +79,15 @@ class DrugEncoder(nn.Module):
             except Exception:
                 fp_array = torch.zeros(2048, dtype=torch.float32)
             fps.append(fp_array)
-
-        fps = torch.stack(fps).to(device)  # (B, 2048)
-        return fps
+        return self._match_dim(torch.stack(fps).to(device))
 
     def encode_transformer(self, smiles_list: List[str], device: torch.device) -> torch.Tensor:
-        """Encode SMILES using ChemBERTa/MolFormer."""
-        if self.transformer is None:
+        self._load_transformer()
+        if self._transformer is None or self._tokenizer is None:
             if self.use_morgan_fallback:
                 return self.encode_morgan(smiles_list, device)
             raise RuntimeError("No drug encoder available")
 
-        # Tokenize
         tokens = self._tokenizer(
             smiles_list,
             padding=True,
@@ -103,142 +95,79 @@ class DrugEncoder(nn.Module):
             max_length=self.max_smiles_len,
             return_tensors="pt",
         ).to(device)
-
         with torch.no_grad():
-            outputs = self.transformer(**tokens)
-            # Use CLS token or mean pooling
+            outputs = self._transformer(**tokens)
             if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
-                embeddings = outputs.pooler_output  # (B, 768)
+                embeddings = outputs.pooler_output
             else:
-                embeddings = outputs.last_hidden_state.mean(dim=1)  # (B, 768)
-
-        return embeddings
+                embeddings = outputs.last_hidden_state.mean(dim=1)
+        return self._match_dim(embeddings.float())
 
     def forward(self, smiles_list: List[str], device: torch.device = None) -> torch.Tensor:
-        """
-        Encode a list of SMILES strings.
-
-        Args:
-          smiles_list: list of SMILES strings
-          device: torch device
-
-        Returns:
-          drug_embeddings: (B, embed_dim)
-        """
         if device is None:
-            device = next(self.parameters()).device if list(self.parameters()) else torch.device("cpu")
-
-        # Try transformer first
+            # DrugEncoder has no parameters; default to CPU unless caller passes device.
+            device = torch.device("cpu")
         try:
             return self.encode_transformer(smiles_list, device)
         except Exception as e:
             if self.use_morgan_fallback:
-                warnings.warn(f"ChemBERTa encoding failed ({e}), using Morgan fingerprints")
+                warnings.warn(f"Transformer encoding failed ({e}); using Morgan fingerprints")
                 return self.encode_morgan(smiles_list, device)
             raise
 
 
 class GeneticPerturbationEncoder(nn.Module):
-    """
-    Encodes genetic perturbations (gene KO/KD/OE) into embeddings.
-
-    Gene identity: gene name → pretrained gene embedding (Geneformer, gene2vec, or learned)
-    Perturbation type: one-hot encoding (CRISPR_KO=0, CRISPRi=1, shRNA=2, OE=3, CRISPRa=4)
-    """
+    """Gene perturbation encoder: (gene_id, perturbation_type) -> embedding."""
 
     def __init__(
         self,
         gene_embed_dim: int = 256,
-        pert_type_dim: int = 5,  # KO, CRISPRi, shRNA, OE, CRISPRa
+        pert_type_dim: int = 5,
         output_dim: int = 768,
         n_genes: int = 20000,
     ):
         super().__init__()
-        # Gene embedding: learned lookup table (can be initialized from Geneformer)
         self.gene_embedding = nn.Embedding(n_genes, gene_embed_dim, padding_idx=0)
-
-        # Perturbation type embedding
         self.pert_type_embedding = nn.Embedding(pert_type_dim, 32)
-
-        # Project combined gene + perturbation type to output dimension
-        combined_dim = gene_embed_dim + 32
         self.projector = nn.Sequential(
-            nn.Linear(combined_dim, 512),
+            nn.Linear(gene_embed_dim + 32, 512),
             nn.SELU(),
             nn.Linear(512, output_dim),
         )
-
-        # Map perturbation type string to index
         self.pert_type_map = {
-            "CRISPR_KO": 0, "CRISPR_KO": 0,
-            "CRISPRi": 1, "KD": 1, "knockdown": 1,
-            "shRNA": 2,
-            "OE": 3, "overexpression": 3, "CRISPRa": 3,
-            "CRISPRa": 4, "activation": 4,
+            "crispr_ko": 0,
+            "ko": 0,
+            "crispri": 1,
+            "kd": 1,
+            "knockdown": 1,
+            "shrna": 2,
+            "oe": 3,
+            "overexpression": 3,
+            "crispra": 4,
+            "activation": 4,
             "unknown": 0,
         }
 
     def load_geneformer_embeddings(self, geneformer_model: str = "ctheodoris/Geneformer"):
-        """
-        Initialize gene embeddings from Geneformer's pretrained gene embeddings.
-        Falls back to random initialization if Geneformer is unavailable.
-        """
-        try:
-            from transformers import AutoModel
-            model = AutoModel.from_pretrained(geneformer_model)
-            # Geneformer stores gene embeddings in the embedding layer
-            # For now, keep learned initialization
-            warnings.warn(
-                "Geneformer embedding extraction not yet automated. "
-                "Using learned embeddings. For best results, pre-extract "
-                "Geneformer gene embeddings and load them manually."
-            )
-        except Exception:
-            pass
-
-    def forward(
-        self,
-        gene_ids: torch.LongTensor,
-        pert_types: List[str],
-    ) -> torch.Tensor:
-        """
-        Args:
-          gene_ids: (B,) integer gene indices
-          pert_types: list of perturbation type strings
+        warnings.warn(
+            "Geneformer embedding extraction is not automated in this implementation; "
+            "using learned gene embeddings. Pre-extracted Geneformer embeddings can be loaded manually."
+        )
 
-        Returns:
-          gene_pert_embedding: (B, output_dim)
-        """
+    def forward(self, gene_ids: torch.LongTensor, pert_types: List[str]) -> torch.Tensor:
         device = gene_ids.device
-
-        # Gene embedding
-        gene_emb = self.gene_embedding(gene_ids)  # (B, gene_embed_dim)
-
-        # Perturbation type
+        gene_emb = self.gene_embedding(gene_ids)
         pert_type_idxs = torch.tensor(
-            [self.pert_type_map.get(pt.lower(), 0) for pt in pert_types],
-            dtype=torch.long, device=device,
+            [self.pert_type_map.get(str(pt).lower(), 0) for pt in pert_types],
+            dtype=torch.long,
+            device=device,
         )
-        pert_type_emb = self.pert_type_embedding(pert_type_idxs)  # (B, 32)
-
-        # Combine and project
-        combined = torch.cat([gene_emb, pert_type_emb], dim=-1)  # (B, gene_embed_dim + 32)
-        return self.projector(combined)  # (B, output_dim)
+        pert_type_emb = self.pert_type_embedding(pert_type_idxs)
+        return self.projector(torch.cat([gene_emb, pert_type_emb], dim=-1))
 
 
 class PerturbationEncoder(nn.Module):
-    """
-    Unified perturbation encoder for both drug and genetic perturbations.
-
-    Architecture:
-      Drug: SMILES → DrugEncoder (ChemBERTa or Morgan) → drug_embed (768-dim)
-      Genetic: (gene_id, pert_type) → GeneticPerturbationEncoder → gene_embed (768-dim)
-      Fusion: z_pert = α·z_drug + (1-α)·z_gene
-        where α = sigmoid(learned_alpha) ∈ [0, 1]
-
-    When only one type of perturbation is provided, the other branch
-    outputs zeros and the fusion weight adapts accordingly.
-    """
+    """Unified perturbation encoder for drugs and genetic perturbations."""
 
     def __init__(
         self,
@@ -251,41 +180,15 @@ class PerturbationEncoder(nn.Module):
     ):
         super().__init__()
         self.output_dim = output_dim
-
-        # Drug encoder
         self.drug_encoder = DrugEncoder(
             transformer_model=drug_encoder_model,
             embed_dim=drug_embed_dim,
             use_morgan_fallback=use_morgan_fallback,
         )
-
-        # Project drug embedding to output dim if needed
-        if drug_embed_dim != output_dim:
-            self.drug_proj = nn.Sequential(
-                nn.Linear(drug_embed_dim, output_dim),
-                nn.SELU(),
-            )
-        else:
-            self.drug_proj = nn.Identity()
-
-        # Genetic perturbation encoder
-        self.gene_encoder = GeneticPerturbationEncoder(
-            gene_embed_dim=gene_embed_dim,
-            output_dim=output_dim,
-            n_genes=n_genes,
-        )
-
-        # Fusion gate: learned α ∈ [0, 1] determines drug vs. gene contribution
-        self.alpha_logit = nn.Parameter(torch.tensor(0.0))  # sigmoid(0) = 0.5
-
-        # Output projection (after fusion)
-        self.output_proj = nn.Sequential(
-            nn.Linear(output_dim, output_dim),
-            nn.SELU(),
-            nn.Linear(output_dim, output_dim),
-        )
-
-        # Dropout for training
+        self.drug_proj = nn.Sequential(nn.Linear(drug_embed_dim, output_dim), nn.SELU()) if drug_embed_dim != output_dim else nn.Identity()
+        self.gene_encoder = GeneticPerturbationEncoder(gene_embed_dim=gene_embed_dim, output_dim=output_dim, n_genes=n_genes)
+        self.alpha_logit = nn.Parameter(torch.tensor(0.0))
+        self.output_proj = nn.Sequential(nn.Linear(output_dim, output_dim), nn.SELU(), nn.Linear(output_dim, output_dim))
         self.dropout = nn.AlphaDropout(0.1)
 
     def forward(
@@ -294,58 +197,31 @@ class PerturbationEncoder(nn.Module):
         gene_ids: Optional[torch.LongTensor] = None,
         pert_types: Optional[List[str]] = None,
     ) -> torch.Tensor:
-        """
-        Encode perturbation into unified embedding.
-
-        At least one of (smiles_list) or (gene_ids + pert_types) must be provided.
-
-        Args:
-          smiles_list: list of SMILES strings, or None
-          gene_ids: (B,) integer gene indices, or None
-          pert_types: list of perturbation type strings, or None
-
-        Returns:
-          z_pert: (B, output_dim) unified perturbation embedding
-        """
-        batch_size = None
-        device = None
-
-        # Determine batch size and device
         if gene_ids is not None:
             batch_size = gene_ids.shape[0]
             device = gene_ids.device
         elif smiles_list is not None:
             batch_size = len(smiles_list)
-            device = next(self.parameters()).device if list(self.parameters()) else torch.device("cpu")
-
-        if batch_size is None:
+            device = self.alpha_logit.device
+        else:
             raise ValueError("Either smiles_list or gene_ids must be provided")
 
-        # Encode drug
         if smiles_list is not None and len(smiles_list) > 0:
-            z_drug = self.drug_encoder(smiles_list, device)  # (B, drug_embed_dim)
-            z_drug = self.drug_proj(z_drug)  # (B, output_dim)
+            z_drug = self.drug_proj(self.drug_encoder(smiles_list, device))
         else:
             z_drug = torch.zeros(batch_size, self.output_dim, device=device)
 
-        # Encode genetic perturbation
         if gene_ids is not None and pert_types is not None:
-            z_gene = self.gene_encoder(gene_ids, pert_types)  # (B, output_dim)
+            z_gene = self.gene_encoder(gene_ids, pert_types)
         else:
             z_gene = torch.zeros(batch_size, self.output_dim, device=device)
 
-        # Learned fusion
-        alpha = torch.sigmoid(self.alpha_logit)  # ∈ [0, 1]
-        z_pert = alpha * z_drug + (1 - alpha) * z_gene
-
-        # Output projection
-        z_pert = self.output_proj(self.dropout(F.selu(z_pert)))
-
-        return z_pert
+        alpha = torch.sigmoid(self.alpha_logit)
+        z_pert = alpha * z_drug + (1.0 - alpha) * z_gene
+        return self.output_proj(self.dropout(F.selu(z_pert)))
 
 
 def create_perturbation_encoder(config) -> PerturbationEncoder:
-    """Factory function to create PerturbationEncoder from config."""
     return PerturbationEncoder(
         drug_encoder_model=config.drug_encoder_model,
         drug_embed_dim=config.drug_embed_dim,