TRACE

This repository is the official implementation of Trajectory Representation and Composition Estimator (TRACE).

TRACE is a multi-task PFN-style Deep Learning model, fine-tuned for cell developmental coordinate prediction in organoid scRNA-seq data.

Note: TRACE Version 1.0 is fine-tuned on an early stage, human neuroepithelial organoid scRNA-seq time course dataset Jain et al., 2025.

Joint-rotational project at Queen Mary University London, University College London and The Alan Turing Institute under Prof. Julien Gautrot, Prof. Yanlan Mao and Dr. Isabel Palacios and Dr. Federico Nanni.

---

Table of Contents

• Background
• Repository Structure
• Installation
• Data Requirements
• Usage
• Known Issues & Limitations
• Contributing
• Citation
• License

---

Background

The brain develops into distinct regions with specialised functions and cell types. Neuroepithelial organoids model this human neurodevelopment, but how and when regionalized populations emerge transcriptionally remains unclear. Mapping these events is hard: scRNA-seq organoid datasets are siloed by batch effects, and standard integration blends datasets in ways that erase the transitional profiles where regional identity first appears. TRACE (Trajectory Representation and Composition Estimator) is a multi-task tabular prior-fitted network fine-tuned on scRNA-seq data for pseudotime and cell type composition prediction. It uses in-context learning for zero-shot inference on unseen organoid data.

---

Repository Structure

TRACE/
├── scripts/
│   ├── train_trace.sh          # SGE batch script — fine-tune TRACE on Myriad HPC
│   └── generate_figures.sh     # SGE batch script — generate pseudotime & composition figures
├── src/
│   ├── trace_icl/              # Installable Python package (import trace_icl)
│   │   ├── config/             #   ExperimentConfig, named presets, path registry
│   │   ├── context/            #   ContextSampler, CellTableBuilder, icl_collate
│   │   ├── data_layer/         #   ProcessedDataset, HVG selection, DPT pipeline
│   │   ├── eval/               #   TracePredictor, metrics (Spearman, Wasserstein, JSD)
│   │   ├── model/              #   TabICLRegressor, dual heads, loss functions
│   │   └── training/           #   Trainer, SupervisedTrainer, Muon optimiser, callbacks
│   ├── data_prep/
│   │   ├── css_pseudotime.R    # CSS + DPT pseudotime computation (run before data_prep.py)
│   │   └── data_prep.py        # Main preprocessing pipeline — h5ad → ProcessedDataset
│   ├── figures/
│   │   └── model/              # Publication figures: section A (pseudotime), B (composition)
│   └── train/
│       ├── model/              # TRACE training entry points (debug + Myriad)
│       └── baseline/           # Benchmark models: ridge, XGBoost RF (pseudotime & composition)
├── tests/
│   ├── unit/                   # 340 unit tests (config, data, context, model, training, eval)
│   └── integration/            # Integration tests (require full anndata install)
├── CHANGES.md
├── LICENSE
├── README.md
└── pyproject.toml

---

Installation

Prerequisites

• Python 3.11+

1. Clone the repository

git clone https://github.com/ChristianLangridge/TRACE.git
cd TRACE

2. Install the package and dependencies

pip install -e ".[dev]"

This installs the trace_icl package in editable mode along with all dependencies needed to run training, evaluation, and tests.

---

Data Requirements

Minimum (inference only)

File	Description	Source
data/training_data/AnnData/neurectoderm_complete.h5ad	Raw training scRNA-seq time course	See Data Requests badge above
data/training_data/AnnData/neurectoderm_with_pseudotime.h5ad	Preprocessed file with CSS + DPT pseudotime added to adata.obs — produced by running src/data_prep/css_pseudotime.R then src/data_prep/data_prep.py on neurectoderm_complete.h5ad	Generated locally (see Usage)
data/TabICLv2_checkpoint/tabicl-regressor-v2-20260212.ckpt	Pre-trained TabICLv2 backbone	soda-inria/tabicl

To request access to the neurectoderm dataset, use the Data Requests badge at the top of this page.

Expected layout

data/
├── training_data/
│   ├── AnnData/
│   │   ├── neurectoderm_complete.h5ad         # raw input
│   │   └── neurectoderm_with_pseudotime.h5ad  # generated by data_prep.py
│   ├── ML_data/                # Preprocessed ML inputs
│   └── matrix_bundle/
└── TabICLv2_checkpoint/
    └── tabicl-regressor-v2-20260212.ckpt

---

Usage

Zero-shot inference (CPU — no GPU required)

TRACE can be initialised and run for inference on a standard CPU. A GPU is only required for fine-tuning.

Download seed_3_final.pt from the latest GitHub Release (produced by Step 2 of fine-tuning) or provide your own checkpoint.

from trace_icl.eval.predictor import TracePredictor

# Load a TRACE checkpoint and the training h5ad
predictor = TracePredictor(
    ckpt_path="path/to/your_final.pt",
    h5ad_path="data/training_data/AnnData/neurectoderm_with_pseudotime.h5ad",
)

# Select cells to query (e.g. day-11 held-out cells)
cell_ids = predictor.select_cells("day11")

# Run stochastic inference (n_runs draws over random context samples)
result = predictor.predict(cell_ids, n_runs=100)

