fix: improve variance precision

[flying-sheep]

Fixes scverse/scanpy#3740

Turns out precision doesn’t matter if you only use numbers that are perfectly representable as floats.

Long story short: doing the power operation in float32 led to imprecise exponentiated values. Dask’s scheduling led to errors accumulating in summing them up. Storing exponentiation results in float64 fixes this, but of course slows things down.

Details

Something’s really off here, sparse-in-dask usually succeeds, but occasionally can fail enormously:

FAILED tests/test_stats.py::test_mean_var_pbmc_dask[dask.array.Array[scipy.sparse.csr_matrix]-2d-ax0] - AssertionError: 
Arrays are not almost equal to 6 decimals

Mismatched elements: 761 / 765 (99.5%)
Max absolute difference among violations: 7.88906984
Max relative difference among violations: 10.68116412
 ACTUAL: array([4.081325e-01, 8.858602e-02, 4.668730e-01, 5.370092e-01,
       9.700749e-02, 5.163984e-01, 3.209537e-01, 4.335481e-01,
       4.517616e-01, 4.680808e-01, 4.769589e-01, 3.681803e-01,...
 DESIRED: array([0.42406 , 0.090823, 0.646268, 0.790074, 0.097231, 0.575962,
       0.362582, 0.469195, 0.462772, 0.486282, 0.502374, 0.383135,
       0.979183, 0.451647, 0.542628, 0.059287, 0.815319, 1.323111,...

going into mean_var and re-executing the same expression shows how completely nondeterministic it is:

>>> np.testing.assert_almost_equal(var.compute(), mean_var_(x.compute(), axis=axis, correction=correction)[1])
[…]
AssertionError: 
Arrays are not almost equal to 7 decimals

Mismatched elements: 381 / 765 (49.8%)
Max absolute difference among violations: 7.88169589
Max relative difference among violations: 11.27194009
 ACTUAL: array([4.2406003e-01, 9.0823245e-02, 6.4626844e-01, 7.9007370e-01,
       9.7230909e-02, 5.7596200e-01, 3.6258189e-01, 4.6919450e-01,
       4.6277188e-01, 4.8628229e-01, 5.0237379e-01, 3.8313501e-01,...
 DESIRED: array([0.42406  , 0.0908232, 0.6462684, 0.7900737, 0.0972309, 0.575962 ,
       0.3625819, 0.4691945, 0.4627719, 0.4862823, 0.5023738, 0.383135 ,
       0.979183 , 0.4516468, 0.5426282, 0.0592871, 0.8153192, 1.3231111,...

>>> np.testing.assert_almost_equal(var.compute(), mean_var_(x.compute(), axis=axis, correction=correction)[1])

>>> np.testing.assert_almost_equal(var.compute(), mean_var_(x.compute(), axis=axis, correction=correction)[1])
[…]
AssertionError: 
Arrays are not almost equal to 7 decimals

Mismatched elements: 364 / 765 (47.6%)
Max absolute difference among violations: 1.37848993
Max relative difference among violations: 1.88310113
 ACTUAL: array([0.42406  , 0.0908232, 0.6462684, 0.7900737, 0.0972309, 0.575962 ,
       0.3625819, 0.4691945, 0.4627719, 0.4862823, 0.5023738, 0.383135 ,
       0.979183 , 0.4516468, 0.5426282, 0.0592871, 0.8153192, 1.3231111,...
 DESIRED: array([0.42406  , 0.0908232, 0.6462684, 0.7900737, 0.0972309, 0.575962 ,
       0.3625819, 0.4691945, 0.4627719, 0.4862823, 0.5023738, 0.383135 ,
       0.979183 , 0.4516468, 0.5426282, 0.0592871, 0.8153192, 1.3231111,...

the nondeterminism seems to come from this line:

fast-array-utils/src/fast_array_utils/stats/_mean_var.py

Line 41 in 92483a4

var = mean_sq - mean_**2

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 98.33%. Comparing base (9df6b5a) to head (0cd5237).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #127      +/-   ##
==========================================
+ Coverage   98.32%   98.33%   +0.01%     
==========================================
  Files          17       17              
  Lines         418      421       +3     
==========================================
+ Hits          411      414       +3     
  Misses          7        7              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

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

[codspeed-hq]

CodSpeed Performance Report

Merging #127 will improve performances by 62.9%

_{Comparing pa/test-accuracy (0cd5237) with main (9df6b5a)}

Summary

⚡ 12 improvements
✅ 148 untouched

Benchmarks breakdown

	Benchmark	BASE	HEAD	Change
⚡	test_stats_benchmark[scipy.sparse.csc_array-1d-all-float32-mean_var]	15 ms	9.3 ms	+60.88%
⚡	test_stats_benchmark[scipy.sparse.csc_array-1d-all-float64-mean_var]	21.4 ms	17.4 ms	+23.37%
⚡	test_stats_benchmark[scipy.sparse.csc_array-1d-all-int32-mean_var]	14.7 ms	9 ms	+62.62%
⚡	test_stats_benchmark[scipy.sparse.csc_array-2d-all-float32-mean_var]	15 ms	9.3 ms	+60.91%
⚡	test_stats_benchmark[scipy.sparse.csc_array-2d-all-float64-mean_var]	21.4 ms	17.4 ms	+23.38%
⚡	test_stats_benchmark[scipy.sparse.csc_array-2d-all-int32-mean_var]	14.7 ms	9 ms	+62.64%
⚡	test_stats_benchmark[scipy.sparse.csr_array-1d-all-float32-mean_var]	14.9 ms	9.3 ms	+60.51%
⚡	test_stats_benchmark[scipy.sparse.csr_array-1d-all-float64-mean_var]	21.4 ms	17.3 ms	+23.62%
⚡	test_stats_benchmark[scipy.sparse.csr_array-1d-all-int32-mean_var]	14.7 ms	9 ms	+62.9%
⚡	test_stats_benchmark[scipy.sparse.csr_array-2d-all-float32-mean_var]	15 ms	9.3 ms	+60.88%
⚡	test_stats_benchmark[scipy.sparse.csr_array-2d-all-float64-mean_var]	21.4 ms	17.4 ms	+23.38%
⚡	test_stats_benchmark[scipy.sparse.csr_array-2d-all-int32-mean_var]	14.7 ms	9 ms	+62.7%