structboost.BAE

class structboost.BAE(n_genes, config=None)[source]

Bases: Module

Boosting Autoencoder for interpretable dimensionality reduction.

Combines a linear encoder (optimized via componentwise boosting) with an MLP decoder (optimized via SGD). The hybrid training procedure alternates between boosting-based encoder updates and gradient-based decoder updates.

Parameters:
  • n_genes (int) – Number of input genes.

  • config (BAEConfig | None) – BAE configuration. If None, uses defaults.

Examples

>>> import anndata as ad
>>> adata = ad.read_h5ad("data.h5ad")
>>> model = BAE(adata.n_vars)
>>> model.fit(adata)  # Stores results in adata.obsm["X_bae"]
>>> latent = adata.obsm["X_bae"]

Methods

__init__(n_genes[, config])

Initialize internal Module state, shared by both nn.Module and ScriptModule.

apply_encoder(weights[, adata, preserve_prior])

Install aggregated encoder weights, replacing the fitted ones.

fit(adata, *[, layer, mandatory_genes, ...])

Fit BAE using hybrid boosting+SGD training.

fit_transform(adata, **fit_kwargs)

Fit model and return latent representation.

forward(x[, obs_covariates])

Forward pass: encode then decode.

from_reference(reference, adata, *[, ...])

Build a model that carries a reference encoder matrix onto new data.

get_encoder_weights([as_numpy])

Get encoder weight matrix.

get_latent(x)

Get latent representation.

get_splitsoftmax_encoder_weights(*[, ...])

Get encoder weights mapped to split-softmax dimensions.

load(path, *[, device])

Load a model written by save().

reconstruct(adata, *[, layer])

Reconstruct expression using the fitted, conditioned decoder.

save(path)

Write the fitted model to a single checkpoint file.

stability_selection(adata, *[, n_runs, ...])

Stability-select genes for each latent dimension of the fitted model.

transfer_diagnostics(adata)

Transfer diagnostics evaluated on adata, which need not be the fit data.

transform(adata, *[, layer])

Transform data to the gene-only latent space.

transform_splitsoftmax(adata)

Transform data to split-softmax representation.

Attributes

SCALED_LATENT_KEY

obsm key for the per-dimension standardized latent, written only by transfer models.

fitted_prior_weights

Fitted transferred block (n_genes, k0), or None if not a transfer.

n_prior_dims

Number of transferred dimensions; 0 for an ordinary model.

novel_weights

Fitted novel block (n_genes, k1), or None if not a transfer.

prior_weights

aligned reference matrix (n_genes, k0), or None.

training_history

Per-iteration reconstruction and checkpoint-selection losses.

training_report

Per-iteration diagnostics, or None if diagnostics was off.

training

classmethod from_reference(reference, adata, *, n_additional_dims=5, prior_mode=None, min_coverage=0.5, join_on=None, config=None)[source]

Build a model that carries a reference encoder matrix onto new data.

Exploratory

Under active development. Alignment, freezing and the diagnostics work, but the prior and novel latent blocks land on incomparable scales (a 232x gap in per-dimension standard deviation was measured), so every Euclidean consumer must be handed obsm["X_bae_scaled"].

The transferable product of a BAE fit is its encoder weight matrix: k0 sparse gene programs. This constructor aligns such a matrix to adata’s gene panel, places it in the first k0 latent dimensions, and appends n_additional_dims zero-initialized dimensions for variance the prior programs cannot explain.

Fitting then runs in two phases. The decoder is first trained against the prior programs alone (fit(..., decoder_warmup_epochs=...)), which is what makes the added dimensions residual: the boosting target is z* = z - lr * dL/dz, so until the decoder has converged against the prior programs the gradient still carries signal those programs could explain, and the new dimensions would merely re-learn them. Boosting then starts, restricted to the new dimensions or anchored on the prior ones according to prior_mode.

