SubcellPert: Benchmarking Subcellular Perturbation Prediction

Jinghao Wang¹, Hanqun Cao¹, Xilin Dang¹, Tianhui Zhou²,
Yiyang Zhao^3,4, Chang-Yu Hsieh³, Chunbin Gu^1,✉, Pheng Ann Heng¹

¹ The Chinese University of Hong Kong   ² Duke University
³ Zhejiang University   ⁴ Shanghai Innovation Institute

TL;DR: SubcellPert benchmarks perturbation-response prediction across 14 imaging datasets and eight baselines, extending virtual cell evaluation from gene counts to intracellular RNA organization.

📄 Abstract

Perturbations change RNA abundance and intracellular organization, but gene-count predictors omit the positions needed to model spatial responses. SubcellPert introduces perturbation-response prediction at single-RNA-molecule resolution, with 14 imaging datasets, 13.88 million cells, and 3.65 billion RNA molecules. Given source observations and a directed condition change, models generate target RNA populations in normalized two-dimensional coordinates.

Two independently trained tasks separate abundance from localization: Oracle predicts coordinates given exact target counts; Joint predicts both. Eight baselines span source resampling, classical transport, and frozen foundation-model adapters. Perturbation and cellular-context holdouts evaluate abundance error, target-gene coverage, and conditional spatial error. This repository provides the data loading, baseline adapters, training, validation, and evaluation code.

🗂️ Codebase Structure

Repository layout

SubcellPert/
├── benchmark/           # Frozen schedules and sampling protocol
├── configs/
│   ├── datasets/        # Dataset definitions
│   ├── models/          # Foundation-model settings and resource paths
│   ├── oracle/          # Localization with target counts
│   └── joint/           # Abundance and localization prediction
├── docs/                # README figures
├── scripts/             # Pretrained resource download and preparation
├── src/subcellpert/     # Data loading, baselines, training, and evaluation
├── LICENSE
└── pyproject.toml

🛠️ Environment Setup

uv is used for environment management in this project. Dependencies are specified in pyproject.toml.

Run these commands from the code repository root:

uv venv --python 3.12 .venv
source .venv/bin/activate
uv pip install -e .
python -m subcellpert.main --help

Data and pretrained weights are downloaded separately. Keep each model's weights, vocabulary, and configuration together under weights/<model>/, using the pretrained model preparation script.

💾 Data Preparation

The dataset repository is HF king159/SubcellPert.

Authenticate with hf auth login or HF_TOKEN. Use a new, empty destination: the Hub downloader does not remove leftover files from older snapshots. On clusters without compute-node internet access, download on a networked login node to shared storage, then validate and run experiments in a compute allocation.

Download the pinned dataset snapshot

export SUBCELLPERT_CODE="$PWD"

python - <<'PY'
import os
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="king159/SubcellPert",
    repo_type="dataset",
    ignore_patterns=[".gitattributes"],
    token=True,
)
PY

🧬 SubcellPert

📐 Evaluation Protocol

The benchmark has two independently fitted settings:

Setting	Prediction task	Evaluation
Oracle	Predict RNA coordinates given target per-cell gene counts	Spatial W1
Joint	Predict gene counts and RNA coordinates from source cells and a condition change	Expression MSE, PCC-delta, gene coverage, and conditional spatial W1

There are 224 run configurations: 14 datasets × eight baselines × two settings. The frozen schedule contains 38 dataset/split instances and 608 planned fits. Most datasets have split seeds 1–3; ATM-RT and DDX54-PD1 have seed 1 only.

🧪 Baselines

Baseline	Configuration name	Pretrained weights
Count-matched resampling	count_matched_resampling	Not required
Gaussian affine transport	gaussian_affine_transport	Not required
Radial quantile transport	radial_quantile_transport	Not required
scGPT	scgpt	Required
scGPT-spatial	scgpt_spatial	Required
scFoundation	scfoundation	Required
SpatialFormer	spatialformer	Required
STATE-SE	state_se	Required

The five foundation-model baselines use frozen encoders with trainable prediction heads. Their networks are implemented directly in src/subcellpert/models/. Resource locations are specified in configs/models/.

Download the pretrained models

The five foundation-model baselines require pretrained weights and auxiliary resources. From the code repository root, with the environment active, run:

python scripts/prepare_pretrained_models.py

The script downloads all five models into weights/, extracts scGPT-spatial, and prepares the inference configuration for SpatialFormer. Existing files are kept. To download only selected models:

python scripts/prepare_pretrained_models.py --models scgpt state_se

Downloads use huggingface_hub (installed with the project dependencies) and the following sources. Hugging Face and GitHub resources use fixed revisions recorded in the script; Figshare downloads use fixed file IDs.

