ezmsg-sigproc 1.2.3__py3-none-any.whl → 1.3.1__py3-none-any.whl

ezmsg/sigproc/sampler.py

@@ -1,81 +1,94 @@
+import asyncio  # Dev/test apparatus
 from collections import deque
 from dataclasses import dataclass, replace, field
 import time
-from typing import Optional, Any, Tuple, List, Union, AsyncGenerator, Generator
+import typing
 
-import ezmsg.core as ez
 import numpy as np
-
+import numpy.typing as npt
+import ezmsg.core as ez
 from ezmsg.util.messages.axisarray import AxisArray, slice_along_axis
 from ezmsg.util.generator import consumer
 
-# Dev/test apparatus
-import asyncio
-
 
 @dataclass(unsafe_hash=True)
 class SampleTriggerMessage:
     timestamp: float = field(default_factory=time.time)
-    period: Optional[Tuple[float, float]] = None
-    value: Any = None
+    """Time of the trigger, in seconds. The Clock depends on the input but defaults to time.time"""
+
+    period: typing.Optional[typing.Tuple[float, float]] = None
+    """The period around the timestamp, in seconds"""
+
+    value: typing.Any = None
+    """A value or 'label' associated with the trigger."""
 
 
 @dataclass
 class SampleMessage:
     trigger: SampleTriggerMessage
+    """The time, window, and value (if any) associated with the trigger."""
+
     sample: AxisArray
+    """The data sampled around the trigger."""
 
 
 @consumer
 def sampler(
-        buffer_dur: float,
-        axis: Optional[str] = None,
-        period: Optional[Tuple[float, float]] = None,
-        value: Any = None,
-        estimate_alignment: bool = True
-) -> Generator[Union[AxisArray, SampleTriggerMessage], List[SampleMessage], None]:
+    buffer_dur: float,
+    axis: typing.Optional[str] = None,
+    period: typing.Optional[typing.Tuple[float, float]] = None,
+    value: typing.Any = None,
+    estimate_alignment: bool = True,
+) -> typing.Generator[
+    typing.List[SampleMessage], typing.Union[AxisArray, SampleTriggerMessage], None
+]:
     """
     A generator function that samples data into a buffer, accepts triggers, and returns slices of sampled
     data around the trigger time.
 
-    Parameters:
-    - buffer_dur (float): The duration of the buffer in seconds. The buffer must be long enough to store the oldest
-        sample to be included in a window. e.g., a trigger lagged by 0.5 seconds with a period of (-1.0, +1.5) will
-        need a buffer of 0.5 + (1.5 - -1.0) = 3.0 seconds. It is best to at least double your estimate if memory allows.
-    - axis (Optional[str]): The axis along which to sample the data.
-        None (default) will choose the first axis in the first input.
-    - period (Optional[Tuple[float, float]]): The period in seconds during which to sample the data.
-        Defaults to None. Only used if not None and the trigger message does not define its own period.
-    - value (Any): The value to sample. Defaults to None.
-    - estimate_alignment (bool): Whether to estimate the sample alignment. Defaults to True.
-        If True, the trigger timestamp field is used to slice the buffer.
-        If False, the trigger timestamp is ignored and the next signal's .offset is used.
-        NOTE: For faster-than-realtime playback -- Signals and triggers must share the same (fast) clock for
-        estimate_alignment to operate correctly.
-
-    Sends:
-    - AxisArray containing streaming data messages
-    - SampleTriggerMessage containing a trigger
-    Yields:
-    - list[SampleMessage]: The list of sample messages.
+    Args:
+        buffer_dur: The duration of the buffer in seconds. The buffer must be long enough to store the oldest
+            sample to be included in a window. e.g., a trigger lagged by 0.5 seconds with a period of (-1.0, +1.5) will
+            need a buffer of 0.5 + (1.5 - -1.0) = 3.0 seconds. It is best to at least double your estimate if memory allows.
+        axis: The axis along which to sample the data.
+            None (default) will choose the first axis in the first input.
+        period: The period in seconds during which to sample the data.
+            Defaults to None. Only used if not None and the trigger message does not define its own period.
+        value: The value to sample. Defaults to None.
+        estimate_alignment: Whether to estimate the sample alignment. Defaults to True.
+            If True, the trigger timestamp field is used to slice the buffer.
+            If False, the trigger timestamp is ignored and the next signal's .offset is used.
+            NOTE: For faster-than-realtime playback -- Signals and triggers must share the same (fast) clock for
+            estimate_alignment to operate correctly.
+
+    Returns:
+        A generator that expects `.send` either an :obj:`AxisArray` containing streaming data messages,
+        or a :obj:`SampleTriggerMessage` containing a trigger, and yields the list of :obj:`SampleMessage` s.
     """
