Figs. 2–3 and S1: paired CITE-seq

Figs. 2–3 and S1: paired CITE-seq#

Open in Colab

The paired RNA + protein analysis of the Genome Research article (Figs. 2–3, Supplemental Fig. S1), run through UniVI’s public API: univi.datasets for the data, fitted preprocessors, UniVITrainer, and univi.evaluation. Model settings are those of the archived notebook UniVI_manuscript_GR-Figure__3__CITE_paired_biological_latent.ipynb (summarized in Supplemental Table S8). Preprocessing follows Supplemental Table S7 using UniVI’s preprocessors rather than the archived notebook’s own code, so results will be close to, but not identical with, the published numbers. For the published numbers themselves, run the archived notebook with UniVI 0.4.7.

Data: Hao et al. (2021) PBMC CITE-seq, 161,764 cells with RNA and 228 surface proteins.

import sys

if "google.colab" in sys.modules:
    %pip install -q "univi[tutorials]>=1.1" "pandas==2.2.3"

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import torch

import univi.datasets as uds
from univi import ModalityConfig, TrainingConfig, UniVIConfig, UniVIMultiModalVAE, UniVITrainer
from univi.evaluation import (compute_modality_mixing, cross_modal_predict, encode_adata, evaluate_alignment,
                              pearson_corr_per_feature)
from univi.preprocessing import ADTPreprocessor, RNAPreprocessor, split_by_label
from univi.utils.seed import set_seed
from univi.workflows import make_loader, stack_embeddings

device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")
set_seed(0)

C:\Users\ashfo\micromamba\envs\univi-release-050\Lib\site-packages\tqdm\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

N_EPOCHS = 3000        # archived notebook: 3000 with early stopping (patience 100)
N_HVG = 2000
MAX_TEST_CELLS = 20000  # subsample the (large) test set for metrics and plots; None = all

Data and split#

The hosted files carry the split used in the article in obs["split"]. Otherwise the split is rebuilt the same way: stratified by celltype.l3, at most 1,200 cells per type in the train/validation pool (80/10), all remaining cells in test.

data = uds.hao_citeseq_pbmc()
rna, adt = data["rna"], data["adt"]
if "split" in rna.obs:
    splits = {k: np.flatnonzero(rna.obs["split"].to_numpy() == k) for k in ("train", "val", "test")}
else:
    splits = split_by_label(rna.obs["celltype.l3"], train_fraction=0.8, val_fraction=0.1, max_per_label=1200, seed=42)
print({k: len(v) for k, v in splits.items()})

{'train': 37862, 'val': 4722, 'test': 119180}

hao_citeseq_pbmc_rna.h5ad: 100%|██████████| 749M/749M [04:08<00:00, 3.01MB/s] 
hao_citeseq_pbmc_adt.h5ad: 100%|██████████| 49.9M/49.9M [00:25<00:00, 1.94MB/s]

Preprocessing (fit on training cells)#

RNA: 2,000 Seurat-v3 HVGs, library-size normalization to 10⁴, log1p, z-scores. Protein: per-cell CLR, z-scores.

rna_prep = RNAPreprocessor(n_hvg=N_HVG, scale=True).fit(rna[splits["train"]])
adt_prep = ADTPreprocessor(scale=True).fit(adt[splits["train"]])
parts = {k: {"rna": rna_prep.transform(rna[i]), "adt": adt_prep.transform(adt[i])} for k, i in splits.items()}

Model and training#

cfg = UniVIConfig(
    latent_dim=30, beta=1.5, gamma=5.0, encoder_dropout=0.1, decoder_dropout=0.0,
    kl_anneal_start=0, kl_anneal_end=0, align_anneal_start=0, align_anneal_end=0,
    modalities=[
        ModalityConfig("rna", parts["train"]["rna"].n_vars, [512, 256, 128], [128, 256, 512], likelihood="gaussian"),
        ModalityConfig("adt", parts["train"]["adt"].n_vars, [128, 64], [64, 128], likelihood="gaussian"),
    ],
)
model = UniVIMultiModalVAE(cfg, loss_mode="v1", v1_recon="avg", normalize_v1_terms=True)
UniVITrainer(
    model, make_loader(parts["train"], batch_size=256, shuffle=True, drop_last=True),
    make_loader(parts["val"], batch_size=1024),
    TrainingConfig(n_epochs=N_EPOCHS, batch_size=256, lr=1e-3, weight_decay=1e-4, device=device,
                   early_stopping=True, patience=100, log_every=50),
).fit();

