Introduction

Prepare your dataset and choose the basic COATI settings before running the other tutorials. The workflow is prepare inputs → choose modality and mass settings → train a vector field → predict trajectories.

This notebook runs with small synthetic observations, so no biological download is required. Replace the demonstration AnnData objects with your prepared RNA and, optionally, ATAC data. Install COATI with the notebooks and visualization extras.

from pathlib import Path
import sys

example_dir = next(parent / "examples" for parent in (Path.cwd(), *Path.cwd().parents)
                   if (parent / "examples" / "synthetic" / "toy_data.py").is_file())
if str(example_dir) not in sys.path:
    sys.path.insert(0, str(example_dir))

from synthetic import toy_data as datasets
from tempfile import mkdtemp
import anndata as ad
import numpy as np
import pandas as pd
import coati as ct

Prepare your dataset

Choose a prepared feature representation for each modality, such as RNA PCA or ATAC LSI. Use the same feature columns, dimensions and preprocessing at every observed time. UMAP is for display; train on the selected feature representation.

COATI recommends the following Scanpy-compatible fields. These are COATI conventions, not names required by Scanpy.

Field

Contents

obs["time"]

Numeric observation time in your chosen units

obs["celltype"]

Supplied cell-type annotations; use unknown for unlabeled cells

obs["barcode"] and obs_names

The same globally unique cell identities

obsm["X_latent"] or obsm["X_pca"]

Prepared features selected explicitly with use_rep

obsm["X_umap"]

Optional existing UMAP coordinates, in matching row order

The example below creates two paired synthetic representations in X_latent. For your own data, replace this cell with rna = ad.read_h5ad("rna.h5ad") and, if needed, atac = ad.read_h5ad("atac.h5ad").

demo = datasets.paired(n=128, seed=7)
primary = np.concatenate(demo.primary)
secondary = np.concatenate(demo.secondary)
cell_times = np.repeat(demo.times, [len(x) for x in demo.primary])
barcodes = pd.Index([f"cell-{i:04d}" for i in range(len(primary))])
obs = pd.DataFrame({"day": cell_times, "cell_type": "Toy population"}, index=barcodes)

rna = ad.AnnData(primary.copy(), obs=obs.copy())
rna.obsm["X_latent"] = primary.copy()
atac = ad.AnnData(secondary.copy(), obs=obs.copy())
atac.obsm["X_latent"] = secondary.copy()
print(f"RNA: {rna.n_obs} cells, {rna.obsm['X_latent'].shape[1]} features")
print(f"ATAC: {atac.n_obs} cells, {atac.obsm['X_latent'].shape[1]} features")

RNA: 384 cells, 2 features
ATAC: 384 cells, 3 features

Standardize your AnnData

Map your source columns explicitly. Here, day becomes time, cell_type becomes celltype, and the existing cell identities become barcode. prepare_anndata returns a copy, retains source columns and preserves row order.

It does not normalize expression, run PCA/LSI, infer cell types or fit a UMAP by default. Prepare those features beforehand. Existing UMAP coordinates are retained when present.

cells = ct.prepare_anndata(
    rna, use_rep="X_latent", time_key="day", celltype_key="cell_type", warn=False,
)
cells.obs[["time", "celltype", "barcode"]].head()
time celltype barcode
cell-0000 0.0 Toy population cell-0000
cell-0001 0.0 Toy population cell-0001
cell-0002 0.0 Toy population cell-0002
cell-0003 0.0 Toy population cell-0003
cell-0004 0.0 Toy population cell-0004 
cells.obs.groupby("time", observed=True).size().rename("Cells").to_frame()
Cells
time
0.0 128
0.5 128
1.0 128

Keep the observation times explicit

Keep the spacing between observations meaningful; do not turn unequal intervals into equally spaced ranks. If your source uses labels such as early, middle and late, supply a numeric mapping. The following is an alternative to the numeric-time input above.

staged = rna.copy()
staged.obs["stage"] = staged.obs["day"].map({0.: "early", .5: "middle", 1.: "late"})
staged_cells = ct.prepare_anndata(
    staged, use_rep="X_latent", time_key="stage", celltype_key="cell_type",
    time_mapping={"early": 0., "middle": .5, "late": 1.}, warn=False,
)
staged_cells.obs[["stage", "time"]].drop_duplicates().reset_index(drop=True)
stage time
0 early 0.0
1 middle 0.5
2 late 1.0

Add paired secondary features when needed

