Supplemental Fig. S3: SHARE-seq mouse skin

Supplemental Fig. S3: SHARE-seq mouse skin#

Open in Colab

Paired RNA + ATAC from late-anagen mouse skin (Ma et al. 2020): 22 cell types across epidermal, hair-follicle, dermal, and immune lineages, with sparser ATAC than 10x Multiome. Settings follow the archived notebook UniVI_manuscript_GR-Supple_____mouse_skin_SHARE-seq_integration.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 re

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 ATACPreprocessor, 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 = 5000   # archived notebook: 5000 with early stopping (patience 200)
N_HVG = 5000
N_LSI = 101       # fit 101 components, drop the first (depth-correlated) one

Data, split, preprocessing#

The hosted files are after quality control, ENCODE blacklist filtering, and removal of the ambiguous “Mix” label. obs["split"] holds a train/validation/test assignment of 21,845 / 3,112 / 6,263 cells, which matches a split stratified by cell_type at 70/10/20; the archived notebook used 80/10/10 (24,969 / 3,112 / 3,139), so the test set here is about twice as large as in the archived analysis. Mitochondrial, ribosomal, Rik, and pseudogene families are excluded from HVG selection; peaks open in 0.25–80% of training cells are kept; TF-IDF uses the Signac convention.

data = uds.load("shareseq_mouse_skin")
rna, atac = data["rna"], data["atac"]
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["cell_type"], train_fraction=0.8, val_fraction=0.1, seed=42)
print({k: len(v) for k, v in splits.items()})

pattern = re.compile(r"^(mt-|Rps|Rpl|Mrps|Mrpl)|Rik$|-ps\d*$")
excluded = [g for g in rna.var_names if pattern.search(g)]
rna_prep = RNAPreprocessor(n_hvg=N_HVG, scale=True, exclude_genes=excluded).fit(rna[splits["train"]])
atac_prep = ATACPreprocessor(n_components=N_LSI, drop_first=True, scale=True, method="signac",
                             min_fraction=0.0025, max_fraction=0.80).fit(atac[splits["train"]])
parts = {k: {"rna": rna_prep.transform(rna[i]), "atac": atac_prep.transform(atac[i])} for k, i in splits.items()}

{'train': 21845, 'val': 3112, 'test': 6263}

shareseq_mouse_skin_rna.h5ad: 100%|██████████| 46.6M/46.6M [00:24<00:00, 1.94MB/s]
shareseq_mouse_skin_atac.h5ad: 100%|██████████| 246M/246M [02:57<00:00, 1.39MB/s]

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

