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

ezmsg/sigproc/ewma.py

@@ -0,0 +1,217 @@
+import functools
+from dataclasses import field
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+import scipy.signal as sps
+from ezmsg.baseproc import BaseStatefulTransformer, BaseTransformerUnit, processor_state
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+from ezmsg.util.messages.util import replace
+
+
+def _tau_from_alpha(alpha: float, dt: float) -> float:
+    """
+    Inverse of _alpha_from_tau. See that function for explanation.
+    """
+    return -dt / np.log(1 - alpha)
+
+
+def _alpha_from_tau(tau: float, dt: float) -> float:
+    """
+    # https://en.wikipedia.org/wiki/Exponential_smoothing#Time_constant
+    :param tau: The amount of time for the smoothed response of a unit step function to reach
+        1 - 1/e approx-eq 63.2%.
+    :param dt: sampling period, or 1 / sampling_rate.
+    :return: alpha, the "fading factor" in exponential smoothing.
+    """
+    return 1 - np.exp(-dt / tau)
+
+
+def ewma_step(sample: npt.NDArray, zi: npt.NDArray, alpha: float, beta: float | None = None):
+    """
+    Do an exponentially weighted moving average step.
+
+    Args:
+        sample: The new sample.
+        zi: The output of the previous step.
+        alpha: Fading factor.
+        beta: Persisting factor. If None, it is calculated as 1-alpha.
+
+    Returns:
+        alpha * sample + beta * zi
+
+    """
+    # Potential micro-optimization:
+    #  Current: scalar-arr multiplication, scalar-arr multiplication, arr-arr addition
+    #  Alternative: arr-arr subtraction, arr-arr multiplication, arr-arr addition
+    # return zi + alpha * (new_sample - zi)
+    beta = beta or (1 - alpha)
+    return alpha * sample + beta * zi
+
+
+class EWMA_Deprecated:
+    """
+    Grabbed these methods from https://stackoverflow.com/a/70998068 and other answers in that topic,
+    but they ended up being slower than the scipy.signal.lfilter method.
+    Additionally, `compute` and `compute2` suffer from potential errors as the vector length increases
+    and beta**n approaches zero.
+    """
+
+    def __init__(self, alpha: float, max_len: int):
+        self.alpha = alpha
+        self.beta = 1 - alpha
+        self.prev: npt.NDArray | None = None
+        self.weights = np.empty((max_len + 1,), float)
+        self._precalc_weights(max_len)
+        self._step_func = functools.partial(ewma_step, alpha=self.alpha, beta=self.beta)
+
+    def _precalc_weights(self, n: int):
+        #   (1-α)^0, (1-α)^1, (1-α)^2, ..., (1-α)^n
+        np.power(self.beta, np.arange(n + 1), out=self.weights)
+
+    def compute(self, arr: npt.NDArray, out: npt.NDArray | None = None) -> npt.NDArray:
+        if out is None:
+            out = np.empty(arr.shape, arr.dtype)
+
+        n = arr.shape[0]
+        weights = self.weights[:n]
+        weights = np.expand_dims(weights, list(range(1, arr.ndim)))
+
+        #   α*P0, α*P1, α*P2, ..., α*Pn
+        np.multiply(self.alpha, arr, out)
+
+        #   α*P0/(1-α)^0, α*P1/(1-α)^1, α*P2/(1-α)^2, ..., α*Pn/(1-α)^n
+        np.divide(out, weights, out)
+
+        #   α*P0/(1-α)^0, α*P0/(1-α)^0 + α*P1/(1-α)^1, ...
+        np.cumsum(out, axis=0, out=out)
+
+        #   (α*P0/(1-α)^0)*(1-α)^0, (α*P0/(1-α)^0 + α*P1/(1-α)^1)*(1-α)^1, ...
+        np.multiply(out, weights, out)
+
+        # Add the previous output
+        if self.prev is None:
+            self.prev = arr[:1]
+
+        out += self.prev * np.expand_dims(self.weights[1 : n + 1], list(range(1, arr.ndim)))
+
+        self.prev = out[-1:]
+
+        return out
+
+    def compute2(self, arr: npt.NDArray) -> npt.NDArray:
+        """
+        Compute the Exponentially Weighted Moving Average (EWMA) of the input array.
+
+        Args:
+            arr: The input array to be smoothed.
+
+        Returns:
+            The smoothed array.
+        """
+        n = arr.shape[0]
+        if n > len(self.weights):
+            self._precalc_weights(n)
+        weights = self.weights[:n][::-1]
+        weights = np.expand_dims(weights, list(range(1, arr.ndim)))
+
+        result = np.cumsum(self.alpha * weights * arr, axis=0)
+        result = result / weights
+
+        # Handle the first call when prev is unset
+        if self.prev is None:
+            self.prev = arr[:1]
+
+        result += self.prev * np.expand_dims(self.weights[1 : n + 1], list(range(1, arr.ndim)))
+
+        # Store the result back into prev
+        self.prev = result[-1]
+
+        return result
+
+    def compute_sample(self, new_sample: npt.NDArray) -> npt.NDArray:
+        if self.prev is None:
+            self.prev = new_sample
+        self.prev = self._step_func(new_sample, self.prev)
+        return self.prev
+
+
+class EWMASettings(ez.Settings):
+    time_constant: float = 1.0
+    """The amount of time for the smoothed response of a unit step function to reach 1 - 1/e approx-eq 63.2%."""
+
+    axis: str | None = None
+
+    accumulate: bool = True
+    """If True, update the EWMA state with each sample. If False, only apply
+    the current EWMA estimate without updating state (useful for inference
+    periods where you don't want to adapt statistics)."""
+
+
+@processor_state
+class EWMAState:
+    alpha: float = field(default_factory=lambda: _alpha_from_tau(1.0, 1000.0))
+    zi: npt.NDArray | None = None
+
+
+class EWMATransformer(BaseStatefulTransformer[EWMASettings, AxisArray, AxisArray, EWMAState]):
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        sample_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((sample_shape, message.axes[axis].gain, message.key))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = self.settings.axis or message.dims[0]
+        self._state.alpha = _alpha_from_tau(self.settings.time_constant, message.axes[axis].gain)
+        sub_dat = slice_along_axis(message.data, slice(None, 1, None), axis=message.get_axis_idx(axis))
+        self._state.zi = (1 - self._state.alpha) * sub_dat
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if np.prod(message.data.shape) == 0:
+            return message
+        axis = self.settings.axis or message.dims[0]
+        axis_idx = message.get_axis_idx(axis)
+        if self.settings.accumulate:
+            # Normal behavior: update state with new samples
+            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,
+            )
+        else:
+            # Process-only: compute output without updating state
+            expected, _ = sps.lfilter(
+                [self._state.alpha],
+                [1.0, self._state.alpha - 1.0],
+                message.data,
+                axis=axis_idx,
+                zi=self._state.zi,
+            )
+        return replace(message, data=expected)
+
+
+class EWMAUnit(BaseTransformerUnit[EWMASettings, AxisArray, AxisArray, EWMATransformer]):
+    SETTINGS = EWMASettings
+
+    @ez.subscriber(BaseTransformerUnit.INPUT_SETTINGS)
+    async def on_settings(self, msg: EWMASettings) -> None:
+        """
+        Handle settings updates with smart reset behavior.
+
+        Only resets state if `axis` changes (structural change).
+        Changes to `time_constant` or `accumulate` are applied without
+        resetting accumulated state.
+        """
+        old_axis = self.SETTINGS.axis
+        self.apply_settings(msg)
+
+        if msg.axis != old_axis:
+            # Axis changed - need full reset
+            self.create_processor()
+        else:
+            # Only accumulate or time_constant changed - keep state
+            self.processor.settings = msg

