def test_lagged_mi():
"""Test estimation of lagged MI."""
n = 10000
cov = 0.4
source = [rn.normalvariate(0, 1) for r in range(n)]
target = [0] + [sum(pair) for pair in zip(
[cov * y for y in source[0:n - 1]],
[(1 - cov) * y for y in
[rn.normalvariate(0, 1) for r in range(n - 1)]])]
source = np.array(source)
target = np.array(target)
settings = {
'discretise_method': 'equal',
'n_discrete_bins': 4,
'history': 1,
'history_target': 1,
'lag_mi': 1,
'source_target_delay': 1}
est_te_k = JidtKraskovTE(settings)
te_k = est_te_k.estimate(source, target)
est_te_d = JidtDiscreteTE(settings)
te_d = est_te_d.estimate(source, target)
est_d = JidtDiscreteMI(settings)
mi_d = est_d.estimate(source, target)
est_k = JidtKraskovMI(settings)
mi_k = est_k.estimate(source, target)
est_g = JidtGaussianMI(settings)
mi_g = est_g.estimate(source, target)
_compare_result(mi_d, te_d, 'JidtDiscreteMI', 'JidtDiscreteTE',
'lagged MI', tol=0.05)
_compare_result(mi_k, te_k, 'JidtKraskovMI', 'JidtKraskovTE',
'lagged MI', tol=0.05)
_compare_result(mi_g, te_k, 'JidtGaussianMI', 'JidtKraskovTE',
'lagged MI', tol=0.05)
def test_one_two_dim_input_kraskov():
"""Test one- and two-dimensional input for Kraskov estimators."""
expected_mi, src_one, s, target_one = _get_gauss_data(expand=False)
src_two = np.expand_dims(src_one, axis=1)
target_two = np.expand_dims(target_one, axis=1)
ar_src_one, s = _get_ar_data(expand=False)
ar_src_two = np.expand_dims(ar_src_one, axis=1)
# MI
mi_estimator = JidtKraskovMI(settings={})
mi_cor_one = mi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtKraskovMI', 'MI')
mi_cor_two = mi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtKraskovMI', 'MI')
_compare_result(mi_cor_one, mi_cor_two,
'JidtKraskovMI one dim', 'JidtKraskovMI two dim', 'MI')
# CMI
cmi_estimator = JidtKraskovCMI(settings={})
mi_cor_one = cmi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtKraskovCMI', 'CMI')
mi_cor_two = cmi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtKraskovCMI', 'CMI')
_compare_result(mi_cor_one, mi_cor_two,
'JidtKraskovMI one dim', 'JidtKraskovMI two dim', 'CMI')
# TE
te_estimator = JidtKraskovTE(settings={'history_target': 1})
mi_cor_one = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_one, expected_mi, 'JidtKraskovTE', 'TE')
mi_cor_two = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_two, expected_mi, 'JidtKraskovTE', 'TE')
_compare_result(mi_cor_one, mi_cor_two,
'JidtKraskovMI one dim', 'JidtKraskovMI two dim', 'TE')
# AIS
ais_estimator = JidtKraskovAIS(settings={'history': 2})
mi_cor_one = ais_estimator.estimate(ar_src_one)
mi_cor_two = ais_estimator.estimate(ar_src_two)
_compare_result(mi_cor_one, mi_cor_two,
'JidtKraskovAIS one dim', 'JidtKraskovAIS two dim',
'AIS (AR process)')
def test_one_two_dim_input_kraskov():
"""Test one- and two-dimensional input for Kraskov estimators."""
expected_mi, src_one, s, target_one = _get_gauss_data(expand=False,
seed=SEED)
src_two = np.expand_dims(src_one, axis=1)
target_two = np.expand_dims(target_one, axis=1)
ar_src_one, s = _get_ar_data(expand=False, seed=SEED)
ar_src_two = np.expand_dims(ar_src_one, axis=1)
# MI
mi_estimator = JidtKraskovMI(settings={})
mi_cor_one = mi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtKraskovMI', 'MI')
mi_cor_two = mi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtKraskovMI', 'MI')
_compare_result(mi_cor_one, mi_cor_two, 'JidtKraskovMI one dim',
'JidtKraskovMI two dim', 'MI')
# CMI
cmi_estimator = JidtKraskovCMI(settings={})
mi_cor_one = cmi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtKraskovCMI', 'CMI')
mi_cor_two = cmi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtKraskovCMI', 'CMI')
_compare_result(mi_cor_one, mi_cor_two, 'JidtKraskovMI one dim',
'JidtKraskovMI two dim', 'CMI')
# TE
te_estimator = JidtKraskovTE(settings={'history_target': 1})
mi_cor_one = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_one, expected_mi, 'JidtKraskovTE', 'TE')
mi_cor_two = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_two, expected_mi, 'JidtKraskovTE', 'TE')
_compare_result(mi_cor_one, mi_cor_two, 'JidtKraskovMI one dim',
'JidtKraskovMI two dim', 'TE')
# AIS
ais_estimator = JidtKraskovAIS(settings={'history': 2})
mi_cor_one = ais_estimator.estimate(ar_src_one)
mi_cor_two = ais_estimator.estimate(ar_src_two)
_compare_result(mi_cor_one, mi_cor_two, 'JidtKraskovAIS one dim',
'JidtKraskovAIS two dim', 'AIS (AR process)')
def test_lagged_mi():
"""Test estimation of lagged MI."""
n = 10000
cov = 0.4
source = [rn.normalvariate(0, 1) for r in range(n)]
target = [0] + [
sum(pair) for pair in zip([cov * y for y in source[0:n - 1]], [
(1 - cov) * y
for y in [rn.normalvariate(0, 1) for r in range(n - 1)]
])
]
source = np.array(source)
target = np.array(target)
settings = {
'discretise_method': 'equal',
'n_discrete_bins': 4,
'history': 1,
'history_target': 1,
'lag_mi': 1,
'source_target_delay': 1
}
est_te_k = JidtKraskovTE(settings)
te_k = est_te_k.estimate(source, target)
est_te_d = JidtDiscreteTE(settings)
te_d = est_te_d.estimate(source, target)
est_d = JidtDiscreteMI(settings)
mi_d = est_d.estimate(source, target)
est_k = JidtKraskovMI(settings)
mi_k = est_k.estimate(source, target)
est_g = JidtGaussianMI(settings)
mi_g = est_g.estimate(source, target)
_compare_result(mi_d,
te_d,
'JidtDiscreteMI',
'JidtDiscreteTE',
'lagged MI',
tol=0.05)
_compare_result(mi_k,
te_k,
'JidtKraskovMI',
'JidtKraskovTE',
'lagged MI',
tol=0.05)
_compare_result(mi_g,
te_k,
'JidtGaussianMI',
'JidtKraskovTE',
'lagged MI',
tol=0.05)
def test_gauss_data():
"""Test bivariate TE estimation from correlated Gaussians."""
# Generate data and add a delay one one sample.
expected_mi, source, source_uncorr, target = _get_gauss_data(seed=SEED)
source = source[1:]
source_uncorr = source_uncorr[1:]
target = target[:-1]
data = Data(np.hstack((source, source_uncorr, target)),
dim_order='sp',
normalise=False,
seed=SEED)
settings = {
'cmi_estimator': 'JidtKraskovCMI',
'n_perm_max_stat': 21,
'n_perm_min_stat': 21,
'n_perm_max_seq': 21,
'n_perm_omnibus': 21,
'max_lag_sources': 1,
'min_lag_sources': 1,
'max_lag_target': 1
}
nw = BivariateTE()
results = nw.analyse_single_target(settings,
data,
target=2,
sources=[0, 1])
te = results.get_single_target(2, fdr=False)['te'][0]
sources = results.get_target_sources(2, fdr=False)
# Assert that only the correlated source was detected.
assert len(sources) == 1, 'Wrong no. inferred sources: {0}.'.format(
len(sources))
assert sources[0] == 0, 'Wrong inferred source: {0}.'.format(sources[0])
# Compare BivarateTE() estimate to JIDT estimate.
est = JidtKraskovTE({
'history_target': 1,
'history_source': 1,
'source_target_delay': 1,
'normalise': False
})
jidt_cmi = est.estimate(source=source, target=target)
print('Estimated TE: {0:0.6f}, estimated TE using JIDT core estimator: '
'{1:0.6f} (expected: ~ {2:0.6f}).'.format(te, jidt_cmi, expected_mi))
assert np.isclose(te, jidt_cmi, atol=0.005), (
'Estimated TE {0:0.6f} differs from JIDT estimate {1:0.6f} (expected: '
'TE {2:0.6f}).'.format(te, jidt_cmi, expected_mi))
assert np.isclose(te, expected_mi, atol=0.05), (
'Estimated TE {0:0.6f} differs from expected TE {1:0.6f}.'.format(
te, expected_mi))
def test_te_gauss_data():
"""Test TE estimation on two sets of Gaussian random data.
The first test is on correlated variables, the second on uncorrelated
variables.
Note that the calculation is based on a random variable (because the
generated data is a set of random variables) - the result will be of the
order of what we expect, but not exactly equal to it in fact, there will
be a large variance around it.
"""
expected_mi, source1, source2, target = _get_gauss_data(expand=False,
seed=SEED)
# add delay of one sample
source1 = source1[1:]
source2 = source2[1:]
target = target[:-1]
settings = {
'discretise_method': 'equal',
'n_discrete_bins': 4,
'history_target': 1
}
# Test Kraskov
mi_estimator = JidtKraskovTE(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtKraskovTE', 'TE (corr.)')
_assert_result(mi_uncor, 0, 'JidtKraskovTE', 'TE (uncorr.)')
# Test Gaussian
mi_estimator = JidtGaussianTE(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtGaussianTE', 'TE (corr.)')
_assert_result(mi_uncor, 0, 'JidtGaussianTE', 'TE (uncorr.)')
# Test Discrete
mi_estimator = JidtDiscreteTE(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtDiscreteTE', 'TE (corr.)',
0.08) # More variability here
_assert_result(mi_uncor, 0, 'JidtDiscreteTE', 'TE (uncorr.)',
0.08) # More variability here
def test_te_gauss_data():
"""Test TE estimation on two sets of Gaussian random data.
The first test is on correlated variables, the second on uncorrelated
variables.
Note that the calculation is based on a random variable (because the
generated data is a set of random variables) - the result will be of the
order of what we expect, but not exactly equal to it; in fact, there will
be a large variance around it.
"""
expected_mi, source1, source2, target = _get_gauss_data(expand=False)
# add delay of one sample
source1 = source1[1:]
source2 = source2[1:]
target = target[:-1]
settings = {'discretise_method': 'equal',
'n_discrete_bins': 4,
'history_target': 1}
# Test Kraskov
mi_estimator = JidtKraskovTE(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtKraskovTE', 'TE (corr.)')
_assert_result(mi_uncor, 0, 'JidtKraskovTE', 'TE (uncorr.)')
# Test Gaussian
mi_estimator = JidtGaussianTE(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtGaussianTE', 'TE (corr.)')
_assert_result(mi_uncor, 0, 'JidtGaussianTE', 'TE (uncorr.)')
# Test Discrete
mi_estimator = JidtDiscreteTE(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtDiscreteTE', 'TE (corr.)')
_assert_result(mi_uncor, 0, 'JidtDiscreteTE', 'TE (uncorr.)')
def test_insufficient_no_points():
"""Test if estimation aborts for too few data points."""
expected_mi, source1, source2, target = _get_gauss_data(n=4)
settings = {
'kraskov_k': 4,
'theiler_t': 0,
'history': 1,
'history_target': 1,
'lag_mi': 1,
'source_target_delay': 1
}
# Test first settings combination with k==N
est = JidtKraskovTE(settings)
with pytest.raises(RuntimeError):
est.estimate(source1, target)
est = JidtKraskovMI(settings)
with pytest.raises(RuntimeError):
est.estimate(source1, target)
est = JidtKraskovCMI(settings)
with pytest.raises(RuntimeError):
est.estimate(source1, target, target)
est = JidtKraskovAIS(settings)
with pytest.raises(RuntimeError):
est.estimate(source1)
# Test a second combination with a Theiler-correction != 0
settings['theiler_t'] = 1
settings['kraskov_k'] = 2
est = JidtKraskovTE(settings)
with pytest.raises(RuntimeError):
est.estimate(source1, target)
est = JidtKraskovMI(settings)
with pytest.raises(RuntimeError):
est.estimate(source1, target)
est = JidtKraskovCMI(settings)
with pytest.raises(RuntimeError):
est.estimate(source1, target, target)
est = JidtKraskovAIS(settings)
with pytest.raises(RuntimeError):
est.estimate(source1)
def test_insufficient_no_points():
"""Test if estimation aborts for too few data points."""
expected_mi, source1, source2, target = _get_gauss_data(n=4)
settings = {
'kraskov_k': 4,
'theiler_t': 0,
'history': 1,
'history_target': 1,
'lag_mi': 1,
'source_target_delay': 1}
# Test first settings combination with k==N
est = JidtKraskovTE(settings)
with pytest.raises(RuntimeError): est.estimate(source1, target)
est = JidtKraskovMI(settings)
with pytest.raises(RuntimeError): est.estimate(source1, target)
est = JidtKraskovCMI(settings)
with pytest.raises(RuntimeError): est.estimate(source1, target, target)
est = JidtKraskovAIS(settings)
with pytest.raises(RuntimeError): est.estimate(source1)
# Test a second combination with a Theiler-correction != 0
settings['theiler_t'] = 1
settings['kraskov_k'] = 2
est = JidtKraskovTE(settings)
with pytest.raises(RuntimeError): est.estimate(source1, target)
est = JidtKraskovMI(settings)
with pytest.raises(RuntimeError): est.estimate(source1, target)
est = JidtKraskovCMI(settings)
with pytest.raises(RuntimeError): est.estimate(source1, target, target)
est = JidtKraskovAIS(settings)
with pytest.raises(RuntimeError): est.estimate(source1)