-
Notifications
You must be signed in to change notification settings - Fork 0
/
run_test.py
28 lines (21 loc) · 908 Bytes
/
run_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import numpy as np
import matplotlib.pyplot as plt
from sklearn.feature_selection import mutual_info_regression, mutual_info_classif
from mine_estimator import mine
def func(x):
return x
def gen_x(data_size):
return np.sign(np.random.normal(1.,1.,[data_size,1]))
def gen_y(x, data_size):
return func(x)+np.random.normal(0.1,np.sqrt(0.2),[data_size,1])
if __name__ == "__main__":
y_sample=gen_x(10000)
x_sample=gen_y(y_sample, 10000)
res, est_hist = mine(x_sample.reshape(-1, ), y_sample.reshape(-1,))
mi_numerical = mutual_info_classif(x_sample.reshape(-1, 1), y_sample.reshape(-1,))[0]
print(f'MINE output {res}')
print(f'scikit-learn output {mi_numerical}')
plt.plot(np.arange(len(est_hist)), np.array(est_hist), label='MINE estimation')
plt.plot(np.arange(len(est_hist)), np.ones(len(est_hist))*mi_numerical, label='True')
plt.legend()
plt.show()