Downstream Analysis

Downstream Analysis#

This tutorial builds on the quickstarts (single / multi-sample) and shows the analyses TERRA embeddings unlock: gene-level embeddings, spatial gene-pair scoring, spatial structure (Earth Mover’s Distance), subsetting, and in-silico perturbation. We run them on the two Human Breast Cancer sections from the multi-sample quickstart.

Note

TERRA requires an NVIDIA GPU for the embedding step.

1. Setup, data, and embeddings#

We download the two breast cancer sections (as in the multi-sample quickstart) and embed both at once, this time also saving the tokenized dataset to disk — the downstream functions operate on that tokenized dataset.

import logging

import anndata as ad
import numpy as np
import pandas as pd
import scanpy as sc
import squidpy as sq
import torch
from datasets import load_from_disk

from terra import (download_pretrained, embed_dataset, get_average_gene_embed,
                   get_emd_distance, get_gene_embed, get_spatial_score,
                   harmonize_tokenize_embed_pipeline, perturb_dataset)
from terra.datasets import read_xenium_10x
from terra.utils.evaluation import get_top_gene_pairs, get_top_gene_score

logging.basicConfig(level="INFO")
sc.settings.set_figure_params(dpi=80, frameon=False)

# Two sections of the 10x Xenium Human Breast Cancer panel (313 genes). The
# reader downloads the small standalone files on first run and caches them.
def load_sample(base_url, out_dir, sample):
    adata = read_xenium_10x(base_url, out_dir)
    # Namespace cell ids by sample so they stay unique after concatenation.
    adata.obs["sample"] = sample
    adata.obs_names = sample + "_" + adata.obs_names.astype(str)
    adata.obs["cell_id"] = adata.obs_names.astype(str)
    # Crop a contiguous window per section, centred on the median cell
    # position (the bounding-box corner is often empty), to keep things fast.
    xy = adata.obsm["spatial"]
    center = np.median(xy, axis=0)
    keep = (np.abs(xy - center) <= 375).all(axis=1)
    return adata[keep].copy()


base = "https://cf.10xgenomics.com/samples/xenium/1.0.1"
rep1 = load_sample(f"{base}/Xenium_FFPE_Human_Breast_Cancer_Rep1/Xenium_FFPE_Human_Breast_Cancer_Rep1",
                   "data/breast_rep1", "rep1")
rep2 = load_sample(f"{base}/Xenium_FFPE_Human_Breast_Cancer_Rep2/Xenium_FFPE_Human_Breast_Cancer_Rep2",
                   "data/breast_rep2", "rep2")

adata = ad.concat([rep1, rep2])
print(adata.obs["sample"].value_counts())
print(adata)

# By default the model is cached in the Hugging Face cache; pass
# local_dir="terra_model" to download it into a folder of your choice.
model_dir = download_pretrained("Lotfollahi-lab/TERRA-112M")

# Batch size for every model-inference step below (the pipeline and the
# downstream functions). Lower it if you run out of GPU memory.
batch_size = 128

adata = harmonize_tokenize_embed_pipeline(
    adata=adata,
    sample_key="sample",      # tokenize each section separately
    batch_key="sample",
    model_folder_path=model_dir,
    cache_directory_path="./terra_cache",
    save_dataset_path="./terra_cache/breast.dataset",
    batch_size=batch_size,
    add_neigh_cell_ids=True,  # store neighbor IDs, needed for the
                              # neighborhood-of-selected-cells perturbation below
)

# The tokenized dataset that the downstream functions consume.
dataset = load_from_disk("./terra_cache/breast.dataset")

# Ensembl IDs of the panel genes (added to adata.var by harmonization).
gene_ids = list(adata.var["ensembl_id"])

2. Gene programs#

get_average_gene_embed returns, for each gene, its mean embedding across the dataset in two spaces: the intrinsic cell-level embedding (driven by a gene’s own expression) and the spatial neighborhood-level embedding (driven by its tissue context). Clustering genes by these embeddings groups co-varying genes into gene programs — clustering the intrinsic embeddings yields intrinsic gene programs, while clustering the spatial embeddings yields spatial gene programs.

avg = get_average_gene_embed(
    dataset=dataset,
    model_folder_path=model_dir,
    cell_gene_ensembl_id=gene_ids,
    neighborhood_gene_ensembl_id=gene_ids,
    batch_size=batch_size,
)

# Cluster genes by each embedding to recover gene programs: the intrinsic
# (cell-level) embedding gives intrinsic programs, the spatial
# (neighborhood-level) embedding gives spatial programs.
programs = {
    "Intrinsic": avg["cell_gene_emb_average_per_data"],
    "Spatial": avg["neighborhood_gene_emb_average_per_data"],
}
palettes = {"Intrinsic": "tab20", "Spatial": "tab20b"}  # distinct per view
for kind, gene_emb in programs.items():
    gene_emb = np.asarray(gene_emb)
    ok = ~np.isnan(gene_emb).any(axis=1)
    genes = sc.AnnData(gene_emb[ok])
    genes.obs_names = np.asarray(gene_ids, dtype=str)[ok]

    sc.pp.neighbors(genes, use_rep="X", n_neighbors=15)
    sc.tl.umap(genes)
    sc.tl.leiden(genes, resolution=0.5, flavor="igraph", n_iterations=2,
                 directed=False, key_added="program")
    sc.pl.umap(genes, color="program", palette=palettes[kind],
               title=f"{kind} gene programs")

