structboost.BAE¶
- class structboost.BAE(n_genes, config=None)[source]¶
Bases:
ModuleBoosting 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:
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
obsm key for the per-dimension standardized latent, written only by transfer models.
Fitted transferred block
(n_genes, k0), or None if not a transfer.Number of transferred dimensions;
0for an ordinary model.Fitted novel block
(n_genes, k1), or None if not a transfer.aligned reference matrix
(n_genes, k0), or None.Per-iteration reconstruction and checkpoint-selection losses.
Per-iteration diagnostics, or None if
diagnosticswas 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 appendsn_additional_dimszero-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 isz* = 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 toprior_mode.- Parameters:
reference (object) – The prior encoder matrix. Accepts a fitted
BAE, anAnnDatacarryingvarm["BAE_encoder_weights"], a gene-indexedpandas.DataFrame, or a path to a.parquet/.csvwritten bywrite_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
fittime 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.
0is 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_dimand the novel dimensions’ stability: a dimension contributing almost nothing is surplus. Evaluate on held-out cells viatransfer_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"; seeBAEConfig. Overridesconfig.prior_modewhen 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.varto 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 invar_namesand accessions invar["gene_ids"]). The pair used is recorded inadata.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_dimis overwritten withk0 + n_additional_dims, since the layout is determined by the reference matrix rather than chosen. Defaults to a freshBAEConfigand is not inherited from the reference: the reference’sseed,max_iterationsand 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:
- save(path)[source]¶
Write the fitted model to a single checkpoint file.
The checkpoint holds everything needed to reproduce
transform()andreconstruct(), to use the model as areferenceforfrom_reference(), and to read backtraining_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 withweights_only=True: loading a structboost checkpoint cannot execute code from the file.- Parameters:
path (str | Path) – Destination file.
.ptis the conventional suffix. Parent directories are created if needed.- Returns:
The path written.
- Return type:
See also
BAE.loadRead a checkpoint back.
structboost.write_encoder_weightsPersist 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:
- Return type:
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()andreconstruct().- 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:
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 againstfitted_prior_weightsto see how far the data moved the programs.- Type:
The anchor
- 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
diagnosticswas off.
- 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:
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)(Optional) Orthogonalizing the targets across latent dimensions
Resetting encoder weights to zero
Fitting encoder via allboost to map X → z*
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 ofadata.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 tolinear_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.
Nonemeans 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_keyraises, rather than silently integrating nothing.The ridge that stabilizes the
"encoder"mechanism when covariates are near-collinear lives onBAEConfigasnuisance_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
BAEnotes on warm starts. Warm-start the latent state fromadata.obsm[init_obsm]instead of from zero. If that representation has a different number of columns thanconfig.latent_dim, the representation wins:latent_dimis overwritten for this fit and aUserWarningis emitted. Mutually exclusive withinit_pca.init_pca (bool) – Exploratory, under active development. Warm-start from a PCA of
adata.Xkeepingconfig.latent_dimcomponents. 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 withinit_obsm.init_pretrain_epochs (int) – Before training, fit the decoder to map the warm-start latent state to
adata.Xfor this many full passes over the cells. Only valid withinit_obsmorinit_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/dzstill 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 inadata. Off by default; call the method directly for control overn_runsandthreshold.
- Returns:
Self for method chaining.
- Return type:
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
Noneto forceadata.X.
- Returns:
Latent representation of shape (n_cells, latent_dim).
- Return type:
np.ndarray
See also
BAE.reconstructReconstruct expression; needs the conditioning columns.
BAE.fit_transformFit 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
Noneto forceadata.X.
- Returns:
Reconstructed expression of shape
(n_cells, n_genes).- Return type:
np.ndarray
See also
BAE.transformGene-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_qualityis high. Defaults have already moved once (thresholdwent from 0.7 to 0.5 in 0.4.0).Records how often each gene is selected across
n_runsfurther 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 control —
expected_false_positivesis deliberatelyNaNrather 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_qualityis 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_runsdefaults 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_sqneeds on the order of1/(1 - beta2) = 1000steps to become a usable variance estimate. An iteration takesceil(n_cells / batch_size)steps, so that transient spans roughly1000 * batch_size / n_cellsof the counted iterations — brief on a large dataset (about 31 ofn_runs=300at 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
fitrestored, 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_lrfor 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 — seedim_match_quality, and fall back tofrequency.max(axis=1)when it is low. A summary is written underadata.uns["bae"]["stability_selection"].
See also
structboost.stability_selectionThe 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 makefit()followed by a reliability check produce a different model thanfit()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 withstability_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 bystable_encoderand stored inadata.varm.adata (AnnData | None) – If given, refresh
varm["BAE_encoder_weights"]andobsm["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 fromweights. DefaultTrue, because the obvious call — installingstable_encoder()— would otherwise delete the transferred programs: that estimator zeroes every entry outside the stable support, and underprior_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. PassFalseonly to overwrite the transferred programs deliberately.
- Returns:
Self, for chaining.
- Return type:
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.fitwrites these for the cells it trained on, wherenovel_variance_shareis 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
fitfor an out-of-sample reading.- Returns:
The same mapping
fitstores inadata.uns["bae_transfer"].- Raises:
ValueError – If the model carries no prior matrix, or the panel does not match.
- Return type:
- 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_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:
- Returns:
Encoder weights of shape (n_genes, 2 * latent_dim).
- Return type: