Python API

coati.anndata_io.prepare_anndata(data, *, time=None, celltype=None, barcode=None, obs=None, time_key=None, celltype_key=None, barcode_key=None, time_mapping=None, use_rep='auto', secondary=None, secondary_key=None, umap=None, compute_umap=False, n_neighbors=30, min_dist=0.35, seed=0, copy=True, warn=True)

Standardize AnnData or wrap prepared arrays without normalizing features.

Recommended fields are obs[‘celltype’], obs[‘time’], obs[‘barcode’], with barcodes also in obs_names and UMAP in obsm[‘X_umap’]. These are COATI conventions compatible with Scanpy, not names mandated by Scanpy. Existing source columns are retained. Copies are returned by default.

Arrays may be a cells-by-features matrix, a list of snapshots, or TemporalData. Snapshot lists require explicit numeric times; TemporalData supplies its times and paired secondary snapshots. Missing labels become ‘unknown’, never inferred cell types. Missing array barcodes are generated. Ambiguous aliases require an explicit source-column key. Categorical stage labels require time_mapping.

compute_umap=True fits a display UMAP on the selected prepared representation. Otherwise existing coordinates are reused; no embedding is fitted implicitly.

class coati.data.TemporalData(primary: object, times: object, secondary: object = None, labels: object = None)

Ordered snapshots at explicit numeric times, optionally paired across modalities.

Secondary rows must represent the same cells in the same order within each time. Pairing across different times is neither required nor assumed.

classmethod from_anndata(adata, *, time_key='time', obsm_key='auto', secondary_obsm_key=None, time_order=None, times=None)

Use X or an explicit embedding. Both modalities share adata.obs row identities.

classmethod from_dataframe(frame, *, time_key, feature_keys, time_order=None, times=None, secondary_keys=None)

Split a table without sorting cells; select feature columns explicitly.

classmethod from_dict(primary, *, times, secondary=None)

Use mapping insertion order; times explicitly define its numeric coordinates.

classmethod from_npz(path, *, times, keys=None, secondary_path=None)

Read numeric arrays only. Explicit keys are recommended for reproducibility.

to_anndata(**kwargs)

Assemble snapshots with canonical time, celltype and barcode metadata.

without_times(times)

Return paired snapshots excluding exact numeric observation times.

Supply this object to fit to exclude the same stages from automatic map fitting and trajectory fitting. Remaining numeric times are not renumbered. At least two observed stages must remain. This does not refit preprocessing.

class coati.config.Config(*, multimodal=None, balanced=None, extra=None, **kwargs)

Keep the original research parameter names; defaults target small CPU examples.

Additional engine parameters may be supplied using extra={...}. No normalization is applied by default. Call fit with normalization parameter paths to reuse the research pipeline’s normalization exactly.

property balanced

Whether population mass is conserved by the trajectory model.

property multimodal

Whether synchronization through a secondary modality is enabled.

with_parameters(**kwargs)

Return a new config with clear names or original engine names overridden.

class coati.config.MapConfig(hidden_dim: int = 128, n_layers: int = 3, time_hidden_dim: int = 32, epochs: int = 300, batch_size: int = 1024, lr: float = 0.001, weight_decay: float = 0.0001, validation_fraction: float = 0.2, patience: int = 30, eval_samples: int = 256, eval_repeats: int = 5, min_validation_cells: int = 20, max_normalized_mse: float = 0.5, max_sinkhorn_ratio: float = 5.0, quality_action: str = 'warn')

Paired map fitting and validation, separate from trajectory settings.

Quality thresholds are configurable diagnostics, not universal acceptance criteria. Validation cells are withheld within each supplied stage, not at new stages.

coati.training.train_vector_field(data, config=None, *, multimodal=None, balanced=None, output_dir=None, map_config=None, T_model=None, map_model=None, primal_norm_path=None, sec_norm_path=None, warm_start_path=None, verbose=False, use_rep='auto', time_key=None, celltype_key=None, barcode_key=None, time_mapping=None, secondary_obsm_key=None, **parameters)

