Models

class cellarium.ml.models.AmortizedOnlineNonNegativeMatrixFactorization(var_names_g: Sequence[str], k_values: list[int], r: int, latent_dim: int, total_n_cells: int, batch_size: int, solver: Literal['hals', 'fista'] = 'fista', encoder_improvement_threshold: float = 0.002, reconstruction_improvement_threshold: float = 0.002, trigger_patience: int = 3, cooldown_patience: int = 5, max_cooldown_epochs: float = 1.0, trailing_window_fraction: float = 0.25, exploration_epochs: int = 2, max_solver_iter_train: int = 50, max_solver_iter_cooldown: int = 200, init: Literal['sklearn_random', 'uniform_random'] = 'uniform_random', transformed_data_mean: None | float = None)[source]

Bases: NonNegativeMatrixFactorization, ValidateMixin

Amortized version of OnlineNonNegativeMatrixFactorization.

The idea is that OnlineNonNegativeMatrixFactorization reproduces the Kotliar cNMF results when using the nmf-torch-hals algorithm. However, this algorithm requires storage of the full loadings matrix, which is not feasible for large datasets, since the loadings matrix has shape (n_cells, n_components, n_replicates). This amortized version is a minimal change, which trains an encoder neural network to predict the loadings matrix from gene expression data. The encoder conditions on both the cell expression and the current factor matrix, using a bilinear dot product to compute loadings affinities. Replicates are kept independent by construction since each replicate’s factors are processed as an independent batch dimension with no cross-replicate attention.

Stopping uses a dual-signal AND criterion based on two independent per-k signals: the encoder loss and the NMF reconstruction loss. Both are measured as trailing-window averages over the last trailing_window_fraction of each monitoring period, which avoids contaminating estimates with the transient spike that follows each catastrophic-forgetting reset. The per-step relative improvement rate is computed for each signal independently, with separate patience counters (_encoder_trigger_patience_counters and _recon_trigger_patience_counters). Cooldown begins only when both signals have failed to improve for trigger_patience consecutive periods — whichever signal is slower is the trailing condition. Cooldown (A/B accumulation without forgetting) ends when both signals again plateau for cooldown_patience consecutive periods, or when max_cooldown_epochs elapses, whichever comes first. Thresholds are independently tunable via encoder_improvement_threshold and reconstruction_improvement_threshold.

Parameters:

var_names_g (Sequence[str])
k_values (list[int])
r (int)
latent_dim (int)
total_n_cells (int)
batch_size (int)
solver (Literal['hals', 'fista'])
encoder_improvement_threshold (float)
reconstruction_improvement_threshold (float)
trigger_patience (int)
cooldown_patience (int)
max_cooldown_epochs (float)
trailing_window_fraction (float)
exploration_epochs (int)
max_solver_iter_train (int)
max_solver_iter_cooldown (int)
init (Literal['sklearn_random', 'uniform_random'])
transformed_data_mean (None | float)

property factors_dict: dict[int, Tensor]: Return the learned factors for each k value.

online_dictionary_update(x_ng: Tensor, k: int, n_iterations: int = 100) → dict[str, Tensor][source]

Algorithm 1 from Mairal et al. [1] for online dictionary learning.

Parameters:

x_ng (Tensor) – The data.
k (int) – The value of k to run.
minibatch_indices_n – The indices of the cells in the current minibatch.
n_iterations (int)

Returns:

Loss for the encoder based on HALS targets. hals_loadings_rnk: The loadings after the HALS update, which are the targets for the encoder. encoder_loadings_rnk: The loadings predicted by the encoder before the update, which

Return type:

loss

forward(x_ng: Tensor, var_names_g: ndarray) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
obs_names_n – The names of the cells in the current minibatch (used when there are local latents).

Returns:

An empty dictionary.

Return type:

dict[str, Tensor | None]

