Skip to content

[auto-optimize]: Speed up deseq2, summary and lfc_shrink - #493

Open
jmschrei wants to merge 2 commits into
scverse:mainfrom
jmschrei:auto-optimize/2026-10-31-iter-C119
Open

jmschrei wants to merge 2 commits into
scverse:mainfrom
jmschrei:auto-optimize/2026-10-31-iter-C119

Conversation

@jmschrei

@jmschrei jmschrei commented Nov 7, 2026

Copy link
Copy Markdown

This optimization was found with auto-optimize, a tool developed by the Programmable Genomics Laboratory at UMass Chan Medical School: it fits the genes of a block together, with scipy's L-BFGS-B stepped for every gene of a block in lockstep, batched IRLS and Wald tests, gene blocks as joblib tasks and threaded whole-matrix passes, while every output stays bitwise identical to main's. Across 38 timed configurations of 20 inputs, this branch runs 0.07x to 17.55x as fast as main. The exception is PyDESeq2's bundled 10-gene example at 8 threads: a loky worker imports PyDESeq2 on its first task, and with 10 genes this branch gives each map a single block, so the timed call still meets cold workers (Behaviour changes). The output was checked against main in 421 comparisons, all of them identical; see How the output was checked.

Most of the gain comes from three changes: stepping L-BFGS-B for every gene of a block together in the gene-wise and MAP dispersion fits, sending the loky workers one task per block of genes instead of one per gene, and the same lockstep L-BFGS-B in the apeGLM shrinkage. Batched IRLS and the lstsq form of fit_lin_mu come next. No OS-specific call is added: the threads use concurrent.futures, and no file is written. Two constants set the threaded passes in misc.py: _CHUNK_COLS = 1024 genes per chunk, chosen by timing balanced chunk widths of 512 to 2,048 genes on the benchmark input on one AMD EPYC machine, and _MIN_THREADED_SIZE = 1 << 18 count entries, below which the serial code runs. Neither is derived from the machine at run time. The batched Wald test takes 1,024 genes at a time, and gene blocks hold batch_size genes, 128 by default as in main. The fork's GitHub Actions ran the suite and an output comparison with main on Linux (AMD EPYC 7763 and 9V74), Windows (AMD family 25 and Intel family 6 model 207) and macOS (Apple M1 and M2 Pro) runners.

It builds on the tests pull request #492, which changes no source. It is one commit changing nine files, +2,957 / -173. It is opened against main with the tests commit beneath it, so until #492 is merged GitHub shows both commits here; the change under review is the last one. The lockstep L-BFGS-B drivers adapt code from SciPy 1.18.1 (scipy/optimize/_lbfgsb_py.py, _minimize_lbfgsb; BSD-3-Clause, with that file's MIT notice by David M. Cooke): the per-gene L-BFGS-B state allocation and setulb loop in _fit_alpha_mle_lockstep and _lbfgsb_step in dispersions.py, and in _nbinom_glm_lockstep in distributions.py. Both files carry the two notices verbatim in a comment block above that code. Nothing else is vendored or translated from another project.

Usage benchmarks

