ENH: add line_bandwidth to ZapLine to subtract only the line band

docs/changes/devel/148.feature.rst

@@ -0,0 +1,2 @@
+Added ``line_bandwidth`` to ``ZapLine`` to subtract the removed components only within a band
+around the line harmonics, leaving the rest of the spectrum unchanged (weaker on strong or broad lines).

docs/zapline.md

@@ -61,6 +61,23 @@ remove what is left channel by channel with
 {class}`~mne_denoise.spectrum_interpolation.SpectrumInterpolation` or a notch
 filter.
 
+## Subtracting only the line band
+
+The removed components also carry broadband activity, so subtracting them
+whole can raise power away from the line on single channels. With
+`line_bandwidth`, an mne-denoise extension, only their part within
+`line_bandwidth` Hz of each targeted harmonic is subtracted, and the spectrum
+outside these bands is unchanged:
+
+```python
+model = ZapLine(sfreq=200.0, line_freq=60.0, line_bandwidth=1.0)  # 59-61 Hz
+```
+
+ZapLine selects the same components as without it. A line broader than the
+band is removed less completely. On short inputs, or epochs that are not
+contiguous, part of the line falls outside the band and is not removed; apply
+it to the continuous recording before epoching.
+
 ## References
 
 ```{footbibliography}

mne_denoise/_filtering.py

@@ -1,4 +1,4 @@
-"""Butterworth filter-design helpers."""
+"""Filter helpers shared across methods."""
 
 from __future__ import annotations
 
@@ -18,3 +18,19 @@ def design_butter_sos(
     ``fs=sfreq`` rather than by the caller.
     """
     return butter(order, freqs, btype=btype, fs=sfreq, output="sos")
+
+
+def line_band_mask(
+    fft_freqs: np.ndarray, freqs: float | np.ndarray, half_width: float
+) -> np.ndarray:
+    """Return the FFT bins within ``half_width`` Hz of any of ``freqs``.
+
+    A frequency whose band is narrower than the frequency resolution, and so
+    holds no bin, gets its nearest bin instead.
+    """
+    freqs = np.atleast_1d(freqs)
+    bins = fft_freqs[:, np.newaxis]
+    band = (bins >= freqs - half_width) & (bins <= freqs + half_width)
+    for column in np.flatnonzero(~band.any(axis=0)):
+        band[np.argmin(np.abs(fft_freqs - freqs[column])), column] = True
+    return band.any(axis=1)

mne_denoise/spectrum_interpolation.py

@@ -12,6 +12,7 @@
 
 from . import _mne
 from ._data import extract_data_from_mne, reconstruct_mne_object
+from ._filtering import line_band_mask
 from ._logging import logger, verbose
 from ._validation import (
     check_matching_sfreq,
@@ -104,13 +105,7 @@ def interpolate_spectrum(
         if not (0.0 < f < nyquist):
             continue
 
-        band = (fft_freqs >= f - bandwidth) & (fft_freqs <= f + bandwidth)
-        if not np.any(band):
-            # Band narrower than the frequency resolution: snap to nearest bin.
-            nearest = int(np.argmin(np.abs(fft_freqs - f)))
-            band = np.zeros_like(band)
-            band[nearest] = True
-
+        band = line_band_mask(fft_freqs, f, bandwidth)
         left = (fft_freqs >= f - bandwidth - neighbour_width) & (
             fft_freqs < f - bandwidth
         )

mne_denoise/zapline/core.py

@@ -12,8 +12,10 @@
     extract_data_from_mne,
     reconstruct_mne_object,
 )
+from .._filtering import line_band_mask
 from .._logging import logger, verbose
 from .._spatial import apply_spatial_transform
+from .._validation import check_positive_real
 
 # Inherit from DSS
 from ..dss._whitening import (
@@ -78,6 +80,9 @@ class ZapLine(DSS):
         Pre-whiten supported MNE data channels before processing.
     noise_cov : mne.Covariance or None, default=None
         Noise covariance used when whiten=True.
+    line_bandwidth : float or None, default=None
+        Half-width in Hz of the band around each targeted harmonic within which
+        the removed components are subtracted. None subtracts them entirely.
     verbose : bool, str, int, or None, default=None
         Logging level.
 
@@ -139,6 +144,7 @@ def __init__(
         crossfade: float = 0.0,
         whiten: bool = False,
         noise_cov=None,
+        line_bandwidth: float | None = None,
         verbose: bool | str | int | None = None,
     ):
         self.sfreq = float(sfreq)
@@ -149,6 +155,7 @@ def __init__(
         self.threshold = threshold
         self.knee_rel_floor = knee_rel_floor
         self.knee_min_ratio = knee_min_ratio
+        self.line_bandwidth = line_bandwidth
         self.adaptive_params = adaptive_params if adaptive_params is not None else {}
         self.verbose = verbose
 
@@ -230,6 +237,7 @@ def fit(
                 "Adaptive mode requires simultaneous fit and transform (local chunks). "
                 "Use fit_transform() instead."
             )
+        self._check_line_bandwidth()
 
         data, extracted_sfreq, _, orig_inst, _, ch_names = extract_data_from_mne(
             X,
@@ -355,6 +363,7 @@ def fit_transform(
         callback = _validate_callback(callback)
         if not self.adaptive:
             return super().fit_transform(X, y=y, callback=callback, **fit_params)
+        self._check_line_bandwidth()
 
         data, extracted_sfreq, mne_type, orig_inst, picks, ch_names = (
             extract_data_from_mne(
@@ -576,9 +585,30 @@ def _apply_standard_cleaning(self, data: np.ndarray) -> np.ndarray:
             return data.copy()
 
         # The smoother is linear and per channel: smooth the sources, not the data
-        sources = self._get_residual(self.filters_ @ data)
+        sources = self._line_part(
+            self._get_residual(self.filters_ @ data), self.bias._harmonic_freqs
+        )
         return data - self._artifact_mixing_ @ sources
 
+    def _check_line_bandwidth(self) -> float | None:
+        """Return line_bandwidth as a float after validation, or None."""
+        if self.line_bandwidth is None:
+            return None
+        return check_positive_real(self.line_bandwidth, name="line_bandwidth")
+
+    def _line_part(self, signal: np.ndarray, harmonics: np.ndarray) -> np.ndarray:
+        """Zero the Fourier bins of signal outside the bands around harmonics."""
+        half_width = self._check_line_bandwidth()
+        if half_width is None:
+            return signal
+        from scipy.fft import irfft, rfft, rfftfreq
+
+        n_times = signal.shape[-1]
+        freqs = rfftfreq(n_times, d=1.0 / self.sfreq)
+        spectrum = rfft(signal, axis=-1)
+        spectrum[..., ~line_band_mask(freqs, harmonics, half_width)] = 0
+        return irfft(spectrum, n=n_times, axis=-1)
+
     def _get_smooth_residual(self, data: np.ndarray, warn: bool = False):
         """Split data into smooth and residual components."""
         # Use self.sfreq directly
@@ -773,9 +803,27 @@ def _run_adaptive(
         finally:
             self._target_freq_ = None
 
+        # Segments select and subtract whole components; restricting their sum
+        # to the line bands over the whole recording keeps segment boundaries
+        # and the segments' spectral checks out of the result.
+        removed = data - current_data
+        if self.line_bandwidth is not None and all_chunk_metadata:
+            harmonics = np.concatenate(
+                [
+                    LineNoiseBias(
+                        freq=chunk["fine_freq"],
+                        sfreq=self.sfreq,
+                        n_harmonics=self.n_harmonics,
+                    )._harmonic_freqs
+                    for chunk in all_chunk_metadata
+                ]
+            )
+            removed = self._line_part(removed, harmonics)
+            current_data = data - removed
+
         return {
             "cleaned": current_data,
-            "removed": data - current_data,
+            "removed": removed,
             "n_removed": sum(c.get("n_removed", 0) for c in all_chunk_metadata),
             "line_freq": line_freqs[0] if line_freqs else 0,
             "line_freqs": tuple(float(freq) for freq in all_freqs_to_process),
@@ -839,6 +887,7 @@ def _process_segment(self, chunk: np.ndarray, start: int = 0) -> dict:
                 "threshold": current_sigma,
                 "adaptive": False,
                 "crossfade": 0.0,
+                "line_bandwidth": None,
                 "verbose": "WARNING",
             }
         )

tests/test_zapline_core.py

@@ -395,3 +395,158 @@ def test_zapline_exposes_only_n_select():
     assert "n_select" in estimator.get_params()
     assert ZapLine(sfreq=500.0, line_freq=60.0, n_select="auto").n_select == "auto"
     assert ZapLine(sfreq=500.0, line_freq=60.0, n_select=3).n_select == 3
+
+
+def _eeg_with_muscle_and_line(sfreq, line_amplitudes, *, duration=60.0, seed=0):
+    """Return 19-channel EEG-like data with muscle on 3 channels and a 60 Hz line.
+
+    ``line_amplitudes`` gives the relative amplitude of each 60 Hz harmonic.
+    """
+    rng = np.random.default_rng(seed)
+    n_channels, n_times = 19, int(sfreq * duration)
+    times = np.arange(n_times) / sfreq
+    brain = np.cumsum(rng.standard_normal((n_channels, n_times)), axis=1)
+    brain = signal.sosfiltfilt(
+        signal.butter(2, 0.5, "hp", fs=sfreq, output="sos"), brain
+    )
+    mixing = np.eye(n_channels) + 0.3 * rng.standard_normal((n_channels, n_channels))
+    data = mixing @ brain
+    muscle_band = signal.butter(
+        4, [20.0, 0.95 * sfreq / 2], "bp", fs=sfreq, output="sos"
+    )
+    data[:3] += 20.0 * signal.sosfiltfilt(
+        muscle_band, rng.standard_normal((3, n_times))
+    )
+    line = sum(
+        amplitude * np.sin(2 * np.pi * 60.0 * (harmonic + 1) * times + harmonic)
+        for harmonic, amplitude in enumerate(line_amplitudes)
+    )
+    return data + 30.0 * rng.uniform(0.5, 1.0, n_channels)[:, None] * line
+
+
+def _welch_db_change(before, after, sfreq):
+    """Return frequencies and per-channel Welch power change in dB."""
+    freqs, psd_before = signal.welch(before, fs=sfreq, nperseg=int(4 * sfreq))
+    _, psd_after = signal.welch(after, fs=sfreq, nperseg=int(4 * sfreq))
+    return freqs, 10 * np.log10(psd_after / psd_before)
+
+
+def _line_reduction_db(before, after, sfreq, freq):
+    """Return the drop in summed Welch power within 0.5 Hz of freq, in dB.
+
+    Unlike ``_power_at``, the 0.25 Hz resolution keeps the broadband power next
+    to the line, which only the whole-component subtraction removes, out of the
+    measure.
+    """
+    freqs, psd_before = signal.welch(before, fs=sfreq, nperseg=int(4 * sfreq))
+    _, psd_after = signal.welch(after, fs=sfreq, nperseg=int(4 * sfreq))
+    near = np.abs(freqs - freq) <= 0.5
+    return 10 * np.log10(psd_before[:, near].sum() / psd_after[:, near].sum())
+
+
+@pytest.mark.parametrize(
+    ("sfreq", "n_harmonics", "line_amplitudes"),
+    [
+        (200.0, None, (1.0,)),  # fractional period; only 60 Hz is below Nyquist
+        (500.0, 2, (1.0, 0.5)),
+    ],
+)
+def test_zapline_line_bandwidth_leaves_spectrum_outside_the_line_unchanged(
+    sfreq, n_harmonics, line_amplitudes
+):
+    """Line-only subtraction keeps the default's line removal without its broadband cost.
+
+    With three components, the default whole-component subtraction moves the
+    muscle activity of some channels onto others; subtracting only each
+    component's part within 1 Hz of the harmonics leaves every other Fourier bin
+    of the input unchanged.
+    """
+    data = _eeg_with_muscle_and_line(sfreq, line_amplitudes)
+    harmonics = 60.0 * np.arange(1, len(line_amplitudes) + 1)
+    params = {
+        "sfreq": sfreq,
+        "line_freq": 60.0,
+        "n_select": 3,
+        "n_harmonics": n_harmonics,
+    }
+    default = ZapLine(**params).fit_transform(data)
+    line_only = ZapLine(**params, line_bandwidth=1.0).fit_transform(data)
+
+    freqs, change = _welch_db_change(data, default, sfreq)
+    away = np.all(np.abs(freqs[:, None] - harmonics) > 1.5, axis=1) & (freqs >= 1.0)
+    assert change[:, away].max() > 3.0
+
+    fft_freqs = np.fft.rfftfreq(data.shape[-1], d=1.0 / sfreq)
+    outside = np.all(np.abs(fft_freqs[:, None] - harmonics) > 1.0, axis=1)
+    spectrum = np.fft.rfft(data, axis=-1)
+    assert_allclose(
+        np.fft.rfft(line_only, axis=-1)[:, outside],
+        spectrum[:, outside],
+        rtol=0,
+        atol=1e-9 * np.abs(spectrum).max(),
+    )
+
+    for freq in harmonics:
+        reduction = _line_reduction_db(data, line_only, sfreq, freq)
+        assert reduction > 20.0
+        assert reduction == pytest.approx(
+            _line_reduction_db(data, default, sfreq, freq), abs=0.5
+        )
+
+
+def test_zapline_line_bandwidth_in_adaptive_mode_keeps_the_segment_selection():
+    """Adaptive segments select as by default; only the line band is subtracted.
+
+    The segments' spectral checks see whole-component subtraction, so they
+    remove the same number of components, and the band is applied to the
+    removed signal of the whole recording, so no segment boundary leaks into
+    the spectrum outside the line bands.
+    """
+    sfreq = 200.0
+    data = _eeg_with_muscle_and_line(sfreq, (1.0,))
+    params = {
+        "sfreq": sfreq,
+        "line_freq": 60.0,
+        "adaptive": True,
+        "adaptive_params": {"min_chunk_len": 20.0},
+    }
+    default = ZapLine(**params)
+    cleaned_default = default.fit_transform(data)
+    line_only = ZapLine(**params, line_bandwidth=1.0)
+    cleaned = line_only.fit_transform(data)
+
+    assert len(line_only.adaptive_results_["chunk_info"]) > 1
+    assert line_only.n_removed_ == default.n_removed_
+
+    # The segments refine the line frequency to within 0.05 Hz.
+    fft_freqs = np.fft.rfftfreq(data.shape[-1], d=1.0 / sfreq)
+    outside = np.abs(fft_freqs - 60.0) > 1.1
+    spectrum = np.fft.rfft(data, axis=-1)
+    assert_allclose(
+        np.fft.rfft(cleaned, axis=-1)[:, outside],
+        spectrum[:, outside],
+        rtol=0,
+        atol=1e-9 * np.abs(spectrum).max(),
+    )
+    reduction = _line_reduction_db(data, cleaned, sfreq, 60.0)
+    assert reduction > 20.0
+    assert reduction == pytest.approx(
+        _line_reduction_db(data, cleaned_default, sfreq, 60.0), abs=0.5
+    )
+
+
+@pytest.mark.parametrize(
+    ("line_bandwidth", "error"),
+    [(0.0, ValueError), (-2.0, ValueError), (np.inf, ValueError), (True, TypeError)],
+)
+def test_zapline_rejects_invalid_line_bandwidth(minimal_data, line_bandwidth, error):
+    """line_bandwidth must be a positive, finite number of hertz, also when set later."""
+    data, sfreq = minimal_data["data"], minimal_data["sfreq"]
+    model = ZapLine(sfreq=sfreq, line_freq=50.0, line_bandwidth=line_bandwidth)
+    with pytest.raises(error, match="line_bandwidth"):
+        model.fit(data)
+
+    fitted = ZapLine(sfreq=sfreq, line_freq=50.0, n_select=1).fit(data)
+    fitted.set_params(line_bandwidth=line_bandwidth)
+    with pytest.raises(error, match="line_bandwidth"):
+        fitted.transform(data)