validate(trainer: Trainer, pl_module: LightningModule, batch_idx: int, x_ng: Tensor, var_names_g: ndarray) → None[source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
trainer (Trainer)
pl_module (LightningModule)
batch_idx (int)

Returns:

An empty dictionary.

Return type:

None

class cellarium.ml.models.CellariumGPT(categorical_token_size_dict: dict[str, int], d_model: int, d_ffn: int, n_heads: int, n_blocks: int, dropout_p: float, use_bias: bool, attention_backend: Literal['flex', 'math', 'mem_efficient', 'torch'], attention_softmax_fp32: bool, loss_scale_dict: dict[str, float], initializer_range: float = 0.02, embeddings_scale: float = 1.0, attention_logits_scale: float = 1.0, output_logits_scale: float = 1.0, mup_base_d_model: int | None = None, mup_base_d_ffn: int | None = None)[source]

Bases: CellariumModel, PredictMixin, ValidateMixin

CellariumGPT model.

Parameters:

categorical_token_size_dict (dict[str, int]) – Categorical token vocabulary sizes. Must include “gene_value” and “gene_id”. Additionally, it can include experimental conditions, such as “assay” and “suspension_type”, and metadata tokens such as “cell_type”, “tissue”, “sex”, “development_stage”, and “disease”.
d_model (int) – Dimensionality of the embeddings and hidden states.
d_ffn (int) – Dimensionality of the inner feed-forward layers.
n_heads (int) – Number of attention heads.
n_blocks (int) – Number of transformer blocks.
dropout_p (float) – Dropout probability.
use_bias (bool) – Whether to use bias in the linear transformations.
attention_backend (Literal['flex', 'math', 'mem_efficient', 'torch']) – Backend for the attention computation.
attention_softmax_fp32 (bool) – Whether to use float32 for softmax computation when torch backend is used.
loss_scale_dict (dict[str, float]) – A dictionary of loss scales for each label type. These are the query tokens that are used to compute the loss.
initializer_range (float) – The standard deviation of the truncated normal initializer.
embeddings_scale (float) – Multiplier for the embeddings.
attention_logits_scale (float) – Multiplier for the attention logits.
output_logits_scale (float) – Multiplier for the output logits.
mup_base_d_model (int | None) – Base dimensionality of the model for muP.
mup_base_d_ffn (int | None) – Base dimensionality of the inner feed-forward layers for muP.

predict(token_value_nc_dict: dict[str, Tensor], token_mask_nc_dict: dict[str, Tensor], prompt_mask_nc: Tensor) → dict[str, ndarray | Tensor][source]

Parameters:

token_value_nc_dict (dict[str, Tensor]) – Dictionary of token value tensors of shape (n, c).
token_mask_nc_dict (dict[str, Tensor]) – Dictionary of token mask tensors of shape (n, c).
prompt_mask_nc (Tensor)

Returns:

Dictionary of logits tensors of shape (n, c, k).

Return type:

dict[str, ndarray | Tensor]

class cellarium.ml.models.CellariumModel[source]

Bases: Module

Base class for Cellarium ML compatible models.

abstractmethod reset_parameters() → None[source]

Reset the model parameters and buffers that were constructed in the __init__ method. Constructed means methods like torch.tensor, torch.empty, torch.zeros, torch.randn, torch.as_tensor, etc. If the parameter or buffer was assigned (e.g. from a torch.Tensor passed to the __init__) then there is no need to reset it.

Return type:: None

class cellarium.ml.models.ContrastiveMLP(n_obs: int, hidden_size: Sequence[int], embed_dim: int, temperature: float = 1.0)[source]

Bases: CellariumModel, PredictMixin

Multilayer perceptron trained with contrastive learning.

Parameters:

n_obs (int) – Number of observations in each entry (network input size).
hidden_size (Sequence[int]) – Dimensionality of the fully-connected hidden layers.
embed_dim (int) – Size of embedding (network output size).
temperature (float) – Parameter governing Normalized Temperature-scaled cross-entropy (NT-Xent) loss.

forward(x_ng: Tensor) → dict[str, Tensor][source]

Parameters:: x_ng (Tensor) – Gene counts matrix.
Returns:: A dictionary with the loss value.
Return type:: dict[str, Tensor]

predict(x_ng: Tensor) → TransformPrediction[source]

Sends (transformed) data through the model and returns outputs.

Parameters:: x_ng (Tensor) – Gene counts matrix.
Returns:: A dictionary with the embedding matrix x_ng and the corresponding variable names var_names_g.
Return type:: TransformPrediction

class cellarium.ml.models.SCANVI(classifier_type: Literal['flat', 'ontology'] = 'flat', classification_weight: float = 50.0, chunk_size: int = 100, classifier_n_hidden: list[int] | None = None, classifier_dropout_rate: float = 0.1, secondary_n_hidden: list[int] | None = None, y_prior_probs: list[float] | None = None, unlabeled_category: str = 'unknown', descendant_tensor: Tensor | None = None, cl_names: list[str] | None = None, class_counts: Series | None = None, frontier_min_cells: int = 50, propagate_class_counts: bool = False, probability_propagation_flag: bool = True, use_torch_compile: bool = False, marginalization_warn_fraction: float = 0.25, **scvi_kwargs)[source]

Bases: SingleCellVariationalInference

Single-cell ANnotation using Variational Inference (scANVI) [1].

Extends scVI with a hierarchical prior on the latent variable z:

\[c \sim p(c), \quad u \sim \mathcal{N}(0, I), \quad z \sim p(z \mid u, c), \quad x \sim p(x \mid z, s)\]

The encoder q(z|x) is identical to scVI. Three new networks are added:

Classifier q(c|z): predicts cell-type probabilities from z.
u-encoder q(u|z,c): encodes a deeper latent u given z and cell type c.
z-prior decoder p(z|u,c): provides the class-conditional prior on z.

The class variable c ranges over a partition of cell-type space. Two classifier modes are supported:

classifier_type="flat" (default): a plain MLP over a flat list of cell types (cell_type_categories), with standard cross-entropy — the classic scANVI behaviour.
classifier_type="ontology": a SOCAM-style ontology-aware MLP that emits logits over all active ontology nodes (leaf and internal) and propagates them up the Cell Ontology. The generative partition is a support-backed frontier (see compute_frontier()) computed from class_counts at construction. A coarse (internal-node) label licenses a restricted marginalization of the ELBO over the frontier leaves in its subtree, so labels at any granularity supervise the model. Optional inverse-frequency class balancing is applied to the (hierarchical) cross-entropy term.

For labeled cells the ELBO is maximized using the known cell type (or, for a coarse label, marginalized over its subtree) together with an auxiliary cross-entropy classification term. For unlabeled cells (label equal to unlabeled_category) the cell-type expectation is marginalized over the whole partition in chunks to control peak VRAM.

The classifier is batch-agnostic (it consumes only z), so with a batch-agnostic encoder configuration one can go x -> z -> c at inference time.

References:

Probabilistic harmonization and annotation of single-cell transcriptomics data with deep generative models (Xu et al., 2021).

Parameters:

classifier_type (Literal['flat', 'ontology']) – "flat" (default) for a plain MLP over cell_type_categories; "ontology" for the SOCAM-style ontology-aware classifier.
classification_weight (float) – Weight α applied to the cross-entropy loss on labeled cells.
chunk_size (int) – Number of partition members processed per chunk during marginalization. Reduce this to lower peak VRAM at the cost of slightly more computation.
classifier_n_hidden (list[int] | None) – Hidden layer widths for the classifier q(c|z) MLP. Defaults to [128].
classifier_dropout_rate (float) – Dropout rate in classifier hidden layers.
secondary_n_hidden (list[int] | None) – Hidden layer widths shared by the u-encoder q(u|z,c) and the z-prior decoder p(z|u,c). Defaults to [128].
y_prior_probs (list[float] | None) – Prior probabilities over the partition p(c). Must sum to 1 and have length equal to the partition size. If None a uniform prior is used (recommended).
unlabeled_category (str) – String label marking unlabeled cells. Default "unknown".
descendant_tensor (Tensor | None) – (ontology mode) (C, C) binary tensor, diagonal 1, where [i, j] == 1 means j is a descendant of i. Rows/columns are indexed by cl_names.
cl_names (list[str] | None) – (ontology mode) Ordered node names matching descendant_tensor.
class_counts (Series | None) – (ontology mode) pandas Series mapping node names to training cell counts. Drives both the frontier cut and (optionally) the class-balancing weights.
frontier_min_cells (int) – (ontology mode) Support threshold for the frontier cut.
propagate_class_counts (bool) – (ontology mode) If True, propagate counts up the ontology before computing balancing weights (see compute_class_weights()).
probability_propagation_flag (bool) – (ontology mode) If True, apply hierarchical propagation to the classifier logits.
use_torch_compile (bool) – (ontology mode) If True, use the ``torch.compile``d propagation functions. Defaults to False (eager) to avoid recompilation across variable batch shapes.
marginalization_warn_fraction (float) – (ontology mode) Warn (once per label) when a coarse label marginalizes over more than this fraction of the frontier — a signal that the label is nearly as expensive as unlabeled while providing little supervision.
**scvi_kwargs – All keyword arguments accepted by SingleCellVariationalInference. In flat mode, cell_type_categories defines the partition and is required.

forward(x_ng: Tensor, var_names_g: ndarray, batch_index_n: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_index_nd: Tensor | None = None, total_mrna_umis_n: Tensor | None = None, cell_type_labels_n: ndarray | None = None) → dict[source]

Compute the scANVI ELBO and auxiliary classification loss.

Parameters:

x_ng (Tensor) – Gene counts matrix [N, G].
var_names_g (ndarray) – Variable names for input validation.
batch_index_n (Tensor) – Integer batch indices [N].
continuous_covariates_nc (Tensor | None) – Continuous covariates [N, C].
categorical_covariate_index_nd (Tensor | None) – Integer categorical covariate codes [N, D].
total_mrna_umis_n (Tensor | None) – Total mRNA UMIs per cell (not log-scaled).
cell_type_labels_n (ndarray | None) – Per-cell string labels [N]. Cells equal to unlabeled_category (default "unknown") are treated as unlabeled. Finer-than- frontier labels are binned to their frontier ancestor; coarse labels marginalize over their subtree. If None, all cells are treated as unlabeled.

Returns:

A dict with loss, reconstruction_loss, kl_divergence_z, kl_divergence_u, kl_divergence_c, kl_divergence_batch, classification_loss, z_nk, and cell_type_logits_nc (classifier logits in its output space).

Return type:

dict

validate(trainer: Trainer, pl_module: LightningModule, batch_idx: int, x_ng: Tensor, var_names_g: ndarray, batch_index_n: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_index_nd: Tensor | None = None, total_mrna_umis_n: Tensor | None = None, cell_type_labels_n: ndarray | None = None) → None[source]

Log val_loss (ELBO) and a SOCAM-style cell-type accuracy on labeled cells.

The accuracy is the mean propagated probability at the resolved ground-truth node (ontology mode) or the mean predicted probability of the true class (flat mode).

Parameters:

trainer (Trainer)
pl_module (LightningModule)
batch_idx (int)
x_ng (Tensor)
var_names_g (ndarray)
batch_index_n (Tensor)
continuous_covariates_nc (Tensor | None)
categorical_covariate_index_nd (Tensor | None)
total_mrna_umis_n (Tensor | None)
cell_type_labels_n (ndarray | None)

Return type:

None

predict(x_ng: Tensor, var_names_g: ndarray, batch_index_n: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_index_nd: Tensor | None = None) → TransformClassifierPrediction[source]

Embed cells and predict cell-type probabilities.

If reconstruct_counts_on_predict is True, falls back to the parent’s count-reconstruction behaviour and omits cell-type probabilities.

Parameters:

x_ng (Tensor) – Gene counts matrix [N, G]
var_names_g (ndarray) – Variable names for input validation.
batch_index_n (Tensor) – Integer batch indices [N]. (Consumed by the encoder; may be a dummy when the encoder is configured batch-agnostic.)
continuous_covariates_nc (Tensor | None) – Continuous covariates [N, C].
categorical_covariate_index_nd (Tensor | None) – Integer categorical covariate codes [N, D].

Returns:

A dict with - x_ng reconstructed counts [N, self.reconstruction_var_names_g] if

self.reconstruct_counts_on_predict is True; otherwise it is the latent embedding matrix [N, n_latent] (misnomer).)

var_names_g Variable names corresponding to the columns of x_ng.
y_probs_nc In ontology mode y_probs_nc is the propagated probability
over all active nodes ([N, n_active], columns = active_cl_names); in flat mode it is the softmax over the partition ([N, n_partition]).
y_logits_nc (logits before softmax [N, n_partition]).

Return type:

TransformClassifierPrediction

class cellarium.ml.models.Geneformer(var_names_g: ndarray, hidden_size: int = 256, num_hidden_layers: int = 6, num_attention_heads: int = 4, intermediate_size: int = 512, hidden_act: str = 'relu', hidden_dropout_prob: float = 0.02, attention_probs_dropout_prob: float = 0.02, max_position_embeddings: int = 2048, type_vocab_size: int = 2, initializer_range: float = 0.02, position_embedding_type: str = 'absolute', layer_norm_eps: float = 1e-12, mlm_probability: float = 0.15)[source]

Bases: CellariumModel, PredictMixin

Geneformer model.

References:

Transfer learning enables predictions in network biology (Theodoris et al.).

Parameters:

var_names_g (ndarray) – The variable names schema for the input data validation.
hidden_size (int) – Dimensionality of the encoder layers and the pooler layer.
num_hidden_layers (int) – Number of hidden layers in the Transformer encoder.
num_attention_heads (int) – Number of attention heads for each attention layer in the Transformer encoder.
intermediate_size (int) – Dimensionality of the “intermediate” (often named feed-forward) layer in the Transformer encoder.
hidden_act (str) – The non-linear activation function (function or string) in the encoder and pooler. If string, "gelu", "relu", "silu" and "gelu_new" are supported.
hidden_dropout_prob (float) – The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
attention_probs_dropout_prob (float) – The dropout ratio for the attention probabilities.
max_position_embeddings (int) – The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
type_vocab_size (int) – The vocabulary size of the token_type_ids passed when calling transformers.BertModel.
initializer_range (float) – The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
position_embedding_type (str) – Type of position embedding. Choose one of "absolute", "relative_key", "relative_key_query". For positional embeddings use "absolute". For more information on "relative_key", please refer to Self-Attention with Relative Position Representations (Shaw et al.). For more information on "relative_key_query", please refer to Method 4 in Improve Transformer Models with Better Relative Position Embeddings (Huang et al.).
layer_norm_eps (float) – The epsilon used by the layer normalization layers.
mlm_probability (float) – Ratio of tokens to mask for masked language modeling loss.

