spatialrna

Spatial RNA graph data processing and analysis for single-molecule spatial transcriptomics (10x Xenium and Atera): transcripts → tile graphs and subgraphs → any graph neural network (train / validate / embed) → clusters → spatial maps.

This continues the lab's spatialrna 0.2 (Ruqian Lyu) with the same scope and adds reading of raw Xenium / Atera outputs, zarr storage, a fast KD-tree graph builder, balanced grid tiling for large sections, and GPU Gaussian mixtures. The ISTIST pipeline builds on it with its own model (CrossSampleGAT), transcript-level annotation, statistics and an HTML report.

Install

pip install -e .                 # readers, graph builders, clustering, maps (no torch)
pip install -e ".[gnn]"          # + SpatialRNA dataset, datasets, train_val
                                 #   (install torch + PyG for your CUDA first)

Python >= 3.10, zarr-python >= 3.

Modules

Quick start

import numpy as np, torch
import spatialrna as sr

df = sr.readers.load_transcripts("/path/to/output")          # x, y, z, qv, gene
print(sr.readers.describe("/path/to/output"))

# tile graphs (here: 100 near-square tiles with equal transcript counts)
# SpatialRNA(root="data/S01.zarr", sample_name="S01", cohort_zarr_path="data/cohort.zarr",
#            num_tiles=100, tile_mode="grid", one_hot_encoding=enc, ...)

# train any model(x, edge_index) on the tiles / subgraphs
from spatialrna import train_val as tv
from spatialrna.datasets import ZarrGraphDataset
graphs = ZarrGraphDataset.from_sample_store("data/S01.zarr", kind="subgraphs")
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss, acc = tv.train_link_prediction(model, graphs, opt, device, num_genes)
emb, tx_id = tv.embed_tile(model, graphs[0], device, num_genes)

# clusters and a map
labels, gmm = sr.cluster.gmm(emb.numpy(), 12)
fig = sr.viz.pixel_map(x, y, labels)

Tiling

SpatialRNA(tile_mode="strips") (default, as in 0.2) cuts num_tiles equal-width bands along tile_by_dim. tile_mode="grid" cuts near-square tiles with equal transcript counts: no empty tiles in irregular tissue and little padding with many tiles (~5–7 % at 500 tiles, vs ~40 % for 500 strips).

Tests

pip install -e ".[dev]"
pytest              # test_zarr_roundtrip needs the [gnn] extra

Note on the name

spatialrna 0.1–0.2 on PyPI is maintained by Ruqian Lyu; this is 0.3.0.dev0. Coordinate with her before uploading to PyPI under the same name.