[2026-09-21 16:17:27,608] [UniVITrainer] [INFO] TrainingConfig:
[2026-09-21 16:17:27,609] [UniVITrainer] [INFO]   n_epochs: 5000
[2026-09-21 16:17:27,610] [UniVITrainer] [INFO]   batch_size: 128
[2026-09-21 16:17:27,611] [UniVITrainer] [INFO]   lr: 0.001
[2026-09-21 16:17:27,614] [UniVITrainer] [INFO]   weight_decay: 0.0001
[2026-09-21 16:17:27,614] [UniVITrainer] [INFO]   device: 'cuda'
[2026-09-21 16:17:27,615] [UniVITrainer] [INFO]   log_every: 100
[2026-09-21 16:17:27,616] [UniVITrainer] [INFO]   grad_clip: None
[2026-09-21 16:17:27,617] [UniVITrainer] [INFO]   num_workers: 0
[2026-09-21 16:17:27,617] [UniVITrainer] [INFO]   seed: 0
[2026-09-21 16:17:27,618] [UniVITrainer] [INFO]   early_stopping: True
[2026-09-21 16:17:27,619] [UniVITrainer] [INFO]   patience: 200
[2026-09-21 16:17:27,619] [UniVITrainer] [INFO]   min_delta: 0.0
[2026-09-21 16:17:27,620] [UniVITrainer] [INFO]   best_epoch_warmup: 110
Training UniVI:   0%|          | 0/5000 [00:00<?, ?it/s][2026-09-21 16:17:33,091] [UniVITrainer] [INFO] [Epoch 001] Train loss=2530.0430 (beta=0.000, gamma=0.000)
[2026-09-21 16:17:33,265] [UniVITrainer] [INFO] [Epoch 001] Val loss=173473.0547 (beta=1.250, gamma=6.350)
Training UniVI:   2%|▏         | 99/5000 [08:14<6:20:15,  4.66s/it, beta=1.250, gamma=4.354, train_loss=2427.2132, val_loss=2924.0523][2026-09-21 16:25:46,695] [UniVITrainer] [INFO] [Epoch 100] Train loss=2427.6872 (beta=1.250, gamma=4.536)
[2026-09-21 16:25:46,839] [UniVITrainer] [INFO] [Epoch 100] Val loss=2994.8659 (beta=1.250, gamma=6.350)
Training UniVI:   2%|▏         | 108/5000 [08:59<6:09:28,  4.53s/it, beta=1.250, gamma=6.169, train_loss=2427.4673, val_loss=2956.5640][2026-09-21 16:26:27,476] [UniVITrainer] [INFO] [Epoch 109] Best tracking warmup ends next epoch (best_epoch_warmup=110).
Training UniVI:   2%|▏         | 109/5000 [09:04<6:08:56,  4.53s/it, beta=1.250, gamma=6.350, train_loss=2426.0652, val_loss=2963.4593][2026-09-21 16:26:32,092] [UniVITrainer] [INFO] [Epoch 110] New best val loss: 2963.4593
Training UniVI:   2%|▏         | 110/5000 [09:09<6:11:05,  4.55s/it, beta=1.250, gamma=6.350, train_loss=2429.3529, val_loss=2918.5341][2026-09-21 16:26:36,658] [UniVITrainer] [INFO] [Epoch 111] New best val loss: 2918.5341
Training UniVI:   2%|▏         | 119/5000 [09:50<6:17:31,  4.64s/it, beta=1.250, gamma=6.350, train_loss=2420.2634, val_loss=2913.4987][2026-09-21 16:27:18,176] [UniVITrainer] [INFO] [Epoch 120] New best val loss: 2913.4987
Training UniVI:   2%|▏         | 121/5000 [09:59<6:15:28,  4.62s/it, beta=1.250, gamma=6.350, train_loss=2423.6501, val_loss=2913.1823][2026-09-21 16:27:27,375] [UniVITrainer] [INFO] [Epoch 122] New best val loss: 2913.1823
Training UniVI:   3%|▎         | 127/5000 [10:27<6:12:26,  4.59s/it, beta=1.250, gamma=6.350, train_loss=2420.1265, val_loss=2907.8961][2026-09-21 16:27:54,871] [UniVITrainer] [INFO] [Epoch 128] New best val loss: 2907.8961
Training UniVI:   3%|▎         | 141/5000 [11:30<6:02:45,  4.48s/it, beta=1.250, gamma=6.350, train_loss=2408.6373, val_loss=2889.7634][2026-09-21 16:28:58,054] [UniVITrainer] [INFO] [Epoch 142] New best val loss: 2889.7634
Training UniVI:   3%|▎         | 167/5000 [13:28<6:06:16,  4.55s/it, beta=1.250, gamma=6.350, train_loss=2401.9770, val_loss=2879.6205][2026-09-21 16:30:56,226] [UniVITrainer] [INFO] [Epoch 168] New best val loss: 2879.6205
Training UniVI:   4%|▍         | 199/5000 [15:53<5:47:59,  4.35s/it, beta=1.250, gamma=6.350, train_loss=2393.0059, val_loss=3021.1689][2026-09-21 16:33:25,250] [UniVITrainer] [INFO] [Epoch 200] Train loss=2394.1094 (beta=1.250, gamma=6.350)
[2026-09-21 16:33:25,386] [UniVITrainer] [INFO] [Epoch 200] Val loss=2909.5222 (beta=1.250, gamma=6.350)
Training UniVI:   6%|▌         | 299/5000 [23:18<5:30:27,  4.22s/it, beta=1.250, gamma=6.350, train_loss=2373.9409, val_loss=3037.3566][2026-09-21 16:40:50,308] [UniVITrainer] [INFO] [Epoch 300] Train loss=2371.2899 (beta=1.250, gamma=6.350)
[2026-09-21 16:40:50,438] [UniVITrainer] [INFO] [Epoch 300] Val loss=3008.2357 (beta=1.250, gamma=6.350)
Training UniVI:   7%|▋         | 367/5000 [28:20<5:25:01,  4.21s/it, beta=1.250, gamma=6.350, train_loss=2360.5831, val_loss=3067.2313][2026-09-21 16:45:48,169] [UniVITrainer] [INFO] Early stopping at epoch 368 (best val loss=2879.6205, best epoch=168)
Training UniVI:   7%|▋         | 367/5000 [28:20<5:57:47,  4.63s/it, beta=1.250, gamma=6.350, train_loss=2360.5831, val_loss=3067.2313]
[2026-09-21 16:45:48,188] [UniVITrainer] [INFO] Restored best model from epoch 168 (val loss=2879.6205)

