ezmsg-sigproc 2.2.0__py3-none-any.whl → 2.3.0__py3-none-any.whl

ezmsg/sigproc/__version__.py

@@ -1,7 +1,14 @@
 # file generated by setuptools-scm
 # don't change, don't track in version control
 
-__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
 
 TYPE_CHECKING = False
 if TYPE_CHECKING:
@@ -9,13 +16,19 @@ if TYPE_CHECKING:
     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.2.0'
-__version_tuple__ = version_tuple = (2, 2, 0)
+__version__ = version = '2.3.0'
+__version_tuple__ = version_tuple = (2, 3, 0)
+
+__commit_id__ = commit_id = None

ezmsg/sigproc/denormalize.py

@@ -0,0 +1,86 @@
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+from ezmsg.sigproc.base import (
+    BaseTransformerUnit,
+    BaseStatefulTransformer,
+    processor_state,
+)
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class DenormalizeSettings(ez.Settings):
+    low_rate: float = 2.0
+    """Low end of probable rate after denormalization (Hz)."""
+
+    high_rate: float = 40.0
+    """High end of probable rate after denormalization (Hz)."""
+
+    distribution: str = "uniform"
+    """Distribution to sample rates from. Options are 'uniform', 'normal', or 'constant'."""
+
+
+@processor_state
+class DenormalizeRateState:
+    gains: npt.NDArray | None = None
+    offsets: npt.NDArray | None = None
+
+
+class DenormalizeTransformer(
+    BaseStatefulTransformer[
+        DenormalizeSettings, AxisArray, AxisArray, DenormalizeRateState
+    ]
+):
+    """
+    Scales data from a normalized distribution (mean=0, std=1) to a denormalized
+    distribution using random per-channel offsets and gains designed to keep the
+    99.9% CIs between 0 and 2x the offset.
+
+    This is useful for simulating realistic firing rates from normalized data.
+    """
+
+    def _reset_state(self, message: AxisArray) -> None:
+        ax_ix = message.get_axis_idx("ch")
+        nch = message.data.shape[ax_ix]
+        arr_size = (nch, 1) if ax_ix == 0 else (1, nch)
+        if self.settings.distribution == "uniform":
+            self.state.offsets = np.random.uniform(2.0, 40.0, size=arr_size)
+        elif self.settings.distribution == "normal":
+            self.state.offsets = np.random.normal(
+                loc=(self.settings.low_rate + self.settings.high_rate) / 2.0,
+                scale=(self.settings.high_rate - self.settings.low_rate) / 6.0,
+                size=arr_size,
+            )
+            self.state.offsets = np.clip(
+                self.state.offsets,
+                a_min=self.settings.low_rate,
+                a_max=self.settings.high_rate,
+            )
+        elif self.settings.distribution == "constant":
+            self.state.offsets = np.full(
+                shape=arr_size,
+                fill_value=(self.settings.low_rate + self.settings.high_rate) / 2.0,
+            )
+        else:
+            raise ValueError(f"Invalid distribution: {self.settings.distribution}")
+        # Input has std == 1
+        # Desired output has range from 0 to 2*self.state.offsets within 99.9% confidence interval
+        # For a standard normal distribution, 99.9% of data is within +/- 3.29 std devs.
+        # So, gain = offset / 3.29 to scale the std dev appropriately.
+        self.state.gains = self.state.offsets / 3.29
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        denorm = message.data * self.state.gains + self.state.offsets
+        return replace(
+            message,
+            data=np.clip(denorm, a_min=0.0, a_max=None),
+        )
+
+
+class DenormalizeRateUnit(
+    BaseTransformerUnit[
+        DenormalizeSettings, AxisArray, AxisArray, DenormalizeTransformer
+    ]
+):
+    SETTINGS = DenormalizeSettings

ezmsg/sigproc/fbcca.py

@@ -0,0 +1,332 @@
+import typing
+import math
+from dataclasses import field
+
+import numpy as np
+
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+from .sampler import SampleTriggerMessage
+from .window import WindowTransformer, WindowSettings
+
+from .base import (
+    BaseTransformer,
+    BaseTransformerUnit,
+    CompositeProcessor,
+    BaseProcessor,
+    BaseStatefulProcessor,
+)
+
+from .kaiser import KaiserFilterSettings
+from .filterbankdesign import (
+    FilterbankDesignSettings,
+    FilterbankDesignTransformer,
+)
+
+
+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

ezmsg/sigproc/filter.py

