diff --git a/src/ezmsg/simbiophys/__init__.py b/src/ezmsg/simbiophys/__init__.py index 8db5aa4..347a038 100644 --- a/src/ezmsg/simbiophys/__init__.py +++ b/src/ezmsg/simbiophys/__init__.py @@ -14,6 +14,14 @@ from .__version__ import __version__ as __version__ +# Baseline Drift +from .baseline_drift import ( + BaselineDriftSettings, + BaselineDriftState, + BaselineDriftTransformer, + BaselineDriftUnit, +) + # Cosine Encoder from .cosine_encoder import ( CosineEncoderSettings, @@ -50,6 +58,14 @@ EEGSynthSettings, ) +# Line Noise +from .line_noise import ( + LineNoiseSettings, + LineNoiseState, + LineNoiseTransformer, + LineNoiseUnit, +) + # Noise from .noise import ( PinkNoise, @@ -106,11 +122,21 @@ # EEG "EEGSynth", "EEGSynthSettings", + # Line Noise + "LineNoiseSettings", + "LineNoiseState", + "LineNoiseTransformer", + "LineNoiseUnit", # Cosine Encoder "CosineEncoderSettings", "CosineEncoderState", "CosineEncoderTransformer", "CosineEncoderUnit", + # Baseline Drift + "BaselineDriftSettings", + "BaselineDriftState", + "BaselineDriftTransformer", + "BaselineDriftUnit", # Dynamic Colored Noise "DynamicColoredNoiseSettings", "DynamicColoredNoiseState", diff --git a/src/ezmsg/simbiophys/baseline_drift.py b/src/ezmsg/simbiophys/baseline_drift.py new file mode 100644 index 0000000..f69209f --- /dev/null +++ b/src/ezmsg/simbiophys/baseline_drift.py @@ -0,0 +1,243 @@ +"""Always-on slow 1/f (pink) baseline drift. + +Adds a low-frequency, velocity-independent baseline wander to a signal so that +even at rest there is non-zero power at very low frequencies. This complements +:mod:`ezmsg.simbiophys.dynamic_colored_noise`, whose Kasdin filter runs at the +full output rate and therefore only holds a true 1/f slope down to a corner of +a few hundred Hz (at 30 kHz, ~700 Hz with 5 poles). Reaching sub-Hz drift that +way would need thousands of poles. + +Instead we generate pink noise with a fixed spectral exponent (β = 1) on a +**low internal grid** (``drift_fs``, e.g. 50 Hz), where a handful of poles put +the 1/f corner well below 1 Hz, then linearly interpolate that slow signal up +to the output rate and add it. Because the drift is dominated by frequencies far +below the drift-grid Nyquist, linear interpolation is effectively exact. + +The drift filter state, interpolation anchors, and phase are all carried across +chunks so the output is continuous and independent of how the stream is chunked. + +Performance: the grid is generated on only ``drift_fs`` samples per second (a +handful per chunk), so the sequential IIR recursion is tiny. Upsampling to the +output rate, scaling, and adding to the input are fused into a single compiled +pass that touches the ``(n_samples, n_channels)`` result exactly once (see +:func:`_pink_drift_add`), which keeps the stage memory-bandwidth bound and +scales well to large channel counts on low-power CPUs. +""" + +import ezmsg.core as ez +import numba +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 .dynamic_colored_noise import compute_kasdin_coefficients + + +def _generate_grid(white: np.ndarray, coeffs: np.ndarray, delay_lines: np.ndarray) -> np.ndarray: + """Advance the Kasdin IIR one step per row of ``white`` (vectorised over channels). + + Used only for the one-off warm-up at reset (``white`` has a few hundred rows), + so it stays in plain NumPy. ``delay_lines`` (shape ``(n_channels, n_poles)``) + is mutated in place; returns the generated grid, shape ``(n_new, n_channels)``. + """ + n_new, n_channels = white.shape + n_poles = coeffs.shape[0] + grid = np.empty((n_new, n_channels), dtype=np.float64) + for k in range(n_new): + # y[n] = x[n] - sum_p(a_p * y[n-p]); coeffs contracts the pole axis. + y = white[k] - delay_lines @ coeffs + # Shift delay lines (newest into slot 0); walk high->low to avoid aliasing. + for p in range(n_poles - 1, 0, -1): + delay_lines[:, p] = delay_lines[:, p - 1] + delay_lines[:, 0] = y + grid[k] = y + return grid + + +@numba.jit(nopython=True, cache=True, fastmath=True) +def _pink_drift_add( + out: np.ndarray, # (n_out, n_channels) result, written in place: data + scale*drift + data: np.ndarray, # (n_out, n_channels) input signal + white: np.ndarray, # (n_new, n_channels) white noise for new grid samples + coeffs: np.ndarray, # (n_poles,) fixed Kasdin coefficients + delay_lines: np.ndarray, # (n_channels, n_poles) filter state, mutated in place + prev: np.ndarray, # (n_channels,) left interp anchor, mutated in place + nxt: np.ndarray, # (n_channels,) right interp anchor, mutated in place + phase: float, # fractional position between prev and nxt, in [0, 1) + step: float, # low-rate grid samples advanced per output sample + scale: float, # drift amplitude +) -> float: + """Interpolate, scale, and add the slow pink drift in a single fused pass. + + Each output sample linearly interpolates the low-rate grid (``prev`` -> + ``nxt``), scales it, and adds it to ``data``, writing ``out`` once. When + ``phase`` crosses 1.0 a new grid sample is produced by advancing the Kasdin + IIR one step. Returns the carried phase so the next chunk resumes seamlessly. + """ + n_out = out.shape[0] + n_channels = out.shape[1] + n_poles = coeffs.shape[0] + + w = 0 # index into the pre-drawn white-noise grid + ph = phase + for i in range(n_out): + # Hoist the interpolation weights so the channel loop is a plain axpy + # (out = data + a*prev + b*nxt), which the compiler can vectorise (SIMD). + a = scale * (1.0 - ph) + b = scale * ph + for c in range(n_channels): + out[i, c] = data[i, c] + a * prev[c] + b * nxt[c] + ph += step + while ph >= 1.0: + ph -= 1.0 + for c in range(n_channels): + acc = 0.0 + for p in range(n_poles): + acc += coeffs[p] * delay_lines[c, p] + y = white[w, c] - acc + for p in range(n_poles - 1, 0, -1): + delay_lines[c, p] = delay_lines[c, p - 1] + delay_lines[c, 0] = y + prev[c] = nxt[c] + nxt[c] = y + w += 1 + return ph + + +class BaselineDriftSettings(ez.Settings): + scale: float = 1.0 + """Amplitude of the added drift (multiplies the raw filter output).""" + + drift_fs: float = 50.0 + """Internal sample rate (Hz) at which the slow pink noise is generated. + Lower rates push the 1/f corner lower for a given pole count; it only needs + to sit comfortably above the drift frequencies of interest.""" + + n_poles: int = 24 + """Number of IIR poles for the drift filter. With ``drift_fs`` = 50 Hz this + places the 1/f corner near ~0.3 Hz, giving sub-Hz to few-Hz wander.""" + + beta: float = 1.0 + """Spectral exponent of the drift. 1.0 is pink (1/f); kept < 2 for stability.""" + + seed: int | None = None + """Random seed for reproducibility. If None, uses system entropy.""" + + +@processor_state +class BaselineDriftState: + delay_lines: npt.NDArray[np.float64] | None = None + """Kasdin filter delay lines, shape (n_channels, n_poles).""" + + coeffs: npt.NDArray[np.float64] | None = None + """Fixed filter coefficients for ``beta``, shape (n_poles,).""" + + prev: npt.NDArray[np.float64] | None = None + """Left interpolation anchor (current low-rate sample), shape (n_channels,).""" + + nxt: npt.NDArray[np.float64] | None = None + """Right interpolation anchor (next low-rate sample), shape (n_channels,).""" + + phase: float = 0.0 + """Fractional position between ``prev`` and ``nxt`` for the next output sample.""" + + rng: np.random.Generator | None = None + """Random number generator.""" + + step: float = 0.0 + """Low-rate grid samples advanced per output sample (drift_fs / output_fs).""" + + +class BaselineDriftTransformer( + BaseStatefulTransformer[BaselineDriftSettings, AxisArray, AxisArray, BaselineDriftState] +): + """Add always-on slow 1/f baseline drift to a signal. + + The drift is generated independently of the input values (it depends only on + the input's shape and sample rate for timing), so it is present even when the + upstream signal is at rest. One independent drift process is generated per + channel of the input. + """ + + 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 + n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 + return hash((n_channels, gain)) + + def _reset_state(self, message: AxisArray) -> None: + n_channels = message.data.shape[1] if message.data.ndim > 1 else 1 + + self._state.rng = np.random.default_rng(self.settings.seed) + self._state.coeffs = compute_kasdin_coefficients(self.settings.beta, self.settings.n_poles) + self._state.delay_lines = np.zeros((n_channels, self.settings.n_poles), dtype=np.float64) + + time_axis = message.axes.get("time") + output_gain = time_axis.gain if time_axis is not None else 1.0 + output_fs = 1.0 / output_gain + self._state.step = self.settings.drift_fs / output_fs + + # Warm up the filter so the delay lines hold representative pink state + # before the first output sample (avoids a startup transient from zero). + # The last two warmup samples seed the interpolation anchors. + n_warmup = 20 * self.settings.n_poles + warm_white = self._state.rng.standard_normal((n_warmup + 2, n_channels)) + grid = _generate_grid(warm_white, self._state.coeffs, self._state.delay_lines) + self._state.prev = grid[-2].copy() + self._state.nxt = grid[-1].copy() + self._state.phase = 0.0 + + def _process(self, message: AxisArray) -> AxisArray: + 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_out, n_channels = data.shape + + # Number of new grid samples this chunk consumes. Carrying phase across + # chunks makes this exact and chunk-order independent. + n_steps = int(np.floor(self._state.phase + n_out * self._state.step)) + white = self._state.rng.standard_normal((n_steps, n_channels)) + + # Fused interpolate + scale + add, writing the result in a single pass. + out = np.empty_like(data) + self._state.phase = _pink_drift_add( + out, + data, + white, + self._state.coeffs, + self._state.delay_lines, # mutated in place + self._state.prev, # mutated in place + self._state.nxt, # mutated in place + self._state.phase, + self._state.step, + self.settings.scale, + ) + + 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 BaselineDriftUnit( + BaseTransformerUnit[ + BaselineDriftSettings, + AxisArray, + AxisArray, + BaselineDriftTransformer, + ] +): + """Unit wrapper for BaselineDriftTransformer.""" + + SETTINGS = BaselineDriftSettings diff --git a/src/ezmsg/simbiophys/line_noise.py b/src/ezmsg/simbiophys/line_noise.py new file mode 100644 index 0000000..1604884 --- /dev/null +++ b/src/ezmsg/simbiophys/line_noise.py @@ -0,0 +1,137 @@ +"""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 diff --git a/src/ezmsg/simbiophys/oscillator.py b/src/ezmsg/simbiophys/oscillator.py index 0131238..76d8435 100644 --- a/src/ezmsg/simbiophys/oscillator.py +++ b/src/ezmsg/simbiophys/oscillator.py @@ -12,6 +12,70 @@ from ezmsg.util.messages.axisarray import AxisArray, LinearAxis, replace +def freq_drift_step_std(drift_rate_per_sec: float, dt: float) -> float: + """Per-sample random-walk std that yields a given drift rate. + + A frequency doing a zero-mean random walk with per-sample increment std + ``s`` drifts with RMS ``s * sqrt(T/dt)`` over an interval ``T``. Choosing + ``s = drift_rate_per_sec * sqrt(dt)`` makes the RMS drift over 1 s equal to + ``drift_rate_per_sec`` (Hz/s), independent of sample rate or chunking. + """ + return drift_rate_per_sec * np.sqrt(dt) + + +def advance_drifting_sine( + n_samples: int, + dt: float, + ang_freq: np.ndarray, + amp: np.ndarray, + phase_state: np.ndarray, + freq_off_state: np.ndarray, + step_std: float, + bound_hz: float, + rng: np.random.Generator, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Generate a sinusoid whose frequency slowly wanders, via phase accumulation. + + The instantaneous frequency is ``base + offset`` where ``offset`` does a + bounded random walk (clamped to ``+/- bound_hz``). Phase is integrated + sample-by-sample so the waveform stays continuous even as the frequency + changes and across chunk boundaries. + + Args: + n_samples: Number of output samples for this chunk. + dt: Sample period (seconds). + ang_freq: Base angular frequency ``2*pi*f``, shape ``(1, k)``. + amp: Amplitude, shape ``(1, k)``. + phase_state: Carried phase from the previous chunk, shape ``(1, k)``. + freq_off_state: Carried frequency offset (Hz), shape ``(1, k)``. + step_std: Per-sample random-walk std (see :func:`freq_drift_step_std`); + 0 disables drift (fixed frequency). + bound_hz: Frequency offset is clamped to ``+/- bound_hz``. + rng: Random generator for the walk. + + Returns: + ``(data, new_phase_state, new_freq_off_state)`` where ``data`` has shape + ``(n_samples, k)`` and the two states have shape ``(1, k)``. + """ + k = ang_freq.shape[1] + if n_samples == 0: + return np.zeros((0, k)), phase_state, freq_off_state + + if step_std > 0.0: + inc = rng.standard_normal((n_samples, k)) * step_std + freq_off = freq_off_state + np.cumsum(inc, axis=0) + np.clip(freq_off, -bound_hz, bound_hz, out=freq_off) + else: + freq_off = np.zeros((n_samples, k)) + + inst_ang = ang_freq + 2.0 * np.pi * freq_off # (n_samples, k) + phase = phase_state + np.cumsum(inst_ang * dt, axis=0) + data = amp * np.sin(phase) + new_phase_state = np.mod(phase[-1:, :], 2.0 * np.pi) + new_freq_off_state = freq_off[-1:, :].copy() + return data, new_phase_state, new_freq_off_state + + class SpiralGeneratorSettings(ClockDrivenSettings): """Settings for :obj:`SpiralGenerator`. @@ -131,6 +195,18 @@ class SinGeneratorSettings(ClockDrivenSettings): phase: float | npt.ArrayLike = 0.0 """The initial phase of the sinusoid, in radians. Scalar or per-channel array.""" + freq_drift_rate: float = 0.0 + """Frequency drift rate in Hz per second (RMS wander over 1 s). When > 0 the + frequency does a slow bounded random walk to emulate e.g. recording-clock + drift. 0 (default) keeps the frequency fixed.""" + + freq_drift_bound: float = 1.5 + """The drifting frequency is clamped to ``+/- freq_drift_bound`` Hz around the + base ``freq``.""" + + freq_drift_seed: int | None = None + """Random seed for the frequency-drift walk. If None, uses system entropy.""" + @processor_state class SinGeneratorState(ClockDrivenState): @@ -141,6 +217,11 @@ class SinGeneratorState(ClockDrivenState): ang_freq: np.ndarray | None = None # 2*pi*freq amp: np.ndarray | None = None phase: np.ndarray | None = None + # Frequency-drift state (only used when freq_drift_rate > 0) + drift_rng: np.random.Generator | None = None + drift_phase: np.ndarray | None = None # accumulated phase, shape (1, k) + drift_freq_off: np.ndarray | None = None # frequency offset (Hz), shape (1, k) + drift_step_std: float = 0.0 class SinProducer(BaseClockDrivenProducer[SinGeneratorSettings, SinGeneratorState]): @@ -190,12 +271,34 @@ def _reset_state(self, time_axis: LinearAxis) -> None: self._state.amp = amp self._state.phase = phase + # Frequency-drift state. When drift is enabled we switch to phase + # accumulation (in _produce) to keep the waveform continuous as the + # frequency wanders; when disabled the fast closed-form path is used. + if self.settings.freq_drift_rate > 0.0: + self._state.drift_rng = np.random.default_rng(self.settings.freq_drift_seed) + self._state.drift_phase = np.array(self._state.phase, dtype=np.float64).reshape(1, -1) + self._state.drift_freq_off = np.zeros_like(self._state.drift_phase) + self._state.drift_step_std = freq_drift_step_std(self.settings.freq_drift_rate, time_axis.gain) + def _produce(self, n_samples: int, time_axis: LinearAxis) -> AxisArray: """Generate sinusoidal waveform for this chunk.""" - # Calculate sinusoid: amp * sin(ang_freq*t + phase) - # t shape: (n_time,) -> (n_time, 1) for broadcasting with (1, n_ch) - t = (np.arange(n_samples) + self._state.counter)[:, np.newaxis] * time_axis.gain - sin_data = self._state.amp * np.sin(self._state.ang_freq * t + self._state.phase) + if self.settings.freq_drift_rate > 0.0: + sin_data, self._state.drift_phase, self._state.drift_freq_off = advance_drifting_sine( + n_samples, + time_axis.gain, + self._state.ang_freq, + self._state.amp, + self._state.drift_phase, + self._state.drift_freq_off, + self._state.drift_step_std, + self.settings.freq_drift_bound, + self._state.drift_rng, + ) + else: + # Fast closed-form path: amp * sin(ang_freq*t + phase) + # t shape: (n_time,) -> (n_time, 1) for broadcasting with (1, n_ch) + t = (np.arange(n_samples) + self._state.counter)[:, np.newaxis] * time_axis.gain + sin_data = self._state.amp * np.sin(self._state.ang_freq * t + self._state.phase) # Tile if all params were scalar but n_ch > 1 if sin_data.shape[1] < self.settings.n_ch: diff --git a/src/ezmsg/simbiophys/system/velocity2ecephys.py b/src/ezmsg/simbiophys/system/velocity2ecephys.py index e917101..dc3263b 100644 --- a/src/ezmsg/simbiophys/system/velocity2ecephys.py +++ b/src/ezmsg/simbiophys/system/velocity2ecephys.py @@ -26,6 +26,7 @@ from ezmsg.sigproc.math.add import Add from ezmsg.util.messages.axisarray import AxisArray +from ..line_noise import LineNoiseSettings, LineNoiseUnit from .velocity2lfp import Velocity2LFP, Velocity2LFPSettings from .velocity2spike import Velocity2Spike, Velocity2SpikeSettings @@ -56,6 +57,24 @@ class VelocityEncoderSettings(ez.Settings): n_sources: int = 8 """Number of cosine-encoded LFP sources.""" + drift_scale: float = 4.0 + """Amplitude of always-on slow 1/f drift added to each LFP source before + mixing (a shared low-frequency field drift across channels). Velocity- + independent, so baseline wander is present even at rest. Set to 0 to disable.""" + + line_noise_freq: float | None = None + """Mains line-noise frequency in Hz (50.0 or 60.0). None (default) disables + line noise (pass-through). Added to every channel of the final signal.""" + + line_noise_amp: float = 10.0 + """Amplitude of the line-noise sinusoid (same units as the output signal).""" + + line_noise_drift_rate: float = 0.002 + """Line-noise frequency drift rate in Hz per second (recording-clock drift).""" + + line_noise_drift_bound: float = 1.5 + """Line-noise frequency is clamped to ``line_noise_freq +/- this`` Hz.""" + class VelocityEncoder(ez.Collection): """Encode cursor velocity into simulated extracellular electrophysiology. @@ -94,6 +113,7 @@ class VelocityEncoder(ez.Collection): SPIKES = Velocity2Spike() LFP = Velocity2LFP() ADD = Add() # Add colored noise and waveforms + LINE_NOISE = LineNoiseUnit() # Optional mains pickup on the combined signal OUTPUT_SIGNAL = ez.OutputTopic(AxisArray) def configure(self) -> None: @@ -113,6 +133,16 @@ def configure(self) -> None: output_ch=self.SETTINGS.output_ch, n_lfp_sources=self.SETTINGS.n_sources, max_velocity=self.SETTINGS.max_velocity, + drift_scale=self.SETTINGS.drift_scale, + seed=self.SETTINGS.seed, + ) + ) + self.LINE_NOISE.apply_settings( + LineNoiseSettings( + freq=self.SETTINGS.line_noise_freq, + amp=self.SETTINGS.line_noise_amp, + drift_rate=self.SETTINGS.line_noise_drift_rate, + drift_bound=self.SETTINGS.line_noise_drift_bound, seed=self.SETTINGS.seed, ) ) @@ -124,5 +154,6 @@ def network(self) -> ez.NetworkDefinition: (self.SPIKES.OUTPUT_SIGNAL, self.ADD.INPUT_SIGNAL_A), (self.COORDS.OUTPUT_SIGNAL, self.LFP.INPUT_SIGNAL), (self.LFP.OUTPUT_SIGNAL, self.ADD.INPUT_SIGNAL_B), - (self.ADD.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), + (self.ADD.OUTPUT_SIGNAL, self.LINE_NOISE.INPUT_SIGNAL), + (self.LINE_NOISE.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), ) diff --git a/src/ezmsg/simbiophys/system/velocity2lfp.py b/src/ezmsg/simbiophys/system/velocity2lfp.py index 0dca076..74d93a6 100644 --- a/src/ezmsg/simbiophys/system/velocity2lfp.py +++ b/src/ezmsg/simbiophys/system/velocity2lfp.py @@ -22,16 +22,38 @@ :mod:`ezmsg.simbiophys.system.velocity2ecephys`: Combined spike + LFP encoding. """ +import functools + import ezmsg.core as ez import numpy as np from ezmsg.sigproc.affinetransform import AffineTransform, AffineTransformSettings from ezmsg.sigproc.math.clip import Clip, ClipSettings from ezmsg.util.messages.axisarray import AxisArray +from ..baseline_drift import BaselineDriftSettings, BaselineDriftUnit from ..cosine_encoder import CosineEncoderSettings, CosineEncoderUnit from ..dynamic_colored_noise import DynamicColoredNoiseSettings, DynamicColoredNoiseUnit +def _make_mixing_weights(n_in: int, output_ch: int, seed: int) -> np.ndarray: + """Build the ``n_in`` -> ``output_ch`` spatial mixing matrix. + + Module-level (not a closure) so it can be pickled as part of the settings + metadata snapshot; ``output_ch``/``seed`` are bound via ``functools.partial`` + in :meth:`Velocity2LFP.configure`. Called with ``n_in`` so the weights are + rebuilt automatically if the number of sources changes at runtime. + """ + rng = np.random.default_rng(seed) + ch_idx = np.arange(output_ch) + weights = np.zeros((n_in, output_ch)) + for i in range(n_in): + freq = (i + 1) / n_in + phase = 2 * np.pi * i / n_in + weights[i, :] = np.sin(2 * np.pi * freq * ch_idx / output_ch + phase) + weights += 0.3 * rng.standard_normal((n_in, output_ch)) + return weights + + class Velocity2LFPSettings(ez.Settings): """Settings for :obj:`Velocity2LFP`.""" @@ -48,6 +70,16 @@ class Velocity2LFPSettings(ez.Settings): max_velocity: float = 315.0 + drift_scale: float = 4.0 + """Amplitude of always-on slow 1/f drift added to each LFP source before + mixing, so it appears as a shared low-frequency field drift across output + channels. Velocity-independent, so baseline wander is present even at rest. + Set to 0 to disable. Roughly ~20% of the per-channel LFP std at defaults.""" + + drift_fs: float = 50.0 + """Internal generation rate (Hz) for the baseline drift. Lower rates push the + 1/f corner to lower frequencies (slower wander) for a given pole count.""" + seed: int = 6767 """Random seed for reproducible preferred directions and mixing matrix.""" @@ -58,12 +90,16 @@ class Velocity2LFP(ez.Collection): This system converts polar velocity coordinates into multi-channel LFP-like signals: 1. **Cosine encoder**: Each of n_lfp_sources has a different preferred - direction. The spectral exponent beta (0-2) is modulated by the cosine + direction. The spectral exponent beta (0-1.95) is modulated by the cosine of the angle between velocity and preferred direction, scaled by speed. - 2. **Clip**: Ensures beta values stay within valid range [0, 2]. + 2. **Clip**: Ensures beta values stay within valid range [0, 1.95]. The + ceiling is kept just below 2.0 to keep the noise filter stable (see + ``configure``). 3. **Colored noise**: Generates 1/f^beta noise where beta is dynamically modulated per source. - 4. **Spatial mixing**: Projects the n_lfp_sources onto output_ch channels + 4. **Baseline drift**: Adds always-on slow 1/f wander to each source, giving + a shared low-frequency field drift once mixed (present even at rest). + 5. **Spatial mixing**: Projects the n_lfp_sources onto output_ch channels using a sinusoidal mixing matrix with random perturbations. Input: @@ -83,6 +119,7 @@ class Velocity2LFP(ez.Collection): BETA_ENCODER = CosineEncoderUnit() CLIP_BETA = Clip() PINK_NOISE = DynamicColoredNoiseUnit() + BASELINE_DRIFT = BaselineDriftUnit() # Always-on slow 1/f wander per source MIX_NOISE = AffineTransform() # Project n_lfp_sources to output_ch sensors OUTPUT_SIGNAL = ez.OutputTopic(AxisArray) @@ -101,7 +138,13 @@ def configure(self) -> None: ) ) - self.CLIP_BETA.apply_settings(ClipSettings(min=0.0, max=2.0)) + # Cap strictly below 2.0. At beta == 2 the Kasdin all-pole filter has a + # pole exactly on the unit circle (y[n] = x[n] + y[n-1], a pure + # integrator), so its delay-line state does an unbounded random walk. + # When sustained high velocity pins beta at the ceiling the state grows + # large, and a subsequent beta change unleashes it as a big transient + # followed by ringing. 1.95 keeps the pole inside the unit circle. + self.CLIP_BETA.apply_settings(ClipSettings(min=0.0, max=1.95)) self.PINK_NOISE.apply_settings( DynamicColoredNoiseSettings( @@ -114,23 +157,29 @@ def configure(self) -> None: ) ) - # Create mixing matrix factory: n_lfp_sources -> output_ch - # Using a callable so the weights are rebuilt automatically if the - # number of input sources changes at runtime (e.g. via live settings). - output_ch = self.SETTINGS.output_ch - seed = self.SETTINGS.seed - - def make_mixing_weights(n_in: int) -> np.ndarray: - rng = np.random.default_rng(seed) - ch_idx = np.arange(output_ch) - weights = np.zeros((n_in, output_ch)) - for i in range(n_in): - freq = (i + 1) / n_in - phase = 2 * np.pi * i / n_in - weights[i, :] = np.sin(2 * np.pi * freq * ch_idx / output_ch + phase) - weights += 0.3 * rng.standard_normal((n_in, output_ch)) - return weights + # Add always-on slow 1/f drift to each source before mixing. The pink + # filter runs at output_fs and cannot produce sub-Hz power (its 1/f corner + # is a few hundred Hz), so the low-frequency wander is generated here. + # Adding it per source (rather than per output channel) means the drift is + # spread across sensors by the mixing matrix -- a shared "field" drift + # rather than independent per-electrode drift -- and its cost scales with + # n_lfp_sources instead of output_ch. + self.BASELINE_DRIFT.apply_settings( + BaselineDriftSettings( + scale=self.SETTINGS.drift_scale, + drift_fs=self.SETTINGS.drift_fs, + beta=1.0, + seed=self.SETTINGS.seed, + ) + ) + # Mixing matrix factory: n_lfp_sources -> output_ch. A picklable + # module-level function (bound with partial) so the settings metadata + # snapshot can serialize it; passing a callable keeps the weights + # rebuildable if the number of sources changes at runtime. + make_mixing_weights = functools.partial( + _make_mixing_weights, output_ch=self.SETTINGS.output_ch, seed=self.SETTINGS.seed + ) self.MIX_NOISE.apply_settings(AffineTransformSettings(weights=make_mixing_weights, axis="ch")) def network(self) -> ez.NetworkDefinition: @@ -138,6 +187,7 @@ def network(self) -> ez.NetworkDefinition: (self.INPUT_SIGNAL, self.BETA_ENCODER.INPUT_SIGNAL), (self.BETA_ENCODER.OUTPUT_SIGNAL, self.CLIP_BETA.INPUT_SIGNAL), (self.CLIP_BETA.OUTPUT_SIGNAL, self.PINK_NOISE.INPUT_SIGNAL), - (self.PINK_NOISE.OUTPUT_SIGNAL, self.MIX_NOISE.INPUT_SIGNAL), + (self.PINK_NOISE.OUTPUT_SIGNAL, self.BASELINE_DRIFT.INPUT_SIGNAL), + (self.BASELINE_DRIFT.OUTPUT_SIGNAL, self.MIX_NOISE.INPUT_SIGNAL), (self.MIX_NOISE.OUTPUT_SIGNAL, self.OUTPUT_SIGNAL), ) diff --git a/tests/unit/test_baseline_drift.py b/tests/unit/test_baseline_drift.py new file mode 100644 index 0000000..080f1cf --- /dev/null +++ b/tests/unit/test_baseline_drift.py @@ -0,0 +1,83 @@ +"""Unit tests for ezmsg.simbiophys.baseline_drift module.""" + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray +from numpy.fft import rfft, rfftfreq + +from ezmsg.simbiophys import ( + BaselineDriftSettings, + BaselineDriftTransformer, +) + + +def _msg(n, fs, offset, n_ch=1): + data = np.zeros((n, n_ch)) if n_ch else np.zeros(n) + return AxisArray(data, dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=fs, offset=offset)}) + + +class TestBaselineDriftTransformer: + def test_output_shape_and_addition(self): + """Drift is added to the input, preserving shape.""" + tr = BaselineDriftTransformer(BaselineDriftSettings(scale=1.0, drift_fs=50.0, seed=0)) + base = np.full((3000, 4), 7.0) + msg = AxisArray(base, dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=30000.0, offset=0.0)}) + out = tr(msg) + assert out.data.shape == (3000, 4) + # Output differs from input (drift added) but stays near the constant baseline. + assert not np.allclose(out.data, base) + assert np.abs(out.data.mean() - 7.0) < 5.0 + + def test_scale_zero_is_passthrough(self): + tr = BaselineDriftTransformer(BaselineDriftSettings(scale=0.0, drift_fs=50.0, seed=0)) + base = np.random.default_rng(1).standard_normal((2000, 2)) + msg = AxisArray(base, dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=30000.0, offset=0.0)}) + out = tr(msg) + np.testing.assert_allclose(out.data, base) + + def test_continuity_chunk_independent(self): + """Chunked processing is bit-identical to a single-shot call.""" + fs = 30000.0 + single = BaselineDriftTransformer(BaselineDriftSettings(scale=1.0, drift_fs=50.0, seed=7)) + whole = single(_msg(5 * int(fs), fs, 0.0)).data + + chunked_tr = BaselineDriftTransformer(BaselineDriftSettings(scale=1.0, drift_fs=50.0, seed=7)) + parts, off = [], 0.0 + for _ in range(5): + parts.append(chunked_tr(_msg(int(fs), fs, off)).data) + off += 1.0 + chunked = np.concatenate(parts) + np.testing.assert_allclose(whole, chunked, atol=1e-12) + + def test_reproducible_with_seed(self): + a = BaselineDriftTransformer(BaselineDriftSettings(scale=1.0, seed=3))(_msg(3000, 30000.0, 0.0)).data + b = BaselineDriftTransformer(BaselineDriftSettings(scale=1.0, seed=3))(_msg(3000, 30000.0, 0.0)).data + np.testing.assert_allclose(a, b) + + def test_has_sub_hz_power(self): + """The always-on drift produces a 1/f spectrum with real sub-Hz power.""" + fs = 30000.0 + tr = BaselineDriftTransformer(BaselineDriftSettings(scale=1.0, drift_fs=50.0, n_poles=24, seed=0)) + parts, off = [], 0.0 + for _ in range(30): # 30 s so we resolve sub-Hz bins + parts.append(tr(_msg(int(fs), fs, off)).data[:, 0]) + off += 1.0 + x = np.concatenate(parts) + x = x - x.mean() + f = rfftfreq(len(x), 1 / fs) + p = np.abs(rfft(x)) ** 2 + # log-log slope over 0.1-10 Hz should be near -1 (pink). + band = (f >= 0.1) & (f <= 10) + slope = np.polyfit(np.log10(f[band]), np.log10(p[band]), 1)[0] + assert -1.4 < slope < -0.7 + # Sub-Hz power must dominate mid-band power (the flat-baseline complaint). + sub = p[(f > 0.1) & (f < 1.0)].mean() + mid = p[(f >= 10) & (f < 100)].mean() + assert sub > 20 * mid + + def test_finite_output(self): + tr = BaselineDriftTransformer(BaselineDriftSettings(scale=10.0, drift_fs=50.0, n_poles=24, seed=0)) + off = 0.0 + for _ in range(50): + out = tr(_msg(3000, 30000.0, off, n_ch=8)) + assert np.isfinite(out.data).all() + off += 0.1 diff --git a/tests/unit/test_dynamic_colored_noise.py b/tests/unit/test_dynamic_colored_noise.py index c996adf..8d5535b 100644 --- a/tests/unit/test_dynamic_colored_noise.py +++ b/tests/unit/test_dynamic_colored_noise.py @@ -155,6 +155,41 @@ def test_continuity_across_chunks(self): # Boundary diffs should be within reasonable range of internal diffs assert diffs[idx] < 5 * max_internal_diff, f"Discontinuity at index {idx}" + def test_beta_below_two_is_bounded_on_drop(self): + """Regression: capping beta below 2.0 prevents transient blow-up. + + At beta == 2.0 the Kasdin filter is a pure integrator (pole on the unit + circle), so its state random-walks unbounded while beta is pinned there; + a subsequent beta drop then releases that state as a large transient + followed by ringing. Keeping beta just below 2.0 (as the LFP pipeline's + clip does) bounds the state and eliminates the artifact. + """ + fs = 100.0 + + def hold_then_drop(hold_beta): + tr = DynamicColoredNoiseTransformer( + DynamicColoredNoiseSettings( + output_fs=30000.0, n_poles=5, smoothing_tau=0.01, initial_beta=1.0, scale=20.0, seed=0 + ) + ) + hold = AxisArray( + np.full((300, 1), hold_beta), dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=fs, offset=0.0)} + ) + steady = tr(hold).data + drop = AxisArray( + np.full((50, 1), 0.2), dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=fs, offset=3.0)} + ) + drop_out = tr(drop).data + return np.abs(drop_out).max(), np.abs(steady[-2000:]).std() + + capped_peak, capped_std = hold_then_drop(1.95) + marginal_peak, _ = hold_then_drop(2.0) + + # Below the boundary the post-drop transient stays within a small factor + # of the steady-state fluctuation and is far smaller than the beta==2 case. + assert capped_peak < 15 * capped_std + assert capped_peak < 0.25 * marginal_peak + def test_varying_beta(self): """Test varying beta values across samples.""" transformer = DynamicColoredNoiseTransformer( diff --git a/tests/unit/test_line_noise.py b/tests/unit/test_line_noise.py new file mode 100644 index 0000000..021b3d4 --- /dev/null +++ b/tests/unit/test_line_noise.py @@ -0,0 +1,73 @@ +"""Unit tests for ezmsg.simbiophys.line_noise module.""" + +import numpy as np +from ezmsg.util.messages.axisarray import AxisArray +from numpy.fft import rfft, rfftfreq + +from ezmsg.simbiophys import LineNoiseSettings, LineNoiseTransformer + + +def _msg(n, fs, offset, n_ch=4, val=0.0): + return AxisArray( + np.full((n, n_ch), val), dims=["time", "ch"], axes={"time": AxisArray.TimeAxis(fs=fs, offset=offset)} + ) + + +class TestLineNoiseTransformer: + def test_disabled_is_passthrough(self): + tr = LineNoiseTransformer(LineNoiseSettings(freq=None)) + m = _msg(1000, 30000.0, 0.0, val=3.0) + out = tr(m) + np.testing.assert_array_equal(out.data, m.data) + + def test_added_to_all_channels_equally(self): + tr = LineNoiseTransformer(LineNoiseSettings(freq=60.0, amp=5.0, seed=1)) + out = tr(_msg(3000, 30000.0, 0.0, n_ch=4)).data + for ch in range(1, 4): + np.testing.assert_allclose(out[:, 0], out[:, ch]) + + def test_amplitude_and_peak_frequency(self): + fs = 30000.0 + tr = LineNoiseTransformer(LineNoiseSettings(freq=60.0, amp=7.0, drift_rate=0.0, seed=1)) + # No drift -> pure 60 Hz tone of amplitude 7. + y = tr(_msg(int(fs), fs, 0.0, n_ch=1)).data[:, 0] + assert abs(y.max() - 7.0) < 0.1 + f = rfftfreq(len(y), 1 / fs) + p = np.abs(rfft(y - y.mean())) ** 2 + assert abs(f[np.argmax(p)] - 60.0) < 0.5 + + def test_continuity_across_chunks(self): + fs = 30000.0 + tr = LineNoiseTransformer(LineNoiseSettings(freq=60.0, amp=1.0, drift_rate=0.1, seed=3)) + parts, off = [], 0.0 + for _ in range(20): + parts.append(tr(_msg(int(fs), fs, off, n_ch=1)).data[:, 0]) + off += 1.0 + y = np.concatenate(parts) + d = np.abs(np.diff(y)) + boundaries = [int(fs) * k - 1 for k in range(1, 20)] + # Boundary steps should be no worse than interior steps (phase is continuous). + assert d[boundaries].max() < 1.5 * np.delete(d, boundaries).max() + + def test_frequency_bounded(self): + fs = 30000.0 + # Aggressive drift so the bound is actually reached. + tr = LineNoiseTransformer(LineNoiseSettings(freq=60.0, amp=1.0, drift_rate=2.0, drift_bound=1.5, seed=5)) + off = 0.0 + for _ in range(1200): # 20 min + tr(_msg(int(fs), fs, off, n_ch=1)) + assert abs(tr._state.freq_off[0, 0]) <= 1.5 + 1e-9 + off += 1.0 + + def test_drift_rate_magnitude(self): + fs = 30000.0 + rate = 0.05 # Hz per second + tr = LineNoiseTransformer(LineNoiseSettings(freq=60.0, amp=1.0, drift_rate=rate, drift_bound=99.0, seed=7)) + offs, off = [], 0.0 + for _ in range(600): # 600 one-second chunks, unbounded to measure the walk itself + tr(_msg(int(fs), fs, off, n_ch=1)) + offs.append(tr._state.freq_off[0, 0]) + off += 1.0 + offs = np.array(offs) + one_sec_drift = np.std(np.diff(offs)) # std of 1 s increments ~ drift_rate + assert 0.5 * rate < one_sec_drift < 1.5 * rate diff --git a/tests/unit/test_oscillator.py b/tests/unit/test_oscillator.py index a6ffd28..88bece2 100644 --- a/tests/unit/test_oscillator.py +++ b/tests/unit/test_oscillator.py @@ -251,3 +251,41 @@ def test_sin_generator_variable_chunk_mode(): t = np.arange(350) / srate expected = np.sin(2 * np.pi * freq * t)[:, np.newaxis] np.testing.assert_allclose(concat.data, expected, rtol=1e-10) + + +def test_sin_generator_freq_drift_bounded_and_continuous(): + """Frequency drift stays bounded and keeps the waveform continuous.""" + fs = 30000.0 + n_time = int(fs) # 1-second chunks + producer = SinProducer( + SinGeneratorSettings( + fs=fs, + n_time=n_time, + n_ch=1, + freq=60.0, + amp=1.0, + freq_drift_rate=2.0, + freq_drift_bound=1.5, + freq_drift_seed=0, + ) + ) + parts = [] + for i in range(120): # 2 minutes + parts.append(producer(AxisArray.TimeAxis(fs=fs, offset=i * 1.0)).data[:, 0]) + # Instantaneous frequency offset must respect the bound. + assert abs(producer._state.drift_freq_off[0, 0]) <= 1.5 + 1e-9 + y = np.concatenate(parts) + # No phase discontinuity at chunk joins. + d = np.abs(np.diff(y)) + boundaries = [n_time * k - 1 for k in range(1, 120)] + assert d[boundaries].max() < 1.5 * np.delete(d, boundaries).max() + + +def test_sin_generator_no_drift_matches_fast_path(): + """freq_drift_rate=0 must reproduce the exact closed-form sinusoid.""" + fs = 1000.0 + producer = SinProducer(SinGeneratorSettings(fs=fs, n_time=100, n_ch=1, freq=10.0, freq_drift_rate=0.0)) + parts = [producer(AxisArray.TimeAxis(fs=fs, offset=i * 0.1)).data[:, 0] for i in range(10)] + y = np.concatenate(parts) + t = np.arange(len(y)) / fs + np.testing.assert_allclose(y, np.sin(2 * np.pi * 10.0 * t), rtol=1e-10, atol=1e-9)