Input class Inputs Configurations Speedup, geometric mean Speedup, median Speedup, range Peak memory, median ratio (this branch / main) Output
GTEx bulk RNA-seq, the benchmark 1 5 8.84x 9.18x 6.74x to 10.54x 0.99x identical
GTEx bulk RNA-seq 10 17 3.10x 2.42x 1.23x to 17.55x 0.99x identical
synthetic 9 12 3.23x 5.17x 0.07x to 15.26x 1.00x identical
lowest supported dependencies 3 4 4.15x 5.10x 1.65x to 7.63x 0.98x identical
all 20 38 3.72x 3.51x 0.07x to 17.55x 0.99x identical
All 38 configurations
Input class Input Threads main (s) This branch (s) Speedup main peak (MB) This branch peak (MB) Output
GTEx bulk RNA-seq, the benchmark 50 lung vs 50 thyroid, every gene (seed 0, int32) 1 65.4 6.28 10.41x 883 849 identical
GTEx bulk RNA-seq, the benchmark 50 lung vs 50 thyroid, every gene (seed 0, int32) 2 38.9 3.69 10.54x 1,142 1,138 identical
GTEx bulk RNA-seq, the benchmark 50 lung vs 50 thyroid, every gene (seed 0, int32) 4 20.8 2.27 9.18x 1,377 1,341 identical
GTEx bulk RNA-seq, the benchmark 50 lung vs 50 thyroid, every gene (seed 0, int32) 8 12.0 1.52 7.94x 1,758 1,747 identical
GTEx bulk RNA-seq, the benchmark 50 lung vs 50 thyroid, every gene (seed 0, int32) 16 8.77 1.30 6.74x 2,572 2,572 identical
GTEx bulk RNA-seq lung (655) vs thyroid (706), every gene, ~tissue 1 135 53.5 2.52x 9,049 8,022 identical
GTEx bulk RNA-seq lung (655) vs thyroid (706), every gene, ~tissue 2 78.5 32.5 2.42x 9,526 9,619 identical
GTEx bulk RNA-seq lung (655) vs thyroid (706), every gene, ~tissue 4 43.5 19.5 2.23x 9,734 8,819 identical
GTEx bulk RNA-seq lung (655) vs thyroid (706), every gene, ~tissue 8 26.3 13.1 2.01x 10,210 10,640 identical
GTEx bulk RNA-seq lung (655) vs thyroid (706), every gene, ~tissue 16 19.8 8.76 2.25x 11,050 11,100 identical
GTEx bulk RNA-seq muscle (881) vs tibial nerve (659), every gene, ~sex + tissue 1 163 71.9 2.26x 9,477 9,128 identical
GTEx bulk RNA-seq muscle (881) vs tibial nerve (659), every gene, ~sex + tissue 8 30.7 18.1 1.69x 11,410 10,790 identical
GTEx bulk RNA-seq muscle (881) vs tibial nerve (659), every gene, ~sex + tissue 16 22.2 11.8 1.88x 12,280 11,750 identical
GTEx bulk RNA-seq lung vs thyroid with a RIN (1,361), ~sex + rin + tissue 8 29.4 17.6 1.68x 10,120 10,550 identical
GTEx bulk RNA-seq lung (655), male vs female, every gene, ~sex 8 19.4 6.14 3.16x 5,425 5,624 identical
GTEx bulk RNA-seq lung (655) vs thyroid (706), every gene, ~tissue, vst() 8 9.95 8.12 1.23x 5,564 5,743 identical
GTEx bulk RNA-seq 5 lung vs 5 thyroid, every gene 1 50.1 2.86 17.55x 245 213 identical
GTEx bulk RNA-seq 5 lung vs 5 thyroid, every gene 8 9.22 0.711 12.96x 1,148 1,066 identical
GTEx bulk RNA-seq 3 lung vs 3 thyroid, 20,000 genes 8 2.79 0.304 9.20x 1,026 1,000 identical
GTEx bulk RNA-seq 250 lung vs 250 thyroid, 5,000 genes 8 1.74 0.627 2.78x 1,320 1,760 identical
GTEx bulk RNA-seq 30 lung vs 30 thyroid, 10,000 genes, ~sex + rin + tissue 8 2.23 0.579 3.86x 1,064 1,063 identical
GTEx bulk RNA-seq 250 muscle vs 250 tibial nerve, every gene, ~sex + tissue 8 18.9 5.37 3.51x 4,563 4,506 identical
synthetic PyDESeq2's bundled data, 100 x 10 1 0.0403 0.021 1.91x 142 142 identical
synthetic PyDESeq2's bundled data, 100 x 10 8 0.0943 1.31 0.07x 974 821 identical
synthetic 8 vs 8, 3,000 genes with degenerate counts 8 1.03 0.235 4.40x 995 999 identical
synthetic 3 groups x 6, 20,000 genes 8 4.19 0.457 9.17x 1,060 1,035 identical
synthetic 40 samples, 15,000 genes, ~age 8 3.39 0.456 7.44x 1,088 1,056 identical
synthetic 1,000 vs 1,000, 20,000 genes 1 51.1 24.0 2.13x 3,215 3,350 identical
synthetic 1,000 vs 1,000, 20,000 genes 8 10.4 6.53 1.60x 4,448 4,620 identical
synthetic 1,000 vs 1,000, 20,000 genes 16 8.27 4.20 1.97x 5,464 6,387 identical
synthetic 2 vs 2, 200 genes 1 0.224 0.031 7.25x 143 145 identical
synthetic 2 vs 2, 60,000 genes 8 11.0 0.72 15.26x 1,162 1,033 identical
synthetic 10 vs 10, 5,000 genes, outliers in 20% 8 1.57 0.264 5.94x 1,009 1,012 identical
synthetic 3 batches x (4 vs 4), 10,000 genes, ~batch + condition 8 2.36 0.328 7.19x 1,031 1,012 identical
lowest supported dependencies 50 lung vs 50 thyroid, every gene (seed 0, int32) 1 64.0 8.38 7.63x 818 784 identical
lowest supported dependencies 50 lung vs 50 thyroid, every gene (seed 0, int32) 8 11.8 1.76 6.69x 1,658 1,627 identical
lowest supported dependencies lung (655) vs thyroid (706), every gene, ~tissue 8 26.1 15.8 1.65x 10,030 9,744 identical
lowest supported dependencies 30 lung vs 30 thyroid, 10,000 genes, ~sex + rin + tissue 8 2.21 0.628 3.52x 996 986 identical
How it was timed

Machine: AMD EPYC 9575F 64-Core Processor, 256 cores; Linux 7.0.0-38-generic; Python 3.12.13; numpy 2.5.4, scipy 1.18.1 with its bundled OpenBLAS, pandas 3.0.6, joblib 1.6.0, scikit-learn 1.9.1, anndata 0.13.4; the four rows marked lowest supported dependencies use numpy 2.2.0, scipy 1.15.0, pandas 2.3.0, scikit-learn 1.6.0 and anndata 0.11.0. Two sockets of 64 cores, two NUMA nodes, 1 TB of memory. Both builds are pure Python: main and this commit were each installed with uv pip install from a git archive export into a fresh Python 3.12.13 environment with the same dependency versions; the lowest-dependency rows install the same two sources with the lower versions.

