Cell-level representation learning¶

The goal: turn raw, multi-batch, multi-modal single-cell data into one clean embedding — a handful of coordinates per cell that mix the batches together while keeping the biology — so you can cluster and analyze all your data as one.

This demo runs end-to-end on 3,000 cells of the built-in Liu subsample (3 batches, RNA + ADT) — one line to load, no downloads. It uses the same settings as the full-data tutorial; only the data is smaller. For the full-data run that reproduces the publication numbers, see m3pkg/tutorials/01_representation_learning.py.

In [1]:

import os
import numpy as np
import matplotlib.pyplot as plt
import scanpy as sc
import m3

OUT = "./tutorial_out_demo/01_representation_learning"
os.makedirs(OUT, exist_ok=True)
sc.settings.verbosity = 1

import os import numpy as np import matplotlib.pyplot as plt import scanpy as sc import m3 OUT = "./tutorial_out_demo/01_representation_learning" os.makedirs(OUT, exist_ok=True) sc.settings.verbosity = 1

1. Load the demo data¶

m3.datasets.liu_demo() returns a ready-to-use m3.Dataset — a subsample of the Liu et al. COVID-19 CITE-seq data (RNA + ADT), with the obs columns already set up. No downloads.

In [2]:

data = m3.datasets.liu_demo(n_cells=3000)
print(data)                                              # cells, batches, modalities
print("celltypes:", data.obs["mergedcelltype"].value_counts().to_dict())

data = m3.datasets.liu_demo(n_cells=3000) print(data) # cells, batches, modalities print("celltypes:", data.obs["mergedcelltype"].value_counts().to_dict())

Dataset(n_cells=3000, batches=['B1', 'B2', 'B3'], modalities=[rna:1000, adt:192])
celltypes: {'CD4_Mem': 361, 'Mono_Classical': 303, 'CD8_Mem': 297, 'B_Mem': 294, 'NK': 246, 'CD8_Va7.2': 202, 'T_Vd2': 200, 'Mono_Nonclassical': 184, 'Plasmablast': 177, 'CD4_Naive': 169, 'CD8_Naive': 151, 'B_Naive': 112, 'Platelet': 106, 'T_gd': 57, 'DC': 54, 'Erythrocyte': 52, 'pDC': 35}

2. Build and train the model¶

m3.M3(...) sets up the model from your data, telling it which obs columns to use:

• condition_keys — the conditions you're comparing across (here disease group + age band).
• celltype_key — the column holding cell-type labels.
• batch_key — the column marking which batch each cell is from (B1/B2/B3). The model corrects for batch, so cells group by biology rather than by batch.
• embedding_dim — how many numbers describe each cell in the output (here 30).

.train() fits the model. donor_prediction=False skips the disease-prediction step — not needed for embeddings, and faster. model.capabilities then shows which outputs are available.

In [3]:

model = m3.M3(
    data,
    condition_keys=["cond_group", "Age_interval"],
    celltype_key="mergedcelltype",
    batch_key="batch",
    embedding_dim=30,
)
model.train(max_epochs=500, donor_prediction=False)
print("capabilities:", model.capabilities)

model = m3.M3( data, condition_keys=["cond_group", "Age_interval"], celltype_key="mergedcelltype", batch_key="batch", embedding_dim=30, ) model.train(max_epochs=500, donor_prediction=False) print("capabilities:", model.capabilities)

Using 'matrix/data'
torch.Size([1000, 1000])
Using 'matrix/data'
torch.Size([1000, 1000])
Using 'matrix/data'
torch.Size([1000, 1000])
Using 'matrix/data'
torch.Size([192, 1000])
Using 'matrix/data'
torch.Size([192, 1000])
Using 'matrix/data'
torch.Size([192, 1000])
Batch counts: {np.int64(0): np.int64(1000), np.int64(1): np.int64(1000), np.int64(2): np.int64(1000)}
Minimum batch size: 1000

Epoch 1, Validation Loss: 5.0042

Epoch 2, Validation Loss: 5.0026

Epoch 3, Validation Loss: 4.9794

...  (493 epochs omitted)  ...

Epoch 497, Validation Loss: 1.8574

Epoch 498, Validation Loss: 1.8512

Epoch 499, Validation Loss: 1.8378

Epoch 500, Validation Loss: 1.8504
Using 'matrix/data'
torch.Size([1000, 1000])
Using 'matrix/data'
torch.Size([1000, 1000])
Using 'matrix/data'
torch.Size([1000, 1000])
Using 'matrix/data'
torch.Size([192, 1000])
Using 'matrix/data'
torch.Size([192, 1000])
Using 'matrix/data'
torch.Size([192, 1000])

capabilities: {'embedding': True, 'reconstruct': True, 'predict_donors': False}

3. Get the embedding¶

model.embedding(part=...) returns the trained coordinates as a cells x embedding_dim array. Ask for the part you want by name:

• "bio" — the biological signal, with batch differences removed. This is the one you cluster and analyze.
• "batch" — the batch signal on its own, if you ever want to inspect it separately.

model.cell_metadata is the per-cell obs table, in the same row order as the embedding, so you can colour plots by cell type / batch / condition.

In [4]:

