User Guide

User Guide#

Overview#

TERRA is a self-supervised foundation model for spatial transcriptomics based on a Joint-Embedding Predictive Architecture (JEPA). A context-aware tokenization serializes each cell together with its spatial neighbors into a sequence of gene tokens — cells ordered by distance from the index cell, genes ranked by abundance within each cell — which preserves gene identity and works across heterogeneous gene panels. During pretraining some of these tokens are masked, and the model predicts their latent representations rather than reconstructing raw counts, inferring the masked molecular and spatial context of neighboring cells. Averaging the resulting gene-token embeddings then yields representations at three scales: genes, cells, and neighborhoods.

TERRA is pretrained on HST-Corpus-112M, a corpus of more than 100 million cells at single-cell resolution spanning human spatial-transcriptomics datasets. The pretrained embeddings transfer to a range of downstream tasks without task-specific retraining.

The inference pipeline#

The user-facing API is exposed at the top level of terra (its implementation lives in terra.inference). The typical workflow has three stages, exposed both as a single convenience function and as individual steps:

  1. Harmonize — map gene names to Ensembl IDs and apply quality control so the input matches the gene vocabulary the model was trained on (harmonize_adata).

  2. Tokenize — build the per-cell neighborhood token sequences against a trained model’s tokenizer (tokenize_adata).

  3. Embed — run the model to retrieve cell- and neighborhood-level embeddings (embed_dataset).

The convenience wrapper harmonize_tokenize_embed_pipeline runs all three:

from terra import download_pretrained, harmonize_tokenize_embed_pipeline

model_dir = download_pretrained("Lotfollahi-lab/TERRA-96M")

adata = harmonize_tokenize_embed_pipeline(
    adata=adata,                       # raw counts in adata.X
    sample_key="sample",
    batch_key="batch",
    model_folder_path=model_dir,
    cache_directory_path="./terra_cache",
)

The resulting cell- and neighborhood-level embeddings are stored in adata.obsm. See the Tutorials for the step-by-step version and downstream analyses, and the API for the full reference.

Running the steps individually#

The wrapper is equivalent to calling the three steps yourself, which is useful when you want to inspect or cache an intermediate result — the harmonized AnnData or the tokenized dataset. The gene-reference files the harmonizer needs live inside the downloaded bundle:

import os

from terra import (download_pretrained, harmonize_adata, tokenize_adata,
                   embed_dataset)

model_dir = download_pretrained("Lotfollahi-lab/TERRA-96M")

# 1. Harmonize: map gene symbols -> Ensembl IDs and apply QC, using the bundle's
#    gene-reference files (min_*_per_* = 0 mirrors the pipeline default).
adata = harmonize_adata(
    adata,                              # raw counts in adata.X
    gene_mapping_dict_file_path=os.path.join(model_dir, "ensembl_dictionary.pkl"),
    gene_occurrence_count_file_path=os.path.join(model_dir, "gene_count_dictionary.pkl"),
    min_genes_per_cell=0,
    min_cells_per_gene=0,
)

# 2. Tokenize: build the per-cell neighborhood token sequences.
dataset = tokenize_adata(
    adata=adata,
    model_folder_path=model_dir,
    cache_directory_path="./terra_cache",
)

# 3. Embed: cell- and neighborhood-level embeddings.
embeddings = embed_dataset(dataset=dataset, model_folder_path=model_dir)
for key, values in embeddings.items():   # cell_emb, neighborhood_emb, spatial_cell_emb
    adata.obsm[key] = values

The tokenized dataset is also what the downstream functions (get_average_gene_embed, get_spatial_score, perturb_dataset, …) consume, so you can persist it with dataset.save_to_disk(...) and reload it later. For multi-sample data, tokenize each sample separately and concatenate the resulting datasets — this is exactly what the pipeline’s sample_key does, and it keeps spatial neighborhoods from crossing sample boundaries; harmonizing and tokenizing the whole object at once treats it as a single sample.

Note

TERRA reports progress through the standard logging module rather than printing to stdout. To see progress messages (for example in a notebook), enable logging once:

import logging
logging.basicConfig(level="INFO")

The command-line entry points (terra-hub, terra.training, terra.inference) configure this for you automatically.

Pretrained models#

Pretrained TERRA models are distributed on the Hugging Face Hub. Each model is a self-contained bundle — the checkpoint, model config, token dictionary, and the gene-reference files needed to reproduce the model’s training-time harmonization (ensembl_dictionary.pkl, gene_count_dictionary.pkl).

Model

Training data

TERRA-96M

A 96M-cell subset of HST-Corpus-112M; the remaining cells are held out for benchmarking and downstream analyses.

TERRA-112M

The full HST-Corpus-112M.

Download a bundle with download_pretrained; the returned folder is what you pass as model_folder_path:

from terra import download_pretrained

model_dir = download_pretrained("Lotfollahi-lab/TERRA-96M")          # latest
model_dir = download_pretrained("Lotfollahi-lab/TERRA-96M", revision="v1.0")  # pin a tag/commit

Because the gene-reference files are part of the bundle, harmonization at inference time reproduces the tokenization the model was trained with — no external file paths are required.

In-silico perturbation#

perturb_dataset edits the gene tokens of a tokenized dataset; re-embedding the result with embed_dataset shows how the perturbation moves cells in latent space. Each edit is one row of a small perturb_df:

Column

Meaning

perturbed_cell_id

the cell to edit, or "all" for every cell

perturbed_ensembl_id

the gene to edit (Ensembl ID), or "all" for the whole panel

perturbation_target

"cell" or "neighborhood" (see below)

perturbation_type

"knockout" or "foldchange" (see below)

foldchange

the multiplier for foldchange rows (use np.nan for knockout)

Perturbation type — what happens to the targeted expression:

  • knockout sets it to zero (simulates removing the gene).

  • foldchange multiplies it by foldchange (e.g. 0.5 to halve, 2.0 to double).

Perturbation target — where the edit is applied. Each cell is tokenized together with its spatial neighbors, so a sequence has two parts:

  • cell edits the cell’s own gene tokens (its intrinsic expression).

  • neighborhood edits the tokens contributed by the cell’s spatial neighbors (its microenvironment), leaving the cell itself untouched.

Scope — perturbed_cell_id and perturbed_ensembl_id each accept a specific value or "all" (every cell / the whole panel). Multiple rows are applied together, so you can combine edits in a single run.

import numpy as np
import pandas as pd
from terra import perturb_dataset, embed_dataset

# Halve all genes and knock out one gene, for every cell, applied to the cell
# itself. Add `perturbation_target="neighborhood"` rows to perturb the context.
perturb_df = pd.DataFrame([
    {"perturbed_cell_id": "all", "perturbed_ensembl_id": "all",
     "perturbation_target": "cell", "perturbation_type": "foldchange", "foldchange": 0.5},
    {"perturbed_cell_id": "all", "perturbed_ensembl_id": "ENSG00000136997",  # a gene from your panel
     "perturbation_target": "cell", "perturbation_type": "knockout", "foldchange": np.nan},
])

perturbed = perturb_dataset(dataset, perturb_df, model_folder_path=model_dir)
emb = embed_dataset(perturbed, model_folder_path=model_dir)
# Compare emb["cell_emb"] / emb["neighborhood_emb"] against the unperturbed run.

Pass return_only_perturbed_cells=True to get back just the edited cells. See the Tutorials (downstream analysis) for a worked before/after comparison.

Citation#

If you use TERRA in your research, please cite the manuscript (in preparation). A BibTeX entry and DOI will be added on publication.

Contents