[2026-09-21 16:12:49,111] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-21 16:12:49,112] [UniVITrainer] [INFO]   n_epochs: 3000
[2026-09-21 16:12:49,113] [UniVITrainer] [INFO]   batch_size: 256
[2026-09-21 16:12:49,113] [UniVITrainer] [INFO]   lr: 0.001
[2026-09-21 16:12:49,114] [UniVITrainer] [INFO]   weight_decay: 0.0001
[2026-09-21 16:12:49,114] [UniVITrainer] [INFO]   device: 'cuda'
[2026-09-21 16:12:49,115] [UniVITrainer] [INFO]   log_every: 50
[2026-09-21 16:12:49,115] [UniVITrainer] [INFO]   grad_clip: None
[2026-09-21 16:12:49,115] [UniVITrainer] [INFO]   num_workers: 0
[2026-09-21 16:12:49,116] [UniVITrainer] [INFO]   seed: 0
[2026-09-21 16:12:49,116] [UniVITrainer] [INFO]   early_stopping: True
[2026-09-21 16:12:49,116] [UniVITrainer] [INFO]   patience: 100
[2026-09-21 16:12:49,117] [UniVITrainer] [INFO]   min_delta: 0.0
[2026-09-21 16:12:49,117] [UniVITrainer] [INFO]   best_epoch_warmup: 0
Training UniVI:   0%|          | 0/3000 [00:00<?, ?it/s][2026-09-21 16:12:54,371] [UniVITrainer] [INFO] [Epoch 001] Train loss=1060.4807 (beta=1.500, gamma=5.000)
[2026-09-21 16:12:54,599] [UniVITrainer] [INFO] [Epoch 001] Val loss=1014.8379 (beta=1.500, gamma=5.000)
Training UniVI:   0%|          | 0/3000 [00:05<?, ?it/s, beta=1.500, gamma=5.000, train_loss=1060.4807, val_loss=1014.8379][2026-09-21 16:12:54,614] [UniVITrainer] [INFO] [Epoch 001] New best val loss: 1014.8379
Training UniVI:   0%|          | 1/3000 [00:10<4:34:46,  5.50s/it, beta=1.500, gamma=5.000, train_loss=986.5674, val_loss=982.0678]  [2026-09-21 16:12:59,659] [UniVITrainer] [INFO] [Epoch 002] New best val loss: 982.0678
Training UniVI:   0%|          | 2/3000 [00:16<4:21:23,  5.23s/it, beta=1.500, gamma=5.000, train_loss=967.4124, val_loss=970.1073][2026-09-21 16:13:05,923] [UniVITrainer] [INFO] [Epoch 003] New best val loss: 970.1073
Training UniVI:   0%|          | 3/3000 [00:21<4:44:51,  5.70s/it, beta=1.500, gamma=5.000, train_loss=957.1987, val_loss=960.8242][2026-09-21 16:13:10,978] [UniVITrainer] [INFO] [Epoch 004] New best val loss: 960.8242
Training UniVI:   0%|          | 4/3000 [00:27<4:32:00,  5.45s/it, beta=1.500, gamma=5.000, train_loss=948.6549, val_loss=955.8263][2026-09-21 16:13:16,293] [UniVITrainer] [INFO] [Epoch 005] New best val loss: 955.8263
Training UniVI:   0%|          | 5/3000 [00:32<4:29:32,  5.40s/it, beta=1.500, gamma=5.000, train_loss=944.0098, val_loss=950.8116][2026-09-21 16:13:21,430] [UniVITrainer] [INFO] [Epoch 006] New best val loss: 950.8116
Training UniVI:   0%|          | 6/3000 [00:37<4:24:57,  5.31s/it, beta=1.500, gamma=5.000, train_loss=937.4438, val_loss=946.8219][2026-09-21 16:13:26,158] [UniVITrainer] [INFO] [Epoch 007] New best val loss: 946.8219
Training UniVI:   0%|          | 7/3000 [00:41<4:15:23,  5.12s/it, beta=1.500, gamma=5.000, train_loss=934.4944, val_loss=942.6856][2026-09-21 16:13:31,032] [UniVITrainer] [INFO] [Epoch 008] New best val loss: 942.6856
Training UniVI:   0%|          | 8/3000 [00:46<4:11:24,  5.04s/it, beta=1.500, gamma=5.000, train_loss=931.2137, val_loss=939.8450][2026-09-21 16:13:35,534] [UniVITrainer] [INFO] [Epoch 009] New best val loss: 939.8450
Training UniVI:   0%|          | 9/3000 [00:51<4:02:54,  4.87s/it, beta=1.500, gamma=5.000, train_loss=927.1744, val_loss=936.9945][2026-09-21 16:13:40,158] [UniVITrainer] [INFO] [Epoch 010] New best val loss: 936.9945
Training UniVI:   0%|          | 10/3000 [00:55<3:59:00,  4.80s/it, beta=1.500, gamma=5.000, train_loss=924.8322, val_loss=934.0199][2026-09-21 16:13:44,484] [UniVITrainer] [INFO] [Epoch 011] New best val loss: 934.0199
Training UniVI:   0%|          | 11/3000 [00:59<3:51:45,  4.65s/it, beta=1.500, gamma=5.000, train_loss=922.8757, val_loss=933.0563][2026-09-21 16:13:48,950] [UniVITrainer] [INFO] [Epoch 012] New best val loss: 933.0563
Training UniVI:   0%|          | 12/3000 [01:04<3:48:52,  4.60s/it, beta=1.500, gamma=5.000, train_loss=920.4907, val_loss=931.1245][2026-09-21 16:13:53,450] [UniVITrainer] [INFO] [Epoch 013] New best val loss: 931.1245
Training UniVI:   0%|          | 13/3000 [01:08<3:47:20,  4.57s/it, beta=1.500, gamma=5.000, train_loss=919.3428, val_loss=929.1396][2026-09-21 16:13:57,773] [UniVITrainer] [INFO] [Epoch 014] New best val loss: 929.1396
Training UniVI:   0%|          | 14/3000 [01:13<3:43:36,  4.49s/it, beta=1.500, gamma=5.000, train_loss=917.1256, val_loss=927.4090][2026-09-21 16:14:02,305] [UniVITrainer] [INFO] [Epoch 015] New best val loss: 927.4090
Training UniVI:   0%|          | 15/3000 [01:17<3:44:06,  4.50s/it, beta=1.500, gamma=5.000, train_loss=914.2662, val_loss=926.7390][2026-09-21 16:14:06,728] [UniVITrainer] [INFO] [Epoch 016] New best val loss: 926.7390
Training UniVI:   1%|          | 16/3000 [01:22<3:42:48,  4.48s/it, beta=1.500, gamma=5.000, train_loss=912.5910, val_loss=924.1765][2026-09-21 16:14:11,607] [UniVITrainer] [INFO] [Epoch 017] New best val loss: 924.1765
Training UniVI:   1%|          | 17/3000 [01:26<3:48:42,  4.60s/it, beta=1.500, gamma=5.000, train_loss=910.4579, val_loss=922.5154][2026-09-21 16:14:16,070] [UniVITrainer] [INFO] [Epoch 018] New best val loss: 922.5154
Training UniVI:   1%|          | 18/3000 [01:31<3:46:34,  4.56s/it, beta=1.500, gamma=5.000, train_loss=908.5837, val_loss=920.2901][2026-09-21 16:14:20,823] [UniVITrainer] [INFO] [Epoch 019] New best val loss: 920.2901
Training UniVI:   1%|          | 20/3000 [01:40<3:48:14,  4.60s/it, beta=1.500, gamma=5.000, train_loss=905.7431, val_loss=919.3707][2026-09-21 16:14:29,787] [UniVITrainer] [INFO] [Epoch 021] New best val loss: 919.3707
Training UniVI:   1%|          | 21/3000 [01:45<3:45:32,  4.54s/it, beta=1.500, gamma=5.000, train_loss=905.0395, val_loss=919.0836][2026-09-21 16:14:34,208] [UniVITrainer] [INFO] [Epoch 022] New best val loss: 919.0836
Training UniVI:   1%|          | 22/3000 [01:49<3:43:39,  4.51s/it, beta=1.500, gamma=5.000, train_loss=904.6275, val_loss=917.9142][2026-09-21 16:14:38,330] [UniVITrainer] [INFO] [Epoch 023] New best val loss: 917.9142
Training UniVI:   1%|          | 24/3000 [01:57<3:35:09,  4.34s/it, beta=1.500, gamma=5.000, train_loss=902.8010, val_loss=917.7271][2026-09-21 16:14:46,744] [UniVITrainer] [INFO] [Epoch 025] New best val loss: 917.7271
Training UniVI:   1%|          | 25/3000 [02:01<3:33:02,  4.30s/it, beta=1.500, gamma=5.000, train_loss=901.4458, val_loss=917.5232][2026-09-21 16:14:51,061] [UniVITrainer] [INFO] [Epoch 026] New best val loss: 917.5232
Training UniVI:   1%|          | 28/3000 [02:16<3:45:42,  4.56s/it, beta=1.500, gamma=5.000, train_loss=899.7483, val_loss=916.6047][2026-09-21 16:15:05,203] [UniVITrainer] [INFO] [Epoch 029] New best val loss: 916.6047
Training UniVI:   1%|          | 29/3000 [02:21<3:45:08,  4.55s/it, beta=1.500, gamma=5.000, train_loss=899.3710, val_loss=915.9121][2026-09-21 16:15:10,279] [UniVITrainer] [INFO] [Epoch 030] New best val loss: 915.9121
Training UniVI:   1%|          | 30/3000 [02:25<3:52:54,  4.71s/it, beta=1.500, gamma=5.000, train_loss=898.7943, val_loss=915.3100][2026-09-21 16:15:14,527] [UniVITrainer] [INFO] [Epoch 031] New best val loss: 915.3100
Training UniVI:   1%|          | 31/3000 [02:29<3:46:02,  4.57s/it, beta=1.500, gamma=5.000, train_loss=898.1455, val_loss=914.6885][2026-09-21 16:15:18,766] [UniVITrainer] [INFO] [Epoch 032] New best val loss: 914.6885
Training UniVI:   1%|          | 32/3000 [02:34<3:41:05,  4.47s/it, beta=1.500, gamma=5.000, train_loss=897.8339, val_loss=914.3098][2026-09-21 16:15:23,186] [UniVITrainer] [INFO] [Epoch 033] New best val loss: 914.3098
Training UniVI:   1%|          | 34/3000 [02:43<3:41:24,  4.48s/it, beta=1.500, gamma=5.000, train_loss=895.5031, val_loss=914.2601][2026-09-21 16:15:32,625] [UniVITrainer] [INFO] [Epoch 035] New best val loss: 914.2601
Training UniVI:   1%|          | 36/3000 [02:52<3:48:30,  4.63s/it, beta=1.500, gamma=5.000, train_loss=896.9494, val_loss=912.6159][2026-09-21 16:15:42,064] [UniVITrainer] [INFO] [Epoch 037] New best val loss: 912.6159
Training UniVI:   1%|▏         | 38/3000 [03:02<3:52:58,  4.72s/it, beta=1.500, gamma=5.000, train_loss=895.1117, val_loss=912.5730][2026-09-21 16:15:51,361] [UniVITrainer] [INFO] [Epoch 039] New best val loss: 912.5730
Training UniVI:   1%|▏         | 39/3000 [03:08<3:49:02,  4.64s/it, beta=1.500, gamma=5.000, train_loss=894.3145, val_loss=911.3434][2026-09-21 16:15:57,165] [UniVITrainer] [INFO] [Epoch 040] New best val loss: 911.3434
Training UniVI:   2%|▏         | 45/3000 [03:36<3:55:37,  4.78s/it, beta=1.500, gamma=5.000, train_loss=890.4429, val_loss=911.2940][2026-09-21 16:16:25,785] [UniVITrainer] [INFO] [Epoch 046] New best val loss: 911.2940
Training UniVI:   2%|▏         | 47/3000 [03:45<3:51:03,  4.69s/it, beta=1.500, gamma=5.000, train_loss=891.3024, val_loss=910.9213][2026-09-21 16:16:34,905] [UniVITrainer] [INFO] [Epoch 048] New best val loss: 910.9213
Training UniVI:   2%|▏         | 49/3000 [03:50<3:45:58,  4.59s/it, beta=1.500, gamma=5.000, train_loss=889.6571, val_loss=911.3861][2026-09-21 16:16:43,455] [UniVITrainer] [INFO] [Epoch 050] Train loss=889.5407 (beta=1.500, gamma=5.000)
[2026-09-21 16:16:43,665] [UniVITrainer] [INFO] [Epoch 050] Val loss=910.7629 (beta=1.500, gamma=5.000)
Training UniVI:   2%|▏         | 49/3000 [03:54<3:45:58,  4.59s/it, beta=1.500, gamma=5.000, train_loss=889.5407, val_loss=910.7629][2026-09-21 16:16:43,678] [UniVITrainer] [INFO] [Epoch 050] New best val loss: 910.7629
Training UniVI:   2%|▏         | 50/3000 [03:59<3:41:51,  4.51s/it, beta=1.500, gamma=5.000, train_loss=889.0186, val_loss=910.0779][2026-09-21 16:16:48,261] [UniVITrainer] [INFO] [Epoch 051] New best val loss: 910.0779
Training UniVI:   2%|▏         | 51/3000 [04:03<3:42:50,  4.53s/it, beta=1.500, gamma=5.000, train_loss=888.1789, val_loss=908.9143][2026-09-21 16:16:52,636] [UniVITrainer] [INFO] [Epoch 052] New best val loss: 908.9143
Training UniVI:   2%|▏         | 55/3000 [04:21<3:43:26,  4.55s/it, beta=1.500, gamma=5.000, train_loss=887.5780, val_loss=908.5998][2026-09-21 16:17:10,619] [UniVITrainer] [INFO] [Epoch 056] New best val loss: 908.5998
Training UniVI:   2%|▏         | 59/3000 [04:38<3:35:37,  4.40s/it, beta=1.500, gamma=5.000, train_loss=886.1465, val_loss=908.3102][2026-09-21 16:17:27,877] [UniVITrainer] [INFO] [Epoch 060] New best val loss: 908.3102
Training UniVI:   2%|▏         | 62/3000 [04:52<3:37:24,  4.44s/it, beta=1.500, gamma=5.000, train_loss=884.7969, val_loss=907.8951][2026-09-21 16:17:41,400] [UniVITrainer] [INFO] [Epoch 063] New best val loss: 907.8951
Training UniVI:   2%|▏         | 65/3000 [05:05<3:41:16,  4.52s/it, beta=1.500, gamma=5.000, train_loss=884.5410, val_loss=907.5010][2026-09-21 16:17:55,100] [UniVITrainer] [INFO] [Epoch 066] New best val loss: 907.5010
Training UniVI:   2%|▎         | 75/3000 [05:49<3:31:27,  4.34s/it, beta=1.500, gamma=5.000, train_loss=882.1519, val_loss=906.6626][2026-09-21 16:18:39,014] [UniVITrainer] [INFO] [Epoch 076] New best val loss: 906.6626
Training UniVI:   3%|▎         | 80/3000 [06:11<3:30:28,  4.32s/it, beta=1.500, gamma=5.000, train_loss=880.6615, val_loss=906.6297][2026-09-21 16:19:00,873] [UniVITrainer] [INFO] [Epoch 081] New best val loss: 906.6297
Training UniVI:   3%|▎         | 82/3000 [06:20<3:33:54,  4.40s/it, beta=1.500, gamma=5.000, train_loss=880.4894, val_loss=906.1990][2026-09-21 16:19:09,927] [UniVITrainer] [INFO] [Epoch 083] New best val loss: 906.1990
Training UniVI:   3%|▎         | 90/3000 [06:55<3:26:48,  4.26s/it, beta=1.500, gamma=5.000, train_loss=879.3212, val_loss=905.1755][2026-09-21 16:19:44,133] [UniVITrainer] [INFO] [Epoch 091] New best val loss: 905.1755
Training UniVI:   3%|▎         | 94/3000 [07:12<3:29:29,  4.33s/it, beta=1.500, gamma=5.000, train_loss=878.5813, val_loss=905.1363][2026-09-21 16:20:01,691] [UniVITrainer] [INFO] [Epoch 095] New best val loss: 905.1363
Training UniVI:   3%|▎         | 99/3000 [07:29<3:26:22,  4.27s/it, beta=1.500, gamma=5.000, train_loss=877.0275, val_loss=907.0146][2026-09-21 16:20:22,650] [UniVITrainer] [INFO] [Epoch 100] Train loss=876.4508 (beta=1.500, gamma=5.000)
[2026-09-21 16:20:22,845] [UniVITrainer] [INFO] [Epoch 100] Val loss=906.7888 (beta=1.500, gamma=5.000)
Training UniVI:   5%|▍         | 149/3000 [11:00<3:19:11,  4.19s/it, beta=1.500, gamma=5.000, train_loss=870.4820, val_loss=907.7504][2026-09-21 16:23:53,646] [UniVITrainer] [INFO] [Epoch 150] Train loss=870.5482 (beta=1.500, gamma=5.000)
[2026-09-21 16:23:53,827] [UniVITrainer] [INFO] [Epoch 150] Val loss=906.8848 (beta=1.500, gamma=5.000)
Training UniVI:   5%|▌         | 158/3000 [11:42<3:18:05,  4.18s/it, beta=1.500, gamma=5.000, train_loss=868.7815, val_loss=904.2874][2026-09-21 16:24:31,162] [UniVITrainer] [INFO] [Epoch 159] New best val loss: 904.2874
Training UniVI:   7%|▋         | 199/3000 [14:14<3:00:27,  3.87s/it, beta=1.500, gamma=5.000, train_loss=865.4438, val_loss=908.4448][2026-09-21 16:27:07,534] [UniVITrainer] [INFO] [Epoch 200] Train loss=864.2717 (beta=1.500, gamma=5.000)
[2026-09-21 16:27:07,720] [UniVITrainer] [INFO] [Epoch 200] Val loss=908.0589 (beta=1.500, gamma=5.000)
Training UniVI:   8%|▊         | 249/3000 [17:23<2:51:00,  3.73s/it, beta=1.500, gamma=5.000, train_loss=860.9412, val_loss=909.6887][2026-09-21 16:30:15,696] [UniVITrainer] [INFO] [Epoch 250] Train loss=860.0894 (beta=1.500, gamma=5.000)
[2026-09-21 16:30:15,884] [UniVITrainer] [INFO] [Epoch 250] Val loss=908.0677 (beta=1.500, gamma=5.000)
Training UniVI:   9%|▊         | 258/3000 [18:00<2:50:03,  3.72s/it, beta=1.500, gamma=5.000, train_loss=859.5278, val_loss=908.0736][2026-09-21 16:30:49,419] [UniVITrainer] [INFO] Early stopping at epoch 259 (best val loss=904.2874, best epoch=159)
Training UniVI:   9%|▊         | 258/3000 [18:00<3:11:21,  4.19s/it, beta=1.500, gamma=5.000, train_loss=859.5278, val_loss=908.0736]
[2026-09-21 16:30:49,428] [UniVITrainer] [INFO] Restored best model from epoch 159 (val loss=904.2874)

