Neural networks

scVI trained straight off a bixverse store, on torch or MLX.

nn

Neural network models trained straight off a bixverse store.

scVI (Lopez et al., 2018) for batch correction, on torch or on MLX. Counts stream from the store through the Rust loader; nothing goes through AnnData. Importing this module needs neither framework; training on one needs its extra, bixverse[torch] or bixverse[mlx].

>>> import bixverse.nn as bvnn
>>> fit = bvnn.train_scvi("counts_cells.bin", cells, hvg, batch)
>>> latent = bvnn.scvi_latent(fit, "counts_cells.bin", cells, batch)

ScviParams dataclass

ScviParams(
    n_hidden: int = 128,
    n_latent: int = 10,
    n_layers: int = 1,
    dropout_rate: float = 0.1,
    gene_likelihood: Likelihood = "zinb",
    lr: float = 0.001,
    weight_decay: float = 1e-06,
    eps: float = 0.01,
    batch_size: int = 128,
    n_epochs_kl_warmup: int = 400,
    max_epochs: int | None = None,
)

scVI hyperparameters.

Defaults are scvi-tools 1.5.1's for SCVI plus its TrainingPlan, so a run with defaults on both sides is the same model and the same optimiser.

Attributes:

Name	Type	Description
n_hidden	int	Nodes per hidden layer.
n_latent	int	Latent dimensions.
n_layers	int	Hidden layers in the encoder and in the decoder.
dropout_rate	float	Encoder dropout. The decoder has none, as in scvi-tools.
gene_likelihood	Likelihood	"zinb" (scvi-tools' default) or "nb".
lr	float	Adam learning rate.
weight_decay	float	Adam L2 penalty (coupled, torch.optim.Adam style).
eps	float	Adam epsilon. scvi-tools uses 0.01, far above the usual 1e-8.
batch_size	int	Cells per training step.
n_epochs_kl_warmup	int	Epochs over which the KL weight ramps linearly from 0 to 1.
max_epochs	int | None	Training epochs. None takes scvi-tools' heuristic, min(round(20000 / n_cells * 400), 400).

ScviFit dataclass

ScviFit(
    model: Any,
    backend: Backend,
    params: ScviParams,
    genes: ndarray,
    batch_categories: ndarray,
    device: str,
    history: ScviHistory = ScviHistory(),
)

A trained scVI model plus what is needed to use it again.

Attributes:

Name	Type	Description
model	Any	The backend module (torch.nn.Module or mlx.nn.Module).
backend	Backend	Which backend model belongs to.
params	ScviParams	The hyperparameters it was trained with.
genes	ndarray	Store gene indices the model was trained on, in column order.
batch_categories	ndarray	Batch labels, in the order of the one-hot columns.
device	str	Torch device string, or "gpu" for MLX.
history	ScviHistory	Per-epoch training means. Empty for a fit wrapped around weights trained elsewhere.

ScviHistory dataclass

ScviHistory(
    loss: list[float] = list(),
    reconstruction: list[float] = list(),
    kl: list[float] = list(),
    kl_weight: list[float] = list(),
)

Per-epoch training means, one entry per epoch.

Attributes:

Name	Type	Description
loss	list[float]	Minibatch loss, reconstruction plus weighted KL.
reconstruction	list[float]	Negative log likelihood per cell.
kl	list[float]	KL divergence of z per cell.
kl_weight	list[float]	The KL weight used that epoch.

train_scvi

train_scvi(
    store: str | CellReader,
    cells: Any = None,
    genes: Any = None,
    batch: Any = None,
    *,
    params: ScviParams | None = None,
    backend: Backend = "torch",
    device: str = "auto",
    seed: int = 42,
) -> ScviFit

Train scVI on raw counts streamed from a bixverse store.

Parameters:

Name	Type	Description	Default
store	str | CellReader	Path to a cell-major .bin store, or a CellReader.	required
cells	Any	Store cell indices to train on. None takes every cell.	None
genes	Any	Store gene indices to model, typically the HVGs. None takes every gene. Library size is the sum over these genes, as in scvi-tools on a gene-subset AnnData.	None
batch	Any	One batch label per entry of cells. None is one batch.	None
params	ScviParams | None	Hyperparameters. None takes scvi-tools' defaults.	None
backend	Backend	"torch" or "mlx" (Apple Silicon).	'torch'
device	str	Torch device, or "auto" for cuda, then mps, then cpu. Ignored by MLX, which runs on the GPU.	'auto'
seed	int	Seeds initialisation, sampling and shuffling.	42

Returns:

Type	Description
ScviFit	The fitted model.

Raises:

Type	Description
ValueError	On bad indices, a batch vector of the wrong length, or a cell with no counts over genes.

scvi_latent

scvi_latent(
    fit: ScviFit,
    store: str | CellReader,
    cells: Any = None,
    batch: Any = None,
) -> ndarray

The scVI embedding: the posterior mean of z, as get_latent_representation.

The encoder never sees the batch, so batch only has to be valid, not meaningful; it is accepted so a held-out set can be scored in one pass.

Parameters:

Name	Type	Description	Default
fit	ScviFit	A fitted model.	required
store	str | CellReader	The store.	required
cells	Any	Store cell indices, None for all. Rows come back in this order.	None
batch	Any	Batch labels aligned with cells, or None for a one-batch fit.	None

Returns:

Type	Description
ndarray	(n_cells, n_latent) float32 array.

scvi_elbo

scvi_elbo(
    fit: ScviFit,
    store: str | CellReader,
    cells: Any = None,
    batch: Any = None,
    *,
    seed: int = 0,
) -> dict[str, float]

Mean per-cell negative ELBO and its two terms, at KL weight 1.

The reconstruction term decodes one draw of z per cell, as scvi-tools' get_elbo does, so it carries Monte Carlo noise.

Parameters:

Name	Type	Description	Default
fit	ScviFit	A fitted model.	required
store	str | CellReader	The store.	required
cells	Any	Store cell indices, None for all.	None
batch	Any	Batch labels aligned with cells.	None
seed	int	Seeds the draw of z.	0

Returns:

Type	Description
dict[str, float]	elbo, reconstruction and kl, each a mean over cells.