For multimodal training, RNA and ATAC must represent the same cells within each snapshot. Align unique barcodes and check that their times agree before attaching secondary features. Cells do not need to be matched across different observation times, and snapshot sizes may differ.

The primary and secondary feature dimensions may differ. Changing use_rep later requires the representation used by that trained model; matching dimension counts alone is not enough.

assert rna.obs_names.is_unique and atac.obs_names.is_unique
atac_aligned = atac[rna.obs_names].copy()
np.testing.assert_array_equal(rna.obs["day"], atac_aligned.obs["day"])

paired_cells = ct.prepare_anndata(
    cells, use_rep="X_latent", secondary=atac_aligned.obsm["X_latent"],
    secondary_key="X_atac", warn=False,
)
print("Primary features:", paired_cells.obsm["X_latent"].shape)
print("Secondary features:", paired_cells.obsm["X_atac"].shape)

Primary features: (384, 2)
Secondary features: (384, 3)

If you only have arrays

COATI can assemble an AnnData from a list of snapshot matrices and their times. Provide one cell-type label and barcode per row when available. The example uses the same rows as above; NPZ, CSV and other adapters are covered in Input formats.

Feature scaling is a separate choice. Reuse already prepared inputs rather than normalizing them twice. The biological reconstruction tutorial includes optional model-space normalization.

array_cells = ct.prepare_anndata(
    demo.primary, time=demo.times,
    celltype=cells.obs["celltype"].astype(str).to_numpy(),
    barcode=cells.obs_names.to_numpy(), warn=False,
)
print("Assembled AnnData:", array_cells.shape)
print("Prepared representation:", array_cells.uns["coati_input"]["use_rep"])

Assembled AnnData: (384, 2)
Prepared representation: X_coati

Choose the model

The two main controls are multimodal and balanced.

multimodal

balanced

Model

False

True

Primary modality, conserved particle mass

True

True

Paired modalities, conserved particle mass

False

False

Primary modality with growth and changing mass

True

False

Paired modalities with growth and changing mass

Balanced training can still use snapshots with different observed cell counts. Unbalanced defaults use snapshot counts relative to the initial snapshot as mass targets; review whether those sampling counts reflect your experiment’s population sizes.

For multimodal fitting, the default map_type="film" first trains a time-conditioned RNA → ATAC map, then freezes it during vector-field training. Map fitting is described in Cross-modal map training.

Basic parameters

Values not supplied to Config use defaults. The following are the basic settings to review; a short example below overrides the training budget and batch size.

Parameter

Default

Meaning

multimodal

False

Enable paired secondary-modality constraints

balanced

True

Keep particle mass fixed; False enables growth

iterations

200

Number of vector-field training iterations

batch_size

64

Sampled cells per snapshot; cannot exceed the smallest snapshot

learning_rate

1e-3

Vector-field optimizer learning rate

secondary_weight

0.35 when multimodal

Secondary share of the energy and manifold penalties

matching_weight

50

Strength of distribution matching at observed stages

energy_weight

1

Strength of the kinetic and growth energy terms

manifold_weight

0

Strength of the manifold-support penalty

dt

0.1

Training solver step; each observation interval must be an integer multiple

seed

0

Random seed

Use multimodal=False to disable secondary training. Setting secondary_weight=0 alone leaves secondary distribution matching enabled. For the full parameter reference, see Configuration.

MULTIMODAL = False
BALANCED = True

config = ct.Config(
    multimodal=MULTIMODAL, balanced=BALANCED,
    iterations=25, batch_size=32, learning_rate=1e-3, dt=.1, seed=7,
)
print("Multimodal:", config.multimodal)
print("Balanced:", config.balanced)

Multimodal: False
Balanced: True

Run a short example

Select the prepared AnnData and train. Change the two controls above to use paired inputs or unbalanced dynamics. In this first example, 25 iterations are used to check that the inputs and training interface work together. The temporary output folder contains the fitted result; choose your own new output folder for a real experiment.

training_cells = paired_cells if config.multimodal else cells
output_dir = Path(mkdtemp(prefix="coati-introduction-")) / "model"
result = ct.train_vector_field(
    training_cells, config, use_rep="X_latent",
    secondary_obsm_key="X_atac" if config.multimodal else None,
    output_dir=output_dir,
)
trajectory = result.predict_trajectory(times=np.linspace(0., 1., 11))
print("Prediction times:", trajectory["times"])
print("Trajectory shape (times, cells, features):", trajectory["primary"].shape)

Prediction times: [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1. ]
Trajectory shape (times, cells, features): (11, 128, 2)

Continue with the tutorials