MRG, ENH: Allow just line freq filtering (#7459)

* ENH: Allow just line freq filtering

* STY: Flake

doc/changes/latest.inc

@@ -27,6 +27,8 @@ Changelog
 
 - Add :func:`mne.chpi.compute_chpi_amplitudes`, :func:`mne.chpi.compute_chpi_locs`, and :func:`mne.chpi.compute_head_pos` to compute head positions from cHPI coil locations by `Eric Larson`_ and `Luke Bloy`_
 
+- Add ``allow_line_only`` option to :func:`mne.chpi.filter_chpi` to allow filtering line frequencies only in files that do not have cHPI information by `Eric Larson`_
+
 - Add :func:`mne.io.Raw.set_meas_date` by `Eric Larson`_
 
 - Add :meth:`mne.Epochs.as_type` to  allow remapping data in MEG channels to virtual magnetometer or gradiometer channels by `Sophie Herbst`_ and `Alex Gramfort`_

mne/chpi.py

@@ -25,6 +25,7 @@
 from scipy import linalg
 import itertools
 
+from .io.base import BaseRaw
 from .io.meas_info import _simplify_info
 from .io.pick import (pick_types, pick_channels, pick_channels_regexp,
                       pick_info)
@@ -214,10 +215,12 @@ def extract_chpi_locs_ctf(raw, verbose=None):
 # ############################################################################
 # Estimate positions from data
 @verbose
-def _get_hpi_info(info, verbose=None):
+def _get_hpi_info(info, allow_empty=False, verbose=None):
     """Get HPI information from raw."""
     if len(info['hpi_meas']) == 0 or \
             ('coil_freq' not in info['hpi_meas'][0]['hpi_coils'][0]):
+        if allow_empty:
+            return np.empty(0), None, np.empty(0)
         raise RuntimeError('Appropriate cHPI information not found in'
                            'info["hpi_meas"] and info["hpi_subsystem"], '
                            'cannot process cHPI')
@@ -435,16 +438,19 @@ def _fit_coil_order_dev_head_trans(dev_pnts, head_pnts):
 
 @verbose
 def _setup_hpi_amplitude_fitting(info, t_window, remove_aliased=False,
-                                 ext_order=1, verbose=None):
+                                 ext_order=1, allow_empty=False, verbose=None):
     """Generate HPI structure for HPI localization."""
     # grab basic info.
-    hpi_freqs, hpi_pick, hpi_ons = _get_hpi_info(info)
+    hpi_freqs, hpi_pick, hpi_ons = _get_hpi_info(info, allow_empty=allow_empty)
     _validate_type(t_window, (str, 'numeric'), 't_window')
     if isinstance(t_window, str):
         if t_window != 'auto':
             raise ValueError('t_window must be "auto" if a string, got %r'
                              % (t_window,))
-        t_window = max(5. / min(hpi_freqs), 1. / np.diff(hpi_freqs).min())
+        if len(hpi_freqs):
+            t_window = max(5. / min(hpi_freqs), 1. / np.diff(hpi_freqs).min())
+        else:
+            t_window = 0.2
     t_window = float(t_window)
     if t_window <= 0:
         raise ValueError('t_window (%s) must be > 0' % (t_window,))
@@ -1015,7 +1021,7 @@ def _chpi_locs_to_times_dig(chpi_locs):
 
 @verbose
 def filter_chpi(raw, include_line=True, t_step=0.01, t_window=None,
-                ext_order=1, verbose=None):
+                ext_order=1, allow_line_only=False, verbose=None):
     """Remove cHPI and line noise from data.
 
     .. note:: This function will only work properly if cHPI was on
@@ -1031,6 +1037,11 @@ def filter_chpi(raw, include_line=True, t_step=0.01, t_window=None,
         Time step to use for estimation, default is 0.01 (10 ms).
     %(chpi_t_window)s
     %(chpi_ext_order)s
+    allow_line_only : bool
+        If True, allow filtering line noise only. The default is False,
+        which only allows the function to run when cHPI information is present.
+
+        .. versionadded:: 0.20
     %(verbose)s
 
     Returns
@@ -1047,6 +1058,7 @@ def filter_chpi(raw, include_line=True, t_step=0.01, t_window=None,
 
     .. versionadded:: 0.12
     """
+    _validate_type(raw, BaseRaw, 'raw')
     if not raw.preload:
         raise RuntimeError('raw data must be preloaded')
     if t_window is None:
