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
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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user