-    msg_in = None
-    msg_out: Optional[list[SampleMessage]] = None
+    msg_out: list[SampleMessage] = []
 
     # State variables (most shared between trigger- and data-processing.
     triggers: deque[SampleTriggerMessage] = deque()
-    last_msg_stats = None
-    buffer = None
+    buffer: typing.Optional[npt.NDArray] = None
+    n_samples: int = 0
+    offset: float = 0.0
+
+    check_inputs = {
+        "fs": None,  # Also a state variable
+        "key": None,
+        "shape": None,
+    }
 
     while True:
         msg_in = yield msg_out
         msg_out = []
+
         if isinstance(msg_in, SampleTriggerMessage):
-            if last_msg_stats is None or buffer is None:
+            # Input is a trigger message that we will use to sample the buffer.
+
+            if buffer is None or check_inputs["fs"] is None:
                 # We've yet to see any data; drop the trigger.
                 continue
-            fs = last_msg_stats["fs"]
-            axis_idx = last_msg_stats["axis_idx"]
 
             _period = msg_in.period if msg_in.period is not None else period
             _value = msg_in.value if msg_in.value is not None else value
@@ -86,50 +99,73 @@ def sampler(
 
             # Check that period is valid
             if _period[0] >= _period[1]:
-                ez.logger.warning(f"Sampling failed: invalid period requested ({_period})")
+                ez.logger.warning(
+                    f"Sampling failed: invalid period requested ({_period})"
+                )
                 continue
 
             # Check that period is compatible with buffer duration.
-            max_buf_len = int(buffer_dur * fs)
-            req_buf_len = int((_period[1] - _period[0]) * fs)
+            max_buf_len = int(np.round(buffer_dur * check_inputs["fs"]))
+            req_buf_len = int(np.round((_period[1] - _period[0]) * check_inputs["fs"]))
             if req_buf_len >= max_buf_len:
-                ez.logger.warning(
-                    f"Sampling failed: {period=} >= {buffer_dur=}"
-                )
+                ez.logger.warning(f"Sampling failed: {period=} >= {buffer_dur=}")
                 continue
 
             trigger_ts: float = msg_in.timestamp
             if not estimate_alignment:
                 # Override the trigger timestamp with the next sample's likely timestamp.
-                trigger_ts = last_msg_stats["offset"] + (last_msg_stats["n_samples"] + 1) / fs
+                trigger_ts = offset + (n_samples + 1) / check_inputs["fs"]
 
-            new_trig_msg = replace(msg_in, timestamp=trigger_ts, period=_period, value=_value)
+            new_trig_msg = replace(
+                msg_in, timestamp=trigger_ts, period=_period, value=_value
+            )
             triggers.append(new_trig_msg)
 
         elif isinstance(msg_in, AxisArray):
-            if axis is None:
-                axis = msg_in.dims[0]
+            # Get properties from message
+            axis = axis or msg_in.dims[0]
             axis_idx = msg_in.get_axis_idx(axis)
             axis_info = msg_in.get_axis(axis)
             fs = 1.0 / axis_info.gain
-            sample_shape = msg_in.data.shape[:axis_idx] + msg_in.data.shape[axis_idx + 1:]
-
-            # If the signal properties have changed in a breaking way then reset buffer and triggers.
-            if last_msg_stats is None or fs != last_msg_stats["fs"] or sample_shape != last_msg_stats["sample_shape"]:
-                last_msg_stats = {
-                    "fs": fs,
-                    "sample_shape": sample_shape,
-                    "axis_idx": axis_idx,
-                    "n_samples": msg_in.data.shape[axis_idx]
-                }
+            sample_shape = (
+                msg_in.data.shape[:axis_idx] + msg_in.data.shape[axis_idx + 1 :]
+            )
+
+            # TODO: We could accommodate change in dim order.
+            # if axis_idx != check_inputs["axis_idx"]:
+            #     msg_in = replace(
+            #         msg_in,
+            #         data=np.moveaxis(msg_in.data, axis_idx, check_inputs["axis_idx"]),
+            #         dims=TODO...
+            #     )
+            #    axis_idx = check_inputs["axis_idx"]
+
+            # If the properties have changed in a breaking way then reset buffer and triggers.
+            b_reset = fs != check_inputs["fs"]
+            b_reset = b_reset or sample_shape != check_inputs["shape"]
+            # TODO: Skip next line if we do np.moveaxis above
+            b_reset = b_reset or axis_idx != check_inputs["axis_idx"]
+            b_reset = b_reset or msg_in.key != check_inputs["key"]
+            if b_reset:
+                check_inputs["fs"] = fs
+                check_inputs["shape"] = sample_shape
+                check_inputs["axis_idx"] = axis_idx
+                check_inputs["key"] = msg_in.key
+                n_samples = msg_in.data.shape[axis_idx]
                 buffer = None
                 if len(triggers) > 0:
                     ez.logger.warning("Data stream changed: Discarding all triggers")
                 triggers.clear()
-            last_msg_stats["offset"] = axis_info.offset  # Should be updated on every message.
+
+            # Save some info for trigger processing
+            offset = axis_info.offset
 
             # Update buffer
-            buffer = msg_in.data if buffer is None else np.concatenate((buffer, msg_in.data), axis=axis_idx)
+            buffer = (
+                msg_in.data
+                if buffer is None
+                else np.concatenate((buffer, msg_in.data), axis=axis_idx)
+            )
 
             # Calculate timestamps associated with buffer.
             buffer_offset = np.arange(buffer.shape[axis_idx], dtype=float)
@@ -153,11 +189,10 @@ def sampler(
                     )
                     triggers.remove(trig)
 
-                # TODO: Speed up with searchsorted?
                 t_start = trig.timestamp + trig.period[0]
                 if t_start >= buffer_offset[0]:
                     start = np.searchsorted(buffer_offset, t_start)
-                    stop = start + int(fs * (trig.period[1] - trig.period[0]))
+                    stop = start + int(np.round(fs * (trig.period[1] - trig.period[0])))
                     if buffer.shape[axis_idx] > stop:
                         # Trigger period fully enclosed in buffer.
                         msg_out.append(
@@ -165,9 +200,16 @@ def sampler(
                                 trigger=trig,
                                 sample=replace(
                                     msg_in,
-                                    data=slice_along_axis(buffer, slice(start, stop), axis_idx),
-                                    axes={**msg_in.axes, axis: replace(axis_info, offset=buffer_offset[start])}
-                                )
+                                    data=slice_along_axis(
+                                        buffer, slice(start, stop), axis_idx
+                                    ),
+                                    axes={
+                                        **msg_in.axes,
+                                        axis: replace(
+                                            axis_info, offset=buffer_offset[start]
+                                        ),
+                                    },
+                                ),
                             )
                         )
                         triggers.remove(trig)