Parameters:
  • reference (object) – The prior encoder matrix. Accepts a fitted BAE, an AnnData carrying varm["BAE_encoder_weights"], a gene-indexed pandas.DataFrame, or a path to a .parquet / .csv written by write_encoder_weights(). A bare array is rejected: it carries no gene identifiers, so aligning it to another panel would be guesswork.

  • adata (AnnData) – Target dataset. Required here rather than at fit time because the encoder’s shape depends on its gene panel, and because a coverage failure should surface before any training happens.

  • n_additional_dims (int) –

    Number of new latent dimensions, default 5. 0 is valid and useful: it adapts the decoder (and, under "anchored", the programs) to the new data without adding capacity.

    This is a user choice, and the default is a pragmatic starting point rather than one derived from the data — how much residual structure a dataset holds is not knowable in advance. Erring high is the cheaper mistake under prior_mode="frozen", where the transferred programs stay bitwise fixed no matter how many dimensions are added, and each novel column’s gene set can be read or ignored independently without refitting. Too small a value silently misses structure instead.

    To check the choice after fitting, read novel_variance_share_per_dim and the novel dimensions’ stability: a dimension contributing almost nothing is surplus. Evaluate on held-out cells via transfer_diagnostics() — in-sample every dimension appears to contribute, because free dimensions always reduce training error.

  • prior_mode (Literal['frozen', 'anchored'] | None) – "frozen" (default) or "anchored"; see BAEConfig. Overrides config.prior_mode when given.

  • min_coverage (float) – Minimum share of a prior dimension’s absolute weight mass that must be present in adata. Below this the transfer raises rather than silently returning an attenuated program.

  • join_on (str | None) – Column of adata.var to align on. By default every identifier set available on both sides is tried and the one matching the most genes wins — a reference keyed on Ensembl accessions therefore aligns to a symbol-indexed dataset without intervention, provided either side carries the other convention (CellRanger and scanpy put symbols in var_names and accessions in var["gene_ids"]). The pair used is recorded in adata.uns["bae_transfer"]["join_key"]. Set this explicitly when a dataset carries several identifier columns and the automatic choice must not be trusted.

  • config (BAEConfig | None) – Base configuration. latent_dim is overwritten with k0 + n_additional_dims, since the layout is determined by the reference matrix rather than chosen. Defaults to a fresh BAEConfig and is not inherited from the reference: the reference’s seed, max_iterations and stopping rule describe how that model was fitted, not how this one should be. Pass the reference’s config explicitly to reuse its boosting hyperparameters.

Returns:

An unfitted model with the aligned prior installed.

Return type:

BAE

save(path)[source]

Write the fitted model to a single checkpoint file.

The checkpoint holds everything needed to reproduce transform() and reconstruct(), to use the model as a reference for from_reference(), and to read back training_history / training_report.

Two things are deliberately excluded. The decoder’s optimizer state is not written, so a loaded model is deployable but cannot resume a training run mid-flight — a fresh fit() still works, since it rebuilds the decoder and optimizer regardless. And the training-set covariate design matrix (ObsCovariateEncoding.encoded) is not written, so a shipped model file carries no cell-level training data and its size does not grow with the training set; only the encoding parameters needed to encode new data are kept.

The payload contains only tensors and plain Python values, which is what lets load() read it with weights_only=True: loading a structboost checkpoint cannot execute code from the file.

Parameters:

path (str | Path) – Destination file. .pt is the conventional suffix. Parent directories are created if needed.

Returns:

The path written.

Return type:

Path

See also

BAE.load

Read a checkpoint back.

structboost.write_encoder_weights

Persist only the gene programs, in a readable, shareable format.

Raises:
  • RuntimeError – If the model has not been fitted.

  • ValueError – If a conditioned fit used an obs column whose categorical levels cannot be persisted (for example datetimes). Converting the column to string, integer or boolean before fitting resolves it.

Parameters:

path (str | Path)

Return type:

Path

Examples

>>> model.fit(adata)
>>> model.save("bae_model.pt")
classmethod load(path, *, device=None)[source]

Load a model written by save().

Parameters:
  • path (str | Path) – Checkpoint file.

  • device (Device | None) – Device to place the model on. If None, the device recorded at save time is used when it is available, and CPU otherwise — the fallback warns rather than relocating silently.

Returns:

The restored model, ready for transform() and reconstruct().

Raises:
  • FileNotFoundError – If path does not exist.

  • ValueError – If the file is not a structboost checkpoint, or was written by a newer checkpoint format than this version understands.

Return type:

BAE

Examples

>>> model = BAE.load("bae_model.pt")
>>> latent = model.transform(adata)
property prior_weights: ndarray | None

aligned reference matrix (n_genes, k0), or None.

This is what was transferred in, not what was fitted. Under prior_mode="frozen" the two are identical; under "anchored" compare against fitted_prior_weights to see how far the data moved the programs.

Type:

The anchor

