[BUG FIX] Fix CircleCI for Fetching and Spectral Connectivity Saving After averaging frequencies (#91)

Author: adam2392

.circleci/config.yml

@@ -24,6 +24,7 @@ jobs:
               if ! git remote -v | grep upstream; then
                 git remote add upstream https://github.com/mne-tools/mne-connectivity.git
               fi
+              git remote set-url upstream https://github.com/mne-tools/mne-connectivity.git
               git fetch upstream
 
         - save_cache:

doc/authors.inc

@@ -3,4 +3,5 @@
 .. _Britta Westner: https://github.com/britta-wstnr
 .. _Alexander Kroner: https://github.com/alexanderkroner
 .. _Richard Höchenberger: https://github.com/hoechenberger
-.. _Alex Rockhill: https://github.com/alexrockhill
+.. _Alex Rockhill: https://github.com/alexrockhill
+.. _Szonja Weigl: https://github.com/weiglszonja

doc/whats_new.rst

@@ -29,7 +29,7 @@ Enhancements
 Bug
 ~~~
 
--
+- Fix the output of :func:`mne_connectivity.spectral_connectivity_epochs` when ``faverage=True``, allowing one to save the Connectivity object, by `Adam Li`_ and `Szonja Weigl`_ :gh:`91`
 
 API
 ~~~

mne_connectivity/spectral/epochs.py

@@ -23,12 +23,47 @@
 from ..utils import fill_doc, check_indices
 
 
+def _compute_freqs(n_times, sfreq, cwt_freqs, mode):
+    from scipy.fft import rfftfreq
+    # get frequencies of interest for the different modes
+    if mode in ('multitaper', 'fourier'):
+        # fmin fmax etc is only supported for these modes
+        # decide which frequencies to keep
+        freqs_all = rfftfreq(n_times, 1. / sfreq)
+    elif mode == 'cwt_morlet':
+        # cwt_morlet mode
+        if cwt_freqs is None:
+            raise ValueError('define frequencies of interest using '
+                             'cwt_freqs')
+        else:
+            cwt_freqs = cwt_freqs.astype(np.float64)
+        if any(cwt_freqs > (sfreq / 2.)):
+            raise ValueError('entries in cwt_freqs cannot be '
+                             'larger than Nyquist (sfreq / 2)')
+        freqs_all = cwt_freqs
+    else:
+        raise ValueError('mode has an invalid value')
+
+    return freqs_all
+
+
+def _compute_freq_mask(freqs_all, fmin, fmax, fskip):
+    # create a frequency mask for all bands
+    freq_mask = np.zeros(len(freqs_all), dtype=bool)
+    for f_lower, f_upper in zip(fmin, fmax):
+        freq_mask |= ((freqs_all >= f_lower) & (freqs_all <= f_upper))
+
+    # possibly skip frequency points
+    for pos in range(fskip):
+        freq_mask[pos + 1::fskip + 1] = False
+    return freq_mask
+
+
 def _prepare_connectivity(epoch_block, times_in, tmin, tmax,
                           fmin, fmax, sfreq, indices,
                           mode, fskip, n_bands,
                           cwt_freqs, faverage):
     """Check and precompute dimensions of results data."""
-    from scipy.fft import rfftfreq
     first_epoch = epoch_block[0]
 
     # get the data size and time scale
@@ -68,25 +103,6 @@ def _prepare_connectivity(epoch_block, times_in, tmin, tmax,
     logger.info('    using t=%0.3fs..%0.3fs for estimation (%d points)'
                 % (tmin_true, tmax_true, n_times))
 
-    # get frequencies of interest for the different modes
-    if mode in ('multitaper', 'fourier'):
-        # fmin fmax etc is only supported for these modes
-        # decide which frequencies to keep
-        freqs_all = rfftfreq(n_times, 1. / sfreq)
-    elif mode == 'cwt_morlet':
-        # cwt_morlet mode
-        if cwt_freqs is None:
-            raise ValueError('define frequencies of interest using '
-                             'cwt_freqs')
-        else:
-            cwt_freqs = cwt_freqs.astype(np.float64)
-        if any(cwt_freqs > (sfreq / 2.)):
-            raise ValueError('entries in cwt_freqs cannot be '
-                             'larger than Nyquist (sfreq / 2)')
-        freqs_all = cwt_freqs
-    else:
-        raise ValueError('mode has an invalid value')
-
     # check that fmin corresponds to at least 5 cycles
     dur = float(n_times) / sfreq
     five_cycle_freq = 5. / dur
@@ -101,17 +117,15 @@ def _prepare_connectivity(epoch_block, times_in, tmin, tmax,
                  'unreliable.' % (np.min(fmin), dur * np.min(fmin), dur,
                                   5. / np.min(fmin), five_cycle_freq))
 
-    # create a frequency mask for all bands
-    freq_mask = np.zeros(len(freqs_all), dtype=bool)
-    for f_lower, f_upper in zip(fmin, fmax):
-        freq_mask |= ((freqs_all >= f_lower) & (freqs_all <= f_upper))
+    # compute frequencies to analyze based on number of samples,
+    # sampling rate, specified wavelet frequencies and mode
+    freqs = _compute_freqs(n_times, sfreq, cwt_freqs, mode)
 
-    # possibly skip frequency points
-    for pos in range(fskip):
-        freq_mask[pos + 1::fskip + 1] = False
+    # compute the mask based on specified min/max and decimation factor
+    freq_mask = _compute_freq_mask(freqs, fmin, fmax, fskip)
 
     # the frequency points where we compute connectivity
-    freqs = freqs_all[freq_mask]
+    freqs = freqs[freq_mask]
     n_freqs = len(freqs)
 
     # get the freq. indices and points for each band
@@ -1107,7 +1121,13 @@ def spectral_connectivity_epochs(data, names=None, method='coh', indices=None,
     if faverage:
         # for each band we return the frequencies that were averaged
         freqs = [np.mean(x) for x in freqs_bands]
+
+        # make sure freq_bands is a list of equal-length lists
+        # XXX: we lose information on which frequency points went into the
+        # computation. If h5netcdf supports numpy objects in the future, then
+        # we can change the min/max to just make it a list of lists.
         freqs_used = freqs_bands
+        freqs_used = [[np.min(band), np.max(band)] for band in freqs_used]
 
     if indices is None:
         # return all-to-all connectivity matrices

mne_connectivity/spectral/tests/test_spectral.py

@@ -15,6 +15,8 @@
     SpectralConnectivity, spectral_connectivity_epochs,
     read_connectivity, spectral_connectivity_time)
 from mne_connectivity.spectral.epochs import _CohEst, _get_n_epochs
+from mne_connectivity.spectral.epochs import (
+    _compute_freq_mask, _compute_freqs)
 
 
 def create_test_dataset(sfreq, n_signals, n_epochs, n_times, tmin, tmax,
@@ -336,6 +338,7 @@ def test_spectral_connectivity(method, mode):
                 assert (n == n2)
                 assert_array_almost_equal(times_data, times2)
 
+        # Test with faverage
         # compute same connections for two bands, fskip=1, and f. avg.
         fmin = (5., 15.)
         fmax = (15., 30.)
@@ -354,21 +357,51 @@ def test_spectral_connectivity(method, mode):
         assert (isinstance(freqs3, list))
         assert (len(freqs3) == len(fmin))
         for i in range(len(freqs3)):
-            assert np.all((freqs3[i] >= fmin[i]) &
-                          (freqs3[i] <= fmax[i]))
+            _fmin = max(fmin[i], min(cwt_freqs))
+            _fmax = min(fmax[i], max(cwt_freqs))
+            assert_allclose(freqs3[i][0], _fmin, atol=1)
+            assert_allclose(freqs3[i][1], _fmax, atol=1)
 
         # average con2 "manually" and we get the same result
+        fskip = 1
         if not isinstance(method, list):
             for i in range(len(freqs3)):
-                freq_idx = np.searchsorted(freqs2, freqs3[i])
-                con2_avg = np.mean(con2.get_data()[:, freq_idx], axis=1)
+                # now we want to get the frequency indices
+                # create a frequency mask for all bands
+                n_times = len(con2.attrs.get('times_used'))
+
+                # compute frequencies to analyze based on number of samples,
+                # sampling rate, specified wavelet frequencies and mode
+                freqs = _compute_freqs(n_times, sfreq, cwt_freqs, mode)
+
+                # compute the mask based on specified min/max and decim factor
+                freq_mask = _compute_freq_mask(
+                    freqs, [fmin[i]], [fmax[i]], fskip)
+                freqs = freqs[freq_mask]
+                freqs_idx = np.searchsorted(freqs2, freqs)
+                con2_avg = np.mean(con2.get_data()[:, freqs_idx], axis=1)
                 assert_array_almost_equal(con2_avg, con3.get_data()[:, i])
         else:
             for j in range(len(con2)):
                 for i in range(len(freqs3)):
-                    freq_idx = np.searchsorted(freqs2, freqs3[i])
-                    con2_avg = np.mean(con2[j].get_data()[:, freq_idx],
-                                       axis=1)
+                    # now we want to get the frequency indices
+                    # create a frequency mask for all bands
+                    n_times = len(con2[0].attrs.get('times_used'))
+
+                    # compute frequencies to analyze based on number of
+                    # samples, sampling rate, specified wavelet frequencies
+                    # and mode
+                    freqs = _compute_freqs(n_times, sfreq, cwt_freqs, mode)
+
+                    # compute the mask based on specified min/max and
+                    # decim factor
+                    freq_mask = _compute_freq_mask(
+                        freqs, [fmin[i]], [fmax[i]], fskip)
+                    freqs = freqs[freq_mask]
+                    freqs_idx = np.searchsorted(freqs2, freqs)
+
+                    con2_avg = np.mean(con2[j].get_data()[
+                                       :, freqs_idx], axis=1)
                     assert_array_almost_equal(
                         con2_avg, con3[j].get_data()[:, i])
 
@@ -551,3 +584,20 @@ def test_time_resolved_spectral_conn_regression(method, mode):
     conn_data = conn.get_data(output='dense')[
         :, row_triu_inds, col_triu_inds, ...]
     assert_array_almost_equal(conn_data, test_conn)
+
+
+def test_save(tmp_path):
+    """Test saving results of spectral connectivity."""
+    rng = np.random.RandomState(0)
+    n_epochs, n_chs, n_times, sfreq, f = 10, 2, 2000, 1000., 20.
+    data = rng.randn(n_epochs, n_chs, n_times)
+    sig = np.sin(2 * np.pi * f * np.arange(1000) / sfreq) * np.hanning(1000)
+    data[:, :, 500:1500] += sig
+    info = create_info(n_chs, sfreq, 'eeg')
+    tmin = -1
+    epochs = EpochsArray(data, info, tmin=tmin)
+
+    conn = spectral_connectivity_epochs(
+        epochs, fmin=(4, 8, 13, 30), fmax=(8, 13, 30, 45),
+        faverage=True)
+    conn.save(tmp_path / 'foo.nc')