def test_partial_directed_coherence():
    c = Connectivity(fourier_coefficients=np.empty((1, )))
    type(c)._MVAR_Fourier_coefficients = PropertyMock(
        return_value=np.arange(1, 5).reshape((2, 2)))
    pdc = c.partial_directed_coherence()
    assert np.allclose(pdc.sum(axis=-2), 1.0)
    assert np.all((pdc >= 0.0) & (pdc <= 1.0))
def test_pairwise_phase_consistency():
    '''Test that incoherent signals are set to zero or below
    and that differences in power are ignored.'''
    np.random.seed(0)
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 200, 5,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    magnitude1 = np.random.uniform(
        0.5, 3, (n_time_samples, n_trials, n_tapers, n_fft_samples))
    angles1 = np.random.uniform(
        0, 2 * np.pi, (n_time_samples, n_trials, n_tapers, n_fft_samples))
    magnitude2 = np.random.uniform(
        0.5, 3, (n_time_samples, n_trials, n_tapers, n_fft_samples))
    angles2 = np.random.uniform(
        0, 2 * np.pi, (n_time_samples, n_trials, n_tapers, n_fft_samples))

    fourier_coefficients[..., 0] = magnitude1 * np.exp(1j * angles1)
    fourier_coefficients[..., 1] = magnitude2 * np.exp(1j * angles2)

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    ppc = this_Conn.pairwise_phase_consistency()

    # set diagonal to zero because its always 1
    diagonal_ind = np.arange(0, n_signals)
    ppc[..., diagonal_ind, diagonal_ind] = 0

    assert np.all(ppc < np.finfo(float).eps)
def test_directed_transfer_function():
    c = Connectivity(fourier_coefficients=np.empty((1, )))
    type(c)._transfer_function = PropertyMock(
        return_value=np.arange(1, 5).reshape((2, 2)))
    dtf = c.directed_transfer_function()
    assert np.allclose(dtf.sum(axis=-1), 1.0)
    assert np.all((dtf >= 0.0) & (dtf <= 1.0))
def test_weighted_phase_lag_index_sets_zero_phase_signals_to_zero():
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 30, 1,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    fourier_coefficients[..., :] = [2 * np.exp(1j * 0), 3 * np.exp(1j * 0)]
    expected_phase_lag_index = np.zeros((2, 2))

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    assert np.allclose(this_Conn.weighted_phase_lag_index().squeeze(),
                       expected_phase_lag_index)
def test_weighted_phase_lag_index_is_same_as_phase_lag_index():
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 30, 1,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    fourier_coefficients[..., :] = [
        1 * np.exp(1j * 3 * np.pi / 4), 1 * np.exp(1j * 5 * np.pi / 4)
    ]

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    assert np.allclose(this_Conn.phase_lag_index(),
                       this_Conn.weighted_phase_lag_index())
def test_imaginary_coherence():
    '''Test that imaginary coherence sets signals with the same phase
    to zero.'''
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 30, 1,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    fourier_coefficients[..., :] = [2 * np.exp(1j * 0), 3 * np.exp(1j * 0)]
    expected_imaginary_coherence = np.zeros((2, 2))

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    assert np.allclose(this_Conn.imaginary_coherence().squeeze(),
                       expected_imaginary_coherence)
def test_phase_locking_value():
    '''Make sure phase locking value ignores magnitudes.'''
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 30, 1,
                                                                    1, 2)
    fourier_coefficients = (np.random.uniform(
        0, 2, (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals)) *
                            np.exp(1j * np.pi / 2))
    expected_phase_locking_value_magnitude = np.ones(
        fourier_coefficients.shape)
    expected_phase_locking_value_angle = np.zeros(fourier_coefficients.shape)
    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)

    assert np.allclose(np.abs(this_Conn.phase_locking_value()),
                       expected_phase_locking_value_magnitude)
    assert np.allclose(np.angle(this_Conn.phase_locking_value()),
                       expected_phase_locking_value_angle)
def test_cross_spectrum(axis):
    '''Test that the cross spectrum is correct for each dimension.'''
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (2, 2, 2, 2,
                                                                    2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    signal_fourier_coefficient = [
        2 * np.exp(1j * np.pi / 2), 3 * np.exp(1j * -np.pi / 2)
    ]
    fourier_ind = [slice(0, 4)] * 5
    fourier_ind[-1] = slice(None)
    fourier_ind[axis] = slice(1, 2)
    fourier_coefficients[fourier_ind] = signal_fourier_coefficient

    expected_cross_spectral_matrix = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals,
         n_signals),
        dtype=np.complex)

    expected_slice = np.array([[4, -6], [-6, 9]], dtype=np.complex)
    expected_ind = [slice(0, 5)] * 6
    expected_ind[-1] = slice(None)
    expected_ind[-2] = slice(None)
    expected_ind[axis] = slice(1, 2)
    expected_cross_spectral_matrix[expected_ind] = expected_slice

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    assert np.allclose(expected_cross_spectral_matrix,
                       this_Conn._cross_spectral_matrix)