property n_prior_dims: int

Number of transferred dimensions; 0 for an ordinary model.

property fitted_prior_weights: ndarray | None

Fitted transferred block (n_genes, k0), or None if not a transfer.

property novel_weights: ndarray | None

Fitted novel block (n_genes, k1), or None if not a transfer.

These are the gene programs the transfer added — the residual structure the reference matrix could not represent. Empty (k1 == 0) when the model was built for decoder-only adaptation.

property training_history: dict[str, list[float]]

Per-iteration reconstruction and checkpoint-selection losses.

property training_report: TrainingReport | None

Per-iteration diagnostics, or None if diagnostics was off.

forward(x, obs_covariates=None)[source]

Forward pass: encode then decode.

Parameters:
  • x (Tensor) – Input tensor of shape (n_cells, n_genes).

  • obs_covariates (Tensor | None) – Optional obs covariate tensor for cVAE conditioning.

Returns:

  • x_recon – Reconstructed input.

  • z – Latent representation.

Return type:

tuple[Tensor, Tensor]

get_latent(x)[source]

Get latent representation.

Parameters:

x (Tensor) – Input tensor of shape (n_cells, n_genes).

Returns:

Latent representation of shape (n_cells, latent_dim).

Return type:

Tensor

fit(adata, *, layer=None, mandatory_genes=None, batch_key=None, batch_integration_mode='__unset__', max_iterations=None, early_stopping_patience=None, enable_early_stopping=None, seed=None, verbose=True, diagnostics=None, init_obsm=None, init_pca=False, init_pretrain_epochs=0, decoder_warmup_epochs=0, stability_selection=False)[source]

Fit BAE using hybrid boosting+SGD training.

The training loop alternates between:

  1. Computing boosting targets via gradient step: z* = z - lr * ∂L_target/∂z, where L_target sums the squared error over cells and averages it over genes (see _compute_boosting_targets)

  2. (Optional) Orthogonalizing the targets across latent dimensions

  3. Resetting encoder weights to zero

  4. Fitting encoder via allboost to map X → z*

  5. Updating the decoder with one shuffled pass over the cells, i.e. ceil(n_cells / batch_size) AdamW steps

Parameters:
  • adata (AnnData) – AnnData object with gene expression. Data should be standardized (z-transformed).

  • layer (str | None) – Read expression from adata.layers[layer] instead of adata.X. The choice is remembered: transform(), reconstruct() and the diagnostics default to the same layer, so a fitted model always reads the representation it learned on. Pass the same layer to linear_ceiling() when comparing reconstruction quality against its achievable maximum.

  • mandatory_genes (list[str] | list[int] | np.ndarray | list | None) –

    Gene names (str) or column indices (int) placed in the unpenalized adjustment block of the boosting fit, so they are never subject to competitive selection. Can be a flat list (applied to all latent dims) or a list of lists (per latent dimension).

    This forces them into the model specification, not into the fitted support: a gene whose contribution is estimated as zero still ends up with a zero encoder weight. Do not rely on this to guarantee that a marker appears in the selected gene set.

  • batch_key (str | list[str] | None) – Obs column, or several, holding the covariate to integrate over. None means no integration is performed.

  • batch_integration_mode (Literal['encoder', 'decoder', 'both']) –

    Which of the two mechanisms to apply. Defaults to "both".

    "decoder"

    The encoded covariate is concatenated to the decoder input, so the decoder can explain covariate-driven variation directly and the latent code does not have to carry it.

    "encoder"

    The encoded covariate is added to the boosting design as a mandatory regressor, so a covariate-correlated gene is not selected because of the covariate.

    "both"

    Both of the above. This is the usual choice.

    The covariate is never an encoder input. "encoder" names the half of the model it protects, not a tensor it is fed to. transform() stays gene-only and needs no covariate labels under any mode, which is what makes a fitted encoder deployable on data carrying no covariate annotation. reconstruct() needs them only under "decoder" and "both".

    Passing a mode without a batch_key raises, rather than silently integrating nothing.

    The ridge that stabilizes the "encoder" mechanism when covariates are near-collinear lives on BAEConfig as nuisance_ridge.

  • max_iterations (int | None) – Maximum training iterations (overrides config).

  • early_stopping_patience (int | None) – Early stopping patience (overrides config).

  • enable_early_stopping (bool | None) – Enable/disable early stopping (overrides config).

  • seed (int | None) – Random seed for reproducibility (overrides config).

  • verbose (bool) – Print training progress.

  • diagnostics (bool | None) – Collect per-iteration training diagnostics (overrides config). Adds full-data forward/backward passes per iteration; see training_report. Also adds the relative encoder weight change (dW, which converges toward 0) and the number of selected genes (n_sel) to the progress bar. Does not change the fitted model.

  • init_obsm (str | None) – Exploratory, under active development; see BAE notes on warm starts. Warm-start the latent state from adata.obsm[init_obsm] instead of from zero. If that representation has a different number of columns than config.latent_dim, the representation wins: latent_dim is overwritten for this fit and a UserWarning is emitted. Mutually exclusive with init_pca.

  • init_pca (bool) – Exploratory, under active development. Warm-start from a PCA of adata.X keeping config.latent_dim components. The data is never rescaled, and it is mean-centered only when it is not already z-transformed, so no transformation is applied on top of one the caller already performed. Mutually exclusive with init_obsm.

  • init_pretrain_epochs (int) – Before training, fit the decoder to map the warm-start latent state to adata.X for this many full passes over the cells. Only valid with init_obsm or init_pca; passing a positive value without one raises. Default 0 (no pre-training); 20 is a conservative value. Measured trade-off: this lowers the initial loss substantially but does not improve the converged loss, and large values noticeably reduce gene-selection precision, because the decoder is tuned to a latent code the sparse encoder cannot exactly reproduce.

  • decoder_warmup_epochs (int) – Transfer models only (from_reference()). Trains the decoder against the frozen prior programs for this many passes before boosting starts. This is what makes the added dimensions residual: until the decoder has converged against the prior, dL/dz still carries signal those programs could explain, and the new dimensions would re-learn it.

  • stability_selection (bool) – Run stability_selection() once after training and store its results in adata. Off by default; call the method directly for control over n_runs and threshold.

