README.md

---
license: mit
datasets:
- mims-harvard/SPATIA_MIST
tags:
- ICML
---
<h1 align="center">
  SPATIA: Multimodal Model for Prediction and Generation of Spatial Cell Phenotypes
</h1>

<p align="center">
  <a href="https://zitniklab.hms.harvard.edu/SPATIA/"><img src="https://img.shields.io/badge/Website-SPATIA-4CAF50?logo=googlechrome&logoColor=white" /></a>
  <a href="https://arxiv.org/abs/2507.04704"><img src="https://img.shields.io/badge/Paper-ICML2026-b31b1b?logo=arxiv&logoColor=white" /></a>
  <a href="https://huggingface.co/datasets/mims-harvard/SPATIA_MIST"><img src="https://img.shields.io/badge/🤗%20HuggingFace-SPATIA__MIST-FFD21E" /></a>
  <a href="https://github.com/mims-harvard/SPATIA"><img src="https://img.shields.io/badge/GitHub-Code-181717?logo=github&logoColor=white" /></a>
  <a href="https://mims-harvard.readthedocs.io/en/latest/"><img src="https://img.shields.io/badge/Docs-readthedocs-blue?logo=readthedocs&logoColor=white" /></a>
</p>

## Overview

**SPATIA** is a multi-scale model for spatial transcriptomics that learns cell-level embeddings by fusing image-derived morphological tokens and transcriptomic tokens via cross-attention, aggregates them at niche and tissue levels, and generates cell morphology images conditioned on predicted state transitions using flow matching.

This repository hosts pretrained model weights, precomputed embeddings, evaluation results, and tutorial notebooks for reproducing the paper's prediction and generation experiments.

---

## Repository Structure

```
mims-harvard/SPATIA/
├── embeddings/
│   └── spatia_embeddings.npy          # Precomputed cell embeddings (98.5 MB)
├── notebooks/
│   ├── prediction_tasks_showcase.ipynb # Clustering & annotation benchmarks
│   └── spatia_generation_inference.ipynb  # Flow-matching image generation
├── pretrained/
│   ├── best_model.pt                  # Pretrained SPATIA-scGPT weights (220 MB)
│   ├── args.json                      # Training configuration
│   ├── vocab.json                     # Gene vocabulary (30k+ genes)
│   └── all_dict_mean_std.csv          # Gene expression normalization statistics
├── results/
│   ├── table2_summary.csv             # Clustering benchmarks (Table 2)
│   └── table4_clustering_summary.csv  # scRNA-seq clustering (Table 4)
├── .gitattributes                     # Git LFS tracking rules
└── README.md
```

---

## Components

### `pretrained/` — Model Weights & Config

Pretrained SPATIA-scGPT checkpoint for cell representation learning.

| File | Description |
|------|-------------|
| `best_model.pt` | Model weights (12-layer transformer, 512-dim, trained on MIST atlas) |
| `args.json` | Full training hyperparameters (lr, batch size, masking ratios, etc.) |
| `vocab.json` | Gene-to-token vocabulary mapping |
| `all_dict_mean_std.csv` | Per-gene mean and std for input normalization |

**Usage:**

```python
import torch

# Load pretrained weights
checkpoint = torch.load("pretrained/best_model.pt", map_location="cpu")
model.load_state_dict(checkpoint)
```

### `embeddings/` — Precomputed Cell Embeddings

`spatia_embeddings.npy` contains cell embeddings extracted from the pretrained SPATIA model on the MIST atlas. These embeddings fuse morphological and transcriptomic information and can be used directly for downstream tasks without requiring a GPU.

**Usage:**

```python
import numpy as np

embeddings = np.load("embeddings/spatia_embeddings.npy")
print(embeddings.shape)  # (num_cells, embedding_dim)
```

### `notebooks/` — Tutorial Notebooks

| Notebook | Description |
|----------|-------------|
| `prediction_tasks_showcase.ipynb` | Reproduces clustering (Leiden) and cell-type annotation benchmarks from Tables 2 & 4. Runs on frozen embeddings — no GPU needed. |
| `spatia_generation_inference.ipynb` | Loads the generation model checkpoint, runs ODE-based flow matching to generate perturbation cell images, and evaluates with FID, SSIM, and morphology metrics. |

**Run locally:**

```bash
pip install jupyter
jupyter notebook notebooks/prediction_tasks_showcase.ipynb
```

### `results/` — Evaluation Results

Precomputed benchmark results from the paper, ready for comparison.

| File | Content |
|------|---------|
| `table2_summary.csv` | Cell clustering (ARI, NMI) across Xenium & CosMx platforms for 6 models (PCA, scGPT, scFoundation, Nicheformer, UCE, SPATIA) |
| `table4_clustering_summary.csv` | scRNA-seq clustering benchmarks (ARI, NMI) for PCA, scGPT, Geneformer |


### `.gitattributes` — Git LFS Configuration

Configures Git Large File Storage for binary files (`.pt`, `.npy`, `.h5`, `.ckpt`, etc.). Required for cloning — install [Git LFS](https://git-lfs.github.com/) before pulling:

```bash
git lfs install
git clone https://huggingface.co/mims-harvard/SPATIA
```

Or download individual files via the HuggingFace Hub:

```python
from huggingface_hub import hf_hub_download

# Download pretrained weights
model_path = hf_hub_download(repo_id="mims-harvard/SPATIA", filename="pretrained/best_model.pt")

# Download embeddings
emb_path = hf_hub_download(repo_id="mims-harvard/SPATIA", filename="embeddings/spatia_embeddings.npy")
```

---

## Quick Start

```python
from huggingface_hub import snapshot_download

# Download everything
snapshot_download(repo_id="mims-harvard/SPATIA", local_dir="./SPATIA")

# Or download only what you need
from huggingface_hub import hf_hub_download

# For prediction tasks (no GPU needed):
hf_hub_download(repo_id="mims-harvard/SPATIA", filename="embeddings/spatia_embeddings.npy", local_dir="./SPATIA")
hf_hub_download(repo_id="mims-harvard/SPATIA", filename="notebooks/prediction_tasks_showcase.ipynb", local_dir="./SPATIA")
```

---

## Related Resources

- **Code**: [github.com/mims-harvard/SPATIA](https://github.com/mims-harvard/SPATIA)
- **Dataset**: [mims-harvard/SPATIA_MIST](https://huggingface.co/datasets/mims-harvard/SPATIA_MIST) (MIST atlas for spatial transcriptomics)
- **Documentation**: [mims-harvard.readthedocs.io](https://mims-harvard.readthedocs.io/en/latest/)
- **Paper**: [arXiv 2507.04704](https://arxiv.org/abs/2507.04704)

---

## Citation

```bibtex
@article{kong2026spatia,
  title={Spatia: Multimodal model for prediction and generation of spatial cell phenotypes},
  author={Kong, Zhenglun and Qiu, Mufan and Boesen, John and Lin, Xiang and Yun, Sukwon and Chen, Tianlong and Kellis, Manolis and Zitnik, Marinka},
  booktitle = {Proceedings of the 43th International Conference on Machine Learning},
  year={2026}
}
```

## Contact

- Zhenglun Kong — [zhenglun_kong@hms.harvard.edu](mailto:zhenglun_kong@hms.harvard.edu)
- Marinka Zitnik — [marinka@hms.harvard.edu](mailto:marinka@hms.harvard.edu)

Zitnik Lab, Harvard Medical School