structboost 0.6.0__tar.gz → 0.6.1__tar.gz

{structboost-0.6.0 → structboost-0.6.1}/CHANGELOG.md

@@ -3,6 +3,34 @@
 Releases follow [semantic versioning](https://semver.org). While the project is
 pre-1.0, a minor bump may break API.
 
+### [0.6.1] - 2026-09-24
+
+**Not breaking.** CSR and dense input give bit-identical fits; CSC input
+changes in the last bits (below).
+
+**Fits need less memory and run faster.** Measured at 300,000 cells by 2,000
+genes: peak memory 10.4 GB to 8.3 GB, wall clock 44.8 s to 25.9 s. Three copies
+of the expression matrix are gone:
+
+- The sparse matrix is densified in row blocks straight into float32, instead of
+  through a full-size intermediate in its source dtype.
+- `fit` hands the training matrix to the step that writes `adata.obsm["X_bae"]`
+  instead of densifying `adata.X` a second time.
+- With batch integration on (`batch_integration_mode` `"encoder"` or `"both"`),
+  the boosting design is allocated once with room for the covariate columns,
+  instead of copying the whole panel with `np.hstack` to append them. Batch
+  integration now costs about the same peak memory as a fit without it, and
+  400,000 cells fit where they ran out of memory before.
+
+A float32 `adata.X` is now trained on in place rather than copied. A fit never
+writes to it; a read-only array is still copied.
+
+**CSC input now gives the same fit as CSR.** `csc.toarray()` returns a
+Fortran-ordered array, which changed the BLAS reduction order and with it the
+fitted weights (2.3e-9 on a 12-iteration fit). The densified matrix is now
+row-major whatever the input, so a CSC-backed fit will not reproduce its earlier
+result, and matches the CSR result for the same data instead.
+
 ### [0.6.0] - 2026-09-09
 
 **Not breaking.** Nothing existing changes behaviour; this adds a way to read a

{structboost-0.6.0 → structboost-0.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: structboost
-Version: 0.6.0
+Version: 0.6.1
 Summary: Boosting Autoencoders (BAE) and componentwise L2 boosting utilities for scverse-style workflows.
 Project-URL: Homepage, https://github.com/NiklasBrunn/structboost
 Project-URL: Documentation, https://niklasbrunn.github.io/structboost

{structboost-0.6.0 → structboost-0.6.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "structboost"
-version = "0.6.0"
+version = "0.6.1"
 description = "Boosting Autoencoders (BAE) and componentwise L2 boosting utilities for scverse-style workflows."
 readme = "README.md"
 license = { file = "LICENSE" }

{structboost-0.6.0 → structboost-0.6.1}/src/structboost/_model.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import copy
 import functools
 import warnings
-from dataclasses import fields, replace
+from dataclasses import dataclass, fields, replace
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal
 
@@ -275,6 +275,38 @@ def _build_allboost_mandatory(
     return np.concatenate([resolved_mandatory, obs_indices])
 
 
+@dataclass(frozen=True)
+class _BoostingDesign:
+    """The matrix ``allboost`` is fitted against, and the parts of it that are fixed.
+
+    Built the same way by :meth:`BAE.fit` and by
+    :meth:`BAE._iteration_support_frequency`, which documents itself as mirroring
+    the fit loop and whose equivalence ``tests/test_stability.py`` pins. Sharing
+    the construction is what keeps that true rather than merely intended.
+
+    Attributes
+    ----------
+    panel
+        ``(n_cells, n_genes)`` expression, float32 — what the encoder reads. When
+        there are nuisance columns this is a column *view* of ``matrix``, so the
+        two are never separate copies of the same numbers.
+    matrix
+        ``(n_cells, n_genes + n_nuisance)`` — what boosting actually sees.
+    col_norms_sq
+        Squared column norms of ``matrix``. Fixed for a whole fit, so ``allboost``
+        is handed them instead of repeating an O(n*p) pass, and a full-size
+        temporary, on every training iteration.
+    mandatory_ridge
+        Per-column ridge for the unpenalized block, non-zero only on the nuisance
+        columns. Mandatory *gene* coefficients are never penalized.
+    """
+
+    panel: np.ndarray
+    matrix: np.ndarray
+    col_norms_sq: np.ndarray
+    mandatory_ridge: np.ndarray
+
+
 class _FitLayer:
     """Sentinel: "whichever layer this model was fitted on"."""
 
@@ -1477,10 +1509,68 @@ class BAE(nn.Module):
 
     @staticmethod
     def _to_tensor(X: np.ndarray | sp.spmatrix, device: torch.device) -> torch.Tensor:
-        """Convert array or sparse matrix to tensor."""
-        if sp.issparse(X):
-            X = X.toarray()
-        return torch.from_numpy(np.asarray(X, dtype=np.float32)).to(device)
+        """Convert array or sparse matrix to tensor.
+
+        A float32 ndarray is shared, not copied, so the tensor may alias the
+        caller's ``adata.X``: nothing downstream may write to it. A read-only
+        array is copied instead, since torch cannot hold one without warning.
+        """
+        from ._utils import densify
+
+        X = densify(X)
+        return torch.from_numpy(X if X.flags.writeable else X.copy()).to(device)
+
+    def _build_boosting_design(
+        self,
+        matrix: np.ndarray | sp.spmatrix,
+        *,
+        covariates: np.ndarray | None = None,
+        n_nuisance: int = 0,
+    ) -> _BoostingDesign:
+        """Read the expression panel into the design matrix boosting is fitted on.
+
+        One allocation holds both blocks and the panel is densified straight into
+        the gene block. ``np.hstack`` instead builds the augmented matrix *from*
+        an already-dense panel, so two full copies are resident at once; and
+        :func:`~structboost._utils.densify` fills the block in row blocks, so no
+        full-size intermediate in the source dtype appears either.
+
+        The returned ``panel`` is then a column view of the design: C-ordered
+        within each row and differing only in row stride, which BLAS takes as a
+        leading dimension. Torch runs the encoder on it without materializing a
+        contiguous copy, and returns bit-identical results — see
+        ``tests/test_densify.py``, which pins that, since a torch that copied
+        instead would silently restore the second panel this avoids.
+
+        Parameters
+        ----------
+        matrix
+            Expression, dense or sparse, as :func:`_expression_matrix` returns it.
+        covariates
+            Encoded obs covariates, float32. Required when ``n_nuisance`` is set.
+        n_nuisance
+            Width of the nuisance block; ``0`` leaves the design gene-only, and
+            ``panel`` is then the design itself rather than a view into it.
+        """
+        from ._utils import densify
+
+        if n_nuisance:
+            full = np.empty((matrix.shape[0], self.n_genes + n_nuisance), dtype=np.float32)
+            panel = densify(matrix, out=full[:, : self.n_genes])
+            full[:, self.n_genes :] = covariates
+        else:
+            panel = full = densify(matrix)
+
+        ridge = np.zeros(full.shape[1], dtype=np.float64)
+        if n_nuisance:
+            ridge[self.n_genes :] = self.config.nuisance_ridge
+
+        return _BoostingDesign(
+            panel=panel,
+            matrix=full,
+            col_norms_sq=column_norms_sq(full),
+            mandatory_ridge=ridge,
+        )
 
     @_isolates_torch_rng(config_fallback=True)
     def fit(
@@ -1660,7 +1750,6 @@ class BAE(nn.Module):
                 raise ValueError("batch_key must name at least one obs column")
             batch_mode = "both" if batch_integration_mode is _MODE_UNSET else batch_integration_mode
         conditions_decoder, regresses_encoder = _BATCH_MODES[batch_mode]
-        nuisance_ridge = self.config.nuisance_ridge
 
         resolved_mandatory = resolve_mandatory_genes(mandatory_genes, adata)
         self._mandatory_genes = mandatory_genes
@@ -1718,10 +1807,35 @@ class BAE(nn.Module):
             if torch.cuda.is_available():
                 torch.cuda.manual_seed_all(seed)
 
-        # Prepare data
-        matrix = _expression_matrix(adata, self._layer)
-        X_np = matrix.toarray() if sp.issparse(matrix) else np.asarray(matrix)
-        X_np = X_np.astype(np.float32)
+        # --- Batch covariates: one encoding, two mechanisms it can drive ---
+        # Resolved before the expression matrix is read, because how many nuisance
+        # columns the boosting design carries decides how wide the array it is
+        # read into has to be.
+        self._batch_encoding = None
+        self._batch_integration_mode = batch_mode
+        D_condition: torch.Tensor | None = None
+        covariates_np: np.ndarray | None = None
+        n_nuisance = 0
+        if batch_columns is not None:
+            from ._utils import encode_obs_covariates
+
+            self._batch_encoding = encode_obs_covariates(adata, batch_columns)
+            # Both mechanisms read the same numbers, so they share one float32
+            # copy of them rather than casting the encoding twice.
+            covariates_np = self._batch_encoding.encoded.astype(np.float32)
+            if conditions_decoder:
+                D_condition = torch.from_numpy(covariates_np).to(self.config.device)
+            if regresses_encoder:
+                n_nuisance = self._batch_encoding.n_columns
+
+        # Prepare data: one allocation for the whole boosting design, with the
+        # panel read straight into its gene block. See `_build_boosting_design`.
+        design = self._build_boosting_design(
+            _expression_matrix(adata, self._layer),
+            covariates=covariates_np,
+            n_nuisance=n_nuisance,
+        )
+        X_np = design.panel
 
         # Warn if data doesn't appear standardized. The result also decides whether
         # a PCA warm start needs to center the data (it must not re-transform data
@@ -1738,9 +1852,6 @@ class BAE(nn.Module):
         # Use full dataset for training (no validation split)
         X_train = self._to_tensor(X_np, self.config.device)
 
-        # Covariance cache for boosting (uses training data only)
-        X_train_np = X_train.cpu().numpy()
-
         # --- Latent state initialization (warm start) ---
         if init_pretrain_epochs < 0:
             raise ValueError(f"init_pretrain_epochs must be >= 0, got {init_pretrain_epochs}")
@@ -1775,23 +1886,6 @@ class BAE(nn.Module):
             else None
         )
 
-        # --- Batch covariates: one encoding, two mechanisms it can drive ---
-        self._batch_encoding = None
-        self._batch_integration_mode = batch_mode
-        D_condition: torch.Tensor | None = None
-        D_nuisance_np: np.ndarray | None = None
-        n_nuisance = 0
-        if batch_columns is not None:
-            from ._utils import encode_obs_covariates
-
-            self._batch_encoding = encode_obs_covariates(adata, batch_columns)
-            if conditions_decoder:
-                condition_np = self._batch_encoding.encoded.astype(np.float32)
-                D_condition = torch.from_numpy(condition_np).to(self.config.device)
-            if regresses_encoder:
-                D_nuisance_np = self._batch_encoding.encoded.astype(np.float32)
-                n_nuisance = self._batch_encoding.n_columns
-
         # Always rebuild so repeated fits cannot retain a stale conditioning shape.
         decoder_input = (
             2 * self.config.latent_dim if self.config.split_softmax else self.config.latent_dim
@@ -1815,17 +1909,8 @@ class BAE(nn.Module):
                 torch.cuda.manual_seed_all(seed)
             self.decoder.reset_parameters()
 
-        # Build augmented sourcemat for allboost
-        if D_nuisance_np is not None:
-            sourcemat_aug = np.hstack([X_train_np, D_nuisance_np])
-        else:
-            sourcemat_aug = X_train_np
-
         # Build combined mandatory indices (genes + obs covariates)
         allboost_mandatory = _build_allboost_mandatory(resolved_mandatory, n_nuisance, self.n_genes)
-        mandatory_ridge = np.zeros(sourcemat_aug.shape[1], dtype=np.float64)
-        if n_nuisance:
-            mandatory_ridge[self.n_genes :] = nuisance_ridge
 
         # A frozen transfer with no added dimensions never calls allboost at all
         # (there is nothing left to select), so building a covariance matrix for
@@ -1839,21 +1924,16 @@ class BAE(nn.Module):
         from ._utils import resolve_precompute_covcache
 
         precompute_covcache = boosts_anything and resolve_precompute_covcache(
-            self.config.boosting_precompute_covcache, sourcemat_aug.shape[1]
+            self.config.boosting_precompute_covcache, design.matrix.shape[1]
         )
         self._precomputed_covcache = precompute_covcache
         if precompute_covcache:
             from ._utils import compute_covariance_cache
 
-            covcache = compute_covariance_cache(sourcemat_aug)
+            covcache = compute_covariance_cache(design.matrix)
         else:
             covcache = None  # Lazy computation during allboost
 
-        # `sourcemat_aug` is fixed for the whole fit, so its column norms are too.
-        # allboost recomputes them on every call otherwise -- an O(n*p) pass and a
-        # full-size temporary, once per training iteration.
-        boosting_col_norms_sq = column_norms_sq(sourcemat_aug)
-
         # Optimizer for decoder only
         decoder_optimizer = torch.optim.AdamW(
             self.decoder.parameters(),
@@ -1960,46 +2040,46 @@ class BAE(nn.Module):
             # model the prior columns are either withheld ("frozen") or boosted
             # from the fixed original matrix as an offset ("anchored").
             fit_targets, beta_init, fit_mandatory = self._transfer_boosting_inputs(
-                targets, sourcemat_aug.shape[1], prior_dims, allboost_mandatory
+                targets, design.matrix.shape[1], prior_dims, allboost_mandatory
             )
             if fit_targets.shape[1] == 0:
                 # Frozen transfer with no additional dimensions: nothing competes
                 # for selection, and only the decoder adapts to the new data.
-                betamat = np.zeros((0, sourcemat_aug.shape[1]), dtype=np.float64)
+                betamat = np.zeros((0, design.matrix.shape[1]), dtype=np.float64)
             elif covcache is None:
                 betamat, covcache = allboost(
-                    sourcemat_aug,
+                    design.matrix,
                     fit_targets,
                     covcache=covcache,
-                    col_norms_sq=boosting_col_norms_sq,
+                    col_norms_sq=design.col_norms_sq,
                     stepno=self.config.boosting_stepno,
                     nu=self.config.boosting_nu,
                     csf=self.config.boosting_csf,
                     independent=self.config.boosting_independent,
                     mandatory_features=fit_mandatory,
-                    mandatory_ridge=mandatory_ridge,
+                    mandatory_ridge=design.mandatory_ridge,
                     beta_init=beta_init,
                     return_covcache=True,
                 )
             else:
                 betamat = allboost(
-                    sourcemat_aug,
+                    design.matrix,
                     fit_targets,
                     covcache=covcache,
-                    col_norms_sq=boosting_col_norms_sq,
+                    col_norms_sq=design.col_norms_sq,
                     stepno=self.config.boosting_stepno,
                     nu=self.config.boosting_nu,
                     csf=self.config.boosting_csf,
                     independent=self.config.boosting_independent,
                     mandatory_features=fit_mandatory,
-                    mandatory_ridge=mandatory_ridge,
+                    mandatory_ridge=design.mandatory_ridge,
                     beta_init=beta_init,
                 )
-            betamat = self._expand_transfer_betamat(betamat, sourcemat_aug.shape[1], prior_dims)
+            betamat = self._expand_transfer_betamat(betamat, design.matrix.shape[1], prior_dims)
 
             # Extract gene weights only; obs weights are nuisance (discarded)
             W_genes = betamat[:, : self.n_genes]
-            batch_weights = betamat[:, self.n_genes :].copy() if D_nuisance_np is not None else None
+            batch_weights = betamat[:, self.n_genes :].copy() if n_nuisance else None
             W = torch.from_numpy(W_genes.astype(np.float32)).to(self.config.device)
             self.encoder.set_weights(W)
 
@@ -2089,8 +2169,10 @@ class BAE(nn.Module):
 
         self._is_fitted = True
 
-        # Store results in AnnData
-        self._store_results(adata)
+        # Store results in AnnData. `X_train` is handed over rather than re-read:
+        # it is the same matrix, and densifying it twice doubles the peak at the
+        # point in a fit where the most else is already resident.
+        self._store_results(adata, X_train)
 
         # Optional stability selection of the encoder's gene sets. Off by default:
         # it adds a fraction of one fit's cost (see BAE.stability_selection).
@@ -2266,41 +2348,40 @@ class BAE(nn.Module):
         if seed is not None:
             torch.manual_seed(seed)
 
-        matrix = _expression_matrix(adata, self._layer)
-        X_np = matrix.toarray() if sp.issparse(matrix) else np.asarray(matrix)
-        X_np = X_np.astype(np.float32)
-        X_train = self._to_tensor(X_np, self.config.device)
-
-        D_condition = None
-        if self._conditions_decoder:
-            D_condition = self._to_tensor(
-                transform_obs_covariates(adata, self._batch_encoding), self.config.device
+        # One encoding drives both mechanisms. Under "both" this used to re-run
+        # `transform_obs_covariates` over the same obs columns for each of them.
+        covariates_np = None
+        if self._conditions_decoder or self._regresses_encoder:
+            covariates_np = np.asarray(
+                transform_obs_covariates(adata, self._batch_encoding), dtype=np.float32
             )
-        # float32, matching `fit`. This path used to promote to float64 on top of
-        # the float32 copy above -- three copies of the expression matrix resident
-        # at once -- buying a precision difference measured at about one gene in
-        # 380, far inside the run-to-run support variation this method documents.
-        sourcemat_aug = X_np
-        n_nuisance = 0
-        if self._regresses_encoder:
-            D_nuisance = transform_obs_covariates(adata, self._batch_encoding)
-            sourcemat_aug = np.hstack([sourcemat_aug, np.asarray(D_nuisance, dtype=np.float32)])
-            n_nuisance = self._batch_encoding.n_columns
+        n_nuisance = self._batch_encoding.n_columns if self._regresses_encoder else 0
+
+        # The same design `fit` builds, float32 throughout: this path used to
+        # promote to float64 on top of the float32 copy -- three copies of the
+        # expression matrix resident at once -- buying a precision difference
+        # measured at about one gene in 380, far inside the run-to-run support
+        # variation this method documents.
+        design = self._build_boosting_design(
+            _expression_matrix(adata, self._layer),
+            covariates=covariates_np,
+            n_nuisance=n_nuisance,
+        )
+        X_train = self._to_tensor(design.panel, self.config.device)
+
+        D_condition = (
+            self._to_tensor(covariates_np, self.config.device) if self._conditions_decoder else None
+        )
 
         resolved_mandatory = resolve_mandatory_genes(self._mandatory_genes, adata)
         allboost_mandatory = _build_allboost_mandatory(resolved_mandatory, n_nuisance, self.n_genes)
-        mandatory_ridge = np.zeros(sourcemat_aug.shape[1], dtype=np.float64)
-        if n_nuisance:
-            mandatory_ridge[self.n_genes :] = self.config.nuisance_ridge
-        # Fixed across every iteration recorded here, exactly as in `fit`.
-        boosting_col_norms_sq = column_norms_sq(sourcemat_aug)
         # Honour the same covariance-cache setting `fit` resolved. Without this,
         # the method documented as mirroring the fit loop would run a different
         # cache strategy from the fit it is analysing.
         if resolve_precompute_covcache(
-            self.config.boosting_precompute_covcache, sourcemat_aug.shape[1]
+            self.config.boosting_precompute_covcache, design.matrix.shape[1]
         ):
-            covcache_initial = compute_covariance_cache(sourcemat_aug)
+            covcache_initial = compute_covariance_cache(design.matrix)
         else:
             covcache_initial = None
 
@@ -2380,22 +2461,22 @@ class BAE(nn.Module):
                     )
                 self.encoder.reset_weights()
                 fit_targets, beta_init, fit_mandatory = self._transfer_boosting_inputs(
-                    targets, sourcemat_aug.shape[1], prior_dims, allboost_mandatory
+                    targets, design.matrix.shape[1], prior_dims, allboost_mandatory
                 )
                 if fit_targets.shape[1] == 0:
-                    betamat = np.zeros((0, sourcemat_aug.shape[1]), dtype=np.float64)
+                    betamat = np.zeros((0, design.matrix.shape[1]), dtype=np.float64)
                 else:
                     result = allboost(
-                        sourcemat_aug,
+                        design.matrix,
                         fit_targets,
                         covcache=covcache,
-                        col_norms_sq=boosting_col_norms_sq,
+                        col_norms_sq=design.col_norms_sq,
                         stepno=self.config.boosting_stepno,
                         nu=self.config.boosting_nu,
                         csf=self.config.boosting_csf,
                         independent=self.config.boosting_independent,
                         mandatory_features=fit_mandatory,
-                        mandatory_ridge=mandatory_ridge,
+                        mandatory_ridge=design.mandatory_ridge,
                         beta_init=beta_init,
                         return_covcache=covcache is None,
                     )
@@ -2403,7 +2484,7 @@ class BAE(nn.Module):
                         betamat, covcache = result
                     else:
                         betamat = result
-                betamat = self._expand_transfer_betamat(betamat, sourcemat_aug.shape[1], prior_dims)
+                betamat = self._expand_transfer_betamat(betamat, design.matrix.shape[1], prior_dims)
 
                 W_genes = betamat[:, : self.n_genes]
                 self.encoder.set_weights(
@@ -3138,7 +3219,7 @@ class BAE(nn.Module):
         explained = 1.0 - ss_error / ss_total if ss_total > 0 else float("nan")
         return residual, explained
 
-    def _store_results(self, adata: AnnData) -> None:
+    def _store_results(self, adata: AnnData, X: torch.Tensor) -> None:
         """Store results in AnnData (scverse convention).
 
         Two of the recorded metrics answer different questions and are easy to
@@ -3162,9 +3243,18 @@ class BAE(nn.Module):
         :func:`~structboost.linear_ceiling` for the achievable maximum. It is
         written by every fit; only ``reconstruction_loss_by_obs`` and
         ``latent_obs_r2_per_dim`` require a covariate.
-        """
-        X = self._to_tensor(_expression_matrix(adata, self._layer), self.config.device)
 
+        Parameters
+        ----------
+        adata
+            Receives the embedding, the encoder weights and the metadata below.
+        X
+            The training matrix ``fit`` is already holding. This used to be read
+            and densified a second time here, putting two ``(n_cells, n_genes)``
+            float32 arrays on the heap at the one moment a fit has the most else
+            resident -- 2.4 GB of them at 300,000 cells by 2,000 genes. There is
+            only ever one caller, and it has the identical tensor in hand.
+        """
         # Recorded so a fitted model is itself a usable prior for a later transfer:
         # `from_reference` needs gene identifiers, which the weight matrix lacks.
         self._var_names = np.asarray(adata.var_names, dtype=object)

{structboost-0.6.0 → structboost-0.6.1}/src/structboost/_utils.py

@@ -6,12 +6,98 @@ from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
 import numpy as np
-from numpy.typing import NDArray
+from numpy.typing import DTypeLike, NDArray
 
 if TYPE_CHECKING:
     from anndata import AnnData
 
 
+#: Bytes of dense output materialized per row block by :func:`densify`. Large
+#: enough that the per-block cost of slicing a CSR matrix is negligible against
+#: the copy itself, small enough that the temporary never rivals the result.
+_DENSIFY_BLOCK_BYTES = 32 * 1024**2
+
+
+def densify(
+    matrix: object,
+    *,
+    dtype: DTypeLike = np.float32,
+    out: NDArray[np.floating] | None = None,
+    block_bytes: int = _DENSIFY_BLOCK_BYTES,
+) -> NDArray[np.floating]:
+    """Dense array of ``dtype`` from a possibly-sparse matrix, block by block.
+
+    ``matrix.toarray().astype(dtype)`` holds the dense matrix twice at once: the
+    source-dtype array returned by ``toarray`` and the converted copy. On a
+    float64 CSR panel that peaks at ``12 * n * p`` bytes to produce a ``4 * n * p``
+    byte result — 18 GB for 500,000 cells x 3,000 genes, which is where a fit
+    stops being possible rather than merely slow. Converting in row blocks
+    straight into the output removes the intermediate entirely: the peak becomes
+    the result plus one block.
+
+    **The bytes produced are identical to the unchunked expression.** ``toarray``
+    scatters stored values into a zero-filled buffer and the ``dtype`` conversion
+    is elementwise, so neither depends on how the rows are grouped. No arithmetic
+    is reassociated, and nothing here rounds differently.
+
+    Copies only when it must, like :func:`numpy.asarray`: an ndarray that already
+    has ``dtype`` is returned as-is. Callers use the result read-only and let it
+    go out of scope, so the array is never aliased into a fitted model.
+
+    Parameters
+    ----------
+    matrix
+        Dense array-like, or any SciPy sparse matrix. Non-CSR sparse formats are
+        converted to CSR first, because row slicing is what the block loop needs
+        and doing it per block on a column-major format costs ``O(nnz)`` each
+        time. That conversion allocates ``nnz``-sized arrays, not ``n * p`` ones.
+    dtype
+        Output dtype. Defaults to ``float32``, the precision a fit runs in.
+        Ignored when ``out`` is given, which carries its own.
+    out
+        Write into this array instead of allocating one, and return it. It may be
+        a *view*, which is the point: a caller building a wider design matrix can
+        hand over a column block and have the panel read straight into place,
+        rather than densifying separately and copying the result in. Unlike the
+        allocating form this always writes, so a dense input is copied.
+    block_bytes
+        Target size of one row block in the output. The block is always at least
+        one row, so a panel wider than this still works.
+
+    Returns
+    -------
+    Dense ``(n_samples, n_features)`` array of ``dtype``, or ``out``.
+    """
+    import scipy.sparse as sp
+
+    dtype = np.dtype(dtype)
+    if out is not None:
+        if tuple(out.shape) != tuple(matrix.shape):
+            raise ValueError(f"out must have shape {tuple(matrix.shape)}, got {tuple(out.shape)}")
+        dtype = out.dtype
+
+    if not sp.issparse(matrix):
+        if out is None:
+            return np.asarray(matrix, dtype=dtype)
+        # Straight elementwise conversion into place: no dense intermediate in
+        # the source dtype, exactly as in the block loop below.
+        out[...] = matrix
+        return out
+
+    n_rows, n_cols = matrix.shape
+    if out is None:
+        out = np.empty((n_rows, n_cols), dtype=dtype)
+    if out.size == 0:
+        return out
+
+    source = matrix if matrix.format == "csr" else matrix.tocsr()
+    rows_per_block = max(1, block_bytes // (n_cols * dtype.itemsize))
+    for start in range(0, n_rows, rows_per_block):
+        stop = min(start + rows_per_block, n_rows)
+        out[start:stop] = source[start:stop].toarray()
+    return out
+
+
 def compute_covariance_cache(
     sourcemat: NDArray[np.floating],
     *,

structboost-0.6.1/tests/test_densify.py

@@ -0,0 +1,236 @@
+"""Contracts for :func:`structboost._utils.densify`.
+
+The helper exists to cut peak memory, so the thing worth pinning is that it buys
+that without moving a single bit: every expression matrix a fit reads goes
+through it, and a changed byte there changes the fitted model. These tests
+therefore compare raw bytes rather than using ``allclose``.
+
+They also pin the property the helper introduced on top of that: a fit no longer
+depends on which sparse format ``adata.X`` happens to be stored in.
+"""
+
+import warnings
+
+import numpy as np
+import pytest
+
+from structboost._utils import densify
+
+sp = pytest.importorskip("scipy.sparse")
+
+#: Every sparse layout an AnnData is realistically carrying, plus both dtypes.
+FORMATS = ("csr", "csc", "coo")
+DTYPES = (np.float32, np.float64)
+
+
+def _reference(matrix, dtype=np.float32):
+    """What the call sites did before ``densify`` existed."""
+    dense = matrix.toarray() if sp.issparse(matrix) else np.asarray(matrix)
+    return dense.astype(dtype)
+
+
+@pytest.fixture
+def matrix():
+    rng = np.random.default_rng(11)
+    x = rng.standard_normal((257, 83))
+    # Zeros are what makes a sparse format worth using, and they are also the
+    # entries `toarray` writes rather than scatters.
+    x[rng.random(x.shape) < 0.6] = 0.0
+    return x
+
+
+def _as_sparse(x, fmt, dtype):
+    return getattr(sp, f"{fmt}_matrix")(x.astype(dtype))
+
+
+class TestByteIdentity:
+    @pytest.mark.parametrize("fmt", FORMATS)
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_matches_unchunked_conversion_exactly(self, matrix, fmt, dtype):
+        sparse = _as_sparse(matrix, fmt, dtype)
+        assert densify(sparse).tobytes() == _reference(sparse).tobytes()
+
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_dense_input_matches(self, matrix, dtype):
+        dense = matrix.astype(dtype)
+        assert densify(dense).tobytes() == _reference(dense).tobytes()
+
+    @pytest.mark.parametrize("block_bytes", [1, 64, 1024, 1 << 20])
+    def test_block_size_never_changes_the_bytes(self, matrix, block_bytes):
+        """The block loop is a pure partition of the rows: no reassociation, so
+        every block size -- including one row at a time -- gives the same array."""
+        sparse = _as_sparse(matrix, "csr", np.float64)
+        whole = densify(sparse)
+        assert densify(sparse, block_bytes=block_bytes).tobytes() == whole.tobytes()
+
+    def test_block_size_below_one_row_still_works(self, matrix):
+        """A panel wider than the block budget must not produce a zero-row step."""
+        sparse = _as_sparse(matrix, "csr", np.float64)
+        assert densify(sparse, block_bytes=0).tobytes() == densify(sparse).tobytes()
+
+
+class TestWritingIntoAView:
+    """``out=`` is what lets the boosting design be one allocation.
+
+    ``fit`` builds ``(n_cells, n_genes + n_nuisance)`` up front and reads the
+    panel into its gene block, instead of densifying separately and copying the
+    result in with ``hstack``.
+    """
+
+    @pytest.mark.parametrize("fmt", (*FORMATS, "dense"))
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_column_block_of_a_wider_array(self, matrix, fmt, dtype):
+        payload = matrix.astype(dtype) if fmt == "dense" else _as_sparse(matrix, fmt, dtype)
+        n_rows, n_cols = matrix.shape
+        design = np.empty((n_rows, n_cols + 3), dtype=np.float32)
+        block = densify(payload, out=design[:, :n_cols])
+
+        assert block.base is design
+        assert not block.flags["C_CONTIGUOUS"]  # a strided view, by construction
+        # Same values as the allocating form, which is what keeps a fit identical.
+        assert block.tobytes() == _reference(payload).tobytes()
+
+    def test_nuisance_columns_are_untouched(self, matrix):
+        n_rows, n_cols = matrix.shape
+        design = np.zeros((n_rows, n_cols + 2), dtype=np.float32)
+        design[:, n_cols:] = 7.5
+        densify(_as_sparse(matrix, "csr", np.float64), out=design[:, :n_cols])
+        assert (design[:, n_cols:] == 7.5).all()
+
+    def test_out_dtype_wins_over_dtype_argument(self, matrix):
+        out = np.empty(matrix.shape, dtype=np.float64)
+        assert densify(_as_sparse(matrix, "csr", np.float64), dtype=np.float32, out=out) is out
+        assert out.dtype == np.float64
+
+    def test_shape_mismatch_is_rejected(self, matrix):
+        """Silently filling the wrong block would corrupt the design matrix."""
+        sparse = _as_sparse(matrix, "csr", np.float64)
+        with pytest.raises(ValueError, match="out must have shape"):
+            densify(sparse, out=np.empty((matrix.shape[0], matrix.shape[1] + 1), np.float32))
+
+    @pytest.mark.parametrize("block_bytes", [1, 4096, 1 << 20])
+    def test_block_size_invariant_when_writing_into_a_view(self, matrix, block_bytes):
+        """The same partition property as the allocating form, but the block
+        loop now writes across a row stride, so it is worth pinning separately."""
+        sparse = _as_sparse(matrix, "csr", np.float64)
+        n_rows, n_cols = matrix.shape
+        design = np.empty((n_rows, n_cols + 3), dtype=np.float32)
+        block = densify(sparse, out=design[:, :n_cols], block_bytes=block_bytes)
+        assert block.tobytes() == densify(sparse).tobytes()
+
+    def test_torch_reads_the_block_without_copying_it(self, matrix):
+        """The assumption the single-allocation design rests on.
+
+        The gene block has a row stride of ``n_genes + n_nuisance``, which BLAS
+        takes as a leading dimension, so the encoder's matmul runs on it directly.
+        If torch ever materialized a contiguous copy instead, the second full
+        panel this construction exists to avoid would silently come back.
+        """
+        torch = pytest.importorskip("torch")
+        n_rows, n_cols = matrix.shape
+        design = np.empty((n_rows, n_cols + 4), dtype=np.float32)
+        block = densify(_as_sparse(matrix, "csr", np.float64), out=design[:, :n_cols])
+
+        strided = torch.from_numpy(block)
+        contiguous = torch.from_numpy(np.ascontiguousarray(block))
+        assert not strided.is_contiguous()
+
+        torch.manual_seed(0)
+        layer = torch.nn.Linear(n_cols, 3, bias=False)
+        with torch.no_grad():
+            from_view, from_copy = layer(strided), layer(contiguous)
+        assert from_view.numpy().tobytes() == from_copy.numpy().tobytes()
+
+
+class TestOutputContract:
+    @pytest.mark.parametrize("fmt", FORMATS)
+    def test_sparse_input_is_c_contiguous(self, matrix, fmt):
+        """Row-major whatever the source format.
+
+        ``csc.toarray()`` returns a Fortran-ordered array, which made the fitted
+        model depend on the storage format (different BLAS tiling) and handed
+        torch a non-contiguous training tensor. One layout for every input.
+        """
+        result = densify(_as_sparse(matrix, fmt, np.float64))
+        assert result.flags["C_CONTIGUOUS"]
+
+    @pytest.mark.parametrize("fmt", FORMATS)
+    @pytest.mark.parametrize("dtype", DTYPES)
+    def test_dtype_and_shape(self, matrix, fmt, dtype):
+        result = densify(_as_sparse(matrix, fmt, dtype))
+        assert result.dtype == np.float32
+        assert result.shape == matrix.shape
+
+    def test_honours_requested_dtype(self, matrix):
+        assert densify(_as_sparse(matrix, "csr", np.float64), dtype=np.float64).dtype == np.float64
+
+    def test_matching_dense_input_is_not_copied(self, matrix):
+        """`np.asarray` semantics: no needless duplicate of an already-dense panel."""
+        dense = matrix.astype(np.float32)
+        assert densify(dense) is dense
+
+    @pytest.mark.parametrize("shape", [(0, 5), (5, 0), (0, 0)])
+    def test_empty(self, shape):
+        result = densify(sp.csr_matrix(np.zeros(shape)))
+        assert result.shape == shape
+        assert result.dtype == np.float32
+
+
+def _fit(adata, **kwargs):
+    """A short, seeded fit; returns the model."""
+    from structboost import BAE, BAEConfig
+
+    model = BAE(adata.n_vars, BAEConfig(latent_dim=3, max_iterations=6, seed=5, batch_size=64))
+    model.fit(adata, verbose=False, **kwargs)
+    return model
+
+
+@pytest.fixture
+def standardized():
+    rng = np.random.default_rng(3)
+    x = rng.standard_normal((120, 40))
+    return (x - x.mean(axis=0)) / x.std(axis=0)
+
+
+class TestStorageFormatIndependence:
+    """The same data must give the same model however AnnData stored it."""
+
+    def test_csr_csc_dense_agree_bitwise(self, standardized):
+        """Regression test: ``csc.toarray()`` is Fortran-ordered, so before
+        ``densify`` the storage format changed the BLAS reduction order and with
+        it the fitted weights, by ~2e-9 on a short fit."""
+        pytest.importorskip("torch")
+        anndata = pytest.importorskip("anndata")
+
+        # AnnData stores CSR and CSC only, which is also the whole population of
+        # formats this property can be violated by.
+        def weights(payload):
+            return _fit(anndata.AnnData(X=payload)).get_encoder_weights().tobytes()
+
+        reference = weights(sp.csr_matrix(standardized))
+        for payload in (sp.csc_matrix(standardized), standardized):
+            assert weights(payload) == reference, type(payload).__name__
+
+
+class TestInputIsLeftAlone:
+    """A float32 ``adata.X`` is trained on in place rather than copied, so the
+    training tensor aliases the user's data. Nothing in a fit may write to it."""
+
+    @pytest.mark.parametrize("fit_kwargs", [{}, {"init_pca": True}, {"batch_key": "batch"}])
+    def test_fit_does_not_modify_adata_x(self, standardized, fit_kwargs):
+        pytest.importorskip("torch")
+        anndata = pytest.importorskip("anndata")
+        x = standardized.astype(np.float32)
+        adata = anndata.AnnData(X=x.copy())
+        adata.obs["batch"] = np.where(np.arange(x.shape[0]) % 2, "a", "b")
+        _fit(adata, **fit_kwargs)
+        assert adata.X.tobytes() == x.tobytes()
+
+    def test_read_only_input_fits_without_a_torch_warning(self, standardized):
+        pytest.importorskip("torch")
+        anndata = pytest.importorskip("anndata")
+        x = standardized.astype(np.float32)
+        x.setflags(write=False)
+        with warnings.catch_warnings():
+            warnings.filterwarnings("error", message=".*not writable.*")
+            _fit(anndata.AnnData(X=x))