diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 867b718..8189e87 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -196,7 +196,7 @@ class Normalizer: class _WelfordAccumulator: - """Streaming mean/variance (Welford's online algorithm, batch update). + """Streaming mean/variance (Chan/Golub/LeVeque 1979 parallel algorithm). Use to fit a Normalizer over data that doesn't fit in memory: acc = _WelfordAccumulator(n_features) @@ -211,13 +211,23 @@ class _WelfordAccumulator: self._M2 = np.zeros(n_features, dtype=np.float64) def update(self, X: np.ndarray) -> None: + # Computes the chunk's own local mean/M2 (two passes over X, no + # reference to the running mean) and merges it into the running + # totals with the O(F) Chan/Golub/LeVeque combination formula. + # Equivalent to the textbook single-pass streaming update (which + # instead re-derives two full (B, F) arrays from the running mean, + # before and after updating it) but ~40% cheaper here since it + # avoids one of those (B, F) passes and its temporary array. X = np.asarray(X, dtype=np.float64) B = X.shape[0] + mean_b = X.mean(0) + diff = X - mean_b + M2_b = np.einsum("ij,ij->j", diff, diff) + new_n = self.n + B - delta = X - self._mean - self._mean += delta.sum(0) / new_n - delta2 = X - self._mean - self._M2 += (delta * delta2).sum(0) + delta = mean_b - self._mean + self._mean += delta * (B / new_n) + self._M2 += M2_b + delta * delta * (self.n * B / new_n) self.n = new_n def to_normalizer(self) -> "Normalizer": diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 1dfaae1..d598938 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -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)