@@ -1058,8 +1070,12 @@ def filter_chpi(raw, include_line=True, t_step=0.01, t_window=None,
     if t_step <= 0:
         raise ValueError('t_step (%s) must be > 0' % (t_step,))
     n_step = int(np.ceil(t_step * raw.info['sfreq']))
-    hpi = _setup_hpi_amplitude_fitting(raw.info, t_window, remove_aliased=True,
-                                       ext_order=ext_order, verbose=False)
+    if include_line and raw.info['line_freq'] is None:
+        raise RuntimeError('include_line=True but raw.info["line_freq"] is '
+                           'None, consider setting it to the line frequency')
+    hpi = _setup_hpi_amplitude_fitting(
+        raw.info, t_window, remove_aliased=True, ext_order=ext_order,
+        allow_empty=allow_line_only, verbose=False)
 
     fit_idxs = np.arange(0, len(raw.times) + hpi['n_window'] // 2, n_step)
     n_freqs = len(hpi['freqs'])

mne/tests/test_chpi.py

@@ -27,13 +27,14 @@
 from mne.datasets import testing
 
 base_dir = op.join(op.dirname(__file__), '..', 'io', 'tests', 'data')
-test_fif_fname = op.join(base_dir, 'test_raw.fif')
 ctf_fname = op.join(base_dir, 'test_ctf_raw.fif')
 hp_fif_fname = op.join(base_dir, 'test_chpi_raw_sss.fif')
 hp_fname = op.join(base_dir, 'test_chpi_raw_hp.txt')
 raw_fname = op.join(base_dir, 'test_raw.fif')
 
 data_path = testing.data_path(download=False)
+sample_fname = op.join(
+    data_path, 'MEG', 'sample', 'sample_audvis_trunc_raw.fif')
 chpi_fif_fname = op.join(data_path, 'SSS', 'test_move_anon_raw.fif')
 pos_fname = op.join(data_path, 'SSS', 'test_move_anon_raw.pos')
 sss_fif_fname = op.join(data_path, 'SSS', 'test_move_anon_raw_sss.fif')
@@ -235,7 +236,7 @@ def test_calculate_chpi_positions_vv():
     _assert_quats(py_quats, mf_quats, dist_tol=0.001, angle_tol=0.7)
 
     # degenerate conditions
-    raw_no_chpi = read_raw_fif(test_fif_fname)
+    raw_no_chpi = read_raw_fif(sample_fname)
     with pytest.raises(RuntimeError, match='cHPI information not found'):
         _calculate_chpi_positions(raw_no_chpi)
     raw_bad = raw.copy()
@@ -494,12 +495,34 @@ def test_calculate_chpi_coil_locs_artemis():
     assert_array_less(1e-13, amps)
 
 
+def assert_suppressed(new, old, suppressed, retained):
+    """Assert that some frequencies are suppressed and others aren't."""
+    __tracebackhide__ = True
+    from scipy.signal import welch
+    picks = pick_types(new.info, meg='grad')
+    sfreq = new.info['sfreq']
+    new = new.get_data(picks)
+    old = old.get_data(picks)
+    f, new = welch(new, sfreq, 'hann', nperseg=1024)
+    _, old = welch(old, sfreq, 'hann', nperseg=1024)
+    new = np.median(new, axis=0)
+    old = np.median(old, axis=0)
+    for freqs, lim in ((suppressed, (10, 60)), (retained, (-3, 3))):
+        for freq in freqs:
+            fidx = np.argmin(np.abs(f - freq))
+            this_new = np.median(new[fidx])
+            this_old = np.median(old[fidx])
+            suppression = -10 * np.log10(this_new / this_old)
+            assert lim[0] < suppression < lim[1], freq
+
+
 @testing.requires_testing_data
-def test_chpi_subtraction():
+def test_chpi_subtraction_filter_chpi():
     """Test subtraction of cHPI signals."""
     raw = read_raw_fif(chpi_fif_fname, allow_maxshield='yes', preload=True)
     raw.info['bads'] = ['MEG0111']
     raw.del_proj()
+    raw_orig = raw.copy().crop(0, 16)
     with catch_logging() as log:
         with pytest.deprecated_call(match='"auto"'):
             filter_chpi(raw, include_line=False, verbose=True)
@@ -512,11 +535,38 @@ def test_chpi_subtraction():
     raw_c.pick_types(
         meg=True, eeg=True, eog=True, ecg=True, stim=True, misc=True)
     assert_meg_snr(raw, raw_c, 143, 624)
-
-    # Degenerate cases
-    raw_nohpi = read_raw_fif(test_fif_fname, preload=True)
+    # cHPI suppressed but not line freqs (or others)
+    assert_suppressed(raw, raw_orig, np.arange(83, 324, 60), [30, 60, 150])
+    raw = raw_orig.copy()
+    with catch_logging() as log:
+        with pytest.deprecated_call(match='"auto"'):
+            filter_chpi(raw, include_line=True, verbose=True)
+    log = log.getvalue()
+    assert '5 cHPI' in log
+    assert '6 line' in log
+    # cHPI and line freqs suppressed
+    suppressed = np.sort(np.concatenate([
+        np.arange(83, 324, 60), np.arange(60, 301, 60),
+    ]))
+    assert_suppressed(raw, raw_orig, suppressed, [30, 150])
+
+    # No HPI information
+    raw = read_raw_fif(sample_fname, preload=True)
+    raw_orig = raw.copy()
+    assert raw.info['line_freq'] is None
+    with pytest.raises(RuntimeError, match='line_freq.*consider setting it'):
+        filter_chpi(raw, t_window=0.2)
+    raw.info['line_freq'] = 60.
     with pytest.raises(RuntimeError, match='cHPI information not found'):
-        filter_chpi(raw_nohpi, t_window=0.2)
+        filter_chpi(raw, t_window=0.2)
+    # but this is allowed
+    with catch_logging() as log:
+        filter_chpi(raw, t_window='auto', allow_line_only=True, verbose=True)
+    log = log.getvalue()
+    assert '0 cHPI' in log
+    assert '1 line' in log
+    # Our one line freq suppressed but not others
+    assert_suppressed(raw, raw_orig, [60], [30, 45, 75])
 
     # When MaxFliter downsamples, like::
     #     $ maxfilter -nosss -ds 2 -f test_move_anon_raw.fif \