# result.pt_pred   — (n_runs, n_cells)     pseudotime predictions
# result.comp_pred — (n_runs, n_cells, K)  composition predictions
# result.pt_true   — ground-truth pseudotime
# result.comp_true — ground-truth soft labels

Inference on novel organoid data

TRACE uses in-context learning: at inference time each query cell is predicted against a set of anchor cells sampled from the training corpus. Novel cells must therefore be merged into the same h5ad as the training data before loading — the anchors always come from the training distribution.

Step 1 — Prepare your h5ad

Your novel h5ad must satisfy the following before merging:

obs column	Format	Purpose
rank-transformed-pseudotime	float ∈ [0, 1]	Pseudotime label (written by css_pseudotime.R)
orig.ident	string matching *_D{N} (e.g. EXP_D14)	Collection day extraction
class3	string cell-type label	Composition soft labels

Gene names (var_names) must overlap the training HVG list at data/training_data/ML_data/r_hvg_names.txt. Genes not in that list are silently dropped; the model operates on the intersection. Expression can be raw counts (auto-normalised to 1e4 CPM + log1p) or already log-normalised.

Run src/data_prep/css_pseudotime.R on a merged Seurat object to compute CSS-integrated pseudotime and write rank-transformed-pseudotime into obs.

Step 2 — Merge with the training h5ad and run inference

import anndata as ad
from trace_icl.eval.predictor import TracePredictor

# Merge training cells (anchor source) with your novel cells
training = ad.read_h5ad("data/training_data/AnnData/neurectoderm_with_pseudotime.h5ad")
novel    = ad.read_h5ad("path/to/your_novel_data.h5ad")
combined = ad.concat([training, novel], join="inner")  # inner join aligns gene space
combined.write_h5ad("path/to/combined.h5ad")

# Load the combined dataset
predictor = TracePredictor(
    ckpt_path="path/to/final.pt",
    h5ad_path="path/to/combined.h5ad",
)

# Predict novel cells by passing their cell IDs directly
novel_ids = list(novel.obs_names)
result = predictor.predict(novel_ids, n_runs=100)

result.pt_pred — (n_runs, n_cells) pseudotime predictions
result.comp_pred — (n_runs, n_cells, K) composition predictions
result.pt_true / result.comp_true — ground-truth values from obs (only meaningful if your novel cells have annotated pseudotime and cell types)

---

Fine-tuning with GPU access (HPC)

Fine-tuning requires an A100-class GPU. The full workflow below uses UCL's Myriad HPC (SGE scheduler).

Step 1 — Preprocess the training data

# Compute CSS + DPT pseudotime in R (outputs rank-transformed pseudotime to adata.obs)
Rscript src/data_prep/css_pseudotime.R

# Build ProcessedDataset from the h5ad (HVG selection, soft labels, centroid distances)
python src/data_prep/data_prep.py

Step 2 — Fine-tune TRACE

qsub scripts/train_trace.sh

This saves final.pt after all training steps, used for inference.

For a local debug run (MPS / small GPU, 1 000 steps):

python src/train/model/trace_debug_run.py

Step 3 — Generate figures

CKPT_PATH=/path/to/final.pt qsub scripts/generate_figures.sh

Outputs section A (pseudotime) and section B (composition) PDFs to results/figures/.

Step 4 — Run baseline comparisons

python src/train/baseline/pseudotime/ridge_pseudotime.py
python src/train/baseline/pseudotime/xgbrf_pseudotime.py
python src/train/baseline/composition/logistic_regression_comp.py
python src/train/baseline/composition/xgbrf_comp.py

---

Known Issues & Limitations

• Integration tests: tests/integration/ segfault on import due to an anndata/scanpy C-extension conflict with the test runner. This is a dependency issue, not a logic bug — all 340 unit tests pass cleanly.
• Prosencephalic progenitors: This cell type is consistently underperforming on the Wasserstein distance metric. This is a structural data limitation (severe class imbalance in the training time course) rather than a model failure.
• GPU requirement for fine-tuning: A100-class GPU recommended for forward passes with 1k-2k gene features. The 512-gene standard tier runs on smaller GPUs. Inference does not require a GPU.

---

Contributing

This is an active research project. If you'd like to contribute:

1. Fork the repository
2. Create a feature branch (git checkout -b feature/your-feature)
3. Commit your changes (git commit -m 'Add: your feature')
4. Push and open a Pull Request
5. Add a dated, detailed comment/annotation in the CHANGES.md file of the change

Please ensure any new scripts avoid hardcoded paths and include basic inline documentation.

---

Citation

If you use this codebase or the TRACE architecture in your work, please cite:

Langridge, C. (2025–2026). TRACE.
Joint-rotational project, Queen Mary University London, University College London
and The Alan Turing Institute.
https://github.com/ChristianLangridge/TRACE

---

License

This project is licensed under the BSD-3-Clause License. See LICENSE for details.

---

Developed as part of a joint rotational PhD project at Queen Mary University London, University College London and The Alan Turing Institute, under Prof. Julien Gautrot, Prof. Yanlan Mao, Dr. Isabel Palacios and Dr. Federico Nanni.