ezmsg/sigproc/ewmfilter.py

@@ -2,9 +2,9 @@ import asyncio
 import typing
 
 import ezmsg.core as ez
+import numpy as np
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
-import numpy as np
 
 from .window import Window, WindowSettings
 

ezmsg/sigproc/extract_axis.py

@@ -0,0 +1,39 @@
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray, replace
+
+
+class ExtractAxisSettings(ez.Settings):
+    axis: str = "freq"
+    reference: str = "time"
+
+
+class ExtractAxisData(BaseTransformer[ExtractAxisSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        targ_ax = message.axes[self.settings.axis]
+        if hasattr(targ_ax, "data"):
+            # Extracted axis is of type CoordinateAxis
+            return replace(
+                message,
+                data=targ_ax.data,
+                dims=targ_ax.dims,
+                axes={k: v for k, v in message.axes.items() if k in targ_ax.dims},
+            )
+            # Note: So far we don't have any transformers where the coordinate axis has its own axes,
+            # but if that happens in the future, we'd need to consider how to handle that.
+
+        else:
+            # Extracted axis is of type LinearAxis
+            #  LinearAxis can only yield a 1d array data which simplifies dims and axes.
+            n = message.data.shape[message.get_axis_idx(self.settings.reference)]
+            return replace(
+                message,
+                data=targ_ax.value(np.arange(n)),
+                dims=[self.settings.reference],
+                axes={self.settings.reference: message.axes[self.settings.reference]},
+            )
+
+
+class ExtractAxisDataUnit(BaseTransformerUnit[ExtractAxisSettings, AxisArray, AxisArray, ExtractAxisData]):
+    SETTINGS = ExtractAxisSettings

ezmsg/sigproc/fbcca.py

@@ -0,0 +1,307 @@
+import math
+import typing
+from dataclasses import field
+
+import ezmsg.core as ez
+import numpy as np
+from ezmsg.baseproc import (
+    BaseProcessor,
+    BaseStatefulProcessor,
+    BaseTransformer,
+    BaseTransformerUnit,
+    CompositeProcessor,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .filterbankdesign import (
+    FilterbankDesignSettings,
+    FilterbankDesignTransformer,
+)
+from .kaiser import KaiserFilterSettings
+from .sampler import SampleTriggerMessage
+from .window import WindowSettings, WindowTransformer
+
+
+class FBCCASettings(ez.Settings):
+    """
+    Settings for :obj:`FBCCATransformer`
+    """
+
+    time_dim: str
+    """
+    The time dim in the data array.
+    """
+
+    ch_dim: str
+    """
+    The channels dim in the data array.
+    """
+
+    filterbank_dim: str | None = None
+    """
+    The filter bank subband dim in the data array. If unspecified, method falls back to CCA
+    None (default): the input has no subbands; just use CCA
+    """
+
+    harmonics: int = 5
+    """
+    The number of additional harmonics beyond the fundamental to use for the 'design' matrix.
+    5 (default): Evaluate 5 harmonics of the base frequency.
+    Many periodic signals are not pure sinusoids, and inclusion of higher harmonics can help evaluate the
+    presence of signals with higher frequency harmonic content
+    """
+
+    freqs: typing.List[float] = field(default_factory=list)
+    """
+    Frequencies (in hz) to evaluate the presence of within the input signal.
+    [] (default): an empty list; frequencies will be found within the input SampleMessages.
+    AxisArrays have no good place to put this metadata, so specify frequencies here if only AxisArrays
+    will be passed as input to the generator.  If the input has a `trigger` attr of type :obj:`SampleTriggerMessage`,
+    the processor looks for the `freqs` attribute within that trigger for a list of frequencies to evaluate.
+    This field is present in the :obj:`SSVEPSampleTriggerMessage` defined in ezmsg.tasks.ssvep from
+    the ezmsg-tasks package.
+    NOTE: Avoid frequencies that have line-noise (60 Hz/50 Hz) as a harmonic.
+    """
+
+    softmax_beta: float = 1.0
+    """
+    Beta parameter for softmax on output --> "probabilities".
+    1.0 (default): Use the shifted softmax transformation to output 0-1 probabilities.
+    If 0.0, the maximum singular value of the SVD for each design matrix is output
+    """
+
+    target_freq_dim: str = "target_freq"
+    """
+    Name for dim to put target frequency outputs on.
+    'target_freq' (default)
+    """
+
+    max_int_time: float = 0.0
+    """
+    Maximum integration time (in seconds) to use for calculation.
+    0 (default): Use all time provided for the calculation.
+    Useful for artificially limiting the amount of data used for the CCA method to evaluate
+    the necessary integration time for good decoding performance
+    """
+
+
+class FBCCATransformer(BaseTransformer[FBCCASettings, AxisArray, AxisArray]):
+    """
+    A canonical-correlation (CCA) signal decoder for detection of periodic activity in multi-channel timeseries
+    recordings. It is particularly useful for detecting the presence of steady-state evoked responses in multi-channel
+    EEG data. Please see Lin et. al. 2007 for a description on the use of CCA to detect the presence of SSVEP in EEG
+    data.
+    This implementation also includes the "Filterbank" extension of the CCA decoding approach which utilizes a
+    filterbank to decompose input multi-channel EEG data into several frequency sub-bands; each of which is analyzed
+    with CCA, then combined using a weighted sum; allowing CCA to more readily identify harmonic content in EEG data.
+    Read more about this approach in Chen et. al. 2015.
+
+    ## Further reading:
+    * [Lin et. al. 2007](https://ieeexplore.ieee.org/document/4015614)
+    * [Nakanishi et. al. 2015](https://doi.org/10.1371%2Fjournal.pone.0140703)
+    * [Chen et. al. 2015](http://dx.doi.org/10.1088/1741-2560/12/4/046008)
+    """
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        """
+        Input: AxisArray with at least a time_dim, and ch_dim
+        Output: AxisArray with time_dim, ch_dim, (and filterbank_dim if specified)
+            collapsed, with a new 'target_freq' dim of length 'freqs'
+        """
+
+        test_freqs: list[float] = self.settings.freqs
+        trigger = message.attrs.get("trigger", None)
+        if isinstance(trigger, SampleTriggerMessage):
+            if len(test_freqs) == 0:
+                test_freqs = getattr(trigger, "freqs", [])
+
+        if len(test_freqs) == 0:
+            raise ValueError("no frequencies to test")
+
+        time_dim_idx = message.get_axis_idx(self.settings.time_dim)
+        ch_dim_idx = message.get_axis_idx(self.settings.ch_dim)
+
+        filterbank_dim_idx = None
+        if self.settings.filterbank_dim is not None:
+            filterbank_dim_idx = message.get_axis_idx(self.settings.filterbank_dim)
+
+        # Move (filterbank_dim), time, ch to end of array
+        rm_dims = [self.settings.time_dim, self.settings.ch_dim]
+        if self.settings.filterbank_dim is not None:
+            rm_dims = [self.settings.filterbank_dim] + rm_dims
+        new_order = [i for i, dim in enumerate(message.dims) if dim not in rm_dims]
+        if filterbank_dim_idx is not None:
+            new_order.append(filterbank_dim_idx)
+        new_order.extend([time_dim_idx, ch_dim_idx])
+        out_dims = [message.dims[i] for i in new_order if message.dims[i] not in rm_dims]
+        data_arr = message.data.transpose(new_order)
+
+        # Add a singleton dim for filterbank dim if we don't have one
+        if filterbank_dim_idx is None:
+            data_arr = data_arr[..., None, :, :]
+            filterbank_dim_idx = data_arr.ndim - 3
+
+        # data_arr is now (..., filterbank, time, ch)
+        # Get output shape for remaining dims and reshape data_arr for iterative processing
+        out_shape = list(data_arr.shape[:-3])
+        data_arr = data_arr.reshape([math.prod(out_shape), *data_arr.shape[-3:]])
+
+        # Create output dims and axes with added target_freq_dim
+        out_shape.append(len(test_freqs))
+        out_dims.append(self.settings.target_freq_dim)
+        out_axes = {
+            axis_name: axis
+            for axis_name, axis in message.axes.items()
+            if axis_name not in rm_dims
+            and not (isinstance(axis, AxisArray.CoordinateAxis) and any(d in rm_dims for d in axis.dims))
+        }
+        out_axes[self.settings.target_freq_dim] = AxisArray.CoordinateAxis(
+            np.array(test_freqs), [self.settings.target_freq_dim]
+        )
+
+        if message.data.size == 0:
+            out_data = message.data.reshape(out_shape)
+            output = replace(message, data=out_data, dims=out_dims, axes=out_axes)
+            return output
+
+        # Get time axis
+        t_ax_info = message.ax(self.settings.time_dim)
+        t = t_ax_info.values
+        t -= t[0]
+        max_samp = len(t)
+        if self.settings.max_int_time > 0:
+            max_samp = int(abs(t_ax_info.values - self.settings.max_int_time).argmin())
+        t = t[:max_samp]
+
+        calc_output = np.zeros((*data_arr.shape[:-2], len(test_freqs)))
+
+        for test_freq_idx, test_freq in enumerate(test_freqs):
+            # Create the design matrix of base frequency and requested harmonics
+            Y = np.column_stack(
+                [
+                    fn(2.0 * np.pi * k * test_freq * t)
+                    for k in range(1, self.settings.harmonics + 1)
+                    for fn in (np.sin, np.cos)
+                ]
+            )
+
+            for test_idx, arr in enumerate(data_arr):  # iterate over first dim; arr is (filterbank x time x ch)
+                for band_idx, band in enumerate(arr):  # iterate over second dim: arr is (time x ch)
+                    calc_output[test_idx, band_idx, test_freq_idx] = cca_rho_max(band[:max_samp, ...], Y)
+
+        # Combine per-subband canonical correlations using a weighted sum
+        # https://iopscience.iop.org/article/10.1088/1741-2560/12/4/046008
+        freq_weights = (np.arange(1, calc_output.shape[1] + 1) ** -1.25) + 0.25
+        calc_output = ((calc_output**2) * freq_weights[None, :, None]).sum(axis=1)
+
+        if self.settings.softmax_beta != 0:
+            calc_output = calc_softmax(calc_output, axis=-1, beta=self.settings.softmax_beta)
+
+        output = replace(
+            message,
+            data=calc_output.reshape(out_shape),
+            dims=out_dims,
+            axes=out_axes,
+        )
+
+        return output
+
+
+class FBCCA(BaseTransformerUnit[FBCCASettings, AxisArray, AxisArray, FBCCATransformer]):
+    SETTINGS = FBCCASettings
+
+
+class StreamingFBCCASettings(FBCCASettings):
+    """
+    Perform rolling/streaming FBCCA on incoming EEG.
+    Decomposes the input multi-channel timeseries data into multiple sub-bands using a FilterbankDesign Transformer,
+    then accumulates data using Window into short-time observations for analysis using an FBCCA Transformer.
+    """
+
+    window_dur: float = 4.0  # sec
+    window_shift: float = 0.5  # sec
+    window_dim: str = "fbcca_window"
+    filter_bw: float = 7.0  # Hz
+    filter_low: float = 7.0  # Hz
+    trans_bw: float = 2.0  # Hz
+    ripple_db: float = 20.0  # dB
+    subbands: int = 12
+
+
+class StreamingFBCCATransformer(CompositeProcessor[StreamingFBCCASettings, AxisArray, AxisArray]):
+    @staticmethod
+    def _initialize_processors(
+        settings: StreamingFBCCASettings,
+    ) -> dict[str, BaseProcessor | BaseStatefulProcessor]:
+        pipeline = {}
+
+        if settings.filterbank_dim is not None:
+            cut_freqs = (np.arange(settings.subbands + 1) * settings.filter_bw) + settings.filter_low
+            filters = [
+                KaiserFilterSettings(
+                    axis=settings.time_dim,
+                    cutoff=(c - settings.trans_bw, cut_freqs[-1]),
+                    ripple=settings.ripple_db,
+                    width=settings.trans_bw,
+                    pass_zero=False,
+                )
+                for c in cut_freqs[:-1]
+            ]
+
+            pipeline["filterbank"] = FilterbankDesignTransformer(
+                FilterbankDesignSettings(filters=filters, new_axis=settings.filterbank_dim)
+            )
+
+        pipeline["window"] = WindowTransformer(
+            WindowSettings(
+                axis=settings.time_dim,
+                newaxis=settings.window_dim,
+                window_dur=settings.window_dur,
+                window_shift=settings.window_shift,
+                zero_pad_until="shift",
+            )
+        )
+
+        pipeline["fbcca"] = FBCCATransformer(settings)
+
+        return pipeline
+
+
+class StreamingFBCCA(BaseTransformerUnit[StreamingFBCCASettings, AxisArray, AxisArray, StreamingFBCCATransformer]):
+    SETTINGS = StreamingFBCCASettings
+
+
+def cca_rho_max(X: np.ndarray, Y: np.ndarray) -> float:
+    """
+    X: (n_time, n_ch)
+    Y: (n_time, n_ref)  # design matrix for one frequency
+    returns: largest canonical correlation in [0,1]
+    """
+    # Center columns
+    Xc = X - X.mean(axis=0, keepdims=True)
+    Yc = Y - Y.mean(axis=0, keepdims=True)
+
+    # Drop any zero-variance columns to avoid rank issues
+    Xc = Xc[:, Xc.std(axis=0) > 1e-12]
+    Yc = Yc[:, Yc.std(axis=0) > 1e-12]
+    if Xc.size == 0 or Yc.size == 0:
+        return 0.0
+
+    # Orthonormal bases
+    Qx, _ = np.linalg.qr(Xc, mode="reduced")  # (n_time, r_x)
+    Qy, _ = np.linalg.qr(Yc, mode="reduced")  # (n_time, r_y)
+
+    # Canonical correlations are the singular values of Qx^T Qy
+    with np.errstate(divide="ignore", over="ignore", invalid="ignore"):
+        s = np.linalg.svd(Qx.T @ Qy, compute_uv=False)
+    return float(s[0]) if s.size else 0.0
+
+
+def calc_softmax(cv: np.ndarray, axis: int, beta: float = 1.0):
+    # Calculate softmax with shifting to avoid overflow
+    # (https://doi.org/10.1093/imanum/draa038)
+    cv = cv - cv.max(axis=axis, keepdims=True)
+    cv = np.exp(beta * cv)
+    cv = cv / np.sum(cv, axis=axis, keepdims=True)
+    return cv