@@ -177,29 +219,40 @@ def sampler(
 
 
 class SamplerSettings(ez.Settings):
+    """
+    Settings for :obj:`Sampler`.
+    See :obj:`sampler` for a description of the fields.
+    """
+
     buffer_dur: float
-    axis: Optional[str] = None
-    period: Optional[
-        Tuple[float, float]
-    ] = None  # Optional default period if unspecified in SampleTriggerMessage
-    value: Any = None  # Optional default value if unspecified in SampleTriggerMessage
+    axis: typing.Optional[str] = None
+    period: typing.Optional[typing.Tuple[float, float]] = None
+    """Optional default period if unspecified in SampleTriggerMessage"""
+
+    value: typing.Any = None
+    """Optional default value if unspecified in SampleTriggerMessage"""
 
     estimate_alignment: bool = True
-    # If true, use message timestamp fields and reported sampling rate to estimate
-    # sample-accurate alignment for samples.
-    # If false, sampling will be limited to incoming message rate -- "Block timing"
-    # NOTE: For faster-than-realtime playback --  Incoming timestamps must reflect
-    # "realtime" operation for estimate_alignment to operate correctly.
+    """
+    If true, use message timestamp fields and reported sampling rate to estimate sample-accurate alignment for samples.
+    If false, sampling will be limited to incoming message rate -- "Block timing"
+    NOTE: For faster-than-realtime playback --  Incoming timestamps must reflect
+    "realtime" operation for estimate_alignment to operate correctly.
+    """
 
 
 class SamplerState(ez.State):
     cur_settings: SamplerSettings
-    gen: Generator[Union[AxisArray, SampleTriggerMessage], List[SampleMessage], None]
+    gen: typing.Generator[
+        typing.Union[AxisArray, SampleTriggerMessage], typing.List[SampleMessage], None
+    ]
 
 
 class Sampler(ez.Unit):
-    SETTINGS: SamplerSettings
-    STATE: SamplerState
+    """An :obj:`Unit` for :obj:`sampler`."""
+
+    SETTINGS = SamplerSettings
+    STATE = SamplerState
 
     INPUT_TRIGGER = ez.InputStream(SampleTriggerMessage)
     INPUT_SETTINGS = ez.InputStream(SamplerSettings)
@@ -212,10 +265,10 @@ class Sampler(ez.Unit):
             axis=self.STATE.cur_settings.axis,
             period=self.STATE.cur_settings.period,
             value=self.STATE.cur_settings.value,
-            estimate_alignment=self.STATE.cur_settings.estimate_alignment
+            estimate_alignment=self.STATE.cur_settings.estimate_alignment,
         )
 
