Metadata-Version: 2.4
Name: mimyr
Version: 0.1.0
Summary: Generative modeling of spatial transcriptomics: cell locations, cell types, and gene expression for unseen tissue sections.
Author: Ajinkya Deshpande, Zhilei Bei, Jian Ma
Author-email: Spencer Krieger <skrieger@andrew.cmu.edu>
License-Expression: MIT
Project-URL: Repository, https://github.com/gkrieg/mimyr
Classifier: Programming Language :: Python :: 3
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy>=2.2.6
Requires-Dist: scipy>=1.15.3
Requires-Dist: pandas>=2.3.3
Requires-Dist: scikit-learn>=1.7.2
Requires-Dist: torch>=2.5.1
Requires-Dist: transformers>=4.30.2
Requires-Dist: scanpy>=1.11.5
Requires-Dist: anndata>=0.10.2
Requires-Dist: shapely>=2.1.2
Requires-Dist: slafdb>=0.5.2
Requires-Dist: tqdm>=4.67.3
Requires-Dist: loguru>=0.7.0
Requires-Dist: wandb>=0.21.0
Requires-Dist: matplotlib>=3.10.9
Requires-Dist: seaborn>=0.13.2
Requires-Dist: huggingface_hub>=0.23
Requires-Dist: filelock
Requires-Dist: pyyaml>=6.0
Dynamic: license-file

