Fix find_bad_channels_maxwell cHPI correction

doc/changes/dev/14142.bugfix.rst

@@ -0,0 +1 @@
+Fix bug in :func:`mne.preprocessing.find_bad_channels_maxwell` where, when ``head_pos`` was provided, every interval was processed using the head positions from the beginning of the recording rather than from its own time window, by `Christian Brodbeck`_.

mne/_ola.py

@@ -23,6 +23,11 @@ class _Interp2:
         arrays that must be interpolated.
     interp : str
         Can be 'zero', 'linear', 'hann', or 'cos2' (same as hann).
+    start : int
+        The control-point index at which feeding begins, for processing a segment
+        that does not start at the first control point. Feeding ``n`` points then
+        yields exactly what a continuous pass would yield for ``start`` to
+        ``start + n``, interpolation phase included.
 
     Notes
     -----
@@ -42,7 +47,9 @@ class _Interp2:
 
     """
 
-    def __init__(self, control_points, values, interp="hann", *, name="Interp2"):
+    def __init__(
+        self, control_points, values, interp="hann", *, name="Interp2", start=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 +83,11 @@ def val(pt):
             values = val
         self.values = values
         self.n_last = None
-        self._position = 0  # start at zero
-        self._left_idx = 0
+        self._position = start
+        # The last control point at or before start is the one in effect there, so
+        # feeding resumes mid-interval with the correct interpolation phase.
+        left_idx = np.searchsorted(self.control_points, start, "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 +104,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
 

mne/preprocessing/maxwell.py

@@ -741,6 +741,7 @@ def _run_maxwell_filter(
     st_fixed,
     st_overlap,
     mc,
+    mc_start=0,
 ):
     # 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,13 +771,13 @@ 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
     # find_bad_channels_maxwell modifies good_mask
-    mc.initialize(_get_this_decomp_trans, info["dev_head_t"], S_recon)
+    mc.initialize(_get_this_decomp_trans, info["dev_head_t"], S_recon, mc_start)
     update_kwargs.update(reg_moments=mc.reg_moments_0)
 
     # Process each valid block of data separately
@@ -865,6 +866,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
@@ -890,24 +895,32 @@ def get_decomp_by_offset(self, offset):
         op_resid -= np.dot(S_decomp[:, n_use_in:], pS_decomp[n_use_in:])
         return op_sss, op_in, op_resid
 
-    def initialize(self, get_decomp, dev_head_t, S_recon):
-        """Secondary initialization."""
+    def initialize(self, get_decomp, dev_head_t, S_recon, start=0):
+        """Secondary initialization.
+
+        :attr:`self.pos` is indexed relative to the start of the recording;
+        ``start`` adds an index offset when processing data in chunks.
+        """
+        self.start = start
         self.smooth = _Interp2(
             self.pos[1],
             self.get_decomp_by_offset,
             interp=self.interp,
             name="MC",
+            start=start,
         )
         _, _, 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.offset = start
         self.get_decomp = get_decomp
         # For the average passes
         self.last_avg_quat = np.nan * np.ones(6)
 
     def get_avg_op(self, *, start, stop):
         """Apply an average transformation over the next interval."""
+        # Start and stop are relative to the start set at .initialize()
+        start, stop = start + self.start, stop + self.start
         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
@@ -2941,7 +2954,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, mc_start=start, **params)
 
             if n_iter == 1 and len(chunk_flats):
                 logger.info(

mne/preprocessing/tests/test_maxwell.py

@@ -54,6 +54,7 @@
     _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,
@@ -2056,6 +2057,40 @@ 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()
+    # One head position per second, translating steadily so that using the positions
+    # from the wrong time window gives a different result.
+    trans = raw.info["dev_head_t"]["trans"]
+    head_pos = np.zeros((11, 10))
+    head_pos[:, 0] = raw._first_time + np.arange(11.0)
+    head_pos[:, 1:4] = rot_to_quat(trans[:3, :3])
+    head_pos[:, 4:7] = trans[:3, 3]
+    head_pos[:, 6] += np.arange(11) * 5e-3  # 5 mm/s
+    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
+    )
+    scores = find_bad_channels_maxwell(raw, head_pos=head_pos, **kwargs)[2]
+    assert scores["bins"].shape == (2, 2)
+    # 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",
     [

mne/tests/test_ola.py

@@ -84,6 +84,23 @@ def test_interp_2pt():
         assert_allclose(out, expected)
 
 
+@pytest.mark.parametrize("interp", ("zero", "linear", "hann"))
+def test_interp_2pt_start(interp):
+    """Test that starting mid-stream matches a continuous pass."""
+    # control points deliberately uneven, so that most starts land mid-interval
+    control_points = [0, 100, 240, 390, 520, 680, 800]
+    values = np.array(control_points, float)
+    want = _Interp2(control_points, values, interp).feed(900)[0]
+    for start in range(0, 900, 7):
+        interper = _Interp2(control_points, values, interp, start=start)
+        assert_allclose(interper.feed(900 - start)[0], want[start:], atol=1e-12)
+    # state must carry across feeds of differing size, as when a segment is read
+    # in buffer-sized blocks
+    interper = _Interp2(control_points, values, interp, start=300)
+    out = np.concatenate([interper.feed(n)[0] for n in (70, 130, 100, 300)])
+    assert_allclose(out, want[300:900], atol=1e-12)
+
+
 @pytest.mark.parametrize("ndim", (1, 2, 3))
 def test_cola(ndim):
     """Test COLA processing."""