forward(x_ng: Tensor, var_names_g: ndarray) → dict[str, Tensor][source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.

Returns:

A dictionary with the loss value.

Return type:

dict[str, Tensor]

predict(x_ng: Tensor, var_names_g: ndarray, output_hidden_states: bool = True, output_attentions: bool = True, output_input_ids: bool = True, output_attention_mask: bool = True, feature_activation: list[str] | None = None, feature_deletion: list[str] | None = None, feature_map: dict[str, int] | None = None) → dict[str, Tensor | ndarray][source]

Send (transformed) data through the model and return outputs. Optionally perform in silico perturbations and masking.

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
output_hidden_states (bool) – Whether to return all hidden-states.
output_attentions (bool) – Whether to return all attentions.
output_input_ids (bool) – Whether to return input ids.
output_attention_mask (bool) – Whether to return attention mask.
feature_activation (list[str] | None) – Specify features whose expression should be set to > max(x_ng) before tokenization (top rank).
feature_deletion (list[str] | None) – Specify features whose expression should be set to zero before tokenization (remove from inputs).
feature_map (dict[str, int] | None) – Specify a mapping for input tokens, to be applied before model.

Returns:

A dictionary with the inference results.

Return type:

dict[str, Tensor | ndarray]

Note

In silico perturbations can be achieved in one of three ways:

Use feature_map to replace a feature token with MASK (1) or PAD (0) (e.g. feature_map={"ENSG0001": 1} will replace var_names_g feature ENSG0001 with a MASK token).
Use feature_deletion to remove a feature from the cell’s inputs, which instead of adding a PAD or MASK token, will allow another feature to take its place (e.g. feature_deletion=["ENSG0001"] will remove var_names_g feature ENSG0001 from the input, and allow a new feature token to take its place).
Use feature_activation to move a feature all the way to the top rank position in the input (e.g. feature_activation=["ENSG0001"] will make var_names_g feature ENSG0001 the first in rank order. Multiple input features will be ranked according to their order in the input list).

Number (2) and (3) are described in the Geneformer paper under “In silico perturbation” in the Methods section.

class cellarium.ml.models.HVGSeuratV3(var_names_g: ndarray, n_top_genes: int | list[int], n_batch: int = 1, flavor: Literal['seurat_v3', 'seurat_v3_paper'] = 'seurat_v3', use_batch_key: bool = False, span: float = 0.3, output_path: str | None = 'hvg_seurat_v3_output.csv', batch_n_cell_minimum: int = 2, n_batch_minimum: int = 1, accumulator_dtype: Literal['float32', 'float64'] = 'float64')[source]

Bases: CellariumModel

Compute highly variable genes using the Seurat v3 method in two Lightning epochs.

Epoch 0 — streams data to accumulate per-batch mean and variance. The sums are of the counts minus a per-batch, per-gene shift (the mean of the first cells seen from that batch), so that the variance is not computed as a difference of two large, nearly equal numbers.

Between epochs — fits a LOESS model of log10(var) ~ log10(mean) per batch to estimate a regularized standard deviation and per-cell clip value.

Epoch 1 — streams data again, clips counts per cell at the batch-level clip_val, and accumulates the sum of squared deviations of the clipped counts from the batch mean.

After epoch 1 (on_train_epoch_end) — computes normalized variance per batch, ranks genes, and writes self.hvg_dfs (and optionally a CSV file at output_path).

Both epochs use only sums of positive terms or of shifted values, so the accumulators can be float32 (see accumulator_dtype) with little loss of accuracy.

The flavor argument controls how genes are ranked across batches when n_batch > 1, matching the two multi-batch modes in Scanpy:

"seurat_v3" (default) — sorts by highly_variable_rank ascending first, then by highly_variable_nbatches descending as tiebreaker. Genes with the lowest median rank across batches are preferred. This matches Scanpy’s flavor='seurat_v3'.
"seurat_v3_paper" — sorts by highly_variable_nbatches descending first, then by highly_variable_rank ascending as tiebreaker. Genes that are highly variable in more batches are preferred, regardless of their exact rank. This matches Scanpy’s flavor='seurat_v3_paper'.

Usage:

model = HVGSeuratV3(var_names_g=gene_names, n_top_genes=2000, n_batch=4,
                    flavor="seurat_v3_paper", span=0.3)
trainer = pl.Trainer(max_epochs=2)
trainer.fit(module, datamodule)
df = model.hvg_df  # pandas DataFrame, Scanpy-compatible columns

Parameters:

var_names_g (np.ndarray) – Array of gene names, length n_genes.
n_top_genes (int | list[int]) – Number of highly variable genes to select. Can be a list of ints to produce multiple gene sets in a single training run.
n_batch (int) – Number of batches (use 1 when no batch information is given).
flavor (Literal['seurat_v3', 'seurat_v3_paper']) – Multi-batch gene ranking strategy. "seurat_v3_paper" (default) prioritises batch consistency; "seurat_v3" prioritises median rank.
use_batch_key (bool) – Whether to expect a batch_index_n batch key in the dataloader output. If False, the model will ignore batch keys and treat all data as a single batch.
span (float) – LOESS span (fraction of data used per local fit). Default 0.3.
output_path (str | None) – If given, the result DataFrame is written to this CSV filepath after training. (Ends with .csv). If None, no file is written, but you really want to write this output, as it would require manually calling _compute_hvg_df() later.
batch_n_cell_minimum (int) – Minimum number of cells required for a batch to be considered valid.
n_batch_minimum (int) – Minimum number of batches in which the gene is in n_top_genes highly variable for a gene to make the final list.
accumulator_dtype (Literal['float32', 'float64']) – The dtype of the accumulators, "float64" (default) or "float32". Use "float32" on devices without float64 support (MPS).

on_train_epoch_start(trainer: Trainer) → None[source]

Cache the current epoch number so forward() can branch on it.

Parameters:: trainer (Trainer)
Return type:: None

on_train_start(trainer: Trainer) → None[source]

Validation: if in a distributed setting, use DDP with broadcast_buffers=False.

Parameters:: trainer (Trainer)
Return type:: None

on_train_epoch_end(trainer: Trainer) → None[source]

Implement the two-step Seurat v3 HVG method using hooks at the end of first and second epochs. After epoch 0, reduce buffers and fit LOESS to set clip_val_bg and reg_std_bg. After epoch 1, reduce buffers and compute hvg_df.

Parameters:: trainer (Trainer)
Return type:: None

class cellarium.ml.models.IncrementalPCA(var_names_g: ndarray, n_components: int, svd_lowrank_niter: int = 2, perform_mean_correction: bool = True)[source]

Bases: CellariumModel, PredictMixin

Distributed and Incremental PCA.

References:

Parameters:

var_names_g (ndarray) – The variable names schema for the input data validation.
n_components (int) – Number of principal components.
svd_lowrank_niter (int) – Number of iterations for the low-rank SVD algorithm.
perform_mean_correction (bool) – If True then the mean correction is applied to the update step. If False then the data is assumed to be centered and the mean correction is not applied to the update step.

forward(x_ng: Tensor, var_names_g: ndarray) → dict[str, Tensor | None][source]

Incrementally update partial SVD with new data.

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.

Returns:

An empty dictionary.

Return type:

dict[str, Tensor | None]

on_train_epoch_end(trainer: Trainer) → None[source]

Merge partial SVD results from parallel processes at the end of the epoch.

Merging SVDs is performed hierarchically. At each merging level, the leading process (even rank) merges its SVD with the trailing process (odd rank). The trailing process discards its SVD and is terminated. The leading process continues to the next level. This process continues until only one process remains. The final SVD is stored on the remaining process.

The number of levels (hierarchy depth) scales logarithmically with the number of processes.

Parameters:: trainer (Trainer)
Return type:: None

property explained_variance_k: Tensor

The amount of variance explained by each of the selected components. The variance estimation uses x_size degrees of freedom.

Equal to n_components largest eigenvalues of the covariance matrix of input data.

property components_kg: Tensor: Principal axes in feature space, representing the directions of maximum variance in the data. Equivalently, the right singular vectors of the centered input data, parallel to its eigenvectors. The components are sorted by decreasing explained_variance_k.

predict(x_ng: Tensor, var_names_g: ndarray) → TransformPrediction[source]

Centering and embedding of the input data x_ng into the principal component space.

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.

Returns:

x_ng: (misnomer) Embedding of the input data into the principal component space.
var_names_g: The list of variable names for the output data.

Return type:

A dictionary with the following keys

class cellarium.ml.models.LogisticRegression(n_obs: int, var_names_g: ndarray, y_categories: ndarray, W_prior_scale: float = 1.0, W_init_scale: float = 1.0, seed: int = 0, log_metrics: bool = True)[source]

Bases: CellariumModel, PredictMixin, ValidateMixin

Logistic regression model.

Parameters:

n_obs (int) – Number of observations.
var_names_g (ndarray) – The variable names schema for the input data validation.
y_categories (ndarray) – The categories for the target data.
W_prior_scale (float) – The scale of the Laplace prior for the weights.
W_init_scale (float) – Initialization scale for the W_gc parameter.
seed (int) – Random seed used to initialize parameters.
log_metrics (bool) – Whether to log the histogram of the W_gc parameter.

forward(x_ng: Tensor, var_names_g: ndarray, y_n: Tensor, y_categories: ndarray) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – The input data.
var_names_g (ndarray) – The variable names for the input data.
y_n (Tensor) – The target data.
y_categories (ndarray) – The categories for the input target data.

Returns:

A dictionary with the loss value.

Return type:

dict[str, Tensor | None]

predict(x_ng: Tensor, var_names_g: ndarray) → TransformPrediction[source]

Predict the target logits.

Parameters:

x_ng (Tensor) – The input data.
var_names_g (ndarray) – The variable names for the input data.

Returns:

A dictionary with the target logits as x_ng (misnomer).

Return type:

TransformPrediction

class cellarium.ml.models.OnePassMeanVarStd(var_names_g: ndarray, algorithm: Literal['naive', 'shifted_data'] = 'naive', n_batch: int = 1, output_path: str | None = 'onepass_mean_var_std_output.csv')[source]

Bases: CellariumModel

Calculate the mean, variance, and standard deviation of the data in one pass (epoch) using running sums and running squared sums.

Tracks per-batch statistics. Use n_batch=1 when there is no meaningful batch structure. After training, batch_mean_bg and batch_var_bg give per-batch per-gene statistics suitable for passing to get_highly_variable_genes.

References:

Algorithms for calculating variance.

Parameters:

var_names_g (ndarray) – The variable names schema for the input data validation.
algorithm (Literal['naive', 'shifted_data']) – "naive" (default) or "shifted_data" (numerically stable).
n_batch (int) – Number of batches. Use 1 to reproduce ``batch_key``=None behavior.
output_path (str | None) – Optional path to save a summary CSV at the end of training, containing the mean and variance for each gene. If None, no CSV will be saved.

forward(x_ng: Tensor, var_names_g: ndarray, batch_index_n: Tensor | None = None) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – Variable names in the input data.
batch_index_n (Tensor | None) – Optional batch indices for each cell, required if n_batch > 1.

Returns:

An empty dictionary.

Return type:

dict[str, Tensor | None]

property batch_mean_bg: Tensor: Per-batch mean, shape (n_batch, n_genes).

property batch_var_bg: Tensor: Per-batch population variance, shape (n_batch, n_genes).

on_train_end(trainer: Trainer) → None[source]

Save a summary output CSV so we need not load the checkpoint downstream.

Parameters:: trainer (Trainer)
Return type:: None

class cellarium.ml.models.OnlineNonNegativeMatrixFactorization(var_names_g: Sequence[str], k_values: list[int], r: int, solver: Literal['fista', 'hals'] = 'fista', init: Literal['sklearn_random', 'uniform_random'] = 'uniform_random', transformed_data_mean: None | float = None, n_cells_total: int | None = None)[source]

Bases: NonNegativeMatrixFactorization

Use the online NMF algorithm of Mairal et al. [1] to factorize the count matrix into a dictionary of gene expression programs and a matrix of cell program loadings.

References:

Online learning for matrix factorization and sparse coding. Mairal, Bach, Ponce, Sapiro. JMLR 2009.

Parameters:

var_names_g (Sequence[str]) – The variable names schema for the input data: should be highly variable genes.
k_values (list[int]) – A list of the number of gene expression programs to infer.
r (int) – Number of NMF replicates (for consensus).
solver (Literal['fista', 'hals']) – The solver for the loadings and dictionary updates, in [“fista”, “hals”]. "fista" (FISTA for non-negative least squares) recomputes each minibatch’s loadings from scratch, so memory does not grow with the number of cells, and it supports multi-process training. "hals" (hierarchical alternating least squares, following nmf-torch) stores the loadings of every cell (see stores_loadings), requires n_cells_total and obs_names_n, and does not support multi-process training.
init (Literal['sklearn_random', 'uniform_random']) – The initialization method to use for the NMF factors, in [“sklearn_random”, “uniform_random”].
transformed_data_mean (None | float) – The mean of the transformed data, used for initialization if an only if init is “sklearn_random”.
n_cells_total (int | None) – The total number of cells in the dataset, used if and only if solver is “hals”.

property stores_loadings: bool: Whether the loadings of every cell are stored across minibatches, as buffers of shape (r, n_cells_total, k). True for the "hals" solver, which warm-starts from them; False for "fista", which recomputes each minibatch’s loadings from scratch.

property factors_dict: dict[int, Tensor]: Return the learned factors for each k value.

online_dictionary_update(x_ng: Tensor, k: int, minibatch_indices_n: Tensor | None) → None[source]

Algorithm 1 from Mairal et al. [1] for online dictionary learning.

Parameters:

x_ng (Tensor) – The data.
k (int) – The value of k to run.
minibatch_indices_n (Tensor | None) – The indices of the cells in the current minibatch.

Return type:

None

forward(x_ng: Tensor, var_names_g: ndarray, obs_names_n: ndarray | None = None) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
obs_names_n (ndarray | None) – The names of the cells in the current minibatch (used when there are local latents).

Returns:

An empty dictionary.

Return type:

dict[str, Tensor | None]

class cellarium.ml.models.PredictMixin[source]

Bases: object

Abstract mixin class for models that can perform prediction.

abstractmethod predict(*args: Any, **kwargs: Any) → Mapping[str, Any][source]

Perform prediction.

Parameters:

args (Any)
kwargs (Any)

Return type:

Mapping[str, Any]

class cellarium.ml.models.ProbabilisticPCA(n_obs: int, var_names_g: ndarray, n_components: int, ppca_flavor: Literal['marginalized', 'linear_vae'], mean_g: Tensor | None = None, W_init_scale: float = 1.0, sigma_init_scale: float = 1.0, seed: int = 0)[source]

Bases: CellariumModel, PredictMixin

Probabilistic PCA implemented in Pyro.

Two flavors of probabilistic PCA are available - marginalized pPCA [1] and linear VAE [2].

References:

Parameters:

n_obs (int) – Number of cells.
var_names_g (ndarray) – The variable names schema for the input data validation.
n_components (int) – Number of principal components.
ppca_flavor (Literal['marginalized', 'linear_vae']) – Type of the PPCA model. Has to be one of marginalized or linear_vae.
mean_g (Tensor | None) – Mean gene expression of the input data. If None then the mean is set to a learnable parameter.
W_init_scale (float) – Scale of the random initialization of the W_kg parameter.
sigma_init_scale (float) – Initialization value of the sigma parameter.
seed (int) – Random seed used to initialize parameters.

forward(x_ng: Tensor, var_names_g: ndarray) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.

Returns:

A dictionary with the loss value.

Return type:

dict[str, Tensor | None]

predict(x_ng: Tensor, var_names_g: ndarray) → TransformPrediction[source]

Centering and embedding of the input data x_ng into the principal component space.

Note

Gradients are disabled, used for inference only.

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.

Returns:

x_ng: (misnomer) Embedding of the input data into the principal component space.
var_names_g: The list of the variable names corresponding to the embedding features.

Return type:

A dictionary with the following keys

property L_k: Tensor: Vector with elements given by the PC eigenvalues.

Note

Gradients are disabled, used for inference only.

property U_gk: Tensor: Principal components corresponding to eigenvalues L_k.

Note

Gradients are disabled, used for inference only.

property W_variance: float: Note

Gradients are disabled, used for inference only.

property sigma_variance: float: Note

Gradients are disabled, used for inference only.

class cellarium.ml.models.TDigest(var_names_g: ndarray)[source]

Bases: CellariumModel

Compute an approximate non-zero histogram of the distribution of each gene in a batch of cells using t-digests.

References:

Computing Extremely Accurate Quantiles Using T-Digests (Dunning et al.).

Parameters:: var_names_g (ndarray) – The variable names schema for the input data validation.

forward(x_ng: Tensor, var_names_g: ndarray) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.

Returns:

An empty dictionary.

Return type:

dict[str, Tensor | None]

property median_g: Tensor: Median of the data.

class cellarium.ml.models.TestMixin[source]

Bases: object

Abstract mixin class for models that can perform testing.

abstractmethod test(trainer: Trainer, pl_module: LightningModule, batch_idx: int, *args: Any, **kwargs: Any) → None[source]

Perform testing.

Parameters:

trainer (Trainer)
pl_module (LightningModule)
batch_idx (int)
args (Any)
kwargs (Any)

Return type:

None

class cellarium.ml.models.ValidateMixin[source]

Bases: object

Mixin class for models that can perform validation.

validate(trainer: Trainer, pl_module: LightningModule, batch_idx: int, *args: Any, **kwargs: Any) → None[source]

Default validation method for models. This method logs the validation loss to TensorBoard. Override this method to customize the validation behavior.

Parameters:

trainer (Trainer)
pl_module (LightningModule)
batch_idx (int)
args (Any)
kwargs (Any)

Return type:

None

class cellarium.ml.models.SingleCellVariationalInference(var_names_g: Sequence[str], encoder: dict[str, list[dict] | dict | bool], decoder: dict[str, list[dict] | dict | bool], n_batch: int = 0, n_latent: int = 10, n_continuous_cov: int = 0, n_cats_per_cov: list[int] | None = None, dropout_rate: float = 0.1, dispersion: Literal['gene', 'gene-batch', 'gene-label', 'gene-cell'] = 'gene', log_variational: bool = True, normalizetotal_variational: bool = False, gene_likelihood: Literal['zinb', 'nb', 'poisson'] = 'nb', latent_distribution: Literal['normal', 'ln'] = 'normal', batch_embedded: bool = False, batch_representation_sampled: bool = False, n_latent_batch: int | None = None, z_kl_weight_max: float = 1.0, batch_kl_weight_max: float = 0.0, input_gene_dropout_rate: float = 0.0, use_batch_norm: Literal['encoder', 'decoder', 'none', 'both'] = 'both', use_layer_norm: Literal['encoder', 'decoder', 'none', 'both'] = 'none', kl_warmup_epochs: int | None = 400, kl_warmup_steps: int | None = None, kl_annealing_start: float = 0.0, use_size_factor_key: bool = False, reconstruct_counts_on_predict: bool = False, reconstruction_var_names_g: ndarray | list | None = None, reconstruction_transform_batch: None | int | str = 0, reconstruction_n_latent_samples: int = 30, reconstruction_use_latent_mean: bool = False, reconstruction_use_importance_sampling: bool = False, reconstructed_library_size: int = 10000, reconstruction_transform_categorical_covariates: list[int] | None = None, use_flow: bool = False, flow_hidden_features: list[int] = [64, 64], cell_type_categories: list[str] | None = None, ontology_distance_matrix: DataFrame | None = None, val_cell_type_classifier_reservoir_size: int = 50000, use_torch_compile: bool = False)[source]

Bases: CellariumModel, PredictMixin, ValidateMixin

Flexible version of single-cell variational inference (scVI) [1] re-implemented in Cellarium ML.

References:

Deep generative modeling for single-cell transcriptomics (Lopez et al.).

Parameters:

var_names_g (Sequence[str]) – The variable names schema for the input data validation.
encoder (dict[str, list[dict] | dict | bool]) – Dict specifying the encoder configuration.
decoder (dict[str, list[dict] | dict | bool]) – Dict specifying the decoder configuration.
n_latent (int) – Dimension of the latent space.
n_batch (int) – Number of total batches in the dataset.
batch_representation_sampled (bool) – True to sample latent batch from a distribution.
n_continuous_cov (int) – Number of continuous covariates.
n_cats_per_cov (list[int] | None) – A list of integers containing the number of categories for each categorical covariate.
dropout_rate (float) – Dropout rate for hidden units in the encoder only.
input_gene_dropout_rate (float) – Gene dropout rate for input data that goes into the encoder during training.
dispersion (Literal['gene', 'gene-batch', 'gene-label', 'gene-cell']) –
Flexibility of the dispersion parameter when gene_likelihood is either "nb" or "zinb". One of the following:
- "gene": parameter is constant per gene across cells.
- "gene-batch": parameter is constant per gene per batch.
- "gene-label": parameter is constant per gene per label.
- "gene-cell": parameter is constant per gene per cell.
log_variational (bool) – If True, use log1p() on input data before encoding for numerical stability (not normalization).
normalizetotal_variational (bool) – If True, normalize input data to a total of 10,000 counts per cell before encoding for numerical stability. Happens before log_variational if both are True.
gene_likelihood (Literal['zinb', 'nb', 'poisson']) – Distribution to use for reconstruction in the generative process. One of the following: * "nb": NegativeBinomial. * "zinb": ZeroInflatedNegativeBinomial. (not implemented) * "poisson": Poisson.
latent_distribution (Literal['normal', 'ln']) – Distribution to use for the latent space. One of the following: * "normal": isotropic normal. * "ln": logistic normal with normal params N(0, 1). (not implemented)
use_batch_norm (Literal['encoder', 'decoder', 'none', 'both']) – Specifies where to use BatchNorm1d in the model. One of the following: * "none": don’t use batch norm in either encoder(s) or decoder. * "encoder": use batch norm only in the encoder(s). * "decoder": use batch norm only in the decoder. * "both": use batch norm in both encoder(s) and decoder.
use_layer_norm (Literal['encoder', 'decoder', 'none', 'both']) –
Specifies where to use LayerNorm in the model. One of the following:
- "none": don’t use layer norm in either encoder(s) or decoder.
- "encoder": use layer norm only in the encoder(s).
- "decoder": use layer norm only in the decoder.
- "both": use layer norm in both encoder(s) and decoder.
Note: only one of use_batch_norm or use_layer_norm should be specified.
use_size_factor_key (bool) – If True, use the obs column as defined by the size_factor_key parameter in the model’s setup_anndata method as the scaling factor in the mean of the conditional distribution. If False, the observed library size (log of the sum of counts per cell) is used. Should be False.
reconstruct_counts_on_predict (bool) – Changes the behavior of predict(). True will reconstruct gene expression count data, False will return the latent representations
reconstruction_var_names_g (ndarray | list | None) – List of var_names to be reconstructed (outputs are dense matrices)
reconstruction_transform_batch (None | int | str) – None will reconstruct in the original data batch. This is like imputation or smoothing. An integer will reconstruct counts in the specified batch index. The string “mean” will reconstruct counts in the first 10 batches and return the mean.
reconstruction_n_latent_samples (int) – Number of latent samples to use for reconstruction. Each latent sample will be used to compute the mean of the generative distribution, and the final output will be the mean of those.
reconstruction_use_latent_mean (bool) – True to use the mean of the latent distribution rather than sampling.
reconstruction_use_importance_sampling (bool) – True to use importance sampling weighted by each latent sample’s likelihood.
reconstructed_library_size (int) – The library size to use for the reconstruction, common to all cells.
use_flow (bool) – If True, use a Neural Spline Flow (NSF) as the prior on the latent space instead of the standard normal N(0, I). The flow is unconditional (batch-blind) and is jointly trained with the encoder/decoder via an MC-KL estimate: E_q[log q(z|x) - log p_flow(z)].
flow_hidden_features (list[int]) – Hidden layer widths for the NSF. Only used when use_flow=True.
cell_type_categories (list[str] | None) – Ordered list of CL ID strings (e.g. ["CL:0000540", ...]) that matches adata.obs[cell_type_col].cat.categories exactly (same order), so that integer codes from .cat.codes map directly to rows of the internal distance buffer. Required if ontology_distance_matrix is provided.
ontology_distance_matrix (DataFrame | None) – Square pandas.DataFrame with CL ID strings as both index and columns, as returned by compute_cl_distance_matrix(). The constructor will slice and reorder this to match cell_type_categories. Enables the frequency-weighted Spearman correlation metric (val_ontology_spearman) during validation. Not saved to checkpoints.
val_cell_type_classifier_reservoir_size (int) – Maximum number of cells to retain per split (train / test) for the logistic regression cell type classifier. Reservoir sampling is used so this bound is respected regardless of validation set size. Train cells are drawn from even-numbered validation batches; test cells from odd-numbered batches. Ignored when cell_type_categories is not provided. Default 50_000.
batch_embedded (bool)
n_latent_batch (int | None)
z_kl_weight_max (float)
batch_kl_weight_max (float)
kl_warmup_epochs (int | None)
kl_warmup_steps (int | None)
kl_annealing_start (float)
reconstruction_transform_categorical_covariates (list[int] | None)
use_torch_compile (bool)

batch_representation_from_batch_index(batch_index_n: Tensor, use_mean_though_sampling: bool = False) → Tensor[source]

Compute a batch representation from batch indices.

If self.batch_embedded is False, the batch representation will be one-hot (like scvi-tools) If self.batch_embedded is True:

If self.batch_representation_sampled is True, the batch representation is sampled from a normal distribution If self.batch_representation_sampled is False, the batch representation is a point estimate

Parameters:

batch_index_n (Tensor)
use_mean_though_sampling (bool)

Return type:

Tensor

categorical_onehot_from_categorical_index(categorical_covariate_index_nd: Tensor | None) → Tensor | None[source]

Compute one-hot encoding of categorical covariates from integer category indices.

Parameters:: categorical_covariate_index_nd (Tensor | None) – a tensor of shape (n, n_categorical_covariates)
Return type:: Tensor | None

inference(x_ng: Tensor, batch_nb: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_np: Tensor | None = None)[source]

High level inference method. Runs the inference (encoder) model.

Parameters:

x_ng (Tensor)
batch_nb (Tensor)
continuous_covariates_nc (Tensor | None)
categorical_covariate_np (Tensor | None)

generative(z_nk: Tensor, library_size_n1: Tensor, batch_nb: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_np: Tensor | None = None) → dict[str, Distribution][source]

Runs the generative model.

Parameters:

z_nk (Tensor)
library_size_n1 (Tensor)
batch_nb (Tensor)
continuous_covariates_nc (Tensor | None)
categorical_covariate_np (Tensor | None)

Return type:

dict[str, Distribution]

forward(x_ng: Tensor, var_names_g: ndarray, batch_index_n: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_index_nd: Tensor | None = None, total_mrna_umis_n: Tensor | None = None)[source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
batch_index_n (Tensor) – Batch indices of input cells as integers.
continuous_covariates_nc (Tensor | None) – Continuous covariates for each cell (c-dimensional).
categorical_covariate_index_nd (Tensor | None) – Categorical covariates for each cell (d-dimensional). Integer membership categorical codes.
total_mrna_umis_n (Tensor | None) – Total mRNA UMIs for each cell (not log scaled) if this should be used.

Returns:

“loss”: The total loss value.
”reconstruction_loss”: The reconstruction loss value.
”kl_divergence_z”: The KL divergence for the latent variable z.
”z_nk”: The latent variable z.

Return type:

A dictionary with keys

predict(x_ng: Tensor, var_names_g: ndarray, batch_index_n: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_index_nd: Tensor | None = None) → TransformPrediction[source]

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
batch_index_n (Tensor) – Batch indices of input cells as integers.
continuous_covariates_nc (Tensor | None) – Continuous covariates for each cell (c-dimensional).
categorical_covariate_index_nd (Tensor | None) – Categorical covariates for each cell (d-dimensional where d is number of categorical variables). Values are integer membership categorical codes.

Returns:

x_ng:
- If self.reconstruct_counts_on_predict is False:
  
  (x_ng is a notational misnomer) Embedding of the input data into the scVI latent space,
  typically referred to as z_nk.
- If self.reconstruct_counts_on_predict is True:
  
  (x_ng is a notational misnomer) Reconstruction of the input data: x_ng', which
  may have a different number of genes depending on self.reconstruction_var_names_g.

Return type:

A dictionary with the following keys

reconstruct(x_ng: Tensor, var_names_g: ndarray, gene_inds: ndarray | list[int], batch_index_n: Tensor, continuous_covariates_nc: Tensor | None = None, categorical_covariate_index_nd: Tensor | None = None, transform_batch: str | int | None = None, transform_categorical_covariates: list[int] | None = None, use_latent_mean: bool = False, n_latent_samples: int = 1000, use_importance_sampling: bool = False, reconstructed_library_size: float = 10000) → TransformPrediction[source]

Reconstruct the data using the VAE, optionally transforming the batch.

Note: scvi-tools uses the following strategy -

for each transform_batch, put the data through the encoder (no dropout)
sample n_latent_samples times to get several z values
take the mean of the generative distribution
obtain tensor shape [n_transform_batches, n_latent_samples, n_cells, n_genes]
take a mean over the batches dimension

(- they optionally use importance sampling based on sampled z likelihoods) - take a (weighted?) mean over the n_latent_samples dimension

Parameters:

x_ng (Tensor) – Gene counts matrix.
var_names_g (ndarray) – The list of the variable names in the input data.
gene_inds (ndarray | list[int]) – The indices of the genes from var_names_g to be reconstructed. Output order preserves this order.
batch_index_n (Tensor) – Batch indices of input cells as integers.
continuous_covariates_nc (Tensor | None) – Continuous covariates for each cell (c-dimensional).
categorical_covariate_index_nd (Tensor | None) – Categorical covariates for each cell (d-dimensional where d is the number of categorical variables). Used for the encoder; also used for the decoder when transform_categorical_covariates is None.
transform_batch (str | int | None) – If not None, transform the batch to this index before reconstruction.
transform_categorical_covariates (list[int] | None) – A list of integer category indices, one per categorical covariate (in the same order as n_cats_per_cov), to fix for all cells during decoding. Must be supplied when transform_batch is not None and the decoder uses categorical covariates — use enumerate_observed_batch_covariate_combinations() to identify valid combinations present in the training data. If None, the per-cell observed categorical covariates are passed to the decoder.
use_latent_mean (bool) – If True, use the mean of the latent distribution instead of sampling.
n_latent_samples (int) – The number of latent samples to use for reconstruction.
use_importance_sampling (bool) – True to use importance sampling for the reconstruction, weighting each sample of the latent by its likelihood.
reconstructed_library_size (float) – The library size to use for the reconstruction, common to all cells.

Returns:

x_ng: Model’s reconstruction of the input data, possibly de-batched. The notational misnomer
here is that “g” no longer stands for all genes, but the genes in gene_inds.

Return type:

A dictionary with the following keys

class cellarium.ml.models.SOCAM(n_obs: int, var_names_g: ndarray, descendant_tensor: Tensor, cl_names: list[str], cl_name_subset: list[str] | None = None, probability_propagation_flag: bool = True, W_prior_scale: float = 0.01, W_init_scale: float = 1.0, seed: int = 0, log_metrics: bool = True, include_ancestors_of_cl_name_subset: bool = True, class_counts: Series | None = None, propagate_class_counts: bool = False)[source]

Bases: CellariumModel, PredictMixin, ValidateMixin

Logistic regression model for cell type ontology classification. [1]

References

[1] Magre et al., “A supervised ontology-aware cell annotation method for single-cell: transcriptomic data”, bioRxiv, 2026. https://doi.org/10.64898/2026.01.13.699356

Parameters:

n_obs (int) – Number of observations in the dataset (used to scale the cross-entropy loss).
var_names_g (ndarray) – The variable-name schema for the input data; used for validation.
output_categories – Total number of target categories expected at prediction/validation time. Used when the trained model has fewer categories than the final output space.
descendant_tensor (Tensor) – Binary (0/1) tensor of shape (n_categories, n_categories) defining the descendant relationships between categories. Row i contains ones for all categories considered descendants of category i (plus self). Used for probability-propagation.
cl_names (list[str]) – Full list of category identifiers matching the rows/columns of descendant_tensor.
cl_name_subset (list[str] | None) – Optional list of category names (from cl_names) to restrict training and prediction to. The list is sorted internally so order does not matter. When None, all categories are used.
probability_propagation_flag (bool) – If True, applies hierarchical probability propagation before predicting the output distribution.
W_prior_scale (float) – Scale (b) parameter of the Laplace prior on the weight matrix W_gc.
W_init_scale (float) – Standard deviation for initializing W_gc.
seed (int) – Random seed used to initialize parameters.
log_metrics (bool) – If True, logs weight histograms (TensorBoard) during training. If True, logs weight histograms (TensorBoard) during training.
class_counts (Series | None) – Optional pandas Series mapping class names to cell counts in the training data. Classes absent from the Series (typically pure-ancestor nodes with no direct cell labels) and classes with a count of zero are both treated as unlabeled and receive a neutral weight of 1.0, which has no effect on the loss since no cells carry those labels. Inverse-frequency weights are computed and normalized to mean 1 over the nonzero-count classes only. Negative counts raise a ValueError. Extra Series entries not in the active category set are ignored. When None, all classes are weighted equally.
propagate_class_counts (bool) – If True and class_counts is provided, propagate raw counts up the ontology before computing weights. Each node’s effective count becomes its own direct count plus the sum of all descendant counts. This down-weights ancestor nodes (they accumulate large effective counts) relative to rare leaf nodes. A pure-ancestor node with zero direct labels but nonzero descendant counts will receive a real inverse-frequency weight rather than the neutral 1.0 it would otherwise get. Has no effect when class_counts is None.
include_ancestors_of_cl_name_subset (bool)

forward(x_ng: Tensor, var_names_g: ndarray, cl_names_n: ndarray) → dict[str, Tensor | None][source]

Parameters:

x_ng (Tensor) – The input data.
var_names_g (ndarray) – The variable names for the input data.
cl_names_n (ndarray) – Array of length n containing a category name string (from self.cl_names) for each cell. When self.cl_name_subset is set, every label must be a member of that subset.

Returns:

A dictionary with the loss value.

Return type:

dict[str, Tensor | None]

validate(trainer: Trainer, pl_module: LightningModule, batch_idx: int, x_ng: Tensor, var_names_g: ndarray, cl_names_n: ndarray) → None[source]

Default validation method for models. This method logs the validation loss to TensorBoard. Override this method to customize the validation behavior.

Parameters:

trainer (Trainer)
pl_module (LightningModule)
batch_idx (int)
x_ng (Tensor)
var_names_g (ndarray)
cl_names_n (ndarray)

Return type:

None

predict(x_ng: Tensor, var_names_g: ndarray) → ClassifierPrediction[source]

Predict the target logits.

Parameters:

x_ng (Tensor) – The input data.
var_names_g (ndarray) – The variable names for the input data.

Returns:

A dictionary with the target logits in y_logits_nc and the predicted probabilities in y_probs_nc (possibly propagated). Output tensors have shape (n, n_active_cats).

Return type:

ClassifierPrediction

class cellarium.ml.models.StreamingHyperplaneGeometricSketch(var_names_g: ndarray, n_bits: int = 12, max_cells_per_bucket: int = 100, min_cells_per_bucket: int = 1, min_metadata_diversity: int = 1, store_cell_data: bool = True, projector: Module | None = None, limit_input_to_top_pcs: int | None = None, seed: int = 0)[source]

Bases: StreamingGeometricSketch

Online geometric sketching via locality-sensitive hashing (LSH).

Streams single-cell gene expression data and retains a geometrically diverse sketch of cells across a single pass. Rare cell states are preserved while over-represented dense clusters are capped at max_cells_per_bucket.

Cells are bucketed by projecting into a low-dimensional space and binarizing the result to produce an integer bucket ID. Each bucket maintains a reservoir of up to max_cells_per_bucket cells via uniform reservoir sampling.

By default the projection is a frozen random Gaussian linear map from gene space to n_bits dimensions (classical random-hyperplane LSH). Optionally a pre-trained module (e.g. a frozen PCA or scVI encoder) can be provided as projector; its output is then passed through the LSH random linear layer, so n_bits controls bucket granularity independently of the encoder’s output dimension.

Only single-device training is supported. A RuntimeError is raised at the start of training if more than one device is detected.

Parameters:

var_names_g (ndarray) – Gene names for input validation.
n_bits (int) – Number of LSH projection bits. Determines up to 2^n_bits buckets.
max_cells_per_bucket (int) – Maximum cells retained per bucket via reservoir sampling.
min_cells_per_bucket (int) – Density threshold applied at the end of training. Buckets that observed fewer cells in total are dropped. Default of 1 disables filtering.
min_metadata_diversity (int) – Diversity threshold applied at the end of training. Buckets that observed fewer unique metadata categories are dropped. Default of 1 disables filtering.
store_cell_data (bool) – If True (default), accumulate sparse cell expression vectors. If False, only cell IDs (obs_names) are stored; calling get_reservoir(return_cell_data=True) will raise.
projector (Module | None) – Optional frozen encoder mapping (N, G) → (N, D). When given, its output feeds the LSH linear layer rather than raw gene expression. The module’s gradients are disabled on assignment.
limit_input_to_top_pcs (int | None) – If specified, limit the input to this number of top principal components.
seed (int) – Random seed for the LSH projection weights and for reservoir sampling. Sampling draws from a generator owned by this model, so results are unaffected by other consumers of the global torch RNG.

forward(x_ng: Tensor, var_names_g: ndarray, obs_names_n: ndarray, metadata_n: ndarray | None = None) → dict[str, Tensor | None][source]

Accumulate a minibatch of cells into the geometric sketch reservoir.

Parameters:

x_ng (Tensor) – Expression matrix of shape (N, G), dense or sparse.
var_names_g (ndarray) – Gene names for the batch; must match self.var_names_g.
obs_names_n (ndarray) – Cell IDs for the batch, shape (N,).
metadata_n (ndarray | None) – Optional integer metadata of shape (N,) for diversity tracking. When provided, each cell’s value is recorded in its bucket’s metadata set, enabling apply_bucket_filters() with min_metadata_diversity > 1.

Returns:

An empty dict (no loss is computed).

Return type:

dict[str, Tensor | None]

update(x_ng: Tensor, obs_names_n: ndarray, metadata_n: ndarray | None = None) → int[source]

Update the reservoir with a minibatch of cells.

Parameters:

x_ng (Tensor) – Expression matrix of shape (N, G), dense or sparse.
obs_names_n (ndarray) – Cell IDs of shape (N,).
metadata_n (ndarray | None) – Optional integer metadata of shape (N,).

Returns:

Number of cells inserted or replaced in this update.

Return type:

int

get_reservoir(return_cell_data: bool = False, max_cells: int | None = None, seed: int | None = None) → dict[str, ndarray | Tensor][source]

Retrieve the current sketch.

Buckets are visited in sorted key order so the pre-downsample ordering is deterministic. This method is non-destructive; call apply_bucket_filters() separately to prune the internal state.

Parameters:

return_cell_data (bool) – If True, include sparse cell expression in the output. Requires store_cell_data=True at construction.
max_cells (int | None) – If given, randomly downsample the result to at most this many cells. When the reservoir is already smaller than max_cells no downsampling occurs.
seed (int | None) – Random seed for the max_cells downsample. None gives a non-reproducible draw.

Returns:

"obs_names" — np.ndarray of cell IDs, always present.
"x_ng" — sparse CSR tensor of shape (N_sketch, G), present when return_cell_data=True.

Return type:

A dict with

apply_bucket_filters(min_cells_per_bucket: int = 1, min_metadata_diversity: int = 1) → None[source]

Prune buckets in-place that fail density or diversity thresholds.

Uses _bucket_total_seen (the total number of cells ever hashed to a bucket, not the capped retained count) as the density measure. This correctly identifies sparse regions even when max_cells_per_bucket caps the retained count below the threshold.

Parameters:

min_cells_per_bucket (int) – Buckets where fewer than this many cells were ever observed are deleted.
min_metadata_diversity (int) – Buckets that observed fewer than this many unique metadata categories (from metadata_n) are deleted. Ignored when min_metadata_diversity <= 1.

Raises:

ValueError – If min_metadata_diversity > 1 but no metadata was ever tracked (i.e. metadata_n was never passed to update() or forward()).

Return type:

None

property total_cells: int: Total cells currently stored in the sketch.

property num_filled_buckets: int: Number of buckets containing at least one cell.

property bucket_fill_fraction: float: Fraction of the 2^n_bits buckets that are occupied.

reset_parameters() → None[source]

Reset all accumulated sketch data and re-seed the LSH projection weights.

Safe to call before or after lazy initialization.

Return type:: None

class cellarium.ml.models.StreamingPlaidGeometricSketch(var_names_g: ndarray, target_voxels: int = 100000, max_cells_per_bucket: int = 1, min_cells_per_bucket: int = 1, min_metadata_diversity: int = 1, store_cell_data: bool = False, projector: Module | None = None, limit_input_to_top_pcs: int | None = None, seed: int = 0, target_n_cells: int | None = None)[source]

Bases: StreamingGeometricSketch

Online geometric sketching [1] via Dynamic Spatial Hashing.

Streams single-cell gene expression data and retains a geometrically diverse sketch of cells across a single pass. It maintains a plaid $epsilon$-cover of the embedding space: an axis-aligned grid whose occupied voxels each hold a reservoir of up to max_cells_per_bucket cells. Whenever the number of occupied voxels exceeds target_voxels the grid is coarsened and colliding voxels are merged exactly (no per-cell coordinates are stored). With max_cells_per_bucket=1 this is the sketch of [1]: one cell per occupied box. With a few cells per voxel, the final sketch can be filled up to exactly target_n_cells as in [1] (see below).

To combat technical artifacts, an end-of-epoch pruning step removes voxels that fail to meet minimum cell count (density) or minimum categorical diversity (e.g., number of unique datasets or patients) thresholds.

Grid geometry. As in the reference implementation of [1], the grid lives in the embedding’s native units: every axis starts with the same scalar voxel size (axes are not rescaled or whitened, so Euclidean distances in the embedding keep their meaning), and an axis whose spread is smaller than the voxel size is effectively never split. The effective dimensionality of the cover therefore adapts to the data instead of growing with the number of embedding dimensions. The starting voxel size is a small fraction of the largest per-axis standard deviation of the first minibatch, so the grid starts finer than needed: coarsening is irreversible, starting too coarse is not recoverable.

Grid origin. The origin is estimated once from the first minibatch (just below its per-axis minimum) and then held fixed. A permanent grid boundary sits at the origin for every voxel size, so an origin at zero would bisect zero-centered embeddings (e.g. PCA output) down the middle on every axis and prevent coarsening from ever merging those cells.

Coarsening. The box length of [1] is a continuous scalar found by bisection on a static dataset. Here the grid can only double its spacing (so that old voxels nest exactly inside new ones), and it does so one axis at a time, always the axis with the currently smallest voxel size (so axes take turns and stay within a factor of 2 of each other). Doubling a single axis merges voxels in pairs at most, so the occupied voxel count falls by at most a factor of 2 per step, in any dimension. A coarsening event takes the smallest number of such steps that brings the count to target_voxels or fewer, so it lands in (target_voxels / 2, target_voxels]. (Doubling all axes at once can collapse the cover from far above target_voxels to a single voxel in one step in moderate-to-high dimension, and that cannot be undone.) The steps of an event are not carried out one by one: the number of steps is found by bisection, with each candidate evaluated by hashing the voxel coordinates shifted right by the per-axis number of doublings, and the voxels are then merged in a single pass.

Storage. All state is a handful of tensors on the device of the embedding (a GPU, if training on one): a row per occupied voxel with its integer coordinates, a 64-bit hash of them (voxels are identified by the hash), the number of cells seen, and the retained cells. Cells are retained by random tags: every cell draws an independent uniform tag and a voxel keeps the max_cells_per_bucket cells with the smallest tags, which is a uniform reservoir sample and merges exactly when voxels merge. Inserting a minibatch and coarsening are sort / scatter kernels with no Python loop over voxels or cells. Cell names (and cell data, if store_cell_data) live host-side, keyed by integer cell ids.

Exactly ``target_n_cells`` cells. The voxel count alone does not give the cell count (the count after a coarsening event is anywhere in (target_voxels / 2, target_voxels]). If target_n_cells is given, get_reservoir() (with max_cells as that number, which is also what is stored in sketch_obs_names at the end of training) does what [1] does at the end: it coarsens a copy of the grid to the coarsest nested grid that still has at least target_n_cells occupied voxels, and takes one cell from each of target_n_cells randomly chosen voxels. If the grid is already too coarse for that (fewer voxels than target_n_cells) it instead cycles through the voxels, round by round, taking each voxel’s next retained cell, which needs max_cells_per_bucket > 1. A good choice is target_voxels of about 2 * target_n_cells and max_cells_per_bucket=2, so that the grid has room to be coarsened into place at the end.

Only single-device training is supported. A RuntimeError is raised at the start of training if more than one device is detected.

References

[1] Hie, B., …, Berger, B. (2019). Geometric sketching compactly summarizes: the single-cell transcriptomic landscape. Cell Systems, 8(6), 483-493.e7.

Parameters:

var_names_g (ndarray) – Gene names for input validation.
target_voxels (int) – Maximum number of occupied voxels. The grid coarsens whenever this is exceeded.
max_cells_per_bucket (int) – Maximum cells retained per voxel via uniform reservoir sampling. The retained cells take 16 bytes per slot of every voxel on the device, so keep this small (1 or 2).
min_cells_per_bucket (int) – Density threshold applied at the end of training via apply_bucket_filters(). Voxels that observed fewer cells in total are dropped. Default of 1 disables filtering.
min_metadata_diversity (int) – Diversity threshold applied at the end of training via apply_bucket_filters(). Voxels that observed fewer unique metadata categories (from metadata_n) are dropped.
store_cell_data (bool) – If True, accumulate sparse cell expression vectors.
projector (Module | None) – Optional frozen encoder mapping (N, G) → (N, D).
limit_input_to_top_pcs (int | None) – If specified, limit the input to this number of top principal components.
seed (int) – Random seed for reservoir sampling and for the final selection of cells. Sampling draws from a generator owned by this model, so results are unaffected by other consumers of the global torch RNG.
target_n_cells (int | None) – If given, the exact size of the sketch stored in sketch_obs_names at the end of training (when the sketch has at least that many cells). Must not exceed target_voxels * max_cells_per_bucket.

get_reservoir(return_cell_data: bool = False, max_cells: int | None = None, seed: int | None = None) → dict[str, ndarray | Tensor][source]

Retrieve the current sketch.

This method is non-destructive; call apply_bucket_filters() separately to prune the internal state. Cells are ordered by voxel (in the order of the voxels’ hashes, which is deterministic).

Parameters:

return_cell_data (bool) – If True, include sparse cell expression in the output. Requires store_cell_data=True at construction.
max_cells (int | None) – If given and the sketch holds more cells than this, return exactly this many cells, spread evenly over the occupied regions as in the geometric sketching paper: coarsen a copy of the grid to the coarsest nested grid that still has at least max_cells voxels and take the retained cell of max_cells randomly chosen voxels. If the grid has fewer voxels than max_cells already, take every voxel’s first retained cell, then the second retained cell of randomly chosen voxels, and so on. A sketch with no more than max_cells cells is returned whole.
seed (int | None) – Random seed for the choice of voxels. None uses the model’s own seed, so repeated calls give the same cells.

Returns:

"obs_names" — np.ndarray of cell IDs, always present.
"x_ng" — sparse CSR tensor of shape (N_sketch, G), present when return_cell_data=True.

Return type:

A dict with

apply_bucket_filters(min_cells_per_bucket: int = 1, min_metadata_diversity: int = 1) → None[source]

Prune voxels in-place that fail density or diversity thresholds.

Uses the total number of cells ever assigned to a voxel (not the capped retained count) as the density measure. This correctly identifies sparse regions even when max_cells_per_bucket caps the retained count below the threshold.

Parameters:

min_cells_per_bucket (int) – Voxels where fewer than this many cells were ever observed are deleted.
min_metadata_diversity (int) – Voxels that observed fewer than this many unique metadata categories (from metadata_n) are deleted. Ignored when min_metadata_diversity <= 1.

Raises:

ValueError – If min_metadata_diversity > 1 but no metadata was ever tracked (i.e. metadata_n was never passed to update() or forward()).

Return type:

None

property total_cells: int: Total cells currently stored in the sketch.

property num_filled_buckets: int: Number of voxels containing at least one cell.

on_load_checkpoint(checkpoint: dict[str, Any]) → None[source]

Pre-register lazily-shaped buffers with their checkpointed shape before load_state_dict runs.

voxel_offset and voxel_size are registered lazily by _lazy_init() on the first minibatch, so a freshly constructed model that hasn’t seen a batch yet won’t have them, causing load_state_dict to reject them as unexpected keys. This hook runs before Lightning’s load_state_dict call (see lightning.pytorch.core.saving._load_state), so it can pre-create same-shaped placeholder buffers for each key to load into.

Parameters:: checkpoint (dict[str, Any])
Return type:: None

class cellarium.ml.models.StreamingOrdinaryLeastSquares(var_names_g: ndarray, n_targets: int, univariate: bool = False, ridge_penalty: float = 1e-06, fit_intercept: bool = False)[source]

Bases: CellariumModel, PredictMixin

Streaming ordinary least squares (OLS) solver.

Accumulates sufficient statistics over minibatches, then solves the normal equations once at the end of the first epoch. Training is stopped after one pass.

Two modes are supported:

Multivariate (univariate=False, default): fits a single joint regression X @ W = Y where all features compete simultaneously. Requires accumulating X^T X of shape (n_features, n_features) — only tractable when n_features is small (e.g. a gene expression matrix filtered to a relevant gene set, or a low-dimensional embedding from obsm).
Univariate (univariate=True): fits n_features × n_targets independent simple linear regressions, one per feature–target pair. Accumulates only the sum of squared feature values (shape (n_features,)), making it tractable for high-dimensional X such as a raw genotype matrix with millions of variants. The ridge penalty is added per-feature to the scalar denominator rather than to a matrix diagonal; it is otherwise equivalent in intent.

An intercept is not fitted by default, which means the fit is forced through the origin. That is rarely what is wanted and fails silently: with a strictly positive feature and non-negative targets, for instance, every coefficient is forced positive regardless of the data. Pass fit_intercept=True unless a through-the-origin fit is genuinely intended; leaving it off emits a warning when any feature is strongly off-center.

The intercept is fitted via the centered normal equations rather than by appending a column of ones, so solve() keeps returning slopes of shape (n_features, n_targets), the var_names_g schema is unchanged, and the ridge penalty applies to the slopes only (an intercept should never be penalized).

Note

The centering correction subtracts n * mean**2 from the accumulated sum of squares. When a feature’s mean is very large relative to its spread (e.g. a year covariate, 2024 +/- 2) this is a catastrophic cancellation, and the precision lost while accumulating in float32 cannot be recovered at solve time. Centering such a feature before it reaches this model, or accumulating in float64, avoids the issue.

Parameters:

var_names_g (ndarray) – The variable names schema for the input data validation.
n_targets (int) – Number of target columns (k in y_nk).
univariate (bool) – If True, run massively parallel univariate regressions instead of a single joint multivariate regression.
ridge_penalty (float) – L2 penalty added to the diagonal of X^T X (multivariate) or to each feature’s sum of squares (univariate) before solving. Recommended for numerical stability. Never applied to the intercept.
fit_intercept (bool) – If True, fit an intercept per target (multivariate) or per feature-target pair (univariate), and expose it as intercept_k / intercept_gk. Defaults to False for backwards compatibility.

forward(x_ng: Tensor, var_names_g: ndarray, y_nk: Tensor) → dict[str, Tensor | None][source]

Accumulate sufficient statistics for a minibatch.

Parameters:

x_ng (Tensor) – Feature matrix of shape (batch_size, n_features).
var_names_g (ndarray) – The variable names for the input data.
y_nk (Tensor) – Target matrix of shape (batch_size, n_targets).

Returns:

An empty dictionary (no loss).

Return type:

dict[str, Tensor | None]

update(x_ng: Tensor, y_nk: Tensor) → None[source]

Update OLS accumulators with a minibatch.

Parameters:

x_ng (Tensor) – Tensor of shape (batch_size, n_features).
y_nk (Tensor) – Tensor of shape (batch_size, n_targets).

Return type:

None

solve(ridge_penalty: float | None = None, return_intercept: bool = False) → Tensor | tuple[Tensor, Tensor][source]

Solve the normal equations for the accumulated data.

When fit_intercept is True this solves the centered normal equations, which yields exactly the slopes of a fit with an intercept while leaving the returned shape unchanged.

Parameters:

ridge_penalty (float | None) – L2 penalty. In multivariate mode it is added to the diagonal of X^T X; in univariate mode it is added to each feature’s sum of squares. Defaults to self.ridge_penalty. Never applied to the intercept.
return_intercept (bool) – If True, also return the fitted intercept: shape (n_targets,) in multivariate mode, (n_features, n_targets) in univariate mode. Zeros when fit_intercept is False.

Returns:

Coefficient matrix of shape (n_features, n_targets), or a (coefficients, intercept) tuple when return_intercept is True.

Return type:

Tensor | tuple[Tensor, Tensor]

on_train_epoch_end(trainer: Trainer) → None[source]

Solve the normal equations at the end of the first (and only) epoch.

In multi-GPU training the accumulators are all-reduced before solving so the solution uses the full dataset rather than a single shard.

Parameters:: trainer (Trainer)
Return type:: None

predict(x_ng: Tensor, var_names_g: ndarray) → TransformPrediction[source]

Apply the solved coefficients to new data.

Parameters:

x_ng (Tensor) – Feature matrix of shape (batch_size, n_features).
var_names_g (ndarray) – The variable names for the input data.

Returns:

A dictionary with x_ng (misnomer) of shape (batch_size, n_targets).

Raises:

NotImplementedError – in univariate mode with fit_intercept=True, where each feature-target pair has its own intercept and there is no single joint prediction to make.

Return type:

TransformPrediction