cellink.tl.external.train_livi

cellink.tl.external.train_livi(adata_or_dd, output_dir, *, individual_col=None, z_dim=15, n_dxc_factors=100, n_persistent_factors=5, n_cis_snps=0, encoder_hidden_dims=None, learning_rate=0.0008, use_size_factor=True, size_factor_key=None, layer_key=None, covariates_keys=None, covariates_dims=None, known_cis_eqtls=None, eqtl_genotypes=None, warmup_epochs_vae=30, warmup_epochs_G=0, max_epochs=200, min_epochs=50, batch_size=512, seed=42, l1_weight=0.001, A_weight=0.001, batch_norm_decoder=False, genetics_seed=None, cell_state_cis=True, num_workers=0, strict=False, drop_last=True, shuffle=True, pin_memory=True, log_every_n_steps=1, enable_progress_bar=True, enable_checkpointing=True, enable_logger=True, gradient_clip_val=None, accumulate_grad_batches=1, limit_train_batches=None, deterministic=False, callbacks=None, run=True, runner=None)

Train a LIVI model on single-cell RNA-seq data.

LIVI decomposes cell × gene expression into:

• Cell-state factors (z): shared across donors, learned by the VAE encoder.

• DxC factors (D): donor × cell-state interaction effects, stored in a learned donor embedding.

• Persistent factors (V): cell-state-independent donor effects.

Parameters:

• adata_or_dd (AnnData or DonorData) – Input data. When DonorData is passed, dd.C (cell-level AnnData) is used and individual_col defaults to dd.donor_id.

• output_dir (str) – Directory where checkpoints and logs are written. Created if absent.

• individual_col (str, optional) – Column in adata.obs carrying donor / individual IDs. Required when passing raw AnnData; auto-inferred from DonorData.

• z_dim (int) – Dimensionality of the cell-state latent space.

• n_dxc_factors (int) – Number of donor × cell-state interaction (DxC) factors.

• n_persistent_factors (int) – Number of persistent (cell-state-independent) donor factors.

• n_cis_snps (int) – Number of known cis-eQTL SNPs included during training. Set to 0 when not using genotype data.

• encoder_hidden_dims (list of int, optional) – Hidden-layer widths of the encoder MLP. Defaults to [512, 256, 64].

• learning_rate (float) – Adam optimizer learning rate.

• use_size_factor (bool) – Use per-cell library-size factors for count normalisation.

• size_factor_key (str, optional) – Key in adata.obs with pre-computed size factors. If None and use_size_factor=True, size factors are the total counts per cell.

• layer_key (str, optional) – Key in adata.layers containing raw integer counts. If None, adata.X is used.

• covariates_keys (list of str, optional) – Categorical covariate columns in adata.obs (e.g. ["pool", "sex"]). These are corrected for as nuisance effects.

• covariates_dims (list of int, optional) – Number of unique categories per covariate (same order as covariates_keys). If None, inferred from data.

• known_cis_eqtls (str or pd.DataFrame, optional) – TSV path or DataFrame (SNPs × genes, binary 0/1) indicating known cis-eQTL associations. n_cis_snps must equal the number of rows.

• eqtl_genotypes (str or pd.DataFrame, optional) – TSV path or DataFrame (individuals × SNPs) with genotype dosages (0/1/2). Required when known_cis_eqtls is provided.

• warmup_epochs_vae (int) – Epochs to train only the VAE before activating donor-level effects.

• warmup_epochs_G (int) – Additional epochs to train only persistent donor effects (V) after VAE warm-up, before enabling DxC effects.

• max_epochs (int) – Maximum training epochs.

• min_epochs (int) – Minimum training epochs.

• batch_size (int) – Cells per training batch.

• seed (int) – Global random seed.

• l1_weight (float) – L1 penalty weight on the DxC decoder weights.

• A_weight (float) – Penalty weight on the assignment matrix A.

• batch_norm_decoder (bool) – Apply batch normalisation in the combined decoder.

• genetics_seed (int, optional) – Separate seed for initialising genetic model parameters.

• cell_state_cis (bool) – If True, learn a cell-state-specific cis-eQTL correction per SNP. If False, learn a single cell-level correction.

• num_workers (int) – DataLoader worker processes.

• strict (bool) – Raise on non-integer values in input data (LIVI expects raw counts).

• drop_last (bool) – Drop the last incomplete batch each epoch.

• shuffle (bool) – Shuffle cells before each epoch.

• pin_memory (bool) – Pin DataLoader host memory for faster GPU transfer.

• log_every_n_steps (int) – Logging interval (steps).

• enable_progress_bar (bool) – Show the PyTorch Lightning training progress bar.

• gradient_clip_val (float, optional) – Maximum gradient norm for clipping.

• accumulate_grad_batches (int) – Gradient accumulation steps.

• limit_train_batches (int or float, optional) – Cap on batches per epoch (handy for benchmarking a fixed number of batches instead of a full pass); forwarded to pytorch_lightning.Trainer.

• enable_checkpointing (bool) – If False, skip the ModelCheckpoint callback entirely (and the return value becomes None) – useful for short benchmark/smoke runs where you don’t want checkpoint I/O.

• enable_logger (bool) – Passed through to pytorch_lightning.Trainer as logger=; set False to skip the default TensorBoardLogger.

• deterministic (bool) – Passed through to pytorch_lightning.Trainer; set True for reproducible (but slower) training.

• callbacks (list, optional) – Extra pytorch_lightning.Callback instances to attach to the Trainer (e.g. a throughput-logging callback), added alongside the ModelCheckpoint this function always sets up.

• run (bool) – If False, log the resolved configuration and return without training (dry-run / debug).

• runner (LIVIRunner, optional) – Runner instance. Uses the global runner when None.

Return type:

str | None

Returns:

str or None Path to the best model checkpoint, or None when run=False.

Examples

Basic training from DonorData:

>>> import cellink as cl
>>> cl.tl.external.configure_livi_runner("/lustre/projects/LIVI")
>>> ckpt = cl.tl.external.train_livi(
...     dd,
...     output_dir="livi_run",
...     z_dim=15,
...     n_dxc_factors=100,
...     covariates_keys=["pool", "sex"],
...     max_epochs=300,
...     warmup_epochs_vae=60,
... )

With cis-eQTL correction:

>>> ckpt = cl.tl.external.train_livi(
...     dd,
...     output_dir="livi_cis_run",
...     n_cis_snps=3000,
...     known_cis_eqtls="known_cis.tsv",
...     eqtl_genotypes="genotypes.tsv",
...     layer_key="counts",
... )