MRG: Add convenience function to "prepare" emptyroom raw file for maxwell_filter

doc/changes/latest.inc

@@ -35,6 +35,7 @@ Enhancements
 
 - It is now possible to compute inverse solutions with restricted source orientations using discrete forward models (:gh:`10464` by `Marijn van Vliet`_)
 
+- The new function :func:`mne.preprocessing.maxwell_filter_prepare_emptyroom` simplifies the preconditioning of an empty-room recording for our Maxwell filtering operations (:gh:`10533` by `Richard Höchenberger`_ and `Eric Larson`_)
 
 Bugs
 ~~~~

doc/preprocessing.rst

@@ -92,6 +92,7 @@ Projections:
    infomax
    equalize_bads
    maxwell_filter
+   maxwell_filter_prepare_emptyroom
    oversampled_temporal_projection
    peak_finder
    read_ica

mne/io/base.py

@@ -660,8 +660,8 @@ def set_annotations(self, annotations, emit_warning=True,
         annotations : instance of mne.Annotations | None
             Annotations to set. If None, the annotations is defined
             but empty.
-        emit_warning : bool
-            Whether to emit warnings when cropping or omitting annotations.
+        %(emit_warning)s
+            The default is True.
         %(on_missing_ch_names)s
         %(verbose)s
 

mne/preprocessing/__init__.py

@@ -20,7 +20,7 @@
 from .infomax_ import infomax
 from .stim import fix_stim_artifact
 from .maxwell import (maxwell_filter, find_bad_channels_maxwell,
-                      compute_maxwell_basis)
+                      compute_maxwell_basis, maxwell_filter_prepare_emptyroom)
 from .realign import realign_raw
 from .xdawn import Xdawn
 from ._csd import compute_current_source_density

mne/preprocessing/maxwell.py

@@ -41,6 +41,144 @@
 # differences between algorithms
 
 
