diff --git a/doc/changes/dev/14142.bugfix.rst b/doc/changes/dev/14142.bugfix.rst new file mode 100644 index 00000000000..f53d1823736 --- /dev/null +++ b/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`_. diff --git a/mne/_ola.py b/mne/_ola.py index e182928f4a5..242609364a5 100644 --- a/mne/_ola.py +++ b/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 = [ diff --git a/mne/preprocessing/maxwell.py b/mne/preprocessing/maxwell.py index 2f013961293..b161e07d674 100644 --- a/mne/preprocessing/maxwell.py +++ b/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( diff --git a/mne/preprocessing/tests/test_maxwell.py b/mne/preprocessing/tests/test_maxwell.py index f74195292c4..989ab780e30 100644 --- a/mne/preprocessing/tests/test_maxwell.py +++ b/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", [ diff --git a/mne/tests/test_ola.py b/mne/tests/test_ola.py index 0528b68f363..cc6ae92608b 100644 --- a/mne/tests/test_ola.py +++ b/mne/tests/test_ola.py @@ -4,7 +4,7 @@ import numpy as np import pytest -from numpy.testing import assert_allclose +from numpy.testing import assert_allclose, assert_array_equal from mne._ola import _COLA, _Interp2, _Storer @@ -84,6 +84,25 @@ def test_interp_2pt(): assert_allclose(out, expected) +@pytest.mark.parametrize("interp", ("zero", "linear", "hann")) +def test_interp_2pt_offset(interp): + """Test that starting mid-stream matches a continuous pass.""" + # control points deliberately uneven, so that most offsets 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 offset in range(0, 900, 7): + interper = _Interp2(control_points, values, interp, offset=offset) + assert_allclose(interper.feed(900 - offset)[0], want[offset:], 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, offset=300) + out = np.concatenate([interper.feed(n)[0] for n in (70, 130, 100, 300)]) + assert_allclose(out, want[300:900], atol=1e-12) + with pytest.raises(ValueError, match="offset must be non-negative"): + _Interp2(control_points, values, interp, offset=-1) + + @pytest.mark.parametrize("ndim", (1, 2, 3)) def test_cola(ndim): """Test COLA processing.""" @@ -126,3 +145,29 @@ def processor(x, *, start, stop): cola.feed(signal[..., n_input : n_input + next_len]) n_input += next_len assert_allclose(out, signal / 2.0, atol=1e-7) + + +def test_cola_offset(): + """Test that COLA shifts the limits it hands to the processor.""" + n_total, n_samples, n_overlap, offset = 1000, 100, 50, 4321 + limits = list() + + def processor(x, *, start, stop): + limits.append((start, stop)) + return (x,) + + signal = np.zeros(n_total) + runs = list() + for use_offset in (0, offset): + limits.clear() + _COLA( + processor, + np.zeros(n_total), + n_total, + n_samples, + n_overlap, + 1000.0, + offset=use_offset, + ).feed(signal) + runs.append(np.array(limits)) + assert_array_equal(runs[1] - offset, runs[0])