ezmsg-sigproc 1.8.2__tar.gz → 2.1.0__tar.gz

{ezmsg_sigproc-1.8.2 → ezmsg_sigproc-2.1.0}/PKG-INFO

@@ -1,11 +1,12 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 1.8.2
+Version: 2.1.0
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>
 License-Expression: MIT
 License-File: LICENSE.txt
 Requires-Python: >=3.10.15
+Requires-Dist: array-api-compat>=1.11.1
 Requires-Dist: ezmsg>=3.6.0
 Requires-Dist: numba>=0.61.0
 Requires-Dist: numpy>=1.26.0

ezmsg_sigproc-2.1.0/docs/ProcessorsBase.md

@@ -0,0 +1,173 @@
+## Processor Base Classes
+
+The `ezmsg.sigproc.base` module contains the base classes for the signal processors. The base classes are designed to allow users to create custom signal processors with minimal errors and minimal repetition of boilerplate code.
+
+> The information below was written at the time of a major refactor to `ezmsg.sigproc.base` to help collate the design decisions and to help with future refactoring.  However, it may be outdated or incomplete. Please refer to the source code for the most accurate information.
+
+> Updates: Added CompositeProducer, BaseProcessorUnit.
+
+
+### Generic TypeVars
+
+| Idx | Class                 | Description                                                                |
+|-----|-----------------------|----------------------------------------------------------------------------|
+| 1   | `MessageInType` (Mi)  | for messages passed to a consumer, processor, or transformer               |
+| 2   | `MessageOutType` (Mo) | for messages returned by a producer, processor, or transformer             |
+| 3   | `SettingsType`        | bound to ez.Settings                                                       |
+| 4   | `StateType` (St)      | bound to ProcessorState which is simply ez.State with a `hash: int` field. |
+
+
+### Protocols
+
+| Idx | Class                 | Parent | State | `__call__` sig           | @state | partial_fit |
+|-----|-----------------------|--------|-------|--------------------------|--------|-------------|
+| 1   | `Processor`           | -      | No    | Any -> Any               | -      | -           |
+| 2   | `Producer`            | -      | No    | None -> Mo               | -      | -           |
+| 3   | `Consumer`            | 1      | No    | Mi -> None               | -      | -           |
+| 4   | `Transformer`         | 1      | No    | Mi -> Mo                 | -      | -           |
+| 5   | `StatefulProcessor`   | -      | Yes   | Any -> Any               | Y      | -           |
+| 6   | `StatefulProducer`    | -      | Yes   | None -> Mo               | Y      | -           |
+| 7   | `StatefulConsumer`    | 5      | Yes   | Mi -> None               | Y      | -           |
+| 8   | `StatefulTransformer` | 5      | Yes   | Mi -> Mo                 | Y      | -           |
+| 9   | `AdaptiveTransformer` | 8      | Yes   | Mi -> Mo                 | Y      | Y           |
+
+Note: `__call__` and `partial_fit` both have asynchronous alternatives: `__acall__` and `apartial_fit` respectively.
+
+
+### Abstract implementations (Base Classes) for standalone processors
+
+| Idx | Class                     | Parent | Protocol | Features                                                                                   |
+|-----|---------------------------|--------|----------|--------------------------------------------------------------------------------------------|
+| 1   | `BaseProcessor`           | -      | 1        | `__init__` for settings; `__call__` (alias: `send`) wraps abstract `_process`.             |
+| 2   | `BaseProducer`            | -      | 2        | Similar to `BaseProcessor`; `next`/`anext` instead of `send`/`asend` aliases. async first! |
+| 3   | `BaseConsumer`            | 1      | 3        | Overrides return type to None                                                              |
+| 4   | `BaseTransformer`         | 1      | 4        | Overrides input and return types                                                           |
+| 5   | `BaseStatefulProcessor`   | 1      | 5        | `state` setter unpickles arg; `stateful_op` wraps `__call__`                               |
+| 6   | `BaseStatefulProducer`    | 2      | 6        | `state` setter and getter; `stateful_op` wraps `__call__` which runs `__acall__`.          |
+| 7   | `BaseStatefulConsumer`    | 5      | 7        | Overrides return type to None                                                              |
+| 8   | `BaseStatefulTransformer` | 5      | 8        | Overrides input and return types                                                           |
+| 9   | `BaseAdaptiveTransformer` | 8      | 9        | Implements `partial_fit`. `__call__` may call `partial_fit` if message has `.trigger`.     |
+| 10  | `BaseAsyncTransformer`    | 8      | 8        | `__acall__` wraps abstract `_aprocess`; `__call__` runs `__acall__`.                       |
+| 11  | `CompositeProcessor`      | 1      | 5        | Methods iterate over sequence of processors created in `_initialize_processors`.           |
+| 12  | `CompositeProducer`       | 2      | 6        | Similar to `CompositeProcessor`, but first processor must be a producer.                   |
+
+NOTES:
+1. Producers do not inherit from `BaseProcessor`, so concrete implementations should subclass `BaseProducer` or `BaseStatefulProducer`.
+2. For concrete implementations of non-producer processors, inherit from the base subclasses of `BaseProcessor` (eg. `BaseConsumer`, `BaseTransformer`) and from base subclasses of `BaseStatefulProcessor`. These two processor classes are primarily used for efficient abstract base class construction.
+3. For most base classes, the async methods simply call the synchronous methods where the processor logic is expected. Exceptions are `BaseProducer` (and its children) and `BaseAsyncTransformer` which are async-first and should be strongly considered for operations that are I/O bound.
+4. For async-first classes, the logic is implemented in the async methods and the sync methods are thin wrappers around the async methods. The wrapper uses a helper method called `run_coroutine_sync` to run the async method in a synchronous context, but this adds some noticeable processing overhead.
+5. If you need to call your processor outside ezmsg (which uses async), and you cannot easily add an async context* in your processing, then you might want to consider duplicating the processor logic in the sync methods. __Note__: Jupyter notebooks are async by default, so you can await async code in a notebook without any extra setup.
+6. `CompositeProcessor` and `CompositeProducer` are stateful, and structurally subclass the `StatefulProcessor` and `StatefulProducer` protocols, but they
+do not inherit from `BaseStatefulProcessor` and `BaseStatefulProducer`. They accomplish statefulness by inheriting from the mixin abstract base class `CompositeStateful`, which implements the state related methods: `get_state_type`, `state.setter`, `state.getter`, `_hash_message`, `_reset_state`, and `stateful_op` (as well as composite processor chain related methods). However, `BaseStatefulProcessor`, `BaseStatefulProducer` implement `stateful_op` method for a single processor in an incompatible way to what is required for composite chains of processors. 
+
+
+### Generic TypeVars for ezmsg Units
+
+| Idx | Class                     | Description                                                                                                      |
+|-----|---------------------------|------------------------------------------------------------------------------------------------------------------|
+| 5   | `ProducerType`            | bound to `BaseProducer` (hence, also `BaseStatefulProducer`, `CompositeProducer`)                                |
+| 6   | `ConsumerType`            | bound to `BaseConsumer`, `BaseStatefulConsumer`                                                                  |
+| 7   | `TransformerType`         | bound to `BaseTransformer`, `BaseStatefulTransformer`, `CompositeProcessor` (hence, also `BaseAsyncTransformer`) |
+| 8   | `AdaptiveTransformerType` | bound to `BaseAdaptiveTransformer`                                                                               |
+
+
+### Abstract implementations (Base Classes) for ezmsg Units using processors:
+
+| Idx | Class                         | Parents | Expected TypeVars         |
+|-----|-------------------------------|---------|---------------------------|
+| 1   | `BaseProcessorUnit`           | -       | -                         |
+| 2   | `BaseProducerUnit`            | -       | `ProducerType`            |
+| 3   | `BaseConsumerUnit`            | 1       | `ConsumerType`            |
+| 4   | `BaseTransformerUnit`         | 1       | `TransformerType`         |
+| 5   | `BaseAdaptiveTransformerUnit` | 1       | `AdaptiveTransformerType` |
+
+Note, it is strongly recommended to use `BaseConsumerUnit`, `BaseTransformerUnit`, or `BaseAdaptiveTransformerUnit` for implementing concrete subclasses rather than `BaseProcessorUnit`.
+
+
+## Implementing a custom standalone processor
+
+1. Create a new settings dataclass: `class MySettings(ez.Settings):`
+2. Create a new state dataclass: 
+```
+@processor_state
+class MyState:
+```
+3. Decide on your base processor class, considering the protocol, whether it should be async-first, and other factors.
+
+```mermaid
+flowchart TD
+    AMP{Multiple Processors?};
+    AMP -->|no| ARI{Receives Input?};
+    AMP -->|yes| ACB{Single Chain / Branching?}
+    ARI -->|no| P(Producer);
+    ARI -->|yes| APO{Produces Output?};
+    ACB -->|branching| NBC[no base class];
+    ACB -->|single chain| ACRI{Receives Input?};
+    P --> PS{Stateful?};
+    APO -->|no| C(Consumer);
+    APO -->|yes| T(Transformer);
+    ACRI -->|no| CompositeProducer;
+    ACRI -->|yes| CompositeProcessor;
+    PS -->|no| BaseProducer;
+    PS -->|yes| BaseStatefulProducer;
+    C --> CS{Stateful?};
+    T --> TS{Stateful?};
+    CS -->|no| BaseConsumer;
+    CS -->|yes| BaseStatefulConsumer;
+    TS -->|no| BaseTransformer;
+    TS -->|yes| TSA{Adaptive?};
+    TSA -->|no| TSAF{Async First?};
+    TSA -->|yes| BaseAdaptiveTransformer;
+    TSAF -->|no| BaseStatefulTransformer;
+    TSAF -->|yes| BaseAsyncTransformer;
+```
+
+4. Implement the child class.
+    * The minimum implementation is `_process` for sync processors, `_aprocess` for async processors, and `_produce` for producers.
+    * For any stateful processor, implement `_reset_state`.
+    * For stateful processors that need to respond to a change in the incoming data, implement `_hash_message`.
+    * For adaptive transformers, implement `partial_fit`.
+    * For chains of processors (`CompositeProcessor`/ `CompositeProducer`), need to implement `_initialize_processors`.
+    * See processors in `ezmsg.sigproc` for examples.
+5. Override non-abstract methods if you need special behaviour. For example:
+    * `WindowTransformer` overrides `__init__` to do some sanity checks on the provided settings.
+    * `TransposeTransformer` and `WindowTransformer` override `__call__` to provide a passthrough shortcut when the settings allow for it.
+    * `ClockProducer` overrides `__call__` in order to provide a synchronous call bypassing the default async behaviour.
+
+
+## Implementing a custom ezmsg Unit
+
+1. Create and test custom standalone processor as above.
+2. Decide which base unit to implement.
+    * Use the "Generic TypeVars for ezmsg Units" table above to determine the expected TypeVar.
+    * Fine the Expected TypeVar in the "ezmsg Units" table. 
+3. Create the derived class.
+
+Often, all that is required is the following (e.g., for a custom transformer):
+
+```Python
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.sigproc.base import BaseTransformer, BaseTransformerUnit
+
+
+class CustomTransformerSettings(ez.Settings):
+    ...
+
+
+class CustomTransformer(BaseTransformer[CustomTransformerSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        # Your processing code here...
+        return message
+
+
+class CustomUnit(BaseTransformerUnit[
+        CustomTransformerSettings,    # SettingsType
+        AxisArray,                    # MessageInType
+        AxisArray,                    # MessageOutType
+        CustomTransformer,            # TransformerType
+    ]):
+        SETTINGS = CustomTransformerSettings
+```
+
+__Note__, the type of ProcessorUnit is based on the internal processor and not the input or output of the unit. Input streams are allowed in ProducerUnits and output streams in ConsumerUnits. For an example of such a use case, see `BaseCounterFirstProducerUnit` and its subclasses. `BaseCounterFirstProducerUnit` has an input stream that receives a flag signal from a clock that triggers a call to the internal producer. 

