FIX: Fix grand_average to support BaseSpectrum objects (#13375)

Author: tsbinns

doc/changes/dev/13375.bugfix.rst

@@ -0,0 +1 @@
+Fix bug with :func:`mne.grand_average` not working with :class:`mne.time_frequency.Spectrum` objects, by `Thomas Binns`_.

mne/channels/channels.py

@@ -120,8 +120,11 @@ def equalize_channels(instances, copy=True, verbose=None):
     ----------
     instances : list
         A list of MNE-Python objects to equalize the channels for. Objects can
-        be of type Raw, Epochs, Evoked, AverageTFR, Forward, Covariance,
+        be of type Raw, Epochs, Evoked, Spectrum, AverageTFR, Forward, Covariance,
         CrossSpectralDensity or Info.
+
+        .. versionchanged:: 1.11
+            Added support for :class:`mne.time_frequency.Spectrum` objects.
     copy : bool
         When dropping and/or re-ordering channels, an object will be copied
         when this parameter is set to ``True``. When set to ``False`` (the
@@ -148,21 +151,24 @@ def equalize_channels(instances, copy=True, verbose=None):
     from ..forward import Forward
     from ..io import BaseRaw
     from ..time_frequency import BaseTFR, CrossSpectralDensity
+    from ..time_frequency.spectrum import BaseSpectrum
 
     # Instances need to have a `ch_names` attribute and a `pick_channels`
     # method that supports `ordered=True`.
     allowed_types = (
         BaseRaw,
         BaseEpochs,
         Evoked,
+        BaseSpectrum,
         BaseTFR,
         Forward,
         Covariance,
         CrossSpectralDensity,
         Info,
     )
     allowed_types_str = (
-        "Raw, Epochs, Evoked, TFR, Forward, Covariance, CrossSpectralDensity or Info"
+        "Raw, Epochs, Evoked, Spectrum, TFR, Forward, Covariance, CrossSpectralDensity "
+        "or Info"
     )
     for inst in instances:
         _validate_type(

mne/time_frequency/spectrum.py

@@ -503,6 +503,10 @@ def method(self):
     def nave(self):
         return self._nave
 
+    @nave.setter
+    def nave(self, nave):
+        self._nave = nave
+
     @property
     def weights(self):
         return self._weights

mne/time_frequency/tests/test_spectrum.py

@@ -10,7 +10,14 @@
 from matplotlib.colors import same_color
 from numpy.testing import assert_allclose, assert_array_equal
 
-from mne import Annotations, BaseEpochs, create_info, make_fixed_length_epochs
+from mne import (
+    Annotations,
+    BaseEpochs,
+    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.multitaper import _psd_from_mt
@@ -200,8 +207,8 @@ def test_combine_spectrum(raw_spectrum, weights):
     spectrum1 = raw_spectrum.copy()
     spectrum2 = raw_spectrum.copy()
     if weights == "nave":
-        spectrum1._nave = 1
-        spectrum2._nave = 2
+        spectrum1.nave = 1
+        spectrum2.nave = 2
         spectrum2._data *= 2
         new_spectrum = combine_spectrum([spectrum1, spectrum2], weights=weights)
         assert_allclose(new_spectrum.data, spectrum1.data * (5 / 3))
@@ -243,6 +250,25 @@ def test_combine_spectrum_error_catch(raw_spectrum):
         combine_spectrum([raw_spectrum, raw_spectrum2], weights="equal")
 
 
+def test_grand_average(raw_spectrum):
+    """Test `grand_average()` works for instances of `BaseSpectrum`."""
+    spectrum1 = raw_spectrum.copy()
+    spectrum2 = raw_spectrum.copy()
+    spectrum2._data *= 2
+    new_spectrum = grand_average([spectrum1, spectrum2])
+    assert_allclose(new_spectrum.data, spectrum1.data * 1.5)
+
+
+def test_equalize_channels(raw_spectrum):
+    """Test equalization of channels for instances of `BaseSpectrum`."""
+    spect1 = raw_spectrum.copy()
+    spect2 = spect1.copy().pick(["MEG 0122", "MEG 0111"])
+    spect1, spect2 = equalize_channels([spect1, spect2])
+
+    assert spect1.ch_names == ["MEG 0111", "MEG 0122"]
+    assert spect2.ch_names == ["MEG 0111", "MEG 0122"]
+
+
 def test_spectrum_reject_by_annot(raw):
     """Test rejecting by annotation.