Fix realign_raw and test_realign (mne-tools#11950)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel McCloy <dan@mccloy.info>
Co-authored-by: Eric Larson <larson.eric.d@gmail.com>

Author: 4 people

doc/changes/devel.rst

@@ -36,6 +36,7 @@ Enhancements
 
 Bugs
 ~~~~
+- Fix bugs with :func:`mne.preprocessing.realign_raw` where the start of ``other`` was incorrectly cropped; and onsets and durations in ``other.annotations`` were left unsynced with the resampled data (:gh:`11950` by :newcontrib:`Qian Chu`)
 - Fix bug where ``encoding`` argument was ignored when reading annotations from an EDF file (:gh:`11958` by :newcontrib:`Andrew Gilbert`)
 - Fix bugs with saving splits for :class:`~mne.Epochs` (:gh:`11876` by `Dmitrii Altukhov`_)
 - Fix bug with multi-plot 3D rendering where only one plot was updated (:gh:`11896` by `Eric Larson`_)

doc/changes/names.inc

@@ -418,6 +418,8 @@
 
 .. _Proloy Das: https://github.com/proloyd
 
+.. _Qian Chu: https://github.com/qian-chu
+
 .. _Qianliang Li: https://www.dtu.dk/english/service/phonebook/person?id=126774
 
 .. _Quentin Barthélemy: https://github.com/qbarthelemy

mne/preprocessing/realign.py

@@ -1,5 +1,6 @@
 # Authors: Eric Larson <larson.eric.d@gmail.com>
-
+#          Qian Chu <qianchu99@gmail.com>
+#
 # License: BSD-3-Clause
 
 import numpy as np
@@ -42,7 +43,8 @@ def realign_raw(raw, other, t_raw, t_other, verbose=None):
     2. Crop the start of ``raw`` or ``other``, depending on which started
        recording first.
     3. Resample ``other`` to match ``raw`` based on the clock drift.
-    4. Crop the end of ``raw`` or ``other``, depending on which stopped
+    4. Realign the onsets and durations in ``other.annotations``.
+    5. Crop the end of ``raw`` or ``other``, depending on which stopped
        recording first (and the clock drift rate).
 
     This function is primarily designed to work on recordings made at the same
@@ -85,25 +87,41 @@ def realign_raw(raw, other, t_raw, t_other, verbose=None):
         f"{raw.times[-1] * dr_ms_s:0.1f} ms)"
     )
 
-    # 2. Crop start of recordings to match using the zero-order term
-    msg = f"Cropping {zero_ord:0.3f} s from the start of "
+    # 2. Crop start of recordings to match
     if zero_ord > 0:  # need to crop start of raw to match other
-        logger.info(msg + "raw")
+        logger.info(f"Cropping {zero_ord:0.3f} s from the start of raw")
         raw.crop(zero_ord, None)
         t_raw -= zero_ord
     else:  # need to crop start of other to match raw
-        logger.info(msg + "other")
-        other.crop(-zero_ord, None)
-        t_other += zero_ord
+        t_crop = zero_ord / first_ord
+        logger.info(f"Cropping {t_crop:0.3f} s from the start of other")
+        other.crop(-t_crop, None)
+        t_other += t_crop
 
     # 3. Resample data using the first-order term
+    nan_ch_names = [
+        ch for ch in other.info["ch_names"] if np.isnan(other.get_data(picks=ch)).any()
+    ]
+    if len(nan_ch_names) > 0:  # Issue warning if any channel in other has nan values
+        warn(
+            f"Channel(s) {', '.join(nan_ch_names)} in `other` contain NaN values. "
+            "Resampling these channels will result in the whole channel being NaN. "
+            "(If realigning eye-tracking data, consider using interpolate_blinks and "
+            "passing interpolate_gaze=True)"
+        )
     logger.info("Resampling other")
     sfreq_new = raw.info["sfreq"] * first_ord
     other.load_data().resample(sfreq_new, verbose=True)
     with other.info._unlock():
         other.info["sfreq"] = raw.info["sfreq"]
 
-    # 4. Crop the end of one of the recordings if necessary
+    # 4. Realign the onsets and durations in other.annotations
+    # Must happen before end cropping to avoid losing annotations
+    logger.info("Correcting annotations in other")
+    other.annotations.onset *= first_ord
+    other.annotations.duration *= first_ord
+
+    # 5. Crop the end of one of the recordings if necessary
     delta = raw.times[-1] - other.times[-1]
     msg = f"Cropping {abs(delta):0.3f} s from the end of "
     if delta > 0:

mne/preprocessing/tests/test_realign.py

@@ -1,4 +1,5 @@
-# Author: Mark Wronkiewicz <wronk@uw.edu>
+# Authors: Mark Wronkiewicz <wronk@uw.edu>
+#          Qian Chu <qianchu99@gmail.com>
 #
 # License: BSD-3-Clause
 
@@ -7,12 +8,12 @@
 from scipy.interpolate import interp1d
 import pytest
 
-from mne import create_info, find_events, Epochs
+from mne import create_info, find_events, Epochs, Annotations
 from mne.io import RawArray
 from mne.preprocessing import realign_raw
 
 
