Repository: mne-tools/mne-python

Issue:
add equalize_event_counts method to EpochsTFR
time-domain Epochs object already has this method so should be fairly straightforward.

requested here:
https://mne.discourse.group/t/equalize-epoch-counts-for-epochstfr/6879/2


Discussion:
I would like to work on this issue

Resolving patch:
diff --git a/mne/epochs.py b/mne/epochs.py
--- a/mne/epochs.py
+++ b/mne/epochs.py
@@ -75,6 +75,7 @@
 from .html_templates import _get_html_template
 from .parallel import parallel_func
 from .time_frequency.spectrum import EpochsSpectrum, SpectrumMixin, _validate_method
+from .time_frequency.tfr import EpochsTFR
 from .utils import (
     ExtendedTimeMixin,
     GetEpochsMixin,
@@ -2479,8 +2480,8 @@ def equalize_event_counts(self, event_ids=None, method="mintime"):
         for eq in event_ids:
             eq_inds.append(self._keys_to_idx(eq))
 
-        event_times = [self.events[e, 0] for e in eq_inds]
-        indices = _get_drop_indices(event_times, method)
+        sample_nums = [self.events[e, 0] for e in eq_inds]
+        indices = _get_drop_indices(sample_nums, method)
         # need to re-index indices
         indices = np.concatenate([e[idx] for e, idx in zip(eq_inds, indices)])
         self.drop(indices, reason="EQUALIZED_COUNT")
@@ -3616,7 +3617,7 @@ def combine_event_ids(epochs, old_event_ids, new_event_id, copy=True):
 
 
 def equalize_epoch_counts(epochs_list, method="mintime"):
-    """Equalize the number of trials in multiple Epoch instances.
+    """Equalize the number of trials in multiple Epochs or EpochsTFR instances.
 
     Parameters
     ----------
@@ -3643,33 +3644,32 @@ def equalize_epoch_counts(epochs_list, method="mintime"):
     --------
     >>> equalize_epoch_counts([epochs1, epochs2])  # doctest: +SKIP
     """
-    if not all(isinstance(e, BaseEpochs) for e in epochs_list):
+    if not all(isinstance(epoch, (BaseEpochs, EpochsTFR)) for epoch in epochs_list):
         raise ValueError("All inputs must be Epochs instances")
 
     # make sure bad epochs are dropped
-    for e in epochs_list:
-        if not e._bad_dropped:
-            e.drop_bad()
-    event_times = [e.events[:, 0] for e in epochs_list]
-    indices = _get_drop_indices(event_times, method)
-    for e, inds in zip(epochs_list, indices):
-        e.drop(inds, reason="EQUALIZED_COUNT")
+    for epoch in epochs_list:
+        if not epoch._bad_dropped:
+            epoch.drop_bad()
+    sample_nums = [epoch.events[:, 0] for epoch in epochs_list]
+    indices = _get_drop_indices(sample_nums, method)
+    for epoch, inds in zip(epochs_list, indices):
+        epoch.drop(inds, reason="EQUALIZED_COUNT")
 
 
-def _get_drop_indices(event_times, method):
+def _get_drop_indices(sample_nums, method):
     """Get indices to drop from multiple event timing lists."""
-    small_idx = np.argmin([e.shape[0] for e in event_times])
-    small_e_times = event_times[small_idx]
+    small_idx = np.argmin([e.shape[0] for e in sample_nums])
+    small_epoch_indices = sample_nums[small_idx]
     _check_option("method", method, ["mintime", "truncate"])
     indices = list()
-    for e in event_times:
+    for event in sample_nums:
         if method == "mintime":
-            mask = _minimize_time_diff(small_e_times, e)
+            mask = _minimize_time_diff(small_epoch_indices, event)
         else:
-            mask = np.ones(e.shape[0], dtype=bool)
-            mask[small_e_times.shape[0] :] = False
+            mask = np.ones(event.shape[0], dtype=bool)
+            mask[small_epoch_indices.shape[0] :] = False
         indices.append(np.where(np.logical_not(mask))[0])
-
     return indices
 
 
diff --git a/mne/time_frequency/tfr.py b/mne/time_frequency/tfr.py
--- a/mne/time_frequency/tfr.py
+++ b/mne/time_frequency/tfr.py
@@ -2764,6 +2764,8 @@ def __init__(
         self.method = method
         self.preload = True
         self.metadata = metadata
+        # we need this to allow equalize_epoch_counts to work with EpochsTFRs
+        self._bad_dropped = True
 
     @property
     def _detrend_picks(self):
@@ -2875,6 +2877,41 @@ def average(self, method="mean", dim="epochs", copy=False):
             self.freqs = freqs
             return self
 
+    @verbose
+    def drop(self, indices, reason="USER", verbose=None):
+        """Drop epochs based on indices or boolean mask.
+
+        .. note:: The indices refer to the current set of undropped epochs
+                  rather than the complete set of dropped and undropped epochs.
+                  They are therefore not necessarily consistent with any
+                  external indices (e.g., behavioral logs). To drop epochs
+                  based on external criteria, do not use the ``preload=True``
+                  flag when constructing an Epochs object, and call this
+                  method before calling the :meth:`mne.Epochs.drop_bad` or
+                  :meth:`mne.Epochs.load_data` methods.
+
+        Parameters
+        ----------
+        indices : array of int or bool
+            Set epochs to remove by specifying indices to remove or a boolean
+            mask to apply (where True values get removed). Events are
+            correspondingly modified.
+        reason : str
+            Reason for dropping the epochs ('ECG', 'timeout', 'blink' etc).
+            Default: 'USER'.
+        %(verbose)s
+
+        Returns
+        -------
+        epochs : instance of Epochs or EpochsTFR
+            The epochs with indices dropped. Operates in-place.
+        """
+        from ..epochs import BaseEpochs
+
+        BaseEpochs.drop(self, indices=indices, reason=reason, verbose=verbose)
+
+        return self
+
 
 def combine_tfr(all_tfr, weights="nave"):
     """Merge AverageTFR data by weighted addition.


Test patch:
diff --git a/mne/tests/test_epochs.py b/mne/tests/test_epochs.py
--- a/mne/tests/test_epochs.py
+++ b/mne/tests/test_epochs.py
@@ -2803,25 +2803,58 @@ def test_subtract_evoked():
 
 
 def test_epoch_eq():
-    """Test epoch count equalization and condition combining."""
+    """Test for equalize_epoch_counts and equalize_event_counts functions."""
+    # load data
     raw, events, picks = _get_data()
-    # equalizing epochs objects
+    # test equalize epoch counts
+    # create epochs with unequal counts
     events_1 = events[events[:, 2] == event_id]
     epochs_1 = Epochs(raw, events_1, event_id, tmin, tmax, picks=picks)
     events_2 = events[events[:, 2] == event_id_2]
     epochs_2 = Epochs(raw, events_2, event_id_2, tmin, tmax, picks=picks)
+    # events 2 has one more event than events 1
     epochs_1.drop_bad()  # make sure drops are logged
+    epochs_2.drop_bad()  # make sure drops are logged
+    # make sure there is a difference in the number of events
+    assert len(epochs_1) != len(epochs_2)
+    # make sure bad epochs are dropped before equalizing epoch counts
     assert_equal(
         len([log for log in epochs_1.drop_log if not log]), len(epochs_1.events)
     )
-    assert epochs_1.drop_log == ((),) * len(epochs_1.events)
-    assert_equal(len([lg for lg in epochs_1.drop_log if not lg]), len(epochs_1.events))
-    assert epochs_1.events.shape[0] != epochs_2.events.shape[0]
+    assert epochs_2.drop_log == ((),) * len(epochs_2.events)
+    # test mintime method
+    events_1[-1, 0] += 60  # hack: ensure mintime drops something other than last trial
+    # now run equalize_epoch_counts with mintime method
     equalize_epoch_counts([epochs_1, epochs_2], method="mintime")
+    # mintime method should give us the smallest difference between timings of epochs
+    alleged_mintime = np.sum(np.abs(epochs_1.events[:, 0] - epochs_2.events[:, 0]))
+    # test that "mintime" works as expected, by systematically dropping each event from
+    # events_2 and ensuring the latencies are actually smallest in the
+    # equalize_epoch_counts case. NB: len(events_2) > len(events_1)
+    for idx in range(events_2.shape[0]):
+        # delete epoch from events_2
+        test_events = np.delete(events_2.copy(), idx, axis=0)
+        assert test_events.shape == epochs_1.events.shape == epochs_2.events.shape
+        # difference (in samples) between epochs_1 event times and the event times we
+        # get from our deletion of row `idx` from events_2
+        latencies = epochs_1.events[:, 0] - test_events[:, 0]
+        got_mintime = np.sum(np.abs(latencies))
+        assert got_mintime >= alleged_mintime
+    # make sure the number of events is equal
     assert_equal(epochs_1.events.shape[0], epochs_2.events.shape[0])
+    # create new epochs with the same event ids as epochs_1 and epochs_2
     epochs_3 = Epochs(raw, events, event_id, tmin, tmax, picks=picks)
     epochs_4 = Epochs(raw, events, event_id_2, tmin, tmax, picks=picks)
+    epochs_3.drop_bad()  # make sure drops are logged
+    epochs_4.drop_bad()  # make sure drops are logged
+    # make sure there is a difference in the number of events
+    assert len(epochs_3) != len(epochs_4)
+    # test truncate method
     equalize_epoch_counts([epochs_3, epochs_4], method="truncate")
+    if len(epochs_3.events) > len(epochs_4.events):
+        assert_equal(epochs_3.events[-2, 0], epochs_3.events.shape[-1, 0])
+    elif len(epochs_3.events) < len(epochs_4.events):
+        assert_equal(epochs_4.events[-2, 0], epochs_4.events[-1, 0])
     assert_equal(epochs_1.events.shape[0], epochs_3.events.shape[0])
     assert_equal(epochs_3.events.shape[0], epochs_4.events.shape[0])
 
diff --git a/mne/time_frequency/tests/test_tfr.py b/mne/time_frequency/tests/test_tfr.py
--- a/mne/time_frequency/tests/test_tfr.py
+++ b/mne/time_frequency/tests/test_tfr.py
@@ -21,6 +21,7 @@
     pick_types,
     read_events,
 )
