Fig. 6 and S6: trimodal TEA-seq with a held-out well#
RNA, surface protein, and chromatin accessibility measured in the same cells (Swanson et al. 2021). UniVI is trained on wells 3, 4 and 6 and evaluated on well 5, with every preprocessing step fit on the training wells. Model settings follow the archived notebook UniVI_manuscript_GR-Figure__6__TEA-seq_tri-modal.ipynb; preprocessing uses UniVI’s fitted preprocessors, so results will be close to, not identical with, the published ones.
import sys
if "google.colab" in sys.modules:
%pip install -q "univi[tutorials]>=1.1" "pandas==2.2.3"
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_foscttm, encode_adata, encode_fused_adata_pair
from univi.preprocessing import ADTPreprocessor, ATACPreprocessor, RNAPreprocessor
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 200)
N_HVG = 2000
N_LSI = 100
HOLDOUT = "GSM5123953_X066-MP0C1W5_leukopak_perm-cells_tea" # well 5
LEIDEN_RESOLUTION = 0.8
Data and well hold-out#
The ATAC modality is a genome-wide tile matrix (500 bp tiles). Tiles open in 0.5–80% of training cells are kept. A random 10% of training-well cells is used for early stopping.
data = uds.load("teaseq_pbmc")
rna, adt, atac = data["rna"], data["adt"], data["atac"]
well = rna.obs["sample_id"].astype(str).to_numpy()
held = np.flatnonzero(well == HOLDOUT)
rest = np.random.default_rng(0).permutation(np.flatnonzero(well != HOLDOUT))
splits = {"train": np.sort(rest[len(rest) // 10:]), "val": np.sort(rest[: len(rest) // 10]), "test": held}
print({k: len(v) for k, v in splits.items()})
rna_prep = RNAPreprocessor(n_hvg=N_HVG, scale=True).fit(rna[splits["train"]])
adt_prep = ADTPreprocessor(scale=True).fit(adt[splits["train"]])
atac_prep = ATACPreprocessor(n_components=N_LSI, drop_first=False, scale=True,
min_fraction=0.005, max_fraction=0.80).fit(atac[splits["train"]])
parts = {k: {"rna": rna_prep.transform(rna[i]), "adt": adt_prep.transform(adt[i]), "atac": atac_prep.transform(atac[i])}
for k, i in splits.items()}
{'train': 20226, 'val': 2247, 'test': 7421}
teaseq_pbmc_rna.h5ad: 100%|██████████| 76.5M/76.5M [00:46<00:00, 1.64MB/s]
teaseq_pbmc_adt.h5ad: 100%|██████████| 4.67M/4.67M [00:06<00:00, 682kB/s]
teaseq_pbmc_atac.h5ad: 100%|██████████| 339M/339M [03:05<00:00, 1.82MB/s]
cfg = UniVIConfig(
latent_dim=30, beta=1.15,
gamma=1.45, # value in the archived notebook; Supplemental Table S8 lists 1.35
encoder_dropout=0.10, decoder_dropout=0.0,
kl_anneal_start=0, kl_anneal_end=50, align_anneal_start=25, align_anneal_end=75,
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"),
ModalityConfig("atac", parts["train"]["atac"].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-4, weight_decay=1e-4, device=device,
early_stopping=True, patience=200, best_epoch_warmup=75, log_every=100),
).fit();
[2026-09-21 16:17:05,767] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-21 16:17:05,768] [UniVITrainer] [INFO] n_epochs: 3000
[2026-09-21 16:17:05,768] [UniVITrainer] [INFO] batch_size: 256
[2026-09-21 16:17:05,768] [UniVITrainer] [INFO] lr: 0.0001
[2026-09-21 16:17:05,769] [UniVITrainer] [INFO] weight_decay: 0.0001
[2026-09-21 16:17:05,769] [UniVITrainer] [INFO] device: 'cuda'
[2026-09-21 16:17:05,770] [UniVITrainer] [INFO] log_every: 100
[2026-09-21 16:17:05,770] [UniVITrainer] [INFO] grad_clip: None
[2026-09-21 16:17:05,770] [UniVITrainer] [INFO] num_workers: 0
[2026-09-21 16:17:05,771] [UniVITrainer] [INFO] seed: 0
[2026-09-21 16:17:05,771] [UniVITrainer] [INFO] early_stopping: True
[2026-09-21 16:17:05,772] [UniVITrainer] [INFO] patience: 200
[2026-09-21 16:17:05,772] [UniVITrainer] [INFO] min_delta: 0.0
[2026-09-21 16:17:05,773] [UniVITrainer] [INFO] best_epoch_warmup: 75
Training UniVI: 0%| | 0/3000 [00:00<?, ?it/s][2026-09-21 16:17:09,793] [UniVITrainer] [INFO] [Epoch 001] Train loss=716.3032 (beta=0.023, gamma=0.000)
[2026-09-21 16:17:09,949] [UniVITrainer] [INFO] [Epoch 001] Val loss=715.1627 (beta=1.150, gamma=1.450)
Training UniVI: 2%|▏ | 73/3000 [04:42<3:01:05, 3.71s/it, beta=1.150, gamma=1.421, train_loss=671.0216, val_loss=646.0187][2026-09-21 16:21:47,857] [UniVITrainer] [INFO] [Epoch 074] Best tracking warmup ends next epoch (best_epoch_warmup=75).
Training UniVI: 2%|▏ | 74/3000 [04:45<3:00:34, 3.70s/it, beta=1.150, gamma=1.450, train_loss=671.1755, val_loss=646.3227][2026-09-21 16:21:51,573] [UniVITrainer] [INFO] [Epoch 075] New best val loss: 646.3227
Training UniVI: 3%|▎ | 76/3000 [04:53<3:00:02, 3.69s/it, beta=1.150, gamma=1.450, train_loss=670.8795, val_loss=646.1594][2026-09-21 16:21:58,955] [UniVITrainer] [INFO] [Epoch 077] New best val loss: 646.1594
Training UniVI: 3%|▎ | 78/3000 [05:00<2:59:41, 3.69s/it, beta=1.150, gamma=1.450, train_loss=670.6944, val_loss=646.0969][2026-09-21 16:22:06,365] [UniVITrainer] [INFO] [Epoch 079] New best val loss: 646.0969
Training UniVI: 3%|▎ | 80/3000 [05:07<2:58:56, 3.68s/it, beta=1.150, gamma=1.450, train_loss=670.5961, val_loss=646.0242][2026-09-21 16:22:13,661] [UniVITrainer] [INFO] [Epoch 081] New best val loss: 646.0242
Training UniVI: 3%|▎ | 81/3000 [05:11<2:59:03, 3.68s/it, beta=1.150, gamma=1.450, train_loss=670.3775, val_loss=645.8127][2026-09-21 16:22:17,394] [UniVITrainer] [INFO] [Epoch 082] New best val loss: 645.8127
Training UniVI: 3%|▎ | 86/3000 [05:30<2:58:55, 3.68s/it, beta=1.150, gamma=1.450, train_loss=670.1635, val_loss=645.6289][2026-09-21 16:22:35,852] [UniVITrainer] [INFO] [Epoch 087] New best val loss: 645.6289
Training UniVI: 3%|▎ | 87/3000 [05:33<2:58:37, 3.68s/it, beta=1.150, gamma=1.450, train_loss=670.1406, val_loss=645.5928][2026-09-21 16:22:39,580] [UniVITrainer] [INFO] [Epoch 088] New best val loss: 645.5928
Training UniVI: 3%|▎ | 92/3000 [05:52<2:58:30, 3.68s/it, beta=1.150, gamma=1.450, train_loss=669.5410, val_loss=645.5178][2026-09-21 16:22:57,989] [UniVITrainer] [INFO] [Epoch 093] New best val loss: 645.5178
Training UniVI: 3%|▎ | 93/3000 [05:55<2:58:34, 3.69s/it, beta=1.150, gamma=1.450, train_loss=669.7000, val_loss=645.4531][2026-09-21 16:23:01,731] [UniVITrainer] [INFO] [Epoch 094] New best val loss: 645.4531
Training UniVI: 3%|▎ | 95/3000 [06:03<2:59:33, 3.71s/it, beta=1.150, gamma=1.450, train_loss=669.7251, val_loss=645.3406][2026-09-21 16:23:09,104] [UniVITrainer] [INFO] [Epoch 096] New best val loss: 645.3406
Training UniVI: 3%|▎ | 99/3000 [06:14<2:59:04, 3.70s/it, beta=1.150, gamma=1.450, train_loss=669.4687, val_loss=645.6017][2026-09-21 16:23:23,799] [UniVITrainer] [INFO] [Epoch 100] Train loss=669.4638 (beta=1.150, gamma=1.450)
[2026-09-21 16:23:23,926] [UniVITrainer] [INFO] [Epoch 100] Val loss=645.4751 (beta=1.150, gamma=1.450)
Training UniVI: 3%|▎ | 100/3000 [06:21<2:58:55, 3.70s/it, beta=1.150, gamma=1.450, train_loss=669.3533, val_loss=645.3172][2026-09-21 16:23:27,622] [UniVITrainer] [INFO] [Epoch 101] New best val loss: 645.3172
Training UniVI: 3%|▎ | 101/3000 [06:25<2:58:45, 3.70s/it, beta=1.150, gamma=1.450, train_loss=669.3846, val_loss=645.2324][2026-09-21 16:23:31,398] [UniVITrainer] [INFO] [Epoch 102] New best val loss: 645.2324
Training UniVI: 4%|▎ | 106/3000 [06:43<2:56:59, 3.67s/it, beta=1.150, gamma=1.450, train_loss=669.1889, val_loss=645.1819][2026-09-21 16:23:49,712] [UniVITrainer] [INFO] [Epoch 107] New best val loss: 645.1819
Training UniVI: 4%|▎ | 107/3000 [06:47<2:57:45, 3.69s/it, beta=1.150, gamma=1.450, train_loss=669.0001, val_loss=645.1483][2026-09-21 16:23:53,550] [UniVITrainer] [INFO] [Epoch 108] New best val loss: 645.1483
Training UniVI: 4%|▎ | 109/3000 [06:55<2:58:41, 3.71s/it, beta=1.150, gamma=1.450, train_loss=669.0559, val_loss=645.1373][2026-09-21 16:24:00,989] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 645.1373
Training UniVI: 4%|▎ | 111/3000 [07:02<2:58:20, 3.70s/it, beta=1.150, gamma=1.450, train_loss=668.8252, val_loss=644.9153][2026-09-21 16:24:08,300] [UniVITrainer] [INFO] [Epoch 112] New best val loss: 644.9153
Training UniVI: 4%|▍ | 115/3000 [07:17<2:56:24, 3.67s/it, beta=1.150, gamma=1.450, train_loss=668.7074, val_loss=644.7350][2026-09-21 16:24:22,954] [UniVITrainer] [INFO] [Epoch 116] New best val loss: 644.7350
Training UniVI: 4%|▍ | 120/3000 [07:34<2:45:04, 3.44s/it, beta=1.150, gamma=1.450, train_loss=668.5257, val_loss=644.7209][2026-09-21 16:24:39,923] [UniVITrainer] [INFO] [Epoch 121] New best val loss: 644.7209
Training UniVI: 4%|▍ | 132/3000 [08:15<2:39:24, 3.33s/it, beta=1.150, gamma=1.450, train_loss=667.9547, val_loss=644.4647][2026-09-21 16:25:20,862] [UniVITrainer] [INFO] [Epoch 133] New best val loss: 644.4647
Training UniVI: 5%|▌ | 156/3000 [09:33<2:34:31, 3.26s/it, beta=1.150, gamma=1.450, train_loss=667.2689, val_loss=644.3986][2026-09-21 16:26:39,485] [UniVITrainer] [INFO] [Epoch 157] New best val loss: 644.3986
Training UniVI: 5%|▌ | 161/3000 [09:50<2:34:55, 3.27s/it, beta=1.150, gamma=1.450, train_loss=667.0680, val_loss=644.3930][2026-09-21 16:26:55,825] [UniVITrainer] [INFO] [Epoch 162] New best val loss: 644.3930
Training UniVI: 5%|▌ | 164/3000 [10:00<2:41:42, 3.42s/it, beta=1.150, gamma=1.450, train_loss=666.9480, val_loss=644.3839][2026-09-21 16:27:06,378] [UniVITrainer] [INFO] [Epoch 165] New best val loss: 644.3839
Training UniVI: 6%|▌ | 167/3000 [10:10<2:38:01, 3.35s/it, beta=1.150, gamma=1.450, train_loss=666.9287, val_loss=644.2520][2026-09-21 16:27:16,221] [UniVITrainer] [INFO] [Epoch 168] New best val loss: 644.2520
Training UniVI: 6%|▌ | 173/3000 [10:30<2:37:06, 3.33s/it, beta=1.150, gamma=1.450, train_loss=666.6891, val_loss=644.2466][2026-09-21 16:27:36,281] [UniVITrainer] [INFO] [Epoch 174] New best val loss: 644.2466
Training UniVI: 6%|▌ | 180/3000 [10:53<2:34:33, 3.29s/it, beta=1.150, gamma=1.450, train_loss=666.6179, val_loss=644.1902][2026-09-21 16:27:59,499] [UniVITrainer] [INFO] [Epoch 181] New best val loss: 644.1902
Training UniVI: 6%|▋ | 190/3000 [11:26<2:33:26, 3.28s/it, beta=1.150, gamma=1.450, train_loss=666.1468, val_loss=644.1394][2026-09-21 16:28:32,210] [UniVITrainer] [INFO] [Epoch 191] New best val loss: 644.1394
Training UniVI: 6%|▋ | 191/3000 [11:29<2:31:53, 3.24s/it, beta=1.150, gamma=1.450, train_loss=666.3388, val_loss=644.0842][2026-09-21 16:28:35,466] [UniVITrainer] [INFO] [Epoch 192] New best val loss: 644.0842
Training UniVI: 7%|▋ | 199/3000 [11:52<2:32:06, 3.26s/it, beta=1.150, gamma=1.450, train_loss=665.7939, val_loss=644.3318][2026-09-21 16:29:01,177] [UniVITrainer] [INFO] [Epoch 200] Train loss=665.8931 (beta=1.150, gamma=1.450)
[2026-09-21 16:29:01,313] [UniVITrainer] [INFO] [Epoch 200] Val loss=644.0813 (beta=1.150, gamma=1.450)
Training UniVI: 7%|▋ | 199/3000 [11:55<2:32:06, 3.26s/it, beta=1.150, gamma=1.450, train_loss=665.8931, val_loss=644.0813][2026-09-21 16:29:01,322] [UniVITrainer] [INFO] [Epoch 200] New best val loss: 644.0813
Training UniVI: 7%|▋ | 200/3000 [11:58<2:31:57, 3.26s/it, beta=1.150, gamma=1.450, train_loss=665.8760, val_loss=644.0324][2026-09-21 16:29:04,540] [UniVITrainer] [INFO] [Epoch 201] New best val loss: 644.0324
Training UniVI: 7%|▋ | 204/3000 [12:11<2:32:55, 3.28s/it, beta=1.150, gamma=1.450, train_loss=665.7469, val_loss=643.9356][2026-09-21 16:29:17,694] [UniVITrainer] [INFO] [Epoch 205] New best val loss: 643.9356
Training UniVI: 7%|▋ | 214/3000 [12:44<2:30:23, 3.24s/it, beta=1.150, gamma=1.450, train_loss=665.4751, val_loss=643.8322][2026-09-21 16:29:50,121] [UniVITrainer] [INFO] [Epoch 215] New best val loss: 643.8322
Training UniVI: 7%|▋ | 216/3000 [12:50<2:29:30, 3.22s/it, beta=1.150, gamma=1.450, train_loss=665.4741, val_loss=643.8173][2026-09-21 16:29:56,606] [UniVITrainer] [INFO] [Epoch 217] New best val loss: 643.8173
Training UniVI: 8%|▊ | 227/3000 [13:26<2:31:18, 3.27s/it, beta=1.150, gamma=1.450, train_loss=665.0639, val_loss=643.8075][2026-09-21 16:30:32,561] [UniVITrainer] [INFO] [Epoch 228] New best val loss: 643.8075
Training UniVI: 8%|▊ | 231/3000 [13:39<2:29:49, 3.25s/it, beta=1.150, gamma=1.450, train_loss=664.8827, val_loss=643.7928][2026-09-21 16:30:45,513] [UniVITrainer] [INFO] [Epoch 232] New best val loss: 643.7928
Training UniVI: 8%|▊ | 241/3000 [14:17<2:54:04, 3.79s/it, beta=1.150, gamma=1.450, train_loss=664.6049, val_loss=643.5679][2026-09-21 16:31:23,367] [UniVITrainer] [INFO] [Epoch 242] New best val loss: 643.5679
Training UniVI: 9%|▊ | 259/3000 [15:16<2:23:15, 3.14s/it, beta=1.150, gamma=1.450, train_loss=664.1696, val_loss=643.5619][2026-09-21 16:32:21,882] [UniVITrainer] [INFO] [Epoch 260] New best val loss: 643.5619
Training UniVI: 9%|▉ | 264/3000 [15:31<2:23:20, 3.14s/it, beta=1.150, gamma=1.450, train_loss=664.0607, val_loss=643.3774][2026-09-21 16:32:37,596] [UniVITrainer] [INFO] [Epoch 265] New best val loss: 643.3774
Training UniVI: 9%|▉ | 275/3000 [16:08<2:23:04, 3.15s/it, beta=1.150, gamma=1.450, train_loss=663.6574, val_loss=643.3462][2026-09-21 16:33:14,072] [UniVITrainer] [INFO] [Epoch 276] New best val loss: 643.3462
Training UniVI: 9%|▉ | 277/3000 [16:14<2:20:37, 3.10s/it, beta=1.150, gamma=1.450, train_loss=663.5388, val_loss=643.2396][2026-09-21 16:33:20,242] [UniVITrainer] [INFO] [Epoch 278] New best val loss: 643.2396
Training UniVI: 9%|▉ | 282/3000 [16:30<2:21:09, 3.12s/it, beta=1.150, gamma=1.450, train_loss=663.4960, val_loss=643.1107][2026-09-21 16:33:36,022] [UniVITrainer] [INFO] [Epoch 283] New best val loss: 643.1107
Training UniVI: 9%|▉ | 283/3000 [16:33<2:19:51, 3.09s/it, beta=1.150, gamma=1.450, train_loss=663.5777, val_loss=643.0530][2026-09-21 16:33:39,089] [UniVITrainer] [INFO] [Epoch 284] New best val loss: 643.0530
Training UniVI: 10%|▉ | 295/3000 [17:10<2:17:43, 3.05s/it, beta=1.150, gamma=1.450, train_loss=663.0957, val_loss=642.9470][2026-09-21 16:34:16,134] [UniVITrainer] [INFO] [Epoch 296] New best val loss: 642.9470
Training UniVI: 10%|▉ | 297/3000 [17:16<2:18:22, 3.07s/it, beta=1.150, gamma=1.450, train_loss=663.0351, val_loss=642.9253][2026-09-21 16:34:22,376] [UniVITrainer] [INFO] [Epoch 298] New best val loss: 642.9253
Training UniVI: 10%|▉ | 299/3000 [17:19<2:20:08, 3.11s/it, beta=1.150, gamma=1.450, train_loss=663.0094, val_loss=643.3835][2026-09-21 16:34:28,510] [UniVITrainer] [INFO] [Epoch 300] Train loss=662.9435 (beta=1.150, gamma=1.450)
[2026-09-21 16:34:28,654] [UniVITrainer] [INFO] [Epoch 300] Val loss=643.3714 (beta=1.150, gamma=1.450)
Training UniVI: 10%|█ | 305/3000 [17:41<2:21:11, 3.14s/it, beta=1.150, gamma=1.450, train_loss=662.7196, val_loss=642.7023][2026-09-21 16:34:47,480] [UniVITrainer] [INFO] [Epoch 306] New best val loss: 642.7023
Training UniVI: 11%|█ | 326/3000 [18:51<2:30:09, 3.37s/it, beta=1.150, gamma=1.450, train_loss=662.0349, val_loss=642.6286][2026-09-21 16:35:57,661] [UniVITrainer] [INFO] [Epoch 327] New best val loss: 642.6286
Training UniVI: 11%|█ | 329/3000 [19:02<2:37:25, 3.54s/it, beta=1.150, gamma=1.450, train_loss=661.9268, val_loss=642.3970][2026-09-21 16:36:08,438] [UniVITrainer] [INFO] [Epoch 330] New best val loss: 642.3970
Training UniVI: 11%|█ | 334/3000 [19:20<2:38:46, 3.57s/it, beta=1.150, gamma=1.450, train_loss=661.7969, val_loss=642.3644][2026-09-21 16:36:26,430] [UniVITrainer] [INFO] [Epoch 335] New best val loss: 642.3644
Training UniVI: 11%|█▏ | 344/3000 [19:57<2:42:38, 3.67s/it, beta=1.150, gamma=1.450, train_loss=661.6026, val_loss=642.1798][2026-09-21 16:37:03,060] [UniVITrainer] [INFO] [Epoch 345] New best val loss: 642.1798
Training UniVI: 12%|█▏ | 348/3000 [20:11<2:40:30, 3.63s/it, beta=1.150, gamma=1.450, train_loss=661.5320, val_loss=642.0118][2026-09-21 16:37:17,221] [UniVITrainer] [INFO] [Epoch 349] New best val loss: 642.0118
Training UniVI: 13%|█▎ | 380/3000 [21:53<2:11:42, 3.02s/it, beta=1.150, gamma=1.450, train_loss=660.6777, val_loss=641.9172][2026-09-21 16:38:59,348] [UniVITrainer] [INFO] [Epoch 381] New best val loss: 641.9172
Training UniVI: 13%|█▎ | 395/3000 [22:39<2:16:05, 3.13s/it, beta=1.150, gamma=1.450, train_loss=660.0608, val_loss=641.8553][2026-09-21 16:39:45,154] [UniVITrainer] [INFO] [Epoch 396] New best val loss: 641.8553
Training UniVI: 13%|█▎ | 397/3000 [22:45<2:14:50, 3.11s/it, beta=1.150, gamma=1.450, train_loss=660.0953, val_loss=641.8245][2026-09-21 16:39:51,304] [UniVITrainer] [INFO] [Epoch 398] New best val loss: 641.8245
Training UniVI: 13%|█▎ | 399/3000 [22:48<2:14:35, 3.10s/it, beta=1.150, gamma=1.450, train_loss=660.1598, val_loss=642.4406][2026-09-21 16:39:57,308] [UniVITrainer] [INFO] [Epoch 400] Train loss=660.3156 (beta=1.150, gamma=1.450)
[2026-09-21 16:39:57,439] [UniVITrainer] [INFO] [Epoch 400] Val loss=642.0370 (beta=1.150, gamma=1.450)
Training UniVI: 14%|█▍ | 423/3000 [24:07<2:17:14, 3.20s/it, beta=1.150, gamma=1.450, train_loss=659.2115, val_loss=641.6327][2026-09-21 16:41:12,885] [UniVITrainer] [INFO] [Epoch 424] New best val loss: 641.6327
Training UniVI: 14%|█▍ | 430/3000 [24:29<2:17:55, 3.22s/it, beta=1.150, gamma=1.450, train_loss=659.0405, val_loss=641.6027][2026-09-21 16:41:35,454] [UniVITrainer] [INFO] [Epoch 431] New best val loss: 641.6027
Training UniVI: 14%|█▍ | 435/3000 [24:44<2:10:53, 3.06s/it, beta=1.150, gamma=1.450, train_loss=658.9521, val_loss=641.5104][2026-09-21 16:41:50,690] [UniVITrainer] [INFO] [Epoch 436] New best val loss: 641.5104
Training UniVI: 16%|█▌ | 486/3000 [27:32<2:11:15, 3.13s/it, beta=1.150, gamma=1.450, train_loss=657.5944, val_loss=641.4186][2026-09-21 16:44:37,858] [UniVITrainer] [INFO] [Epoch 487] New best val loss: 641.4186
Training UniVI: 17%|█▋ | 499/3000 [28:07<2:05:54, 3.02s/it, beta=1.150, gamma=1.450, train_loss=657.4353, val_loss=641.7952][2026-09-21 16:45:16,718] [UniVITrainer] [INFO] [Epoch 500] Train loss=657.3714 (beta=1.150, gamma=1.450)
[2026-09-21 16:45:16,851] [UniVITrainer] [INFO] [Epoch 500] Val loss=642.3189 (beta=1.150, gamma=1.450)
Training UniVI: 20%|█▉ | 599/3000 [33:23<2:18:50, 3.47s/it, beta=1.150, gamma=1.450, train_loss=655.0741, val_loss=642.5529][2026-09-21 16:50:32,866] [UniVITrainer] [INFO] [Epoch 600] Train loss=654.7128 (beta=1.150, gamma=1.450)
[2026-09-21 16:50:33,012] [UniVITrainer] [INFO] [Epoch 600] Val loss=641.6089 (beta=1.150, gamma=1.450)
Training UniVI: 21%|██ | 621/3000 [34:40<2:15:35, 3.42s/it, beta=1.150, gamma=1.450, train_loss=654.4082, val_loss=641.2729][2026-09-21 16:51:46,477] [UniVITrainer] [INFO] [Epoch 622] New best val loss: 641.2729
Training UniVI: 23%|██▎ | 676/3000 [37:47<2:08:48, 3.33s/it, beta=1.150, gamma=1.450, train_loss=652.7708, val_loss=641.2381][2026-09-21 16:54:53,576] [UniVITrainer] [INFO] [Epoch 677] New best val loss: 641.2381
Training UniVI: 23%|██▎ | 699/3000 [39:01<2:10:20, 3.40s/it, beta=1.150, gamma=1.450, train_loss=651.8999, val_loss=641.5966][2026-09-21 16:56:10,542] [UniVITrainer] [INFO] [Epoch 700] Train loss=651.8528 (beta=1.150, gamma=1.450)
[2026-09-21 16:56:10,658] [UniVITrainer] [INFO] [Epoch 700] Val loss=642.5086 (beta=1.150, gamma=1.450)
Training UniVI: 24%|██▎ | 708/3000 [39:35<2:09:58, 3.40s/it, beta=1.150, gamma=1.450, train_loss=651.9442, val_loss=641.0636][2026-09-21 16:56:41,421] [UniVITrainer] [INFO] [Epoch 709] New best val loss: 641.0636
Training UniVI: 25%|██▍ | 738/3000 [41:06<1:47:05, 2.84s/it, beta=1.150, gamma=1.450, train_loss=651.1563, val_loss=641.0210][2026-09-21 16:58:12,417] [UniVITrainer] [INFO] [Epoch 739] New best val loss: 641.0210
Training UniVI: 26%|██▋ | 794/3000 [43:56<2:03:25, 3.36s/it, beta=1.150, gamma=1.450, train_loss=649.6361, val_loss=640.8471][2026-09-21 17:01:02,327] [UniVITrainer] [INFO] [Epoch 795] New best val loss: 640.8471
Training UniVI: 27%|██▋ | 799/3000 [44:10<2:08:03, 3.49s/it, beta=1.150, gamma=1.450, train_loss=649.6590, val_loss=641.6116][2026-09-21 17:01:19,645] [UniVITrainer] [INFO] [Epoch 800] Train loss=649.4795 (beta=1.150, gamma=1.450)
[2026-09-21 17:01:19,772] [UniVITrainer] [INFO] [Epoch 800] Val loss=641.9411 (beta=1.150, gamma=1.450)
Training UniVI: 30%|██▉ | 899/3000 [49:37<1:41:39, 2.90s/it, beta=1.150, gamma=1.450, train_loss=647.3629, val_loss=641.8805][2026-09-21 17:06:46,304] [UniVITrainer] [INFO] [Epoch 900] Train loss=647.4621 (beta=1.150, gamma=1.450)
[2026-09-21 17:06:46,427] [UniVITrainer] [INFO] [Epoch 900] Val loss=641.7296 (beta=1.150, gamma=1.450)
Training UniVI: 33%|███▎ | 994/3000 [54:15<1:36:11, 2.88s/it, beta=1.150, gamma=1.450, train_loss=646.1260, val_loss=641.8454][2026-09-21 17:11:21,071] [UniVITrainer] [INFO] Early stopping at epoch 995 (best val loss=640.8471, best epoch=795)
Training UniVI: 33%|███▎ | 994/3000 [54:15<1:49:29, 3.27s/it, beta=1.150, gamma=1.450, train_loss=646.1260, val_loss=641.8454]
[2026-09-21 17:11:21,083] [UniVITrainer] [INFO] Restored best model from epoch 795 (val loss=640.8471)
Held-out well: alignment between all three modalities#
test = parts["test"]
z = {m: encode_adata(model, a, modality=m, device=device, latent="modality_mean") for m, a in test.items()}
pd.Series({f"FOSCTTM {a}–{b}": compute_foscttm(z[a], z[b])
for a, b in [("rna", "adt"), ("rna", "atac"), ("adt", "atac")]}).round(3)
FOSCTTM rna–adt 0.075
FOSCTTM rna–atac 0.049
FOSCTTM adt–atac 0.081
dtype: float64
joint = stack_embeddings(model, [("well5", m, a) for m, a in test.items()], device=device)
sc.pp.neighbors(joint, use_rep="X_univi", n_neighbors=30)
sc.tl.leiden(joint, resolution=LEIDEN_RESOLUTION, key_added="leiden", flavor="igraph", n_iterations=2)
sc.tl.umap(joint, random_state=0)
sc.pl.umap(joint, color=["modality", "leiden"], wspace=0.35)
pd.crosstab(joint.obs["leiden"], joint.obs["modality"], normalize="index").round(2)
| modality | adt | atac | rna |
|---|---|---|---|
| leiden | |||
| 0 | 0.43 | 0.07 | 0.50 |
| 1 | 0.51 | 0.22 | 0.27 |
| 2 | 0.30 | 0.35 | 0.34 |
| 3 | 0.72 | 0.13 | 0.15 |
| 4 | 0.38 | 0.32 | 0.29 |
| 5 | 0.31 | 0.32 | 0.36 |
| 6 | 0.07 | 0.52 | 0.41 |
| 7 | 0.11 | 0.43 | 0.46 |
| 8 | 0.48 | 0.26 | 0.26 |
| 9 | 0.13 | 0.46 | 0.41 |
| 10 | 0.21 | 0.41 | 0.38 |
| 11 | 0.33 | 0.37 | 0.30 |
| 12 | 0.34 | 0.29 | 0.37 |
| 13 | 0.31 | 0.36 | 0.33 |
| 14 | 0.29 | 0.34 | 0.37 |
| 15 | 0.40 | 0.28 | 0.32 |
| 16 | 0.32 | 0.32 | 0.36 |
Marker concordance on the fused embedding (Supplemental Fig. S6)#
encode_fused_adata_pair(model, adata_by_mod=test, device=device, fused_obsm_key="X_univi_fused")
r, p = test["rna"], test["adt"]
sc.pp.neighbors(r, use_rep="X_univi_fused", n_neighbors=30)
sc.tl.umap(r, random_state=0)
p.obsm["X_umap"] = r.obsm["X_umap"]
genes = [g for g in ["MS4A1", "BANK1", "PAX5", "CD74", "LEF1", "BCL11B", "LYZ", "ITGAX"] if g in r.var_names]
proteins = [q for q in ["CD19", "IgD", "IgM", "CD3", "CD4", "CD45RA", "CD14", "CD11c", "HLA-DR"] if q in p.var_names]
sc.pl.umap(r, color=genes, ncols=4, vmax="p99", cmap="viridis")
sc.pl.umap(p, color=proteins, ncols=4, vmax="p99", cmap="viridis")