3. Cell-level gene embeddings#

get_gene_embed returns a per-cell embedding for specific genes, in both the intrinsic (cell-level) and spatial (neighborhood-level) spaces. Replace the example Ensembl IDs with genes of interest from your panel. We store the embeddings in adata.obsm, then map one gene’s embedding across the tissue.

Note

The released models use Ensembl release 111 — look gene symbols up against that release for correct IDs.

genes_of_interest = gene_ids[:2]  # <-- replace with your genes of interest

gene_embed = get_gene_embed(
    dataset=dataset,
    model_folder_path=model_dir,
    cell_gene_ensembl_id=genes_of_interest,
    neighborhood_gene_ensembl_id=genes_of_interest,
    batch_size=batch_size,
)
for key, values in gene_embed.items():
    adata.obsm[key] = values
print(list(gene_embed))

# Map one gene's per-cell embedding across the tissue. Cells that don't express
# the gene are zeroed out by `get_gene_embed`, so we mask them and summarize the
# embedding of the expressing cells with its first principal component (a 1D
# score of how the model contextualizes the gene).
goi = genes_of_interest[0]
match = adata.var_names[adata.var["ensembl_id"] == goi]
symbol = match[0] if len(match) else goi

views = {
    "intrinsic": adata.obsm[f"cell_emb_gene{goi}"],
    "spatial": adata.obsm[f"neighborhood_emb_gene{goi}"],
}
for view, emb in views.items():
    expressed = np.abs(emb).sum(axis=1) > 0          # zeroed rows = gene absent
    score = np.full(adata.n_obs, np.nan)
    if expressed.sum() > 1:
        X = emb[expressed] - emb[expressed].mean(axis=0)
        pc1 = np.linalg.svd(X, full_matrices=False)[2][0]
        score[expressed] = X @ pc1
    col = f"{symbol} ({view})"
    adata.obs[col] = score
    for sample in adata.obs["sample"].unique():
        sq.pl.spatial_scatter(adata[adata.obs["sample"] == sample],
                              color=col, shape=None, size=6,
                              title=f"{col} — {sample}")

4. Spatial gene-pair score#

get_spatial_score compares each gene’s cell vs. neighborhood representation; the ratio highlights genes whose signal is spatially structured. get_top_gene_score / get_top_gene_pairs rank the strongest genes and gene pairs.

score = get_spatial_score(
    dataset=dataset,
    model_folder_path=model_dir,
    cell_gene_ensembl_id=gene_ids,
    neighborhood_gene_ensembl_id=gene_ids,
    batch_size=batch_size,
)
gene_pair_score = score["cos_sim_neighborhood"] / score["cos_sim_cell"]

df_scores = get_top_gene_score(
    gene_pair_score,
    cell_gene_ensembl_id=gene_ids,
    gene_df=adata.var,
    gene_counts=score["cell_count_neighborhood"],
    min_count=20,
)
df_scores.sort_values("gene_score", ascending=False).head(20)

df_pairs = get_top_gene_pairs(
    torch.from_numpy(gene_pair_score),
    count_cell=torch.from_numpy(score["cell_count_cell"]),
    count_neb=torch.from_numpy(score["cell_count_neighborhood"]),
    cos_sim_cell=torch.from_numpy(score["cos_sim_cell"]),
    cos_sim_neb=torch.from_numpy(score["cos_sim_neighborhood"]),
    gene_df=adata.var,
    cell_gene_ids=gene_ids,
    neighborhood_gene_ids=gene_ids,
    min_count=20,
    k=100,
)
df_pairs.head(20)

5. Spatial structure (Earth Mover’s Distance)#

For each cell, get_emd_distance compares the gene–gene relationships in the cell’s own (intrinsic) embedding against those in its spatial neighborhood, and summarizes the difference as an Earth Mover’s Distance (EMD). A high EMD marks cells whose local tissue context most reshapes their gene relationships — tissue boundaries and structured microenvironments — while a low EMD marks cells coherent with their surroundings. Plotting it in space shows how spatially organized each region is.

# get_emd_distance returns (per-cell EMD distance, per-cell gene matrix); the
# per-cell distance is what we map in space.
emd_dist, _ = get_emd_distance(
    dataset=dataset,
    model_folder_path=model_dir,
    cell_gene_ensembl_id=gene_ids,
    neighborhood_gene_ensembl_id=gene_ids,
    batch_size=batch_size,
)
adata.obs["emd_dist"] = emd_dist
# Plot each section separately (the two sections share a coordinate frame).
for sample in adata.obs["sample"].unique():
    sq.pl.spatial_scatter(adata[adata.obs["sample"] == sample],
                          color="emd_dist", shape=None, size=6, title=sample)

6. Analyze a subset of cells#

