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_discrete():
"""Test one- and two-dimensional input for discrete 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)
settings = {
'discretise_method': 'equal',
'n_discrete_bins': 4,
'history_target': 1,
'history': 2
}
# MI
mi_estimator = JidtDiscreteMI(settings=settings)
mi_cor_one = mi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtDiscreteMI', 'MI',
0.08) # More variability here
mi_cor_two = mi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtDiscreteMI', 'MI',
0.08) # More variability here
_compare_result(mi_cor_one, mi_cor_two, 'JidtDiscreteMI one dim',
'JidtDiscreteMI two dim', 'MI')
# CMI
cmi_estimator = JidtDiscreteCMI(settings=settings)
mi_cor_one = cmi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtDiscreteCMI', 'CMI',
0.08) # More variability here
mi_cor_two = cmi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtDiscreteCMI', 'CMI',
0.08) # More variability here
_compare_result(mi_cor_one, mi_cor_two, 'JidtDiscreteMI one dim',
'JidtDiscreteMI two dim', 'CMI')
# TE
te_estimator = JidtDiscreteTE(settings=settings)
mi_cor_one = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_one, expected_mi, 'JidtDiscreteTE', 'TE',
0.08) # More variability here
mi_cor_two = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_two, expected_mi, 'JidtDiscreteTE', 'TE',
0.08) # More variability here
_compare_result(mi_cor_one, mi_cor_two, 'JidtDiscreteMI one dim',
'JidtDiscreteMI two dim', 'TE')
# AIS
ais_estimator = JidtDiscreteAIS(settings=settings)
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, 'JidtDiscreteAIS one dim',
'JidtDiscreteAIS 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_mi_gauss_data():
"""Test MI estimators on correlated Gauss data.
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()
# Test Kraskov
mi_estimator = JidtKraskovMI(settings={})
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtKraskovMI', 'CMI (no cond.)')
_assert_result(mi_uncor, 0, 'JidtKraskovMI', 'CMI (uncorr., no cond.)')
# Test Gaussian
mi_estimator = JidtGaussianMI(settings={})
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtGaussianMI', 'CMI (no cond.)')
_assert_result(mi_uncor, 0, 'JidtGaussianMI', 'CMI (uncorr., no cond.)')
# Test Discrete
settings = {'discretise_method': 'equal', 'n_discrete_bins': 5}
mi_estimator = JidtDiscreteMI(settings=settings)
mi_cor = mi_estimator.estimate(source1, target)
mi_uncor = mi_estimator.estimate(source2, target)
_assert_result(mi_cor, expected_mi, 'JidtDiscreteMI', 'CMI (no cond.)',
0.08) # More variability here
_assert_result(mi_uncor, 0, 'JidtDiscreteMI', 'CMI (uncorr., no cond.)',
0.08) # More variability here
def test_discrete_mi_memerror():
"""Test exception handling for memory exhausted exceptions."""
var1, var2 = _get_mem_binary_data()
# Check that we catch instantiation error for an enormous history:
caughtException = False
settings = {
'n_discrete_bins': 1000000000
}
result = -1
try:
mi_estimator = JidtDiscreteMI(settings=settings)
result = mi_estimator.estimate(var1, var2)
print('Result of MI calc (which should not have completed) was ',
result)
except ex.JidtOutOfMemoryError:
caughtException = True
print('ex.JidtOutOfMemoryError caught as required')
assert caughtException, 'No exception instantiating MI calculator with 10^18 bins'
# Check that we instantiate correctly for a small history, even after
# the error encountered above:
caughtException = False
settings = {'n_discrete_bins': 2}
try:
mi_estimator = JidtDiscreteMI(settings=settings)
mi_estimator.estimate(var1, var2)
print('Subsequent JIDT calculation worked OK')
except ex.JidtOutOfMemoryError:
caughtException = True
assert not (
caughtException), 'Unable to instantiate MI calculator with 2 bins'
def test_one_two_dim_input_discrete():
"""Test one- and two-dimensional input for discrete 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)
settings = {'discretise_method': 'equal',
'n_discrete_bins': 4,
'history_target': 1,
'history': 2}
# MI
mi_estimator = JidtDiscreteMI(settings=settings)
mi_cor_one = mi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtDiscreteMI', 'MI')
mi_cor_two = mi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtDiscreteMI', 'MI')
_compare_result(mi_cor_one, mi_cor_two,
'JidtDiscreteMI one dim', 'JidtDiscreteMI two dim', 'MI')
# CMI
cmi_estimator = JidtDiscreteCMI(settings=settings)
mi_cor_one = cmi_estimator.estimate(src_one, target_one)
_assert_result(mi_cor_one, expected_mi, 'JidtDiscreteCMI', 'CMI')
mi_cor_two = cmi_estimator.estimate(src_two, target_two)
_assert_result(mi_cor_two, expected_mi, 'JidtDiscreteCMI', 'CMI')
_compare_result(mi_cor_one, mi_cor_two,
'JidtDiscreteMI one dim', 'JidtDiscreteMI two dim', 'CMI')
# TE
te_estimator = JidtDiscreteTE(settings=settings)
mi_cor_one = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_one, expected_mi, 'JidtDiscreteTE', 'TE')
mi_cor_two = te_estimator.estimate(src_one[1:], target_one[:-1])
_assert_result(mi_cor_two, expected_mi, 'JidtDiscreteTE', 'TE')
_compare_result(mi_cor_one, mi_cor_two,
'JidtDiscreteMI one dim', 'JidtDiscreteMI two dim', 'TE')
# AIS
ais_estimator = JidtDiscreteAIS(settings=settings)
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,
'JidtDiscreteAIS one dim', 'JidtDiscreteAIS two dim',
'AIS (AR process)')
def test_invalid_settings_input():
"""Test handling of wrong inputs for settings dictionary."""
# Wrong input type for settings dict.
with pytest.raises(TypeError):
JidtDiscreteMI(settings=1)
with pytest.raises(TypeError):
JidtDiscreteCMI(settings=1)
with pytest.raises(TypeError):
JidtDiscreteAIS(settings=1)
with pytest.raises(TypeError):
JidtDiscreteTE(settings=1)
with pytest.raises(TypeError):
JidtGaussianMI(settings=1)
with pytest.raises(TypeError):
JidtGaussianCMI(settings=1)
with pytest.raises(TypeError):
JidtGaussianAIS(settings=1)
with pytest.raises(TypeError):
JidtGaussianTE(settings=1)
with pytest.raises(TypeError):
JidtKraskovMI(settings=1)
with pytest.raises(TypeError):
JidtKraskovCMI(settings=1)
with pytest.raises(TypeError):
JidtKraskovAIS(settings=1)
with pytest.raises(TypeError):
JidtKraskovTE(settings=1)
# Test if settings dict is initialised correctly.
e = JidtDiscreteMI()
assert type(
e.settings) is dict, 'Did not initialise settings as dictionary.'
e = JidtDiscreteCMI()
assert type(
e.settings) is dict, 'Did not initialise settings as dictionary.'
e = JidtGaussianMI()
assert type(
e.settings) is dict, 'Did not initialise settings as dictionary.'
e = JidtGaussianCMI()
assert type(
e.settings) is dict, 'Did not initialise settings as dictionary.'
e = JidtKraskovMI()
assert type(
e.settings) is dict, 'Did not initialise settings as dictionary.'
e = JidtKraskovCMI()
assert type(
e.settings) is dict, 'Did not initialise settings as dictionary.'
# History parameter missing for AIS and TE estimation.
with pytest.raises(RuntimeError):
JidtDiscreteAIS(settings={})
with pytest.raises(RuntimeError):
JidtDiscreteTE(settings={})
with pytest.raises(RuntimeError):
JidtGaussianAIS(settings={})
with pytest.raises(RuntimeError):
JidtGaussianTE(settings={})
with pytest.raises(RuntimeError):
JidtKraskovAIS(settings={})
with pytest.raises(RuntimeError):
JidtKraskovTE(settings={})
def test_local_values():
"""Test estimation of local values and their return type."""
expected_mi, source, s, target = _get_gauss_data(expand=False)
ar_proc, s = _get_ar_data(expand=False)
settings = {
'discretise_method': 'equal',
'n_discrete_bins': 4,
'history_target': 1,
'history': 2,
'local_values': True
}
# MI - Discrete
mi_estimator = JidtDiscreteMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtDiscreteMI', 'MI',
0.08) # More variability here
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Gaussian
mi_estimator = JidtGaussianMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtGaussianMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Kraskov
mi_estimator = JidtKraskovMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtKraskovMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# CMI - Discrete
cmi_estimator = JidtDiscreteCMI(settings=settings)
mi = cmi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtDiscreteCMI', 'CMI',
0.08) # More variability here
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Gaussian
mi_estimator = JidtGaussianCMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtGaussianCMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Kraskov
mi_estimator = JidtKraskovCMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtKraskovCMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# TE - Discrete
te_estimator = JidtDiscreteTE(settings=settings)
mi = te_estimator.estimate(source[1:], target[:-1])
_assert_result(np.mean(mi), expected_mi, 'JidtDiscreteTE', 'TE',
0.08) # More variability here
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# TE - Gaussian
mi_estimator = JidtGaussianTE(settings=settings)
mi = mi_estimator.estimate(source[1:], target[:-1])
_assert_result(np.mean(mi), expected_mi, 'JidtGaussianTE', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# TE - Kraskov
mi_estimator = JidtKraskovTE(settings=settings)
mi = mi_estimator.estimate(source[1:], target[:-1])
_assert_result(np.mean(mi), expected_mi, 'JidtKraskovTE', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# AIS - Kraskov
ais_estimator = JidtKraskovAIS(settings=settings)
mi_k = ais_estimator.estimate(ar_proc)
assert type(mi_k) is np.ndarray, 'Local values are not a numpy array.'
# AIS - Discrete
ais_estimator = JidtDiscreteAIS(settings=settings)
mi_d = ais_estimator.estimate(ar_proc)
assert type(mi_d) is np.ndarray, 'Local values are not a numpy array.'
# TODO should we compare these?
# _compare_result(np.mean(mi_k), np.mean(mi_d),
# 'JidtKraskovAIS', 'JidtDiscreteAIS', 'AIS (AR process)')
# AIS - Gaussian
ais_estimator = JidtGaussianAIS(settings=settings)
mi_g = ais_estimator.estimate(ar_proc)
assert type(mi_g) is np.ndarray, 'Local values are not a numpy array.'
_compare_result(np.mean(mi_k), np.mean(mi_g), 'JidtKraskovAIS',
'JidtGaussianAIS', 'AIS (AR process)')
def test_local_values():
"""Test estimation of local values and their return type."""
expected_mi, source, s, target = _get_gauss_data(expand=False)
ar_proc, s = _get_ar_data(expand=False)
settings = {'discretise_method': 'equal',
'n_discrete_bins': 4,
'history_target': 1,
'history': 2,
'local_values': True}
# MI - Discrete
mi_estimator = JidtDiscreteMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtDiscreteMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Gaussian
mi_estimator = JidtGaussianMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtGaussianMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Kraskov
mi_estimator = JidtKraskovMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtKraskovMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# CMI - Discrete
cmi_estimator = JidtDiscreteCMI(settings=settings)
mi = cmi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtDiscreteCMI', 'CMI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Gaussian
mi_estimator = JidtGaussianCMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtGaussianCMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# MI - Kraskov
mi_estimator = JidtKraskovCMI(settings=settings)
mi = mi_estimator.estimate(source, target)
_assert_result(np.mean(mi), expected_mi, 'JidtKraskovCMI', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# TE - Discrete
te_estimator = JidtDiscreteTE(settings=settings)
mi = te_estimator.estimate(source[1:], target[:-1])
_assert_result(np.mean(mi), expected_mi, 'JidtDiscreteTE', 'TE')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# TE - Gaussian
mi_estimator = JidtGaussianTE(settings=settings)
mi = mi_estimator.estimate(source[1:], target[:-1])
_assert_result(np.mean(mi), expected_mi, 'JidtGaussianTE', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# TE - Kraskov
mi_estimator = JidtKraskovTE(settings=settings)
mi = mi_estimator.estimate(source[1:], target[:-1])
_assert_result(np.mean(mi), expected_mi, 'JidtKraskovTE', 'MI')
assert type(mi) is np.ndarray, 'Local values are not a numpy array.'
# AIS - Kraskov
ais_estimator = JidtKraskovAIS(settings=settings)
mi_k = ais_estimator.estimate(ar_proc)
assert type(mi_k) is np.ndarray, 'Local values are not a numpy array.'
# AIS - Discrete
ais_estimator = JidtDiscreteAIS(settings=settings)
mi_d = ais_estimator.estimate(ar_proc)
assert type(mi_d) is np.ndarray, 'Local values are not a numpy array.'
# TODO should we compare these?
# _compare_result(np.mean(mi_k), np.mean(mi_d),
# 'JidtKraskovAIS', 'JidtDiscreteAIS', 'AIS (AR process)')
# AIS - Gaussian
ais_estimator = JidtGaussianAIS(settings=settings)
mi_g = ais_estimator.estimate(ar_proc)
assert type(mi_g) is np.ndarray, 'Local values are not a numpy array.'
_compare_result(np.mean(mi_k), np.mean(mi_g),
'JidtKraskovAIS', 'JidtGaussianAIS', 'AIS (AR process)')
# Generate Gaussian test data
covariance = 0.4
n = 10000
expected_mi = np.log(1 / (1 - np.power(covariance, 2)))
source_cor = np.random.normal(0, 1, size=n) # correlated src
source_uncor = np.random.normal(0, 1, size=n) # uncorrelated src
target = (covariance * source_cor +
(1 - covariance) * np.random.normal(0, 1, size=n))
# JIDT Discrete estimators
settings = {'discretise_method': 'equal', 'num_discrete_bins': 5}
est = JidtDiscreteCMI(settings)
cmi = est.estimate(source_cor, target, source_uncor)
print('Estimated CMI: {0:.5f}, expected CMI: {1:.5f}'.format(cmi, expected_mi))
est = JidtDiscreteMI(settings)
mi = est.estimate(source_cor, target)
print('Estimated MI: {0:.5f}, expected MI: {1:.5f}'.format(mi, expected_mi))
settings['history_target'] = 1
est = JidtDiscreteTE(settings)
te = est.estimate(source_cor[1:n], target[0:n - 1])
print('Estimated TE: {0:.5f}, expected TE: {1:.5f}'.format(te, expected_mi))
settings['history'] = 1
est = JidtDiscreteAIS(settings)
ais = est.estimate(target)
print('Estimated AIS: {0:.5f}, expected AIS: ~0'.format(ais))
# JIDT Kraskov estimators
settings = {}
est = JidtKraskovCMI(settings)
cmi = est.estimate(source_cor, target, source_uncor)