Held-out test cells#

test = parts["test"]
z_rna = encode_adata(model, test["rna"], modality="rna", device=device, latent="modality_mean")
z_atac = encode_adata(model, test["atac"], modality="atac", device=device, latent="modality_mean")
lab = test["rna"].obs["cell_type"].astype(str).to_numpy()
m = evaluate_alignment(Z1=z_rna, Z2=z_atac, labels_source=lab, labels_target=lab, recall_ks=(10,))
bi = m.get("bidirectional_transfer", {})
pd.Series({"FOSCTTM": m["foscttm_mean"], "Recall@10": m["recall_at_k"]["10"]["mean"],
           "modality mixing (k=30)": compute_modality_mixing(np.vstack([z_rna, z_atac]), np.repeat(["rna", "atac"], len(z_rna)), k=30),
           "label transfer acc (RNA to ATAC)": m["label_transfer_acc"],
           "macro-F1, worse direction": m["worst_direction_macro_f1"]}).round(3)

FOSCTTM                             0.117
Recall@10                           0.072
modality mixing (k=30)              0.157
label transfer acc (RNA to ATAC)    0.730
macro-F1, worse direction           0.609
dtype: float64

joint = stack_embeddings(model, [("test", "rna", test["rna"]), ("test", "atac", test["atac"])], 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", "cell_type"], wspace=0.5, legend_fontsize=6)
../../_images/f02e2de172348db404a996721ce8c2918120f60c4bba25bd2794226418fd36c8.png

Cross-modal reconstruction, per feature (Supplemental Fig. S3C–D)#

RNA to ATAC is evaluated in the LSI representation; ATAC to RNA on z-scored log-normalized expression.

r_atac = pearson_corr_per_feature(np.asarray(test["atac"].X),
                                  cross_modal_predict(model, test["rna"], src_mod="rna", tgt_mod="atac", device=device))
r_rna = pearson_corr_per_feature(np.asarray(test["rna"].X),
                                 cross_modal_predict(model, test["atac"], src_mod="atac", tgt_mod="rna", device=device))
fig, axes = plt.subplots(1, 2, figsize=(9, 3))
for ax, r, title in zip(axes, [r_atac, r_rna], ["RNA → ATAC (LSI components)", "ATAC → RNA (genes)"]):
    ax.hist(r[np.isfinite(r)], bins=40)
    ax.set(title=f"{title}; mean r = {np.nanmean(r):.3f}", xlabel="per-feature Pearson r")
plt.tight_layout()
plt.show()
../../_images/8339e44d42826db085864d46cfe08e07310e7eab33aad2ef46c445e3ed03fbe1.png
Contents