@@ -263,6 +263,14 @@ class FilterByDesignTransformer(
             axis = self.state.filter.settings.axis
             fs = 1 / message.axes[axis].gain
             coefs = design_fun(fs)
+
+            # Convert BA to SOS if requested
+            if coefs is not None and self.settings.coef_type == "sos":
+                if isinstance(coefs, tuple) and len(coefs) == 2:
+                    # It's BA format, convert to SOS
+                    b, a = coefs
+                    coefs = scipy.signal.tf2sos(b, a)
+
             self.state.filter.update_coefficients(
                 coefs, coef_type=self.settings.coef_type
             )
@@ -282,6 +290,14 @@ class FilterByDesignTransformer(
         axis = message.dims[0] if self.settings.axis is None else self.settings.axis
         fs = 1 / message.axes[axis].gain
         coefs = design_fun(fs)
+
+        # Convert BA to SOS if requested
+        if coefs is not None and self.settings.coef_type == "sos":
+            if isinstance(coefs, tuple) and len(coefs) == 2:
+                # It's BA format, convert to SOS
+                b, a = coefs
+                coefs = scipy.signal.tf2sos(b, a)
+
         new_settings = FilterSettings(
             axis=axis, coef_type=self.settings.coef_type, coefs=coefs
         )

ezmsg/sigproc/filterbankdesign.py

@@ -0,0 +1,136 @@
+import typing
+
+import ezmsg.core as ez
+import numpy as np
+import numpy.typing as npt
+
+from ezmsg.util.messages.util import replace
+from ezmsg.util.messages.axisarray import AxisArray
+
+from .base import (
+    BaseStatefulTransformer,
+    processor_state,
+)
+
+from .filterbank import (
+    FilterbankTransformer,
+    FilterbankSettings,
+    FilterbankMode,
+    MinPhaseMode,
+)
+
+from .kaiser import KaiserFilterSettings, kaiser_design_fun
+
+
+class FilterbankDesignSettings(ez.Settings):
+    filters: typing.Iterable[KaiserFilterSettings]
+
+    mode: FilterbankMode = FilterbankMode.CONV
+    """
+    "conv", "fft", or "auto". If "auto", the mode is determined by the size of the input data.
+      fft mode is more efficient for long kernels. However, fft mode uses non-overlapping windows and will
+      incur a delay equal to the window length, which is larger than the largest kernel.
+      conv mode is less efficient but will return data for every incoming chunk regardless of how small it is
+      and thus can provide shorter latency updates.
+    """
+
+    min_phase: MinPhaseMode = MinPhaseMode.NONE
+    """
+    If not None, convert the kernels to minimum-phase equivalents. Valid options are
+      'hilbert', 'homomorphic', and 'homomorphic-full'. Complex filters not supported.
+      See `scipy.signal.minimum_phase` for details.
+    """
+
+    axis: str = "time"
+    """The name of the axis to operate on. This should usually be "time"."""
+
+    new_axis: str = "kernel"
+    """The name of the new axis corresponding to the kernel index."""
+
+
+@processor_state
+class FilterbankDesignState:
+    filterbank: FilterbankTransformer | None = None
+    needs_redesign: bool = False
+
+
+class FilterbankDesignTransformer(
+    BaseStatefulTransformer[
+        FilterbankDesignSettings, AxisArray, AxisArray, FilterbankDesignState
+    ],
+):
+    """
+    Transformer that designs and applies a filterbank based on Kaiser windowed FIR filters.
+    """
+
+    @classmethod
+    def get_message_type(cls, dir: str) -> type[AxisArray]:
+        if dir in ("in", "out"):
+            return AxisArray
+        else:
+            raise ValueError(f"Invalid direction: {dir}. Must be 'in' or 'out'.")
+
+    def update_settings(
+        self, new_settings: typing.Optional[FilterbankDesignSettings] = None, **kwargs
+    ) -> None:
+        """
+        Update settings and mark that filter coefficients need to be recalculated.
+
+        Args:
+            new_settings: Complete new settings object to replace current settings
+            **kwargs: Individual settings to update
+        """
+        # Update settings
+        if new_settings is not None:
+            self.settings = new_settings
+        else:
+            self.settings = replace(self.settings, **kwargs)
+
+        # Set flag to trigger recalculation on next message
+        if self.state.filterbank is not None:
+            self.state.needs_redesign = True
+
+    def _calculate_kernels(self, fs: float) -> list[npt.NDArray]:
+        kernels = []
+        for filter in self.settings.filters:
+            output = kaiser_design_fun(
+                fs,
+                cutoff=filter.cutoff,
+                ripple=filter.ripple,
+                width=filter.width,
+                pass_zero=filter.pass_zero,
+                wn_hz=filter.wn_hz,
+            )
+
+            kernels.append(np.array([1.0]) if output is None else output[0])
+        return kernels
+
+    def __call__(self, message: AxisArray) -> AxisArray:
+        if self.state.filterbank is not None and self.state.needs_redesign:
+            self._reset_state(message)
+            self.state.needs_redesign = False
+        return super().__call__(message)
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = message.dims[0] if self.settings.axis is None else self.settings.axis
+        gain = message.axes[axis].gain if hasattr(message.axes[axis], "gain") else 1
+        axis_idx = message.get_axis_idx(axis)
+        samp_shape = message.data.shape[:axis_idx] + message.data.shape[axis_idx + 1 :]
+        return hash((message.key, samp_shape, gain))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis_obj = message.axes[self.settings.axis]
+        assert isinstance(axis_obj, AxisArray.LinearAxis)
+        fs = 1 / axis_obj.gain
+        kernels = self._calculate_kernels(fs)
+        new_settings = FilterbankSettings(
+            kernels=kernels,
+            mode=self.settings.mode,
+            min_phase=self.settings.min_phase,
+            axis=self.settings.axis,
+            new_axis=self.settings.new_axis,
+        )
+        self.state.filterbank = FilterbankTransformer(settings=new_settings)
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        return self.state.filterbank(message)