-    def initialize(self) -> None:
+    async def initialize(self) -> None:
         self.STATE.cur_settings = self.SETTINGS
         self.construct_generator()
 
@@ -228,27 +281,36 @@ class Sampler(ez.Unit):
     async def on_trigger(self, msg: SampleTriggerMessage) -> None:
         _ = self.STATE.gen.send(msg)
 
-    @ez.subscriber(INPUT_SIGNAL)
+    @ez.subscriber(INPUT_SIGNAL, zero_copy=True)
     @ez.publisher(OUTPUT_SAMPLE)
-    async def on_signal(self, msg: AxisArray) -> AsyncGenerator:
+    async def on_signal(self, msg: AxisArray) -> typing.AsyncGenerator:
         pub_samples = self.STATE.gen.send(msg)
         for sample in pub_samples:
             yield self.OUTPUT_SAMPLE, sample
 
 
 class TriggerGeneratorSettings(ez.Settings):
-    period: Tuple[float, float]  # sec
-    prewait: float = 0.5  # sec
-    publish_period: float = 5.0  # sec
+    period: typing.Tuple[float, float]
+    """The period around the trigger event."""
+
+    prewait: float = 0.5
+    """The time before the first trigger (sec)"""
+
+    publish_period: float = 5.0
+    """The period between triggers (sec)"""
 
 
 class TriggerGenerator(ez.Unit):
-    SETTINGS: TriggerGeneratorSettings
+    """
+    A unit to generate triggers every `publish_period` interval.
+    """
+
+    SETTINGS = TriggerGeneratorSettings
 
     OUTPUT_TRIGGER = ez.OutputStream(SampleTriggerMessage)
 
     @ez.publisher(OUTPUT_TRIGGER)
-    async def generate(self) -> AsyncGenerator:
+    async def generate(self) -> typing.AsyncGenerator:
         await asyncio.sleep(self.SETTINGS.prewait)
 
         output = 0

ezmsg/sigproc/scaler.py

@@ -1,11 +1,13 @@
 from dataclasses import replace
-from typing import Generator, Optional
+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, GenAxisArray
+from ezmsg.util.generator import consumer
+
+from .base import GenAxisArray
 
 
 def _tau_from_alpha(alpha: float, dt: float) -> float:
@@ -27,24 +29,39 @@ def _alpha_from_tau(tau: float, dt: float) -> float:
 
 
 @consumer
