Speed up _WelfordAccumulator's per-chunk update
CI / Format (ruff format) (push) Successful in 27s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 35s
CI / Type check (ty) (push) Successful in 38s
CI / Format (ruff format) (pull_request) Successful in 33s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 33s
CI / Tests (pull_request) Successful in 1m21s
CI / Tests (push) Successful in 1m27s

The streaming update re-derived two full (B, F) arrays from the
running mean (once before updating it, once after) plus an elementwise
product — five passes over each chunk and three temporary arrays, to
maintain a mean/variance that's tiny in width (COND_DIM=15 at most).

Reformulate as Chan/Golub/LeVeque's parallel-variance algorithm:
compute the chunk's own local mean/M2 (independent of the running
state) and merge it in with an O(F) combination formula. Same
streaming interface and output (identical to ~1e-14, float64 rounding
noise), ~40% faster per update() call on a benchmark chunk.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-07-29 13:55:07 +02:00
co-authored by Claude Sonnet 5
parent 47a6c9db1f
commit 759b67a9e1
2 changed files with 84 additions and 5 deletions
+69
View File
@@ -15,6 +15,7 @@ from giant.data.transforms import (
sorted_membership,
travel_direction,
_vectorized_map_lookup,
_WelfordAccumulator,
)
@@ -491,3 +492,71 @@ def test_vectorized_map_lookup_raises_keyerror_on_missing_value():
values = np.array([1, 2, 3])
with pytest.raises(KeyError):
_vectorized_map_lookup(values, mapping)
# ── _WelfordAccumulator ──────────────────────────────────────────────────────
def test_welford_accumulator_matches_direct_mean_std_over_many_chunks():
rng = np.random.default_rng(5)
F = 4
chunks = [rng.standard_normal((rng.integers(1, 50), F)) * 10 + 3 for _ in range(20)]
full = np.concatenate(chunks, axis=0)
acc = _WelfordAccumulator(F)
for chunk in chunks:
acc.update(chunk)
norm = acc.to_normalizer()
assert norm.mean is not None and norm.std is not None
np.testing.assert_allclose(norm.mean, full.mean(axis=0), rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(norm.std, full.std(axis=0), rtol=1e-5, atol=1e-5)
assert acc.n == full.shape[0]
def test_welford_accumulator_single_chunk():
rng = np.random.default_rng(6)
X = rng.standard_normal((100, 3)) * 5 - 2
acc = _WelfordAccumulator(3)
acc.update(X)
norm = acc.to_normalizer()
assert norm.mean is not None and norm.std is not None
np.testing.assert_allclose(norm.mean, X.mean(axis=0), rtol=1e-5)
np.testing.assert_allclose(norm.std, X.std(axis=0), rtol=1e-5)
def test_welford_accumulator_matches_naive_running_mean_reference():
"""The chunk-local-mean + Chan-merge formula must agree with the naive
textbook streaming update (subtract the *running* mean before and after
updating it) that it replaces, within float64 rounding tolerance."""
rng = np.random.default_rng(7)
F = 3
chunks = [rng.standard_normal((rng.integers(1, 40), F)) for _ in range(15)]
def naive_update(mean, M2, n, X):
X = np.asarray(X, dtype=np.float64)
B = X.shape[0]
new_n = n + B
delta = X - mean
mean = mean + delta.sum(0) / new_n
delta2 = X - mean
M2 = M2 + (delta * delta2).sum(0)
return mean, M2, new_n
naive_mean = np.zeros(F)
naive_M2 = np.zeros(F)
naive_n = 0
for chunk in chunks:
naive_mean, naive_M2, naive_n = naive_update(
naive_mean, naive_M2, naive_n, chunk
)
acc = _WelfordAccumulator(F)
for chunk in chunks:
acc.update(chunk)
assert acc.n == naive_n
np.testing.assert_allclose(acc._mean, naive_mean, rtol=1e-9, atol=1e-9)
np.testing.assert_allclose(acc._M2, naive_M2, rtol=1e-9, atol=1e-9)