+from mne.epochs import equalize_epoch_counts
 from mne.io import read_raw_fif
 from mne.tests.test_epochs import assert_metadata_equal
 from mne.time_frequency import tfr_array_morlet, tfr_array_multitaper
@@ -47,6 +48,28 @@
 raw_ctf_fname = data_path / "test_ctf_raw.fif"
 
 
+def _create_test_epochstfr():
+    n_epos = 3
+    ch_names = ["EEG 001", "EEG 002", "EEG 003", "EEG 004"]
+    n_picks = len(ch_names)
+    ch_types = ["eeg"] * n_picks
+    n_freqs = 5
+    n_times = 6
+    data = np.random.rand(n_epos, n_picks, n_freqs, n_times)
+    times = np.arange(6)
+    srate = 1000.0
+    freqs = np.arange(5)
+    events = np.zeros((n_epos, 3), dtype=int)
+    events[:, 0] = np.arange(n_epos)
+    events[:, 2] = np.arange(5, 5 + n_epos)
+    event_id = {k: v for v, k in zip(events[:, 2], ["ha", "he", "hu"])}
+    info = mne.create_info(ch_names, srate, ch_types)
+    tfr = mne.time_frequency.EpochsTFR(
+        info, data, times, freqs, events=events, event_id=event_id
+    )
+    return tfr
+
+
 def test_tfr_ctf():
     """Test that TFRs can be calculated on CTF data."""
     raw = read_raw_fif(raw_ctf_fname).crop(0, 1)
@@ -723,6 +746,14 @@ def test_init_EpochsTFR():
         del tfr
 
 
+def test_equalize_epochs_tfr_counts():
+    """Test equalize_epoch_counts for EpochsTFR."""
+    tfr = _create_test_epochstfr()
+    tfr2 = tfr.copy()
+    tfr2 = tfr2[:-1]
+    equalize_epoch_counts([tfr, tfr2])
+
+
 def test_dB_computation():
     """Test dB computation in plot methods (gh 11091)."""
     ampl = 2.0