def test_power():
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 1, 1, 1,
                                                                    2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    fourier_coefficients[..., :] = [
        2 * np.exp(1j * np.pi / 2), 3 * np.exp(1j * -np.pi / 2)
    ]

    expected_power = np.zeros((n_time_samples, n_fft_samples, n_signals))

    expected_power[..., :] = [4, 9]

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    assert np.allclose(expected_power, this_Conn.power())
def connectivity_to_xarray(m, method='coherence_magnitude', signal_names=None,
                           squeeze=False, **kwargs):
    """
    calculate connectivity using `method`. Returns an xarray
    with dimensions of ['Time', 'Frequency', 'Source', 'Target']
    or ['Time', 'Frequency'] if squeeze=True

    Parameters
    -----------
    signal_names : iterable of strings
        Sames of time series used to name the 'Source' and 'Target' axes of
        xarray.
    squeeze : bool
        Whether to only take the first and last source and target time series.
        Only makes sense for one pair of signals and symmetrical measures

    """
    if ((method in ['power', 'group_delay', 'canonical_coherence']) or ('directed' in method)):
        raise NotImplementedError(f'{method} is not supported by xarray interface')
    connectivity = Connectivity.from_multitaper(m)
    if method == 'canonical_coherence':
        connectivity_mat, labels = getattr(connectivity, method)(**kwargs)
    else:
        connectivity_mat = getattr(connectivity, method)(**kwargs)
    # Only one couple (only makes sense for symmetrical metrics)
    if (m.time_series.shape[-1] == 2) and squeeze:
        logger.warning(
            f'Squeeze is on, but there are {m.time_series.shape[-1]} pairs!')
        connectivity_mat = connectivity_mat[:, :, 0, -1]
        xar = xr.DataArray(connectivity_mat,
                           coords=[connectivity.time,
                                   connectivity.frequencies],
                           dims=['Time', 'Frequency'])

    else:  # Name the source and target axes
        if signal_names is None:
            signal_names = np.arange(m.time_series.shape[-1])

        xar = xr.DataArray(connectivity_mat,
                           coords=[connectivity.time, connectivity.frequencies,
                                   signal_names, signal_names],
                           dims=['Time', 'Frequency', 'Source', 'Target'])

    xar.name = method

    for attr in dir(m):
        if (attr[0] == '_') or (attr == 'time_series'):
            continue
        # If we don't add 'mt_', get:
        # TypeError: '.dt' accessor only available for DataArray with
        # datetime64 timedelta64 dtype
        # or for arrays containing cftime datetime objects.
        xar.attrs['mt_' + attr] = getattr(m, attr)

    return xar
def test_n_observations(expectation_type, expected_n_observations):
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 2, 3, 4,
                                                                    5)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    this_Conn = Connectivity(
        fourier_coefficients=fourier_coefficients,
        expectation_type=expectation_type,
    )
    assert this_Conn.n_observations == expected_n_observations
def test_debiased_squared_phase_lag_index():
    '''Test that incoherent signals are set to zero or below.'''
    np.random.seed(0)
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 200, 5,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    angles1 = np.random.uniform(
        0, 2 * np.pi, (n_time_samples, n_trials, n_tapers, n_fft_samples))
    angles2 = np.random.uniform(
        0, 2 * np.pi, (n_time_samples, n_trials, n_tapers, n_fft_samples))

    fourier_coefficients[..., 0] = np.exp(1j * angles1)
    fourier_coefficients[..., 1] = np.exp(1j * angles2)

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)

    assert np.all(
        this_Conn.debiased_squared_phase_lag_index() < np.finfo(float).eps)
def test_phase_lag_index_sets_angles_up_to_pi_to_same_value():
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 30, 1,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)
    fourier_coefficients[..., 0] = (np.random.uniform(
        0.1, 2, (n_time_samples, n_trials, n_tapers, n_fft_samples)) *
                                    np.exp(1j * np.pi / 2))
    fourier_coefficients[..., 1] = (np.random.uniform(
        0.1, 2, (n_time_samples, n_trials, n_tapers, n_fft_samples)) *
                                    np.exp(1j * np.pi / 4))

    expected_phase_lag_index = np.zeros((2, 2))
    expected_phase_lag_index[0, 1] = 1
    expected_phase_lag_index[1, 0] = -1

    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)

    assert np.allclose(this_Conn.phase_lag_index().squeeze(),
                       expected_phase_lag_index)
def test_coherency():
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 30, 1,
                                                                    1, 2)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    fourier_coefficients[..., :] = [
        2 * np.exp(1j * np.pi / 2), 3 * np.exp(1j * -np.pi / 2)
    ]
    this_Conn = Connectivity(fourier_coefficients=fourier_coefficients)
    expected_coherence_magnitude = np.array([[np.nan, 1], [1, np.nan]])
    expected_phase = np.zeros((2, 2)) * np.nan
    expected_phase[0, 1] = np.pi
    expected_phase[1, 0] = -np.pi

    assert np.allclose(np.abs(this_Conn.coherency().squeeze()),
                       expected_coherence_magnitude,
                       equal_nan=True)
    assert np.allclose(np.angle(this_Conn.coherency().squeeze()),
                       expected_phase,
                       equal_nan=True)
def test_expectation(expectation_type, expected_shape):
    n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals = (1, 2, 3, 4,
                                                                    5)
    fourier_coefficients = np.zeros(
        (n_time_samples, n_trials, n_tapers, n_fft_samples, n_signals),
        dtype=np.complex)

    this_Conn = Connectivity(
        fourier_coefficients=fourier_coefficients,
        expectation_type=expectation_type,
    )
    expectation_function = this_Conn._expectation
    assert np.allclose(expected_shape,
                       expectation_function(fourier_coefficients).shape)