diff --git a/src/ezmsg/sigproc/ewma.py b/src/ezmsg/sigproc/ewma.py index be046c5..44c518e 100644 --- a/src/ezmsg/sigproc/ewma.py +++ b/src/ezmsg/sigproc/ewma.py @@ -172,7 +172,29 @@ class EWMASettings(ez.Settings): """If True, return the input unchanged (identity) without touching the EWMA. Unlike a very large time_constant -- which still applies a (stale) baseline estimate -- passthrough leaves the data untouched. May be toggled - at runtime without resetting the filter state.""" + at runtime without resetting the filter state; see ``reset_on_resume`` for + what that means for the first message after the gap.""" + + reset_on_resume: bool = False + """Whether switching ``passthrough`` back off discards the filter state. + + The filter sees none of the samples that go by during passthrough, so ``zi`` + describes an exponentially-weighted window that ended when passthrough was + switched on -- and the state cannot tell a 10 ms blip from a 10 minute + outage, since both resume identically. For a scaler that means z-scoring + post-gap data against pre-gap statistics. + + False (the default) resumes from the preserved state, which is right for a + short blip and keeps an estimate that may have taken many ``time_constant``\\ s + to converge. True rebuilds from the first post-gap message instead: with the + bias correction below, that first output is exactly the first sample, and the + estimate re-converges over ``time_constant``. Prefer True where passthrough + may be left on long enough for the signal to drift, which is the case + :obj:`ezmsg.sigproc.binned_aggregate.BinnedAggregateTransformer` always + assumes. + + Empty chunks are not gaps -- they carry no samples past the filter -- so they + never trigger this.""" mlx_metal_chunk_sizes: tuple[int, ...] = (32, 1024) """Allowable compile-time chunk sizes for EWMA Metal kernels. The smallest @@ -191,17 +213,31 @@ class EWMAState: class EWMATransformer(BaseStatefulTransformer[EWMASettings, AxisArray, AxisArray, EWMAState]): # `accumulate` is read live in `_process` to gate state updates and - # `passthrough` is read live in `__call__`/`__acall__`; other fields are - # cached into state (alpha, zi) during `_reset_state`. - NONRESET_SETTINGS_FIELDS = frozenset({"accumulate", "passthrough", "mlx_metal_chunk_sizes"}) + # `passthrough`/`reset_on_resume` are read live in `__call__`/`__acall__`; + # other fields are cached into state (alpha, zi) during `_reset_state`. + NONRESET_SETTINGS_FIELDS = frozenset({"accumulate", "passthrough", "reset_on_resume", "mlx_metal_chunk_sizes"}) + + def _skip(self, message: AxisArray) -> bool: + """Whether to bypass the filter, flagging a reset if this is a real gap. + + The two bypass conditions have to be kept apart: passthrough lets samples + past unfiltered and so leaves a hole in ``zi``'s history, while an empty + chunk carries nothing past and leaves the history intact. Only the former + is a gap, so only the former can invalidate the state. + """ + if self.settings.passthrough: + if self.settings.reset_on_resume and np.prod(message.data.shape) != 0: + self._request_reset() + return True + return bool(np.prod(message.data.shape) == 0) def __call__(self, message: AxisArray) -> AxisArray: - if self.settings.passthrough or np.prod(message.data.shape) == 0: + if self._skip(message): return message return super().__call__(message) async def __acall__(self, message: AxisArray) -> AxisArray: - if self.settings.passthrough or np.prod(message.data.shape) == 0: + if self._skip(message): return message return await super().__acall__(message) diff --git a/src/ezmsg/sigproc/scaler.py b/src/ezmsg/sigproc/scaler.py index 0279261..6f66e75 100644 --- a/src/ezmsg/sigproc/scaler.py +++ b/src/ezmsg/sigproc/scaler.py @@ -104,18 +104,30 @@ class AdaptiveStandardScalerTransformer( ] ): # `accumulate` can be live-propagated into the child EWMAs (see - # `update_settings` below) and `passthrough` is read live in - # `__call__`/`__acall__`; `time_constant` and `axis` are baked into + # `update_settings` below) and `passthrough`/`reset_on_resume` are read live + # in `__call__`/`__acall__`; `time_constant` and `axis` are baked into # the children during `_reset_state`. - NONRESET_SETTINGS_FIELDS = frozenset({"accumulate", "passthrough"}) + NONRESET_SETTINGS_FIELDS = frozenset({"accumulate", "passthrough", "reset_on_resume"}) + + def _skip(self, message: AxisArray) -> bool: + """Whether to bypass scaling, flagging a reset if this is a real gap. + + See :obj:`EWMASettings.reset_on_resume`. A reset here is enough for both + children: `_reset_state` rebuilds them from scratch. + """ + if not self.settings.passthrough: + return False + if self.settings.reset_on_resume and np.prod(message.data.shape) != 0: + self._request_reset() + return True def __call__(self, message: AxisArray) -> AxisArray: - if self.settings.passthrough: + if self._skip(message): return message return super().__call__(message) async def __acall__(self, message: AxisArray) -> AxisArray: - if self.settings.passthrough: + if self._skip(message): return message return await super().__acall__(message) diff --git a/tests/unit/test_ewma.py b/tests/unit/test_ewma.py index ea077fe..7a73915 100644 --- a/tests/unit/test_ewma.py +++ b/tests/unit/test_ewma.py @@ -349,6 +349,71 @@ def test_ewma_passthrough_toggle_preserves_state(): assert np.allclose(out_toggled.data, out_ref.data) +def test_ewma_reset_on_resume_discards_state(): + """reset_on_resume=True rebuilds from the first message after the gap.""" + proc_toggled = EWMATransformer(settings=EWMASettings(time_constant=0.1, reset_on_resume=True)) + proc_fresh = EWMATransformer(settings=EWMASettings(time_constant=0.1, reset_on_resume=True)) + + msg1 = _make_ewma_test_msg(np.ones((10, 2))) + _ = proc_toggled(msg1) + zi_before = proc_toggled._state.zi.copy() + + # During the gap the state is left alone -- only the hash is invalidated, so + # the rebuild happens on the next real message rather than eagerly here. + proc_toggled.settings = dc_replace(proc_toggled.settings, passthrough=True) + msg2 = _make_ewma_test_msg(np.ones((10, 2)) * 100.0) + assert proc_toggled(msg2) is msg2 + assert np.allclose(proc_toggled._state.zi, zi_before) + assert proc_toggled._hash == -1 + + # Resume: output matches a transformer that has never seen anything, not one + # that carries msg1's estimate. + proc_toggled.settings = dc_replace(proc_toggled.settings, passthrough=False) + msg3 = _make_ewma_test_msg(np.ones((10, 2)) * 2.0) + assert np.allclose(proc_toggled(msg3).data, proc_fresh(msg3).data) + + +def test_ewma_reset_on_resume_via_update_settings(): + """The reset is driven from __call__, so update_settings works too. + + `passthrough` and `reset_on_resume` are both in NONRESET_SETTINGS_FIELDS -- + toggling either queues no reset by itself, which is correct: a toggle with no + messages in between leaves no gap. + """ + proc = EWMATransformer(settings=EWMASettings(time_constant=0.1, reset_on_resume=True)) + + _ = proc(_make_ewma_test_msg(np.ones((10, 2)))) + hash_before = proc._hash + + proc.update_settings(EWMASettings(time_constant=0.1, reset_on_resume=True, passthrough=True)) + assert proc._hash == hash_before + + msg = _make_ewma_test_msg(np.ones((10, 2)) * 100.0) + assert proc(msg) is msg + assert proc._hash == -1 + + +def test_ewma_reset_on_resume_ignores_empty_messages(): + """An empty chunk is not a gap -- it passes no samples, so it cannot invalidate zi. + + This is the one case the passthrough and empty short-circuits must not share. + """ + proc = EWMATransformer(settings=EWMASettings(time_constant=0.1, reset_on_resume=True)) + + _ = proc(make_msg()) + hash_before = proc._hash + zi_before = proc._state.zi.copy() + + check_empty_result(proc(make_empty_msg())) + assert proc._hash == hash_before + assert np.allclose(proc._state.zi, zi_before) + + # ... and an empty chunk *during* passthrough is not a gap either. + proc.settings = dc_replace(proc.settings, passthrough=True) + assert proc(make_empty_msg()) is not None + assert proc._hash == hash_before + + class TestEWMAUpdateSettings: """Live settings updates via BaseProcessor.update_settings.""" diff --git a/tests/unit/test_scaler.py b/tests/unit/test_scaler.py index 8c6bc24..03bb33c 100644 --- a/tests/unit/test_scaler.py +++ b/tests/unit/test_scaler.py @@ -376,6 +376,36 @@ def test_passthrough_toggle_preserves_state(self): assert out.data.shape == (50, 4) assert not np.any(np.isnan(out.data)) + def test_reset_on_resume_discards_state(self): + """reset_on_resume=True rebuilds both child EWMAs after the gap.""" + scaler = AdaptiveStandardScalerTransformer( + settings=AdaptiveStandardScalerSettings(time_constant=0.1, reset_on_resume=True) + ) + + np.random.seed(42) + _ = scaler(_make_scaler_test_msg(np.random.randn(100, 4))) + samps_before = scaler._state.samps_ewma + zi_before = samps_before._state.zi.copy() + + # The gap itself leaves the children untouched; only the hash is dropped. + scaler.update_settings( + AdaptiveStandardScalerSettings(time_constant=0.1, reset_on_resume=True, passthrough=True) + ) + msg = _make_scaler_test_msg(np.random.randn(50, 4) + 1000.0) + assert scaler(msg) is msg + assert scaler._state.samps_ewma is samps_before + assert np.allclose(scaler._state.samps_ewma._state.zi, zi_before) + assert scaler._hash == -1 + + # Resume: _reset_state builds fresh children, so the pre-gap statistics + # are gone rather than being applied to post-gap data. + scaler.update_settings( + AdaptiveStandardScalerSettings(time_constant=0.1, reset_on_resume=True, passthrough=False) + ) + out = scaler(_make_scaler_test_msg(np.random.randn(50, 4))) + assert scaler._state.samps_ewma is not samps_before + assert not np.any(np.isnan(out.data)) + class TestAdaptiveStandardScalerUpdateSettings: """Live settings updates via update_settings.""" @@ -539,7 +569,9 @@ def test_scaler_bias_correction_survives_chunking(): data = rng.normal(0.0, 1.0, size=(151, 2)) def mk(d, offset): - return AxisArray(d, dims=["time", "ch"], axes=frozendict({"time": AxisArray.TimeAxis(fs=fs, offset=offset / fs)})) + return AxisArray( + d, dims=["time", "ch"], axes=frozendict({"time": AxisArray.TimeAxis(fs=fs, offset=offset / fs)}) + ) single = AdaptiveStandardScalerTransformer(time_constant=tau, axis="time")(mk(data, 0)).data chunked = AdaptiveStandardScalerTransformer(time_constant=tau, axis="time")