Latent alignment on held-out cells (Fig. 2)#

test = parts["test"]
if MAX_TEST_CELLS is not None and test["rna"].n_obs > MAX_TEST_CELLS:
    keep = np.sort(np.random.default_rng(0).choice(test["rna"].n_obs, MAX_TEST_CELLS, replace=False))
    test = {m: a[keep].copy() for m, a in test.items()}
z_rna = encode_adata(model, test["rna"], modality="rna", device=device, latent="modality_mean")
z_adt = encode_adata(model, test["adt"], modality="adt", device=device, latent="modality_mean")

rows = {}
for level in ["celltype.l1", "celltype.l2"]:
    lab = test["rna"].obs[level].astype(str).to_numpy()
    m = evaluate_alignment(Z1=z_rna, Z2=z_adt, labels_source=lab, labels_target=lab, recall_ks=(1, 10, 50))
    rows[level] = {"FOSCTTM": m["foscttm_mean"], "Recall@1": m["recall_at_k"]["1"]["mean"],
                   "Recall@10": m["recall_at_k"]["10"]["mean"], "Recall@50": m["recall_at_k"]["50"]["mean"],
                   "label transfer acc (RNA to ADT)": m["label_transfer_acc"],
                   "macro-F1, worse direction": m["worst_direction_macro_f1"]}
