ezmsg-sigproc 1.2.2__py3-none-any.whl → 2.10.0__py3-none-any.whl

ezmsg/sigproc/__init__.py

@@ -1 +1 @@
-from .__version__ import __version__
+from .__version__ import __version__ as __version__

ezmsg/sigproc/__version__.py

@@ -1 +1,34 @@
-__version__ = "1.2.2"
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '2.10.0'
+__version_tuple__ = version_tuple = (2, 10, 0)
+
+__commit_id__ = commit_id = None

ezmsg/sigproc/activation.py

@@ -0,0 +1,78 @@
+import ezmsg.core as ez
+import scipy.special
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .spectral import OptionsEnum
+
+
+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/adaptive_lattice_notch.py

@@ -0,0 +1,212 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+import scipy.signal
+from ezmsg.baseproc import BaseStatefulTransformer, processor_state
+from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis
+from ezmsg.util.messages.util import replace
+
+
+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),
+            },
+        )

ezmsg/sigproc/affinetransform.py

@@ -0,0 +1,235 @@
+"""Affine transformations via matrix multiplication: y = Ax or y = Ax + B.
+
+For full matrix transformations where channels are mixed (off-diagonal weights),
+use :obj:`AffineTransformTransformer` or the `AffineTransform` unit.
+
+For simple per-channel scaling and offset (diagonal weights only), use
+:obj:`LinearTransformTransformer` from :mod:`ezmsg.sigproc.linear` instead,
+which is more efficient as it avoids matrix multiplication.
+"""
+
+import os
+from pathlib import Path
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.baseproc import (
+    BaseStatefulTransformer,
+    BaseTransformer,
+    BaseTransformerUnit,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import AxisArray, AxisBase
+from ezmsg.util.messages.util import replace
+
+
+class AffineTransformSettings(ez.Settings):
+    """
+    Settings for :obj:`AffineTransform`.
+    """
+
+    weights: np.ndarray | str | Path
+    """An array of weights or a path to a file with weights compatible with np.loadtxt."""
+
+    axis: str | None = None
+    """The name of the axis to apply the transformation to. Defaults to the leading (0th) axis in the array."""
+
+    right_multiply: bool = True
+    """Set False to transpose the weights before applying."""
+
+
+@processor_state
+class AffineTransformState:
+    weights: npt.NDArray | None = None
+    new_axis: AxisBase | None = None
+
+
+class AffineTransformTransformer(
+    BaseStatefulTransformer[AffineTransformSettings, AxisArray, AxisArray, AffineTransformState]
+):
+    """Apply affine transformation via matrix multiplication: y = Ax or y = Ax + B.
+
+    Use this transformer when you need full matrix transformations that mix
+    channels (off-diagonal weights), such as spatial filters or projections.
+
+    For simple per-channel scaling and offset where each output channel depends
+    only on its corresponding input channel (diagonal weight matrix), use
+    :obj:`LinearTransformTransformer` instead, which is more efficient.
+
+    The weights matrix can include an offset row (stacked as [A|B]) where the
+    input is automatically augmented with a column of ones to compute y = Ax + B.
+    """
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        # Override __call__ so we can shortcut if weights are None.
+        if self.settings.weights is None or (
+            isinstance(self.settings.weights, str) and self.settings.weights == "passthrough"
+        ):
+            return message
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        return hash(message.key)
+
+    def _reset_state(self, message: AxisArray) -> None:
+        weights = self.settings.weights
+        if isinstance(weights, str):
+            weights = Path(os.path.abspath(os.path.expanduser(weights)))
+        if isinstance(weights, Path):
+            weights = np.loadtxt(weights, delimiter=",")
+        if not self.settings.right_multiply:
+            weights = weights.T
+        if weights is not None:
+            weights = np.ascontiguousarray(weights)
+
+        self._state.weights = weights
+
+        axis = self.settings.axis or message.dims[-1]
+        if axis in message.axes and hasattr(message.axes[axis], "data") and weights.shape[0] != weights.shape[1]:
+            in_labels = message.axes[axis].data
+            new_labels = []
+            n_in, n_out = weights.shape
+            if len(in_labels) != n_in:
+                ez.logger.warning(f"Received {len(in_labels)} for {n_in} inputs. Check upstream labels.")
+            else:
+                b_filled_outputs = np.any(weights, axis=0)
+                b_used_inputs = np.any(weights, axis=1)
+                if np.all(b_used_inputs) and np.all(b_filled_outputs):
+                    new_labels = []
+                elif np.all(b_used_inputs):
+                    in_ix = 0
+                    new_labels = []
+                    for out_ix in range(n_out):
+                        if b_filled_outputs[out_ix]:
+                            new_labels.append(in_labels[in_ix])
+                            in_ix += 1
+                        else:
+                            new_labels.append("")
+                elif np.all(b_filled_outputs):
+                    new_labels = np.array(in_labels)[b_used_inputs]
+
+            self._state.new_axis = replace(message.axes[axis], data=np.array(new_labels))
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        data = message.data
+
+        if data.shape[axis_idx] == (self._state.weights.shape[0] - 1):
+            # The weights are stacked A|B where A is the transform and B is a single row
+            #  in the equation y = Ax + B. This supports NeuroKey's weights matrices.
+            sample_shape = data.shape[:axis_idx] + (1,) + data.shape[axis_idx + 1 :]
+            data = np.concatenate((data, np.ones(sample_shape).astype(data.dtype)), axis=axis_idx)
+
+        if axis_idx in [-1, len(message.dims) - 1]:
+            data = np.matmul(data, self._state.weights)
+        else:
+            data = np.moveaxis(data, axis_idx, -1)
+            data = np.matmul(data, self._state.weights)
+            data = np.moveaxis(data, -1, axis_idx)
+
+        replace_kwargs = {"data": data}
+        if self._state.new_axis is not None:
+            replace_kwargs["axes"] = {**message.axes, axis: self._state.new_axis}
+
+        return replace(message, **replace_kwargs)
+
+
+class AffineTransform(BaseTransformerUnit[AffineTransformSettings, AxisArray, AxisArray, AffineTransformTransformer]):
+    SETTINGS = AffineTransformSettings
+
+
+def affine_transform(
+    weights: np.ndarray | str | Path,
+    axis: str | None = None,
+    right_multiply: bool = True,
+) -> AffineTransformTransformer:
+    """
+    Perform affine transformations on streaming data.
+
+    Args:
+        weights: An array of weights or a path to a file with weights compatible with np.loadtxt.
+        axis: The name of the axis to apply the transformation to. Defaults to the leading (0th) axis in the array.
+        right_multiply: Set False to transpose the weights before applying.
+
+    Returns:
+        :obj:`AffineTransformTransformer`.
+    """
+    return AffineTransformTransformer(
+        AffineTransformSettings(weights=weights, axis=axis, right_multiply=right_multiply)
+    )
+
+
+def zeros_for_noop(data: npt.NDArray, **ignore_kwargs) -> npt.NDArray:
+    return np.zeros_like(data)
+
+
+class CommonRereferenceSettings(ez.Settings):
+    """
+    Settings for :obj:`CommonRereference`
+    """
+
+    mode: str = "mean"
+    """The statistical mode to apply -- either "mean" or "median"."""
+
+    axis: str | None = None
+    """The name of the axis to apply the transformation to."""
+
+    include_current: bool = True
+    """Set False to exclude each channel from participating in the calculation of its reference."""
+
+
+class CommonRereferenceTransformer(BaseTransformer[CommonRereferenceSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        if self.settings.mode == "passthrough":
+            return message
+
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+
+        func = {"mean": np.mean, "median": np.median, "passthrough": zeros_for_noop}[self.settings.mode]
+
+        ref_data = func(message.data, axis=axis_idx, keepdims=True)
+
+        if not self.settings.include_current:
+            # Typical `CAR = x[0]/N + x[1]/N + ... x[i-1]/N + x[i]/N + x[i+1]/N + ... + x[N-1]/N`
+            # and is the same for all i, so it is calculated only once in `ref_data`.
+            # However, if we had excluded the current channel,
+            # then we would have omitted the contribution of the current channel:
+            # `CAR[i] = x[0]/(N-1) + x[1]/(N-1) + ... x[i-1]/(N-1) + x[i+1]/(N-1) + ... + x[N-1]/(N-1)`
+            # The majority of the calculation is the same as when the current channel is included;
+            # we need only rescale CAR so the divisor is `N-1` instead of `N`, then subtract the contribution
+            # from the current channel (i.e., `x[i] / (N-1)`)
+            #  i.e., `CAR[i] = (N / (N-1)) * common_CAR - x[i]/(N-1)`
+            # We can use broadcasting subtraction instead of looping over channels.
+            N = message.data.shape[axis_idx]
+            ref_data = (N / (N - 1)) * ref_data - message.data / (N - 1)
+            # Note: I profiled using AffineTransformTransformer; it's ~30x slower than this implementation.
+
+        return replace(message, data=message.data - ref_data)
+
+
+class CommonRereference(
+    BaseTransformerUnit[CommonRereferenceSettings, AxisArray, AxisArray, CommonRereferenceTransformer]
+):
+    SETTINGS = CommonRereferenceSettings
+
+
+def common_rereference(
+    mode: str = "mean", axis: str | None = None, include_current: bool = True
+) -> CommonRereferenceTransformer:
+    """
+    Perform common average referencing (CAR) on streaming data.
+
+    Args:
+        mode: The statistical mode to apply -- either "mean" or "median"
+        axis: The name of hte axis to apply the transformation to.
+        include_current: Set False to exclude each channel from participating in the calculation of its reference.
+
+    Returns:
+        :obj:`CommonRereferenceTransformer`
+    """
+    return CommonRereferenceTransformer(
+        CommonRereferenceSettings(mode=mode, axis=axis, include_current=include_current)
+    )