Train the vector field with explicit modality and mass-conservation modes.

Without a config, defaults are single-modality and balanced. An explicit multimodal mode defaults to secondary_weight=0.35; an explicit unbalanced mode enables growth and count-based mass matching. A supplied Config is copied and overrides use clear names (iterations, batch_size, learning_rate) or existing engine names. AnnData is accepted directly: canonical metadata is prepared on a copy, and use_rep selects the original training features. prepare_anndata is the first step for arrays and categorical stage labels. Existing fit calls continue to work unchanged.

coati.training.fit(data, config=None, *, output_dir=None, T_model=None, map_model=None, map_config=None, primal_norm_path=None, sec_norm_path=None, warm_start_path=None, verbose=False)

Fit validated snapshots and return a portable Result.

Output directories must be empty; repeated notebook runs never erase previous logs. Normalization is opt-in, using the original engine’s saved scale dictionaries. With synchronization enabled, a FiLM map is fitted first unless a map is supplied. Config(map_type="mlp") fits an ordinary MLP without time input. Kernel modes remain explicit alternatives. map_config controls paired neural-map training. A supplied T_model accepts forward(q) or forward(q, t) in training coordinates and is frozen on a copy. Only supplied snapshots are used; omitted stages are never loaded for either map fitting or trajectory fitting.

class coati.result.Result(model, data, metadata, output_dir, T_model=None, map_model=None)

compute_growth(points=None, *, time=None, raw=True, chunk_size=1024)

Evaluate instantaneous d log(weight)/dt for an unbalanced model.

The rate follows the saved engine’s actual ODE: growth_net(x,t) times model.alpha_growth. raw controls input coordinate units, not rate units. Time units are exactly those supplied during training. This is distinct from cumulative log_weights or total population mass.

compute_velocity(points=None, *, time=None, raw=True, chunk_size=1024)

Evaluate the learned field without integrating trajectories.

points must use the same features, order and preprocessing as training. time is a scalar or one numeric time per cell. With points=None, snapshot rows are concatenated in data order and use their observed times. Explicit points require explicit times. raw=True returns input-unit displacement per unit of the numeric time supplied to training.

compute_velocity_embedding(adata, *, use_rep='auto', time_key='time', basis='umap', time=None, time_mapping=None, n_neighbors=30, ridge=0.001)

Return projected vectors for AnnData rows, without modifying AnnData.

use_rep=None selects adata.X. Numeric observation times must match training units; categorical stages require an explicit label-to-time mapping.

evaluate()

Endpoint Sinkhorn before/after transport; descriptive toy sanity check, not a biological benchmark.

classmethod load(directory, *, device='cpu', T_model=None, map_model=None)

Reload trajectories and built-in maps; external maps must be supplied again.

plot(*, ax=None, max_paths=30)

Plot the first two input coordinates; does not fit an embedding.

plot_dynamics(adata, *, use_rep='auto', time_key='time', time=None, time_mapping=None, basis='umap', celltype_key='celltype', labels=True, include_growth=None, palette=None, n_neighbors=30, ridge=0.001, axes=None, **plot_options)

Plot cell types and COATI streams; add growth for unbalanced models.

Returns a flat array of Axes. Each cell uses its observation time unless time is supplied. Subset AnnData for a one-stage view. See coati.dynamics.

plot_growth(adata, *, use_rep='auto', time_key='time', time=None, time_mapping=None, basis='umap', celltype_key='celltype', labels=True, show_velocity=False, n_neighbors=30, ridge=0.001, ax=None, **plot_options)

Plot instantaneous growth on UMAP, optionally with streams and labels.

Set show_velocity=True to overlay COATI flow. Growth uses original model features; color is the actual d log(weight)/dt, not cumulative log weight.

plot_trajectory(**kwargs)

Plot integrated paths; explicit name for plot().