![MIMYR Method](https://raw.githubusercontent.com/gkrieg/mimyr/main/mimyr.jpg)

# MIMYR

MIMYR is a generative framework for spatial transcriptomics. Given neighboring
tissue sections registered to the Allen Common Coordinate Framework (CCF), it
generates an unseen section: where the cells are, what type each cell is, and
what genes each cell expresses.

It has three components, run in sequence:

1. **Location model**: a conditional diffusion model (DDPM) that samples cell
   coordinates in the plane of the new section, with optional backward guidance
   from a KDE fit to a neighboring reference section.
2. **Cell-type model**: a classifier that predicts each cell's type from its CCF
   coordinates.
3. **Expression model**: an autoregressive transformer that generates each cell's
   expression profile, conditioned on its location and cell type.

---

## Installation

MIMYR requires Python 3.10 or newer.

```bash
conda create --name mimyr python=3.10
conda activate mimyr
```

If you have an NVIDIA GPU, first install the PyTorch build that matches your CUDA
version (see [pytorch.org](https://pytorch.org/get-started/locally/)); otherwise
`pip` installs the default build.

Then install MIMYR:

```bash
pip install mimyr
```

This installs the `mimyr` package, its dependencies, and the gene and cell-type
metadata the models need. To run `main.py`, the demo, or the scripts that reproduce
the paper, clone the repository and install from source instead:

```bash
git clone https://github.com/gkrieg/mimyr.git
cd mimyr
pip install -e .
```

## Quickstart

From the repository root:

```bash
python main.py
```

This runs the full pipeline (location, cell type, and expression) with the
pretrained models on a small demo dataset in `demo_data/`. It generates one
section from its two neighbors and scores it against the real section, using the
model and settings behind the paper's Fig. 2. The first run downloads that model's
checkpoints (about 475 MB) from the Hugging Face Hub. The metrics are written to
`results/output.csv` and the generated section to `artifacts/`.

The demo data is the first rq1 test section (`sec_14`) and the two reference
sections next to it (`sec_13`, `sec_15`), each randomly subsampled to 5% of its
cells (3,562 to 4,719 cells per section). Because the sections are subsampled, its
scores are not comparable to the paper's. `demo_data/make_demo_data.py` rebuilds it
from the full rq1 data.

---

## Python API

`main.py` reproduces the paper's experiments. To use MIMYR on your own sections,
call it from Python. Sections are AnnData objects in the [input format](#input-format).

### Generate a section

```python
import anndata as ad
import mimyr

left = ad.read_h5ad("demo_data/sec_13.h5ad")
right = ad.read_h5ad("demo_data/sec_15.h5ad")

model = mimyr.Mimyr.from_pretrained("merscope")      # paper Fig. 2 model, from the Hub
section = model.generate(reference=[left, right], position=10.9, seed=0)
section.write_h5ad("generated.h5ad")
```

`position` is where the new section goes, in mm along the section axis: `x_ccf`
for coronal models, `z_ccf` for the sagittal `zhuang3-finetuned`. By default the
section has the references' extent and their mean number of cells; set `n_cells`
and `bounds` (e.g. `{"y_ccf": (1.0, 7.5), "z_ccf": (0.5, 10.5)}`) to change that.
The result has one row per generated cell, with:

- positions in `obs["x_ccf"]`, `obs["y_ccf"]`, `obs["z_ccf"]`;
- cell types in `obs["class"]`, `obs["subclass"]`, `obs["supertype"]` and `obs["cluster"]`;
- expression over the model's 2,000-gene panel in `X`, log1p-scaled like the input.

Expression is pulled toward the nearest same-type cell in the references
(reference guidance) at the strength the paper used. `reference_guidance=0` turns
it off, and then `reference` is optional if you pass `n_cells` and `bounds`.

Each stage can be swapped for a baseline: `location="closest_slice"`,
`cell_types="majority_baseline"` or `expression="lookup"`.

### Evaluate against a held-out section

```python
held_out = ad.read_h5ad("demo_data/sec_14.h5ad")
section = model.generate(reference=[left, right], position=held_out, seed=0)
mimyr.evaluate(section, held_out)
# {'soft_accuracy@0.03': ..., 'soft_spearman_correlation_radius@0.05': ..., 'soft_f1_radius@0.05': ..., ...}
```

Passing a section as `position` generates it again at the same position, with the
same extent and number of cells. Its expression is never read. `evaluate` takes
`metrics=` (any name in `mimyr.evaluator.METRICS`) and `sample=` (the percentage of
real cells scored).

### Train on your own data

```python
model = mimyr.Mimyr.from_pretrained("merscope")
model.train(train=train_sections, val=val_sections, output_dir="runs/my_lab", epochs=20)
model.save("my_model/")

model = mimyr.Mimyr.from_pretrained("my_model/")       # later
```

`train` fine-tunes the cell-type classifier and the expression model, starting
from the current weights; pass `components=("expression",)` to train just one.
The pretrained location model covers the whole mouse brain in CCF space, so it
works for any CCF-registered mouse data. To fine-tune it on your sections as well,
add `"location"` to `components` (`location_epochs`, default 400). Location
training supports coronal models only. It normalizes coordinates with fixed
constants for the CCF mouse brain, so the data must be in CCF mm.

The models use a fixed vocabulary: the 2,000-gene panel and the 5,274 Allen Brain
Cell Atlas clusters. Data must use these gene symbols and cluster names; other
genes are ignored. `mimyr.validate(sections)` reports what's missing or outside the
vocabulary. To train the expression model on several GPUs, run the script with
`torchrun --nproc_per_node=4 train.py`. Weights & Biases logging is off unless
`WANDB_MODE` is set. Batches load in the main process; `num_workers=4` moves loading
into worker processes. Each worker re-imports your script, so put its code under
`if __name__ == "__main__":` first.

`generate`, `train` and `evaluate` prepare raw sections themselves. To prepare
large sections once and reuse them across calls, use `mimyr.prepare(sections)`.

---

## Pretrained models and data

The pretrained checkpoints are hosted on the Hugging Face Hub at
[gkrieg13/mimyr](https://huggingface.co/gkrieg13/mimyr), and the benchmark data at
[gkrieg13/mimyr-data](https://huggingface.co/datasets/gkrieg13/mimyr-data). Both are
downloaded on first use and cached by `huggingface_hub` (`~/.cache/huggingface` by
default; set `HF_HOME` to move it).

`--pretrained` selects one of the models from the paper. It provides any checkpoint
you do not pass explicitly, and the inference settings used for that model in the
paper:

| `--pretrained` | Paper | Use with | Settings it applies |
|---|---|---|---|
| `merscope` (default) | Fig. 2 | `--data_mode rq1` or `demo` | reference guidance 75, `--metric_filter_by_gt` |
| `zhuang1` | Fig. 3, zero-shot | `--data_mode rq3_v2` | reference guidance 75, `--metric_filter_by_gt`, `--omit_x` |
| `zhuang2-finetuned` | Fig. 3 | `--data_mode rq3_v2` | reference guidance 75, `--metric_filter_by_gt`, `--omit_x` |
| `zhuang3-finetuned` | Fig. 4 | `--data_mode rq4` | reference guidance 50, `--metric_filter_by_gt` |

Settings are resolved in this order, each overriding the previous one: the built-in
defaults, the `--pretrained` model's settings, a `--config` file, and flags on the
command line. For example, `--no-expression_weak_teacher_forcing` turns reference
guidance off, and `--no-omit_x` and `--no-metric_filter_by_gt` undo the other two
settings.

Benchmark data is downloaded only for the data mode you run (the demo needs none) and
extracted into `--data_dir` (default `data/`):

| Archive | Size | Contents | Data modes |
|---|---|---|---|
| `rq1-merscope.zip` | 1.3 GB | 53 MERSCOPE sections from one brain | `rq1` |
| `Zhuang-ABCA-1.zip` | 1.0 GB | Zhuang-ABCA-1 (147 sections) | `base_zhuang1`, `rq4_noref` |
| `Zhuang-ABCA-2.zip` | 0.4 GB | Zhuang-ABCA-2 (66 sections) | `rq3`, `rq3_v2*`, `rq4_rq3` |
| `Zhuang-ABCA-3.zip` | 0.5 GB | Zhuang-ABCA-3 (23 sagittal sections) | `rq4*` |
| `Zhuang-ABCA-4.zip` | 0.05 GB | Zhuang-ABCA-4 (3 sections) | none |

These sections come from the Allen Brain Cell Atlas (Yao et al., *Nature* 2023; Zhang
et al., *Nature* 2023); please cite those papers if you use them.

### Input format

Each section is an `.h5ad` file (`.slaf` is also supported) with:

- `obs["x_ccf"]`, `obs["y_ccf"]`, `obs["z_ccf"]`: CCF coordinates in mm, which MIMYR
  stacks into `obsm["aligned_spatial"]`
- `obs["cluster"]`: cell-type labels (choose a different column with `--data_label`).
  Labels are mapped to tokens using the bundled Allen Brain Cell atlas vocabulary,
  so they must use ABC-atlas cluster names (e.g. `0001 CLA-EPd-CTX Car3 Glut_1`).
  The Python API fills in `class`, `subclass` and `supertype` from the cluster
  when they are missing.
- `X`: the gene expression matrix, with gene symbols as `var_names`.
  If any value is above 10, the matrix is taken to be counts and log1p-transformed
  (it is not normalized first).
- optionally `obs["technology"]` (`M550` for MERSCOPE 550-gene data, `M1100` for
  the Zhuang MERFISH data; default `M550`), `obs["species"]` and
  `obs["disease_state"]` (default `mouse`, `healthy`)

---

## Running inference

`main.py` is the entry point. It loads a dataset according to `--data_mode`,
generates each test section from its neighboring reference sections, scores the
result against the real section, and writes the metrics to a CSV.

Each of the three stages is controlled by its own flag:

| Flag | Default | Values |
|---|---|---|
| `--location_inference_type` | `model` | `model` (DDPM), `closest_slice` (copy the nearest reference section) |
| `--cluster_inference_type` | `model` | `model` (classifier), `majority_baseline` |
| `--expression_inference_type` | `model` | `model` (transformer), `lookup` (nearest same-type reference cell), `averaging` |

Every stage also accepts two special values:

- `skip` uses the **ground truth** for that stage instead of predicting it
  (useful for evaluating one stage in isolation).
- `end` stops the pipeline after the previous stage.

The defaults run every stage with its model. To run the full pipeline on the rq1
benchmark instead of the demo:

```bash
python main.py --data_mode rq1 --out_csv results/rq1.csv
```

To evaluate the expression model alone, given the true locations and cell types:

```bash
python main.py --location_inference_type skip --cluster_inference_type skip
```

The baselines used in the paper:

```bash
python main.py --data_mode rq1 \
  --location_inference_type closest_slice \
  --cluster_inference_type majority_baseline \
  --expression_inference_type lookup \
  --metrics soft_spearman_correlation,soft_f1 --metric_sampling 100
```

To reproduce the paper's results from the published models and data, use the
scripts in [`reproducibility/`](https://github.com/gkrieg/mimyr/blob/main/reproducibility/README.md), one per figure or
supplementary analysis.

### Data modes

`--data_mode` selects the dataset and the train/validation/test/reference split:

| Mode | Dataset | Paths it reads |
|---|---|---|
| `demo` (default) | one rq1 test section and its two neighbors, subsampled to 5% | `--demo_data_dir` (default `demo_data/`) |
| `rq1` | coronal sections from one brain, fixed held-out sections | `--data_dir/subclass_z1_d338_0_rotated/` |
| `rq3`, `rq3_v2`, `rq3_v2_d1` … `rq3_v2_d4` | Zhuang-ABCA-2 (coronal); the `_d*` modes vary the distance to the reference sections | `--zhuang_data_dir/Zhuang-ABCA-2/`, `--data_dir/ccf-ABCA-2.csv` |
| `rq4`, `rq4_rq3`, `rq4_noref` | Zhuang-ABCA-3 (sagittal); `rq4_noref` uses all of Zhuang-ABCA-1 as the reference pool | `--zhuang_data_dir/Zhuang-ABCA-3/` (and `-1`, `-2`), `--data_dir/ccf-ABCA-*.csv` |
| `rq5` | diseased (Trem2 / 5xFAD) sections; not distributed, so supply your own | `--diseased_data_dir` |
| `base_zhuang1` | Zhuang-ABCA-1, for training the base models | `--zhuang_data_dir/Zhuang-ABCA-1/` |

### Main options

- `--pretrained`: which pretrained model to use (see above)
- `--location_model_checkpoint`, `--cluster_model_checkpoint`, `--expression_model_checkpoint`:
  use these checkpoint files instead of the `--pretrained` model's
- `--guidance_signal`: strength of KDE backward guidance for the location model
  (default `0.0`, off); `--kde_bandwidth` sets the KDE bandwidth (default `0.01`)
- `--metrics`: comma-separated metrics (default
  `soft_accuracy,soft_spearman_correlation,soft_f1`). Available:
  `soft_accuracy`, `soft_spearman_correlation`, `soft_correlation_top_100`, `soft_f1`,
  `gridized_l1_distance`, `gridized_kl_divergence`, `celltype_abundance`,
  `composition_entropy`, `composition_pairwise`, `delauney_colocalization`.
  Results go in one column per radius, e.g. `soft_f1_radius@0.05`
- `--metric_sampling`: percentage of ground-truth cells the metrics are computed
  on (default `100`)
- `--metric_gene_set_file`: restrict expression metrics to the genes listed in a
  file, one per line (e.g. `gene_overlap.txt`); `--metric_filter_by_gt` restricts
  them to genes expressed in the ground truth
- `--eval_workers`: number of sections scored in parallel (default `1`)
- `--device`: `cuda` (default) or `cpu`

Expression decoding can be tuned with `--top_k`,
`--expression_select_max_index_margin`, and the weak teacher forcing options
(`--expression_weak_teacher_forcing`, `--wtf_*`), which pull generation toward the
nearest same-type reference cell. The `--pretrained` model sets the guidance
strength; see `python main.py --help` for the full list.

### Output

- `--out_csv` (default `results/output.csv`): one row per test section, containing
  the full configuration and the metrics. Rows are appended, and a section is
  **skipped** if a row with the same configuration is already in the CSV.
- `--artifact_dir` (default `artifacts/`): one timestamped folder per test section
  with `config.json`, `results.json`, `pred.pkl` (the generated AnnData), and plots
  of the predicted and ground-truth locations and cell types.

---

## Training

### Expression model

`--run_mode train` trains the expression transformer with distributed
data-parallel training. Launch it with `torchrun`:

```bash
torchrun --nproc_per_node=<NUM_GPUS> main.py \
  --run_mode train \
  --skip_combined_fit \
  --data_mode rq1 \
  --expression_output_dir /path/to/save/checkpoints \
  --expression_epochs 265 \
  --expression_batch_size 128
```

This fine-tunes the `--pretrained` model's expression transformer (by default
`merscope`, saved at epoch 165) for 100 more epochs. To try training on the demo data,
run `torchrun --nproc_per_node=1 main.py --run_mode train --skip_combined_fit --expression_epochs 166`.

- Training and validation data come from the `--data_mode` split, the same as for
  inference. The preprocessed data is cached in `--expression_output_dir`, so later
  runs with the same output directory skip preprocessing.
- A checkpoint is saved every `--expression_save_frequency` epochs (default `10`).
- Training starts from the weights in `--expression_model_checkpoint` and
  continues its epoch count from the `epochN` in the checkpoint's filename.
  `--expression_epochs` is the **last epoch number**, not a count: starting from the
  `merscope` model (epoch 165), `--expression_epochs 265` trains epochs 166 to 265.
- `--expression_random_init` uses only the checkpoint's architecture and starts from
  random weights. To train from scratch without a checkpoint, pass
  `--expression_model_checkpoint ""` together with
  `--expression_model_size small|medium|large`.
- Training logs to Weights & Biases. Set `WANDB_MODE=offline` to log locally.

Other useful options: `--expression_lr` (default `5e-5`), `--expression_lambda_val`
(weight of the regression loss, default `1.0`), `--expression_max_len` (default
`512`), `--expression_epoch_samples` (cells per epoch; `-1`, the default, uses all of
them), `--expression_xyz_noise` / `--expression_xyz_noise_magnitude` (coordinate
noise), and `--expression_adata2` (an extra `.h5ad`, e.g. scRNA-seq, merged into
training).

### Cell-type and location models

Without `--skip_combined_fit`, `--run_mode train` first trains the cell-type
classifier on the `.h5ad`/`.slaf` files in `--training_slice_directory` and
`--val_slice_directory` (default `data/cleaned_versions`). These files also need an
`obs["token"]` column. The best model is saved to `./best_model.pt`.

`main.py` does not train the location model (the Python API does; see above).
`train_location_zhuang1.py` and
`train_celltype_zhuang1.py` retrain the location and cell-type models on
Zhuang-ABCA-1; they currently contain cluster-specific data paths.

### Configuration files

Any option can be put in a YAML file and passed with `--config`. Each key is an
argument name (for example `data_dir:` or `expression_model_checkpoint:`), and flags
given on the command line override the file:

```bash
python main.py --config config.yaml
torchrun --nproc_per_node=4 main.py --config configs/train_base.yaml --run_mode train
```

`config.yaml` and the files in `configs/` hold settings from the lab's experiments,
and [`reproducibility/training_configs/`](https://github.com/gkrieg/mimyr/blob/main/reproducibility/README.md#training-configurations)
the ones that trained the published checkpoints. They contain absolute paths on the
lab's cluster, so copy one and change the paths before using it. Because they set
their checkpoints and guidance explicitly, they are not affected by `--pretrained`.

---

## Project structure

```
MIMYR/
├── pyproject.toml               # Package definition (pip install -e .)
├── main.py                      # Entry point: inference and training
├── demo_data/                   # Small demo dataset (default data_mode)
├── hub/                         # Hugging Face model and dataset cards
├── scripts/                     # Maintainer scripts for publishing to the Hub
├── config.yaml                  # Example configuration
├── configs/                     # Per-experiment configurations
├── reproducibility/             # Scripts that regenerate the paper's results
├── gene_overlap.txt             # Shared gene set for rq1/rq3 evaluation
├── train_location_zhuang1.py    # Retrain the location model on Zhuang-ABCA-1
├── train_celltype_zhuang1.py    # Retrain the cell-type model on Zhuang-ABCA-1
└── src/mimyr/                   # The installable `mimyr` package
    ├── pretrained.py            # Pretrained models and data on the Hugging Face Hub
    ├── data_loader.py           # SliceDataLoader: data modes, splits, preprocessing
    ├── inference.py             # The three-stage generation pipeline
    ├── evaluator.py             # Runs the selected metrics
    ├── metrics.py               # Metric implementations
    ├── gene_exp_model.py        # Cell-type label to token mapping
    └── models/
        ├── diffusion_model.py        # Location model (DDPM)
        ├── celltype_model.py         # Cell-type classifier
        ├── biological_model.py       # KDE used for location guidance
        ├── combined_model.py         # Loads the location and cell-type models
        └── generative_transformer/   # Expression model
            ├── Mimyr.py              # Generation and inference helpers
            ├── finetune_mimyr.py     # Training loop (DDP)
            ├── data_util.py          # Data harmonization and loaders
            ├── metadata/             # Bundled gene and cell-type metadata
            ├── model/                # Transformer implementation
            ├── reference/            # Gene symbol reference list
            └── utils/                # Tokenizer
```

## Citation

If you use MIMYR, please cite:

> Deshpande A, Bei Z, Ma J, Krieger S. MIMYR: Generative modeling of missing tissue
> in spatial transcriptomics. *bioRxiv* [Preprint]. 2025 Nov 27:2025.11.24.690239.
> doi: [10.1101/2025.11.24.690239](https://doi.org/10.1101/2025.11.24.690239).
> PMID: 41394599; PMCID: PMC12697367.

```bibtex
@article{deshpande2025mimyr,
  title   = {{MIMYR}: Generative modeling of missing tissue in spatial transcriptomics},
  author  = {Deshpande, A. and Bei, Z. and Ma, J. and Krieger, S.},
  journal = {bioRxiv},
  year    = {2025},
  doi     = {10.1101/2025.11.24.690239},
  note    = {Preprint. PMID: 41394599; PMCID: PMC12697367}
}
```

## License

MIMYR is released under the [MIT License](https://github.com/gkrieg/mimyr/blob/main/LICENSE).