{ezmsg_sigproc-1.8.2 → ezmsg_sigproc-2.1.0}/pyproject.toml

@@ -4,28 +4,29 @@ description = "Timeseries signal processing implementations in ezmsg"
 authors = [
   { name = "Griffin Milsap", email = "griffin.milsap@gmail.com" },
   { name = "Preston Peranich", email = "pperanich@gmail.com" },
-  { name = "Chadwick Boulay", email = "chadwick.boulay@gmail.com" }
+  { name = "Chadwick Boulay", email = "chadwick.boulay@gmail.com" },
 ]
 license = "MIT"
 readme = "README.md"
 requires-python = ">=3.10.15"
 dynamic = ["version"]
 dependencies = [
-    "ezmsg>=3.6.0",
-    "numba>=0.61.0",
-    "numpy>=1.26.0",
-    "pywavelets>=1.6.0",
-    "scipy>=1.13.1",
-    "sparse>=0.15.4",
+  "array-api-compat>=1.11.1",
+  "ezmsg>=3.6.0",
+  "numba>=0.61.0",
+  "numpy>=1.26.0",
+  "pywavelets>=1.6.0",
+  "scipy>=1.13.1",
+  "sparse>=0.15.4",
 ]
 
 [project.optional-dependencies]
 test = [
-    "flake8>=7.1.1",
-    "frozendict>=2.4.4",
-    "pytest-asyncio>=0.24.0",
-    "pytest-cov>=5.0.0",
-    "pytest>=8.3.3",
+  "flake8>=7.1.1",
+  "frozendict>=2.4.4",
+  "pytest-asyncio>=0.24.0",
+  "pytest-cov>=5.0.0",
+  "pytest>=8.3.3",
 ]
 
 [build-system]