mixing = compute_modality_mixing(np.vstack([z_rna, z_adt]), np.repeat(["rna", "adt"], len(z_rna)), k=20)
print(f"modality mixing (k=20): {mixing:.3f}")
pd.DataFrame(rows).round(3)

modality mixing (k=20): 0.319
celltype.l1 celltype.l2
FOSCTTM 0.025 0.025
Recall@1 0.011 0.011
Recall@10 0.063 0.063
Recall@50 0.193 0.193
label transfer acc (RNA to ADT) 0.993 0.951
macro-F1, worse direction 0.971 0.571 
joint = stack_embeddings(model, [("test", "rna", test["rna"]), ("test", "adt", test["adt"])], device=device)
sc.pp.neighbors(joint, use_rep="X_univi", n_neighbors=30)
sc.tl.umap(joint, random_state=0)
sc.pl.umap(joint, color=["modality", "celltype.l1", "celltype.l2"], wspace=0.5, legend_fontsize=6)
../../_images/5cd6db95c4d39eb868d9ad7f3377898e0ce67d475fc6ab715f71b3c9385d6cf5.png

Cross-modal reconstruction of markers (Fig. 3, Supplemental Fig. S1)#

Protein predicted from RNA and RNA predicted from protein, shown on the reference UMAP of the RNA embedding next to the observed values.

