def TE(Sy, Sx, bins, T, T_func):
"""Computes transfer entropy from X to Y (TE_XY)
Arguments:
Sy {list of float} -- timestamps of user Y
Sx {list of float} -- timestamps of user X
bins {list of int} -- bin widths in seconds, bin[0] = bin width at time t0, bin[1] = bin width at time t-1, bin[2] = bin width at time t-2, ...
T {int} -- time boundary in hours
T_func {function} -- which function to use for TE computation
Returns:
float -- transfer entropy from X to Y
float -- bias estimation
"""
Sx.sort()
Sy.sort()
k = len(bins)
prob_map, total = binning.compute_binning(Sy, Sx, bins, T)
binsX_indices = list(range(k, 2 * k))
pY_t_tmk_X_tm1_tml = binning.eliminate_bin(prob_map, [binsX_indices[0]])
pY_tm1_tmk_X_tm1_tml = binning.eliminate_bin(prob_map,
[0] + [binsX_indices[0]])
pY_t_tmk = binning.eliminate_bin(prob_map, binsX_indices)
pY_tm1_tmk = binning.eliminate_bin(prob_map, [0] + binsX_indices)
pt_bias = panzeri_treves_bias(prob_map, total, len(bins), len(bins))
te = T_func(pY_t_tmk, pY_tm1_tmk, pY_tm1_tmk_X_tm1_tml, pY_t_tmk_X_tm1_tml,
k)
return te, pt_bias
def test_bin_1(self):
Sy = [1, 5, 9, 17, 22, 29]
Sx = [1, 5, 9, 17, 22, 29]
bins = [1, 2, 3]
T = 30
res, total = compute_binning(Sy, Sx, bins, T)
self.assertEqual(total, 26)
self.assertEqual(res[(0, 0, 0, 0, 0, 0)], 3.0 / 26)
self.assertEqual(res[(0, 1, 0, 0, 1, 0)], 7.0 / 26)
self.assertEqual(res[(0, 1, 1, 0, 1, 1)], 2.0 / 26)
self.assertEqual(res[(1, 0, 0, 1, 0, 0)], 2.0 / 26)
self.assertEqual(res[(0, 0, 1, 0, 0, 1)], 9.0 / 26)
self.assertEqual(res[(1, 0, 1, 1, 0, 1)], 3.0 / 26)
pY_t_tmk = eliminate_bin(res, [3, 4, 5])
pY_tm1_tmk = eliminate_bin(res, [0, 3, 4, 5]) #ignore bin[0]
self.assertEqual(pY_t_tmk[((1, 0, 1))], 3.0 / 26.0)
self.assertEqual(pY_t_tmk[((0, 1, 1))], 2.0 / 26.0)
self.assertEqual(pY_t_tmk[((0, 1, 0))], 7.0 / 26.0)
self.assertEqual(pY_t_tmk[((0, 0, 1))], 9.0 / 26.0)
self.assertEqual(pY_t_tmk[((0, 0, 0))], 3.0 / 26.0)
self.assertEqual(pY_t_tmk[((1, 0, 0))], 2.0 / 26.0)
self.assertEqual(pY_tm1_tmk[((0, 1))], 12.0 / 26.0)
self.assertEqual(pY_tm1_tmk[((1, 1))], 2.0 / 26.0)
self.assertEqual(pY_tm1_tmk[((1, 0))], 7.0 / 26.0)
self.assertEqual(pY_tm1_tmk[((0, 0))], 5.0 / 26.0)
def test_bin_8(self):
St = [
0, 1, 2, 3, 4, 7, 9, 11, 12, 21, 24, 25, 28, 32, 36, 41, 43, 46,
47, 48, 50, 53
]
bins = [1, 1, 1]
T = 60
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 59)
self.assertEqual(res[(1, 0, 1, 0, 0, 0)], 4.0 / 59)
self.assertEqual(res[(0, 0, 0, 0, 0, 0)], 15.0 / 59)
self.assertEqual(res[(1, 1, 1, 0, 0, 0)], 4.0 / 59)
self.assertEqual(res[(1, 1, 0, 0, 0, 0)], 3.0 / 59)
self.assertEqual(res[(0, 0, 1, 0, 0, 0)], 10.0 / 59)
self.assertEqual(res[(1, 0, 0, 0, 0, 0)], 9.0 / 59)
self.assertEqual(res[(0, 1, 0, 0, 0, 0)], 10.0 / 59)
self.assertEqual(res[(0, 1, 1, 0, 0, 0)], 4.0 / 59)
res = eliminate_bin(res, [3, 4, 5])
self.assertEqual(res[((1, 1, 1))], 4.0 / 59)
self.assertEqual(res[((0, 1, 1))], 4.0 / 59)
self.assertEqual(res[((0, 0, 1))], 10.0 / 59)
self.assertEqual(res[((1, 0, 0))], 9.0 / 59)
self.assertEqual(res[((0, 1, 0))], 10.0 / 59)
self.assertEqual(res[((1, 0, 1))], 4.0 / 59)
self.assertEqual(res[((1, 1, 0))], 3.0 / 59)
self.assertEqual(res[((0, 0, 0))], 15.0 / 59)
def test_bin_2(self):
St = [1, 5, 9, 17, 22, 29]
bins = [1]
T = 30
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 31)
self.assertEqual(res[(1, 0)], 6.0 / 31.0)
self.assertEqual(res[(0, 0)], 25.0 / 31.0)
res = eliminate_bin(res, [1])
self.assertEqual(res[((0, ))], 25.0 / 31.0)
self.assertEqual(res[((1, ))], 6.0 / 31.0)
def panzeri_treves_bias(prob_map, denom, k, l):
"""Computes the PT bias correction
Arguments:
Returns:
float -- PT bias correction
"""
freq_map = {
comb: prob_map[comb] * denom
for comb in prob_map.keys() if prob_map[comb] != 0
}
freq_map = binning.eliminate_bin(freq_map, range(k, k + 1))
freq_map_Y = binning.eliminate_bin(freq_map, range(k, k + l))
N = sum(freq_map.values())
Nb_sum1 = sum([f[0] for f in freq_map_Y.keys()])
Nb_sum1_1 = 0
for f in freq_map_Y.keys():
Nb_sum1_1 += f[0]
Nb_sum2 = sum([f[0] for f in freq_map.keys()])
return (-1 / (2 * N * np.log(2))) * (Nb_sum1 - Nb_sum2)
def test_bin_4(self):
St = [10, 15, 20, 25, 30, 35, 40]
bins = [1, 4]
T = 50
k = l = 2
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 47)
self.assertEqual(res[(0, 0, 0, 0)], 12.0 / 47.0)
self.assertEqual(res[(0, 1, 0, 0)], 28.0 / 47.0)
self.assertEqual(res[(1, 0, 0, 0)], 7.0 / 47.0)
res = eliminate_bin(res, [2, 3])
self.assertEqual(res[((0, 0))], 12.0 / 47.0)
self.assertEqual(res[((1, 0))], 7.0 / 47.0)
self.assertEqual(res[((0, 1))], 28.0 / 47.0)
def test_bin_5(self):
St = [1, 5, 5, 5, 9, 9, 17, 17, 22, 29]
bins = [1, 4]
T = 30
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 27)
self.assertEqual(res[(1, 1, 0, 0)], 2.0 / 27)
self.assertEqual(res[(0, 0, 0, 0)], 5.0 / 27)
self.assertEqual(res[(0, 1, 0, 0)], 17.0 / 27)
self.assertEqual(res[(1, 0, 0, 0)], 3.0 / 27)
res = eliminate_bin(res, [2, 3])
self.assertEqual(res[((0, 1))], 17.0 / 27.0)
self.assertEqual(res[((1, 1))], 2.0 / 27.0)
self.assertEqual(res[((0, 0))], 5.0 / 27.0)
self.assertEqual(res[((1, 0))], 3.0 / 27.0)
def test_bin_7(self):
St = [
6, 7, 9, 11, 19, 20, 21, 22, 23, 25, 29, 35, 36, 37, 42, 49, 52,
53, 57, 58
]
bins = [5, 3, 9]
T = 60
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 45)
self.assertEqual(res[(1, 0, 1, 0, 0, 0)], 15.0 / 45.0)
self.assertEqual(res[(0, 1, 1, 0, 0, 0)], 6.0 / 45.0)
self.assertEqual(res[(1, 1, 1, 0, 0, 0)], 24.0 / 45.0)
res = eliminate_bin(res, [3, 4, 5])
self.assertEqual(res[((0, 1, 1))], 6.0 / 45.0)
self.assertEqual(res[((1, 0, 1))], 15.0 / 45.0)
self.assertEqual(res[((1, 1, 1))], 24.0 / 45.0)
def test_bin_6(self):
St = [1, 2, 9, 10, 14, 15, 19, 21, 22, 27, 28, 29]
bins = [4, 1]
T = 30
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 27)
self.assertEqual(res[(1, 1, 0, 0)], 7.0 / 27.0)
self.assertEqual(res[(0, 0, 0, 0)], 2.0 / 27.0)
self.assertEqual(res[(0, 1, 0, 0)], 2.0 / 27.0)
self.assertEqual(res[(1, 0, 0, 0)], 16.0 / 27.0)
res = eliminate_bin(res, [2, 3])
self.assertEqual(res[((1, 0))], 16.0 / 27.0)
self.assertEqual(res[((1, 1))], 7.0 / 27.0)
self.assertEqual(res[(0, 1)], 2.0 / 27.0)
self.assertEqual(res[((0, 0))], 2.0 / 27.0)
def test_bin_3(self):
St = [10, 20, 30, 40]
bins = [4, 4, 4, 4]
T = 50
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 36)
self.assertEqual(res[(1, 0, 0, 1, 0, 0, 0, 0)], 6.0 / 36.0)
self.assertEqual(res[(0, 1, 0, 1, 0, 0, 0, 0)], 6.0 / 36.0)
self.assertEqual(res[(0, 0, 1, 0, 0, 0, 0, 0)], 9.0 / 36.0)
self.assertEqual(res[(0, 1, 0, 0, 0, 0, 0, 0)], 9.0 / 36.0)
self.assertEqual(res[(1, 0, 1, 0, 0, 0, 0, 0)], 6.0 / 36.0)
res = eliminate_bin(res, [4, 5, 6, 7])
self.assertEqual(res[((0, 0, 1, 0))], 9.0 / 36.0)
self.assertEqual(res[((0, 1, 0, 0))], 9.0 / 36.0)
self.assertEqual(res[((1, 0, 1, 0))], 6.0 / 36.0)
self.assertEqual(res[((1, 0, 0, 1))], 6.0 / 36.0)
self.assertEqual(res[((0, 1, 0, 1))], 6.0 / 36.0)
def test_bin_9(self):
St = [
0, 2, 3, 5, 13, 14, 15, 19, 22, 24, 26, 27, 31, 33, 34, 35, 36, 44,
46, 48, 50, 52, 53
]
bins = [1, 9, 7]
T = 60
res, total = compute_binning(St, [], bins, T)
self.assertEqual(total, 45)
self.assertEqual(res[(0, 1, 1, 0, 0, 0)], 29.0 / 45)
self.assertEqual(res[(1, 1, 1, 0, 0, 0)], 14.0 / 45)
self.assertEqual(res[(1, 1, 0, 0, 0, 0)], 2.0 / 45)
res = eliminate_bin(res, [3, 4, 5])
self.assertEqual(res[((0, 1, 1))], 29.0 / 45)
self.assertEqual(res[((1, 1, 1))], 14.0 / 45)
self.assertEqual(res[((1, 1, 0))], 2.0 / 45)
def test_bin_11(self):
St = [13, 16, 17, 18]
bins = [1, 2, 3]
T = 20
res, total = compute_binning(St, St, bins, T)
self.assertEqual(total, 16)
self.assertEqual(res[(0, 0, 0, 0, 0, 0)], 8.0 / 16)
self.assertEqual(res[(0, 1, 0, 0, 1, 0)], 2.0 / 16)
self.assertEqual(res[(0, 1, 1, 0, 1, 1)], 2.0 / 16)
self.assertEqual(res[(1, 0, 0, 1, 0, 0)], 1.0 / 16)
self.assertEqual(res[(1, 0, 1, 1, 0, 1)], 1.0 / 16)
self.assertEqual(res[(1, 1, 1, 1, 1, 1)], 2.0 / 16)
res = eliminate_bin(res, [3, 4, 5])
self.assertEqual(res[(1, 1, 1)], 2.0 / 16)
self.assertEqual(res[(1, 0, 1)], 1.0 / 16)
self.assertEqual(res[(1, 0, 0)], 1.0 / 16)
self.assertEqual(res[(0, 1, 1)], 2.0 / 16)
self.assertEqual(res[(0, 1, 0)], 2.0 / 16)
self.assertEqual(res[(0, 0, 0)], 8.0 / 16)