emb_bio = model.embedding(part="bio")
emb_batch = model.embedding(part="batch")
meta = model.cell_metadata.copy()
print("bio:", emb_bio.shape, "batch:", emb_batch.shape)

emb_bio = model.embedding(part="bio") emb_batch = model.embedding(part="batch") meta = model.cell_metadata.copy() print("bio:", emb_bio.shape, "batch:", emb_batch.shape)

bio: (3000, 28) batch: (3000, 2)

4. Save the embedding¶

Save the embedding and its metadata to disk. The three formats below hold the same data — pick whichever suits your downstream tools:

• .npy — the raw array (plus a CSV of the metadata).
• .h5 — the array in an HDF5 file.
• .h5ad — an AnnData (embedding as X, metadata as obs) — handy for scanpy.

In [5]:

import h5py
import anndata as ad

np.save(os.path.join(OUT, "embedding_bio.npy"), emb_bio)
meta.to_csv(os.path.join(OUT, "cell_metadata.csv"), index=False)
with h5py.File(os.path.join(OUT, "embedding_bio.h5"), "w") as f:
    f["data"] = emb_bio
ad.AnnData(X=emb_bio, obs=meta.astype(str).reset_index(drop=True)).write_h5ad(
    os.path.join(OUT, "embedding_bio.h5ad"))
print("saved -> .npy / .h5 / .h5ad in", OUT)

import h5py import anndata as ad np.save(os.path.join(OUT, "embedding_bio.npy"), emb_bio) meta.to_csv(os.path.join(OUT, "cell_metadata.csv"), index=False) with h5py.File(os.path.join(OUT, "embedding_bio.h5"), "w") as f: f["data"] = emb_bio ad.AnnData(X=emb_bio, obs=meta.astype(str).reset_index(drop=True)).write_h5ad( os.path.join(OUT, "embedding_bio.h5ad")) print("saved -> .npy / .h5 / .h5ad in", OUT)

/opt/homebrew/Caskroom/miniforge/base/lib/python3.13/functools.py:934: ImplicitModificationWarning: Transforming to str index.
  return dispatch(args[0].__class__)(*args, **kw)
... storing 'sample_id' as categorical

... storing 'Donor' as categorical

... storing 'cond_group' as categorical

... storing 'Age_interval' as categorical

... storing 'mergedcelltype' as categorical

... storing 'batch' as categorical

saved -> .npy / .h5 / .h5ad in ./tutorial_out_demo/01_representation_learning

5. Visualise with UMAP¶

A UMAP of the "bio" embedding. If integration worked, cells should group by cell type (biology kept) while the batches mix together (batch difference removed) — colour by each to check.

In [6]:

# Build an AnnData from the bio embedding, then colour the UMAP by cell type, batch,
# and condition.
adn = sc.AnnData(X=emb_bio, obs=meta.reset_index(drop=True))
adn.obs["batch"] = adn.obs["batch"].astype("category")
adn.obs["mergedcelltype"] = adn.obs["mergedcelltype"].astype(str).astype("category")
adn.obs["cond_group"] = adn.obs["cond_group"].astype(str).astype("category")
sc.pp.neighbors(adn, use_rep="X", n_neighbors=15)
sc.tl.umap(adn)
fig, axes = plt.subplots(1, 3, figsize=(18, 5))
for ax, key, title in zip(
    axes,
    ["mergedcelltype", "batch", "cond_group"],
    ["Cell type (biology preserved)", "Batch (mixed)", "Condition"],
):
    sc.pl.umap(adn, color=key, ax=ax, show=False, title=title,
               legend_loc="right margin" if key != "mergedcelltype" else "on data",
               legend_fontsize=6)
fig.tight_layout()
fig.savefig(os.path.join(OUT, "umap_bio.png"), dpi=130, bbox_inches="tight")
plt.show()

# Build an AnnData from the bio embedding, then colour the UMAP by cell type, batch, # and condition. adn = sc.AnnData(X=emb_bio, obs=meta.reset_index(drop=True)) adn.obs["batch"] = adn.obs["batch"].astype("category") adn.obs["mergedcelltype"] = adn.obs["mergedcelltype"].astype(str).astype("category") adn.obs["cond_group"] = adn.obs["cond_group"].astype(str).astype("category") sc.pp.neighbors(adn, use_rep="X", n_neighbors=15) sc.tl.umap(adn) fig, axes = plt.subplots(1, 3, figsize=(18, 5)) for ax, key, title in zip( axes, ["mergedcelltype", "batch", "cond_group"], ["Cell type (biology preserved)", "Batch (mixed)", "Condition"], ): sc.pl.umap(adn, color=key, ax=ax, show=False, title=title, legend_loc="right margin" if key != "mergedcelltype" else "on data", legend_fontsize=6) fig.tight_layout() fig.savefig(os.path.join(OUT, "umap_bio.png"), dpi=130, bbox_inches="tight") plt.show()

/opt/homebrew/Caskroom/miniforge/base/lib/python3.13/functools.py:934: ImplicitModificationWarning: Transforming to str index.
  return dispatch(args[0].__class__)(*args, **kw)

... storing 'sample_id' as categorical

... storing 'Donor' as categorical

... storing 'Age_interval' as categorical

Done. You now have an integrated cell embedding, ready for clustering and downstream analysis. For the full-data run, see tutorials/01_representation_learning.py.