Improve computation time and memory usage across all statistics

[CyrilJl]

Summary

Performance review of all batch statistics, targeting computation time and memory usage. Every change is validated by the existing test suite (59 passed) plus end-to-end correctness checks against the NumPy reference functions (np.corrcoef, np.average, np.var with NaNs, etc.). Along the way the review surfaced and fixed three correctness bugs.

Benchmarks

Streaming 200,000 × 100 float64 in 20,000-row chunks (best of 5, NumPy 2.5.1):

Workload	Before	After	Gain
BatchCorr (single matrix)	364 ms	169 ms	2.2×
BatchNanPeakToPeak	175 ms	66 ms	2.7×
BatchWeightedMean	106 ms	59 ms	1.8×
BatchSum (integer data)	25 ms	13 ms	1.9×
BatchNanMax	86 ms	52 ms	1.7×
BatchCov (single matrix)	214 ms	161 ms	1.3×
BatchVar	39 ms	31 ms	1.25×
BatchCov peak memory	48 MB	32 MB	−33%

Optimizations

• BatchCorr (single-matrix mode): the variances are mathematically the diagonal of the covariance, so the two internal BatchVar streams were redundant work on every update. They are now skipped entirely and __call__ uses sqrt(diag(cov)). Also stops fancy-index-copying the batch when it contains no NaNs.
• BatchCov: when batch2 is batch1 (the common single-argument path), the NaN scan and the mean-centering each ran twice on the same array — now once. The / n normalization moved off a batch-sized temporary onto the small (d1, d2) result.
• BatchVar: einsum("j,ij->j", p+u, v) factors into (p+u) * v.sum(0), and that column sum is already computed for the mean update in the same call — one full O(n·d) pass removed per update.
• BatchNanMax / BatchNanMin: np.fmax.reduce / np.fmin.reduce replace np.nanmax / np.nanmin, which internally copy the whole batch and allocate a NaN mask. Also removes the all-NaN warning-suppression block, since fmax/fmin return NaN silently.
• BatchNanPeakToPeak: the max and min trackers each ran two isfinite passes per update (4 total); they now share a single pass.
• BatchWeightedMean: the sum of weights was computed by materializing broadcast_to(weights, batch.shape) * 1 — a full batch-sized allocation. It now sums the zero-stride broadcast view directly.
• BatchStat: the NaN scan is skipped for non-inexact dtypes (integers/booleans can't hold NaN), which also avoids silent integer overflow in np.add.reduce.

Bug fixes found during the review

• NaN filtering with axis != 0 was broken: the mask is computed along the reduction axis but was applied to axis 0 — IndexError on non-square batches, silently wrong results (NaN output, inflated n_samples) on square ones. Fixed with np.compress along the actual reduction axis; negative axes are normalized. Multi-axis reductions over data containing NaNs now raise a clear NotImplementedError pointing to the BatchNan* classes instead of crashing or returning garbage.
• All-NaN slices poisoned BatchNanMax/BatchNanMin: an all-NaN column in a later batch propagated NaN through np.maximum/np.minimum and permanently corrupted the running extremum, e.g. [[1, 2]] then [[nan, 3]] returned [nan, 3] instead of [1, 3]. np.fmax/np.fmin ignore NaN, fixing this by construction.
• BatchNanMax/BatchNanMin crashed on integer inputs at call time (cannot convert float NaN to integer): the empty-slice masking now only runs when empty slices exist, casting to float only when NaN placeholders are actually needed.

Also documents the linear memory growth of BatchWeightedSum's list mode (when axis 0 is not reduced) in its docstring.

Not changed

The centered covariance update was deliberately kept over the allocation-free raw-Gram formulation (X₁ᵀX₂ − n·m₁m₂), which would reintroduce catastrophic cancellation — numerical stability wins over the last temporary.

🤖 Generated with Claude Code