diff --git a/src/ezmsg/sigproc/adaptive_lnc.py b/src/ezmsg/sigproc/adaptive_lnc.py new file mode 100644 index 00000000..5bfdf14b --- /dev/null +++ b/src/ezmsg/sigproc/adaptive_lnc.py @@ -0,0 +1,564 @@ +"""Adaptive line-noise cancellation (LNC). + +An LMS adaptive filter that estimates and subtracts a line-frequency (e.g. +50/60 Hz mains) interferer, and optionally its harmonics, from a multichannel +signal. Quadrature references (sin/cos) at the line frequency -- and at each +harmonic ``k`` x the fundamental, phase-locked to it -- drive per-channel +adaptive weights that reconstruct the line noise; the summed reconstruction is +subtracted from the signal (Widrow & Stearns, *Adaptive Signal Processing*, LMS +noise cancellation with quadrature references). + +Because we observe the signal only in *samples*, the line shows up at a +normalised frequency of ``f_mains / fs`` cycles/sample. Both the device clock +(``fs``) and the mains frequency are imprecise and drift, so that ratio is +unknown and time-varying. The transformer therefore tracks it adaptively: + + * A numerically-controlled oscillator (NCO) generates the reference at an + angular frequency ``omega`` (rad/sample); harmonics are generated at + ``k * omega``, so a single oscillator serves every harmonic. + * A per-channel LMS adapts amplitude and phase at each harmonic frequency. + * A frequency-locked loop (FLL) nudges ``omega`` from the rotation of the + fundamental's adaptive weights, **pooled across all channels** for + robustness -- the + line frequency is a single global quantity, so every channel constrains + the same estimate. The pooled estimator (a power-weighted cross-product) + is invariant to static per-channel phase, so it needs no assumption about + inter-channel phase relationships. + +Both loops are parameterised by a **time constant in seconds**, so their +behaviour is independent of the (possibly wildly variable) chunk size: + + * ``adapt_time_constant`` is the settling time of the per-channel LMS; + internally ``mu = 2 / (adapt_time_constant * fs)``. + * ``freq_time_constant`` is the settling time of the FLL; the per-update gain + is derived from the elapsed time of each update interval, + ``beta = 1 - exp(-dt / freq_time_constant)`` (the same tau<->alpha mapping + EWMA uses), so chunking never changes the dynamics. It should exceed + ``adapt_time_constant`` so the two loops do not fight. + +The frequency estimate accumulates over an internal window of one mains period +before each update -- not a fixed chunk size, but a measurement window that +fills across however the caller chunks the stream (4, 30, 60, mixed samples). +Output is emitted for every chunk immediately; only the frequency *update* +waits for a full window, which keeps the rotation estimate clean even when +chunks are only a handful of samples. As a result the loop behaves identically +for a single offline buffer (it converges within that one buffer) and for a +fast sequence of tiny chunks, and tracking-on output is chunk-invariant. + +The sequential A/D also gives each channel a known sub-degree per-channel phase +offset, but at the line frequency it is negligible and it cancels in frequency +tracking, so it is *not* handled here -- the dedicated +``SamplingDelayAlignmentTransformer`` owns that correction broadband, upstream, +for every cross-channel step. + +Statefulness: expensive setup happens once in :meth:`_reset_state`; each chunk +runs only the per-sample recursion plus one pooled frequency update. With the +FLL disabled (``freq_time_constant`` None/<=0) the transformer reduces to a +fixed-frequency canceller whose output is independent of chunking. +""" + +from typing import Any + +import ezmsg.core as ez +import numpy as np +import numpy.typing as npt +import scipy.signal +from array_api_compat import array_namespace +from ezmsg.baseproc import ( + BaseStatefulTransformer, + BaseTransformerUnit, + processor_state, +) +from ezmsg.util.messages.axisarray import AxisArray +from ezmsg.util.messages.util import replace + +# Optional Apple-Silicon GPU backend. The canceller is an LTI SOS notch +# cascade (see `design_lnc_sos`), so on MLX arrays we dispatch to the Metal +# `sosfilt` kernel; everything else runs through scipy on the array's own +# namespace (numpy, cupy, ...). +try: # pragma: no cover - exercised only when mlx is installed + import mlx.core as _mx + + from ezmsg.sigproc.util.sosfilt_mlx_metal import sosfilt_mlx_metal as _sosfilt_mlx +except Exception: # pragma: no cover + _mx = None + _sosfilt_mlx = None + + +def _is_mlx(arr: object) -> bool: + """True if ``arr`` is an MLX array (and the MLX backend is available).""" + return _mx is not None and isinstance(arr, _mx.array) + + +def _namespace(arr: object) -> tuple[Any, bool]: + """Return ``(xp, is_mlx)`` for ``arr``: the MLX module for MLX arrays, else + the array's Array-API namespace (numpy, cupy, ...).""" + if _is_mlx(arr): + return _mx, True + return array_namespace(arr), False + + +def _to_numpy(arr: object) -> np.ndarray: + """Bring a small backend array to host numpy (for the scalar FLL math).""" + return np.asarray(arr) + + +class AdaptiveLNCSettings(ez.Settings): + """Settings for :class:`AdaptiveLNCTransformer`.""" + + line_freq: float = 60.0 + """Nominal line (mains) frequency in Hz (60 in N. America, 50 in EU). Seeds + the NCO; the true normalised frequency is then tracked by the FLL.""" + + num_harmonics: int = 1 + """Number of line-frequency harmonics to cancel, including the fundamental. + 1 = fundamental only. Harmonic ``k`` is generated phase-locked at ``k`` x + the single tracked fundamental, so one NCO/FLL serves them all; each + harmonic gets its own per-channel LMS weight. e.g. 5 also cancels + 120/180/240/300 Hz. Absent harmonics simply drive their weights to ~0.""" + + adapt_time_constant: float = 0.1 + """Settling time (seconds) of the amplitude/phase canceller (the LMS). + Smaller = faster tracking of amplitude/phase changes and a wider, noisier + notch; larger = a narrower, cleaner notch that adapts more slowly. Converted + internally to the LMS step size ``mu = 2 / (adapt_time_constant * fs)``. + Read live each chunk.""" + + freq_time_constant: float | None = 0.5 + """Settling time (seconds) of the frequency tracker (the FLL). ``None`` (or + <= 0) freezes the frequency at the nominal ``line_freq`` -- a fixed- + reference LMS. Should be larger than ``adapt_time_constant`` so the two + loops do not fight. Independent of chunk size (the per-update gain is + derived from elapsed time). Read live each chunk.""" + + control: float = 1.0 + """Cancellation gate in [0, 1]. 1.0 fully subtracts the estimate, 0.0 is + passthrough (weights keep adapting either way). Read live each chunk.""" + + cancel_method: str = "notch" + """How to remove the line. ``"notch"`` (default) applies the SOS notch + cascade -- a perfect null that also removes any signal at the line + frequency. ``"subtract"`` estimates the common-mode line (pooled global + phase + per-channel amplitude) once per window and subtracts *only* it, + preserving signal that is independent across channels at the line frequency. + ``"subtract"`` assumes a 1-D channel axis. (Per-channel sampling-delay + alignment is handled separately, upstream, by + ``SamplingDelayAlignmentTransformer``.)""" + + axis: str = "time" + """Name of the axis to filter along.""" + + +@processor_state +class AdaptiveLNCState: + """State for :class:`AdaptiveLNCTransformer`.""" + + omega: float = 0.0 + """Current NCO angular frequency in rad/sample (tracked by the FLL).""" + + phase: float = 0.0 + """NCO phase (rad) for the first sample of the next chunk; accumulates so + the demodulation reference is continuous across chunk boundaries.""" + + block_len: int = 0 + """FLL update window in samples (one mains period).""" + + samples_in_block: int = 0 + """Samples accumulated toward the next FLL update; carried across chunks so + the update grid is global, not per-chunk.""" + + zi: npt.NDArray | None = None + """SOS biquad filter state, carried across chunks. Backend-native layout: + ``(n_harm, 2, *sample_shape)`` for scipy, ``(n_harm, *sample_shape, 2)`` for + the MLX kernel.""" + + sos: npt.NDArray | None = None + """Cached notch coefficients, already on the working backend. Rebuilt only + when ``omega`` or ``mu`` changes (i.e. once per FLL window, not per chunk), + so streaming avoids redundant host->device transfers.""" + + sos_key: tuple | None = None + """``(omega, mu, n_harm)`` the cached ``sos`` was built for.""" + + z_acc_real: npt.NDArray | None = None + """Demodulation accumulator ``sum(removed * cos(phase))`` over the current + window (per channel); the line phasor's real part.""" + + z_acc_imag: npt.NDArray | None = None + """Demodulation accumulator ``sum(removed * sin(phase))`` over the current + window (per channel).""" + + z_phasor_prev: npt.NDArray | None = None + """Pooled per-channel line phasor (numpy complex) from the previous window; + the FLL measures rotation against it.""" + + line_phasor: complex = 0j + """Observable: pooled global line phasor at the last window end.""" + + # --- "subtract" mode state (per-harmonic; lists so MLX immutability is fine) + sub_z_real: list | None = None + """Per-harmonic demod accumulators ``sum(x * cos(k*phase))`` this window.""" + + sub_z_imag: list | None = None + """Per-harmonic demod accumulators ``sum(x * sin(k*phase))`` this window.""" + + sub_amp: list | None = None + """Committed per-harmonic per-channel line amplitude (backend arrays); + ``None`` until the first window completes.""" + + sub_phase_off: list | None = None + """Committed per-harmonic global phase ``theta_k`` (scalar floats), for + reconstruction.""" + + sub_yc_smooth: list | None = None + """Per-harmonic EWMA-smoothed per-channel line phasor (numpy complex). + Smoothing (time constant ``adapt_time_constant``) keeps the line estimate + from chasing transient in-band signal.""" + + +class AdaptiveLNCTransformer( + BaseStatefulTransformer[ + AdaptiveLNCSettings, + AxisArray, + AxisArray, + AdaptiveLNCState, + ] +): + """Quadrature-LMS line-noise canceller, implemented as its exact LTI + equivalent (an SOS notch cascade) with frequency tracking. + + **Cancellation.** A fixed-frequency quadrature-LMS canceller is exactly a + 2nd-order notch (Glover 1977; see :func:`design_lnc_sos`). Each harmonic is + one biquad notch at ``k * omega``; the cascade is applied with ``sosfilt`` + carrying state ``zi`` across chunks. The gated output is:: + + y_notch = sosfilt(design_lnc_sos(omega, mu, num_harmonics), x) + removed = x - y_notch # the line estimate + y = x - control * removed + + This replaces the former per-sample LMS recursion: it is vectorised, scales + with channel count, and dispatches to a GPU ``sosfilt`` (MLX/Metal) on MLX + arrays -- so the whole transformer is Array-API compatible. + + **Frequency tracking (FLL).** Once per window (one mains period) the NCO + frequency is nudged by the pooled rotation of the *demodulated line phasor* + between windows:: + + Z_c = sum_n removed_{c,n} * exp(-j * phase_n) # per-channel, per window + cross = sum_c Z_c * conj(Z_c_prev) # power-weighted, pooled + omega += beta * angle(cross) / window_len # rad/sample + + Demodulating the removed line gives the same frequency-error signal the old + weight-rotation detector did, without materialising adaptive weights. Static + per-channel phase cancels in ``cross``, so no inter-channel phase assumption + is needed. The window fills on a global grid carried in state, so the loop + is chunk-size independent: a single offline buffer converges within itself + and streamed output is chunk-invariant. ``mu = 2/(adapt_time_constant*fs)`` + and ``beta = 1 - exp(-window_dt/freq_time_constant)``. + """ + + # These are read fresh every chunk inside `_process`, so changing them must + # NOT reset the filter state. line_freq seeds omega, axis/delay model and + # num_harmonics shape the state, so those do force a reset. + NONRESET_SETTINGS_FIELDS = frozenset({"adapt_time_constant", "freq_time_constant", "control"}) + + def _hash_message(self, message: AxisArray) -> int: + ax_idx = message.get_axis_idx(self.settings.axis) + sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] + return hash((message.key, message.axes[self.settings.axis].gain, sample_shape)) + + def _reset_state(self, message: AxisArray) -> None: + ax_idx = message.get_axis_idx(self.settings.axis) + sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :] + xp, is_mlx = _namespace(message.data) + + fs = 1.0 / message.axes[self.settings.axis].gain + # Seed the NCO at the nominal normalised frequency; the FLL refines it. + self._state.omega = 2.0 * np.pi * self.settings.line_freq / fs + self._state.phase = 0.0 + + # Frequency-update window: one mains period, inferred from line_freq/fs. + self._state.block_len = max(1, int(round(fs / self.settings.line_freq))) + self._state.samples_in_block = 0 + + # SOS filter state: one biquad per harmonic. The two backends use + # different delay layouts, so `zi` is created per backend; the demod + # accumulators share their shape and go through `xp`. + n_harm = max(1, int(self.settings.num_harmonics)) + if is_mlx: + self._state.zi = _mx.zeros((n_harm,) + sample_shape + (2,)) + else: + self._state.zi = np.zeros((n_harm, 2) + sample_shape) + self._state.z_acc_real = xp.zeros(sample_shape) + self._state.z_acc_imag = xp.zeros(sample_shape) + self._state.z_phasor_prev = None + self._state.sos = None # built lazily on the working backend + self._state.sos_key = None + self._state.line_phasor = 0j + + # "subtract" mode: per-harmonic demod accumulators; estimate not yet known. + self._state.sub_z_real = [xp.zeros(sample_shape) for _ in range(n_harm)] + self._state.sub_z_imag = [xp.zeros(sample_shape) for _ in range(n_harm)] + self._state.sub_amp = None + self._state.sub_phase_off = None + self._state.sub_yc_smooth = None + + @staticmethod + def _cat(pieces: list, is_mlx: bool, xp: Any) -> npt.NDArray: + """Concatenate segment outputs along the time axis (backend-portable).""" + if len(pieces) == 1: + return pieces[0] + return _mx.concatenate(pieces, axis=0) if is_mlx else xp.concat(pieces, axis=0) + + def _ensure_sos(self, mu: float, n_harm: int, is_mlx: bool) -> npt.NDArray: + """Return the notch coefficients on the working backend, rebuilding only + when ``(omega, mu, n_harm)`` changes -- i.e. once per FLL window, not per + chunk. Avoids re-running the design and (on MLX) re-transferring the + coefficients to the device for every segment within a window.""" + st = self._state + key = (st.omega, mu, n_harm) + if st.sos is None or st.sos_key != key: + sos = design_lnc_sos(st.omega, mu, n_harm) # numpy float64 + st.sos = _mx.array(sos.astype(np.float32)) if is_mlx else sos + st.sos_key = key + return st.sos + + def _sosfilt(self, sos: npt.NDArray, x_seg: npt.NDArray, is_mlx: bool) -> npt.NDArray: + """Apply the notch cascade over one segment (time on axis 0), carrying + ``self._state.zi`` across calls. Dispatches to the MLX Metal kernel for + MLX arrays, else scipy.""" + if is_mlx: + x_t = _mx.moveaxis(x_seg, 0, x_seg.ndim - 1) # kernel wants time last + y_t, self._state.zi = _sosfilt_mlx(sos, x_t, zi=self._state.zi) + return _mx.moveaxis(y_t, y_t.ndim - 1, 0) + y, self._state.zi = scipy.signal.sosfilt(sos, x_seg, axis=0, zi=self._state.zi) + return y + + def _accumulate_demod(self, removed: npt.NDArray, seg_len: int, xp: Any) -> None: + """Demodulate the removed line against the fundamental NCO and add to + the current window's phasor accumulators (real/imag, per channel).""" + st = self._state + ph = st.phase + st.omega * xp.arange(seg_len) + bshape = (seg_len,) + (1,) * (removed.ndim - 1) + cos = xp.reshape(xp.cos(ph), bshape) + sin = xp.reshape(xp.sin(ph), bshape) + st.z_acc_real = st.z_acc_real + xp.sum(removed * cos, axis=0) + st.z_acc_imag = st.z_acc_imag + xp.sum(removed * sin, axis=0) + + def _fll_step(self, z_fund: np.ndarray, beta: float) -> None: + """Nudge ``omega`` by the pooled rotation of the fundamental complex line + phasor ``z_fund`` (numpy, per channel) since the previous window. Static + per-channel phase cancels in the cross-product, so no alignment needed.""" + st = self._state + if st.z_phasor_prev is not None: + cross = np.sum(z_fund * np.conj(st.z_phasor_prev)) + if cross != 0: + st.omega = st.omega + beta * float(np.angle(cross)) / st.block_len + st.z_phasor_prev = z_fund + + def _freq_update(self, beta: float) -> None: + """Notch-mode window boundary: run the FLL off the fundamental demod of + the removed line, refresh the observable phasor, reset accumulators.""" + st = self._state + # e^{-j phase} = cos - j sin, so Z = acc_real - j acc_imag (per channel). + z = _to_numpy(st.z_acc_real).reshape(-1) - 1j * _to_numpy(st.z_acc_imag).reshape(-1) + self._fll_step(z, beta) + st.line_phasor = complex(np.sum(z)) # observable pooled line phasor + + # Clear accumulators for the next window. + st.z_acc_real = st.z_acc_real * 0 + st.z_acc_imag = st.z_acc_imag * 0 + + # ----- "subtract" mode -------------------------------------------------- # + def _accumulate_subtract(self, x_seg: npt.NDArray, seg_len: int, n_harm: int, xp: Any) -> None: + """Demodulate the *input* at every harmonic into the per-harmonic window + accumulators (the rank-1 pool will separate the common-mode line from + independent in-band signal).""" + st = self._state + base = st.phase + st.omega * xp.arange(seg_len) + bshape = (seg_len,) + (1,) * (x_seg.ndim - 1) + for kidx in range(n_harm): + ph = (kidx + 1) * base + cos = xp.reshape(xp.cos(ph), bshape) + sin = xp.reshape(xp.sin(ph), bshape) + st.sub_z_real[kidx] = st.sub_z_real[kidx] + xp.sum(x_seg * cos, axis=0) + st.sub_z_imag[kidx] = st.sub_z_imag[kidx] + xp.sum(x_seg * sin, axis=0) + + def _commit_subtract(self, n_harm: int, beta: float, alpha: float, tracking: bool, xp: Any) -> None: + """Window boundary: demod -> EWMA-smooth (time constant + ``adapt_time_constant``, via ``alpha``) the per-harmonic per-channel line + phasor, then commit the rank-1 estimate (global phase + per-channel + amplitude) for the next window's reconstruction, optionally step the FLL, + and reset the accumulators.""" + st = self._state + scale = 2.0 / st.block_len # demod sum -> sinusoid amplitude (window is ~1 period) + smoothed = [] + amps, offs = [], [] + z0 = None + for kidx in range(n_harm): + z = scale * (_to_numpy(st.sub_z_real[kidx]).reshape(-1) - 1j * _to_numpy(st.sub_z_imag[kidx]).reshape(-1)) + if kidx == 0: + z0 = z + if st.sub_yc_smooth is not None: + z = (1.0 - alpha) * st.sub_yc_smooth[kidx] + alpha * z + smoothed.append(z) + # Global common-mode phase: amplitude-weighted pool (rank-1 direction). + theta = float(np.angle(np.sum(z * np.abs(z)))) if np.any(z) else 0.0 + a = np.clip(np.real(z * np.exp(-1j * theta)), 0.0, None) # per-channel amplitude + amps.append(xp.asarray(a)) + offs.append(theta) # global phase per harmonic (scalar) + st.sub_yc_smooth = smoothed + st.sub_amp = amps + st.sub_phase_off = offs + + if tracking and z0 is not None: + self._fll_step(z0, beta) + + st.sub_z_real = [z * 0 for z in st.sub_z_real] + st.sub_z_imag = [z * 0 for z in st.sub_z_imag] + + def _reconstruct(self, seg_len: int, n_harm: int, xp: Any) -> npt.NDArray | None: + """Reconstruct the committed common-mode line over one segment, using the + previous window's estimate. ``None`` before the first commit.""" + st = self._state + if st.sub_amp is None: + return None + base = st.phase + st.omega * xp.arange(seg_len) + line = None + for kidx in range(n_harm): + k = kidx + 1 + # arg[n] = k*base[n] + theta_k (global phase; same across channels) + arg = k * xp.reshape(base, (seg_len, 1)) + st.sub_phase_off[kidx] + term = st.sub_amp[kidx] * xp.cos(arg) + line = term if line is None else line + term + return line + + def _process(self, message: AxisArray) -> AxisArray: + ax_idx = message.get_axis_idx(self.settings.axis) + x_data = message.data + xp, is_mlx = _namespace(x_data) + moved = ax_idx != 0 + if moved: + x_data = xp.moveaxis(x_data, ax_idx, 0) + + n = x_data.shape[0] + st = self._state + dtype = x_data.dtype + fs = 1.0 / message.axes[self.settings.axis].gain + + # Time constants -> gains (independent of chunk size and fs). + # mu = 2 / (tau_adapt * fs); beta = 1 - exp(-window_dt / tau_freq). + mu = 2.0 / (self.settings.adapt_time_constant * fs) + ctrl = self.settings.control + n_harm = max(1, int(self.settings.num_harmonics)) + tau_freq = self.settings.freq_time_constant + tracking = tau_freq is not None and tau_freq > 0 + beta = 1.0 - np.exp(-(st.block_len / fs) / tau_freq) if tracking else 0.0 + + if self.settings.cancel_method == "subtract": + # Reconstruct-and-subtract: always walk the window grid (the per- + # window estimate updates there); the FLL steps omega only if + # tracking. The first window subtracts nothing while it learns. + # Line-estimate EWMA gain from adapt_time_constant (per window). + alpha = 1.0 - np.exp(-(st.block_len / fs) / self.settings.adapt_time_constant) + pieces = [] + pos = 0 + while pos < n: + seg_len = min(st.block_len - st.samples_in_block, n - pos) + x_seg = x_data[pos : pos + seg_len] + line_seg = self._reconstruct(seg_len, n_harm, xp) + pieces.append(x_seg if line_seg is None else x_seg - ctrl * line_seg) + self._accumulate_subtract(x_seg, seg_len, n_harm, xp) + st.phase = st.phase + st.omega * seg_len + st.samples_in_block += seg_len + pos += seg_len + if st.samples_in_block >= st.block_len: + self._commit_subtract(n_harm, beta, alpha, tracking, xp) + st.samples_in_block = 0 + y_out = self._cat(pieces, is_mlx, xp) + elif not tracking: + # Notch, frozen NCO: constant frequency, whole chunk in one sosfilt. + sos = self._ensure_sos(mu, n_harm, is_mlx) + y_notch = self._sosfilt(sos, x_data, is_mlx) + y_out = x_data - ctrl * (x_data - y_notch) + st.phase = st.phase + st.omega * n + else: + # Notch, tracking: walk segments bounded by the frequency-update + # window; when a window fills the FLL fires and omega may step. + pieces = [] + pos = 0 + while pos < n: + seg_len = min(st.block_len - st.samples_in_block, n - pos) + x_seg = x_data[pos : pos + seg_len] + sos = self._ensure_sos(mu, n_harm, is_mlx) + y_notch = self._sosfilt(sos, x_seg, is_mlx) + removed = x_seg - y_notch + pieces.append(x_seg - ctrl * removed) + self._accumulate_demod(removed, seg_len, xp) + st.phase = st.phase + st.omega * seg_len + st.samples_in_block += seg_len + pos += seg_len + if st.samples_in_block >= st.block_len: + self._freq_update(beta) + st.samples_in_block = 0 + y_out = self._cat(pieces, is_mlx, xp) + + if not is_mlx: + y_out = xp.astype(y_out, dtype) # scipy upcasts float32 -> float64 + if moved: + y_out = xp.moveaxis(y_out, 0, ax_idx) + return replace(message, data=y_out) + + +class AdaptiveLNC( + BaseTransformerUnit[ + AdaptiveLNCSettings, + AxisArray, + AxisArray, + AdaptiveLNCTransformer, + ] +): + SETTINGS = AdaptiveLNCSettings + + +def design_lnc_sos(omega: float, mu: float, num_harmonics: int = 1) -> npt.NDArray: + """SOS notch cascade equivalent to the quadrature-LMS line canceller. + + A fixed-frequency quadrature-LMS canceller (unit references, step ``mu``, + ``control=1``) is *exactly* a linear time-invariant 2nd-order notch (Glover + 1977 / Widrow). For the fundamental its input->output (``y = x - nr``) + transfer function is:: + + H(z) = (z^2 - 2cos(w)z + 1) / (z^2 - (2 - mu)cos(w)z + (1 - mu)) + + i.e. zeros exactly on the unit circle at ``e^{+/-j w}`` (a perfect notch) + and poles at radius ``sqrt(1 - mu)``. Harmonic ``k`` is the same notch at + ``k * omega``; the cascade of the per-harmonic sections approximates the + parallel multi-reference LMS to a small fraction of a dB (exact for a + single harmonic). Verified against the per-sample LMS to ~1e-12 (float64). + + This is the building block for an Array-API / GPU (MLX ``sosfilt``) + implementation that replaces the per-sample Python recursion. + + Parameters + ---------- + omega : float + Fundamental angular frequency in rad/sample (the FLL-tracked value). + mu : float + LMS step size, i.e. ``2 / (adapt_time_constant * fs)``. + num_harmonics : int + Number of harmonics (sections), ``k = 1 .. num_harmonics``. + + Returns + ------- + np.ndarray + ``(num_harmonics, 6)`` SOS array (``[b0, b1, b2, a0, a1, a2]`` per row, + ``a0 = 1``), float64. Apply with ``scipy.signal.sosfilt`` (or the MLX + ``sosfilt_mlx_metal``) carrying ``zi`` across chunks. + """ + n_harm = max(1, int(num_harmonics)) + sos = np.empty((n_harm, 6), dtype=np.float64) + for i in range(n_harm): + c = np.cos((i + 1) * omega) + sos[i] = (1.0, -2.0 * c, 1.0, 1.0, -(2.0 - mu) * c, 1.0 - mu) + return sos diff --git a/tests/unit/test_adaptive_lnc.py b/tests/unit/test_adaptive_lnc.py new file mode 100644 index 00000000..74e6c2b0 --- /dev/null +++ b/tests/unit/test_adaptive_lnc.py @@ -0,0 +1,452 @@ +"""Correctness lock for the adaptive line-noise canceller. + +The :class:`AdaptiveLNCTransformer` is a *stateful, chunked* LMS line-noise +filter with an optional frequency-locked loop (FLL). This test pins two layers: + +* **The fixed-frequency core** (FLL disabled, ``freq_time_constant=None``) against an + independent, deliberately slow sample-by-sample reference + (:func:`_reference_lnc`). With the loop frozen the NCO is a constant-frequency + oscillator, so streaming in arbitrary chunks must reproduce the single-pass + reference and must cancel a stationary line. +* **The frequency tracker** (FLL enabled) against a signal whose line sits at a + deliberately *offset* normalised frequency (simulating an off/drifting device + clock or mains): ``omega`` must converge to the true frequency and cancel it + far better than the frozen filter can. + +The reference here is intentionally minimal and self-contained. +""" + +from __future__ import annotations + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray, LinearAxis + +from ezmsg.sigproc.adaptive_lnc import ( + AdaptiveLNCSettings, + AdaptiveLNCTransformer, +) + +FS = 30000.0 +LINE_FREQ = 60.0 + + +# --------------------------------------------------------------------------- # +# Reference implementation (correct, not fast) # +# --------------------------------------------------------------------------- # +def _reference_lnc( + x: np.ndarray, + adapt_rate: float, + control: float = 1.0, + line_freq: float = LINE_FREQ, + fs: float = FS, +) -> np.ndarray: + """Sample-by-sample fixed-frequency LMS line-noise canceller, one channel. + + Quadrature references sin/cos at ``line_freq`` come from a continuous-phase + oscillator; a 2-tap weight adapts via LMS to reconstruct the line, which is + gated by ``control`` and subtracted:: + + nr = w_sin*sin + w_cos*cos + err = nr - x[n] + w -= err * ref * adapt_rate # for each of sin, cos + y[n] = x[n] - control * nr + """ + x = np.asarray(x, dtype=np.float32) + n = x.shape[0] + omega = 2.0 * np.pi * line_freq / fs + phases = omega * np.arange(n) + ref_sin = np.sin(phases).astype(np.float32) + ref_cos = np.cos(phases).astype(np.float32) + + corrected = np.empty(n, dtype=np.float32) + w_sin = np.float32(0.0) + w_cos = np.float32(0.0) + mu = np.float32(adapt_rate) + ctrl = np.float32(control) + + for i in range(n): + rs = ref_sin[i] + rc = ref_cos[i] + nr = w_sin * rs + w_cos * rc + err = nr - x[i] + w_sin = w_sin - err * (rs * mu) + w_cos = w_cos - err * (rc * mu) + corrected[i] = x[i] - nr * ctrl + return corrected + + +def _reference_lnc_multichannel(X: np.ndarray, adapt_rate: float, control: float = 1.0, **kw) -> np.ndarray: + """Apply :func:`_reference_lnc` per channel. ``X`` is (n_samples, n_ch).""" + X = np.atleast_2d(np.asarray(X, dtype=np.float32)) + out = np.empty_like(X) + for c in range(X.shape[1]): + out[:, c] = _reference_lnc(X[:, c], adapt_rate, control, **kw) + return out + + +# --------------------------------------------------------------------------- # +# Helpers # +# --------------------------------------------------------------------------- # +def _make_axisarray(data: np.ndarray, offset: float = 0.0) -> AxisArray: + """Wrap (n_samples, n_ch) data as a time/ch AxisArray at FS.""" + return AxisArray( + data=data, + dims=["time", "ch"], + axes={"time": LinearAxis(offset=offset, gain=1.0 / FS)}, + key="test_lnc", + ) + + +def _stream_in_chunks(proc: AdaptiveLNCTransformer, data: np.ndarray, chunk_sizes: list[int]) -> np.ndarray: + """Feed ``data`` (n_samples, n_ch) through ``proc`` in the given chunks.""" + outputs = [] + start = 0 + for size in chunk_sizes: + chunk = data[start : start + size] + msg = _make_axisarray(chunk, offset=start / FS) + outputs.append(proc(msg).data) + start += size + assert start == data.shape[0], "chunk sizes must cover the whole signal" + return np.concatenate(outputs, axis=0) + + +def _line_signal(n: int, n_ch: int, omega: float, amp: float, seed: int) -> np.ndarray: + """(n, n_ch): 10 Hz 'neural' signal + a line at angular freq ``omega`` + (rad/sample) with per-channel phase + broadband noise.""" + rng = np.random.default_rng(seed) + k = np.arange(n) + t = k / FS + out = np.zeros((n, n_ch), dtype=np.float32) + for c in range(n_ch): + sig = 50.0 * np.sin(2 * np.pi * 10 * t + 0.3 * c) + noise = amp * np.sin(omega * k + 0.7 + 0.5 * c) + out[:, c] = sig + noise + rng.normal(0, 5, n) + return out + + +def _synthetic_mixture(n: int, n_ch: int = 1, seed: int = 0) -> np.ndarray: + """Mixture with the line exactly at the nominal LINE_FREQ/FS frequency.""" + omega = 2.0 * np.pi * LINE_FREQ / FS + return _line_signal(n, n_ch, omega, amp=200.0, seed=seed) + + +def _band_magnitude(sig_col: np.ndarray, omega: float) -> float: + """Magnitude of ``sig_col`` at angular frequency ``omega`` over its 2nd + half (after convergence), via a windowed DFT bin at that exact frequency.""" + half = sig_col.shape[0] // 2 + seg = sig_col[half:] + win = np.hanning(seg.shape[0]) + k = np.arange(seg.shape[0]) + return float(np.abs(np.sum(seg * win * np.exp(-1j * omega * k)))) + + +def adaptive_lnc( + line_freq: float = 60.0, + num_harmonics: int = 1, + adapt_time_constant: float = 0.1, + freq_time_constant: float | None = 0.5, + control: float = 1.0, + cancel_method: str = "notch", + axis: str = "time", +) -> AdaptiveLNCTransformer: + """Construct an :class:`AdaptiveLNCTransformer` with the given parameters.""" + return AdaptiveLNCTransformer( + settings=AdaptiveLNCSettings( + line_freq=line_freq, + num_harmonics=num_harmonics, + adapt_time_constant=adapt_time_constant, + freq_time_constant=freq_time_constant, + control=control, + cancel_method=cancel_method, + axis=axis, + ) + ) + + +# --------------------------------------------------------------------------- # +# Fixed-frequency core (FLL disabled) # +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize( + "chunk_sizes", + [ + [6000], # single chunk + [1] * 6000, # one sample at a time + [30] * 200, # uniform small chunks + [333, 1, 1000, 2666, 1000, 1000], # irregular chunks + ], + ids=["whole", "single-sample", "uniform-30", "irregular"], +) +def test_streaming_matches_reference(chunk_sizes): + """With the FLL frozen, chunked streaming reproduces the single-pass + reference to floating-point noise, for any chunking.""" + n = int(sum(chunk_sizes)) + mixture = _synthetic_mixture(n, n_ch=4) + adapt_rate = 1e-3 # raw LMS step for the reference + + expected = _reference_lnc_multichannel(mixture, adapt_rate, control=1.0) + + # Transformer takes a time constant; mu = 2 / (tau * fs), so tau here + # reproduces the reference's mu exactly. Tracking frozen. + proc = adaptive_lnc( + line_freq=LINE_FREQ, + adapt_time_constant=2.0 / (adapt_rate * FS), + freq_time_constant=None, + ) + got = _stream_in_chunks(proc, mixture, chunk_sizes) + + np.testing.assert_allclose(got, expected, rtol=1e-4, atol=1e-2) + + +def test_chunking_is_invariant(): + """With the FLL frozen, different chunkings yield the same output.""" + n = 6000 + mixture = _synthetic_mixture(n, n_ch=3, seed=1) + + out_whole = _stream_in_chunks(adaptive_lnc(freq_time_constant=None), mixture, [n]) + out_split = _stream_in_chunks(adaptive_lnc(freq_time_constant=None), mixture, [7, 13, 480] + [500] * 11) + np.testing.assert_allclose(out_whole, out_split, rtol=1e-4, atol=1e-2) + + +def test_control_zero_is_passthrough(): + """control=0 subtracts nothing: output equals input exactly.""" + n = 3000 + mixture = _synthetic_mixture(n, n_ch=2, seed=2) + proc = adaptive_lnc(control=0.0, freq_time_constant=None) + got = _stream_in_chunks(proc, mixture, [100] * 30) + np.testing.assert_array_equal(got, mixture.astype(np.float32)) + + +def test_cancels_stationary_line(): + """A line exactly at the nominal frequency is cancelled by the frozen + filter (>20 dB after convergence).""" + n = int(FS * 1.0) + mixture = _synthetic_mixture(n, n_ch=1, seed=3) + omega = 2.0 * np.pi * LINE_FREQ / FS + proc = adaptive_lnc(line_freq=LINE_FREQ, freq_time_constant=None) + corrected = _stream_in_chunks(proc, mixture, [300] * (n // 300)) + + before = _band_magnitude(mixture[:, 0], omega) + after = _band_magnitude(corrected[:, 0], omega) + reduction_db = 20 * np.log10(after / before) + assert reduction_db < -20.0, f"only {reduction_db:.1f} dB of rejection" + + +def test_line_freq_change_reseeds_nco(): + """Changing line_freq reseeds the NCO frequency (state reset).""" + n = 1000 + mixture = _synthetic_mixture(n, n_ch=1, seed=4) + proc = adaptive_lnc(line_freq=60.0, freq_time_constant=None) + _stream_in_chunks(proc, mixture, [n]) + omega_60 = proc._state.omega + + proc.update_settings(AdaptiveLNCSettings(line_freq=50.0, freq_time_constant=None)) + _stream_in_chunks(proc, mixture, [n]) + omega_50 = proc._state.omega + + assert omega_60 == pytest.approx(2 * np.pi * 60.0 / FS) + assert omega_50 == pytest.approx(2 * np.pi * 50.0 / FS) + + +# --------------------------------------------------------------------------- # +# Frequency tracking (FLL enabled) # +# --------------------------------------------------------------------------- # +@pytest.mark.parametrize("n_ch", [1, 16], ids=["single", "pooled-16ch"]) +def test_tracks_offset_frequency(n_ch): + """When the real line sits off the nominal frequency, the FLL converges + omega to it and cancels far better than the frozen filter.""" + n = int(FS * 2.0) # 2 s + omega_nominal = 2.0 * np.pi * LINE_FREQ / FS + omega_true = omega_nominal * 1.01 # ~0.6 Hz mains / clock offset + mixture = _line_signal(n, n_ch, omega_true, amp=200.0, seed=7) + chunks = [200] * (n // 200) + + tracking = adaptive_lnc(line_freq=LINE_FREQ, adapt_time_constant=0.03, freq_time_constant=0.1) + out_track = _stream_in_chunks(tracking, mixture, chunks) + + frozen = adaptive_lnc(line_freq=LINE_FREQ, freq_time_constant=None) + out_frozen = _stream_in_chunks(frozen, mixture, chunks) + + # omega converged close to the true line frequency. + assert tracking._state.omega == pytest.approx(omega_true, rel=2e-3) + + # ...and the residual at the true frequency is far smaller than frozen. + res_track = _band_magnitude(out_track[:, 0], omega_true) + res_frozen = _band_magnitude(out_frozen[:, 0], omega_true) + assert res_track < 0.2 * res_frozen + + +def test_single_shot_tracks_within_one_chunk(): + """A single offline buffer cleans itself: the FLL walks it in blocks and + converges omega to the offset line frequency using only that buffer.""" + n = int(FS * 2.0) + omega_nominal = 2.0 * np.pi * LINE_FREQ / FS + omega_true = omega_nominal * 1.01 + mixture = _line_signal(n, n_ch=8, omega=omega_true, amp=200.0, seed=9) + + proc = adaptive_lnc(line_freq=LINE_FREQ, adapt_time_constant=0.03, freq_time_constant=0.1) + _stream_in_chunks(proc, mixture, [n]) # one shot + assert proc._state.omega == pytest.approx(omega_true, rel=2e-3) + + +def test_tracking_is_chunk_invariant(): + """With the global block grid, tracking-on output no longer depends on how + the stream is chunked (matches to floating-point noise).""" + n = int(FS * 1.0) + omega_true = (2.0 * np.pi * LINE_FREQ / FS) * 1.01 + mixture = _line_signal(n, n_ch=4, omega=omega_true, amp=200.0, seed=10) + + split = [37, 211, 752] + [500] * 58 + assert sum(split) == n + out_whole = _stream_in_chunks(adaptive_lnc(freq_time_constant=0.1), mixture, [n]) + out_split = _stream_in_chunks(adaptive_lnc(freq_time_constant=0.1), mixture, split) + np.testing.assert_allclose(out_whole, out_split, rtol=1e-3, atol=1e-1) + + +def test_pooling_does_not_require_phase_alignment(): + """The pooled estimator must track even when channels carry the line at + very different phases (it uses a static-phase-invariant cross-product).""" + n = int(FS * 2.0) + omega_true = (2.0 * np.pi * LINE_FREQ / FS) * 0.99 + # Wildly different per-channel phases (0.5*c radians across 16 ch). + mixture = _line_signal(n, 16, omega_true, amp=200.0, seed=11) + proc = adaptive_lnc(line_freq=LINE_FREQ, adapt_time_constant=0.03, freq_time_constant=0.1) + _stream_in_chunks(proc, mixture, [200] * (n // 200)) + assert proc._state.omega == pytest.approx(omega_true, rel=3e-3) + + +# --------------------------------------------------------------------------- # +# Harmonics # +# --------------------------------------------------------------------------- # +def test_num_harmonics_cancels_harmonic_content(): + """A non-sinusoidal line (fundamental + 3rd harmonic): num_harmonics=1 + leaves the harmonic, num_harmonics=3 removes it too.""" + n = int(FS * 2.0) + k = np.arange(n) + omega1 = 2.0 * np.pi * LINE_FREQ / FS + rng = np.random.default_rng(21) + # fundamental + strong 3rd harmonic (like the real 'bio' recording) + line = 200.0 * np.sin(omega1 * k) + 120.0 * np.sin(3 * omega1 * k + 0.5) + sig = 50.0 * np.sin(2 * np.pi * 10 * k / FS) + mixture = (sig + line + rng.normal(0, 5, n)).astype(np.float32)[:, None] + + chunks = [200] * (n // 200) + fund = _stream_in_chunks(adaptive_lnc(LINE_FREQ, num_harmonics=1, adapt_time_constant=0.03), mixture, chunks) + harm = _stream_in_chunks(adaptive_lnc(LINE_FREQ, num_harmonics=3, adapt_time_constant=0.03), mixture, chunks) + + def red(y, hz): + return 20 * np.log10( + _band_magnitude(y[:, 0], 2 * np.pi * hz / FS) / _band_magnitude(mixture[:, 0], 2 * np.pi * hz / FS) + ) + + # Both kill the fundamental. + assert red(fund, 60) < -20.0 + assert red(harm, 60) < -20.0 + # Only the harmonic-aware filter kills the 3rd harmonic. + assert red(fund, 180) > -3.0 # essentially untouched + assert red(harm, 180) < -20.0 # strongly suppressed + + +# --------------------------------------------------------------------------- # +# Common-mode reconstruct-and-subtract (cancel_method="subtract") # +# --------------------------------------------------------------------------- # +def _common_mode_scene(n, nch, seed): + """(x, line, sig): a common-mode 60 Hz line (sub-degree tau_c phases, + per-channel amplitude) plus an *independent* per-channel 60 Hz burst to + preserve. (Sampling-delay phases are tiny at 60 Hz; alignment, if needed, + is a separate upstream stage.)""" + rng = np.random.default_rng(seed) + t = np.arange(n) / FS + slot = np.arange(nch) % 32 + tau = slot * (64.0 / 66.0e6) + w = 2 * np.pi * LINE_FREQ + line = (200.0 * (1 + 0.2 * rng.standard_normal(nch)))[None, :] * np.cos(w * t[:, None] + 0.7 + w * tau[None, :]) + psi = rng.uniform(0, 2 * np.pi, nch) + env = np.exp(-(((t - n / FS / 2) / 0.2) ** 2))[:, None] # mid-record burst + sig = (25.0 * np.cos(w * t[:, None] + psi[None, :])) * env + x = (line + sig + 5.0 * rng.standard_normal((n, nch))).astype(np.float32) + return x, line.astype(np.float32), sig.astype(np.float32) + + +def _fraction(y, ref): + """Mean per-channel projection of ``y`` onto a known component ``ref``.""" + return float(np.mean([(y[:, c] @ ref[:, c]) / (ref[:, c] @ ref[:, c]) for c in range(ref.shape[1])])) + + +def test_cancel_method_defaults_to_notch(): + assert adaptive_lnc().settings.cancel_method == "notch" + + +def test_subtract_preserves_inband_signal_that_notch_destroys(): + """Both remove the common-mode line, but subtract keeps the independent + in-band signal the notch erases. Each mode uses its natural + adapt_time_constant: a notch wide enough to catch the line, and a slow + line-estimate for subtract (so it ignores the transient signal).""" + n, nch = int(FS * 2.0), 16 + x, line, sig = _common_mode_scene(n, nch, seed=7) + chunks = [250] * (n // 250) + + def run(method, atc): + proc = adaptive_lnc( + LINE_FREQ, + adapt_time_constant=atc, + freq_time_constant=None, + cancel_method=method, + ) + return _stream_in_chunks(proc, x, chunks) + + y_notch = run("notch", 0.1) # ~3 Hz notch — covers the burst + y_sub = run("subtract", 0.5) # slow line estimate — ignores the burst + + # Both strongly remove the (ground-truth) common-mode line. + assert _fraction(y_notch, line) < 0.2 + assert _fraction(y_sub, line) < 0.2 + + # The discriminator: notch erases the in-band signal; subtract preserves it. + assert _fraction(y_notch, sig) < 0.25 + assert _fraction(y_sub, sig) > 0.6 + + +# --------------------------------------------------------------------------- # +# Array-API / MLX backend # +# --------------------------------------------------------------------------- # +def test_mlx_backend_matches_numpy(): + """The MLX (GPU) backend reproduces the numpy result: same cancellation and + the FLL converges to the same frequency.""" + mx = pytest.importorskip("mlx.core") + n = int(FS * 1.0) + omega_true = (2.0 * np.pi * LINE_FREQ / FS) * 1.01 + mixture = _line_signal(n, 8, omega_true, amp=200.0, seed=12).astype(np.float32) + mixture += (100.0 * np.sin(3 * omega_true * np.arange(n))[:, None]).astype(np.float32) + chunks = [200] * (n // 200) + + def make(): + return adaptive_lnc(LINE_FREQ, num_harmonics=3, adapt_time_constant=0.05, freq_time_constant=0.1) + + p_np = make() + y_np = _stream_in_chunks(p_np, mixture, chunks) + + p_mx = make() + outs, start, last = [], 0, None + for size in chunks: + d = mx.array(mixture[start : start + size]) + msg = AxisArray( + data=d, + dims=["time", "ch"], + axes={"time": LinearAxis(offset=start / FS, gain=1.0 / FS)}, + key="test_lnc", + ) + last = p_mx(msg).data + outs.append(np.array(last)) + start += size + y_mx = np.concatenate(outs, axis=0) + + # Output stays on the MLX backend (same array type and dtype as the input). + assert isinstance(last, mx.array), f"expected mlx output, got {type(last)}" + assert last.dtype == mx.float32 + + # float32 (MLX) vs float64 (scipy) sosfilt on a near-unit-circle notch: + # compare at signal scale rather than element-wise near the cancelled zeros. + rel = np.max(np.abs(y_mx - y_np)) / np.max(np.abs(mixture)) + assert rel < 1e-2, f"MLX vs numpy rel diff {rel:.2e}" + assert p_mx._state.omega == pytest.approx(p_np._state.omega, rel=1e-3) diff --git a/tests/unit/test_adaptive_lnc_sos_equivalence.py b/tests/unit/test_adaptive_lnc_sos_equivalence.py new file mode 100644 index 00000000..85e725ec --- /dev/null +++ b/tests/unit/test_adaptive_lnc_sos_equivalence.py @@ -0,0 +1,109 @@ +"""Locks the LMS <-> LTI-notch equivalence underpinning the SOS/MLX rewrite. + +A fixed-frequency quadrature-LMS line canceller is exactly a 2nd-order notch +(Glover 1977 / Widrow). :func:`design_lnc_sos` encodes that mapping; this test +proves it reproduces the per-sample LMS so the recursion can be replaced by +``sosfilt`` (and thus an Array-API / MLX backend) without changing behaviour. + +* Single harmonic: **exact** (to float64 round-off). +* Multi-harmonic: the SOS cascade matches the parallel LMS dB reduction to a + fraction of a dB (the cascade is sequential, the LMS parallel; they differ + only by tiny cross-harmonic coupling). +""" + +from __future__ import annotations + +import numpy as np +import pytest +import scipy.signal as sps + +from ezmsg.sigproc.adaptive_lnc import design_lnc_sos + +FS = 30000.0 +LINE_FREQ = 60.0 + + +def _lms_output(x: np.ndarray, tau: float, num_harmonics: int = 1) -> np.ndarray: + """Independent per-sample quadrature-LMS canceller (fixed frequency, + control=1), one channel -- the ground truth the SOS form must reproduce. + + This is the original recursion, kept here (not via the transformer, which + now uses the SOS form) so the test still validates the equivalence. + """ + mu = 2.0 / (tau * FS) + omega = 2.0 * np.pi * LINE_FREQ / FS + k = np.arange(1, num_harmonics + 1) + w_sin = np.zeros(num_harmonics) + w_cos = np.zeros(num_harmonics) + y = np.empty_like(x) + for n in range(x.shape[0]): + ph = k * (omega * n) # continuous NCO phase, harmonic k at k*phase + rs, rc = np.sin(ph), np.cos(ph) + nr = np.sum(w_sin * rs + w_cos * rc) + err = nr - x[n] + w_sin -= mu * rs * err + w_cos -= mu * rc * err + y[n] = x[n] - nr + return y + + +def _sos_output(x: np.ndarray, tau: float, num_harmonics: int = 1) -> np.ndarray: + omega = 2.0 * np.pi * LINE_FREQ / FS + mu = 2.0 / (tau * FS) + sos = design_lnc_sos(omega, mu, num_harmonics=num_harmonics) + return sps.sosfilt(sos, x) + + +def _line_db(x: np.ndarray, y: np.ndarray, hz: float) -> float: + def mag(s): + h = len(s) // 2 + seg = s[h:] * np.hanning(len(s) - h) + f = np.fft.rfftfreq(len(seg), 1 / FS) + return np.abs(np.fft.rfft(seg))[np.argmin(np.abs(f - hz))] + + return 20 * np.log10(mag(y) / mag(x)) + + +@pytest.mark.parametrize("tau", [0.02, 0.05, 0.1]) +def test_single_harmonic_sos_matches_lms_exactly(tau): + """The notch biquad reproduces the per-sample LMS to float64 round-off.""" + rng = np.random.default_rng(0) + k = np.arange(30000) + x = ( + 200 * np.sin(2 * np.pi * LINE_FREQ / FS * k) + 40 * np.sin(2 * np.pi * 10 / FS * k) + rng.normal(0, 8, k.size) + ).astype(np.float64) + + y_lms = _lms_output(x, tau) + y_sos = _sos_output(x, tau) + rel = np.max(np.abs(y_lms - y_sos)) / np.max(np.abs(x)) + assert rel < 1e-9, f"LMS vs SOS rel diff {rel:.2e} at tau={tau}" + + +def test_pole_radius_matches_mu(): + """Notch poles sit at radius sqrt(1 - mu), inside the unit circle.""" + omega = 2.0 * np.pi * LINE_FREQ / FS + mu = 2.0 / (0.1 * FS) + sos = design_lnc_sos(omega, mu, num_harmonics=1) + # a = [1, -(2-mu)cos w, (1-mu)] -> pole product = 1-mu = r^2 + a2 = sos[0, 5] + assert a2 == pytest.approx(1 - mu) + assert np.sqrt(a2) < 1.0 # stable + # zeros on the unit circle (perfect notch): b = [1, -2cos w, 1] + assert sos[0, 2] == pytest.approx(1.0) + + +def test_multi_harmonic_cascade_matches_lms_db(): + """Cascade SOS and parallel LMS cancel each harmonic to within ~0.5 dB.""" + rng = np.random.default_rng(1) + k = np.arange(30000) + x = ( + 200 * np.sin(2 * np.pi * 60 / FS * k) + + 100 * np.sin(2 * np.pi * 180 / FS * k + 0.5) + + 40 * np.sin(2 * np.pi * 10 / FS * k) + + rng.normal(0, 8, k.size) + ).astype(np.float64) + + y_lms = _lms_output(x, 0.05, num_harmonics=3) + y_sos = _sos_output(x, 0.05, num_harmonics=3) + for hz in (60.0, 180.0): + assert _line_db(x, y_lms, hz) == pytest.approx(_line_db(x, y_sos, hz), abs=0.5)