plot_velocity(adata=None, *, embedding=None, points=None, time=None, use_rep='auto', time_key='time', time_mapping=None, basis='umap', n_neighbors=30, ridge=0.001, color=None, **plot_options)

Draw streamlines, grid arrows or cell arrows on an existing embedding.

AnnData example: result.plot_velocity(adata, use_rep=”auto”, time_key=”time”, color=”cell_type”, style=”stream”). NumPy inputs instead use embedding, points and time. Plotting never changes AnnData or fits an embedding.

predict(initial=None, *, times=None, method='rk4', raw=True, start_time=None)

Return {times, primary, energy, log_weights?} on an increasing grid.

initial is in input units when raw=True. Returned energy remains in training units. times starts at the observed initial time by default; its spacing controls fixed-step solvers. start_time allows an explicit supplied cohort to start later within the observed range. It does not select, transport or relabel initial cells automatically.

predict_trajectory(initial=None, **kwargs)

Integrate the trained vector field; explicit name for predict().

project_secondary(prediction=None, *, raw=True, chunk_size=1024)

Project a primary trajectory using the trained map and return secondary coordinates.

Pass a prediction in input units when raw=True (the default). Kernel prediction uses all paired observations as a deterministic basis; training retains its original randomly sampled basis. chunk_size limits query memory, not basis size.

project_velocity(embedding, *, points=None, time=None, raw=True, n_neighbors=30, ridge=0.001, basis='umap')

Project corresponding representation vectors into an existing 2D embedding.

write_dynamics(adata, *, copy=False, growth_key='coati_growth', **kwargs)

Store velocity vectors and, for unbalanced models, instantaneous growth.

Returns AnnData. copy=True preserves the caller’s input. Observation identities and row order are retained; no UMAP is fitted here.

write_velocity(adata, *, key='coati_velocity', copy=False, **kwargs)

Explicitly store primary and embedding velocities in AnnData.obsm.

Returns the input AnnData or a copy. Reordered cells retain their own representation, times and embedding; training row order is not assumed.

Example data preparation

Synthetic generators and Gaussian-specific helpers are maintained in the repository examples. They are not imported or installed with coati. The package consumes prepared arrays, AnnData or files through TemporalData and the input adapters.

Projected velocity

class coati.velocity.VelocityEmbedding(coordinates: ndarray, velocity: ndarray, time: ndarray | None = None, basis: str = 'umap', projection: str = 'local_linear')

Corresponding 2D positions and projected velocities, in unchanged row order.

plot(**kwargs)

Return a Matplotlib Axes using cells, grid arrows or streamlines.

coati.velocity.project_velocity(points, velocity, embedding, *, n_neighbors=30, ridge=0.001, time=None, basis='umap')

Project a vector using local weighted least-squares coordinate maps.

All three arrays must contain the same cells in the same order. Local maps use neighbors in the model’s input representation, without refitting UMAP. ridge is relative to the local displacement covariance; ridge=0 uses a minimum-norm solution. A degenerate input neighborhood produces zero flow. Velocities are not normalized, so a zero learned field remains zero.

Paired map training

coati.maps.train_map(data, *, kind, config, directory, device, seed, primary_scale=1.0, secondary_scale=1.0, verbose=False)

Fit, validate, save and freeze a default map before trajectory fitting.

No data discovery or full-atlas fallback is performed. All statistics and map fitting use training rows at supplied times. Validation rows select the checkpoint.

coati.maps.load_map(path, *, device='cpu')

class coati.maps.PairedMap(input_dim, output_dim, *, kind='film', hidden_dim=128, n_layers=3, time_hidden_dim=32)

FiLM or ordinary MLP with training-only feature standardization.

Inputs and outputs are in trajectory-training coordinates. FiLM accepts scalar times or times broadcastable to the leading dimensions of q. MLP ignores time.

