ezmsg-sigproc 1.2.2__py3-none-any.whl → 1.3.1__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,16 @@
-__version__ = "1.2.2"
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '1.3.1'
+__version_tuple__ = version_tuple = (1, 3, 1)

ezmsg/sigproc/activation.py

@@ -0,0 +1,75 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import scipy.special
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
+
+from .spectral import OptionsEnum
+from .base import GenAxisArray
+
+
+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,
+}
+
+
+@consumer
+def activation(
+    function: typing.Union[str, ActivationFunction],
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    if type(function) is ActivationFunction:
+        func = ACTIVATIONS[function]
+    else:
+        # str type. There's probably an easier way to support either enum or str argument. Oh well this works.
+        function: str = 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]
+
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+        msg_out = replace(msg_in, data=func(msg_in.data))
+
+
+class ActivationSettings(ez.Settings):
+    function: str = ActivationFunction.NONE
+
+
+class Activation(GenAxisArray):
+    SETTINGS = ActivationSettings
+
+    def construct_generator(self):
+        self.STATE.gen = activation(function=self.SETTINGS.function)

ezmsg/sigproc/affinetransform.py

@@ -0,0 +1,234 @@
+from dataclasses import replace
+import os
+from pathlib import Path
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer
+
+from .base import GenAxisArray
+
+
+@consumer
+def affine_transform(
+    weights: typing.Union[np.ndarray, str, Path],
+    axis: typing.Optional[str] = None,
+    right_multiply: bool = True,
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    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 tranpose the weights before applying.
+
+    Returns:
+        A primed generator object that yields an :obj:`AxisArray` object for every
+        :obj:`AxisArray` it receives via `send`.
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # Check parameters
+    if isinstance(weights, str):
+        if weights == "passthrough":
+            weights = None
+        else:
+            weights = Path(os.path.abspath(os.path.expanduser(weights)))
+    if isinstance(weights, Path):
+        weights = np.loadtxt(weights, delimiter=",")
+    if not right_multiply:
+        weights = weights.T
+    if weights is not None:
+        weights = np.ascontiguousarray(weights)
+
+    # State variables
+    # New axis with transformed labels, if required
+    new_axis: typing.Optional[AxisArray.Axis] = None
+
+    # Reset if any of these change.
+    check_input = {"key": None}
+    # We assume key change catches labels change; we don't want to check labels every message
+    # We don't need to check if input size has changed because weights multiplication will fail if so.
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        if weights is None:
+            msg_out = msg_in
+            continue
+
+        axis = axis or msg_in.dims[-1]  # Note: Most nodes default do dim[0]
+        axis_idx = msg_in.get_axis_idx(axis)
+
+        b_reset = msg_in.key != check_input["key"]
+        if b_reset:
+            # First sample or key has changed. Reset the state.
+            check_input["key"] = msg_in.key
+            # Determine if we need to modify the transformed axis.
+            if (
+                axis in msg_in.axes
+                and hasattr(msg_in.axes[axis], "labels")
+                and weights.shape[0] != weights.shape[1]
+            ):
+                in_labels = msg_in.axes[axis].labels
+                new_labels = []
+                n_in = weights.shape[1 if right_multiply else 0]
+                n_out = weights.shape[0 if right_multiply else 1]
+                if len(in_labels) != n_in:
+                    # Something upstream did something it wasn't supposed to. We will drop the labels.
+                    ez.logger.warning(
+                        f"Received {len(in_labels)} for {n_in} inputs. Check upstream labels."
+                    )
+                else:
+                    b_used_inputs = np.any(weights, axis=0 if right_multiply else 1)
+                    b_filled_outputs = np.any(weights, axis=1 if right_multiply else 0)
+                    if np.all(b_used_inputs) and np.all(b_filled_outputs):
+                        # All inputs are used and all outputs are used, but n_in != n_out.
+                        # Mapping cannot be determined.
+                        new_labels = []
+                    elif np.all(b_used_inputs):
+                        # Strange scenario: New outputs are filled with empty data.
+                        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):
+                        # Transform is dropping some of the inputs.
+                        new_labels = np.array(in_labels)[b_used_inputs].tolist()
+                new_axis = replace(msg_in.axes[axis], labels=new_labels)
+
+        data = msg_in.data
+
+        if data.shape[axis_idx] == (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(msg_in.dims) - 1]:
+            data = np.matmul(data, weights)
+        else:
+            data = np.moveaxis(data, axis_idx, -1)
+            data = np.matmul(data, weights)
+            data = np.moveaxis(data, -1, axis_idx)
+
+        replace_kwargs = {"data": data}
+        if new_axis is not None:
+            replace_kwargs["axes"] = {**msg_in.axes, axis: new_axis}
+        msg_out = replace(msg_in, **replace_kwargs)
+
+
+class AffineTransformSettings(ez.Settings):
+    """
+    Settings for :obj:`AffineTransform`.
+    See :obj:`affine_transform` for argument details.
+    """
+
+    weights: typing.Union[np.ndarray, str, Path]
+    axis: typing.Optional[str] = None
+    right_multiply: bool = True
+
+
+class AffineTransform(GenAxisArray):
+    """:obj:`Unit` for :obj:`affine_transform`"""
+
+    SETTINGS = AffineTransformSettings
+
+    def construct_generator(self):
+        self.STATE.gen = affine_transform(
+            weights=self.SETTINGS.weights,
+            axis=self.SETTINGS.axis,
+            right_multiply=self.SETTINGS.right_multiply,
+        )
+
+
+def zeros_for_noop(data: npt.NDArray, **ignore_kwargs) -> npt.NDArray:
+    return np.zeros_like(data)
+
+
+@consumer
+def common_rereference(
+    mode: str = "mean", axis: typing.Optional[str] = None, include_current: bool = True
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    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:
+        A primed generator object that yields an :obj:`AxisArray` object
+        for every :obj:`AxisArray` it receives via `send`.
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    if mode == "passthrough":
+        include_current = True
+
+    func = {"mean": np.mean, "median": np.median, "passthrough": zeros_for_noop}[mode]
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+
+        if axis is None:
+            axis = msg_in.dims[-1]
+            axis_idx = -1
+        else:
+            axis_idx = msg_in.get_axis_idx(axis)
+
+        ref_data = func(msg_in.data, axis=axis_idx, keepdims=True)
+
+        if not 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 = msg_in.data.shape[axis_idx]
+            ref_data = (N / (N - 1)) * ref_data - msg_in.data / (N - 1)
+            # Side note: I profiled using affine_transform and it's about 30x slower than this implementation.
+
+        msg_out = replace(msg_in, data=msg_in.data - ref_data)
+
+
+class CommonRereferenceSettings(ez.Settings):
+    """
+    Settings for :obj:`CommonRereference`
+    See :obj:`common_rereference` for argument details.
+    """
+
+    mode: str = "mean"
+    axis: typing.Optional[str] = None
+    include_current: bool = True
+
+
+class CommonRereference(GenAxisArray):
+    """
+    :obj:`Unit` for :obj:`common_rereference`.
+    """
+
+    SETTINGS = CommonRereferenceSettings
+
+    def construct_generator(self):
+        self.STATE.gen = common_rereference(
+            mode=self.SETTINGS.mode,
+            axis=self.SETTINGS.axis,
+            include_current=self.SETTINGS.include_current,
+        )

ezmsg/sigproc/aggregate.py

@@ -0,0 +1,158 @@
+from dataclasses import replace
+import typing
+
+import numpy as np
+import numpy.typing as npt
+import ezmsg.core as ez
+from ezmsg.util.generator import consumer
+from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
+
+from .spectral import OptionsEnum
+from .base import GenAxisArray
+
+
+class AggregationFunction(OptionsEnum):
+    """Enum for aggregation functions available to be used in :obj:`ranged_aggregate` operation."""
+
+    NONE = "None (all)"
+    MAX = "max"
+    MIN = "min"
+    MEAN = "mean"
+    MEDIAN = "median"
+    STD = "std"
+    NANMAX = "nanmax"
+    NANMIN = "nanmin"
+    NANMEAN = "nanmean"
+    NANMEDIAN = "nanmedian"
+    NANSTD = "nanstd"
+    ARGMIN = "argmin"
+    ARGMAX = "argmax"
+
+
+AGGREGATORS = {
+    AggregationFunction.NONE: np.all,
+    AggregationFunction.MAX: np.max,
+    AggregationFunction.MIN: np.min,
+    AggregationFunction.MEAN: np.mean,
+    AggregationFunction.MEDIAN: np.median,
+    AggregationFunction.STD: np.std,
+    AggregationFunction.NANMAX: np.nanmax,
+    AggregationFunction.NANMIN: np.nanmin,
+    AggregationFunction.NANMEAN: np.nanmean,
+    AggregationFunction.NANMEDIAN: np.nanmedian,
+    AggregationFunction.NANSTD: np.nanstd,
+    AggregationFunction.ARGMIN: np.argmin,
+    AggregationFunction.ARGMAX: np.argmax,
+}
+
+
+@consumer
+def ranged_aggregate(
+    axis: typing.Optional[str] = None,
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = None,
+    operation: AggregationFunction = AggregationFunction.MEAN,
+):
+    """
+    Apply an aggregation operation over one or more bands.
+
+    Args:
+        axis: The name of the axis along which to apply the bands.
+        bands: [(band1_min, band1_max), (band2_min, band2_max), ...]
+            If not set then this acts as a passthrough node.
+        operation: :obj:`AggregationFunction` to apply to each band.
+
+    Returns:
+        A primed generator object ready to yield an AxisArray for each .send(axis_array)
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # State variables
+    slices: typing.Optional[typing.List[typing.Tuple[typing.Any, ...]]] = None
+    out_axis: typing.Optional[AxisArray.Axis] = None
+    ax_vec: typing.Optional[npt.NDArray] = None
+
+    # Reset if any of these changes. Key not checked because continuity between chunks not required.
+    check_inputs = {"gain": None, "offset": None}
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+        if bands is None:
+            msg_out = msg_in
+        else:
+            axis = axis or msg_in.dims[0]
+            target_axis = msg_in.get_axis(axis)
+
+            b_reset = target_axis.gain != check_inputs["gain"]
+            b_reset = b_reset or target_axis.offset != check_inputs["offset"]
+            if b_reset:
+                check_inputs["gain"] = target_axis.gain
+                check_inputs["offset"] = target_axis.offset
+
+                # If the axis we are operating on has changed (e.g., "time" or "win" axis always changes),
+                #  or the key has changed, then recalculate slices.
+
+                ax_idx = msg_in.get_axis_idx(axis)
+
+                ax_vec = (
+                    target_axis.offset
+                    + np.arange(msg_in.data.shape[ax_idx]) * target_axis.gain
+                )
+                slices = []
+                mids = []
+                for start, stop in bands:
+                    inds = np.where(np.logical_and(ax_vec >= start, ax_vec <= stop))[0]
+                    mids.append(np.mean(inds) * target_axis.gain + target_axis.offset)
+                    slices.append(np.s_[inds[0] : inds[-1] + 1])
+                out_ax_kwargs = {
+                    "unit": target_axis.unit,
+                    "offset": mids[0],
+                    "gain": (mids[1] - mids[0]) if len(mids) > 1 else 1.0,
+                }
+                if hasattr(target_axis, "labels"):
+                    out_ax_kwargs["labels"] = [f"{_[0]} - {_[1]}" for _ in bands]
+                out_axis = replace(target_axis, **out_ax_kwargs)
+
+            agg_func = AGGREGATORS[operation]
+            out_data = [
+                agg_func(slice_along_axis(msg_in.data, sl, axis=ax_idx), axis=ax_idx)
+                for sl in slices
+            ]
+
+            msg_out = replace(
+                msg_in,
+                data=np.stack(out_data, axis=ax_idx),
+                axes={**msg_in.axes, axis: out_axis},
+            )
+            if operation in [AggregationFunction.ARGMIN, AggregationFunction.ARGMAX]:
+                # Convert indices returned by argmin/argmax into the value along the axis.
+                out_data = []
+                for sl_ix, sl in enumerate(slices):
+                    offsets = np.take(msg_out.data, [sl_ix], axis=ax_idx)
+                    out_data.append(ax_vec[sl][offsets])
+                msg_out.data = np.concatenate(out_data, axis=ax_idx)
+
+
+class RangedAggregateSettings(ez.Settings):
+    """
+    Settings for ``RangedAggregate``.
+    See :obj:`ranged_aggregate` for details.
+    """
+
+    axis: typing.Optional[str] = None
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = None
+    operation: AggregationFunction = AggregationFunction.MEAN
+
+
+class RangedAggregate(GenAxisArray):
+    """
+    Unit for :obj:`ranged_aggregate`
+    """
+
+    SETTINGS = RangedAggregateSettings
+
+    def construct_generator(self):
+        self.STATE.gen = ranged_aggregate(
+            axis=self.SETTINGS.axis,
+            bands=self.SETTINGS.bands,
+            operation=self.SETTINGS.operation,
+        )

ezmsg/sigproc/bandpower.py

@@ -0,0 +1,74 @@
+from dataclasses import field
+import typing
+
+import numpy as np
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import consumer, compose
+
+from .spectrogram import spectrogram, SpectrogramSettings
+from .aggregate import ranged_aggregate, AggregationFunction
+from .base import GenAxisArray
+
+
+@consumer
+def bandpower(
+    spectrogram_settings: SpectrogramSettings,
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = [
+        (17, 30),
+        (70, 170),
+    ],
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Calculate the average spectral power in each band.
+
+    Args:
+        spectrogram_settings: Settings for spectrogram calculation.
+        bands: (min, max) tuples of band limits in Hz.
+
+    Returns:
+        A primed generator object ready to yield an AxisArray for each .send(axis_array)
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    f_spec = spectrogram(
+        window_dur=spectrogram_settings.window_dur,
+        window_shift=spectrogram_settings.window_shift,
+        window=spectrogram_settings.window,
+        transform=spectrogram_settings.transform,
+        output=spectrogram_settings.output,
+    )
+    f_agg = ranged_aggregate(
+        axis="freq", bands=bands, operation=AggregationFunction.MEAN
+    )
+    pipeline = compose(f_spec, f_agg)
+
+    while True:
+        msg_in: AxisArray = yield msg_out
+        msg_out = pipeline(msg_in)
+
+
+class BandPowerSettings(ez.Settings):
+    """
+    Settings for ``BandPower``.
+    See :obj:`bandpower` for details.
+    """
+
+    spectrogram_settings: SpectrogramSettings = field(
+        default_factory=SpectrogramSettings
+    )
+    bands: typing.Optional[typing.List[typing.Tuple[float, float]]] = field(
+        default_factory=lambda: [(17, 30), (70, 170)]
+    )
+
+
+class BandPower(GenAxisArray):
+    """:obj:`Unit` for :obj:`bandpower`."""
+
+    SETTINGS = BandPowerSettings
+
+    def construct_generator(self):
+        self.STATE.gen = bandpower(
+            spectrogram_settings=self.SETTINGS.spectrogram_settings,
+            bands=self.SETTINGS.bands,
+        )

ezmsg/sigproc/base.py

@@ -0,0 +1,38 @@
+import traceback
+import typing
+
+import ezmsg.core as ez
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.generator import GenState
+
+
+class GenAxisArray(ez.Unit):
+    STATE = GenState
+
+    INPUT_SIGNAL = ez.InputStream(AxisArray)
+    OUTPUT_SIGNAL = ez.OutputStream(AxisArray)
+    INPUT_SETTINGS = ez.InputStream(ez.Settings)
+
+    async def initialize(self) -> None:
+        self.construct_generator()
+
+    # Method to be implemented by subclasses to construct the specific generator
+    def construct_generator(self):
+        raise NotImplementedError
+
+    @ez.subscriber(INPUT_SETTINGS)
+    async def on_settings(self, msg: ez.Settings) -> None:
+        self.apply_settings(msg)
+        self.construct_generator()
+
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
+    @ez.publisher(OUTPUT_SIGNAL)
+    async def on_signal(self, message: AxisArray) -> typing.AsyncGenerator:
+        try:
+            ret = self.STATE.gen.send(message)
+            if ret.data.size > 0:
+                yield self.OUTPUT_SIGNAL, ret
+        except (StopIteration, GeneratorExit):
+            ez.logger.debug(f"Generator closed in {self.address}")
+        except Exception:
+            ez.logger.info(traceback.format_exc())