Fix find_bad_channels_maxwell and maxwell_filter cHPI indexing

doc/changes/dev/14142.bugfix.rst

@@ -0,0 +1 @@
+Fix bug in :func:`mne.preprocessing.find_bad_channels_maxwell` and :func:`mne.preprocessing.maxwell_filter`: data is processed in chunks, but when ``head_pos`` was provided, the head positions from the beginning of the recording were used for each chunk rather than from the chunk's time window. In ``find_bad_channels_maxwell`` this affected every interval after the first, and in ``maxwell_filter`` every segment following a segment skipped due to ``skip_by_annotation``, as well as every chunk containing no head position update at all, by `Christian Brodbeck`_.

mne/_ola.py

@@ -23,6 +23,12 @@ class _Interp2:
         arrays that must be interpolated.
     interp : str
         Can be 'zero', 'linear', 'hann', or 'cos2' (same as hann).
+    offset : int
+        The position of the first point that will be fed, in the same units as
+        ``control_points``. Use it to process a segment that does not start at the
+        beginning of the signal the control points refer to: feeding ``n`` points then
+        yields exactly what a continuous pass would yield for ``offset`` to
+        ``offset + n``, interpolation phase included.
 
     Notes
     -----
@@ -42,7 +48,9 @@ class _Interp2:
 
     """
 
-    def __init__(self, control_points, values, interp="hann", *, name="Interp2"):
+    def __init__(
+        self, control_points, values, interp="hann", *, name="Interp2", offset=0
+    ):
         # set up interpolation
         self.control_points = np.array(control_points, int).ravel()
         if not np.array_equal(np.unique(self.control_points), self.control_points):
@@ -76,8 +84,13 @@ def val(pt):
             values = val
         self.values = values
         self.n_last = None
-        self._position = 0  # start at zero
-        self._left_idx = 0
+        self._position = offset = _ensure_int(offset, "offset")
+        if offset < 0:
+            raise ValueError(f"offset must be non-negative, got {offset}")
+        # The last control point at or before offset is the one in effect there, so
+        # feeding resumes mid-interval with the correct interpolation phase.
+        left_idx = np.searchsorted(self.control_points, offset, "right") - 1
+        self._left_idx = max(left_idx, 0)
         self._left = self._right = self._use_interp = None
         self.name = name
         known_types = ("cos2", "linear", "zero", "hann")
@@ -94,9 +107,13 @@ def feed_generator(self, n_pts):
         logger.debug(f"    ~ {self.name} Feed {n_pts} ({self._position}-{stop})")
         used = np.zeros(n_pts, bool)
         if self._left is None:  # first one
-            logger.debug(f"    ~   {self.name} Eval @ 0 ({self.control_points[0]})")
-            self._left = self.values(self.control_points[0])
-            if len(self.control_points) == 1:
+            left_idx = self._left_idx
+            logger.debug(
+                f"    ~   {self.name} Eval @ {left_idx} "
+                f"({self.control_points[left_idx]})"
+            )
+            self._left = self.values(self.control_points[left_idx])
+            if left_idx == len(self.control_points) - 1:  # nothing to interpolate to
                 self._right = self._left
         n_used = 0
 
@@ -244,6 +261,10 @@ class _COLA:
         The window to use. Default is "hann".
     tol : float
         The tolerance for COLA checking.
+    offset : int
+        The index of the first sample that will be fed. Use it to process a segment
+        that does not start at the beginning of the signal that ``process`` is based
+        on: ``offset`` is added to the ``start`` and ``stop`` handed to ``process``.
 
     Notes
     -----
@@ -278,8 +299,10 @@ def __init__(
         tol=1e-10,
         *,
         name="COLA",
+        offset=0,
         verbose=None,
     ):
+        self._offset = _ensure_int(offset, "offset")
         n_samples = _ensure_int(n_samples, "n_samples")
         n_overlap = _ensure_int(n_overlap, "n_overlap")
         n_total = _ensure_int(n_total, "n_total")
@@ -389,12 +412,12 @@ def feed(self, *datas, verbose=None, **kwargs):
                 this_window = np.pad(
                     self._window, (0, this_len - len(this_window)), "constant"
                 )
-                for offset in range(self._step, len(this_window), self._step):
-                    n_use = len(this_window) - offset
-                    this_window[offset:] += self._window[:n_use]
+                for shift in range(self._step, len(this_window), self._step):
+                    n_use = len(this_window) - shift
+                    this_window[shift:] += self._window[:n_use]
             if self._idx == 0:
-                for offset in range(self._n_samples - self._step, 0, -self._step):
-                    this_window[:offset] += self._window[-offset:]
+                for n_use in range(self._n_samples - self._step, 0, -self._step):
+                    this_window[:n_use] += self._window[-n_use:]
             this_proc = [in_[..., :this_len].copy() for in_ in self._in_buffers]
             logger.debug(
                 f"    * {self.name}[:] Processing {start}:{stop} "
@@ -406,7 +429,12 @@ def feed(self, *datas, verbose=None, **kwargs):
                 raise RuntimeError("internal indexing error")
             start = self._store.idx
             stop = self._store.idx + this_len
-            outs = self._process(*this_proc, start=start, stop=stop, **kwargs)
+            outs = self._process(
+                *this_proc,
+                start=start + self._offset,
+                stop=stop + self._offset,
+                **kwargs,
+            )
             if self._out_buffers is None:
                 max_len = np.max(self.stops - self.starts)
                 self._out_buffers = [

mne/preprocessing/maxwell.py

@@ -741,6 +741,7 @@ def _run_maxwell_filter(
     st_fixed,
     st_overlap,
     mc,
+    raw_offset=0,  # time offset of ``raw`` relative to ``mc``
 ):
     # Eventually find_bad_channels_maxwell could be sped up by moving this
     # outside the loop (e.g., in the prep function) but regularization depends
@@ -770,8 +771,8 @@ def _run_maxwell_filter(
     if not 0.0 < st_duration <= max_samps + 1.0:
         raise ValueError(
             f"st_duration ({st_duration / sfreq:0.1f}s) must be between 0 and the "
-            "longest contiguous duration of the data "
-            "({max_samps / sfreq:0.1f}s)."
+            f"longest contiguous duration of the data "
+            f"({max_samps / sfreq:0.1f}s)."
         )
 
     # This must be initialized inside _run_maxwell_filter because
@@ -781,6 +782,10 @@ def _run_maxwell_filter(
 
     # Process each valid block of data separately
     for onset, end in zip(onsets, ends):
+        # head positions are indexed relative to the recording, but onset and end are
+        # relative to raw, which can itself be a chunk of the recording
+        segment_offset = raw_offset + onset
+        mc.set_offset(segment_offset)
         n = end - onset
         assert n > 0
         tsss_valid = n >= st_duration
@@ -808,6 +813,7 @@ def _run_maxwell_filter(
             sfreq,
             window,
             name="tSSS-COLA",
+            offset=segment_offset,
         )
 
         # Generate time points to break up data into equal-length windows
@@ -865,6 +871,10 @@ class _MoveComp:
     """Perform movement compensation."""
 
     def __init__(self, pos, head_frame, raw, interp, reconstruct):
+        #   pos[0]: (n_pos, 4, 4): the dev_head_t transformation matrices
+        #   pos[1]: (n_pos,): sample indices into the recording, starting at 0
+        #   pos[2]: (n_pos, 9): rotation quaternion (:3), translation (3:6),
+        #       goodness of fit, error and velocity (6:9)
         self.pos = pos
         self.sfreq = raw.info["sfreq"]
         self.interp = interp
@@ -891,23 +901,40 @@ def get_decomp_by_offset(self, offset):
         return op_sss, op_in, op_resid
 
     def initialize(self, get_decomp, dev_head_t, S_recon):
-        """Secondary initialization."""
-        self.smooth = _Interp2(
-            self.pos[1],
-            self.get_decomp_by_offset,
-            interp=self.interp,
-            name="MC",
-        )
+        """Secondary initialization.
+
+        Call :meth:`set_offset` before feeding data.
+        """
         _, _, pS_decomp, self.reg_moments_0, _ = get_decomp(dev_head_t, t=0.0)
         self.n_good = pS_decomp.shape[1]
         self.S_recon = S_recon
-        self.offset = 0
         self.get_decomp = get_decomp
         # For the average passes
         self.last_avg_quat = np.nan * np.ones(6)
+        self.smooth = None  # set_offset positions us in the recording
+
+    def set_offset(self, offset):
+        """Position at the given sample of the recording to process a segment there.
+
+        ``pos`` is indexed relative to the start of the recording, so a segment that
+        does not begin there has to be told where it does, both to read the right head
+        positions and to resume interpolation with the right phase.
+        """
+        self.offset = offset
+        self.smooth = _Interp2(
+            self.pos[1],
+            self.get_decomp_by_offset,
+            interp=self.interp,
+            name="MC",
+            offset=offset,
+        )
 
     def get_avg_op(self, *, start, stop):
-        """Apply an average transformation over the next interval."""
+        """Apply an average transformation over the next interval.
+
+        ``start`` and ``stop`` are relative to the start of the recording, like
+        ``offset``.
+        """
         n_positions, avg_quat = _trans_lims(self.pos, start, stop)[1:]
         if not np.allclose(avg_quat, self.last_avg_quat, atol=1e-7):
             self.last_avg_quat = avg_quat
@@ -931,6 +958,7 @@ def get_avg_op(self, *, start, stop):
         return self.op_in_avg, self.op_resid_avg, n_positions
 
     def feed(self, data, good_mask, st_only):
+        assert self.smooth is not None  # set_offset must be called first
         n_samp = data.shape[1]
         pos_data, n_pos = _trans_lims(
             self.pos, self.offset, self.offset + data.shape[-1]
@@ -967,7 +995,8 @@ def feed(self, data, good_mask, st_only):
 
 def _trans_lims(pos, start, stop):
     """Get all trans and limits we need."""
-    pos_idx = np.arange(*np.searchsorted(pos[1], [start, stop]))
+    start_idx, stop_idx = np.searchsorted(pos[1], [start, stop])
+    pos_idx = np.arange(start_idx, stop_idx)
     used = np.zeros(stop - start, bool)
     quats = np.empty((9, stop - start))
     n_positions = len(pos_idx)
@@ -979,7 +1008,9 @@ def _trans_lims(pos, start, stop):
             rel_stop = rel_stop - start
             if rel_start == rel_stop:
                 continue  # our first pos occurs on first time sample
-            this_quat = pos[2][max(pos_idx[0] - 1 if len(pos_idx) else 0, 0)]
+            # the last position at or before start is the one in effect there, also
+            # when the window contains no position at all (pos_idx is empty)
+            this_quat = pos[2][max(start_idx - 1, 0)]
             n_positions += 1
         else:
             rel_start = pos[1][pos_idx[ti]] - start
@@ -2941,7 +2972,7 @@ def find_bad_channels_maxwell(
             chunk_raw._data[:] = orig_data
             delta = chunk_raw.get_data(these_picks)
             with use_log_level(_verbose_safe_false()):
-                _run_maxwell_filter(chunk_raw, copy=False, **params)
+                _run_maxwell_filter(chunk_raw, copy=False, raw_offset=start, **params)
 
             if n_iter == 1 and len(chunk_flats):
                 logger.info(

mne/preprocessing/tests/test_maxwell.py

@@ -51,9 +51,11 @@
     _sh_negate,
     _sh_real_to_complex,
     _sss_basis_basic,
+    _trans_lims,
     _trans_sss_basis,
 )
 from mne.rank import _compute_rank_int, _get_rank_sss, compute_rank
+from mne.transforms import rot_to_quat
 from mne.utils import (
     _record_warnings,
     assert_meg_snr,
@@ -208,6 +210,21 @@ def read_crop(fname, lims=(0, None)):
     return raw.copy().crop(*lims)
 
 
+def _linear_head_pos(raw, n_pos):
+    """Get head positions one second apart, translating steadily along z.
+
+    The steady motion makes reading the positions from the wrong time window
+    give a different result.
+    """
+    trans = raw.info["dev_head_t"]["trans"]
+    head_pos = np.zeros((n_pos, 10))
+    head_pos[:, 0] = raw._first_time + np.arange(float(n_pos))
+    head_pos[:, 1:4] = rot_to_quat(trans[:3, :3])
+    head_pos[:, 4:7] = trans[:3, 3]
+    head_pos[:, 6] += np.arange(n_pos) * 5e-3  # 5 mm/s
+    return head_pos
+
+
 # For backward compat and to be most like MaxFilter, we make "maxwell_filter"
 # the one that behaves like MaxFilter. _maxwell_filter is left to
 # be the advanced/better one.
@@ -1792,6 +1809,56 @@ def test_mf_skips():
     assert_allclose(data_sc, data_cs, atol=1e-20)
 
 
+@pytest.mark.slowtest
+@testing.requires_testing_data
+@pytest.mark.parametrize("st_duration", (None, 2.0))
+def test_mf_skips_head_pos(st_duration):
+    """Test that segments after a skip use their own head positions."""
+    raw = read_raw_fif(raw_fname, allow_maxshield="yes")
+    raw.pick("meg", exclude=()).crop(0, 16).load_data()
+    head_pos = _linear_head_pos(raw, 17)
+    # A 2 s skip, leaving segments of 3 s and 11 s. The latter must stay above the 10 s
+    # chunk size that st_duration=None falls back to.
+    raw.set_annotations(
+        mne.Annotations(
+            onset=[raw._first_time + 3.0],
+            duration=[2.0],
+            description=["bad_acq_skip"],
+            orig_time=raw.info["meas_date"],
+        )
+    )
+    kwargs = dict(
+        origin=(0.0, 0.0, 0.04),
+        regularize=None,
+        bad_condition="ignore",
+        st_duration=st_duration,
+    )
+    # use the default mc_interp="hann" rather than the "zero" of this module's
+    # maxwell_filter, so that each sample blends two head positions and picking the
+    # wrong one on either side of the interval shows up
+    raw_sss = _maxwell_filter_ola(raw, head_pos=head_pos, **kwargs)
+    # Processing the second segment on its own must give the same result as processing
+    # it as part of the whole recording.
+    raw_crop = raw.copy().crop(5.0).set_annotations(None)
+    raw_crop_sss = _maxwell_filter_ola(raw_crop, head_pos=head_pos[5:], **kwargs)
+    for picks in ("meg", "chpi"):  # chpi holds the head positions written back out
+        assert_allclose(
+            raw_sss.get_data(picks, tmin=5.0), raw_crop_sss.get_data(picks), atol=1e-20
+        )
+
+
+def test_trans_lims_sparse_pos():
+    """Test windows containing no head position update."""
+    # head positions sampled more coarsely than buffer_size_sec leave whole windows
+    # without an update, and those must hold the last position, not the first one
+    pos = [None, np.array([0, 100, 200]), np.zeros((3, 9))]
+    pos[2][:, 5] = [0.0, 0.1, 0.2]  # z translation
+    for start, want in ((0, 0.0), (50, 0.0), (100, 0.1), (150, 0.1), (250, 0.2)):
+        quats, n_positions, _ = _trans_lims(pos, start, start + 50)
+        assert_allclose(np.unique(quats[5]), [want], err_msg=f"start={start}")
+        assert n_positions == 1
+
+
 @pytest.mark.slowtest
 @testing.requires_testing_data
 @pytest.mark.parametrize(
@@ -2056,6 +2123,36 @@ def test_find_bads_maxwell_flat():
     assert noisy == want_noisy
 
 
+@pytest.mark.slowtest
+@testing.requires_testing_data
+def test_find_bads_maxwell_head_pos():
+    """Test that each interval uses its own head positions."""
+    raw = read_raw_fif(raw_fname, allow_maxshield="yes")
+    raw.pick("meg", exclude=()).crop(0, 10)  # two 5 s intervals
+    raw.load_data()
+    head_pos = _linear_head_pos(raw, 11)
+    kwargs = dict(
+        origin=(0.0, 0.0, 0.04),
+        regularize=None,
+        bad_condition="ignore",
+        min_count=1,
+        return_scores=True,
+        h_freq=None,  # keep the two calls below operating on identical data
+    )
+    flats, scores = find_bad_channels_maxwell(raw, head_pos=head_pos, **kwargs)[1:]
+    assert scores["bins"].shape == (2, 2)
+    # flats found in interval 1 would be excluded from interval 2 but not from the
+    # cropped run below, making the two good_masks (and hence the scores) differ
+    assert flats == []
+    # Processing the second interval on its own must give the same scores as
+    # processing it as part of the whole recording.
+    raw_crop = raw.copy().crop(5.0)
+    pos_crop = head_pos[5:]
+    scores_crop = find_bad_channels_maxwell(raw_crop, head_pos=pos_crop, **kwargs)[2]
+    assert scores_crop["bins"].shape == (1, 2)
+    assert_allclose(scores_crop["scores_noisy"][:, 0], scores["scores_noisy"][:, 1])
+
+
 @pytest.mark.parametrize(
     "regularize, n, int_order",
     [