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))