sparse_mean_variance_axis now uses all cores

[Intron7]

This functions now uses all cores for mean & var calculations.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 76.27%. Comparing base (ee8505b) to head (d982936).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3015      +/-   ##
==========================================
+ Coverage   75.53%   76.27%   +0.74%     
==========================================
  Files         117      117              
  Lines       12950    12795     -155     
==========================================
- Hits         9782     9760      -22     
+ Misses       3168     3035     -133     

Files	Coverage Δ
scanpy/preprocessing/_utils.py	97.36% <100.00%> (+42.70%)	⬆️

... and 4 files with indirect coverage changes

[flying-sheep]

Really cool!

Please use longer variable names. nthr→num_threads, s→sums_minor or so, s0→???, m0→???

Of course not for things like loop counters, but you know.

Also would be great to have comments, like “calculate sums and sum of squares along the minor axis” (if “minor” is correct here) then “go over minor axis again to calculate means and variances from the sums”.

PS: Benchmarks should be ready in a little bit, so you could add one to check how much this speeds things up! I’ll keep you posted

[Intron7]

Some small benchmarks for 32 cores with CSR.shape=(196943, 20867):

axis	old	new
minor	804 ms	96 ms
major	520 ms	40 ms

[scverse-benchmark]

Benchmark changes

Change	Before [ee8505b]	After [d982936]	Ratio	Benchmark (Parameter)
-	508±2ms	31.9±1ms	0.06	preprocessing.SparseDenseSuite.time_mean_var('lung93k')
+	1.09±0.04ms	1.22±0.04ms	1.12	preprocessing.SparseDenseSuite.time_mean_var('pbmc68k_reduced')
+	241M	330M	1.37	preprocessing.peakmem_pca
+	5.86±0.01ms	6.86±0.03ms	1.17	preprocessing.time_calculate_qc_metrics

Comparison: https://github.com/scverse/scanpy/compare/ee8505b1c1578af0c50defdb3cf64ec18713669e..d9829365d7e23aea5680990ea8570d0a384291d3
Last changed: Tue, 23 Apr 2024 09:27:27 +0000

More details: https://github.com/scverse/scanpy/pull/3015/checks?check_run_id=24144155267

[flying-sheep]

I tentatively added a benchmark that runs just on _get_mean_var.

Locally I don’t see any difference though, what’s wrong? Too small data? Numba not set up with correct number of threads?

/edit: also I think the machine is not sufficiently tuned. The original run (before I added the mean_var benchmarks) said “No changes in benchmarks.”

[flying-sheep]

Looks good! Would be nice to add a sufficiently large dataset to the benchmarks to demonstrate the benefits.