Model	Download source	Destination under weights/
scGPT	whole-human HF mirror; official model zoo	scgpt-human/
scGPT-spatial	official V1 archive	scgpt-spatial/scGPT_spatial_v1/
scFoundation	HF checkpoint mirror and official gene index	scfoundation/
SpatialFormer	official single-input 5k checkpoint and upstream inference resources	spatialformer/
STATE-SE	official SE-600M and MGI mouse–human homology report	state-se/

Required resource layout:

weights/
├── scgpt-human/
│   ├── best_model.pt
│   ├── args.json
│   └── vocab.json
├── scgpt-spatial/scGPT_spatial_v1/
│   ├── best_model.pt
│   ├── args.json
│   ├── vocab.json
│   ├── batch_mapping_dict.json
│   └── all_dict_mean_std.csv
├── scfoundation/
│   ├── models.ckpt
│   └── OS_scRNA_gene_index.19264.tsv
├── spatialformer/
│   ├── stepstep=0176000-traintrain_total_loss=-2.7789-valval_total_loss=0.0000.ckpt
│   ├── xenium_5k_pandavid_gene_embeddings.pkl
│   ├── config/_config_train_large_pair.json
│   ├── tokenizer/tokenv5.json
│   └── data/gene_median.pkl
└── state-se/
    ├── model.safetensors
    ├── protein_embeddings.pt
    ├── config.yaml
    └── HOM_MouseHumanSequence.rpt

🚀 Training

Run from the code repository root with SUBCELLPERT_ROOT pointing to the prepared data snapshot. Each command fits a baseline, performs validation where the frozen split provides it, and evaluates held-out responses.

🔎 Start with a smoke run

This small Oracle example uses a classical baseline and needs no pretrained weights:

python -m subcellpert.main \
  --project-root "$SUBCELLPERT_ROOT" \
  --config configs/oracle/pf-astro/count_matched_resampling.yaml \
  --seed 1 --fit-group seed1 --device cpu --smoke \
  --output "$SUBCELLPERT_CODE/outputs/smoke/oracle/pf-astro/resampling/seed1"

--smoke bounds the workload for a diagnostic run; its outputs are not benchmark results. Use --dry-run to inspect the resolved configuration and selected jobs without training.

🏋️ Run a full fit

After the smoke run, remove --smoke and choose a fresh output directory:

python -m subcellpert.main \
  --project-root "$SUBCELLPERT_ROOT" \
  --config configs/oracle/pf-astro/count_matched_resampling.yaml \
  --seed 1 --fit-group seed1 --device cuda \
  --output "$SUBCELLPERT_CODE/outputs/reproduction/oracle/pf-astro/resampling/seed1"

Classical baselines support --device cpu as well. To run Joint, select the corresponding YAML under configs/joint/ and use a separate output directory. All available dataset, baseline, setting, and split-seed choices are listed in configs/README.md.

For example, after preparing its pretrained resources, train the Joint scGPT adapter:

python -m subcellpert.main \
  --project-root "$SUBCELLPERT_ROOT" \
  --config configs/joint/pf-astro/scgpt.yaml \
  --seed 1 --fit-group seed1 --device cuda \
  --output "$SUBCELLPERT_CODE/outputs/reproduction/joint/pf-astro/scgpt/seed1"

The five foundation-model configurations specify five epochs, BF16, and a global training batch of 256. Encoder microbatches may be smaller without changing the optimizer batch. Scientific settings are defined in the YAML files.

Use a fresh output directory for every new fit. The CLI replaces generated files unless --resume is supplied. Resume a stopped fit with its original command plus --resume, keeping the code, data, weights, and training settings unchanged.

📊 Validation and Results

Foundation-model fits select checkpoints using the validation split. ATM-RT and DDX54-PD1 have no validation split and use the final epoch. Test responses are reserved for evaluation. Joint and Oracle are fitted independently.

Each fit saves its resolved configuration in run_config.json, its checkpoint in cell_model.pt, and evaluation results in result.json. Per-response predictions and metrics are stored under responses/. A completed benchmark fit has status: "complete" and scientific_result: true in result.json.

Report Joint expression MSE and PCC-delta and Oracle spatial W1 separately. The all-gene panel is primary; the top-100 DEG panel is supplementary. Shared-gene spatial scores and target-gene coverage are diagnostics: a missing predicted gene does not have zero spatial error. Compare spatial scores only over matching gene support.

📄 License

Original SubcellPert code is released under the MIT license. Integrated foundation-model code retains its original terms; attribution and license texts are included in the same LICENSE file. Pretrained weights and source datasets retain their respective terms.