-def scaler(time_constant: float = 1.0, axis: Optional[str] = None) -> Generator[AxisArray, AxisArray, None]:
+def scaler(
+    time_constant: float = 1.0, axis: typing.Optional[str] = None
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Create a generator function that applies the
+    adaptive standard scaler from https://riverml.xyz/latest/api/preprocessing/AdaptiveStandardScaler/
+    This is faster than :obj:`scaler_np` for single-channel data.
+
+    Args:
+        time_constant: Decay constant `tau` in seconds.
+        axis: The name of the axis to accumulate statistics over.
+
+    Returns:
+        A primed generator object that expects `.send(axis_array)` and yields a
+        standardized, or "Z-scored" version of the input.
+    """
     from river import preprocessing
-    axis_arr_in = AxisArray(np.array([]), dims=[""])
-    axis_arr_out = AxisArray(np.array([]), dims=[""])
+
+    msg_out = AxisArray(np.array([]), dims=[""])
     _scaler = None
     while True:
-        axis_arr_in = yield axis_arr_out
-        data = axis_arr_in.data
+        msg_in: AxisArray = yield msg_out
+        data = msg_in.data
         if axis is None:
-            axis = axis_arr_in.dims[0]
+            axis = msg_in.dims[0]
             axis_idx = 0
         else:
-            axis_idx = axis_arr_in.get_axis_idx(axis)
+            axis_idx = msg_in.get_axis_idx(axis)
             if axis_idx != 0:
                 data = np.moveaxis(data, axis_idx, 0)
 
         if _scaler is None:
-            alpha = _alpha_from_tau(time_constant, axis_arr_in.axes[axis].gain)
+            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
             _scaler = preprocessing.AdaptiveStandardScaler(fading_factor=alpha)
 
         result = []
@@ -57,17 +74,39 @@ def scaler(time_constant: float = 1.0, axis: Optional[str] = None) -> Generator[
 
         result = np.stack(result)
         result = np.moveaxis(result, 0, axis_idx)
-        axis_arr_out = replace(axis_arr_in, data=result)
+        msg_out = replace(msg_in, data=result)
 
 
 @consumer
-def scaler_np(time_constant: float = 1.0, axis: Optional[str] = None) -> Generator[AxisArray, AxisArray, None]:
-    # The only dependency is numpy.
-    # This is faster for multi-channel data but slower for single-channel data.
-    axis_arr_in = AxisArray(np.array([]), dims=[""])
-    axis_arr_out = AxisArray(np.array([]), dims=[""])
-    means = vars_means = vars_sq_means = None
-    alpha = None
+def scaler_np(
+    time_constant: float = 1.0, axis: typing.Optional[str] = None
+) -> typing.Generator[AxisArray, AxisArray, None]:
+    """
+    Create a generator function that applies an adaptive standard scaler.
+    This is faster than :obj:`scaler` for multichannel data.
+
+    Args:
+        time_constant: Decay constant `tau` in seconds.
+        axis: The name of the axis to accumulate statistics over.
+
+    Returns:
+        A primed generator object that expects `.send(axis_array)` and yields a
+        standardized, or "Z-scored" version of the input.
+    """
+    msg_out = AxisArray(np.array([]), dims=[""])
+
+    # State variables
+    alpha: float = 0.0
+    means: typing.Optional[npt.NDArray] = None
+    vars_means: typing.Optional[npt.NDArray] = None
+    vars_sq_means: typing.Optional[npt.NDArray] = None
+
+    # Reset if input changes
+    check_input = {
+        "gain": None,  # Resets alpha
+        "shape": None,
+        "key": None,  # Key change implies buffered means/vars are invalid.
+    }
 
     def _ew_update(arr, prev, _alpha):
         if np.all(prev == 0):
@@ -77,51 +116,58 @@ def scaler_np(time_constant: float = 1.0, axis: Optional[str] = None) -> Generat
         return prev + _alpha * (arr - prev)
 
     while True:
-        axis_arr_in = yield axis_arr_out
+        msg_in: AxisArray = yield msg_out
 
-        data = axis_arr_in.data
-        if axis is None:
-            axis = axis_arr_in.dims[0]
-            axis_idx = 0
-        else:
-            axis_idx = axis_arr_in.get_axis_idx(axis)
-        data = np.moveaxis(data, axis_idx, 0)
+        axis = axis or msg_in.dims[0]
+        axis_idx = msg_in.get_axis_idx(axis)
 
-        if alpha is None:
-            alpha = _alpha_from_tau(time_constant, axis_arr_in.axes[axis].gain)
+        if msg_in.axes[axis].gain != check_input["gain"]:
+            alpha = _alpha_from_tau(time_constant, msg_in.axes[axis].gain)
+            check_input["gain"] = msg_in.axes[axis].gain
 
-        if means is None or means.shape != data[0].shape:
+        data: npt.NDArray = np.moveaxis(msg_in.data, axis_idx, 0)
+        b_reset = data.shape[1:] != check_input["shape"]
+        b_reset |= msg_in.key != check_input["key"]
+        if b_reset:
+            check_input["shape"] = data.shape[1:]
+            check_input["key"] = msg_in.key
             vars_sq_means = np.zeros_like(data[0], dtype=float)
             vars_means = np.zeros_like(data[0], dtype=float)
             means = np.zeros_like(data[0], dtype=float)
 
-        result = []
-        for sample in data:
+        result = np.zeros_like(data)
+        for sample_ix in range(data.shape[0]):
+            sample = data[sample_ix]
             # Update step
             vars_means = _ew_update(sample, vars_means, alpha)
             vars_sq_means = _ew_update(sample**2, vars_sq_means, alpha)
             means = _ew_update(sample, means, alpha)
             # Get step
-            varis = vars_sq_means - vars_means ** 2
-            y = ((sample - means) / (varis**0.5))
-            y[np.isnan(y)] = 0.0
-            result.append(y)
+            varis = vars_sq_means - vars_means**2
+            y = (sample - means) / (varis**0.5)
+            result[sample_ix] = y
 
-        result = np.stack(result, axis=0)
+        result[np.isnan(result)] = 0.0
         result = np.moveaxis(result, 0, axis_idx)
-        axis_arr_out = replace(axis_arr_in, data=result)
+        msg_out = replace(msg_in, data=result)
 
 
 class AdaptiveStandardScalerSettings(ez.Settings):
+    """
+    Settings for :obj:`AdaptiveStandardScaler`.
+    See :obj:`scaler_np` for a description of the parameters.
+    """
+
     time_constant: float = 1.0
-    axis: Optional[str] = None
+    axis: typing.Optional[str] = None
 
 
 class AdaptiveStandardScaler(GenAxisArray):
-    SETTINGS: AdaptiveStandardScalerSettings
+    """Unit for :obj:`scaler_np`"""
+
+    SETTINGS = AdaptiveStandardScalerSettings
 
     def construct_generator(self):
         self.STATE.gen = scaler_np(
-            time_constant=self.SETTINGS.time_constant,
-            axis=self.SETTINGS.axis
+            time_constant=self.SETTINGS.time_constant, axis=self.SETTINGS.axis
         )

ezmsg/sigproc/signalinjector.py

@@ -1,17 +1,15 @@
+from dataclasses import replace
 import typing
 
 import ezmsg.core as ez
-
-from dataclasses import replace
 from ezmsg.util.messages.axisarray import AxisArray
-
 import numpy as np
 import numpy.typing as npt
 
 
 class SignalInjectorSettings(ez.Settings):
-    time_dim: str = 'time' # Input signal needs a time dimension with units in sec.
-    frequency: typing.Optional[float] = None # Hz
+    time_dim: str = "time"  # Input signal needs a time dimension with units in sec.
+    frequency: typing.Optional[float] = None  # Hz
     amplitude: float = 1.0
     mixing_seed: typing.Optional[int] = None
 
@@ -24,8 +22,8 @@ class SignalInjectorState(ez.State):
 
 
 class SignalInjector(ez.Unit):
-    SETTINGS: SignalInjectorSettings
-    STATE: SignalInjectorState
+    SETTINGS = SignalInjectorSettings
+    STATE = SignalInjectorState
 
     INPUT_FREQUENCY = ez.InputStream(typing.Optional[float])
     INPUT_AMPLITUDE = ez.InputStream(float)
@@ -48,7 +46,6 @@ class SignalInjector(ez.Unit):
     @ez.subscriber(INPUT_SIGNAL)
     @ez.publisher(OUTPUT_SIGNAL)
     async def inject(self, msg: AxisArray) -> typing.AsyncGenerator:
-
         if self.STATE.cur_shape != msg.shape:
             self.STATE.cur_shape = msg.shape
             rng = np.random.default_rng(self.SETTINGS.mixing_seed)
@@ -58,10 +55,10 @@ class SignalInjector(ez.Unit):
         if self.STATE.cur_frequency is None:
             yield self.OUTPUT_SIGNAL, msg
         else:
-            out_msg = replace(msg, data = msg.data.copy())
+            out_msg = replace(msg, data=msg.data.copy())
             t = out_msg.ax(self.SETTINGS.time_dim).values[..., np.newaxis]
             signal = np.sin(2 * np.pi * self.STATE.cur_frequency * t)
             mixed_signal = signal * self.STATE.mixing * self.STATE.cur_amplitude
             with out_msg.view2d(self.SETTINGS.time_dim) as view:
                 view[...] = view + mixed_signal.astype(view.dtype)
-            yield self.OUTPUT_SIGNAL, out_msg
+            yield self.OUTPUT_SIGNAL, out_msg