def run_iaf_experiment(sample_size, num_iters, num_trials=10, verbose=False):
for trial in range(num_trials):
print('=======================================')
name = 'result_blr_low_rank_problem_notlazy_iaf_sample_size_' + str(sample_size) +\
'_trial_' + str(trial)
print(name)
model = lm.IafMap(base_dist=base_dist,
log_l=log_l,
sample_size=sample_size,
num_iters=num_iters,
rank=num_features,
num_stages=4,
depth=2,
whitening_bij=whitening_bij,
verbose=verbose,
step_size=1e-3)
losses, traces, traces_is, neg_elbos, steps = model.train()
if verbose:
training_data = np.array([losses, traces, traces_is, neg_elbos, steps])
np.savetxt(name + '_training_data', training_data)
# compute final diagnostics with many samples
h_is, h_q0 = lm.compute_h_is(base_dist, log_l, 2 * num_features, model.transformed_dist.bijector, whitening_bij)
final_q0_trace = np.trace(h_q0)
final_is_trace = np.trace(h_is)
final_elbo = lm.compute_elbo(base_dist, log_l, 2 * num_features, model.transformed_dist.bijector, whitening_bij)
final_var_diag = lm.variance_diag(base_dist, log_l, 2 * num_features, model.transformed_dist.bijector,whitening_bij)
diagnostics = [final_q0_trace, final_is_trace, final_elbo, final_var_diag]
np.savetxt(name + '_diagnostics', np.array(diagnostics))
return model
def run_lazy_experiment(lazy_rank, sample_size, num_iters, num_lazy_layers, num_trials=10, modeltype='linear', verbose=False):
for trial in range(num_trials):
print('=======================================')
name = 'result_blr_low_rank_problem_lazy_'+modeltype +\
'_rank_'+str(lazy_rank) + \
'_num_lazy_layers_'+str(num_lazy_layers) + \
'_sample_size_' + str(sample_size) + \
'_trial_'+str(trial)
print(name)
model = lm.LazyMap(base_dist=base_dist,
num_lazy_layers=num_lazy_layers,
rank=lazy_rank,
num_iters=num_iters,
log_l=log_l,
whitening_bij=whitening_bij,
h_sample_size=num_features,
sample_size=sample_size,
num_stages=4,
depth=2,
act=modeltype,
verbose=verbose,
step_size=1e-3)
losses, traces, traces_is, neg_elbos, steps, eigvals_list = model.train()
if verbose:
training_data = np.array([losses, traces, traces_is, neg_elbos, steps])
np.savetxt(name + '_training_data', training_data)
# compute final diagnostics with many samples
h_is, h_q0 = lm.compute_h_is(base_dist, log_l, 2*num_features, model.transformed_dist.bijector, whitening_bij)
final_q0_trace = np.trace(h_q0)
final_is_trace = np.trace(h_is)
[eigvals, _] = np.linalg.eigh(h_q0)
eigvals_list.append(eigvals[::-1])
np.savetxt(name + '_eigvals', np.array(eigvals_list))
final_elbo = lm.compute_elbo(base_dist, log_l, 2*num_features, model.transformed_dist.bijector, whitening_bij)
final_var_diag = lm.variance_diag(base_dist, log_l, 2*num_features, model.transformed_dist.bijector, whitening_bij)
diagnostics = [final_q0_trace, final_is_trace, final_elbo, final_var_diag]
np.savetxt(name + '_diagnostics', np.array(diagnostics))
return model
bij = model.bij
traces_is_id_all = []
traces_q0_id_all = []
traces_is_bij_all = []
traces_q0_bij_all = []
#
for ss in [500]:
traces_is_id = []
traces_q0_id = []
traces_is_bij = []
traces_q0_bij = []
for i in range(100):
print('at sample i = '+str(i))
H_is, H_q0 = lm.compute_h_is(base_dist, log_l, ss, tfb.Identity(), whitening_bij=whitening_bij)
traces_is_id.append(np.trace(H_is))
traces_q0_id.append(np.trace(H_q0))
H_is, H_q0 = lm.compute_h_is(base_dist, log_l, ss, bij, whitening_bij=whitening_bij)
traces_is_bij.append(np.trace(H_is))
traces_q0_bij.append(np.trace(H_q0))
plt.figure()
bins=np.histogram(np.hstack((traces_is_id, traces_q0_id)), bins=40)[1] #get the bin edges
plt.hist(traces_is_id, bins, alpha=0.5, label='tr($H$)')
plt.hist(traces_q0_id, bins, alpha=0.5, label='tr($H^B$)')
plt.legend(loc=0,fontsize=18)
plt.xlabel('Traces',fontsize = 18)
plt.ylabel('Counts',fontsize = 18)