@@ -42,10 +43,7 @@ version-file = "src/ezmsg/sigproc/__version__.py"
 packages = ["src/ezmsg"]
 
 [tool.uv]
-dev-dependencies = [
-    "pre-commit>=3.8.0",
-    "ruff>=0.6.7",
-]
+dev-dependencies = ["pre-commit>=3.8.0", "ruff>=0.6.7"]
 
 [tool.pytest.ini_options]
 norecursedirs = "tests/helpers"

{ezmsg_sigproc-1.8.2 → ezmsg_sigproc-2.1.0}/src/ezmsg/sigproc/__version__.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '1.8.2'
-__version_tuple__ = version_tuple = (1, 8, 2)
+__version__ = version = '2.1.0'
+__version_tuple__ = version_tuple = (2, 1, 0)

ezmsg_sigproc-2.1.0/src/ezmsg/sigproc/activation.py

@@ -0,0 +1,84 @@
+import scipy.special
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .spectral import OptionsEnum
+from .base import BaseTransformer, BaseTransformerUnit
+
+
+class ActivationFunction(OptionsEnum):
+    """Activation (transformation) function."""
+
+    NONE = "none"
+    """None."""
+
+    SIGMOID = "sigmoid"
+    """:obj:`scipy.special.expit`"""
+
+    EXPIT = "expit"
+    """:obj:`scipy.special.expit`"""
+
+    LOGIT = "logit"
+    """:obj:`scipy.special.logit`"""
+
+    LOGEXPIT = "log_expit"
+    """:obj:`scipy.special.log_expit`"""
+
+
+ACTIVATIONS = {
+    ActivationFunction.NONE: lambda x: x,
+    ActivationFunction.SIGMOID: scipy.special.expit,
+    ActivationFunction.EXPIT: scipy.special.expit,
+    ActivationFunction.LOGIT: scipy.special.logit,
+    ActivationFunction.LOGEXPIT: scipy.special.log_expit,
+}
+
+
+class ActivationSettings(ez.Settings):
+    function: str | ActivationFunction = ActivationFunction.NONE
+    """An enum value from ActivationFunction or a string representing the activation function.
+         Possible values are: SIGMOID, EXPIT, LOGIT, LOGEXPIT, "sigmoid", "expit", "logit", "log_expit".
+         SIGMOID and EXPIT are equivalent. See :obj:`scipy.special.expit` for more details."""
+
+
+class ActivationTransformer(BaseTransformer[ActivationSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        if type(self.settings.function) is ActivationFunction:
+            func = ACTIVATIONS[self.settings.function]
+        else:
+            # str type handling
+            function = self.settings.function.lower()
+            if function not in ActivationFunction.options():
+                raise ValueError(
+                    f"Unrecognized activation function {function}. Must be one of {ACTIVATIONS.keys()}"
+                )
+            function = list(ACTIVATIONS.keys())[
+                ActivationFunction.options().index(function)
+            ]
+            func = ACTIVATIONS[function]
+
+        return replace(message, data=func(message.data))
+
+
+class Activation(
+    BaseTransformerUnit[ActivationSettings, AxisArray, AxisArray, ActivationTransformer]
+):
+    SETTINGS = ActivationSettings
+
+
+def activation(
+    function: str | ActivationFunction,
+) -> ActivationTransformer:
+    """
+    Transform the data with a simple activation function.
+
+    Args:
+        function: An enum value from ActivationFunction or a string representing the activation function.
+         Possible values are: SIGMOID, EXPIT, LOGIT, LOGEXPIT, "sigmoid", "expit", "logit", "log_expit".
+         SIGMOID and EXPIT are equivalent. See :obj:`scipy.special.expit` for more details.
+
+    Returns: :obj:`ActivationTransformer`
+
+    """
+    return ActivationTransformer(ActivationSettings(function=function))

ezmsg_sigproc-2.1.0/src/ezmsg/sigproc/adaptive_lattice_notch.py

@@ -0,0 +1,231 @@
+import numpy as np
+import numpy.typing as npt
+import scipy.signal
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis
+from ezmsg.util.messages.util import replace
+
+from .base import processor_state, BaseStatefulTransformer
+
+
+class AdaptiveLatticeNotchFilterSettings(ez.Settings):
+    """Settings for the Adaptive Lattice Notch Filter."""
+
+    gamma: float = 0.995
+    """Pole-zero contraction factor"""
+    mu: float = 0.99
+    """Smoothing factor"""
+    eta: float = 0.99
+    """Forgetting factor"""
+    axis: str = "time"
+    """Axis to apply filter to"""
+    init_notch_freq: float | None = None
+    """Initial notch frequency. Should be < nyquist."""
+    chunkwise: bool = False
+    """Speed up processing by updating the target freq once per chunk only."""
+
+
+@processor_state
+class AdaptiveLatticeNotchFilterState:
+    """State for the Adaptive Lattice Notch Filter."""
+
+    s_history: npt.NDArray | None = None
+    """Historical `s` values for the adaptive filter."""
+
+    p: npt.NDArray | None = None
+    """Accumulated product for reflection coefficient update"""
+
+    q: npt.NDArray | None = None
+    """Accumulated product for reflection coefficient update"""
+
+    k1: npt.NDArray | None = None
+    """Reflection coefficient"""
+
+    freq_template: CoordinateAxis | None = None
+    """Template for the frequency axis on the output"""
+
+    zi: npt.NDArray | None = None
+    """Initial conditions for the filter, updated after every chunk"""
+
+
+class AdaptiveLatticeNotchFilterTransformer(
+    BaseStatefulTransformer[
+        AdaptiveLatticeNotchFilterSettings,
+        AxisArray,
+        AxisArray,
+        AdaptiveLatticeNotchFilterState,
+    ]
+):
+    """
+    Adaptive Lattice Notch Filter implementation as a stateful transformer.
+
+    https://biomedical-engineering-online.biomedcentral.com/articles/10.1186/1475-925X-13-170
+
+    The filter automatically tracks and removes frequency components from the input signal.
+    It outputs the estimated frequency (in Hz) and the filtered sample.
+    """
+
+    def _hash_message(self, message: AxisArray) -> int:
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
+        return hash((message.key, message.axes[self.settings.axis].gain, sample_shape))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        ax_idx = message.get_axis_idx(self.settings.axis)
+        sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
+
+        fs = 1 / message.axes[self.settings.axis].gain
+        init_f = (
+            self.settings.init_notch_freq
+            if self.settings.init_notch_freq is not None
+            else 0.07178314656435313 * fs
+        )
+        init_omega = init_f * (2 * np.pi) / fs
+        init_k1 = -np.cos(init_omega)
+
+        """Reset filter state to initial values."""
+        self._state = AdaptiveLatticeNotchFilterState()
+        self._state.s_history = np.zeros((2,) + sample_shape, dtype=float)
+        self._state.p = np.zeros(sample_shape, dtype=float)
+        self._state.q = np.zeros(sample_shape, dtype=float)
+        self._state.k1 = init_k1 + np.zeros(sample_shape, dtype=float)
+        self._state.freq_template = CoordinateAxis(
+            data=np.zeros((0,) + sample_shape, dtype=float),
+            dims=[self.settings.axis]
+            + message.dims[:ax_idx]
+            + message.dims[ax_idx + 1 :],
+            unit="Hz",
+        )
+
+        # Initialize the initial conditions for the filter
+        self._state.zi = np.zeros((2, np.prod(sample_shape)), dtype=float)
+        # Note: we could calculate it properly, but as long as we are initializing s_history with zeros,
+        #  it will always be zero.
+        # a = [1, init_k1 * (1 + self.settings.gamma), self.settings.gamma]
+        # b = [1]
+        # s = np.reshape(self._state.s_history, (2, -1))
+        # for feat_ix in range(np.prod(sample_shape)):
+        #     self._state.zi[:, feat_ix] = scipy.signal.lfiltic(b, a, s[::-1, feat_ix], x=None)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        x_data = message.data
+        ax_idx = message.get_axis_idx(self.settings.axis)
+
+        # TODO: Time should be moved to -1th axis, not the 0th axis
+        if message.dims[0] != self.settings.axis:
+            x_data = np.moveaxis(x_data, ax_idx, 0)
+
+        # Access settings once
+        gamma = self.settings.gamma
+        eta = self.settings.eta
+        mu = self.settings.mu
+        fs = 1 / message.axes[self.settings.axis].gain
+
+        # Pre-compute constants
+        one_minus_eta = 1 - eta
+        one_minus_mu = 1 - mu
+        gamma_plus_1 = 1 + gamma
+        omega_scale = fs / (2 * np.pi)
+
+        # For the lattice filter with constant k1:
+        # s_n = x_n - k1*(1+gamma)*s_n_1 - gamma*s_n_2
+        # This is equivalent to an IIR filter with b=1, a=[1, k1*(1+gamma), gamma]
+
+        # For the output filter:
+        # y_n = s_n + 2*k1*s_n_1 + s_n_2
+        # We can treat this as a direct-form FIR filter applied to s_out
+
+        if self.settings.chunkwise:
+            # Process each chunk using current filter parameters
+            # Reshape input and prepare output arrays
+            _s = self._state.s_history.reshape((2, -1))
+            _x = x_data.reshape((x_data.shape[0], -1))
+            s_n = np.zeros_like(_x)
+            y_out = np.zeros_like(_x)
+
+            # Apply static filter for each feature dimension
+            for ix, k in enumerate(self._state.k1.flatten()):
+                # Filter to get s_n (notch filter state)
+                a_s = [1, k * gamma_plus_1, gamma]
+                s_n[:, ix], self._state.zi[:, ix] = scipy.signal.lfilter(
+                    [1], a_s, _x[:, ix], zi=self._state.zi[:, ix]
+                )
+
+                # Apply output filter to get y_out
+                b_y = [1, 2 * k, 1]
+                y_out[:, ix] = scipy.signal.lfilter(b_y, [1], s_n[:, ix])
+
+            # Update filter parameters using final values from the chunk
+            s_n_reshaped = s_n.reshape((s_n.shape[0],) + x_data.shape[1:])
+            s_final = s_n_reshaped[-1]  # Current s_n
+            s_final_1 = s_n_reshaped[-2]  # s_n_1
+            s_final_2 = (
+                s_n_reshaped[-3] if len(s_n_reshaped) > 2 else self._state.s_history[0]
+            )  # s_n_2
+
+            # Update p and q using final values
+            self._state.p = eta * self._state.p + one_minus_eta * (
+                s_final_1 * (s_final + s_final_2)
+            )
+            self._state.q = eta * self._state.q + one_minus_eta * (
+                2 * (s_final_1 * s_final_1)
+            )
+
+            # Update reflection coefficient
+            new_k1 = -self._state.p / (self._state.q + 1e-8)  # Avoid division by zero
+            new_k1 = np.clip(new_k1, -1, 1)  # Clip to prevent instability
+            self._state.k1 = mu * self._state.k1 + one_minus_mu * new_k1  # Smoothed
+
+            # Calculate frequency from updated k1 value
+            omega_n = np.arccos(-self._state.k1)
+            freq = omega_n * omega_scale
+            freq_out = np.full_like(x_data.reshape(x_data.shape), freq)
+
+            # Update s_history for next chunk
+            self._state.s_history = s_n_reshaped[-2:].reshape((2,) + x_data.shape[1:])
+
+            # Reshape y_out back to original dimensions
+            y_out = y_out.reshape(x_data.shape)
+
+        else:
+            # Perform filtering, sample-by-sample
+            y_out = np.zeros_like(x_data)
+            freq_out = np.zeros_like(x_data)
+            for sample_ix, x_n in enumerate(x_data):
+                s_n_1 = self._state.s_history[-1]
+                s_n_2 = self._state.s_history[-2]
+
+                s_n = x_n - self._state.k1 * gamma_plus_1 * s_n_1 - gamma * s_n_2
+                y_out[sample_ix] = s_n + 2 * self._state.k1 * s_n_1 + s_n_2
+
+                # Update filter parameters
+                self._state.p = eta * self._state.p + one_minus_eta * (
+                    s_n_1 * (s_n + s_n_2)
+                )
+                self._state.q = eta * self._state.q + one_minus_eta * (
+                    2 * (s_n_1 * s_n_1)
+                )
+
+                # Update reflection coefficient
+                new_k1 = -self._state.p / (
+                    self._state.q + 1e-8
+                )  # Avoid division by zero
+                new_k1 = np.clip(new_k1, -1, 1)  # Clip to prevent instability
+                self._state.k1 = mu * self._state.k1 + one_minus_mu * new_k1  # Smoothed
+
+                # Compute normalized angular frequency using equation 13 from the paper
+                omega_n = np.arccos(-self._state.k1)
+                freq_out[sample_ix] = omega_n * omega_scale  # As Hz
+
+                # Update for next iteration
+                self._state.s_history[-2] = s_n_1
+                self._state.s_history[-1] = s_n
+
+        return replace(
+            message,
+            data=y_out,
+            axes={
+                **message.axes,
+                "freq": replace(self._state.freq_template, data=freq_out),
+            },
+        )