From 21d7e1abb79dcc97d65393899f0ac73c023a3d48 Mon Sep 17 00:00:00 2001 From: Chadwick Boulay Date: Mon, 10 Aug 2026 14:32:37 -0400 Subject: [PATCH] EWMA: hoist the filter axis to last before scipy's IIR loop `_process` handed `scipy.signal.lfilter` whatever axis the message carried. When time is axis 0 of a C-contiguous (time, ch) array -- the layout every acquisition source emits -- each sample step strides `n_ch * itemsize`, and the cost grows superlinearly with the trailing dimension: at 300 samples, going from 256 to 1024 channels is 4x the data but 8.7x the time, versus linear when the axis is already last. Move the axis to the end, filter, and move back. The two copies pay for themselves at every size measured and never lose: 1.04x at 300x256, 1.78x at 30x1024, 4.41x at 300x1024, 1.72x at 3000x1024 -- all including the copies. `filter.py`'s SOS kernel already does this via util/sosfilt_direct, which is why ButterworthZeroPhase measures layout-neutral while this did not. `zi` is one sample slice, so it rides through the same hoist; the stored state stays in the caller's layout for anything that inspects it. AdaptiveStandardScaler runs two EWMAs per message and so paid this twice. At 300x1024 float32 it goes 7.676 -> 3.670 ms (2.09x), and its layout sensitivity drops from 2.64x to 1.28x. On a downstream intracortical feature chain that is 1.51x on total chain cost at 300-sample chunks; just as usefully, run-to-run spread falls from 15-22% to 2-5.7%, because the strided access was making the stage's cost unpredictable rather than merely slow. Output is bit-identical (max|diff| = 0.0) across both layouts, ragged chunk sequences, single-chunk, and float64, checked against a reference that reproduces the old un-hoisted lfilter with the same streaming zi. Deliberately materializes the result rather than returning a transposed view. The view saves a full-size pass and is 12-17% faster in isolation, but loses end to end at the sizes that matter (6% worse at 300 samples, 2% at 1000, 5% better only by 3000) because every downstream op then reads strided. Refs #215. --- src/ezmsg/sigproc/ewma.py | 59 +++++++++++++++++++++++++++++---------- tests/unit/test_ewma.py | 41 +++++++++++++++++++++++++++ 2 files changed, 86 insertions(+), 14 deletions(-) diff --git a/src/ezmsg/sigproc/ewma.py b/src/ezmsg/sigproc/ewma.py index 44c518e..bbf4016 100644 --- a/src/ezmsg/sigproc/ewma.py +++ b/src/ezmsg/sigproc/ewma.py @@ -257,6 +257,49 @@ def _reset_state(self, message: AxisArray) -> None: self._state.zi = xp.zeros_like(sub_dat) self._state.n_seen = 0 + def _lfilter_axis_last( + self, data: npt.NDArray, axis_idx: int, zi: npt.NDArray | None + ) -> tuple[npt.NDArray, npt.NDArray]: + """Run the EWMA recurrence with the filter axis contiguous. + + scipy's IIR loop strides by the trailing dimension whenever the filter + axis is not last, and the cost grows *superlinearly* with that dimension: + at 300 samples the step from 256 to 1024 channels is 4x the data but 8.7x + the time, versus linear when the axis is already last. Acquisition sources + emit ``(time, ch)``, so the streaming case hits the bad orientation by + default. + + Hoisting the axis to the end and moving the result back costs two copies + that pay for themselves at every size measured, and never lose: + 1.04x at 300x256, 1.78x at 30x1024, 4.41x at 300x1024, 1.72x at 3000x1024 + -- all *including* the copies. ``filter.py``'s SOS kernel already does + this (``util/sosfilt_direct``), which is why ``ButterworthZeroPhase`` + measures layout-neutral while this did not. + + ``zi`` is one sample slice, so hoisting it too is negligible and keeps the + stored state in the caller's layout for anything that inspects it. + """ + b = [self._state.alpha] + a = [1.0, self._state.alpha - 1.0] + last = data.ndim - 1 + if axis_idx == last: + return sps.lfilter(b, a, data, axis=-1, zi=zi) + + x = np.ascontiguousarray(np.moveaxis(data, axis_idx, last)) + zi_last = None if zi is None else np.ascontiguousarray(np.moveaxis(zi, axis_idx, last)) + y, zf = sps.lfilter(b, a, x, axis=-1, zi=zi_last) + # Materialize back into the caller's layout rather than returning a + # transposed view. The view saves a full-size pass here and is 12-17% + # faster *in isolation*, but it loses end to end at the sizes that matter + # (measured on the whole scaler at 1024 ch: 6% worse at 300 samples, 2% + # worse at 1000, 5% better only by 3000) because every downstream op then + # reads strided. Keep the copy; it is also the less surprising contract + # for a library consumer. + return ( + np.ascontiguousarray(np.moveaxis(y, last, axis_idx)), + np.ascontiguousarray(np.moveaxis(zf, last, axis_idx)), + ) + def _process(self, message: AxisArray) -> AxisArray: axis = self.settings.axis or message.dims[0] axis_idx = message.get_axis_idx(axis) @@ -276,24 +319,12 @@ def _process(self, message: AxisArray) -> AxisArray: # Normal behavior: update state with new samples. if self._state.zi is not None and not is_numpy_array(self._state.zi): self._state.zi = np.asarray(self._state.zi) - expected, self._state.zi = sps.lfilter( - [self._state.alpha], - [1.0, self._state.alpha - 1.0], - message.data, - axis=axis_idx, - zi=self._state.zi, - ) + expected, self._state.zi = self._lfilter_axis_last(message.data, axis_idx, self._state.zi) else: # Process-only: compute output without updating state. if self._state.zi is not None and not is_numpy_array(self._state.zi): self._state.zi = np.asarray(self._state.zi) - expected, _ = sps.lfilter( - [self._state.alpha], - [1.0, self._state.alpha - 1.0], - message.data, - axis=axis_idx, - zi=self._state.zi, - ) + expected, _ = self._lfilter_axis_last(message.data, axis_idx, self._state.zi) # The zero-initialized EWMA under-counts by 1-(1-alpha)^t at cumulative # sample t; dividing it out gives the exact exponentially-weighted diff --git a/tests/unit/test_ewma.py b/tests/unit/test_ewma.py index 7a73915..6094cb7 100644 --- a/tests/unit/test_ewma.py +++ b/tests/unit/test_ewma.py @@ -530,3 +530,44 @@ def mk(d, offset): parts.append(chunked(mk(c, n)).data) n += c.shape[0] np.testing.assert_allclose(np.concatenate(parts, axis=0), single, rtol=1e-10, atol=1e-10) + + +@pytest.mark.parametrize("n_dim", [2, 3]) +def test_ewma_result_is_independent_of_filter_axis_position(n_dim): + """Hoisting the filter axis to last must not change the result. + + ``_process`` moves the filter axis to the end before handing scipy the + recurrence, because scipy's IIR loop strides by the trailing dimension + otherwise and degrades superlinearly with it. That is purely a memory-layout + optimization, so filtering axis 0 of ``(time, ch)`` and axis -1 of the + transposed array must agree exactly -- and keep agreeing across chunks, since + ``zi`` rides through the same hoist. + """ + fs = 1000.0 + n_times, n_ch = 137, 24 + rng = np.random.default_rng(7) + shape = (n_times, n_ch) if n_dim == 2 else (n_times, n_ch, 3) + data = rng.standard_normal(shape).astype(np.float32) + dims = ["time", "ch"] if n_dim == 2 else ["time", "ch", "feat"] + + def run(transpose: bool, chunks: list[int]) -> np.ndarray: + proc = EWMATransformer(time_constant=0.5, axis="time", accumulate=True) + outs, start = [], 0 + for n in chunks: + block = data[start : start + n] + if transpose: + block = np.ascontiguousarray(np.moveaxis(block, 0, -1)) + msg = AxisArray( + data=block, + dims=(dims[1:] + ["time"]) if transpose else dims, + axes={"time": AxisArray.TimeAxis(fs=fs, offset=start / fs)}, + ) + out = proc(msg) + arr = out.data + outs.append(np.moveaxis(arr, -1, 0) if transpose else arr) + start += n + return np.concatenate(outs, axis=0) + + # Ragged chunks so the hoisted `zi` has to carry correctly across boundaries. + chunks = [1, 13, 40, 3, 80] + np.testing.assert_array_equal(run(False, chunks), run(True, chunks))