rna_t, adt_t = test["rna"], test["adt"]
adt_t.layers["from_rna"] = cross_modal_predict(model, rna_t, src_mod="rna", tgt_mod="adt", device=device)
rna_t.layers["from_adt"] = cross_modal_predict(model, adt_t, src_mod="adt", tgt_mod="rna", device=device)
rna_t.obsm["X_univi"] = z_rna
sc.pp.neighbors(rna_t, use_rep="X_univi", n_neighbors=30)
sc.tl.umap(rna_t, random_state=0)
adt_t.obsm["X_umap"] = rna_t.obsm["X_umap"]

genes = [g for g in ["CD79A", "MS4A1", "LYZ", "FCGR3A", "NKG7", "GNLY", "TRAC", "CD3D"] if g in rna_t.var_names]
proteins = [p for p in ["CD19", "CD20", "CD14", "CD16", "CD56-1", "CD3-1", "CD4", "CD8a"] if p in adt_t.var_names]
for adata, feats, layer, label in [(rna_t, genes, "from_adt", "ADT→RNA"), (adt_t, proteins, "from_rna", "RNA→ADT")]:
    sc.pl.umap(adata, color=feats, ncols=4, vmax="p99", cmap="viridis", title=[f"{f} observed" for f in feats])
    sc.pl.umap(adata, color=feats, ncols=4, layer=layer, vmax="p99", cmap="viridis",
               title=[f"{f} {label}" for f in feats])