-@pytest.mark.parametrize("ratio_other", (1.0, 0.999, 1.001))  # drifts
+@pytest.mark.parametrize("ratio_other", (0.9, 0.999, 1, 1.001, 1.1))  # drifts
 @pytest.mark.parametrize("start_raw, start_other", [(0, 0), (0, 3), (3, 0)])
 @pytest.mark.parametrize("stop_raw, stop_other", [(0, 0), (0, 3), (3, 0)])
 def test_realign(ratio_other, start_raw, start_other, stop_raw, stop_other):
@@ -22,6 +23,8 @@ def test_realign(ratio_other, start_raw, start_other, stop_raw, stop_other):
     duration = 50
     stop_raw = duration - stop_raw
     stop_other = duration - stop_other
+    signal_len = 0.2
+    box_len = 0.5
     signal = np.zeros(int(round((duration + 1) * sfreq)))
     orig_events = np.round(
         np.arange(max(start_raw, start_other) + 2, min(stop_raw, stop_other) - 2)
@@ -30,8 +33,10 @@ def test_realign(ratio_other, start_raw, start_other, stop_raw, stop_other):
     signal[orig_events] = 1.0
     n_events = len(orig_events)
     times = np.arange(len(signal)) / sfreq
-    stim = np.convolve(signal, np.ones(int(round(0.02 * sfreq))))[: len(times)]
-    signal = np.convolve(signal, np.hanning(int(round(0.2 * sfreq))))[: len(times)]
+    stim = np.convolve(signal, np.ones(int(round(box_len * sfreq))))[: len(times)]
+    signal = np.convolve(signal, np.hanning(int(round(signal_len * sfreq))))[
+        : len(times)
+    ]
 
     # construct our sampled versions of these signals (linear interp is fine)
     sfreq_raw = sfreq
@@ -45,45 +50,84 @@ def test_realign(ratio_other, start_raw, start_other, stop_raw, stop_other):
     data_raw = np.array(
         [
             interp1d(times, d, kind)(raw_times)
-            for d, kind in ((signal, "linear"), (stim, "nearest"))
+            for d, kind in (
+                (stim, "nearest"),
+                (signal, "linear"),
+            )
         ]
     )
     data_other = np.array(
         [
             interp1d(times, d, kind)(other_times)
-            for d, kind in ((signal, "linear"), (stim, "nearest"))
+            for d, kind in (
+                (stim, "nearest"),
+                (signal, "linear"),
+            )
         ]
     )
-    info_raw = create_info(["raw_data", "raw_stim"], sfreq, ["eeg", "stim"])
-    info_other = create_info(["other_data", "other_stim"], sfreq, ["eeg", "stim"])
-    raw = RawArray(data_raw, info_raw, first_samp=111)
+    info_raw = create_info(["raw_stim", "raw_signal"], sfreq, ["stim", "eeg"])
+    info_other = create_info(["other_stim", "other_signal"], sfreq, ["stim", "eeg"])
+    raw = RawArray(data_raw, info_raw, first_samp=111)  # first_samp shouldn't matter
     other = RawArray(data_other, info_other, first_samp=222)
+    raw.set_meas_date((0, 0))  # meas_date shouldn't matter
+    other.set_meas_date((100, 0))
 
-    # naive processing
-    evoked_raw, events_raw, _, events_other = _assert_similarity(raw, other, n_events)
-    if start_raw == start_other:  # can just naively crop
-        a, b = data_raw[0], data_other[0]
-        n = min(len(a), len(b))
-        corr = np.corrcoef(a[:n], b[:n])[0, 1]
-        min_, max_ = (0.99999, 1.0) if sfreq_raw == sfreq_other else (0.8, 0.9)
-        assert min_ <= corr <= max_
+    # find events and do basic checks
+    evoked_raw, events_raw, _, events_other = _assert_similarity(
+        raw, other, n_events, ratio_other
+    )
+
+    # construct annotations
+    onsets_raw = (events_raw[:, 0] - raw.first_samp) / raw.info["sfreq"]
+    dur_raw = [box_len] * len(onsets_raw)
+    desc_raw = ["raw_box"] * len(onsets_raw)
+    annot_raw = Annotations(onsets_raw, dur_raw, desc_raw)
+    raw.set_annotations(annot_raw)
+
+    onsets_other = (events_other[:, 0] - other.first_samp) / other.info["sfreq"]
+    dur_other = [box_len * ratio_other] * len(onsets_other)
+    desc_other = ["other_box"] * len(onsets_other)
+    annot_other = Annotations(onsets_other, dur_other, desc_other)
+    other.set_annotations(annot_other)
+
+    # onsets/offsets correspond to 0/1 transition in boxcar signals
+    _assert_boxcar_annot_similarity(raw, other)
 
     # realign
-    t_raw = (events_raw[:, 0] - raw.first_samp) / other.info["sfreq"]
+    t_raw = (events_raw[:, 0] - raw.first_samp) / raw.info["sfreq"]
     t_other = (events_other[:, 0] - other.first_samp) / other.info["sfreq"]
     assert duration - 10 <= len(events_raw) < duration
     raw_orig, other_orig = raw.copy(), other.copy()
     realign_raw(raw, other, t_raw, t_other)
 
-    # old events should still work for raw and produce the same result
-    evoked_raw_2, _, _, _ = _assert_similarity(
-        raw, other, n_events, events_raw=events_raw
+    # old events should still work for raw and produce the same evoked data
+    evoked_raw_2, events_raw, _, events_other = _assert_similarity(
+        raw, other, n_events, ratio_other, events_raw=events_raw
     )
     assert_allclose(evoked_raw.data, evoked_raw_2.data)
     assert_allclose(raw.times, other.times)
+
     # raw data now aligned
-    corr = np.corrcoef(raw.get_data([0])[0], other.get_data([0])[0])[0, 1]
-    assert 0.99 < corr <= 1.0
+    corr = np.corrcoef(raw.get_data("data"), other.get_data("data"))
+    assert 0.99 < corr[0, 1] <= 1.0
+
+    # onsets derived from stim and annotations are the same
+    atol = 2 / sfreq
+    assert_allclose(
+        raw.annotations.onset, events_raw[:, 0] / raw.info["sfreq"], atol=atol
+    )
+    assert_allclose(
+        other.annotations.onset, events_other[:, 0] / other.info["sfreq"], atol=atol
+    )
+
+    # onsets/offsets still correspond to 0/1 transition in boxcar signals
+    _assert_boxcar_annot_similarity(raw, other)
+
+    # onsets and durations now aligned
+    onsets_raw, dur_raw, onsets_other, dur_other = _annot_to_onset_dur(raw, other)
+    assert len(onsets_raw) == len(onsets_other) == len(events_raw)
+    assert_allclose(onsets_raw, onsets_other, atol=atol)
+    assert_allclose(dur_raw, dur_other, atol=atol)
 
     # Degenerate conditions -- only test in one run
     test_degenerate = (
@@ -103,17 +147,44 @@ def test_realign(ratio_other, start_raw, start_other, stop_raw, stop_other):
         realign_raw(raw_orig, other_orig, raw_times + rand_times * 1000, other_times)
 
 
-def _assert_similarity(raw, other, n_events, events_raw=None):
+def _assert_similarity(raw, other, n_events, ratio_other, events_raw=None):
     if events_raw is None:
-        events_raw = find_events(raw)
-    events_other = find_events(other)
-    assert len(events_raw) == n_events
-    assert len(events_other) == n_events
+        events_raw = find_events(raw, output="onset")
+    events_other = find_events(other, output="onset")
+    assert len(events_raw) == len(events_other) == n_events
     kwargs = dict(baseline=None, tmin=0, tmax=0.2)
     evoked_raw = Epochs(raw, events_raw, **kwargs).average()
     evoked_other = Epochs(other, events_other, **kwargs).average()
     assert evoked_raw.nave == evoked_other.nave == len(events_raw)
     assert len(evoked_raw.data) == len(evoked_other.data) == 1  # just EEG
-    corr = np.corrcoef(evoked_raw.data[0], evoked_other.data[0])[0, 1]
-    assert 0.9 <= corr <= 1.0
+    if 0.99 <= ratio_other <= 1.01:  #  when drift is not too large
+        corr = np.corrcoef(evoked_raw.data[0], evoked_other.data[0])[0, 1]
+        assert 0.9 <= corr <= 1.0
     return evoked_raw, events_raw, evoked_other, events_other
+
+
+def _assert_boxcar_annot_similarity(raw, other):
+    onsets_raw, dur_raw, onsets_other, dur_other = _annot_to_onset_dur(raw, other)
+
+    n_events = len(onsets_raw)
+    onsets_samp_raw = raw.time_as_index(onsets_raw)
+    offsets_samp_raw = raw.time_as_index(onsets_raw + dur_raw)
+    assert_allclose(raw.get_data("stim")[0, onsets_samp_raw - 2], [0] * n_events)
+    assert_allclose(raw.get_data("stim")[0, onsets_samp_raw + 2], [1] * n_events)
+    assert_allclose(raw.get_data("stim")[0, offsets_samp_raw - 2], [1] * n_events)
+    assert_allclose(raw.get_data("stim")[0, offsets_samp_raw + 2], [0] * n_events)
+    onsets_samp_other = other.time_as_index(onsets_other)
+    offsets_samp_other = other.time_as_index(onsets_other + dur_other)
+    assert_allclose(other.get_data("stim")[0, onsets_samp_other - 2], [0] * n_events)
+    assert_allclose(other.get_data("stim")[0, onsets_samp_other + 2], [1] * n_events)
+    assert_allclose(other.get_data("stim")[0, offsets_samp_other - 2], [1] * n_events)
+    assert_allclose(other.get_data("stim")[0, offsets_samp_other + 2], [0] * n_events)
+
+
+def _annot_to_onset_dur(raw, other):
+    onsets_raw = raw.annotations.onset - raw.first_time
+    dur_raw = raw.annotations.duration
+
+    onsets_other = other.annotations.onset - other.first_time
+    dur_other = other.annotations.duration
+    return onsets_raw, dur_raw, onsets_other, dur_other