Inputs: Real inputs are recount3 GTEx v8 gene counts for lung, thyroid, skeletal muscle and tibial nerve (gene_sums divided by each sample's STAR average mapped length and rounded), obtained from the GTEx Portal through recount3 on 10/28/2026; synthetic inputs are negative binomial counts with log-normal gene means, a dispersion of 0.05 + 1/mean with log-normal noise and log-normal size factors, drawn with numpy's default_rng from fixed seeds 1 to 8, plus PyDESeq2's bundled synthetic data and an 8 vs 8 GTEx cut with all-zero, zero-inflated, outlier and 1e8-scaled genes (seed 20261028). Every row calls DeseqDataSet(counts=..., metadata=..., design=..., n_cpus=t, quiet=True), deseq2(), DeseqStats(dds, contrast=..., n_cpus=t, quiet=True), summary() and lfc_shrink(coeff=...) with default options, on counts in a pandas DataFrame:

  • The benchmark: 50 lung and 50 thyroid samples drawn with seed 0, every gene (63,856), int32 counts, ~tissue, thyroid vs lung.
  • Full GTEx experiments (every sample of the tissues, every gene, int64 counts): lung (655) vs thyroid (706), ~tissue, and the same with vst() alone; muscle (881) vs tibial nerve (659), ~sex + tissue; every lung and thyroid sample with a RIN (1,361), ~sex + rin + tissue with RIN centred; lung alone, male vs female, ~sex.
  • GTEx cuts: 5 vs 5 (a seeded draw) on every gene; 3 vs 3 on 20,000 genes; 250 vs 250 on 5,000 genes; 30 vs 30 on 10,000 genes, ~sex + rin + tissue; 250 muscle vs 250 nerve on every gene, ~sex + tissue.
  • Synthetic: PyDESeq2's bundled 100 x 10 data; the degenerate 8 vs 8 cut above; 3 groups of 6 on 20,000 genes; 40 samples, ~age with age continuous and a numeric contrast; 1,000 vs 1,000 on 20,000 genes; 2 vs 2 on 200 and on 60,000 genes; 10 vs 10 on 5,000 genes with one outlier count in 20% of genes; 3 batches of 4 vs 4 on 10,000 genes, ~batch + condition.

How it was timed:

  • Isolation: Each process was confined to n_cpus physical cores, one logical CPU each, within whole L3 domains (8 cores each) of NUMA node 0, holding every core of those domains so that no other run shared their L3 cache; the timed rows ran one at a time. The output-only runs of How the output was checked used up to 16 cores of NUMA node 1, the other socket, during most of the timing. Timed again at the end with nothing else of this testing running, all five rows agreed with the table within 2% on both builds. Other users' jobs left the machine 91 to 99% idle.
  • Repeats: One process per build and configuration, three for the full experiments, the two builds alternating which ran first; a call that took under 1 s was repeated five times in its process. The table gives the median.
  • What the time includes: Untimed first, loading the counts into a DataFrame and one run of the same calls on every 21st gene (the whole input below 300 genes), which imports everything and starts the loky workers; the time then covers DeseqDataSet(...) through lfc_shrink() (for the vst() row, DeseqDataSet(...) and vst()), measured with time.perf_counter.
  • Peak memory: Each process ran in its own systemd scope, and the table gives the scope's memory.peak over the timed calls (reset just before them): the process and its loky workers, each page once, and any page cache charged during the calls.
  • Output check: Every array a user can read after the calls (each column of obs and var, each varm, obsm, layers and uns entry, the results table and the shrunk table) was compared with main's byte for byte.

Outside the table, the 347 output-only configurations that completed on both builds ran two at a time, so their times are not controlled: they add up to 3,399 s on main and 806 s on this branch with current dependencies (263 configurations), 313 s and 51 s at the lowest supported versions (46), and 386 s and 75 s under the Haswell kernel (38).

Changes

In the order they were made:

  1. Gene blocks as joblib tasks (src/pydeseq2/default_inference.py): The per-gene maps of DefaultInference (lin_reg_mu, irls, alpha_mle, wald_test, lfc_shrink_nbinom_glm) send one joblib task per block of batch_size genes (128 by default) instead of one task per gene. Each task receives its block's per-gene inputs as one C-contiguous (genes, samples) array, runs a batched fit or the unchanged per-gene function on each row, and returns stacked results that the method concatenates in gene order.
    • With batch_size="auto" the blocks hold 128 genes and joblib batches the blocks itself; a value joblib rejects is still passed to joblib, which raises as on main.
    • Parallel is built with max_nbytes=None, so joblib pickles the blocks instead of memory-mapping large arguments.
  2. Lockstep L-BFGS-B for the dispersion fits (src/pydeseq2/default_inference.py, src/pydeseq2/dispersions.py): The gene-wise, MAP and refit dispersions of a block (dispersions.py: _fit_alpha_mle_lockstep, _lbfgsb_step, _alpha_mle_loss_grad; default_inference.py: _alpha_mle_block) are fitted by a driver that keeps one L-BFGS-B state per gene and calls scipy's setulb for each gene as scipy.optimize.minimize would, while the loss and gradient of all genes waiting for an evaluation are computed in one vectorized pass. The per-gene state allocation and setulb loop are adapted from scipy's _minimize_lbfgsb.
    • Private scipy APIs: scipy.optimize._lbfgsb.setulb, the signature of scipy.optimize._lbfgsb_py._minimize_lbfgsb (read with inspect for the defaults of maxcor, ftol, gtol, maxfun, maxiter and maxls, so a scipy that changes them is followed) and scipy.optimize._lbfgsb_py.HAS_ILP64 (int32 when it is absent).
    • When setulb's documented signature is not the one the driver calls, when an import fails, or when _minimize_lbfgsb lacks one of the six parameters, _alpha_mle_block runs the original per-gene fit_alpha_mle. The check runs once per process.
    • Inputs the driver does not take (a design that is not 2-D float64, initial dispersions that are not float64, bounds that are not finite with min_disp < max_disp, prior_reg without prior_disp_var, the BFGS optimizer) run the per-gene path. Genes that do not converge go to grid_fit_alpha with the same arguments as before.
  3. Lockstep L-BFGS-B for LFC shrinkage (src/pydeseq2/default_inference.py, src/pydeseq2/distributions.py): The apeGLM fits of lfc_shrink_nbinom_glm (distributions.py: _nbinom_glm_lockstep, _nbinom_glm_rows, _is_real_scalar; default_inference.py: _nbinom_glm_block) use the same kind of driver: one unbounded setulb state per gene with ftol = gtol = 1e-8 as nbinomGLM sets them, and the objective and gradient of all waiting genes in one vectorized pass.
    • It runs under the same setulb check. When that check fails, and when an input is not a 2-D float64 C- or F-contiguous design, integer or float64 counts, float64 size and offset, real scalar prior scales or a valid shrink index, _nbinom_glm_block runs nbinomGLM per gene.
    • Genes of a two-coefficient design that do not converge still go through grid_fit_shrink_beta one at a time, and the inverse Hessians come from one stacked np.linalg.inv.
  4. Batched IRLS and the trend filter (src/pydeseq2/dds.py, src/pydeseq2/default_inference.py, src/pydeseq2/glm.py): IRLS (glm.py: _irls_rows, _irls_rows_steps, _solve_rows, _solve_takes_stacks; default_inference.py: _irls_block) runs for all genes of a block together with an active set, as stacked products and solves laid out as the per-gene computation lays out its operands. A gene whose step leaves [-max_beta, max_beta] or reaches maxiter is refitted from the start by irls_solver. Inputs the batched code does not handle, or any exception, send the whole block to irls_solver per gene, which raises or warns as on main.
    • scipy.linalg.solve is called once on the stack when it accepts stacks (checked once at run time; scipy before 1.16 does not) and the design has more than one column, else per gene.
      With unique gene names, DeseqDataSet._fit_parametric_dispersion_trend keeps its targets and covariates as arrays plus an array of kept positions, and drops the genes with non-finite covariates with one np.isfinite mask.
    • Duplicated names keep the label code, where the non-finite genes are dropped with one mask and one drop.
  5. Linear mu_hat without scikit-learn's per-call checks (src/pydeseq2/default_inference.py, src/pydeseq2/glm.py): fit_lin_mu (glm.py: fit_lin_mu, _fit_lin_mu_rows, _lstsq_takes_stacks, _linear_regression_cond) makes the one call LinearRegression(fit_intercept=False).fit makes, scipy.linalg.lstsq(X, y, cond=..., check_finite=False), then X @ coef + 0.0, without the estimator's input validation. cond is the value the installed scikit-learn passes: its tol from 1.9 on, max(X.shape) * eps for 1.6 to 1.8. A design or quotient that is not a finite, non-empty float64 ndarray, or a scikit-learn version that cannot be read, uses LinearRegression as before.
    • _lin_reg_mu_block fits a whole block with one stacked lstsq call where scipy accepts stacks (scipy before 1.16 does not), keeping one matrix-vector product per gene. A non-finite value, an np.errstate whose under is not "ignore", LinAlgError or ValueError sends the block to fit_lin_mu per gene.
  6. Wald test over all genes in the calling process (src/pydeseq2/default_inference.py, src/pydeseq2/glm.py): DefaultInference.wald_test (glm.py: _wald_test_rows, _wald_test_block) computes the test in the calling process, 1,024 genes at a time, with stacked products, a stacked np.linalg.inv, p-values from scipy.special.ndtr(-x) (the function norm.sf evaluates) and element-wise forms of the four alt_hypothesis branches, all evaluated as wald_test evaluates them. Inputs it does not handle (arrays that are not float64 of the expected shapes; a contrast Series is accepted), or any exception, run the per-gene joblib path, which raises main's exception.
  7. Independent filtering on sorted p-values (src/pydeseq2/ds.py): DeseqStats._independent_filtering (ds.py: _sorted_tested_p_values, _bh_sorted) sorts the tested p-values once together with their base means and gene positions. For each of the 50 cutoffs, the sorted p-values of the genes at or above it are adjusted with the Benjamini-Hochberg arithmetic of scipy's false_discovery_control (p * (m / rank), a reverse running minimum, a clip to [0, 1]) and the values below alpha are counted. Only the chosen cutoff's adjusted p-values are scattered back to gene order; no genes x 50 frame is built.
    • p-values or base means that are not float64 Series indexed by var_names, or a p-value outside [0, 1], keep the per-cutoff false_discovery_control loop.
  8. Gene-wise means and dispersions in one map (src/pydeseq2/dds.py, src/pydeseq2/default_inference.py, src/pydeseq2/dispersions.py): The terms of the dispersion loss that do not depend on the dispersion are computed once per block, and the shared terms once per evaluation. Each block computes its mu_hat (lin_reg_mu, or irls on designs with more distinct rows than coefficients) and then its gene-wise dispersions in the same task (default_inference.py: DefaultInference._genewise_dispersions, _genewise_block, _fits_genewise_in_one_map; dds.py: DeseqDataSet.fit_genewise_dispersions), so fit_genewise_dispersions makes one map instead of two. The _mu_hat layer is filled from the returned blocks as before.
    • This runs only for a DefaultInference that overrides none of lin_reg_mu, irls, alpha_mle and _genewise_dispersions. Any exception, or no gene, runs the two public calls, which raise as on main.
  9. Threaded whole-matrix passes in the calling process (src/pydeseq2/dds.py, src/pydeseq2/dispersions.py, src/pydeseq2/misc.py, src/pydeseq2/preprocessing.py): The median-of-ratios size factors (preprocessing.py: _threadable_counts, _column_bounds, _deseq2_norm_fit_threads, _deseq2_norm_transform_threads) and the Cook's distances (dds.py: DeseqDataSet.calculate_cooks, _cooks_threads; dispersions.py: _robust_moments_plan, _robust_moments_disp_columns) run on chunks of about 1,024 genes, or of samples for the medians, on threads, since numpy releases the GIL in these loops. The thread count is n_cpus read as joblib's n_jobs and capped at joblib.cpu_count() (misc.py: _thread_count, _chunk_bounds, _balanced_width, _thread_pool, _map_threads, _FPFlag); it is 1 for an inference without an integer n_cpus.
    • The serial code runs from the start for counts with fewer than 2^18 entries or 1,024 or fewer genes (non-zero genes for the Cook's distances), counts that are not int32 or int64 (or, for the size factors, not contiguous), the iterative and poscounts size factors, one thread, and a floating-point flag or exception in any chunk, so warnings follow the caller's error state as on main.
    • _robust_moments_plan computes the design-only part of the robust dispersions once; the public robust_method_of_moments_disp and trimmed_cell_variance are unchanged.
  10. Fused maps, threaded moments and positional writes (src/pydeseq2/dds.py, src/pydeseq2/default_inference.py, src/pydeseq2/dispersions.py, src/pydeseq2/distributions.py, src/pydeseq2/ds.py): Several independent changes, each falling back to main's code for inputs and DefaultInference methods it does not replace:
  • MAP dispersions and LFCs in one map (dds.py: DeseqDataSet.deseq2, _fit_MAP_and_LFC, _MAP_and_LFC_map, _stashed_LFC, fit_MAP_dispersions, fit_LFC, _same_bytes; default_inference.py: DefaultInference._MAP_dispersions_and_LFC, _map_lfc_block, _fits_MAP_and_LFC_in_one_map): in deseq2(), and in the Cook's refit when the one-map refit does not apply, each block fits its MAP dispersions and then runs IRLS on the dispersions fit_MAP_dispersions will store (the clipped MAP value, or the gene-wise value for outlier genes), so the LFCs need no second map. Both public methods still run and write every field with their own code. fit_LFC uses the stored fits only when the dispersions, size factors and design it reads equal byte for byte what IRLS was given and the inference, X, non_zero_idx, min_mu and beta_tol are the same objects; otherwise it calls irls as before. This runs only for a DefaultInference that overrides none of alpha_mle, irls and _MAP_dispersions_and_LFC; any exception runs the two public maps.
  • Cook's refit in one map (dds.py: DeseqDataSet._refit_in_one_map, _refit_trend_and_prior; default_inference.py: DefaultInference._refit_dispersions_and_LFC, _refit_block, _fits_refit_in_one_map): the refit's gene-wise dispersions, MAP dispersions and IRLS run per block in one map, since the refit takes its trend and prior from the main fit and every later step is per gene. The fields copied back are built with the same pandas statements as fit_genewise_dispersions, fit_MAP_dispersions and fit_LFC, and fit_LFC runs again if the dispersions IRLS was given are not byte-equal to what it would read. A DefaultInference that overrides any of lin_reg_mu, irls, alpha_mle, _genewise_dispersions and _refit_dispersions_and_LFC, any exception, or no gene to fit runs the three steps as before.
  • Cook's refit zeroing in one step (dds.py: DeseqDataSet._refit_without_outliers): replace_cooks is zeroed for all replaceable samples and refitted genes in one np.ix_ assignment instead of one gene at a time. A layer that is not a 2-D ndarray, or an obs["replaceable"] that is not a bool mask of the samples, keeps the per-gene loop.
  • Method-of-moments dispersions on threads (dds.py: DeseqDataSet._fit_MoM_dispersions, _mom_dispersions_threads): when the inference's fit_rough_dispersions and fit_moments_dispersions are PyDESeq2's own and more than one thread is allowed, the regression over all non-zero genes is one scipy.linalg.lstsq call in the calling thread, the call LinearRegression.fit makes (with check_finite=False); the gathering of the non-zero genes, the moments and rough estimates, their minimum and the clip run on column chunks on threads. Fewer than 2^18 normed-count entries, 1,024 or fewer non-zero genes, a design that is not float64 with more samples than columns, an unreadable scikit-learn version, a non-finite or all-zero column, a floating-point flag, a warning from lstsq, or an exception run the original code.
  • Trend refits by position: with unique gene names, each refit of the dispersion trend drops the genes far from the curve by position, without per-gene label lookups or Series.drop; the Series given to dispersion_trend_gamma_glm are rebuilt from the kept positions with the same values, dtype, name and index labels. Predictions that are not an ndarray of the kept genes' shape (from a custom inference) switch to the label code from that pass on.
  • lfc_shrink's standard errors and writes (ds.py: DeseqStats.lfc_shrink): the shrunk SEs are one np.sqrt(np.abs(inv_hessians[:, i, i])) when the inverse Hessians are a 3-D float64 array, and LFC, SE and the convergence flags are written by position when coeff_idx is an int, the gene names are unique and the LFC and SE indexes equal var_names. Other forms, such as a custom inference returning a list of Hessians, keep the original code.
  • Lockstep driver details: the apeGLM objective computes counts + size once per block, and both lockstep L-BFGS-B drivers read setulb's task word with scalar reads.
  1. Licence notices, dependency range, hooks and tests (src/pydeseq2/dds.py, src/pydeseq2/default_inference.py, src/pydeseq2/dispersions.py, src/pydeseq2/distributions.py, src/pydeseq2/ds.py, src/pydeseq2/glm.py, src/pydeseq2/misc.py, src/pydeseq2/preprocessing.py, tests/test_threaded_passes.py): - Notices: dispersions.py (above _SETULB_SIGNATURE) and distributions.py (above _nbinom_glm_lockstep) carry SciPy's BSD-3-Clause license and the David M. Cooke MIT notice of scipy/optimize/_lbfgsb_py.py verbatim.
  • scipy 1.15 to 1.17: the batched fits run when setulb.__doc__, with its whitespace removed, starts with the signature the driver calls, which holds on scipy 1.15 to 1.18 (1.15 to 1.17 print it without spaces).
  • longdouble bounds: the log bounds of the dispersion fit are converted with float(), as minimize converts bounds to float64, so an np.longdouble min_disp or max_disp gives main's result.
  • Repository hooks: the new code passes the repository's pre-commit hooks (ruff, ruff format, mypy). The except Exception blocks, each guarding a fast path whose fallback reproduces main's error, carry # noqa: BLE001, and the remaining mypy findings carry # type: ignore[<code>] as elsewhere in the package.
  • Tests for the threaded passes: tests/test_threaded_passes.py builds inputs just above 2^18 count entries with more than 1,024 genes (24 samples x 11,200 genes, and 260 samples x 1,100 genes) and checks that every output of the public calls with n_cpus=2 equals the n_cpus=1 call's, NaN positions included: two groups through deseq2() and step by step, a continuous covariate, 260 samples, vst() with and without the design, duplicated gene names, and a custom inference returning float32 dispersions. No other test in the suite has an input of this size, so none reaches the threaded passes.
  1. Memory of the refit and the moments prediction (src/pydeseq2/dds.py): - The counts_to_refit AnnData that _refit_without_outliers replaces has its obsm cleared, so its copies of _mu_LFC and _hat_diagonals are freed at once instead of by the cyclic garbage collector, which can run after the next allocation peak.
  • The threaded method-of-moments dispersions compute the prediction X @ coef once over the whole matrix in the calling thread, as LinearRegression.predict makes it, instead of per chunk, where BLAS rounded the chunk edges differently from the whole-matrix product and the results moved with the thread count.

Behaviour changes

  • Outputs: on Linux x86-64, every output compared is bitwise identical to main's under OpenBLAS's SkylakeX and Haswell kernels (How the output was checked). On the fork's GitHub Actions, 403 of 408 comparisons with main were bitwise identical: every output on Apple M1 (19 runners), AMD EPYC 7763 and 9V74 (Linux), and AMD family 25 and Intel family 6 model 207 (Windows), at n_cpus 1 and 2. On the one Apple M2 Pro runner, 27 of 44 synthetic outputs were bitwise identical and the rest differed by at most 1.96e-12 relative (Cook's distances; LFCs and p-values by at most 1.8e-13); PyDESeq2's example data was bitwise identical there, and this branch gave the same bits at n_cpus 1 and 2. The cause of the M2 Pro difference was not investigated.
  • Progress output with quiet=False: the same lines in the same order, but some are printed after the single map they now describe. "Fitting dispersions..." is printed after the one map of fit_genewise_dispersions, and its time covers mu_hat. Within deseq2() and refit(), "Fitting MAP dispersions..." includes the LFC fit and "Fitting LFCs..." only collects it, and the Cook's refit prints its three pairs after its one map. With joblib_verbosity > 0, joblib counts blocks of batch_size genes, deseq2() makes fewer maps, and the Wald test prints no joblib progress.
  • Threads and workers: n_cpus also sizes short-lived ThreadPoolExecutor pools in the calling process (read as joblib's n_jobs, capped at joblib.cpu_count()) for the size factors, the method-of-moments dispersions and the Cook's distances, used only with at least 2^18 count entries and more than 1,024 genes. The Wald test runs in the calling process with its BLAS threads, so its floating-point warnings, if any, are issued there (main: on loky workers with one BLAS thread each when n_cpus > 1). With n_cpus=1 no pool is used.
  • Failures: on any exception, the one-map fits fall back to the separate public calls, the threaded passes to the serial code, and batched IRLS and the batched Wald test to main's per-gene functions, so a call that fails in one of them raises after doing its work twice, from the code it falls back to. The batched L-BFGS-B fits and the stacked lstsq choose the per-gene code from their inputs (and, for lstsq, on LinAlgError or ValueError); any other exception raised inside them propagates from the batched code.
  • First calls on inputs with few genes: a loky worker imports PyDESeq2 when it runs its first task. On PyDESeq2's bundled 10-gene data at n_cpus=8, main's joblib sends each map as ten one-gene tasks, so its first call reaches most workers; this branch sends each map as one block on one worker, so workers that have not run a task yet are still cold in later calls, and those calls can take longer than main's (Usage benchmarks). Inputs with at least batch_size * n_cpus non-zero genes give each map at least one task per worker.
  • pandas 3.1.0 deprecation warning: main's trend fit calls Series.drop(inplace=True) on every pass, which warns under pandas 3.1.0. On this branch that call remains only on the label-based path (duplicated gene names, or a custom inference whose predictions are not an ndarray of the kept genes' shape), so with unique names and DefaultInference the warning no longer appears.
  • Unusual input: duplicated gene names, non-float64 designs, dispersions, size factors or offsets, and counts that are neither integer nor float64 take main's label-based or per-gene code. np.longdouble min_disp and max_disp give main's results; a longdouble min_mu or prior_LFC_var raises main's TypeError.
  • Custom Inference subclasses: the one-map fits run only on a DefaultInference that overrides none of the methods they replace, and the threaded method-of-moments dispersions only when fit_rough_dispersions and fit_moments_dispersions are PyDESeq2's own; a subclass that overrides one, and any other Inference, get main's call sequence. A subclass that only calls super() gets the block implementations. The threaded size-factor and Cook's passes call no inference method and run for any Inference with an integer n_cpus.
  • Dependency versions: the batched L-BFGS-B fits run on every supported scipy, 1.15 to 1.18; a scipy whose setulb changes its signature would send every fit back to scipy.optimize.minimize per gene. Stacked solve and lstsq need scipy 1.16 or later; on 1.15 IRLS and _lin_reg_mu_block solve gene by gene. The cond given to lstsq follows the installed scikit-learn. Independent filtering reproduces the arithmetic of scipy's false_discovery_control and would not follow a change to it.
  • Open pull requests (read from the code; their merged outputs were not run on this branch): Fix prior curvature in shrinkage Hessian #470, Preserve dispersion regularization in optimizer fallback #471, Honor the requested coefficient in shrinkage grid fallback #472, Store dispersion trend coefficients as a dict so the dataset can be written to h5ad #477 and Speed up iterative size factor fitting by not slicing AnnData in the objective #479 merge onto this branch without conflicts; Add tximport-style transcript-length normalization and pytximport support #451 conflicts in dds.py, default_inference.py and misc.py, and feat: add likelihood ratio test (LRT) to DeseqStats #452 already conflicts with main. Fix prior curvature in shrinkage Hessian #470 (the prior's curvature in nbinomGLM's Hessian), Preserve dispersion regularization in optimizer fallback #471 (the dispersion prior passed to the grid fallback of fit_alpha_mle) and Honor the requested coefficient in shrinkage grid fallback #472 (shrink_index passed to grid_fit_shrink_beta) fix per-gene functions that the batched code does not call, so their fixes would not take effect here; whichever of them lands second needs the same change in _nbinom_glm_rows, _fit_alpha_mle_lockstep or _nbinom_glm_lockstep.

Tests and CI

The tests suite: 373 passed, 0 skipped and 0 failed on this branch. The repository's two workflows and its pre-commit hooks were run locally on this commit's tree, on Linux x86-64, with this commit's parent, the tests pull request's head, run the same way as the control. Each hatch-test environment was created fresh with uvx hatch env create and run with the workflow's command, uvx hatch run <env>:run-cov -v --color=yes -n 4, on 8 dedicated cores; the workflow's -n auto gives 4 workers on GitHub's 4-core runners.

Job Result
get-environments (uvx hatch env show --json) passed
hatch-test.py3.12-stable with pandas 3.0.6: 373 passed; as resolved today (pandas 3.1.0): 322 passed, 49 failed, 2 errors, the same 51 tests that fail on the parent
hatch-test.py3.14-stable with pandas 3.0.6: 373 passed; as resolved today (pandas 3.1.0): 322 passed, 49 failed, 2 errors, the same 51 tests that fail on the parent
hatch-test.py3.14-pre (PRE-RELEASE DEPENDENCIES) 322 passed, 49 failed, 2 errors, the same 51 tests that fail on the parent (the job is continue-on-error)
generate coverage report passed: 91.6% of 2,309 lines with pandas 3.0.6 (194 not reached)
Upload coverage not run: it needs the repository's Codecov token
Tests pass in all hatch environments fails on the parent and on this commit alike while pandas 3.1.0 is resolved
Check Build: uv build passed
Check Build: uvx twine check --strict dist/*.whl passed
pre-commit, all 13 hooks at their pinned revisions passed

The 51 failures under pandas 3.1.0 come from main's test setup, not from this change: 50 are Pandas4Warnings about the future.infer_string option, which anndata 0.13.4 reads and the suite's filterwarnings = error turns into errors, and 1 is the deprecated inplace keyword of DataFrame.rename in tests/test_pydeseq2.py. The list of failing tests is identical on the parent and on this commit in all three environments. Outside the workflows, the suite also passes in full on both builds (373 tests here, 365 on the parent) at the lowest supported versions (numpy 2.2.0, pandas 2.3.0, scipy 1.15.0, scikit-learn 1.6.0, anndata 0.11.0, formulaic 1.0.2, with pyparsing below 3.3, as matplotlib 3.10.0 warns on import under 3.3 and the suite makes warnings errors) and with pandas 2.3.3. On the fork's GitHub Actions, with the workflow's matrix run on ubuntu-latest, windows-latest and macos-latest, the 373 tests failed exactly main's 51 under pandas 3.1.0 in every environment and none under pandas below 3.1, the build passed, and no pydeseq2-* temporary folder was left in any of 18 test and 24 comparison jobs. The Codecov upload needs the repository's token. Line coverage is 91.6% on this commit (194 of 2,309 lines not reached) against 94.2% on the parent.

How the output was checked

The output was checked against main in 421 comparisons, all of them identical. They covered 26 inputs from 2 assays (10 GTEx count matrices and 16 synthetic count matrices), among them 3 malformed count matrices and 4 edge-case designs and count matrices. Each comparison is a pair of fresh processes of main and this branch on one input, at one thread count, with one set of options: 1, 2, 4, 8 and 16 threads, the default calls and nine option variants, under two sets of dependency versions and two of OpenBLAS's CPU kernels. Each run saved every array a user can read after deseq2(), summary() and lfc_shrink(), or after vst(): each column of obs and var, each varm, obsm, layers and uns entry, the results table of every contrast and the shrunk table, together with the warnings raised, the progress text printed under quiet=False (numbers masked) and the exception, if one was raised. The 17,497 arrays were compared byte for byte.

Bugs in main found along the way. The tests pull request's suite found bugs in main, filed as #491. Each bug's reproduction was run again against this branch's install, and each prints the same output as on main except the pandas 3.1.0 deprecation warning from Series.drop(inplace=True) in _fit_parametric_dispersion_trend, which does not occur on this branch when the gene names are unique and the trend's columns are float64, because the trend fit drops genes by position there.

Builds, commands, results and limits
Pair Built (main) Built (this branch) Commands Identical Identical failure Differ Why
Current dependencies the tip of the tests pull request (main at 2ca3b74 plus test files only), installed with uv pip install from a git archive export into a fresh Python 3.12.13 environment: numpy 2.5.4, scipy 1.18.1, pandas 3.0.6, joblib 1.6.0, scikit-learn 1.9.1, anndata 0.13.4 this branch merged onto that tip, installed the same way 327 297 30 0 every fast path runs; OpenBLAS picks its SkylakeX kernel on this CPU
Haswell kernel the same install, run with OPENBLAS_CORETYPE=Haswell the same 38 38 0 0 the kernel OpenBLAS picks on AVX2 CPUs without AVX-512
Lowest supported dependencies the same source with numpy 2.2.0, scipy 1.15.0, pandas 2.3.0, scikit-learn 1.6.0, anndata 0.11.0, formulaic 1.0.2 and joblib 1.6.0 the same 56 50 6 0 the batched fits on scipy 1.15, per-gene solve and lstsq, scikit-learn 1.6's cond
all 421 385 36 0

What ran

Suite Commands What
Timed 34 the usage table's rows with current dependencies: the defaults, and vst() on one full experiment
Defaults at every thread count 97 all 26 inputs at 1, 2, 4, 8 and 16 threads, apart from the timed configurations; the 7 malformed and edge-case inputs among them
Option variants 196 at 1 and 8 threads on 11 inputs of up to 100 samples, and at 8 on the 250 + 250 cut: refit_cooks=False, fit_type="mean", poscounts size factors, low_memory=True, greaterAbs with lfc_null=0.5 and adapt=False, no independent or Cook's filtering, quiet=False, and vst() without and with the design; iterative size factors on 4 inputs of at most 5,000 genes; refit_cooks=False and vst() on two full experiments
Haswell kernel 38 the defaults on 11 inputs at 1 and 8 threads and on 4 more (two full experiments among them) at 8, and four option variants on 3 inputs at 8 threads
Lowest supported dependencies 56 4 timed rows; the defaults on 18 inputs, at 1 and 8 threads on 10, at 8 on the 250 + 250 cut and at 1 on the 7 malformed and edge-case inputs; and eight option variants on 3 inputs at 8 threads

"Identical" means every array, warning, printed line and exception is byte for byte the same. "Identical failure" means both builds raised the same exception with the same message: the three malformed count matrices (ValueError naming the NaN, the fraction or the negative value), the all-zero counts (scikit-learn's ValueError from LinearRegression, reached on both builds), and the two designs naming a level the data lacks. No configuration differs. On the ~sex + rin + tissue experiment (1,361 samples), main's own results depend on the thread count (between 4 and 8 threads its trend coefficients move by about 1e-8 relative and 0.18% of the genes move by more than 1e-2), and this branch reproduces main's results at each of 1, 2, 4, 8 and 16 threads.

Sanitizers. PyDESeq2 is pure Python, so no sanitizer build applies, and no race detector was run on the threaded passes. Those passes write disjoint slices of preallocated arrays; their outputs were compared with main's at 2, 4, 8 and 16 threads.

Threads. At n_cpus=1 neither build starts a thread pool or a worker process. In the 20 one-thread configurations with the default options that ran to completion, the process had a single thread on both builds and each used one CPU-second per wall-second (median 1.00 on both), and this branch used a median 3.6 times less CPU time than main (from 0.96 to 19 times), so the gain at one thread comes from doing less work. At n_cpus=8 both builds run 8 loky workers; at its peak main's process had 17 threads and this branch's 17 or 25, the extra ones being its thread pools, and the median CPU time per wall-second over the 20 eight-thread configurations was 6.1 on main and 6.3 on this branch.

Memory. Over the 38 timed configurations, this branch's peak memory is a median 0.99 of main's, from 0.84 to 1.33. It is higher by more than 5% in two: 250 vs 250 samples with 5,000 genes at 8 threads (+33%) and 1,000 vs 1,000 synthetic samples at 16 threads (+17%). On all 1,361 lung and thyroid samples it is 0.89 of main's at 1 thread and 1.04 at 8 threads.

Not covered. The comparisons here ran on Linux x86-64 with scipy's bundled OpenBLAS; the fork's GitHub Actions compared outputs on Windows and macOS runners (Behaviour changes), and other CPUs and BLAS builds (MKL, Accelerate) remain untested. Dependency versions between the two sets were tested only on the benchmark inputs (outputs under scipy 1.15.0, 1.15.3, 1.16.3 and 1.17.1 were bitwise identical to main's under the same version); numpy 2.3 and 2.4, joblib other than 1.6.0, and Python 3.13 and 3.14 were not compared (the unit tests ran on 3.12 and 3.14). Options not exercised: size_factors_fit_type="iterative" above 5,000 genes, designs with interaction terms, ref_level, a design given as a DataFrame, control_genes, min_replicates other than 7, joblib backends other than loky, joblib_verbosity > 0, a batch_size other than the default, and custom Inference subclasses beyond the suite's test of one. Inputs were human bulk RNA-seq from four GTEx tissues only, at most 1,540 samples and 63,856 genes, and synthetic counts up to 2,000 x 20,000. The merges with open pull requests were not run on this branch.

🤖 Generated with Claude Code

Jacob Schreiber and others added 2 commits November 6, 2026 18:48
Add a pytest suite for the numerical core: size factors and
normalization, the variance-stabilizing transform, dispersion
estimation, GLM fitting, Cook's distances and refitting, Wald tests,
p-value adjustment and apeGLM shrinkage, through DeseqDataSet,
DeseqStats, DefaultInference and the module-level functions behind
them.

The 300 tests sit in nine modules under tests/, with shared helpers in
tests/helpers.py (a negative binomial simulator that records its true
parameters, hand-worked inputs and reference implementations),
fixtures in tests/conftest.py, DESeq2 1.50.2 reference outputs and the
R scripts that wrote them under tests/data/unit_a/ and
tests/data/unit_b/, tests/README_NEW_TESTS.md on running and extending
the suite, and tests/TARGETS.md mapping every public class, method and
function to its tests. No source, build or workflow file changes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Fit genes in blocks rather than one at a time, so that
DeseqDataSet.deseq2(), DeseqStats.summary() and DeseqStats.lfc_shrink()
take less time and return the existing results:

- DefaultInference sends one task per block of batch_size genes to the
  loky workers instead of one per gene. fit_genewise_dispersions() fits
  mu_hat and the gene-wise dispersions in one map, deseq2() fits the MAP
  dispersions and LFCs in one map, and the Cook's refit fits its genes
  in one map.
- The gene-wise and MAP dispersion fits and the apeGLM shrinkage step
  scipy's L-BFGS-B routine (setulb) for every gene of a block in
  lockstep and evaluate the objective of all waiting genes at once. The
  driver adapts the loop of scipy's _minimize_lbfgsb and carries SciPy's
  BSD-3 and the wrapper's MIT notices. It runs when the installed
  setulb has the expected signature (scipy 1.15 to 1.18); otherwise
  each gene goes through scipy.optimize.minimize as before.
- IRLS, the linear-regression starting means and the Wald test run on
  stacked arrays; fit_lin_mu makes the one lstsq call LinearRegression
  makes, without the estimator.
- The size factors, the method-of-moments dispersions and the Cook's
  distances are computed on column chunks across up to n_cpus threads,
  and independent filtering sorts the p-values once.

Every fast path checks its inputs and falls back to the existing code,
which raises the existing errors. Public signatures, defaults and the
Inference interface are unchanged; an Inference subclass that overrides
a method gets the existing call sequence. tests/test_threaded_passes.py
checks every threaded pass against n_cpus=1.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant