API

AnnData interface

class cloneot.CloneOTModel(adata, condition_key: str, source: str, target: str, clone_key: str | None = None, layer: str | None = None, n_top_genes: int = 2000, config: CloneOTConfig | None = None, vae_config: VAEConfig | None = None, device: str | None = None)

Bases: object

CloneOT on an AnnData object.

Parameters:

• adata (AnnData) – Cells of both conditions. adata.X (or adata.layers[layer]) holds raw counts or log-normalized expression.

• condition_key (str) – Column of adata.obs with the condition of every cell.

• source (str) – Values of condition_key for the cells before and after the perturbation.

• target (str) – Values of condition_key for the cells before and after the perturbation.

• clone_key (Optional[str]) – Column of adata.obs with the clone of every lineage-traced cell (NaN or empty for cells without a label). Without clone labels, CloneOT is entropic optimal transport with the squared Euclidean cost.

• layer (Optional[str]) – Layer to use instead of adata.X.

• n_top_genes (int) – Highly variable genes to model (Seurat v3 selection on counts, treating the two conditions as batches).

• config (Optional[CloneOTConfig]) – Settings of CloneOT and of the VAE (defaults: the settings of the paper).

• vae_config (Optional[VAEConfig]) – Settings of CloneOT and of the VAE (defaults: the settings of the paper).

• device (Optional[str]) – Torch device.

log_expression(adata)

Log-normalized expression of the model genes (normalization uses all genes of adata).

Return type:

ndarray

latent(adata)

VAE latent representation (posterior means) of the cells of adata.

Return type:

ndarray

train(validation=None, representation: CloneOTModel | None = None, n_response_genes: int = 50, verbose: bool = True)

Learn the representation, the transport cost and the neural map.

Parameters:

• representation (Optional[CloneOTModel]) – A trained model whose genes and VAE are reused, so that only the transport cost and the map are learned (e.g. to compare CloneOT with and without lineage on the same representation).

• validation (AnnData, optional) – Held-out source and target cells with clone labels (e.g. held-out cells of training clones). If given, the map checkpoint with the best validation criterion is kept: the mean over validation clones of l2, 1 - r^2 and the energy score on the response genes, each relative to predicting every clone by the mean training target cell.

• n_response_genes (int) – Number of response genes used for validation.

• verbose (bool) – Show progress bars.

Return type:

CloneOTModel

response_genes(n_genes: int = 50)

The n_genes model genes with the largest t-statistics between the training target and source cells (the evaluation genes of the paper). The ranking is computed once and saved with the model.

Return type:

list[str]

predict(adata=None, latent: bool = False)

Predicted response of every source cell of adata (default: the training source cells).

Returns an AnnData with the predicted log-normalized expression of the model genes and the obs of the source cells, or, with latent=True, the predicted latent states as an array.

observed(adata=None)

Log-normalized expression of the model genes of the target cells of adata (default: training cells), as an AnnData to compare with predict().

save(folder)

Save the model (genes, VAE and CloneOT) to folder; the training AnnData is not saved.

Return type:

None

classmethod load(folder, adata=None, device: str | None = None)

Load a saved model. adata (optional) is only needed to predict the training source cells by default.

Return type:

CloneOTModel

Solver

class cloneot.CloneOT(config: CloneOTConfig | None = None, device: str | None = None, **overrides)

Bases: object

Neural inverse optimal transport from lineage tracing.

Parameters:

• config (Optional[CloneOTConfig]) – Settings (defaults: the settings of the paper).

• device (Optional[str]) – Torch device (‘cuda’ if available, else ‘cpu’).

• **overrides – Fields of CloneOTConfig to override, e.g. CloneOT(seed=1).

property learned: bool

Whether the cost is a neural network (otherwise the squared Euclidean cost).

fit(source, target, source_clones: Sequence | None = None, target_clones: Sequence | None = None, verbose: bool = False)

Learn the transport cost from the labeled clones and solve the entropic OT problem under it.

Parameters:

• source (array of shape (n_source, d) and (n_target, d)) – Cell states before and after the perturbation (e.g. VAE latent codes).

• target (array of shape (n_source, d) and (n_target, d)) – Cell states before and after the perturbation (e.g. VAE latent codes).

• source_clones (Optional[Sequence]) – Clone label of every source and target cell; -1, None, NaN or ‘’ for cells without a label. Without labels, CloneOT is entropic optimal transport with the squared Euclidean cost.

• target_clones (Optional[Sequence]) – Clone label of every source and target cell; -1, None, NaN or ‘’ for cells without a label. Without labels, CloneOT is entropic optimal transport with the squared Euclidean cost.

• verbose (bool) – Show a progress bar.

Return type:

CloneOT

fit_map(decoder: Callable | None = None, evaluate: Callable | None = None, eval_every: int = 500, verbose: bool = False)

Fit the neural map T(x) = x + F(x) to the barycentric projection E_pi[y | x] of the learned coupling.

Parameters:

• decoder (Optional[Callable]) – Maps cell states to expression (e.g. a VAE decoder). If given, F is fitted so that decoder(T(x)) matches the expected decoded descendant E_pi[decoder(y) | x].

• evaluate (Optional[Callable]) – evaluate(transport) -> validation score (lower is better), called every eval_every steps with the current map; the best checkpoint is kept.

Return type:

CloneOT

transport(source)

Predicted response T(x) of every source cell (rows of source), from its state alone.

Return type:

ndarray

conditional(source)

Matrix pi(y_j | x_i) of the learned coupling over the training target cells.

Return type:

ndarray

barycentric_projection(source, values=None)

E_pi[v(y) | x] over the training target cells (v: the target states by default).

Return type:

ndarray

cost(source, target)

Learned cost matrix c(x_i, y_j) (the squared Euclidean cost if no cost was learned).

Return type:

ndarray

save(path)

Save the configuration, the networks and the training data needed by conditional() and cost().

Return type:

None

class cloneot.CloneOTConfig(eps: float = 0.1, hidden: int = 128, layers: int = 3, activation: str = 'silu', clone_embed: int = 16, lr_potential: float = 0.001, batch: int = 4096, cost: str = 'neural', h: str = 'icnn', flows: bool = True, icnn_hidden: int = 64, icnn_layers: int = 2, icnn_init: float = 1.0, flow_hidden: int = 64, flow_layers: int = 2, flow_steps: int = 4, lr_cost: float = 0.001, lr_flow: float = 0.0001, cost_beta2: float = 0.99, cost_batch: int = 1024, cost_samples: int = 8, warmup: int = 500, steps: int = 2000, potential_steps: int = 5, polish: int = 2000, map_lr: float = 0.0003, map_steps: int = 3000, map_batch: int = 256, seed: int = 0)

Bases: object

Settings of CloneOT. The defaults are the settings of the paper, shared by all simulations and single-cell tasks.

The fields of CloneOTConfig and their defaults (the settings of the paper):

@dataclass
class CloneOTConfig:
    """Settings of CloneOT. The defaults are the settings of the paper, shared by all simulations and single-cell
    tasks."""
    eps: float = 0.1                 # entropic temperature, in units of the median m of 1/2 |y - x|^2
    # dual potentials and map
    hidden: int = 128                # width of the potentials and of the map
    layers: int = 3                  # hidden layers of the potentials and of the map
    activation: str = 'silu'
    clone_embed: int = 16            # dimension of the clone embedding of the constrained potential (0: none)
    lr_potential: float = 1e-3
    batch: int = 4096                # target cells per potential step (all target cells if fewer)
    # learned cost c(x, y) = h(phi(x) - psi(y))
    cost: str = 'neural'             # 'neural' (learned) or 'euclidean' (fixed 1/2 |x - y|^2, i.e. entropic W2 OT)
    h: str = 'icnn'                  # 'icnn': g(d) + g(-d) + alpha/2 |d|^2; 'quadratic': alpha/2 |d|^2
    flows: bool = True               # neural ODE flows phi and psi (False: identity maps)
    icnn_hidden: int = 64
    icnn_layers: int = 2
    icnn_init: float = 1.            # scale of the ICNN output weights at initialization
    flow_hidden: int = 64
    flow_layers: int = 2
    flow_steps: int = 4              # Runge-Kutta steps of the flows
    lr_cost: float = 1e-3            # learning rate of h
    lr_flow: float = 1e-4            # learning rate of the flows
    cost_beta2: float = 0.99         # Adam second-moment decay of the cost
    cost_batch: int = 1024           # target cells per cost gradient
    cost_samples: int = 8            # source cells sampled per target cell from each coupling
    # schedule
    warmup: int = 500                # potential steps before the cost moves
    steps: int = 2000                # cost updates
    potential_steps: int = 5         # steps of each potential per cost update
    polish: int = 2000               # steps of the free potential at the final cost
    # neural map
    map_lr: float = 3e-4
    map_steps: int = 3000
    map_batch: int = 256             # source cells per map step
    seed: int = 0

