ENH: Support concatenate_epochs() for EpochsSpectrum

doc/changes/dev/13748.newfeature.rst

@@ -0,0 +1 @@
+Add support for :func:`mne.concatenate_epochs` with :class:`~mne.time_frequency.EpochsSpectrum` instances, by ``aman-coder03``. (:gh:`13747`)

mne/epochs.py

@@ -4693,6 +4693,65 @@ def _concatenate_epochs(
     )
 
 
+def _concatenate_epochs_spectrum(epochs_list, add_offset=True):
+    """Concatenate a list of EpochsSpectrum instances."""
+    for ii, ep in enumerate(epochs_list):
+        if type(ep).__name__ != "EpochsSpectrum":
+            raise TypeError(
+                f"epochs_list[{ii}] must be an instance of EpochsSpectrum, "
+                f"got {type(ep)}"
+            )
+    ref = epochs_list[0]
+    for ii, ep in enumerate(epochs_list[1:], 1):
+        if not np.array_equal(ep.freqs, ref.freqs):
+            raise ValueError(f"epochs_list[{ii}] freqs do not match epochs_list[0]")
+        _ensure_infos_match(ep.info, ref.info, f"epochs_list[{ii}]")
+
+    data = np.concatenate([ep.data for ep in epochs_list], axis=0)
+
+    shift = np.int64(10 * ref.info["sfreq"])
+    events_offset = int(np.max(epochs_list[0].events[:, 0])) + shift
+    all_events = [epochs_list[0].events.copy()]
+    for ep in epochs_list[1:]:
+        evs = ep.events.copy()
+        if add_offset:
+            evs[:, 0] += events_offset
+            events_offset += int(np.max(ep.events[:, 0])) + shift
+        all_events.append(evs)
+    events = np.concatenate(all_events, axis=0)
+
+    event_id = deepcopy(ref.event_id)
+    for ep in epochs_list[1:]:
+        event_id.update(ep.event_id)
+
+    selection = np.concatenate([ep.selection for ep in epochs_list])
+    drop_log = sum([ep.drop_log for ep in epochs_list], ())
+
+    metadatas = [ep.metadata for ep in epochs_list]
+    n_have = sum(m is not None for m in metadatas)
+    if n_have == 0:
+        metadata = None
+    elif n_have != len(metadatas):
+        raise ValueError(
+            f"{n_have} of {len(metadatas)} EpochsSpectrum instances have metadata, "
+            "all or none must have metadata"
+        )
+    else:
+        pd = _check_pandas_installed(strict=False)
+        metadata = pd.concat(metadatas) if pd is not False else sum(metadatas, list())
+
+    state = ref.__getstate__()
+    state["data"] = data
+    state["events"] = events
+    state["event_id"] = event_id
+    state["selection"] = selection
+    state["drop_log"] = drop_log
+    state["metadata"] = metadata
+    out = type(epochs_list[0]).__new__(type(epochs_list[0]))
+    out.__setstate__(state)
+    return out
+
+
 @verbose
 def concatenate_epochs(
     epochs_list, add_offset=True, *, on_mismatch="raise", verbose=None
@@ -4725,6 +4784,8 @@ def concatenate_epochs(
     -----
     .. versionadded:: 0.9.0
     """
+    if epochs_list and type(epochs_list[0]).__name__ == "EpochsSpectrum":
+        return _concatenate_epochs_spectrum(epochs_list, add_offset=add_offset)
     (
         info,
         data,

mne/time_frequency/tests/test_spectrum.py

@@ -14,13 +14,14 @@
 from mne import (
     Annotations,
     BaseEpochs,
+    EpochsArray,
     create_info,
     grand_average,
     make_fixed_length_epochs,
 )
 from mne.channels import equalize_channels
 from mne.io import RawArray
-from mne.time_frequency import read_spectrum
+from mne.time_frequency import EpochsSpectrum, read_spectrum
 from mne.time_frequency.multitaper import _psd_from_mt
 from mne.time_frequency.spectrum import (
     EpochsSpectrumArray,
@@ -760,3 +761,44 @@ def test_plot_spectrum_dB(raw_spectrum, dB, amplitude):
         want = (20 if amplitude else 10) * np.log10(want)
 
     assert want == got, f"expected {want}, got {got}"
+
+
+def test_concatenate_epochs_spectrum():
+    """Test concatenate_epochs() works for EpochsSpectrum instances."""
+    from mne.epochs import concatenate_epochs
+
+    info = create_info(ch_names=["EEG1", "EEG2", "EEG3"], sfreq=256.0, ch_types="eeg")
+    data = np.random.randn(20, 3, 512)
+    events = np.column_stack([np.arange(20) * 512, np.zeros(20, int), np.ones(20, int)])
+    epochs = EpochsArray(data, info, events=events, tmin=0)
+    sp1 = epochs[:10].compute_psd()
+    sp2 = epochs[10:20].compute_psd()
+
+    # basic concatenation
+    out = concatenate_epochs([sp1, sp2])
+    assert isinstance(out, EpochsSpectrum)
+    assert out.shape[0] == 20
+    assert out.shape[1] == 3
+    assert_array_equal(out.freqs, sp1.freqs)
+    assert len(out.events) == 20
+
+    # event offset: second block's events should be shifted forward
+    assert np.all(out.events[10:, 0] > out.events[:10, 0])
+
+    # add_offset=False: event times should be preserved as-is
+    out_no_offset = concatenate_epochs([sp1, sp2], add_offset=False)
+    assert_array_equal(out_no_offset.events[:10, 0], sp1.events[:, 0])
+    assert_array_equal(out_no_offset.events[10:, 0], sp2.events[:, 0])
+
+    # data integrity
+    assert_array_equal(out.data[:10], sp1.data)
+    assert_array_equal(out.data[10:], sp2.data)
+
+    # mismatched freqs should raise
+    sp_bad_freqs = epochs[:10].compute_psd(fmax=50)
+    with pytest.raises(ValueError, match="freqs do not match"):
+        concatenate_epochs([sp1, sp_bad_freqs])
+
+    # passing a non-EpochsSpectrum should raise
+    with pytest.raises(TypeError, match="must be an instance of EpochsSpectrum"):
+        concatenate_epochs([sp1, epochs[:10]])