Speed up iterative size factor fitting by not slicing AnnData in the objective

[Sizerta]

Reference issues or PRs

Relates to #453.

What does your PR implement? Be specific.

Diagnosis of #453

When every gene contains at least one zero, fit_size_factors uses the iterative method. If size factors have not been computed, fit_genewise_dispersions() calls fit_size_factors, making size factor estimation the main source of runtime in this case. Adding 1 to all counts removes the zeros, allowing the median-of-ratios method to be used instead, which explains the workaround described in the issue.

In a simulated dataset with 12 samples, 300 genes, and 60% zeros, size factor estimation took 5.1 seconds, whereas genewise dispersion fitting took less than 0.5 seconds.

Profiling _fit_iterate_size_factors showed that a substantial portion of the runtime comes from repeated operations within the objective function passed to Powell's optimizer. At each evaluation, the function slices the AnnData object twice (self[:, self.non_zero_genes]) and indexes self.var by gene names. These inputs remain unchanged throughout the optimization, so the repeated operations are unnecessary. Over a full fit, Powell evaluates the objective several hundred to a few thousand times.

Implementation

The counts, unscaled means, and dispersions for non-zero genes are now extracted once per iteration and passed to the objective function using minimize(..., args=...). The objective's arithmetic and operation order remain unchanged, preserving bit-identical results.

Verification

• Size factors are bit-identical to those from main (np.array_equal) on simulated datasets with 50–90% zeros, including a design with a continuous covariate.
• Runtime comparison for size factor estimation on simulated negative binomial counts with additional zeros:

Dataset	Zeros	main	This PR	Speedup
12 samples × 300 genes	60%	5.1 s	3.1 s	1.7×
12 samples × 300 genes	90%	9.8 s	3.0 s	3.2×
16 samples × 500 genes, ~condition + age	50%	5.4 s	3.2 s	1.7×
60 samples × 2,000 genes	50%	34.9 s	26.8 s	1.3×

The test suite passes, including test_iterative_size_factors, which compares this implementation with the output of DESeq2's estimateSizeFactorsIterate. No additional test was added because the numerical results are unchanged.

Potential follow-ups (not included)

1. For larger datasets, a substantial part of the remaining runtime is spent recomputing gammaln terms in nb_nll that do not depend on the size factors. Precomputing these terms once per iteration reduces runtime for the 60-sample, 2,000-gene dataset to 15.4 seconds (a 2.3× speedup relative to main). However, this changes the numerical results slightly, with relative differences in size factors of up to 6 × 10⁻⁸, and would partially duplicate the likelihood calculation in nb_nll. This optimization could be considered separately.

2. The warning emitted when switching to the iterative method could mention size_factors_fit_type="poscounts" as an alternative for datasets in which many genes contain zeros.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.00%. Comparing base (2ca3b74) to head (41d4758).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #479      +/-   ##
==========================================
- Coverage   86.05%   86.00%   -0.05%     
==========================================
  Files          15       15              
  Lines        1276     1279       +3     
==========================================
+ Hits         1098     1100       +2     
- Misses        178      179       +1     

Files with missing lines	Coverage Δ
src/pydeseq2/dds.py	87.50% <100.00%> (-0.13%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.