Representation

class cloneot.VAE(n_genes: int, config: VAEConfig | None = None)

Bases: Module

latent(x: Tensor, batch: int = 4096)

Posterior means.

Return type:

Tensor

fit(expression: Tensor, domain: ndarray, verbose: bool = False)

Train on expression (cells x genes, log-normalized), with batches balanced between the two domains (domain: 0 for source and 1 for target cells). The checkpoint with the lowest held-out loss is kept.

Return type:

VAE

class cloneot.VAEConfig(hidden: list[int] = <factory>, latent_dim: int = 50, alpha: float = 0.0001, lr: float = 0.001, weight_decay: float = 1e-05, batch_size: int = 256, steps: int = 25000, eval_every: int = 500, validation_fraction: float = 0.15, seed: int = 0)

Bases: object

Evaluation

Clone-level evaluation of predicted perturbation responses.

For every clone with predicted profiles {p_i} (one per source cell) and observed target cells {y_j}:

• l2: distance between the predicted and the observed clone-mean expression;

• r2: squared Pearson correlation between the two clone means across genes;

• energy: the energy score mean ||p - y|| - 1/2 mean ||p - p'||, a strictly proper scoring rule of the predicted distribution.

The clone-specific correlation removes the response that all clones share: it correlates the predicted and observed deviations of every clone mean from the average over clones, across all clones and genes.

cloneot.metrics.energy_score(predicted: ndarray, observed: ndarray)

Energy score mean ||p - y|| - 1/2 mean ||p - p'|| of the predicted profiles (rows) against the observed cells (lower is better).

Return type:

float

cloneot.metrics.clone_scores(predicted, observed, predicted_clones: Sequence, observed_clones: Sequence, min_cells: int = 1)

Per-clone l2, r2 and energy score.

Parameters:

• predicted (array (n_predictions, n_genes)) – Predicted profile of every source cell.

• observed (array (n_observed, n_genes)) – Observed target cells.

• predicted_clones (Sequence) – Clone of every predicted and observed cell. Clones present on both sides with at least min_cells cells on each side are scored.

• observed_clones (Sequence) – Clone of every predicted and observed cell. Clones present on both sides with at least min_cells cells on each side are scored.

Return type:

DataFrame

cloneot.metrics.clone_specific_correlation(predicted, observed, predicted_clones: Sequence, observed_clones: Sequence)

Pearson correlation between predicted and observed clone-specific deviations over all clones and genes.

Return type:

float

cloneot.metrics.evaluate(predicted, observed, clone_key: str, genes: Sequence[str] | None = None, min_cells: int = 1)

Clone-level evaluation of AnnData objects: per-clone scores and the clone-specific correlation.

Parameters:

• predicted (AnnData) – Predictions (e.g. from cloneot.CloneOTModel.predict()), with the clone of every source cell in obs[clone_key].

• observed (AnnData) – Observed target cells (log-normalized expression) with obs[clone_key].

• genes (Optional[Sequence[str]]) – Genes to evaluate (e.g. the differentially expressed genes from cloneot.metrics.response_genes()). By default, all genes shared by the two objects.

Return type:

tuple[DataFrame, float]

cloneot.metrics.response_genes(source, target, n_genes: int = 50)

The n_genes genes with the largest t-statistics between target and source cells (log-normalized AnnData), the evaluation genes of the paper.

Return type:

list[str]

Datasets

The two 2D datasets of the paper, and public lineage-tracing data.

