-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathlfi_experiment.py
188 lines (155 loc) · 7.49 KB
/
lfi_experiment.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
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import argparse
import os
import random
import time
import yaml
import gpytorch
import matplotlib.pyplot as plt
import torch
from torch.distributions import Uniform, MultivariateNormal
from dbo.af import PlainUCB
from dbo.gp import LFIGPModel
from dbo.sampling import MCMCSampler, density_estimator, ScipyKDE
from dbo.util import cd, make_grid, save_object, plot_estimates, configure_matplotlib, load_abc_samples
from dbo.diagnostics import sample_divergence, gs_divergence
from dbo.lfi import RLObjective
def run_experiment(dimensionality, n_it, n_samples_per_it, n_chains, noise_sd,
output_directory, n_test=2000, n_burn=400, show_plots=False):
problem_prior = MultivariateNormal(torch.zeros(dimensionality), torch.eye(dimensionality))
objective = RLObjective(n_dim=dimensionality)
gp_model = LFIGPModel(n_dim=dimensionality, prior=problem_prior)
gp_model.covar_module.base_kernel.lengthscale = torch.tensor([.75, .4])
gp_model.covar_module.outputscale = 40.
f_norm = 3
obs_min = -1000.
objective_samples = load_abc_samples()
true_posterior = ScipyKDE(objective_samples)
gp_model.likelihood.noise = noise_sd**2
acq_fun = PlainUCB(gp_model, f_bound=f_norm, sigma_out=0)
af_sampler = MCMCSampler(acq_fun, problem_prior, n_chains=n_chains, use_jit=True, sampler='EMCEE')
n_points = 100
x_lb = problem_prior.mean.mean().item() - 3 * problem_prior.variance.mean().sqrt().item()
x_ub = problem_prior.mean.mean().item() + 3 * problem_prior.variance.mean().sqrt().item()
r_test = None
x_test = None
if dimensionality == 2 and show_plots:
x_test = make_grid(x_lb, x_ub, (x_ub - x_lb) / n_points)
r_test = true_posterior.log_prob(x_test).cpu()
divergences = torch.zeros(n_it)
gs_divergences = torch.zeros(n_it)
regret_bound = torch.zeros(n_it)
# BO loop
for t in range(n_it):
print(f"Sampling... {t + 1}/{n_it}")
x_t = af_sampler.sample(n_samples_per_it, factor=100, n_burn=n_burn).view(n_samples_per_it, dimensionality)
y_t = torch.zeros(n_samples_per_it)
for i in range(n_samples_per_it):
y_t[i] = objective(x_t[i].view(1, -1)).view([])
# Debug
print("Computing divergence...")
test_samples = af_sampler.last_samples
gs_divergences[t] = gs_divergence(objective_samples, test_samples)
divergences[t] = sample_divergence(objective_samples, test_samples)
regret_bound[t] = 2 * acq_fun.beta_t * gp_model(test_samples).variance.sqrt().mean()
print(f"Divergence/bound: {divergences[t]} < {regret_bound[t]}")
print(f"gsKL divergence/bound: {gs_divergences[t]}/{regret_bound[t]}")
if dimensionality == 2 and show_plots:
plt.figure(figsize=(4, 8))
af_test = af_sampler.un_dist.log_prob(x_test).cpu()
v_min = -100
v_max = 50
plt.subplot(211)
plot_estimates(r_test, x_lb, x_ub, vmin=v_min, vmax=v_max)
plt.subplot(212)
plot_estimates(af_test, x_lb, x_ub, vmin=v_min, vmax=v_max)
plt.plot(test_samples[:, 0], test_samples[:, 1], 'm.', label="MCMC")
plt.plot(x_t[:, 0], x_t[:, 1], 'ko', label="Evaluation points")
plt.legend()
plt.axis([x_lb, x_ub, x_lb, x_ub])
plt.show()
# Update models
gp_model.update(x_t, y_t.clamp(min=obs_min))
acq_fun.update()
gp_sampler = MCMCSampler(gp_model.mean, problem_prior, n_chains=n_chains, use_jit=True, sampler='EMCEE')
final_samples = gp_sampler.sample(n_test, n_burn=n_burn)
final_de = density_estimator(final_samples)
if dimensionality == 2 and show_plots:
m_test = final_de.log_prob(x_test).exp().cpu()
v_max = max(m_test.max().item(), r_test.max().item())
plt.figure(figsize=(4, 7))
plt.subplot(211)
plot_estimates(r_test.exp(), x_lb, x_ub, vmin=0, vmax=v_max)
plt.title("Target density")
plt.subplot(212)
plot_estimates(m_test, x_lb, x_ub, vmin=0, vmax=v_max)
plt.title("Model density")
plt.show()
t_steps = torch.arange(n_it).cpu() + 1
plt.figure(figsize=(4, 3))
plt.plot(t_steps * n_samples_per_it, regret_bound.cumsum(-1).cpu() / t_steps, 'b-', label="Regret bound")
plt.plot(t_steps * n_samples_per_it, divergences.cumsum(-1).cpu() / t_steps, 'k-', label="Averaged KL")
plt.plot(t_steps * n_samples_per_it, divergences.cpu(), 'k+', label="KL divergence")
plt.plot(t_steps * n_samples_per_it, gs_divergences.cumsum(-1).cpu() / t_steps, 'g-', label="Averaged gsKL")
plt.plot(t_steps * n_samples_per_it, gs_divergences.cpu(), 'g+', label="gsKL divergence")
plt.xlabel("Evaluations")
plt.legend()
plt.tight_layout()
with cd(output_directory):
plt.savefig("result.png", dpi=300)
torch.save(regret_bound, "lfi-bounds.pth")
torch.save(divergences, "lfi-divergences.pth")
torch.save(gs_divergences, "lfi-gs_divergences.pth")
torch.save(final_samples, "lfi-final_samples.pth")
save_object("lfi-gp.pkl", gp_model)
save_object("lfi-objective.pkl", objective)
if show_plots:
plt.show()
def main(args):
if args.output_directory is None:
args.output_directory = os.path.join("experiments", time.strftime("lfi-%Y-%m-%d-%H%M%S"))
print(f"Writing files to {args.output_directory}")
if not os.path.exists(args.output_directory):
os.makedirs(args.output_directory)
output_directory = args.output_directory
if args.seed is None:
seed = random.SystemRandom().getrandbits(32)
else:
seed = args.seed
torch.manual_seed(seed)
with cd(output_directory):
with open("lfi-seed.dat", "w") as f:
f.write("{}\n".format(seed))
with open("lfi-args.yaml", "w") as f:
yaml.dump(args, f)
dim = 2
log_lik_noise_sd = 4 # TODO: Estimate this from simulations
n_mcmc_chains = 1
n_samples_per_iteration = args.n_samples_per_iteration
n_iterations = args.n_iterations
n_repeats = args.n_repeats
if dim == 2:
configure_matplotlib()
for r in range(n_repeats):
if n_repeats > 1:
print("##########################################")
print(f"Running trial {r + 1} out of {n_repeats}...")
print("##########################################\n\n")
run_dir = os.path.join(output_directory, f"run-{r + 1:02}")
if not os.path.exists(run_dir):
os.makedirs(run_dir)
else:
run_dir = output_directory
run_experiment(dim, n_iterations, n_samples_per_iteration, n_mcmc_chains, log_lik_noise_sd,
run_dir, show_plots=args.show_plots)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Likelihood-free inference with KL-UCB")
parser.add_argument("-r", "--n-repeats", help="Number of repetitions", type=int, default=1)
parser.add_argument("-t", "--n-iterations", help="Number of iterations", type=int, default=10)
parser.add_argument("-o", "--output-directory", help="Output directory", default=None)
parser.add_argument("-n", "--n-samples-per-iteration", help="Number of samples to evaluate per iteration",
type=int, default=20)
parser.add_argument("-s", "--seed", help="Random number generator seed", default=None, type=int)
parser.add_argument("-p", "--show-plots", help="Show intermediate plots", action='store_true')
cl_args = parser.parse_args()
main(cl_args)
print("Done")