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 math

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, AxisBase
from ezmsg.util.messages.util import replace

from .oscillator import advance_drifting_sine, freq_drift_step_std

MLX_BROADCAST_MIN_ELEMENTS = 200_000
"""Minimum signal size for keeping line-noise addition on MLX.

Below this crossover, MLX dispatch costs more than the existing fused NumPy
round-trip. Above it, transferring only the single generated sine column and
broadcasting on-device avoids converting the full sample-by-channel signal.
"""




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 _chunk_axis(self, message: AxisArray) -> AxisBase | None:
        """The axis the stream grows along, however the producer named it."""
        dim = message.stream_dim or next((d for d in self.STREAMING_DIMS if d in message.dims), None)
        return message.axes.get(dim)

    def _hash_message(self, message: AxisArray) -> int:
        # Deliberately narrower than the default. Every state array is (1, 1) and
        # broadcasts across however many channels arrive, so the channel count and
        # fingerprint the default folds in would only restart the phase
        # accumulator for a sinusoid that did not change. The sample period is the
        # one thing this depends on.
        return hash(getattr(self._chunk_axis(message), "gain", None))

    def _reset_state(self, message: AxisArray) -> None:
        if self.settings.freq is None:
            return
        time_axis = self._chunk_axis(message)
        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)
        was_1d = message.data.ndim == 1
        n_samples = message.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,
        )

        use_mlx_broadcast = xp.__name__ == "mlx.core" and math.prod(message.data.shape) >= MLX_BROADCAST_MIN_ELEMENTS
        if use_mlx_broadcast:
            # Keep the large signal on MLX and transfer only the common-mode
            # (n_samples, 1) sine. This preserves the NumPy RNG/state semantics
            # while avoiding a round-trip of every signal channel.
            data = message.data[:, np.newaxis] if was_1d else message.data
            out = data + xp.asarray(sine)
        else:
            data = np.asarray(message.data, dtype=np.float64)
            if was_1d:
                data = data[:, np.newaxis]
            out = data + sine  # sine broadcasts over channels

        if was_1d:
            out = out[:, 0]
        out_data = out if use_mlx_broadcast else xp.asarray(out)
        return replace(message, data=out_data)





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

    SETTINGS = LineNoiseSettings