Returns:

Self for method chaining.

Return type:

BAE

Notes

The warm start is applied once, on the first iteration only: it sets the boosting targets, so the encoder learns the supplied representation and the decoder is then trained against the encoder’s output.

With config.disentanglement="correlation", the soft penalty is applied inside step 1 rather than as a target transformation. Early stopping and best-state restoration then use reconstruction loss plus the weighted disentanglement penalty. training_history["train_loss"] remains the unregularized decoder MSE; training_history["selection_loss"] records the checkpoint-selection objective.

transform(adata, *, layer=<fit-time layer>)[source]

Transform data to the gene-only latent space.

Parameters:
  • adata (AnnData) – AnnData object with the same genes used during fitting. Batch or other obs labels are deliberately not required because they are not part of the deployable encoder.

  • layer (str | None | _FitLayer) – Where to read expression from. Defaults to the layer the model was fitted on, so a model applied to new data reads the same representation it was trained on. Pass a name to override, or None to force adata.X.

Returns:

Latent representation of shape (n_cells, latent_dim).

Return type:

np.ndarray

See also

BAE.reconstruct

Reconstruct expression; needs the conditioning columns.

BAE.fit_transform

Fit and return the latent in one call.

reconstruct(adata, *, layer=<fit-time layer>)[source]

Reconstruct expression using the fitted, conditioned decoder.

Unlike transform(), this method requires the decoder-conditioning obs columns used during fitting and rejects unseen categorical levels.

Parameters:
  • adata (AnnData) – AnnData object with expression and, when applicable, the fitted batch columns, when the mode conditions the decoder.

  • layer (str | None | _FitLayer) – Where to read expression from. Defaults to the layer the model was fitted on. Pass a name to override, or None to force adata.X.

Returns:

Reconstructed expression of shape (n_cells, n_genes).

Return type:

np.ndarray

See also

BAE.transform

Gene-only latent projection; needs no obs labels.

stability_selection(adata, *, n_runs=300, threshold=0.5, seed=None, verbose=True, continue_optimizer=False)[source]

Stability-select genes for each latent dimension of the fitted model.

Exploratory

Under active development. This provides no formal error control, and its per-dimension frequencies are only interpretable when dim_match_quality is high. Defaults have already moved once (threshold went from 0.7 to 0.5 in 0.4.0).

Records how often each gene is selected across n_runs further training iterations continued from the fitted state, then restores the model, so the call is non-destructive. It answers: would these genes still be selected if the optimizer had stopped somewhere else on its loss plateau?