You can restrict the tokenized dataset to a set of cell IDs (e.g. a niche or cell type from a quickstart) and repeat any analysis on it.

# Restrict the tokenized dataset to a set of cell IDs. Read the cell_id column
# once and build the row indices, then use `select` -- far faster than `filter`,
# which scans and rewrites every (large) token row just to test one field.
subset_ids = set(map(str, adata.obs["cell_id"][:1000]))
keep_idx = [i for i, cid in enumerate(dataset["cell_id"]) if str(cid) in subset_ids]
subset = dataset.select(keep_idx)
print(subset)

7. In-silico perturbation#

perturb_dataset edits gene tokens, and re-embedding with embed_dataset shows the effect in latent space. A perturbation is a perturb_df with one row per edit, controlled along four axes:

  • type — knockout zeros a gene, foldchange scales it;

  • target — cell (the cell’s own expression) or neighborhood (its spatial context);

  • genes — specific Ensembl IDs or "all";

  • cells — specific cell IDs (matched on each cell’s cell_id) or "all".

Below we run four representative perturbations and compare how each shifts the embedding.

def perturb_and_embed(perturb_df):
    """Apply a perturbation to the tokenized dataset and return the new cell- and
    neighborhood-level embeddings (the original dataset is left unchanged)."""
    perturbed = perturb_dataset(dataset=dataset, perturb_df=perturb_df,
                                model_folder_path=model_dir, nproc=1)
    return embed_dataset(dataset=perturbed, model_folder_path=model_dir,
                         batch_size=batch_size)


results = {}

# 1) Knock out several genes in a selected subset of cells (e.g. a niche or cell
#    type of interest): a cell-target, specific-gene knockout on specific cells.
#    One row per (cell, gene) pair.
selected = adata.obs["cell_id"][:1000].tolist()
ko_genes = gene_ids[:200]
perturb_df = pd.DataFrame([
    {"perturbed_cell_id": cell, "perturbed_ensembl_id": gene,
     "perturbation_target": "cell", "perturbation_type": "knockout",
     "foldchange": np.nan}
    for cell in selected for gene in ko_genes
])
results["knockout genes (selected cells)"] = perturb_and_embed(perturb_df)

# 2) Double several genes in every cell: an all-cell, specific-gene, cell-target
#    fold-change.
perturb_df = pd.DataFrame({
    "perturbed_cell_id": "all",
    "perturbed_ensembl_id": gene_ids[:200],
    "perturbation_target": "cell",
    "perturbation_type": "foldchange",
    "foldchange": 2.0,
})
results["fold-change genes (all cells)"] = perturb_and_embed(perturb_df)

# 3) Halve the whole panel in each cell's neighborhood: a neighborhood-target
#    fold-change that perturbs the spatial context rather than the cell itself.
perturb_df = pd.DataFrame([{
    "perturbed_cell_id": "all",
    "perturbed_ensembl_id": "all",
    "perturbation_target": "neighborhood",
    "perturbation_type": "foldchange",
    "foldchange": 0.5,
}])
results["fold-change all (neighborhood)"] = perturb_and_embed(perturb_df)

# 4) Perturb the neighborhood *role* of a selected subset of cells: wherever one
#    of these cells appears in another cell's neighborhood, knock out its
#    contribution. Unlike (3), this targets specific cells, so it needs the
#    neighbor IDs stored via add_neigh_cell_ids=True.
selected = adata.obs["cell_id"][:200].tolist()
perturb_df = pd.DataFrame({
    "perturbed_cell_id": selected,
    "perturbed_ensembl_id": "all",
    "perturbation_target": "neighborhood",
    "perturbation_type": "knockout",
    "foldchange": np.nan,
})
results["neighborhood of selected cells"] = perturb_and_embed(perturb_df)

# How far each perturbation moved the embeddings (mean L2 shift vs. the
# unperturbed run), for both the cell and neighborhood representations. Edits on
# a selected subset of cells move the averages less than the whole-panel edits,
# which change every cell.
for emb_key in ["cell_emb", "neighborhood_emb"]:
    base = adata.obsm[emb_key]
    print(emb_key)
    for name, emb in results.items():
        shift = np.linalg.norm(emb[emb_key] - base, axis=1).mean()
        print(f"  {name:34s} mean shift {shift:.3f}")

# Visualize each perturbation in the space it acts on: the cell-target
# fold-change in the cell embedding, the neighborhood-target knockout in the
# neighborhood embedding. Original vs perturbed in a shared UMAP.
for name, emb_key in [("fold-change genes (all cells)", "cell_emb"),
                      ("neighborhood of selected cells", "neighborhood_emb")]:
    joint = sc.AnnData(np.vstack([adata.obsm[emb_key], results[name][emb_key]]))
    joint.obs["state"] = ["original"] * adata.n_obs + ["perturbed"] * adata.n_obs
    sc.pp.neighbors(joint, use_rep="X", n_neighbors=15)
    sc.tl.umap(joint)
    sc.pl.umap(joint, color="state", title=f"{name} — {emb_key}")

Next steps#

See the API for the full reference of every function used here, and the User Guide for the concepts behind the embeddings.

Contents