-
Notifications
You must be signed in to change notification settings - Fork 0
/
run_BOLFI_single.py
153 lines (115 loc) · 4.81 KB
/
run_BOLFI_single.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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
# Run a single BOLFI reconstruction
import sys
import pickle
import numpy as np
import scipy.stats as sps
import matplotlib.pyplot as plt
from functools import partial
import elfi
from abc_reconstruction.model import Model
from abc_reconstruction.prior import BoundedNormal_x, BoundedNormal_y
from abc_reconstruction.utils import PriorPosition, minimize, ConstraintLCBSC
from elfi.methods.bo.gpy_regression import GPyRegression
from elfi.methods.utils import ModelPrior
def run_BOLFI_single(index, true_x, true_y, folder):
### Setup
model = Model('XENON1T_ABC_all_pmts_on.ini')
model.change_defaults(s2_electrons = 25)
prior_mean = PriorPosition()
pattern = model(true_x, true_y)
pax_pos = model.get_latest_pax_position()
prior_pos = prior_mean(pattern)
r_bound = 47.8
pmt_mask = model.pmt_mask[:127].astype(int)
### Build Priors
px = elfi.Prior(BoundedNormal_x, r_bound, prior_pos, 64)
py = elfi.Prior(BoundedNormal_y, px, r_bound, prior_pos, 64)
### Build Model
model=elfi.tools.vectorize(model)
Y = elfi.Simulator(model, px, py, observed=pattern)
def likelihood_chisquare(y, n, w=None):
if w is not None:
y = y[:,w.astype(bool)]
n = n[:,w.astype(bool)]
n = np.clip(n, 1e-10, None)
y = np.clip(y, 1e-10, None)
res = 2 * np.sum(y - n + n * np.log(n/y), axis=1)
lres = np.log(res)
#if lres > 10:
# lres = np.ones(lres.shape) * 9
return lres
def chisquare(y, n, w=None):
if w is not None:
y = y[:,w.astype(bool)]
n = n[:,w.astype(bool)]
y = np.clip(y, 1e-1, None)
#print('y shape', y.shape)
#print('n shape', n.shape)
chisq, p = sps.chisquare(n, y, axis=1)
return np.array(np.log(chisq))
def k2_test(y, n, w=None):
if w is not None:
y = y[:,w.astype(bool)]
n = n[:,w.astype(bool)]
#d, p = sps.ks_2samp(n, y) # , axis=1)
# ks_2samp does not have axis arg
ds = [sps.ks_2samp(n[0], y[i])[0] for i in range(y.shape[0])]
return np.array(ds)
def sqrt_euclid(y, n, w=None):
if w is not None:
y = y[:,w.astype(bool)]
n = n[:,w.astype(bool)]
d = np.sum(np.sqrt(np.abs(y - n)), axis=1)
return d
#likelihood_chisquare_masked = partial(likelihood_chisquare, w=pmt_mask)
#log_d = elfi.Distance(likelihood_chisquare_masked, Y)
#chisquare_masked = partial(chisquare, w=pmt_mask)
#log_d = elfi.Distance(chisquare_masked, Y)
#k2_test_masked = partial(k2_test, w=pmt_mask)
#d = elfi.Distance(k2_test_masked, Y)
#log_d = elfi.Operation(np.log, d)
#sqrt_euclid_masked = partial(sqrt_euclid, w=pmt_mask)
#d = elfi.Distance(sqrt_euclid_masked, Y)
#log_d = elfi.Operation(np.log, d)
d = elfi.Distance('euclidean', Y, w=pmt_mask)
log_d = elfi.Operation(np.log, d)
### Setup BOLFI
bounds = {'px':(-r_bound, r_bound), 'py':(-r_bound, r_bound)}
target_model = GPyRegression(log_d.model.parameter_names,
bounds=bounds)
acquisition_method = ConstraintLCBSC(target_model,
prior=ModelPrior(log_d.model),
noise_var=[0.1, 0.1],
exploration_rate=10)
bolfi = elfi.BOLFI(log_d, batch_size=1, initial_evidence=20, update_interval=1,
# bounds=bounds, # Not used when using target_model
target_model=target_model,
# acq_noise_var=[0.1, 0.1], # Not used when using acq method
acquisition_method=acquisition_method,
)
### Run BOLFI
post = bolfi.fit(n_evidence=200)
bolfi.plot_discrepancy()
plt.savefig(folder + 'bolfi_disc_%d.png' % index, dpi = 150)
plt.close()
result_BOLFI = bolfi.sample(1000, info_freq=1000)
samples = result_BOLFI.samples_array
means = result_BOLFI.sample_means
modes = sps.mode(samples).mode[0]
medians = np.median(samples, axis=0)
pax_pos['truth'] = {'x': true_x, 'y': true_y}
pax_pos['BOLFI_mean'] = {'x': means['px'], 'y': means['py']}
pax_pos['BOLFI_mode'] = {'x': modes[0], 'y': modes[1]}
pax_pos['BOLFI_median'] = {'x': medians[0], 'y': medians[1]}
return pax_pos
if __name__ == '__main__':
if len(sys.argv) == 3:
i = int(sys.argv[1])
folder = sys.argv[2]
print('Running BOLFI index %d, storing results in %s' % (i, folder))
true_pos = np.loadtxt('data/truepos')
result = run_BOLFI_single(i, true_pos[i][0], true_pos[i][1], folder)
with open(folder + "bolfi_result_%d.pkl" % i, 'wb') as f:
pickle.dump(result, f)
else:
print("Usage: run_BOLFI_single.py index output_folder")