That is the variance source that dominates here. The encoder support does not converge even when the reconstruction loss does — on measured data the loss plateaus at ~95% of the achievable linear ceiling while consecutive iterations share only about a third of their selected genes. A single fit reports one arbitrary position on that walk.

On simulated data with exact ground truth this gives the lowest false-discovery rate of the available readouts (0.26, against 0.31 for cell subsampling and 0.38 for a single fit), and its frequencies are empirically well calibrated: genes selected in 90-100% of iterations are markers 88% of the time. It provides no formal error controlexpected_false_positives is deliberately NaN rather than a number that would look like a guarantee, because training iterations are neither independent nor exchangeable.

Scope. This measures stability conditional on the learned representation. It does not capture the variability from re-initializing and refitting the autoencoder to a different local optimum, which full refits would. For high-stakes marker claims, a handful of full refits remains a worthwhile cross-check.

Parameters:
  • adata (AnnData) – The data the model was fitted on (same genes; obs columns required only if the model used conditioning or nuisance covariates).

  • n_runs (int) – How many further training iterations to average over. The default of 300 is set by the support autocorrelation, which decays slowly (still ~0.45 at lag 100 on measured data), so short windows give highly correlated, near-duplicate samples.

  • threshold (float) –

    Selection-frequency cutoff for the stable support. Default 0.5, lowered from 0.7 in 0.4.0.

    0.7 was measured to be too aggressive whenever the latent representation is still moving: across three real datasets it removed 21-52% of recovered marker genes relative to the fitted encoder, and 0-7% even when the representation had settled. At 0.5 the worst loss over the same six runs was 6%. Raising it back toward 0.7-0.9 buys precision and is reasonable when dim_match_quality is high; see that field on the result before doing so.

    Note the standalone structboost.stability_selection() keeps 0.7, because its Meinshausen-Buhlmann bound is undefined at or below 0.5.

  • seed (int | None) – Seeds torch for the continued training iterations.

  • verbose (bool) – Show a progress bar. On by default: n_runs defaults to 300 further training steps, which is a long silence.

  • continue_optimizer (bool) –

    Carry the decoder’s AdamW moment estimates over from fit() instead of starting them at zero. Default False, which preserves the behaviour every earlier result was measured under.

    The reset is not free. A fresh AdamW restarts its step counter, so bias correction begins again and exp_avg_sq needs on the order of 1/(1 - beta2) = 1000 steps to become a usable variance estimate. An iteration takes ceil(n_cells / batch_size) steps, so that transient spans roughly 1000 * batch_size / n_cells of the counted iterations — brief on a large dataset (about 31 of n_runs=300 at 16,000 cells) but most of the window on a small one (about 250 of 300 at 2,000 cells), and it falls at the end where the support is furthest from equilibrium. Continuing removes it, and the smaller the dataset the more it is worth doing.

    The state is the one from the iteration fit restored, not from its last iteration, so the moments belong to the decoder actually returned. It is held in memory only: save() deliberately excludes optimizer state, so a model read back from a checkpoint has none and this argument warns and falls back rather than silently changing what it measures.

    A caller who lowered config.decoder_lr for this phase keeps that change; only the moment estimates are carried across.

Returns:

StabilitySelectionResult – Stores adata.varm["BAE_iteration_frequency"], shape (n_genes, latent_dim). Each iteration’s dimensions are matched to the fitted model’s before counting, so a dimension index keeps its meaning — see dim_match_quality, and fall back to frequency.max(axis=1) when it is low. A summary is written under adata.uns["bae"]["stability_selection"].

See also

structboost.stability_selection

The standalone allboost-level function, which resamples cells in the Meinshausen-Buhlmann scheme. Available for supervised boosting problems that have no training loop to iterate over.

apply_encoder(weights, adata=None, *, preserve_prior=True)[source]

Install aggregated encoder weights, replacing the fitted ones.

Deliberately separate from stability_selection(), which stays non-destructive. A diagnostic that silently swapped the encoder would make fit() followed by a reliability check produce a different model than fit() alone, and would compound if called twice. It is also not a strictly better encoder but a choice: "masked_cond_mean" buys precision (0.74 vs 0.62 for the fitted encoder on simulated data) at the cost of recall (0.27 vs 0.31), and that trade belongs to the caller — tune it with stability_selection(threshold=...).

