diff --git a/frites/conn/conn_spec.py b/frites/conn/conn_spec.py index 2c63abf05..22dc6f3d5 100644 --- a/frites/conn/conn_spec.py +++ b/frites/conn/conn_spec.py @@ -22,7 +22,15 @@ ############################################################################### def _coh(w, kernel, foi_idx, x_s, x_t, kw_para, average): - """Pairwise coherence.""" + """Pairwise coherence. + + `w` is either (n_trials, n_roi, n_freqs, n_times) or, when `average` is + True, (n_trials, n_roi, n_tapers, n_freqs, n_times). In that case the + cross- and auto-spectra are averaged over tapers *before* the ratio is + taken (multitaper estimator). Averaging the per-taper coherence instead + is wrong: for a single taper, single trial and no smoothing + |s_xy|^2 == s_xx * s_yy, so every per-taper coherence is identically 1. + """ # auto spectra (faster that w * w.conj()) s_auto = w.real ** 2 + w.imag ** 2 @@ -31,17 +39,18 @@ def _coh(w, kernel, foi_idx, x_s, x_t, kw_para, average): # define the pairwise coherence def pairwise_coh(w_x, w_y): - # computes the coherence - s_xy = w[:, w_y, :, :] * np.conj(w[:, w_x, :, :]) + # computes the (smoothed) cross- and auto-spectra + s_xy = w[:, w_y, ...] * np.conj(w[:, w_x, ...]) s_xy = _smooth_spectra(s_xy, kernel) - s_xx = s_auto[:, w_x, :, :] - s_yy = s_auto[:, w_y, :, :] - # average over tapers - out = np.abs(s_xy) ** 2 / (s_xx * s_yy) + s_xx = s_auto[:, w_x, ...] + s_yy = s_auto[:, w_y, ...] + # average the spectra over tapers (not the coherence) if average: - out = np.mean(out, axis=1) + s_xy, s_xx, s_yy = s_xy.mean(1), s_xx.mean(1), s_yy.mean(1) + # computes the coherence + out = np.abs(s_xy) ** 2 / (s_xx * s_yy) # mean inside frequency sliding window (if needed) - if isinstance(foi_idx, np.ndarray): + if foi_idx is not None: return _foi_average(out, foi_idx) else: return out @@ -54,22 +63,27 @@ def pairwise_coh(w_x, w_y): def _plv(w, kernel, foi_idx, x_s, x_t, kw_para, average): - """Pairwise phase-locking value.""" + """Pairwise phase-locking value. + + With a taper axis (`average=True`), the unit phase vectors are averaged + over tapers before taking the modulus, i.e. tapers act as additional + phase samples (like the temporal smoothing does). + """ # define the pairwise plv def pairwise_plv(w_x, w_y): # computes the plv - s_xy = w[:, w_y, :, :] * np.conj(w[:, w_x, :, :]) + s_xy = w[:, w_y, ...] * np.conj(w[:, w_x, ...]) # complex exponential of phase differences exp_dphi = s_xy / np.abs(s_xy) # smooth e^(-i*\delta\phi) exp_dphi = _smooth_spectra(exp_dphi, kernel) + # average the phase vectors over tapers + if average: + exp_dphi = np.mean(exp_dphi, axis=1) # computes plv out = np.abs(exp_dphi) - # average over tapers - if average: - out = np.mean(out, axis=1) # mean inside frequency sliding window (if needed) - if isinstance(foi_idx, np.ndarray): + if foi_idx is not None: return _foi_average(out, foi_idx) else: return out @@ -82,13 +96,17 @@ def pairwise_plv(w_x, w_y): def _cs(w, kernel, foi_idx, x_s, x_t, kw_para, average): - """Pairwise cross-spectra.""" + """Pairwise cross-spectra s_xy = w_x * conj(w_y). + + Note the phase convention: arg(s_xy) = phase(x) - phase(y), where x is + the first (source) and y the second (target) region of the pair. + """ # define the pairwise cross-spectra def pairwise_cs(w_x, w_y): # computes the cross-spectra - out = w[:, w_x, :, :] * np.conj(w[:, w_y, :, :]) + out = w[:, w_x, ...] * np.conj(w[:, w_y, ...]) out = _smooth_spectra(out, kernel) - # average over tapers + # average the cross-spectra over tapers if average: out = np.mean(out, axis=1) if foi_idx is not None: @@ -132,9 +150,18 @@ def conn_spec( * 'coh' : Coherence * 'plv' : Phase-Locking Value (PLV) - * 'sxy' : Cross-spectrum - - By default, the coherenc is used. + * 'sxy' : Cross-spectrum (complex). The output dtype is promoted + to a complex type if needed. The phase convention is + arg(sxy) = phase(source) - phase(target), sources and targets + being the first and second region of each pair. + + By default, the coherence is used. + + .. note:: + These are single-trial estimates: the temporal (`sm_times`) and + frequency (`sm_freqs`) smoothing, and the tapers in 'multitaper' + mode, are the only averaging involved. Without any of them the + single-trial coherence and PLV are identically equal to one. freqs : array_like Array of central frequencies of shape (n_freqs,). roi : array_like | None @@ -150,15 +177,20 @@ def conn_spec( shapes (n_foi, 2) defining where each band of interest start and finish. sm_times : float | .5 - Number of points to consider for the temporal smoothing in seconds. By - default, a 500ms smoothing is used. + Temporal smoothing in seconds. By default, a 500ms smoothing is + used. The kernel length in samples is rounded (after decimation) and + is at least one point (no smoothing). sm_freqs : int | 1 Number of points for frequency smoothing. By default, 1 is used which is equivalent to no smoothing - kernel : {'square', 'hanning'} - Kernel type to use. Choose either 'square' or 'hanning' + sm_kernel : {'square', 'hanning'} + Kernel type to use. Choose either 'square' or 'hanning'. For the + 'hanning' kernel, `sm_times` and `sm_freqs` are the number of non-zero + taps. mode : {'morlet', 'multitaper'} - Spectrum estimation mode can be either: 'multitaper' or 'morlet'. + Spectrum estimation mode can be either: 'multitaper' or 'morlet'. In + 'multitaper' mode the cross- and auto-spectra are averaged over the + DPSS tapers (MNE uses floor(mt_bandwidth - 1) tapers). n_cycles : array_like | 7. Number of cycles to use for each frequency. If a float or an integer is used, the same number of cycles is going to be used for all frequencies @@ -181,6 +213,12 @@ def conn_spec( n_jobs : int | 1 Number of jobs to use for parallel computing (use -1 to use all jobs). The parallel loop is set at the pair level. + dtype : numpy dtype | np.float32 + Output dtype. Promoted to np.complex64 for the complex 'sxy' metric + when a real dtype is given. + mean_trials : bool | False + Average the connectivity across trials (output without the `trials` + dimension). kw_links : dict | {} Additional arguments for selecting links to compute are passed to the function :func:`frites.conn.conn_links` @@ -203,6 +241,13 @@ def conn_spec( 'sxy': (_cs, "Cross-spectrum") }[metric] + # the cross-spectrum is complex : a real container would silently drop the + # imaginary part (numpy only emits a ComplexWarning) + if (metric == 'sxy') and not np.issubdtype(dtype, np.complexfloating): + logger.info(f"Cross-spectrum is complex, dtype promoted from " + f"{np.dtype(dtype).name} to complex64") + dtype = np.complex64 + # _________________________________ INPUTS ________________________________ # inputs conversion kw_links.update({'directed': False, 'net': False}) @@ -239,9 +284,18 @@ def conn_spec( # Create smoothing kernel kernel = _create_kernel(sm_times, sm_freqs, kernel=sm_kernel) - # average over tapers - tapers_average = False - # tapers_average = mode == 'multitaper' + # in multitaper mode, _tf_decomp keeps the taper axis and the metric + # functions average the spectra over tapers + tapers_average = mode == 'multitaper' + + # single-trial coherence / plv are identically 1 without any averaging + _kernels = kernel if isinstance(kernel, list) else [kernel] + if (metric in ('coh', 'plv')) and (not tapers_average) and all( + k.size == 1 for k in _kernels): + logger.warning( + f"Single-trial {f_name} without temporal (sm_times) or frequency " + "(sm_freqs) smoothing is identically equal to one. Increase " + "sm_times / sm_freqs or use mode='multitaper'") # define arguments for parallel computing mesg = f'Estimating pairwise {f_name} for trials %s' diff --git a/frites/conn/conn_tf.py b/frites/conn/conn_tf.py index 4d65dbb59..bce27cd5f 100644 --- a/frites/conn/conn_tf.py +++ b/frites/conn/conn_tf.py @@ -55,7 +55,13 @@ def _tf_decomp(data, sf, freqs, mode='morlet', n_cycles=7.0, mt_bandwidth=None, Returns ------- out : array_like - Time-frequency transform of shape (n_epochs, n_chans, n_freqs, n_times) + Complex time-frequency coefficients. In 'morlet' mode the shape is + (n_epochs, n_chans, n_freqs, n_times). In 'multitaper' mode the taper + axis is kept, (n_epochs, n_chans, n_tapers, n_freqs, n_times), so that + the spectra (and not the complex coefficients) can be averaged over + tapers downstream. DPSS tapers are orthogonal, hence averaging the + complex coefficients across tapers cancels most of the signal and does + not give the multitaper estimator. """ if mode == 'morlet': out = tfr_array_morlet( @@ -67,6 +73,9 @@ def _tf_decomp(data, sf, freqs, mode='morlet', n_cycles=7.0, mt_bandwidth=None, # Frequency center if isinstance(mt_bandwidth, (list, tuple, np.ndarray)): # Arrays freqs, n_cycles, mt_bandwidth should have the same size + # (a scalar n_cycles is broadcasted over frequencies) + n_cycles = np.broadcast_to(np.asarray(n_cycles, dtype=float), + (len(freqs),)) assert len(freqs) == len(n_cycles) == len(mt_bandwidth) out = [] for f_c, n_c, mt in zip(freqs, n_cycles, mt_bandwidth): @@ -74,15 +83,17 @@ def _tf_decomp(data, sf, freqs, mode='morlet', n_cycles=7.0, mt_bandwidth=None, data, sf, [f_c], n_cycles=float(n_c), time_bandwidth=mt, output='complex', decim=decim, n_jobs=n_jobs, **kw_mt ) - out.append(_out.mean(2)) - out = np.concatenate(out, axis=2) + out.append(_out) + # concatenate along the frequency axis (tapers axis is kept) + out = np.concatenate(out, axis=3) elif isinstance(mt_bandwidth, (type(None), int, float)): out = tfr_array_multitaper( data, sf, freqs, n_cycles=n_cycles, time_bandwidth=mt_bandwidth, output='complex', decim=decim, n_jobs=n_jobs, **kw_mt) - # mean across tapers - out = out.mean(axis=2) + else: + raise TypeError("mt_bandwidth should be None, a number or an " + "array of numbers (one per frequency)") else: raise ValueError('Method should be either "morlet" or "multitaper"') @@ -126,7 +137,12 @@ def _create_kernel(sm_times, sm_freqs, kernel='hanning'): if kernel == 'square': return np.full((sm_freqs, sm_times), 1. / (sm_times * sm_freqs)) elif kernel == 'hanning': - hann_t, hann_f = np.hanning(sm_times), np.hanning(sm_freqs) + # np.hanning(n) has zero end points, so np.hanning(2) is all zeros + # (NaN kernel after normalisation) and np.hanning(3) is a delta (no + # smoothing at all). Use the n interior points of a (n + 2) window so + # that sm_times / sm_freqs is the number of non-zero taps. + hann_t = np.hanning(sm_times + 2)[1:-1] + hann_f = np.hanning(sm_freqs + 2)[1:-1] hann = hann_f.reshape(-1, 1) * hann_t.reshape(1, -1) return hann / np.sum(hann) else: @@ -171,6 +187,9 @@ def _smooth_spectra(spectra, kernel, scale=False, decim=1): def __smooth_spectra(spectra, kernel, axes): """Single kernel smoothing.""" # fill potentially missing dimensions + if kernel.ndim > spectra.ndim: + raise ValueError(f"Smoothing kernel has more dimensions " + f"({kernel.ndim}) than the spectra ({spectra.ndim})") while kernel.ndim != spectra.ndim: kernel = kernel[np.newaxis, ...] diff --git a/frites/conn/tests/test_conn_spec.py b/frites/conn/tests/test_conn_spec.py index ebe1f6771..e1d341e34 100644 --- a/frites/conn/tests/test_conn_spec.py +++ b/frites/conn/tests/test_conn_spec.py @@ -1,6 +1,9 @@ import numpy as np +import xarray as xr -from frites.conn.conn_tf import _tf_decomp +from mne.time_frequency import tfr_array_morlet, tfr_array_multitaper + +from frites.conn.conn_tf import _tf_decomp, _create_kernel from frites.conn.conn_spec import conn_spec @@ -13,16 +16,18 @@ class TestConnSpec: sfreq, freqs = 200, np.arange(1, 51, 1) n_freqs = len(freqs) n_cycles = freqs / 2 + # MNE uses floor(time_bandwidth - 1) tapers (default time_bandwidth=4) + n_tapers = 3 times = np.arange(0, n_times // sfreq, 1 / sfreq) eta = np.random.normal(0, 1, size=(n_epochs, n_roi, n_times)) def test_tf_decomp(self, ): - # Test output shape + # Test output shape (the taper axis is kept in multitaper mode) for mode in ["morlet", "multitaper"]: out = _tf_decomp(self.eta, self.sfreq, self.freqs, mode=mode, n_cycles=self.n_cycles, n_jobs=1) - self.__assert_shape(out.shape, conn=False) + self.__assert_shape(out.shape, conn=False, mode=mode) # For multitaper test both single and array mt_bandwidth out1 = _tf_decomp(self.eta, self.sfreq, self.freqs, mode="multitaper", @@ -31,6 +36,13 @@ def test_tf_decomp(self, ): n_cycles=self.n_cycles, mt_bandwidth=[4] * self.n_freqs, n_jobs=1) np.testing.assert_array_equal(out1, out2) + # scalar n_cycles with array mt_bandwidth (broadcasted) + out3 = _tf_decomp(self.eta, self.sfreq, self.freqs, mode="multitaper", + n_cycles=3., mt_bandwidth=[4] * self.n_freqs, + n_jobs=1) + out4 = _tf_decomp(self.eta, self.sfreq, self.freqs, mode="multitaper", + n_cycles=3., mt_bandwidth=4, n_jobs=1) + np.testing.assert_allclose(out3, out4, rtol=1e-9, atol=1e-12) ################################################################## # Compare the auto-spectra with groundtruth ################################################################## @@ -41,12 +53,12 @@ def test_tf_decomp(self, ): out = _tf_decomp(x, self.sfreq, self.freqs, mode=mode, n_cycles=self.n_cycles, n_jobs=1) - out = (out * np.conj(out)).real + out = self.__power(out, mode) if mode == "morlet": val, atol = 20, 2 else: - val, atol = 5.8, 0.35 + val, atol = 16.6, 1 idx_f = self.__get_freqs_indexes(28, 32) actual = out.mean(axis=(0, -1))[:, idx_f].mean(1) np.testing.assert_allclose( @@ -57,12 +69,12 @@ def test_tf_decomp(self, ): out = _tf_decomp(x, self.sfreq, self.freqs, mode=mode, n_cycles=self.n_cycles, n_jobs=1) - out = (out * np.conj(out)).real + out = self.__power(out, mode) if mode == "morlet": val, atol = 11, 1 else: - val, atol = 3.2, 0.3 + val, atol = 9.2, 0.6 actual1 = out.mean( axis=(0, -1))[:, self.__get_freqs_indexes(8, 12)].mean(1) actual2 = out.mean( @@ -123,14 +135,89 @@ def test_conn_spec(self,): np.testing.assert_allclose(actual_2, val * np.ones_like(actual_2), atol=0.1) + def test_multitaper_reference(self): + """Multitaper sxy / coh / plv against a hand-written estimator. + + The reference averages the per-taper cross- and auto-spectra (not the + complex coefficients, and not the per-taper coherence). + """ + x, freqs, sfreq = self.__get_coupled_pair() + n_cycles, tw = 5, 4 + kw = dict(sfreq=sfreq, freqs=freqs, sm_times=0, n_cycles=n_cycles, + roi='roi', times='times', verbose=False, n_jobs=1, + mode='multitaper', mt_bandwidth=tw) + + # reference : (trials, chans, tapers, freqs, times) + w = tfr_array_multitaper(x.data, sfreq, freqs, n_cycles=n_cycles, + time_bandwidth=tw, output='complex', + verbose=False) + s_xy = (w[:, 0] * np.conj(w[:, 1])).mean(1) + s_xx = (np.abs(w[:, 0]) ** 2).mean(1) + s_yy = (np.abs(w[:, 1]) ** 2).mean(1) + coh_ref = np.abs(s_xy) ** 2 / (s_xx * s_yy) + plv_ref = np.abs((w[:, 1] * np.conj(w[:, 0]) / np.abs( + w[:, 0] * w[:, 1])).mean(1)) + + sxy = conn_spec(x, metric='sxy', dtype=np.complex128, **kw) + coh = conn_spec(x, metric='coh', dtype=np.float64, **kw) + plv = conn_spec(x, metric='plv', dtype=np.float64, **kw) + np.testing.assert_allclose(sxy.sel(roi='x-y').data, s_xy, rtol=1e-6) + np.testing.assert_allclose(coh.sel(roi='x-y').data, coh_ref, + rtol=1e-6) + np.testing.assert_allclose(plv.sel(roi='x-y').data, plv_ref, + rtol=1e-6) + # the estimator is not degenerate : coupled pair > independent pair + assert 0.1 < float(coh.sel(roi='x-z').mean()) < float( + coh.sel(roi='x-y').mean()) < 0.99 + assert float(plv.max()) <= 1. + 1e-6 + + def test_sxy_complex(self): + """The cross-spectrum is complex and follows x * conj(y).""" + x, freqs, sfreq = self.__get_coupled_pair() + n_cycles = 5 + kw = dict(sfreq=sfreq, freqs=freqs, sm_times=0, n_cycles=n_cycles, + roi='roi', times='times', verbose=False, n_jobs=1) + # default (real) dtype is promoted, imaginary part is not lost + sxy = conn_spec(x, metric='sxy', **kw) + assert np.issubdtype(sxy.dtype, np.complexfloating) + assert np.abs(sxy.data.imag).max() > 0 + # morlet cross-spectrum is exactly w_x * conj(w_y) + w = tfr_array_morlet(x.data, sfreq, freqs, n_cycles=n_cycles, + output='complex', verbose=False) + sxy = conn_spec(x, metric='sxy', dtype=np.complex128, **kw) + np.testing.assert_allclose(sxy.sel(roi='x-y').data, + w[:, 0] * np.conj(w[:, 1]), rtol=1e-6) + + def test_hanning_kernel(self): + """Small hanning kernels are neither NaN nor a delta.""" + for n in [1, 2, 3, 5]: + k = _create_kernel(n, 1, kernel='hanning') + assert k.shape == (1, n) + assert np.isfinite(k).all() + np.testing.assert_allclose(k.sum(), 1.) + assert (k > 0).all() + # sm_times=2 samples used to be NaN everywhere + kw = dict(sfreq=self.sfreq, freqs=self.freqs, n_jobs=1, verbose=False, + n_cycles=self.n_cycles, times=self.times, + sm_kernel='hanning') + out = conn_spec(self.eta, sm_times=2. / self.sfreq, metric='coh', + **kw) + assert np.isfinite(out.data).all() + # smoothed single-trial coherence of independent noise is well below 1 + out = conn_spec(self.eta, sm_times=2., metric='coh', **kw) + assert float(out.mean()) < 0.6 + ################################################################## # Assertion private methods ################################################################## - def __assert_shape(self, shape, conn=True): + def __assert_shape(self, shape, conn=True, mode='morlet'): if conn: assert shape == (self.n_epochs, self.n_edges, self.n_freqs, self.n_times) + elif mode == 'multitaper': + assert shape == (self.n_epochs, self.n_roi, self.n_tapers, + self.n_freqs, self.n_times) else: assert shape == (self.n_epochs, self.n_roi, self.n_freqs, self.n_times) @@ -171,9 +258,33 @@ def __get_freqs_indexes(self, f_low, f_high): """ Get the indexes of a range of frequencies in the freqs array """ return np.logical_and(self.freqs >= f_low, self.freqs <= f_high) + @staticmethod + def __power(w, mode): + """Power from the complex coefficients (averaged over tapers).""" + pw = (w * np.conj(w)).real + return pw.mean(2) if mode == 'multitaper' else pw + + @staticmethod + def __get_coupled_pair(): + """Three channels : x and y share a lagged 20Hz component, z is + independent noise.""" + rng = np.random.RandomState(1) + sfreq, n_trials, n_times = 200., 30, 600 + times = np.arange(n_times) / sfreq + s = np.sin(2 * np.pi * 20 * times) + x = rng.randn(n_trials, 3, n_times) + x[:, 0] += 3 * s + x[:, 1] += 3 * np.roll(s, 3) + x = xr.DataArray(x, dims=('trials', 'roi', 'times'), + coords=(np.arange(n_trials), ['x', 'y', 'z'], times)) + return x, np.array([10., 20., 30.]), sfreq + if __name__ == "__main__": """ Run the tests """ test = TestConnSpec() test.test_tf_decomp() test.test_conn_spec() + test.test_multitaper_reference() + test.test_sxy_complex() + test.test_hanning_kernel()