Source code for ezmsg.simbiophys.line_noise

"""Additive mains (line) noise with slow frequency drift.

Adds a sinusoidal 50/60 Hz interference component to every channel of a signal,
emulating power-line pickup on a recording. The peak frequency wanders slowly
(a bounded random walk) to mimic drift of the recording device's clock relative
to the mains. When ``freq`` is None the stage is a pass-through.

The drifting sinusoid is generated by phase accumulation (see
:func:`ezmsg.simbiophys.oscillator.advance_drifting_sine`) so it stays
continuous across chunks and as the frequency changes.
"""

import ezmsg.core as ez
import numpy as np
import numpy.typing as npt
from array_api_compat import get_namespace
from ezmsg.baseproc import (
    BaseStatefulTransformer,
    BaseTransformerUnit,
    processor_state,
)
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.util.messages.util import replace

from .oscillator import advance_drifting_sine, freq_drift_step_std




class LineNoiseSettings(ez.Settings):
    freq: float | None = None
    """Mains frequency in Hz (typically 50.0 or 60.0). None disables the stage
    (pass-through)."""

    amp: float = 1.0
    """Amplitude of the added sinusoid (same units as the input signal), shared
    across all channels."""

    drift_rate: float = 0.002
    """Frequency drift rate in Hz per second (RMS wander over 1 s), emulating
    recording-clock drift. Set to 0 for a perfectly fixed frequency."""

    drift_bound: float = 1.5
    """The line frequency is clamped to ``freq +/- drift_bound`` Hz."""

    seed: int | None = None
    """Random seed for the frequency-drift walk. If None, uses system entropy."""





@processor_state
class LineNoiseState:
    ang_freq: npt.NDArray[np.float64] | None = None
    """Base angular frequency ``2*pi*freq``, shape (1, 1)."""

    amp: npt.NDArray[np.float64] | None = None
    """Amplitude, shape (1, 1)."""

    phase: npt.NDArray[np.float64] | None = None
    """Accumulated phase carried across chunks, shape (1, 1)."""

    freq_off: npt.NDArray[np.float64] | None = None
    """Current frequency offset (Hz) carried across chunks, shape (1, 1)."""

    step_std: float = 0.0
    """Per-sample random-walk std for the drift."""

    rng: np.random.Generator | None = None
    """Random number generator for the drift walk."""

    dt: float = 0.0
    """Sample period (seconds)."""





class LineNoiseTransformer(BaseStatefulTransformer[LineNoiseSettings, AxisArray, AxisArray, LineNoiseState]):
    """Add a slowly frequency-drifting mains sinusoid to every channel.

    The same single-frequency component is added to all channels (mains pickup is
    a common source). Timing (sample rate) is taken from the input; the number of
    channels may vary. When ``freq`` is None the input passes through unchanged.
    """

    def _hash_message(self, message: AxisArray) -> int:
        time_axis = message.axes.get("time")
        gain = time_axis.gain if time_axis is not None else 0.0
        return hash(gain)

    def _reset_state(self, message: AxisArray) -> None:
        if self.settings.freq is None:
            return
        time_axis = message.axes.get("time")
        self._state.dt = time_axis.gain if time_axis is not None else 1.0
        self._state.ang_freq = np.array([[2.0 * np.pi * self.settings.freq]], dtype=np.float64)
        self._state.amp = np.array([[self.settings.amp]], dtype=np.float64)
        self._state.phase = np.zeros((1, 1), dtype=np.float64)
        self._state.freq_off = np.zeros((1, 1), dtype=np.float64)
        self._state.step_std = freq_drift_step_std(self.settings.drift_rate, self._state.dt)
        self._state.rng = np.random.default_rng(self.settings.seed)

    def _process(self, message: AxisArray) -> AxisArray:
        if self.settings.freq is None:
            return message  # pass-through

        xp = get_namespace(message.data)
        data = np.asarray(message.data, dtype=np.float64)
        was_1d = data.ndim == 1
        if was_1d:
            data = data[:, np.newaxis]
        n_samples = data.shape[0]

        sine, self._state.phase, self._state.freq_off = advance_drifting_sine(
            n_samples,
            self._state.dt,
            self._state.ang_freq,
            self._state.amp,
            self._state.phase,
            self._state.freq_off,
            self._state.step_std,
            self.settings.drift_bound,
            self._state.rng,
        )

        out = data + sine  # sine is (n_samples, 1), broadcasts over channels
        if was_1d:
            out = out[:, 0]
        out_data = xp.asarray(out) if xp is not np else out
        return replace(message, data=out_data)





class LineNoiseUnit(
    BaseTransformerUnit[
        LineNoiseSettings,
        AxisArray,
        AxisArray,
        LineNoiseTransformer,
    ]
):
    """Unit wrapper for LineNoiseTransformer."""

    SETTINGS = LineNoiseSettings