../../_images/bca354ab1878ccc197154b3c9eb7f6e077fe10ceeee540717b1922c0f6242c9b.png ../../_images/351ca03e80aaeb68bcba2cec7f7ca02f9bb448bd2eae4ae1d5062a2948a8e296.png ../../_images/8510082a423e17891f2f18fef6e312788a5f08f840b1874dc3cfe91d918e917a.png ../../_images/3f85dad5020fd94a8f898a512788a41f17f51e42b18dc30a4879abe2eac5a67f.png

Mean marker values per celltype.l2 subtype, observed and predicted (Supplemental Fig. S1L–M):

def mean_by_group(adata, feats, layer=None):
    X = adata[:, feats].layers[layer] if layer else adata[:, feats].X
    X = X.toarray() if hasattr(X, "toarray") else np.asarray(X)
    return pd.DataFrame(X, columns=feats).groupby(adata.obs["celltype.l2"].astype(str).to_numpy()).mean()

for adata, feats, layer, title in [(rna_t, genes, "from_adt", "RNA markers"), (adt_t, proteins, "from_rna", "ADT markers")]:
    obs_m, pred_m = mean_by_group(adata, feats), mean_by_group(adata, feats, layer)
    fig, axes = plt.subplots(1, 2, figsize=(9, 0.22 * len(obs_m) + 1.5), sharey=True)
    for ax, table, name in zip(axes, [obs_m, pred_m], ["observed", "cross-modal prediction"]):
        ax.imshow(table.values, aspect="auto", cmap="viridis")
        ax.set_xticks(range(len(feats)), feats, rotation=90)
        ax.set_yticks(range(len(table)), table.index, fontsize=7)
        ax.set_title(f"{title}: {name}")
    plt.tight_layout()
    plt.show()
../../_images/fada5c7f7c9d76d3d80899ffb0bcc09f25bb7605579b1f363b09f2486d431acd.png ../../_images/3803fb4f98bef7ff3b55d89749ee383206a00c66e26797f7f256a0910a19747d.png 
summary = {
    "ADT from RNA, median per-protein r": np.nanmedian(pearson_corr_per_feature(np.asarray(adt_t.X), adt_t.layers["from_rna"])),
    "RNA from ADT, median per-gene r": np.nanmedian(pearson_corr_per_feature(np.asarray(rna_t.X), rna_t.layers["from_adt"])),
}
pd.Series(summary).round(3)

ADT from RNA, median per-protein r    0.531
RNA from ADT, median per-gene r       0.301
dtype: float32
Contents