forward(q, t=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Molecular trajectory readouts

coati.downstream.find_gene_peaks(gene, peak_gene_links, *, gene_key='gene', peak_key='peak', region_type=None, max_peaks=None)

Look up candidate peaks in a supplied DataFrame or CSV annotation table.

Matching is exact and case-sensitive. Optional region_type uses the table’s region_type column. Rank by absolute TSS distance when available, otherwise retain table order; no curve or outcome is used to choose peaks. Duplicate peak IDs are collapsed after sorting. Return canonical gene/peak columns and the original link metadata. These annotations do not prove regulation.

coati.downstream.compute_gene_peak_dynamics(gene, prediction, *, gene_decoder, peak_decoder, peak_gene_links, region_type=None, max_peaks=None, selection=None, weighting='uniform')

Read one gene and automatically look up its candidate peaks along paths.

Fit RNA and ATAC FeatureDecoders once in their corresponding model spaces. The decoders must contain the named gene and all selected candidate peaks; missing features raise instead of being silently skipped. prediction must contain primary RNA and secondary ATAC paths. Both use the same selection and weighting, with support screened separately in each modality. A gene name selects existing observations and annotations; no de novo linking, molecular prediction model or new vector-field training is added.

class coati.downstream.GenePeakDynamics(gene: str, rna: FeatureDynamics, atac: FeatureDynamics, links: DataFrame)

One gene and its candidate peaks, keeping each modality’s units separate.

plot(*, layout='combined', normalization='zscore', show_observed=False, observed=None, show_iqr=False)

Overlay separately standardized RNA/ATAC profiles, or use separate figures.

coati.downstream.plot_gene_peak_dynamics(dynamics, *, layout='combined', normalization='zscore', show_observed=False, observed=None, show_iqr=False)

Plot RNA and linked peaks on one coordinate system, or on separate panels.

combined overlays the gene and every linked peak on one Axes; each temporal mean is standardized separately as (value - temporal mean) / temporal std, with ddof=0 and finite trajectory time points. separate creates a gene figure and a peak figure, using the same normalization. normalization=None keeps native units only in separate mode. This display transform leaves the stored readouts unchanged and does not apply an additional log/count transform.

Return Axes arrays keyed by gene and peaks. In combined mode all entries refer to the same Axes. show_observed is optional and requires a table with modality (RNA/ATAC), feature, time, mean and optionally cohort. Observed means use the matching model profile’s center and scale, not their own fitted normalization. show_iqr transforms the bounds with that same scale; it describes across-path variation, not a molecular confidence interval.

Unsupported times remain NaN. A constant mean becomes a zero line with undefined IQR/observed normalization omitted. Profiles with no supported times remain missing. Standardized shapes do not establish temporal precedence.

coati.downstream.read_h5ad_features(path, features, *, obs_names=None, feature_key=None, layer=None, use_raw=False, binary=False, block_size=256)

Read selected H5AD features in bounded blocks without loading the atlas.

Supports dense, CSR and CSC matrices. Reorders cells by requested unique barcodes, rather than assuming RNA/ATAC file row order. binary=True converts nonnegative observed values to presence/absence before returning them. No normalization is applied. Raw gene counts and log-normalized expression require different explicit interpretation; select the appropriate layer.

coati.downstream.fit_feature_decoder(reference, values=None, *, features=None, use_rep='auto', feature_key=None, layer=None, n_neighbors=30, binary=False, units='Observed feature units', max_distance=None, support_quantile=0.99, seed=0)

Fit local readouts from AnnData, or from aligned latent/features arrays.

AnnData uses its selected representation and named columns from X or layer. Arrays require values and feature names in matching rows. To handle large on-disk atlases, read_h5ad_features reads only requested molecular features. No scaling, time interpolation, feature discovery or peak-gene linking is performed implicitly. Fit RNA and ATAC decoders in their own model spaces.

class coati.downstream.FeatureDecoder(reference: ndarray, observed: ndarray, features: tuple, neighbor_model: object, support_cutoff: float, units: str, binary: bool, barcodes: ndarray)

Distance-weighted local feature interpolation with an empirical support check.

predict(points, *, max_distance=None, chunk_size=1024)

Return interpolated values, nearest distances and a supported-point mask.

points may be cells-by-dimensions or time-by-paths-by-dimensions. Unsupported values are NaN. The cutoff is an empirical neighborhood screen, not a statistical confidence interval or guarantee of accuracy. Pass max_distance=np.inf only to explicitly disable this screen.

validate(points, observed)

Score independent molecular observations against a fixed training-mean baseline.

All finite test points are scored, including points outside the empirical support cutoff; coverage is reported separately. For binary targets, MSE is a Brier score and the baseline is the training open fraction. A skill score at or below zero does not improve that constant baseline.

coati.downstream.select_lineage(prediction, reference, *, terminal_celltypes, use_rep='auto', time_key='time', celltype_key='celltype', modality='primary', n_neighbors=15, min_probability=0.6, max_distance=None, support_quantile=0.99, seed=0)

Select paths whose endpoints match target annotations at the terminal time.

Uses only observed cells at prediction[‘times’][-1]. Missing labels become ‘unknown’ and do not count as a target. Returns every path’s target vote, nearest distance and support flag, even when no paths pass. Do not replace an empty selection with all paths or lower thresholds to manufacture a fate.

coati.downstream.compute_feature_dynamics(decoder, prediction, *, modality='primary', selection=None, weighting='uniform', max_distance=None)

Read observed features along paths and summarize each time/feature.

weighting=’growth’ uses normalized exp(log_weights) within the selected, supported cohort at each time. Uniform weighting gives each path equal weight. IQR describes variation across interpolated paths, not uncertainty of gene expression or experimental cell-population estimates. Binary observations yield local open probabilities; do not threshold an interpolated peak value at zero to claim a fraction of open cells.

coati.downstream.plot_feature_dynamics(dynamics, *, axes=None, color='#0072b2', show_iqr=True)

Plot per-feature means and across-path IQR; missing support leaves gaps.

coati.downstream.plot_feature_heatmap(dynamics, *, ax=None, standardize=False, cmap=None)

Plot feature-by-time means; optional row z-scores are visualization only.

Cell-type mass decomposition

coati.downstream.decompose_mass(prediction, reference, *, celltype, start_time=None, end_time=None, use_rep='auto', time_key='time', celltype_key='celltype', time_mapping=None, n_neighbors=1)

Separate a cell-type proportion change into incoming, outgoing and reweighting.

Track the same ordered particles at two observed times. prediction must be in the same feature units as reference (Result.predict_trajectory(raw=True)). Labels are transferred independently from stage-matched reference cells by hard distance-weighted kNN in the primary training representation, not UMAP. Missing annotations become ‘unknown’. One neighbor matches the paper figure; larger counts are an explicit analysis choice. Ties use sorted label order.

Normalize exp(log_weights) over the FULL supplied population at each time. Incoming uses later weights, outgoing uses earlier weights, and reweighting uses weight differences for retained particles. Do not preselect a lineage or a cell type. Missing log_weights means uniform balanced weights.

Endpoints default to the first/last prediction times and must match observed reference stages. The observed_change column is a cross-sectional atlas proportion difference, not a measured transition of the predicted cohort. Relative reweighting is not an absolute proliferation or death estimate.

class coati.downstream.MassDecomposition(summary: DataFrame, assignments: DataFrame, transitions: DataFrame, metadata: dict)

Relative mass components, endpoint assignments and transition summaries.

Incoming and outgoing are nonnegative fractions; net_change equals incoming - outgoing + reweighting. assignments retains every particle, including particles outside the requested type. No paths are discarded.

plot(*, ax=None, show_observed=False, title=None)

Plot signed components in percentage points; optionally show atlas change.

coati.downstream.plot_mass_decomposition(decomposition, *, ax=None, show_observed=False, title=None)

Display incoming, negative outgoing and retained reweighting in percentage points.