In both 2D datasets, source and target are unpaired samples of two classes, and optimal transport with the squared Euclidean cost sends points to the wrong class:

• t_shape: the T-shape dataset. A vertical bar on the left moves to the right and a horizontal bar on top moves down. The points form corresponding subsets of two source and two target points at one position of a bar.

• rotated_moons: the rotated-moons dataset. The target is the two moons rotated by 90 degrees, and the only correspondence is between each source moon and its rotated moon.

Each function returns a dictionary with the source and target points, their classes (source_class, target_class), the labels of the corresponding subsets (source_clone, target_clone, which play the role of clones) and the true target position fate(x) of a source point.

cloneot.datasets.t_shape(n: int = 400, cells: int = 2, noise: float = 0.05, seed: int = 0)

T-shape dataset: n source and n target points in corresponding subsets of cells points, half on each bar.

Return type:

dict

cloneot.datasets.rotated_moons(n: int = 400, angle: float = 90.0, noise: float = 0.05, seed: int = 0)

Rotated two moons: the target is an independent sample of the two moons rotated by angle degrees about their centre. The two moons are the corresponding subsets.

Return type:

dict

cloneot.datasets.label_fraction(source_clone, target_clone, fraction: float, seed: int = 0)

Keep the labels of a random fraction of the subsets (clones) and mark all other points as unlabeled (-1).

cloneot.datasets.class_accuracy(prediction, true_class, reference: dict)

Fraction of predictions whose nearest reference target cell belongs to the true class of the source cell.

Return type:

float

cloneot.datasets.download(url: str, path, sha256: str | None = None)

Download url to path (skipped if a file with the expected checksum exists).

Return type:

None

cloneot.datasets.watermelon(data_dir: str = 'data', download_files: bool = True)

Watermelon lineage tracing of PC9 lung cancer cells treated with osimertinib (Oren et al., Nature 2021; GEO GSE150949).

Cells were profiled before treatment (day 0) and at days 3, 7 and 14; day-14 persister cells were sorted by a proliferation-history reporter into cycling, moderate and non-cycling fractions. The first call downloads 125 MB from GEO and caches <data_dir>/watermelon.h5ad.

Returns:

• AnnData with raw counts in X and obs columns clone (single-barcode clone label, NA otherwise),

• lineage_block, day, condition (‘day0’, ‘day3’, ‘day7’, ‘day14’), fraction (‘cycling’, ‘moderate’,

• ’non-cycling’ or ‘unsorted’ for day 14; ‘unsorted’ before), sample and replicate.

cloneot.datasets.holdout(adata, clone_key: str, condition_key: str, fraction: float = 0.2, unit: str = 'cells', min_cells: int = 2, seed: int = 0)

Boolean mask of held-out cells for evaluation.

• unit='cells' (seen clones): a fraction of the cells of every clone with at least min_cells cells in each condition is held out, so the test clones also appear in training.

• unit='clones' (unseen clones): all cells of a random fraction of the clones are held out.

Cells without a clone label are never held out.

Return type:

ndarray

Plotting

Plotting functions of CloneOT.

cloneot.plotting.transport(source, target, prediction, correct: Sequence[bool] | None = None, ax=None, rays: int = 4, title: str | None = None, legend: bool = True)

Two-dimensional transport plot: source cells, target cells, predictions, and rays from every rays-th source cell to its prediction (red if correct is False for that cell).

cloneot.plotting.clone_deviations(predicted, observed, clone_key: str, genes: Sequence[str] | None = None, ax=None, color: str = '#51A9D6')

Scatter of predicted against observed clone-specific deviations (clone mean minus the mean over clones) for every clone and gene, annotated with the clone-specific correlation r.

cloneot.plotting.paired_scores(baseline, scores, metric: str = 'l2', labels: Sequence[str] = ('Baseline', 'CloneOT'), ax=None, color: str = '#51A9D6', other: str = '#C7BBAF')

Per-clone comparison of two methods on the same clones: the metric of baseline (x-axis) against that of scores (y-axis), one point per clone, from data frames returned by cloneot.metrics.evaluate().

Clones on which the second method is better are colored, and the panel reports their fraction and the p-value of the two-sided Wilcoxon signed-rank test.