The decoder is left untouched. stable_encoder() preserves the latent scale, and on simulated data the aggregate improved reconstruction relative to the fitted encoder (59% of the linear ceiling against 55%), so no refit is required.

Parameters:
  • weights (np.ndarray) – Shape (n_genes, latent_dim) — the orientation returned by stable_encoder and stored in adata.varm.

  • adata (AnnData | None) – If given, refresh varm["BAE_encoder_weights"] and obsm["X_bae"] so the stored results match the installed encoder rather than the superseded one.

  • preserve_prior (bool) – On a model built by from_reference(), keep the transferred block rather than taking it from weights. Default True, because the obvious call — installing stable_encoder() — would otherwise delete the transferred programs: that estimator zeroes every entry outside the stable support, and under prior_mode="frozen" the prior columns have selection frequency zero by construction (they cannot vary, so there is nothing to be stable about). The result would be an encoder whose prior block is all zeros, with no error raised. Pass False only to overwrite the transferred programs deliberately.

Returns:

Self, for chaining.

Return type:

BAE

Examples

>>> res = model.stability_selection(adata)
>>> model.apply_encoder(res.stable_encoder(), adata)
transfer_diagnostics(adata)[source]

Transfer diagnostics evaluated on adata, which need not be the fit data.

fit writes these for the cells it trained on, where novel_variance_share is positive by construction: k free dimensions reduce in-sample reconstruction error whether or not the data contains anything the prior programs missed.

Passing held-out cells is what turns the number into evidence. Capacity that merely fits noise does not generalize, so a novel dimension carrying real structure keeps its share out of sample while one that does not collapses. This needs nothing but the model and your own data — no access to the reference dataset, and no marker ground truth — which matters because a prior encoder matrix is often all that is shared.

Parameters:

adata (AnnData) – Cells to evaluate on, over the same gene panel the model was aligned to. Hold these out of fit for an out-of-sample reading.

Returns:

The same mapping fit stores in adata.uns["bae_transfer"].

Raises:

ValueError – If the model carries no prior matrix, or the panel does not match.

Return type:

dict[str, object]

SCALED_LATENT_KEY: str = 'X_bae_scaled'

obsm key for the per-dimension standardized latent, written only by transfer models. See _refit_latent_scaling.

transform_splitsoftmax(adata)[source]

Transform data to split-softmax representation.

Computes the split-softmax compositional representation h in Delta^{2d-1} from the encoder output z. Each latent dimension z_i is paired with its negation -z_i (interleaved) and softmax-normalized:

h = softmax((z_1, -z_1, …, z_d, -z_d))

Can be called on models trained with or without split_softmax=True. When split_softmax was not enabled during training, a warning is emitted and the transformation is applied post-hoc.

Parameters:

adata (AnnData) – AnnData object.

Returns:

Split-softmax representation of shape (n_cells, 2 * latent_dim).

Raises:

RuntimeError – If model is not fitted.

Return type:

np.ndarray

fit_transform(adata, **fit_kwargs)[source]

Fit model and return latent representation.

Parameters:
  • adata (AnnData) – AnnData object.

  • **fit_kwargs – Arguments passed to fit().

Returns:

Latent representation.

Return type:

np.ndarray

get_encoder_weights(as_numpy=True)[source]

Get encoder weight matrix.

Parameters:

as_numpy (bool) – If True, return numpy array; else torch tensor.

Returns:

Encoder weights of shape (n_genes, latent_dim).

Return type:

ndarray | Tensor

get_splitsoftmax_encoder_weights(*, clip_negative=True, as_numpy=True)[source]

Get encoder weights mapped to split-softmax dimensions.

Returns the effective per-gene weights for each of the 2 * latent_dim split-softmax dimensions. For split-softmax dimension 2i (positive direction of latent dim i), the weights are W[:, i]. For dimension 2i+1 (negative direction), the weights are -W[:, i].

The columns are interleaved in the same order as the split-softmax output: (W_1, -W_1, W_2, -W_2, …, W_d, -W_d).

Can be called on models trained with or without split_softmax=True. When split_softmax was not enabled during training, a warning is emitted.

Parameters:
  • clip_negative (bool) – If True (default), clip all negative weights to zero. This retains only the genes that positively contribute to each split-softmax dimension, improving interpretability.

  • as_numpy (bool) – If True, return numpy array; else torch tensor.

Returns:

Encoder weights of shape (n_genes, 2 * latent_dim).

Return type:

ndarray | Tensor