"""SciPlex3 drug perturbation dataset for population-level training. Loads the Srivatsan & Trapnell 2020 SciPlex3 dataset, which contains single-cell RNA-seq profiles of cells treated with small molecule drugs at multiple doses and time points. Data structure -------------- SciPlex3 contains: - 189 drug perturbations (including vehicle control) - 5 dose levels: 0, 10, 100, 1000, 10000 nM - 2 time points: 24h, 72h - 3 cell lines: MCF7 (breast), A549 (lung), K562 (CML) - ChEMBL IDs for 155 drugs - Protein target annotations (87 target categories) Training pairs are constructed as: source = vehicle control cells at the same dose/time/cell_line target = drug-treated cells This ensures the source/target difference is purely the drug effect, not batch or technical variation. Parameters ---------- h5ad_path : str Path to SciPlex3 h5ad file. n_hvg : int Number of highly variable genes to select (default 2000). min_cells_per_cond : int Minimum cells per (drug, dose, time) condition (default 30). max_source_cells : int Cells sampled per source (vehicle) population (default 64). max_target_cells : int Cells sampled per target (drug-treated) population (default 64). cell_lines : list of str, optional Which cell lines to include. Default: all (MCF7, A549, K562). doses : list of float, optional Which doses to include. Default: all non-zero doses. times : list of float, optional Which time points to include. Default: all. include_control : bool If True, include vehicle-only conditions as training samples (no drug). Default: False. target_sum : float Normalisation target (default 1e4). seed : int Random seed. Usage ----- ds = Sciplex3Dataset("SrivatsanTrapnell2020_sciplex3.h5ad") sample = ds[0] # source_cells: [Ns, G] vehicle control # target_cells: [Nt, G] drug-treated # drug_emb: [D_d] drug molecular embedding # dose: scalar drug concentration # perturbation: [G] drug target genes (multi-hot) # chembl_id: str ChEMBL identifier # target: str protein target name # pathway: str pathway annotation """ from __future__ import annotations import logging import os import csv from typing import Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd import torch from torch.utils.data import Dataset logger = logging.getLogger(__name__) class Sciplex3Dataset(Dataset): """SciPlex3 drug perturbation dataset for population-level training.""" def __init__( self, h5ad_path: str, n_hvg: int = 2000, min_cells_per_cond: int = 30, max_source_cells: int = 64, max_target_cells: int = 64, cell_lines: Optional[List[str]] = None, doses: Optional[List[float]] = None, times: Optional[List[float]] = None, include_control: bool = False, target_sum: float = 1e4, seed: int = 42, # Drug encoder settings drug_encoding: str = "morgan", drug_emb_dim: int = 128, # Pre-computed drug embeddings (optional, avoids importing DrugEncoder) precomputed_drug_embeddings: Optional[Dict[str, torch.Tensor]] = None, # Drug-target mapping: external gene targets per drug drug_target_mapping: Optional[Dict[str, List[str]]] = None, # Preprocessed data path (faster than loading from h5ad) preprocessed_path: Optional[str] = None, # Align gene dimension to this value (pad with zeros if needed) target_num_genes: Optional[int] = None, # Path to ChEMBL SMILES CSV (chembl_id,smiles) drug_smiles_csv: str = "", ) -> None: super().__init__() self.max_source_cells = max_source_cells self.max_target_cells = max_target_cells self.seed = seed self.rng = np.random.default_rng(seed) self.target_num_genes = target_num_genes self._smiles_cache: Dict[str, str] = {} self._load_smiles_cache(drug_smiles_csv) # Load from preprocessed file (fast path) or h5ad (slow path) if preprocessed_path and os.path.exists(preprocessed_path): self._load_preprocessed(preprocessed_path, drug_emb_dim, precomputed_drug_embeddings) # Try to load gene symbol mapping (ENSEMBL → symbol) map_path = preprocessed_path.replace('.pt', '_gene_map.json') if os.path.exists(map_path): self._load_gene_mapping(map_path) return print(f"[Sciplex3Dataset] Loading {os.path.basename(h5ad_path)} ...") try: import anndata as ad except ImportError: raise ImportError("pip install anndata") adata = ad.read_h5ad(h5ad_path) print(f" raw shape: {adata.shape}") # ------------------------------------------------------------------ # 1. Filter by metadata BEFORE extracting expression matrix # SciPlex3 has 799K cells × 110K genes — must filter before dense load # ------------------------------------------------------------------ obs = adata.obs.copy() # Filter by cell line if cell_lines: obs = obs[obs['cell_line'].isin(cell_lines)].copy() print(f" Filtered to cell lines: {cell_lines} → {len(obs)} cells") # Filter by dose (exclude dose=0 = vehicle) if doses: obs = obs[obs['dose_value'].isin(doses)].copy() else: obs = obs[obs['dose_value'] > 0].copy() if times: obs = obs[obs['time'].isin(times)].copy() # Subset adata to filtered cells keep_idx = obs.index.values print(f" After filtering: {len(keep_idx)} cells") print(f" Cell lines: {obs['cell_line'].value_counts().to_dict()}") print(f" Doses: {sorted(obs['dose_value'].unique())}") print(f" Drugs: {obs['perturbation'].nunique()}") adata = adata[keep_idx].copy() obs = adata.obs.copy() obs['_cell_idx'] = np.arange(len(obs)) # NOTE on drug-name whitespace: 11 SciPlex3 perturbations carry a # trailing space in the h5ad ('Busulfan ', 'Mesna ', ...). We deliberately # KEEP the raw name in cond['drug_name'] because the split pair_ids were # generated from the raw (spaced) names, so pair_id<->split matching needs # the raw form. Stripping is applied ONLY at the drug_order / protein-index # lookup boundaries (see _get_batch_protein_labels and eval query builder). # ------------------------------------------------------------------ # 2. Normalize + HVG selection (on filtered data) # Work with sparse matrices to avoid 300+ GiB dense allocation. # ------------------------------------------------------------------ import scipy.sparse as sp # Capture gene names before deleting adata var_names_all = list(adata.var.index) if hasattr(adata, 'var') else [f"Gene_{i}" for i in range(adata.shape[1])] X_sp = adata.X # sparse: [n_cells, n_genes] if not sp.issparse(X_sp): X_sp = sp.csr_matrix(X_sp) elif not sp.isspmatrix_csr(X_sp): X_sp = X_sp.tocsr() # Free the anndata object (keep only the sparse matrix) del adata # --- Sparse library-size normalization (no full densification) --- row_sums = np.asarray(X_sp.sum(axis=1)).ravel().clip(min=1.0) # [n_cells] # Scale factor per cell; apply via row-wise multiplication in sparse form scale = (target_sum / row_sums).astype(np.float32) # [n_cells] X_sp_norm = X_sp.multiply(scale[:, np.newaxis]) # still sparse [n_cells, n_genes] del X_sp # free raw sparse if n_hvg > 0: # Compute per-gene variance in sparse normalized space without densifying # E[x^2] - E[x]^2; works for nonnegative (pre-log) counts n_cells = X_sp_norm.shape[0] mean_g = np.asarray(X_sp_norm.mean(axis=0)).ravel() # [n_genes] # E[x^2]: sum of squares per column X_sq = X_sp_norm.copy() X_sq.data **= 2 mean_sq_g = np.asarray(X_sq.mean(axis=0)).ravel() # [n_genes] del X_sq var_g = mean_sq_g - mean_g ** 2 # [n_genes] n_select = min(n_hvg, X_sp_norm.shape[1]) hvg_idx = np.argsort(var_g)[::-1][:n_select].copy() gene_names = np.array(var_names_all)[hvg_idx] # Densify only the HVG columns: [n_cells, n_hvg] # Slice column-wise in CSC for efficiency X_sp_csc = X_sp_norm.tocsc() del X_sp_norm X_hvg_pre = np.asarray(X_sp_csc[:, hvg_idx].todense(), dtype=np.float32) del X_sp_csc else: # Densify full matrix (only safe for small datasets) gene_names = np.array(var_names_all) X_hvg_pre = np.asarray(X_sp_norm.todense(), dtype=np.float32) del X_sp_norm # Apply log1p after slimming (much cheaper on [n_cells, n_hvg]) X_hvg = np.log1p(X_hvg_pre) del X_hvg_pre self.num_genes = len(gene_names) self.gene_names = gene_names self._gene_to_idx: Dict[str, int] = {g: i for i, g in enumerate(gene_names)} print(f" Using {self.num_genes} genes") # Align gene dimension to target_num_genes if specified if target_num_genes is not None and self.num_genes != target_num_genes: self._align_genes(target_num_genes) # ------------------------------------------------------------------ # 3. Separate vehicle and drug-treated cells # ------------------------------------------------------------------ vehicle_mask = obs['perturbation'].str.lower().str.contains('vehicle', na=False) | \ (obs['dose_value'] == 0) # If include_control=True, also use dose=0 as target (reconstruction task) vehicle_obs = obs[vehicle_mask].copy() drug_obs = obs[~vehicle_mask].copy() # If no vehicle cells found, use all cells as drug-treated and # create pseudo-vehicle from the same condition's minimum dose if len(vehicle_obs) == 0: logger.warning("No vehicle cells found. Using minimum dose as pseudo-vehicle.") min_dose = obs['dose_value'].min() vehicle_obs = obs[obs['dose_value'] == min_dose].copy() drug_obs = obs[obs['dose_value'] > min_dose].copy() print(f" Vehicle cells: {len(vehicle_obs)}, Drug cells: {len(drug_obs)}") # ------------------------------------------------------------------ # 4. Build drug-target mapping # ------------------------------------------------------------------ # From SciPlex3 target annotations (protein-level → gene-level) # Plus optional external mapping (DGIdb) self._drug_target_map = self._build_drug_target_map( drug_obs=drug_obs, external_mapping=drug_target_mapping, gene_names=list(gene_names), ) # ------------------------------------------------------------------ # 5. Build training conditions # ------------------------------------------------------------------ # Each condition = (drug, dose, time, cell_line) tuple conditions: List[Dict] = [] # Group drug-treated cells by (perturbation, dose, time, cell_line) drug_groups = drug_obs.groupby(['perturbation', 'dose_value', 'time', 'cell_line']) for (drug_name, dose, time, cell_line), group in drug_groups: if len(group) < min_cells_per_cond: continue # Find matching vehicle cells: same cell_line, same dose=0, same plate if possible vehicle_same_cl = vehicle_obs[vehicle_obs['cell_line'] == cell_line] if len(vehicle_same_cl) < min_cells_per_cond: # Fall back to any vehicle vehicle_same_cl = vehicle_obs # Build drug target vector in HVG space pert_vec = self._build_perturbation_vector(drug_name, drug_target_mapping) # Get ChEMBL ID chembl_id = group['chembl-ID'].iloc[0] if 'chembl-ID' in group.columns else "" if pd.isna(chembl_id): chembl_id = "" # Get target annotation target_name = group['target'].iloc[0] if 'target' in group.columns else "" pathway_l1 = group['pathway_level_1'].iloc[0] if 'pathway_level_1' in group.columns else "" pathway_l2 = group['pathway_level_2'].iloc[0] if 'pathway_level_2' in group.columns else "" conditions.append({ "pert_name": f"{drug_name}_{dose}nM", "drug_name": drug_name, "cell_line": cell_line, "dose": float(dose), "time": float(time) if not pd.isna(time) else 24.0, "drug_cell_idx": group['_cell_idx'].values, "vehicle_cell_idx": vehicle_same_cl['_cell_idx'].values, "pert_vec": pert_vec, "chembl_id": chembl_id, "target": target_name, "pathway_l1": pathway_l1, "pathway_l2": pathway_l2, }) self._conditions = conditions self._X = X_hvg # [n_cells, n_hvg] print(f" Conditions (≥{min_cells_per_cond} cells): {len(conditions)}") # Add drug_name → SMILES entries to cache (via ChEMBL ID lookup) for cond in self._conditions: name = cond['drug_name'] chembl = cond.get('chembl_id', '') if chembl and chembl in self._smiles_cache: self._smiles_cache[name] = self._smiles_cache[chembl] # Align gene dimension after conditions are built (updates pert_vecs) if self.target_num_genes is not None and self.num_genes != self.target_num_genes: self._align_genes(self.target_num_genes) # Drug name list for encoding drug_names = list(set(c['drug_name'] for c in conditions)) print(f" Unique drugs: {len(drug_names)}") # ------------------------------------------------------------------ # 6. Pre-compute drug embeddings # ------------------------------------------------------------------ self.drug_emb_dim = drug_emb_dim self._drug_name_to_idx = {name: i for i, name in enumerate(sorted(drug_names))} # SMILES lookup table (populated by _load_smiles_cache in __init__) # Add drug_name → SMILES entries via ChEMBL ID lookup for cond in conditions: name = cond['drug_name'] chembl = cond.get('chembl_id', '') if chembl and chembl in self._smiles_cache: self._smiles_cache[name] = self._smiles_cache[chembl] # Pre-computed drug embeddings self._drug_embeddings: Optional[torch.Tensor] = None if precomputed_drug_embeddings is not None: self._load_precomputed_embeddings(precomputed_drug_embeddings) else: # Compute embeddings using a lightweight Morgan fingerprint encoder self._compute_drug_embeddings() # ------------------------------------------------------------------ def __len__(self) -> int: return len(self._conditions) def __getitem__(self, idx: int) -> Dict: cond = self._conditions[idx] # Sample source (vehicle) cells ns = min(self.max_source_cells, len(cond['vehicle_cell_idx'])) src_idx = self.rng.choice(cond['vehicle_cell_idx'], size=ns, replace=False) source_cells = torch.from_numpy(self._X[src_idx]).float() # [ns, G] # Sample target (drug-treated) cells nt = min(self.max_target_cells, len(cond['drug_cell_idx'])) tgt_idx = self.rng.choice(cond['drug_cell_idx'], size=nt, replace=False) target_cells = torch.from_numpy(self._X[tgt_idx]).float() # [nt, G] perturbation = torch.from_numpy(cond['pert_vec']).float() # [G] drug_smiles = self._get_drug_smiles(cond['drug_name']) # str return { "source_cells": source_cells, "target_cells": target_cells, "perturbation": perturbation, "drug_smiles": drug_smiles, "dose": torch.tensor(cond['dose'], dtype=torch.float32), "cell_line": cond['cell_line'], "chembl_id": cond['chembl_id'], "target": cond['target'], "pathway_l1": cond['pathway_l1'], "pathway_l2": cond['pathway_l2'], "metadata": { "pert_name": cond['pert_name'], "drug_name": cond['drug_name'], "num_source_cells": ns, "num_target_cells": nt, }, } # ------------------------------------------------------------------ # ------------------------------------------------------------------ def _get_drug_smiles(self, drug_name: str) -> str: """Get SMILES string for a drug by name.""" # Try direct lookup smiles = self._smiles_cache.get(drug_name, "") if smiles: return smiles # Try ChEMBL ID lookup idx = self._drug_name_to_idx.get(drug_name, 0) cond = self._conditions[0] if self._conditions else {} # Try to find ChEMBL ID from conditions for c in self._conditions: if c['drug_name'] == drug_name: chembl_id = c.get('chembl_id', '') if chembl_id: smiles = self._smiles_cache.get(chembl_id, "") if smiles: return smiles break return "" def _get_drug_embedding(self, drug_name: str) -> torch.Tensor: """Get pre-computed drug embedding by name (legacy, for backward compat).""" if self._drug_embeddings is None: device = torch.device("cpu") return torch.zeros(self.drug_emb_dim, device=device) idx = self._drug_name_to_idx.get(drug_name, 0) return self._drug_embeddings[idx] def _load_preprocessed( self, path: str, drug_emb_dim: int = 128, precomputed_drug_embeddings: Optional[Dict[str, torch.Tensor]] = None, ) -> None: """Load from preprocessed .pt file (much faster than h5ad).""" print(f" Loading preprocessed data from {path} ...") data = torch.load(path, weights_only=False) self._X = data['X'].numpy() if isinstance(data['X'], torch.Tensor) else data['X'] self.gene_names = data['gene_names'] self.num_genes = len(self.gene_names) self._gene_to_idx = {g: i for i, g in enumerate(self.gene_names)} obs = data['obs'] print(f" Loaded: {self._X.shape[0]} cells, {self.num_genes} genes") # Try to load sidecar gene mapping (ENSEMBL → symbol) map_path = path.replace('.pt', '_gene_map.json') if os.path.exists(map_path): self._load_gene_mapping(map_path) # ------------------------------------------------------------------ # Separate vehicle and drug-treated cells # ------------------------------------------------------------------ vehicle_mask = obs['perturbation'].str.lower().str.contains('vehicle', na=False) | \ (obs['dose_value'] == 0) vehicle_obs = obs[vehicle_mask].copy() drug_obs = obs[~vehicle_mask].copy() if len(vehicle_obs) == 0: min_dose = obs['dose_value'].min() vehicle_obs = obs[obs['dose_value'] == min_dose].copy() drug_obs = obs[obs['dose_value'] > min_dose].copy() print(f" Vehicle cells: {len(vehicle_obs)}, Drug cells: {len(drug_obs)}") # ------------------------------------------------------------------ # Build conditions # ------------------------------------------------------------------ drug_target_mapping = getattr(self, '_drug_target_mapping', None) # ------------------------------------------------------------------ # Build conditions # ------------------------------------------------------------------ # Pre-build drug_target_map from target categories in the data self._drug_target_map = {} for drug_name in drug_obs['perturbation'].unique(): target_cat = drug_obs[drug_obs['perturbation'] == drug_name]['target'].iloc[0] gene_targets = self._target_category_to_genes( target_cat, list(self.gene_names), self._gene_to_idx, getattr(self, '_symbol_to_ensembl', None) ) self._drug_target_map[drug_name] = gene_targets conditions: List[Dict] = [] drug_groups = drug_obs.groupby(['perturbation', 'dose_value', 'time', 'cell_line']) for (drug_name, dose, time, cell_line), group in drug_groups: if len(group) < self.max_source_cells: continue vehicle_same_cl = vehicle_obs[vehicle_obs['cell_line'] == cell_line] if len(vehicle_same_cl) < self.max_source_cells: vehicle_same_cl = vehicle_obs pert_vec = self._build_perturbation_vector(drug_name, None) chembl_id = group['chembl-ID'].iloc[0] if 'chembl-ID' in group.columns else "" if pd.isna(chembl_id): chembl_id = "" target_name = group['target'].iloc[0] if 'target' in group.columns else "" pathway_l1 = group['pathway_level_1'].iloc[0] if 'pathway_level_1' in group.columns else "" pathway_l2 = group['pathway_level_2'].iloc[0] if 'pathway_level_2' in group.columns else "" conditions.append({ "pert_name": f"{drug_name}_{dose}nM", "drug_name": drug_name, "cell_line": cell_line, "dose": float(dose), "time": float(time) if not pd.isna(time) else 24.0, "drug_cell_idx": group['_cell_idx'].values, "vehicle_cell_idx": vehicle_same_cl['_cell_idx'].values, "pert_vec": pert_vec, "chembl_id": chembl_id, "target": target_name, "pathway_l1": pathway_l1, "pathway_l2": pathway_l2, }) self._conditions = conditions print(f" Conditions: {len(conditions)}") # Add drug_name → SMILES entries to cache (via ChEMBL ID lookup) for cond in self._conditions: name = cond['drug_name'] chembl = cond.get('chembl_id', '') if chembl and chembl in self._smiles_cache: self._smiles_cache[name] = self._smiles_cache[chembl] # Align gene dimension after conditions are built (updates pert_vecs) if self.target_num_genes is not None and self.num_genes != self.target_num_genes: self._align_genes(self.target_num_genes) drug_names = list(set(c['drug_name'] for c in conditions)) print(f" Unique drugs: {len(drug_names)}") # Drug embeddings self.drug_emb_dim = drug_emb_dim self._drug_name_to_idx = {name: i for i, name in enumerate(sorted(drug_names))} # _smiles_cache already populated by _load_smiles_cache() in __init__ self._drug_embeddings: Optional[torch.Tensor] = None if precomputed_drug_embeddings is not None: self._load_precomputed_embeddings(precomputed_drug_embeddings) else: self._compute_drug_embeddings() def _compute_drug_embeddings(self) -> None: """Compute Morgan fingerprint embeddings for all unique drugs.""" drug_names = sorted(self._drug_name_to_idx.keys()) n_drugs = len(drug_names) self._drug_embeddings = torch.zeros(n_drugs, self.drug_emb_dim) # Build name → ChEMBL ID mapping from conditions metadata name_to_chembl = self._build_name_chembl_map() try: from rdkit import Chem from rdkit.Chem import AllChem # Test that RDKit actually works (may fail with NumPy 2.x) Chem.MolFromSmiles("C") except (ImportError, AttributeError, Exception) as e: logger.warning("RDKit not available (%s) — drug embeddings will be zero vectors.", e) return success_count = 0 for i, name in enumerate(drug_names): smiles = "" # Strategy 1: Look up via ChEMBL ID (most reliable) chembl_id = name_to_chembl.get(name, "") if chembl_id: smiles = self._smiles_cache.get(chembl_id, "") # Strategy 2: Try drug name directly as SMILES (works for simple names like "C") if not smiles: smiles = name mol = Chem.MolFromSmiles(smiles) if mol is None: if chembl_id: logger.warning("No SMILES for drug '%s' (ChEMBL: %s, tried: %s)", name, chembl_id, smiles[:40] if smiles else "") continue fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048) fp_arr = np.array(list(fp), dtype=np.float32) # Simple MLP projection: 2048 → 512 → emb_dim fp_tensor = torch.from_numpy(fp_arr).float() hidden = torch.nn.functional.linear( fp_tensor, torch.randn(512, 2048) * 0.01, ) hidden = torch.nn.functional.gelu(hidden) emb = torch.nn.functional.linear( hidden, torch.randn(self.drug_emb_dim, 512) * 0.01, ) self._drug_embeddings[i] = emb.detach() success_count += 1 print(f" Pre-computed drug embeddings: {self._drug_embeddings.shape} " f"({success_count}/{n_drugs} drugs with fingerprints)") def _load_precomputed_embeddings(self, embeddings: Dict[str, torch.Tensor]) -> None: """Load pre-computed drug embeddings from a dictionary.""" drug_names = sorted(self._drug_name_to_idx.keys()) n_drugs = len(drug_names) self._drug_embeddings = torch.zeros(n_drugs, self.drug_emb_dim) for i, name in enumerate(drug_names): if name in embeddings: emb = embeddings[name] if emb.dim() == 1: emb = emb[:self.drug_emb_dim] elif emb.dim() == 2: emb = emb[0, :self.drug_emb_dim] self._drug_embeddings[i] = emb print(f" Loaded pre-computed drug embeddings: {self._drug_embeddings.shape}") # ------------------------------------------------------------------ def set_smiles(self, smiles_dict: Dict[str, str]) -> None: """Update SMILES for drugs and recompute embeddings. Parameters ---------- smiles_dict : dict mapping drug_name → SMILES string """ self._smiles_cache.update(smiles_dict) self._compute_drug_embeddings() # ------------------------------------------------------------------ def _load_gene_mapping(self, map_path: str) -> None: """Load ENSEMBL→gene symbol mapping from JSON file.""" import json try: with open(map_path) as f: mapping = json.load(f) # Build reverse map: symbol → list of ENSEMBL IDs self._symbol_to_ensembl: Dict[str, List[str]] = {} for eid, sym in mapping.items(): sym = sym.upper() if sym not in self._symbol_to_ensembl: self._symbol_to_ensembl[sym] = [] self._symbol_to_ensembl[sym].append(eid) print(f" Loaded gene mapping: {len(mapping)} ENSEMBL IDs → {len(self._symbol_to_ensembl)} symbols") except Exception as e: print(f" Warning: Failed to load gene mapping ({e})") self._symbol_to_ensembl = {} def _load_smiles_cache(self, csv_path: str) -> None: """Load SMILES from a ChEMBL CSV into _smiles_cache. CSV format: chembl_id,smiles Keys in _smiles_cache: both chembl_id and (later) drug_name. """ if not csv_path or not os.path.exists(csv_path): return import csv as csv_mod try: with open(csv_path, newline="", encoding="utf-8") as f: reader = csv_mod.DictReader(f) for row in reader: cid = row.get("chembl_id", "").strip() smi = row.get("smiles", "").strip() if cid and smi: self._smiles_cache[cid] = smi n_loaded = sum(1 for k in self._smiles_cache if k.startswith("CHEMBL")) print(f" Loaded {n_loaded} SMILES from {os.path.basename(csv_path)}") except Exception as e: print(f" Warning: Failed to load SMILES CSV ({e})") def _build_name_chembl_map(self) -> Dict[str, str]: """Build drug_name → ChEMBL ID mapping from loaded conditions.""" mapping: Dict[str, str] = {} for cond in getattr(self, '_conditions', []): name = cond.get('drug_name', '') chembl = cond.get('chembl_id', '') if name and chembl and not pd.isna(chembl): mapping[name] = str(chembl).strip() return mapping def _align_genes(self, target_num_genes: int) -> None: """Pad or truncate gene dimension to match target_num_genes. When SciPlex3 has a different number of genes than the CRISPRi dataset (e.g., due to different HVG selection), this method pads the expression matrix with zeros to match the target dimension. Parameters ---------- target_num_genes : int Target number of genes (from the model / other dataset). """ if self.num_genes == target_num_genes: return if self.num_genes > target_num_genes: # Truncate: keep first target_num_genes genes print(f" Aligning genes: truncating {self.num_genes} → {target_num_genes}") self._X = self._X[:, :target_num_genes] self.gene_names = self.gene_names[:target_num_genes] self.num_genes = target_num_genes self._gene_to_idx = {g: i for i, g in enumerate(self.gene_names)} else: # Pad with zeros pad_width = target_num_genes - self.num_genes print(f" Aligning genes: padding {self.num_genes} → {target_num_genes} (+{pad_width} zeros)") self._X = np.pad(self._X, ((0, 0), (0, pad_width)), mode='constant') self.gene_names = np.concatenate([ self.gene_names, np.array([f"PAD_{i}" for i in range(pad_width)]), ]) self.num_genes = target_num_genes self._gene_to_idx = {g: i for i, g in enumerate(self.gene_names)} # Update perturbation vectors in conditions (if they exist) if hasattr(self, '_conditions') and self._conditions: for cond in self._conditions: old_vec = cond['pert_vec'] if len(old_vec) < target_num_genes: new_vec = np.zeros(target_num_genes, dtype=np.float32) new_vec[:len(old_vec)] = old_vec cond['pert_vec'] = new_vec elif len(old_vec) > target_num_genes: cond['pert_vec'] = old_vec[:target_num_genes] # ------------------------------------------------------------------ def _build_drug_target_map( self, drug_obs, external_mapping: Optional[Dict[str, List[str]]], gene_names: List[str], ) -> Dict[str, np.ndarray]: """Build per-drug target gene vectors in HVG space.""" drug_target_map: Dict[str, np.ndarray] = {} # From SciPlex3 target annotations for drug_name in drug_obs['perturbation'].unique(): target_cat = drug_obs[drug_obs['perturbation'] == drug_name]['target'].iloc[0] gene_targets = self._target_category_to_genes( target_cat, gene_names, self._gene_to_idx, getattr(self, '_symbol_to_ensembl', None) ) drug_target_map[drug_name] = gene_targets # Merge with external mapping (DGIdb) if external_mapping: for drug_name, gene_list in external_mapping.items(): targets = np.zeros(len(gene_names), dtype=np.float32) for g in gene_list: if g in self._gene_to_idx: targets[self._gene_to_idx[g]] = 1.0 drug_target_map[drug_name] = targets return drug_target_map def _build_perturbation_vector( self, drug_name: str, external_mapping: Optional[Dict] ) -> np.ndarray: """Build multi-hot perturbation vector for a drug.""" pert_vec = np.zeros(self.num_genes, dtype=np.float32) # From internal mapping if drug_name in self._drug_target_map: pert_vec = self._drug_target_map[drug_name].copy() # From external mapping if external_mapping and drug_name in external_mapping: for g in external_mapping[drug_name]: if g in self._gene_to_idx: pert_vec[self._gene_to_idx[g]] = 1.0 return pert_vec # ------------------------------------------------------------------ # Utility methods # ------------------------------------------------------------------ @staticmethod def _normalize_counts(X: np.ndarray, target_sum: float = 1e4) -> np.ndarray: """Library-size normalization + log1p.""" from ..data.gene_selection import normalize_counts return normalize_counts(X, target_sum=target_sum, log1p=True) @staticmethod def _select_hvg(X: np.ndarray, n_genes: int = 2000) -> np.ndarray: """Select highly variable genes.""" from ..data.gene_selection import select_hvg_by_variance return select_hvg_by_variance(X, n_genes=n_genes) @staticmethod def _target_category_to_genes( target_category: str, gene_names: List[str], gene_to_idx: Optional[Dict[str, int]] = None, symbol_to_ensembl: Optional[Dict[str, List[str]]] = None, ) -> np.ndarray: """Map SciPlex3 protein target category to gene targets. This is a simplified mapping. For production use, replace with a curated mapping from ChEMBL/DGIdb. Parameters ---------- target_category : str SciPlex3 target annotation, e.g. "Aurora Kinase", "HDAC", "CDK" gene_names : list of str All gene names in the HVG space. Returns ------- gene_targets : [G] binary vector """ # Curated mapping: target category → known target genes # This should be loaded from a proper database in production TARGET_GENE_MAP = { "Aurora Kinase": ["AURKA", "AURKB", "AURKC"], "EGFR": ["EGFR"], "HDAC": ["HDAC1", "HDAC2", "HDAC3", "HDAC4", "HDAC5", "HDAC6", "HDAC7", "HDAC8", "HDAC9", "HDAC10", "HDAC11"], "DNA alkylator": [], "Bcl-2": ["BCL2", "BCL2L1"], "PKC": ["PRKCA", "PRKCB", "PRKCG", "PRKCD", "PRKCE", "PRKCH", "PRKCI", "PRKCL", "PRKCM", "PRKCN", "PRKDC", "PRKD1", "PRKD2", "PRKD3"], "CDK": ["CDK1", "CDK2", "CDK4", "CDK6", "CDK7", "CDK8", "CDK9", "CDK12", "CDK13"], "Histone Methyltransferase": ["EZH2", "DOT1L", "PRMT1", "PRMT5", "SETD2", "SETDB1", "SUV39H1", "SUV39H2"], "DNA/RNA Synthesis": ["TOP1", "TOP2A", "TOP2B", "RRM1", "RRM2", "TYMS", "DHFR"], "IGF-1R": ["IGF1R"], "PARP": ["PARP1", "PARP2", "PARP3", "PARP4", "PARP6", "PARP8", "PARP10", "PARP11", "PARP12", "PARP14", "PARP15"], "Sirtuin": ["SIRT1", "SIRT2", "SIRT3", "SIRT4", "SIRT5", "SIRT6", "SIRT7"], "TNF-alpha": ["TNF", "TNFRSF1A", "TNFRSF2"], "JAK": ["JAK1", "JAK2", "JAK3", "TYK2"], "Topoisomerase": ["TOP1", "TOP2A", "TOP2B"], "CCR": ["CCR1", "CCR2", "CCR3", "CCR4", "CCR5", "CCR6", "CCR7", "CCR8", "CCR9", "CCR10"], "HSP (e.g. HSP90)": ["HSP90AA1", "HSP90AB1", "HSPA1A", "HSPA1B", "HSPA5", "HSPA8", "HSPB1", "HSPD1", "HSPE1"], "HIF": ["HIF1A", "EPAS1", "HIF3A"], "Microtubule Associated": ["TUBB", "TUBA1A", "TUBA1B", "TUBA1C", "TUBB2A", "TUBB2B", "TUBB3", "TUBB4B"], "VEGF": ["VEGFA", "VEGFB", "VEGFC", "VEGFD"], "NF-kB": ["NFKB1", "NFKB2", "RELA", "RELB", "REL"], "PI3K": ["PIK3CA", "PIK3CB", "PIK3CD", "PIK3CG", "PIK3R1", "PIK3R2", "PIK3R3"], "AKT": ["AKT1", "AKT2", "AKT3"], "mTOR": ["MTOR"], "RAR": ["RARA", "RARB", "RARG"], "RXR": ["RXRA", "RXRB", "RXRG"], "BET": ["BRD2", "BRD3", "BRD4", "BRDT"], "IAP": ["BIRC2", "BIRC3", "BIRC5", "BIRC6", "BIRC7", "XIAP"], "Telomerase": ["TERT", "TERC"], "Mdm2": ["MDM2", "MDM4"], "TRAIL": ["TNFSF10", "TRAILR1", "TRAILR2", "TRAILR3", "TRAILR4"], "TGF-beta/Smad": ["TGFB1", "TGFB2", "TGFB3", "SMAD1", "SMAD2", "SMAD3", "SMAD4", "SMAD5", "SMAD6", "SMAD7"], "Wnt": ["CTNNB1", "APC", "AXIN1", "AXIN2", "GSK3B", "LRP5", "LRP6", "FZD1", "FZD2", "FZD3", "FZD4", "FZD5", "FZD6", "FZD7", "FZD8", "FZD9", "FZD10"], "Notch": ["NOTCH1", "NOTCH2", "NOTCH3", "NOTCH4", "JAG1", "JAG2", "DLL1", "DLL3", "DLL4"], "Hedgehog": ["SHH", "IHH", "DHH", "PTCH1", "PTCH2", "SMO", "GLI1", "GLI2", "GLI3"], "p53": ["TP53"], "Src": ["SRC", "FYN", "YES1", "FGR", "LCK", "HCK", "LYN", "BLK", "HCK", "MATK"], "BCL-2": ["BCL2", "BCL2L1", "MCL1"], "JAK/STAT": ["JAK1", "JAK2", "JAK3", "TYK2", "STAT1", "STAT2", "STAT3", "STAT4", "STAT5A", "STAT5B", "STAT6"], "PKA": ["PRKACA", "PRKACB", "PRKACG", "PRKAR1A", "PRKAR1B", "PRKAR2A", "PRKAR2B"], "PIP3": ["PIK3CA", "PIK3CB", "PIK3CD", "PIK3CG", "PTEN"], "GSK-3": ["GSK3A", "GSK3B"], "RAF": ["ARAF", "BRAF", "RAF1"], "MEK": ["MAP2K1", "MAP2K2"], "ERK": ["MAPK1", "MAPK3"], "p38": ["MAPK11", "MAPK12", "MAPK13", "MAPK14"], "JNK": ["MAPK8", "MAPK9", "MAPK10"], "Caspase": ["CASP1", "CASP2", "CASP3", "CASP4", "CASP5", "CASP6", "CASP7", "CASP8", "CASP9", "CASP10"], "Proteasome": ["PSMA1", "PSMA2", "PSMA3", "PSMA4", "PSMA5", "PSMA6", "PSMA7", "PSMB1", "PSMB2", "PSMB3", "PSMB4", "PSMB5", "PSMB6", "PSMB7", "PSMC1", "PSMC2", "PSMC3", "PSMC4", "PSMC5", "PSMC6", "PSMD1", "PSMD2", "PSMD3", "PSMD4", "PSMD5", "PSMD6", "PSMD7", "PSMD8", "PSMD9", "PSMD10", "PSMD11", "PSMD12", "PSMD13", "PSMD14"], "Calcineurin": ["PPP3CA", "PPP3CB", "PPP3CC", "PPP3R1", "PPP3R2"], "Carbonic Anhydrase": ["CA1", "CA2", "CA3", "CA4", "CA5A", "CA6", "CA7", "CA8", "CA9", "CA10", "CA11", "CA12", "CA13", "CA14"], "Glucocorticoid Receptor": ["NR3C1"], "Estrogen Receptor": ["ESR1", "ESR2"], "Androgen Receptor": ["AR"], "Progesterone Receptor": ["PGR"], "Vitamin D Receptor": ["VDR"], "Retinoic acid receptor": ["RARA", "RARB", "RARG"], "FXR": ["NR1H4"], "LXR": ["NR1H2", "NR1H3"], "PPAR": ["PPARA", "PPARD", "PPARG"], "Serotonin Receptor": ["HTR1A", "HTR1B", "HTR1D", "HTR1E", "HTR1F", "HTR2A", "HTR2B", "HTR2C", "HTR3A", "HTR3B", "HTR3C", "HTR3D", "HTR3E", "HTR4", "HTR5A", "HTR6", "HTR7"], "Dopamine Receptor": ["DRD1", "DRD2", "DRD3", "DRD4", "DRD5"], "Histamine Receptor": ["HRH1", "HRH2", "HRH3", "HRH4"], "Adrenergic Receptor": ["ADRA1A", "ADRA1B", "ADRA1D", "ADRA2A", "ADRA2B", "ADRA2C", "ADRB1", "ADRB2", "ADRB3"], "Sigma Receptor": ["SIGMAR1", "SIGMAR2"], "Cannabinoid Receptor": ["CNR1", "CNR2"], "Adenosine Receptor": ["ADORA1", "ADORA2A", "ADORA2B", "ADORA3"], "Imidazoline Receptor": ["ADRA1A"], "GABA-A Receptor": ["GABRA1", "GABRA2", "GABRA3", "GABRA4", "GABRA5", "GABRA6", "GABRB1", "GABRB2", "GABRB3", "GABRG1", "GABRG2", "GABRG3", "GABRD", "GABRE", "GABRP", "GABRQ", "GABRR1", "GABRR2", "GABRR3"], "Glutamate Receptor": ["GRIA1", "GRIA2", "GRIA3", "GRIA4", "GRID1", "GRID2", "GRIN1", "GRIN2A", "GRIN2B", "GRIN2C", "GRIN2D", "GRIN3A", "GRIN3B", "GRIK1", "GRIK2", "GRIK3", "GRIK4", "GRIK5", "GRM1", "GRM2", "GRM3", "GRM4", "GRM5", "GRM6", "GRM7", "GRM8"], "Muscarinic Acetylcholine Receptor": ["CHRM1", "CHRM2", "CHRM3", "CHRM4", "CHRM5"], "Nicotinic Acetylcholine Receptor": ["CHRNA1", "CHRNA2", "CHRNA3", "CHRNA4", "CHRNA5", "CHRNA6", "CHRNA7", "CHRNA9", "CHRNA10", "CHRNB1", "CHRNB2", "CHRNB3", "CHRNB4", "CHRND", "CHRNE", "CHRNG"], "Opioid Receptor": ["OPRD1", "OPRK1", "OPRL1", "OPRM1"], "Angiotensin Receptor": ["AGTR1", "AGTR2"], "Chemokine Receptor": ["CCR1", "CCR2", "CCR3", "CCR4", "CCR5", "CCR6", "CCR7", "CCR8", "CCR9", "CCR10", "CXCR1", "CXCR2", "CXCR3", "CXCR4", "CXCR5", "CXCR6", "CXCR7", "XCR1", "CX3CR1"], "Cytokine Receptor": ["IL1R1", "IL2RA", "IL2RB", "IL2RG", "IL3RA", "IL4R", "IL5RA", "IL6R", "IL7R", "CSF1R", "CSF2R", "CSF3R", "EPOR", "FLT3", "KIT", "MPL"], "Tachykinin Receptor": ["TACR1", "TACR2", "TACR3"], "Platelet-derived growth factor receptor": ["PDGFRA", "PDGFRB"], "Fibroblast growth factor receptor": ["FGFR1", "FGFR2", "FGFR3", "FGFR4"], "Vascular endothelial growth factor receptor": ["FLT1", "KDR", "FLT4"], "Folate Receptor": ["FOLR1", "FOLR2", "FOLR3"], "NMDA Receptor": ["GRIN1", "GRIN2A", "GRIN2B", "GRIN2C", "GRIN2D"], "AMPA Receptor": ["GRIA1", "GRIA2", "GRIA3", "GRIA4"], "Kainate Receptor": ["GRIK1", "GRIK2", "GRIK3", "GRIK4", "GRIK5"], "Neurotrophin Receptor": ["NTRK1", "NTRK2", "NTRK3"], "Trk Receptor": ["NTRK1", "NTRK2", "NTRK3"], "Neuropeptide Receptor": [], "Protein kinase": [], "DNA polymerase": [], "RNA polymerase": [], "Ribosome": [], "Tubulin": ["TUBB", "TUBA1A", "TUBA1B", "TUBA1C", "TUBB2A", "TUBB2B", "TUBB3", "TUBB4B"], "Farnesyltransferase": ["FNTA", "FNTB"], "Geranylgeranyltransferase": ["PGGT1B", "RABGGTA", "RABGGTB"], "Dihydrofolate reductase": ["DHFR"], "Bacterial cell wall": [], "Fungal cell wall": [], "Bacterial ribosome": [], "Fungal ribosome": [], "Mitochondrial ribosome": ["MRPL1", "MRPL2", "MRPL3", "MRPL4", "MRPL5", "MRPL6", "MRPL7", "MRPL8", "MRPL9", "MRPL10", "MRPL11", "MRPL12", "MRPL13", "MRPL14", "MRPL15", "MRPL16", "MRPL17", "MRPL18", "MRPL19", "MRPL20", "MRPL21", "MRPL22", "MRPL23", "MRPL24", "MRPL27", "MRPL28", "MRPL30", "MRPL32", "MRPL33", "MRPL34", "MRPL35", "MRPL36", "MRPL37", "MRPL38", "MRPL39", "MRPL40", "MRPL41", "MRPL42", "MRPL43", "MRPL44", "MRPL45", "MRPL48", "MRPL49", "MRPL50", "MRPS2", "MRPS5", "MRPS6", "MRPS7", "MRPS9", "MRPS10", "MRPS11", "MRPS12", "MRPS14", "MRPS15", "MRPS16", "MRPS17", "MRPS18A", "MRPS18B", "MRPS18C", "MRPS21", "MRPS22", "MRPS23", "MRPS24", "MRPS25", "MRPS26", "MRPS27", "MRPS28", "MRPS30", "MRPS31", "MRPS33", "MRPS34", "MRPS35"], } pert_vec = np.zeros(len(gene_names), dtype=np.float32) target_genes = TARGET_GENE_MAP.get(target_category, []) # Build index if not provided if gene_to_idx is None: gene_to_idx = {g: i for i, g in enumerate(gene_names)} # If gene_names are ENSEMBL IDs, try to convert target symbols gene_name_set = set(gene_names) if symbol_to_ensembl and any(str(g).startswith('ENSG') for g in gene_names[:10]): converted_targets = [] for g in target_genes: g_upper = g.upper() if g_upper in symbol_to_ensembl: for eid in symbol_to_ensembl[g_upper]: if eid in gene_name_set: converted_targets.append(eid) target_genes = converted_targets for g in target_genes: if g in gene_to_idx: pert_vec[gene_to_idx[g]] = 1.0 return pert_vec # ------------------------------------------------------------------ # Properties # ------------------------------------------------------------------ @property def condition_names(self) -> List[str]: return [c['pert_name'] for c in self._conditions] @property def unique_drugs(self) -> List[str]: return sorted(set(c['drug_name'] for c in self._conditions)) @property def cell_line_list(self) -> List[str]: return sorted(set(c['cell_line'] for c in self._conditions)) # --------------------------------------------------------------------------- # Collate function for SciPlex3 data # --------------------------------------------------------------------------- def sciplex_collate_fn(batch: List[Dict], num_genes: Optional[int] = None) -> Dict[str, torch.Tensor]: """Collate function for SciPlex3Dataset batches. Handles variable-size cell populations by padding to the max size in the batch. Optionally pads gene dimension to match `num_genes` for cross-dataset training (e.g., CRISPRi + SciPlex3). Parameters ---------- batch : list of dicts num_genes : int, optional Target gene dimension. If provided and larger than batch genes, pads with zeros. """ B = len(batch) max_ns = max(b['source_cells'].size(0) for b in batch) max_nt = max(b['target_cells'].size(0) for b in batch) G = batch[0]['source_cells'].size(1) device = batch[0]['source_cells'].device # Target gene dimension target_G = num_genes if num_genes is not None else G gene_pad = max(0, target_G - G) source_cells = torch.zeros(B, max_ns, target_G, device=device) target_cells = torch.zeros(B, max_nt, target_G, device=device) source_mask = torch.zeros(B, max_ns, device=device) target_mask = torch.zeros(B, max_nt, device=device) perturbation = torch.zeros(B, target_G, device=device) drug_smiles = [b['drug_smiles'] for b in batch] # list of str dose = torch.stack([b['dose'] for b in batch]) # [B] for i, b in enumerate(batch): ns = b['source_cells'].size(0) nt = b['target_cells'].size(0) # Copy gene data (existing genes) source_cells[i, :ns, :G] = b['source_cells'] target_cells[i, :nt, :G] = b['target_cells'] # Copy perturbation vector pert = b['perturbation'] perturbation[i, :pert.size(0)] = pert source_mask[i, :ns] = 1.0 target_mask[i, :nt] = 1.0 return { "source_cells": source_cells, "target_cells": target_cells, "source_mask": source_mask, "target_mask": target_mask, "perturbation": perturbation, "drug_smiles": drug_smiles, "dose": dose, "metadata": [b['metadata'] for b in batch], }