"""Preprocess SciPlex3 data for training. Memory-efficient: uses backed mode for initial load, filters metadata first, keeps sparse matrix through HVG selection, only extracts dense for final 2000 HVGs. Usage: python scripts/preprocess_sciplex3.py \ --h5ad /data/boom/Protein/regulatory_field/data/raw/scPerturb/rna_protein/SrivatsanTrapnell2020_sciplex3.h5ad \ --output data/processed/sciplex3_k562_24h.pt \ --cell_lines K562 \ --doses 10 100 1000 10000 \ --times 24 \ --n_hvg 2000 """ import argparse import gc import os import sys import numpy as np import torch def preprocess(args): """Preprocess SciPlex3 and save compact representation.""" try: import anndata as ad import scipy.sparse as sp except ImportError: print("Need anndata and scipy") sys.exit(1) import scanpy as sc print(f"Loading {args.h5ad}...") # Use backed mode to avoid loading full matrix into memory # We only need obs (metadata) for filtering first try: adata = ad.read_h5ad(args.h5ad, backed='r') print(f" Loaded in backed mode: {adata.shape}") # Materialize to memory (backed mode → in-memory) adata = adata.to_memory() print(f" Materialized to memory") except Exception as e: print(f" Backed mode failed ({e}), loading directly") adata = ad.read_h5ad(args.h5ad) print(f" Raw shape: {adata.shape}") # Filter by metadata BEFORE any dense operations obs = adata.obs.copy() if args.cell_lines: obs = obs[obs['cell_line'].isin(args.cell_lines)] if args.doses: obs = obs[obs['dose_value'].isin(args.doses)] elif not args.include_vehicle: obs = obs[obs['dose_value'] > 0] if args.times: obs = obs[obs['time'].isin(args.times)] keep_idx = obs.index.values n_keep = len(keep_idx) print(f" After filtering: {n_keep} cells (from {adata.shape[0]})") # Check if X is sparse before subsetting is_sparse = sp.issparse(adata.X) if is_sparse: print(f" Source matrix is sparse ({adata.X.nnz / 1e6:.1f}M non-zeros)") # Subset to filtered cells adata = adata[keep_idx].copy() print(f" Subset shape: {adata.shape}") # Check memory usage after subset if sp.issparse(adata.X): print(f" Sparse matrix: {adata.X.nnz / 1e6:.1f}M non-zeros, " f"{adata.X.data.nbytes / 1e6:.1f}MB data") else: print(f" Dense matrix: {adata.X.nbytes / 1e9:.2f} GB — WARNING: this will use lots of RAM!") # Normalize + log1p (operates on sparse matrix) print(" Normalizing...") sc.pp.normalize_total(adata, target_sum=1e4) sc.pp.log1p(adata) # HVG selection on SPARSE matrix — memory efficient print(f" Selecting {args.n_hvg} HVGs (sparse mode)...") try: # Try scanpy's HVG (works on sparse, uses seurat_v3) sc.pp.highly_variable_genes(adata, n_top_genes=args.n_hvg, flavor="seurat_v3") hvg_mask = adata.var['highly_variable'].values hvg_idx = np.where(hvg_mask)[0] print(f" Scanpy HVG selected: {len(hvg_idx)} genes") except Exception as e: print(f" Scanpy HVG failed ({e}), using variance-based sampling") # For variance-based: sample cells to estimate variance n_sample = min(10000, adata.shape[0]) sample_idx = np.random.choice(adata.shape[0], n_sample, replace=False) if sp.issparse(adata.X): X_sample = adata.X[sample_idx].toarray() else: X_sample = adata.X[sample_idx] gene_var = np.var(X_sample, axis=0) del X_sample gc.collect() top_idx = np.argsort(gene_var)[-args.n_hvg:] hvg_idx = np.sort(top_idx) print(f" Variance-based HVG: {len(hvg_idx)} genes") # Extract ONLY the HVG columns as dense (2000 genes × N cells) print(f" Extracting HVG expression matrix...") n_cells = adata.shape[0] n_hvg = len(hvg_idx) print(f" Extracting {n_hvg} genes × {n_cells} cells...") if sp.issparse(adata.X): X_hvg = adata.X[:, hvg_idx].toarray().astype(np.float32) else: X_hvg = np.array(adata.X[:, hvg_idx], dtype=np.float32) gene_names = list(adata.var.index[hvg_idx]) # Fix: use ensembl_id if var.index is non-standard (e.g. "nan:X") if hasattr(adata.var, 'columns') and 'ensembl_id' in adata.var.columns: ensembl_ids = adata.var['ensembl_id'].values.copy() # Skip header row if present (CSV parsing artifact) if len(ensembl_ids) > 0 and str(ensembl_ids[0]) == 'id gene_short_name': ensembl_ids = ensembl_ids[1:] # Pad to original length if header was removed if len(ensembl_ids) < len(adata.var): ensembl_ids = np.array(list(ensembl_ids) + ['unknown'] * (len(adata.var) - len(ensembl_ids))) gene_names = [str(ensembl_ids[i]) for i in hvg_idx] print(f" HVG matrix: {X_hvg.shape}, {X_hvg.nbytes / 1e9:.2f} GB") del adata # free memory gc.collect() # Extract metadata (obs was already filtered) obs['_cell_idx'] = np.arange(len(obs)) # Save os.makedirs(os.path.dirname(args.output) or '.', exist_ok=True) save_dict = { 'X': X_hvg, 'gene_names': gene_names, 'obs': obs, 'cell_lines': args.cell_lines, 'doses': args.doses, 'times': args.times, } torch.save(save_dict, args.output) print(f" Saved to {args.output}") file_size_gb = os.path.getsize(args.output) / 1e9 print(f" File size: {file_size_gb:.2f} GB") # Summary print(f"\nSummary:") print(f" Cells: {len(obs)}") print(f" Genes: {len(gene_names)}") 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()}") print(f" Perturbations: {obs['perturbation'].value_counts().head(10).to_dict()}") def main(): parser = argparse.ArgumentParser() parser.add_argument("--h5ad", required=True) parser.add_argument("--output", required=True) parser.add_argument("--cell_lines", nargs="*") parser.add_argument("--doses", nargs="*", type=float) parser.add_argument("--times", nargs="*", type=float) parser.add_argument("--n_hvg", type=int, default=2000) parser.add_argument("--include_vehicle", action="store_true") args = parser.parse_args() preprocess(args) if __name__ == "__main__": main()