Virtual Cell Challenge

Predicting cellular responses to CRISPRi genetic perturbations in human H1 embryonic stem cells for the Virtual Cell Challenge 2025.

Combines causal GRN propagation, variance-weighted training, and corrected evaluation metrics — informed by ICML 2025 findings on mode collapse in perturbation prediction.

Architecture

Causal Edge Inference

The GNN supports two edge sources (--edge_source):

• coexpression — Pearson correlation from control cells (correlation, not causation)
• causal — RENGE-style directed inference: when CRISPRi knockdown of gene A causes gene B to change significantly (Wilcoxon p<0.05, |log₂FC|>0.5), that's a directed edge A→B. This is causal signal extracted directly from perturbation data.

Loss Functions

$$\mathcal{L} = \alpha \cdot \mathcal{L}{\text{recon}} + \beta \cdot \mathcal{L}{\text{disc}} + \gamma \cdot \mathcal{L}_{\text{DE}}$$

WMSE (from ICML 2025): weights each gene by its perturbation-induced variance, so predicting zero-delta for strongly affected genes is heavily penalized.

Evaluation Metrics

WMSE and R²_w(Δ) are computed automatically by evaluate_perturbations().

Project Structure

virtual-cell-challenge/
├── src/
│   ├── models/
│   │   ├── transformer.py          # PerturbationTransformer
│   │   ├── gnn.py                  # BiologicalGNN (register_buffer gene targeting)
│   │   ├── hybrid.py               # Gated transformer+GNN hybrid
│   │   ├── ensemble.py             # Metric-aware weighted ensemble
│   │   ├── foundation_adapter.py   # Backbone adapter (scratch / embeddings)
│   │   ├── pretrained_embeddings.py # Gene embedding loader with coverage guard
│   │   └── baseline.py             # Statistical baselines
│   ├── training/
│   │   ├── loss.py                 # CompositeLoss, WMSELoss, ContrastiveLoss, TripletLoss
│   │   ├── train.py                # Training loop (AMP, augmentation, W&B)
│   │   ├── finetune.py             # 3-phase progressive fine-tuning
│   │   └── hyperparam_search.py    # Optuna HPO (on real data)
│   ├── evaluation/
│   │   └── metrics.py              # MAE, WMSE, R²_w(Δ), DES, PDS
│   ├── data/
│   │   ├── dataset.py              # PerturbationDataset (global pert_to_id)
│   │   ├── augmentation.py         # Noise, Mixup, interpolation (Beta(5,1))
│   │   ├── preprocessing.py        # scanpy QC → normalize → log1p → HVG subset
│   │   ├── biological_knowledge.py # GRN edges: STRING, KEGG, co-expression, causal
│   │   └── mock_data.py            # Synthetic data for testing
│   └── utils/
│       ├── config.py               # VCCConfig (single source of truth)
│       └── seed.py                 # Reproducibility
├── scripts/
│   ├── generate_submission.py      # Submission builder (seeded noise, checkpoint context)
│   ├── run_full_pipeline.py        # End-to-end pipeline
│   ├── download_data.py            # Dataset fetcher
│   └── verify_submission_pipeline.py
├── notebooks/
│   ├── 01_eda.py                   # Exploratory Data Analysis
│   └── 02_pipeline_demo.py         # Pipeline demonstration
├── docs/                           # Research notes and analysis
│   ├── final.md                    # VCC 2025 winner strategy analysis
│   ├── fix1.md                     # Bug identification document 1
│   ├── fix2.md                     # Bug identification document 2
│   ├── logic.md                    # Technical + literature verification of all fixes
│   ├── achieve.md                  # Publication roadmap
│   ├── verdict.md                  # Honest assessment of codebase vs winners
│   └── masterplan.md               # Project plan
└── requirements.txt

Quick Start

Install

pip install -r requirements.txt

Required: torch, anndata, scanpy, scikit-learn, scipy, numpy, pandas Optional: torch_geometric, wandb, optuna

Train

# Transformer with WMSE loss (prevents mode collapse)
python -m src.training.train --model transformer --reconstruction wmse \
    --epochs 50 --data_path data/raw/train.h5ad

# GNN with causal edges + WMSE
python -m src.training.train --model gnn --edge_source causal --reconstruction wmse \
    --epochs 50 --data_path data/raw/train.h5ad

# GNN with co-expression edges + standard loss + triplet discrimination
python -m src.training.train --model gnn --edge_source coexpression \
    --discrimination triplet --use_embeddings --augment 3 \
    --data_path data/raw/train.h5ad

Training Flags

Generate Submission

python scripts/generate_submission.py --model gnn \
    --checkpoint checkpoints/best_gnn.pt \
    --data_path data/raw/train.h5ad \
    --test_perts_path data/test_perturbations.txt \
    --output submission.h5ad --n_cells 100 --n_control 500

The checkpoint contains pert_to_id, gene_names, and edge_index — no need to reconstruct them from data.

Key Design Decisions

Global perturbation vocabulary: A single pert_to_id mapping is built from all perturbations before train/val split, ensuring embedding IDs are consistent across splits.

Unknown perturbation handling: padding_idx = n_perturbations reserves a zeroed, frozen embedding for any perturbation not seen during training.

DES with log₂FC conversion: Scanpy log1p uses natural log. LFC values are divided by ln(2) before applying the |log₂FC| > 0.5 threshold to match competition specification.

No batch correction on adata.X: Harmony only corrects PCA embeddings by design. regress_out and ComBat introduce artifacts (Luecken et al. 2024). Batch is better handled as a model covariate.

Semi-hard negative mining: Triplet loss uses semi-hard negatives (farther than positive but within margin) when available, falling back to hard negatives. Prevents training instability from hardest mining in early epochs.

Documentation

Research notes and analysis are in docs/:

License

MIT