+@verbose
+def maxwell_filter_prepare_emptyroom(
+    raw_er, *, raw, bads='from_raw', annotations='from_raw', meas_date='keep',
+    emit_warning=False, verbose=None
+):
+    """Prepare an empty-room recording for Maxwell filtering.
+
+    Empty-room data by default lacks certain properties that are required to
+    ensure running :func:`~mne.preprocessing.maxwell_filter` will process the
+    empty-room recording the same way as the experimental data. This function
+    preconditions an empty-room raw data instance accordingly so it can be used
+    for Maxwell filtering. Please see the ``Notes`` section for details.
+
+    Parameters
+    ----------
+    raw_er : instance of Raw
+        The empty-room recording. It will not be modified.
+    raw : instance of Raw
+        The experimental recording, typically this will be the reference run
+        used for Maxwell filtering.
+    bads : 'from_raw' | 'union' | 'keep'
+        How to populate the list of bad channel names to be injected into
+        the empty-room recording. If ``'from_raw'`` (default) the list of bad
+        channels will be overwritten with that of ``raw``. If ``'union'``, will
+        use the union of bad channels in ``raw`` and ``raw_er``. Note that
+        this may lead to additional bad channels in the empty-room in
+        comparison to the experimental recording. If ``'keep'``, don't alter
+        the existing list of bad channels.
+
+        .. note::
+           Non-MEG channels are silently dropped from the list of bads.
+    annotations : 'from_raw' | 'union' | 'keep'
+        Whether to copy the annotations over from ``raw`` (default),
+        use the union of the annotations, or to keep them unchanged.
+    meas_date : 'keep' | 'from_raw'
+        Whether to transfer the measurement date from ``raw`` or to keep
+        it as is (default). If you intend to manually transfer annotations
+        from ``raw`` **after** running this function, you should set this to
+        ``'from_raw'``.
+    %(emit_warning)s
+        Unlike :meth:`raw.set_annotations <mne.io.Raw.set_annotations>`, the
+        default here is ``False``, as empty-room recordings are often shorter
+        than raw.
+    %(verbose)s
+
+    Returns
+    -------
+    raw_er_prepared : instance of Raw
+        A copy of the passed empty-room recording, ready for Maxwell filtering.
+
+    Notes
+    -----
+    This function will:
+
+    * Compile the list of bad channels according to the ``bads`` parameter.
+    * Inject the device-to-head transformation matrix from the experimental
+      recording into the empty-room recording.
+    * Set the following properties of the empty-room recording to match the
+      experimental recording:
+
+      * Montage
+      * ``raw.first_time`` and ``raw.first_samp``
+
+    * Adjust annotations according to the ``annotations`` parameter.
+    * Adjust the measurement date according to the ``meas_date`` parameter.
+
+    .. versionadded:: 1.1
+    """  # noqa: E501
+    _validate_type(item=raw_er, types=BaseRaw, item_name='raw_er')
+    _validate_type(item=raw, types=BaseRaw, item_name='raw')
+    _validate_type(item=bads, types=str, item_name='bads')
+    _check_option(
+        parameter='bads', value=bads,
+        allowed_values=['from_raw', 'union', 'keep']
+    )
+    _validate_type(item=annotations, types=str, item_name='annotations')
+    _check_option(
+        parameter='annotations', value=annotations,
+        allowed_values=['from_raw', 'union', 'keep']
+    )
+    _validate_type(item=meas_date, types=str, item_name='meas_date')
+    _check_option(
+        parameter='meas_date', value=annotations,
+        allowed_values=['from_raw', 'keep']
+    )
+
+    raw_er_prepared = raw_er.copy()
+    del raw_er  # just to be sure
+
+    # handle bads; only keep MEG channels
+    if bads == 'from_raw':
+        bads = raw.info['bads']
+    elif bads == 'union':
+        bads = sorted(
+            set(raw.info['bads'] + raw_er_prepared.info['bads'])
+        )
+    elif bads == 'keep':
+        bads = raw_er_prepared.info['bads']
+
+    bads = [ch_name for ch_name in bads
+            if ch_name.startswith('MEG')]
+    raw_er_prepared.info['bads'] = bads
+
+    # handle dev_head_t
+    raw_er_prepared.info['dev_head_t'] = raw.info['dev_head_t']
+
+    # handle montage
+    montage = raw.get_montage()
+    raw_er_prepared.set_montage(montage)
+
+    # handle first_samp
+    raw_er_prepared.annotations.onset += (
+        raw.first_time - raw_er_prepared.first_time
+    )
+    raw_er_prepared._cropped_samp = raw._cropped_samp
+
+    # handle annotations
+    if annotations != 'keep':
+        er_annot = raw_er_prepared.annotations
+        if annotations == 'from_raw':
+            er_annot.delete(np.arange(len(er_annot)))
+        er_annot.append(
+            raw.annotations.onset,
+            raw.annotations.duration,
+            raw.annotations.description,
+            raw.annotations.ch_names
+        )
+        if raw_er_prepared.info['meas_date'] is None:
+            er_annot.onset -= raw_er_prepared.first_time
+        raw_er_prepared.set_annotations(er_annot, emit_warning)
+
+    # handle measurement date
+    if meas_date == 'from_raw':
+        raw_er_prepared.set_meas_date(raw.info['meas_date'])
+
+    return raw_er_prepared
+
+
 # Changes to arguments here should also be made in find_bad_channels_maxwell
 @verbose
 def maxwell_filter(raw, origin='auto', int_order=8, ext_order=3,

mne/preprocessing/tests/test_maxwell.py

@@ -24,7 +24,9 @@
                     BaseRaw, read_raw_ctf)
 from mne.io.constants import FIFF
 from mne.preprocessing import (maxwell_filter, find_bad_channels_maxwell,
-                               annotate_amplitude, compute_maxwell_basis)
+                               annotate_amplitude, compute_maxwell_basis,
+                               maxwell_filter_prepare_emptyroom,
+                               annotate_movement)
 from mne.preprocessing.maxwell import (
     _get_n_moments, _sss_basis_basic, _sh_complex_to_real,
     _sh_real_to_complex, _sh_negate, _bases_complex_to_real, _trans_sss_basis,
@@ -1443,3 +1445,112 @@ def test_compute_maxwell_basis(regularize, n):
     xform = S[:, :n_use_in] @ pS[:n_use_in]
     got = xform @ raw.pick_types(meg=True, exclude='bads').get_data()
     assert_allclose(got, want)
+
+
+@testing.requires_testing_data
+@pytest.mark.parametrize('bads', ('from_raw', 'union', 'keep'))
+def test_prepare_emptyroom_bads(bads):
+    """Test prepare_emptyroom."""
+    raw = read_raw_fif(raw_fname, allow_maxshield='yes', verbose=False)
+    names = [name for name in raw.ch_names if 'EEG' not in name]
+    raw.pick_channels(names)
+    raw_er = read_raw_fif(erm_fname, allow_maxshield='yes', verbose=False)
+    raw_er.pick_channels(names)
+    assert raw.ch_names == raw_er.ch_names
+    assert raw_er.info['dev_head_t'] is None
+    assert raw.info['dev_head_t'] is not None
+    raw_er.set_montage(None)
+
+    if bads == 'from_raw':
+        raw_bads_orig = ['MEG0113', 'MEG2313']
+        raw_er_bads_orig = []
+    elif bads == 'union':
+        raw_bads_orig = ['MEG0113']
+        raw_er_bads_orig = ['MEG2313']
+    elif bads == 'keep':
+        raw_bads_orig = []
+        raw_er_bads_orig = ['MEG0113', 'MEG2313']
+
+    raw.info['bads'] = raw_bads_orig
+    raw_er.info['bads'] = raw_er_bads_orig
+
+    raw_er_prepared = maxwell_filter_prepare_emptyroom(
+        raw_er=raw_er,
+        raw=raw,
+        bads=bads
+    )
+    assert raw_er_prepared.info['bads'] == ['MEG0113', 'MEG2313']
+    assert raw_er_prepared.info['dev_head_t'] == raw.info['dev_head_t']
+
+    montage_expected = raw.copy().pick_types(meg=True).get_montage()
+    assert raw_er_prepared.get_montage().dig == montage_expected.dig
+
+    # Ensure the originals were not modified
+    assert raw.info['bads'] == raw_bads_orig
+    assert raw_er.info['bads'] == raw_er_bads_orig
+    assert raw_er.info['dev_head_t'] is None
+    assert raw_er.get_montage() is None
+
+
+@testing.requires_testing_data
+@pytest.mark.slowtest  # lots of params
+@pytest.mark.parametrize('set_annot_when', ('before', 'after'))
+@pytest.mark.parametrize('raw_meas_date', ('orig', None))
+@pytest.mark.parametrize('raw_er_meas_date', ('orig', None))
+def test_prepare_emptyroom_annot_first_samp(set_annot_when, raw_meas_date,
+                                            raw_er_meas_date):
+    """Test prepare_emptyroom."""
+    raw = read_raw_fif(raw_fname, allow_maxshield='yes', verbose=False)
+    raw_er = read_raw_fif(erm_fname, allow_maxshield='yes', verbose=False)
+    names = raw.ch_names[:3]  # make it faster
+    raw.pick_channels(names)
+    raw_er.pick_channels(names)
+    assert raw.ch_names == raw_er.ch_names
+    assert raw.info['meas_date'] != raw_er.info['meas_date']
+    if raw_meas_date is None:
+        raw.set_meas_date(None)
+    if raw_er_meas_date is None:
+        raw_er.set_meas_date(None)
+    # to make life easier, make it the same duration
+    n_rep = max(int(np.ceil(len(raw.times) / len(raw_er.times))), 1)
+    raw_er = mne.concatenate_raws([raw_er] * n_rep).crop(0, raw.times[-1])
+    assert_allclose(raw.times, raw_er.times)
+    raw_er_first_samp_orig = raw_er.first_samp
+    assert len(raw.annotations) == 0
+    pos = mne.chpi.read_head_pos(pos_fname)
+    annot, _ = annotate_movement(raw, pos, 1.)
+    # Add an annotation right at the beginning and end to make sure nothing
+    # gets cropped
+    onset = raw.times[[0, -1]]
+    duration = 1. / raw.info['sfreq']
+    annot.append(onset + raw.first_time * (raw.info['meas_date'] is not None),
+                 duration, ['BAD_CUSTOM'])
+    want_annot = 7  # 5 from annotate_movement plus our first and last samps
+    if set_annot_when == 'before':
+        raw.set_annotations(annot)
+        meas_date = 'keep'
+        want_date = raw_er.info['meas_date']
+    else:
+        assert set_annot_when == 'after'
+        meas_date = 'from_raw'
+        want_date = raw.info['meas_date']
+    raw_er_prepared = maxwell_filter_prepare_emptyroom(
+        raw_er=raw_er, raw=raw, meas_date=meas_date, emit_warning=True)
+    assert raw_er.first_samp == raw_er_first_samp_orig
+    assert raw_er_prepared.info['meas_date'] == want_date
+    assert raw_er_prepared.first_samp == raw.first_samp
+
+    # Ensure (movement) annotations carry over regardless of whether they're
+    # set before or after preparation
+    assert len(annot) == want_annot
+    if set_annot_when == 'after':
+        raw.set_annotations(annot)
+        raw_er_prepared.set_annotations(annot)
+    assert len(raw.annotations) == want_annot
+    prop_bad = np.isnan(
+        raw.get_data([0], reject_by_annotation='nan')).mean()
+    assert 0.3 < prop_bad < 0.4
+    assert len(raw_er_prepared.annotations) == want_annot
+    prop_bad_er = np.isnan(
+        raw_er_prepared.get_data([0], reject_by_annotation='nan')).mean()
+    assert_allclose(prop_bad, prop_bad_er)

mne/utils/docs.py

@@ -887,6 +887,11 @@ def _reflow_param_docstring(docstring, has_first_line=True, width=75):
 ``'max'``, and ``'auto'``.
 """
 
+docdict['emit_warning'] = """
+emit_warning : bool
+    Whether to emit warnings when cropping or omitting annotations.
+"""
+
 docdict['epochs_preload'] = """
     Load all epochs from disk when creating the object
     or wait before accessing each epoch (more memory