def rvs(self, num_samps):
res_idx = system_res(range(len(self.lw)), self.lw)
rval = np.zeros((num_samps, self.centers.shape[1]))
for i in range(num_samps):
rval[i], _, _, _, _, _ = self.kern.proposal(
self.centers[res_idx[i]], -4, **{})
return rval
def smc_iteration(beg, mid, end, br_old, target_ef_fact):
(br_new, inc_w) = search_bridge_param(target_ef_fact, beg, mid, br_old)
samps_idx = (np.array(system_res(range(population_size), resampled_size=end - mid, weights=inc_w))
+ beg)
rval[mid:end] = rval[samps_idx]
lpost[mid:end] = lpost[samps_idx]
lprior[mid:end] = lprior[samps_idx]
proposal_obj.set_batch(rval[samps_idx])
# pre = (rval[mid:end].copy(), lprior[mid:end].copy(), lpost[mid:end].copy())
acc = mcmc_rejuvenate(rval[mid:end], lprior[mid:end], lpost[mid:end], br_new)
mean_acc = np.mean(acc)
acceptance_rates.append(mean_acc)
proposal_obj.next_iteration()
proposal_obj.update_step_size([mean_acc])
return (br_new, inc_w, mean_acc)
def smc_iteration(beg, mid, end, br_old, target_ef_fact):
(br_new, inc_w) = search_bridge_param(target_ef_fact, beg, mid, br_old)
samps_idx = (np.array(
system_res(range(population_size),
resampled_size=end - mid,
weights=inc_w)) + beg)
rval[mid:end] = rval[samps_idx]
lpost[mid:end] = lpost[samps_idx]
lprior[mid:end] = lprior[samps_idx]
proposal_obj.set_batch(rval[samps_idx])
# pre = (rval[mid:end].copy(), lprior[mid:end].copy(), lpost[mid:end].copy())
acc = mcmc_rejuvenate(rval[mid:end], lprior[mid:end], lpost[mid:end],
br_new)
mean_acc = np.mean(acc)
acceptance_rates.append(mean_acc)
proposal_obj.next_iteration()
proposal_obj.update_step_size([mean_acc])
return (br_new, inc_w, mean_acc)
D, target_log_pdf, target_grad
) #get_AdaptiveLangevin(D, target_log_pdf, target_grad, prec=True, step_size=1.)
start = np.zeros(D)
num_iter = 100
samples, log_target_densities, unadj_samp, unadj_log_target, logw, unw_samptimes = mini_rb_pmc(
sampler_is, start, num_iter, pop_size, D, time_budget=100000)
mom_gen = np.mean((np.array([(samples**i).mean(0)
for i in range(1, max_moment)]) - moments)**2)
mcmc_samps = mini_mcmc(sampler_mh, start, num_iter, D)
#the weights we get back are not Rao-Blackwellized, which is what we do now.
#beware: this only works if the proposal is not adapted during sampling!!
#logw = logsumexp(np.array([sampler_is.proposal_log_pdf(i, unadj_samp) for i in unadj_samp]), 0)
res_idx = system_res(range(len(logw)), logw, resampled_size=10 * len(logw))
samples = unadj_samp[res_idx]
mom_unadj = np.mean(
(np.array([(unadj_samp**i).mean(0)
for i in range(1, max_moment)]) - moments)**2)
mom_w = np.mean((np.array(
[(unadj_samp**i * exp(logw - logsumexp(logw))[:, np.newaxis]).sum(0)
for i in range(1, max_moment)]) - moments)**2)
mom_mcmc = np.mean(
(np.array([(mcmc_samps[0]**i).mean(0)
for i in range(1, max_moment)]) - moments)**2)
if False:
plt.scatter(samples.T[0],
samples.T[1],
c='r',
def mini_rb_pmc(
transition_kernel,
start,
num_iter,
pop_size,
D,
recompute_log_pdf=False,
time_budget=None,
):
# PMC results
print(pop_size)
assert (num_iter % pop_size == 0)
# following not implemented yet
assert (recompute_log_pdf == False)
proposals = np.zeros((num_iter // pop_size, pop_size, D)) + np.nan
logweights = np.zeros((num_iter // pop_size, pop_size)) + np.nan
prop_target_logpdf = np.zeros((num_iter // pop_size, pop_size)) + np.nan
prop_prob_logpdf = np.zeros((num_iter // pop_size, pop_size)) + np.nan
samples = np.zeros((num_iter, D)) + np.nan
log_pdf = np.zeros(num_iter) + np.nan
# timings for output and time limit
times = np.zeros(num_iter)
# for adaptive transition kernels
avg_accept = 0.
current = np.array([start] * pop_size)
current_log_pdf = np.zeros(pop_size) + np.nan
acc_prob = np.zeros(pop_size) + np.nan
logger.info("Starting PMC using %s in D=%d dimensions" % \
(transition_kernel.__class__.__name__, D,))
it = 0
for stage in range(num_iter // pop_size):
start_it = stage * pop_size
# stop sampling if time budget exceeded
if time_budget is not None and not np.isnan(times[start_it]):
if times[start_it] > times[0] + time_budget:
logger.info(
"Time limit of %ds exceeded. Stopping MCMC at iteration %d."
% (time_budget, it))
break
# print progress
if False and stage > 1:
log_str = "PMC iteration %d/%d, current log_pdf: %.6f, avg acceptance: %.3f" % (
it + 1, num_iter, np.nan if log_pdf[it - 1] is None else
log_pdf[it - 1], avg_accept)
logger.info(log_str)
range_it = range(start_it, start_it + pop_size)
for it in range_it:
prop_idx = it - start_it
if np.isnan(current_log_pdf[prop_idx]):
cur_lpdf = None
else:
cur_lpdf = current_log_pdf[prop_idx]
times[it] = time.time()
# marginal sampler: make transition kernel re-compute log_pdf of current state
if recompute_log_pdf:
current_log_pdf = None
# generate proposal and acceptance probability
logger.debug("Performing GRIS sample %d" % it)
proposals[stage, prop_idx], prop_target_logpdf[
stage, prop_idx], current_log_pdf[prop_idx], prop_prob_logpdf[
stage,
prop_idx], backw_logpdf, current_kwargs = transition_kernel.proposal(
current[prop_idx], cur_lpdf, **{})
#logweights[stage, prop_idx] = prop_target_logpdf[stage, prop_idx] - prop_prob_logpdf[stage, prop_idx]
#Rao-Blackwellize over all used proposals
all_prop_logpdfs = np.array([
transition_kernel.proposal_log_pdf(current[it - start_it],
proposals[stage, :])
for it in range_it
])
prop_prob_logpdf[stage, :] = logsumexp(all_prop_logpdfs, 0)
logweights[stage, :] = prop_target_logpdf[stage, :] - prop_prob_logpdf[
stage, :]
res_idx = system_res(range(pop_size), logweights[stage, :])
samples[range_it], log_pdf[range_it] = proposals[
stage, res_idx], prop_prob_logpdf[stage, res_idx]
ess = compute_ess(logweights[stage, :], normalize=True)
if ess / float(pop_size) > 0.5:
current = proposals[stage, :]
current_log_pdf = prop_target_logpdf[stage, :]
else:
current = proposals[stage, res_idx]
current_log_pdf = prop_prob_logpdf[stage, res_idx]
#print(ess, float(pop_size)/ess)
transition_kernel.next_iteration()
transition_kernel.update(np.vstack(proposals[:stage + 1, :]), pop_size,
np.hstack(logweights[:stage + 1, :]))
res_idx = system_res(range(pop_size * stage) * 10,
weights=logweights[:stage, :].flatten())
unw_samp = np.vstack(proposals[:pop_size * stage])
unw_logtarg = np.hstack(prop_target_logpdf[:pop_size * stage])
# recall it might be less than last iterations due to time budget
return samples[:it], log_pdf[:it], unw_samp, unw_logtarg, np.hstack(
logweights[:pop_size * stage]), times[:it]
def mini_smc(
num_samples, # will give size of sample in final iteration
population_size,
prior, # some distribution object that our algorithm will have no problem with
log_targ, # actual target
proposal_obj,
targ_ef_bridge=0.5,
targ_ef_stop=0.9,
ef_tolerance=0.02,
reweight=False,
across=False,
estim_evid=False,
ess=False):
# ToDO: - adaptive resampling (only use importance resampling if ESS < threshold)
# - reweight earlier iterations for actual target
# - use weighted approximation instead of approx after resampling for final iteration
"""
Sample from a geometric sequence of target distributions between prior and target,
reweighting samples from early iterations for estimating the actual target.
Uses a geometric bridge for now
= Parameters =
num_samples - size of sample in final iteration
population_size - size of particle system except for final iteration
prior - some easy target distribution, prior will likely do well
log_target - actual target
proposal_obj - object with which proposals are generated
targ_ef_bridge - target efficiency factor for bridge, i.e. what should eff/num_particles be in a bridge step
targ_ef_stop - target efficiency factor with respect to final target
ef_tolerance - efficiency factor tolerance
reweight - False (only use last iteration), True (reweight for actual target and weight iterations by ESS)
across - Resampling across iterations after reweighting? True (resample. only if reweight = True)
estim_evid - return estimate of evidence/normalizing constant of log_target
ess - Return ESS of last iteration? Defaults to False.
"""
logger.info("Starting SMC using %s" % \
(proposal_obj.get_name()))
if not reweight:
assert (not across)
log_target = lambda x: np.apply_along_axis(
lambda y: np.atleast_1d(log_targ(y)), 1, x)
# will be returned
step_sizes = [proposal_obj.step_size]
acceptance_rates = []
initial_guesses = prior.rvs(population_size)
population_size = initial_guesses.shape[0]
dim = initial_guesses.shape[1]
lprior = np.empty(num_samples + population_size * 3)
lprior[:population_size] = prior.logpdf(initial_guesses)
lpost = np.empty(num_samples + population_size * 3)
lpost[:population_size] = log_target(initial_guesses).flatten()
rval = np.r_[initial_guesses,
np.zeros((num_samples + population_size * 3, dim))]
def ensure(size):
old = len(lprior)
if old < size:
lprior.resize(old * 2)
lpost.resize(old * 2)
rval.resize((old * 2, rval.shape[1]))
def seq_value(br, prior_value, posterior_value):
"""
Compute logprobability from lprior and lpost according to
bridge parameter/temperature br
= Parameters =
br - bridge parameter/temperature
prior_value - value according to br = 0
posterior_value - value according to br = 1
"""
return prior_value * (1. - br) + posterior_value * br
def incr_weight(idx_from, idx_to, br_old, br_new, return_ef=False):
inc_w = (
seq_value(
br_new, lprior[idx_from:idx_to], lpost[idx_from:idx_to]
) # use lpost[idx_beg:idx_mid] here for approximating actual target
-
seq_value(br_old, lprior[idx_from:idx_to], lpost[idx_from:idx_to]))
assert (not np.any(np.isnan(rval)))
if return_ef:
norm_inc_w = inc_w - logsumexp(inc_w)
ESS = exp(2 * logsumexp(norm_inc_w) - logsumexp(2 * norm_inc_w))
EF = ESS / (idx_to - idx_from)
return (inc_w, EF)
return inc_w
def mcmc_rejuvenate(cur_sample, cur_lprior, cur_lpost, br):
"""
Make an MCMC move using proposal_obj, overwriting input (except br)
Returns the acceptance probabilities.
= Parameters =
cur_sample - the particles to be moved (will be overwritten with new state after move)
cur_lpost - the logposteriors according to final target distribution in distribution sequence (will be overwritten)
cur_lprior - the logpriors according to first target distribution in distribution sequence (will be overwritten)
br - the bridge parameter/temperature which determines how posterior and prior are mixed for current target distribution in the sequence
= Return =
Acceptance probabilites
"""
# set the target to be the intermediary distribution
save_target_logpdf = proposal_obj.target_log_pdf
proposal_obj.target_log_pdf = lambda x: seq_value(
br, prior.logpdf(x), save_target_logpdf(x))
if proposal_obj.__dict__.has_key('target_grad'):
save_target_grad = proposal_obj.target_grad
proposal_obj.target_grad = lambda x: seq_value(
br, prior.logpdf_grad(x), save_target_grad(x))
tmp = [
proposal_obj.proposal(
cur_sample[idx], seq_value(br, cur_lprior[idx],
cur_lpost[idx]))
for idx in range(len(cur_sample))
]
# reset the target to be the actual posterior
proposal_obj.target_log_pdf = save_target_logpdf
if proposal_obj.__dict__.has_key('target_grad'):
proposal_obj.target_grad = save_target_grad
# (prop, lprob_move_forw, lprob_move_back) = [np.array(l) for l in
# zip(*tmp)]
(prop, lprob_bridge_forw, _, lprob_move_forw, lprob_move_back,
current_kwargs) = [np.array(l) for l in zip(*tmp)]
# compute log_target for proposals
lprior_forw = prior.logpdf(prop).flatten()
lpost_forw = (lprob_bridge_forw.flatten() -
(1 - br) * lprior_forw.flatten()) / br
assert (np.allclose(lpost_forw, log_target(prop).flatten()))
# compute all acceptance probabilites
assert (not (np.any(np.isinf(lprob_move_forw))
or np.any(np.isnan(lprob_move_forw))))
assert (not (np.any(np.isnan(lprior_forw))))
assert (not (np.any(np.isnan(lpost_forw))))
assert (not (np.any(np.isinf(lprob_move_back))
or np.any(np.isnan(lprob_move_back))))
assert (not (np.any(np.isinf(cur_lprior))
or np.any(np.isnan(cur_lprior))))
assert (not (np.any(np.isinf(cur_lpost))
or np.any(np.isnan(cur_lpost))))
mh_ratio = (lprob_move_back + seq_value(br, lprior_forw, lpost_forw) -
lprob_move_forw -
seq_value(br, cur_lprior.flatten(), cur_lpost.flatten()))
assert (mh_ratio.shape == lpost_forw.shape)
acc = exp(np.min(np.c_[np.zeros_like(mh_ratio), mh_ratio], 1))
assert (not (np.any(np.isnan(acc)) or np.any(np.isinf(acc))))
move = np.random.rand(len(acc)) < acc
assert (np.mean(acc) != 0)
cur_sample[:] = prop * np.atleast_2d(move).T + cur_sample * (
1 - np.atleast_2d(move).T)
cur_lpost[:] = lpost_forw * move + cur_lpost * (1 - move)
cur_lprior[:] = prior.logpdf(
cur_sample) # lprior_forw*move + cur_lprior*(1-move)
return acc
def search_bridge_param(target_ef_fact,
idx_beg,
idx_end,
br_old,
eps=ef_tolerance):
high = 1.0
low = br_old
max_eval = 9
old_EF = 0
logger.debug('Start bridge search')
for i in range(max_eval + 1):
mid = low + (high - low) / 2
(inc_w, EF) = incr_weight(idx_beg, idx_end, br_old, mid, True)
logger.debug("EF: %.4f" % EF)
d = EF - target_ef_fact
if i == max_eval or np.abs(EF - old_EF) < eps:
return (mid, inc_w)
old_EF = EF
if d < -eps:
high = mid
elif d > eps:
low = mid
else:
return (mid, inc_w)
def smc_iteration(beg, mid, end, br_old, target_ef_fact):
(br_new, inc_w) = search_bridge_param(target_ef_fact, beg, mid, br_old)
samps_idx = (np.array(
system_res(range(population_size),
resampled_size=end - mid,
weights=inc_w)) + beg)
rval[mid:end] = rval[samps_idx]
lpost[mid:end] = lpost[samps_idx]
lprior[mid:end] = lprior[samps_idx]
proposal_obj.set_batch(rval[samps_idx])
# pre = (rval[mid:end].copy(), lprior[mid:end].copy(), lpost[mid:end].copy())
acc = mcmc_rejuvenate(rval[mid:end], lprior[mid:end], lpost[mid:end],
br_new)
mean_acc = np.mean(acc)
acceptance_rates.append(mean_acc)
proposal_obj.next_iteration()
proposal_obj.update_step_size([mean_acc])
return (br_new, inc_w, mean_acc)
br = [0.0]
evid = 0
old_EF = 0
j = 1
while True:
step_sizes += [proposal_obj.step_size]
idx_beg = (j - 1) * population_size
idx_mid = idx_beg + population_size
idx_end = idx_mid + population_size
ensure(idx_end)
(br_new, inc_w, mean_acc) = smc_iteration(idx_beg, idx_mid, idx_end,
br[j - 1], targ_ef_bridge)
br.append(br_new)
evid = evid + logsumexp(inc_w) - log(inc_w.size)
norm_inc_w = inc_w - logsumexp(inc_w)
j = j + 1
ESS = exp(2 * logsumexp(norm_inc_w) - logsumexp(2 * norm_inc_w))
logger.debug(
"At bridge distribution #%d, ESS: %.2f, mean acc: %.4f, step_size: %.4e"
% (j, ESS, mean_acc, proposal_obj.step_size))
# test how good we are with respect to the actual distribution of interest
(inc_w_final, EF_final) = incr_weight(idx_mid, idx_end, br[j - 1], 1,
True)
if (np.abs(EF_final - old_EF) < ef_tolerance
or # we're not improving much
np.abs(EF_final - targ_ef_stop) <
ef_tolerance): # we reached our desired efficiency factor
break
old_EF = EF_final
idx_beg = (len(br) - 1) * population_size
idx_mid = idx_beg + population_size
idx_end = idx_mid + num_samples
ensure(idx_end)
(br_new, inc_w, mean_acc) = smc_iteration(idx_beg, idx_mid, idx_end,
br[-1], 1)
br.append(br_new)
evid = evid + logsumexp(inc_w) - log(inc_w.size)
norm_inc_w = inc_w - logsumexp(inc_w)
ESS = exp(2 * logsumexp(norm_inc_w) - logsumexp(2 * norm_inc_w))
logger.debug("Final approx, ESS: %.2f, mean acc: %.4f" % (ESS, mean_acc))
if reweight == False:
(rval, lpost) = (rval[idx_mid:idx_end], lpost[idx_mid:idx_end])
else:
# reweight for actual target
power = 1. - np.repeat(br, population_size)
rew = (lpost[:idx_mid] - lprior[:idx_mid]) * power
# weight each iteration by ESS wrt actual target
rew_resh = rew.reshape((len(br), population_size))
logess = ((2 * logsumexp(rew_resh, 1) - logsumexp(2 * rew_resh, 1)))
print(exp(logess))
logess_sampsize_ratio = logess - log(population_size)
rew = rew + np.repeat(logess_sampsize_ratio.flatten(), population_size)
smp_idx = np.array(
system_res(range(idx_mid),
resampled_size=population_size,
weights=rew))
if across:
smp_idx_acr = np.array(
system_res(range(idx_end),
resampled_size=idx_end,
weights=np.r_[rew, np.ones(idx_end - idx_mid)]))
(rval_acr, lpost_acr) = (rval[smp_idx_acr], lpost[smp_idx_acr])
(rval, lpost) = (np.r_[rval[smp_idx], rval[idx_mid:idx_end]],
np.r_[lpost[smp_idx], lpost[idx_mid:idx_end]])
all_rval = [rval, lpost, np.array(step_sizes), np.array(acceptance_rates)]
if ess:
all_rval.append(ESS)
if across:
all_rval.extend((rval_acr, lpost_acr))
if evid:
all_rval.append(evid)
return all_rval
]
sampler_is = get_StaticLangevin(D, target_log_pdf, target_grad)#get_AdaptiveLangevin(D, target_log_pdf, target_grad)
sampler_mh = get_StaticLangevin(D, target_log_pdf, target_grad)#get_AdaptiveLangevin(D, target_log_pdf, target_grad, prec=True, step_size=1.)
start = np.zeros(D)
num_iter = 100
samples, log_target_densities, unadj_samp, unadj_log_target, logw, unw_samptimes = mini_rb_pmc(sampler_is, start, num_iter, pop_size, D, time_budget=100000)
mom_gen = np.mean((np.array([(samples**i).mean(0) for i in range(1, max_moment)]) - moments)**2)
mcmc_samps = mini_mcmc(sampler_mh, start, num_iter, D)
#the weights we get back are not Rao-Blackwellized, which is what we do now.
#beware: this only works if the proposal is not adapted during sampling!!
#logw = logsumexp(np.array([sampler_is.proposal_log_pdf(i, unadj_samp) for i in unadj_samp]), 0)
res_idx = system_res(range(len(logw)), logw, resampled_size=10*len(logw))
samples = unadj_samp[res_idx]
mom_unadj = np.mean((np.array([(unadj_samp**i).mean(0) for i in range(1, max_moment)]) - moments)**2)
mom_w = np.mean((np.array([(unadj_samp**i * exp(logw - logsumexp(logw))[:,np.newaxis]).sum(0) for i in range(1, max_moment)]) -moments)**2)
mom_mcmc = np.mean((np.array([(mcmc_samps[0]**i).mean(0) for i in range(1, max_moment)]) - moments)**2)
if False:
plt.scatter(samples.T[0], samples.T[1], c='r', marker='*', zorder=4, s=5)
# fig.suptitle("%s - importance resampled" % (sampler_is.__class__.__name__,))
plt.show()
plt.scatter(unadj_samp.T[0], unadj_samp.T[1], c = logw - logsumexp(logw), cmap = plt.get_cmap('Blues'), alpha=0.5, zorder=2) #)visualize_scatter_2d()
# plt.suptitle("%s - unadjusted Langevin" % (sampler_is.__class__.__name__,))
# plt.scatter(mcmc_samps[0].T[0], mcmc_samps[0].T[1], c='b',marker='*')
plt.show()
Log.get_logger().info('===='+str(sampler_mh.step_size)+' '+str(mcmc_samps[2].mean())+'====')
#the following two should be 0 ideally
def mini_pmc(transition_kernel, start, num_iter, pop_size, recompute_log_pdf=False, time_budget=None, weighted_update=True, rao_blackwell_generation=True, resample_at_end=False):
assert(len(start.shape) <= 2)
assert(num_iter % pop_size == 0)
# following not implemented yet
assert(recompute_log_pdf == False)
if len(start.shape) == 2:
prev = start
prev_logp = np.array([None] * start.shape[0])
D = start.shape[1]
else:
prev = np.array([start] * pop_size)
prev_logp = np.array([None] * pop_size)
D = len(start)
# PMC results
proposals = np.zeros((num_iter // pop_size, pop_size, D)) + np.nan
logweights = np.zeros((num_iter // pop_size, pop_size)) + np.nan
prop_target_logpdf = np.zeros((num_iter // pop_size, pop_size)) + np.nan
prop_prob_logpdf = np.zeros((num_iter // pop_size, pop_size)) + np.nan
samples = np.zeros((num_iter, D)) + np.nan
log_pdf = np.zeros(num_iter) + np.nan
# timings for output and time limit
times = np.zeros(num_iter)
logger.info("Starting PMC using %s in D=%d dimensions using %d particles and %d iterations" % \
(transition_kernel.get_name(), D, pop_size, num_iter / pop_size))
it = 0
time_last_printed = time.time()
for stage in range(num_iter // pop_size):
log_str = "PMC stage %d/%d" % (stage + 1, num_iter // pop_size)
current_time = time.time()
if current_time > time_last_printed + 5:
logger.info(log_str)
time_last_printed = current_time
else:
logger.debug(log_str)
start_it = stage * pop_size
# stop sampling if time budget exceeded
if time_budget is not None and not np.isnan(times[start_it]):
if times[start_it] > times[0] + time_budget:
logger.info("Time limit of %ds exceeded. Stopping MCMC at iteration %d." % (time_budget, it))
break
# print progress
# if stage > 1:
# log_str = "PMC iteration %d/%d, current log_pdf: %.6f" % (it + 1, num_iter,
# np.nan if log_pdf[it - 1] is None else log_pdf[it - 1])
# logger.debug(log_str)
range_it = range(start_it, start_it + pop_size)
for it in range_it:
# print(it)
prop_idx = it - start_it
times[it] = time.time()
# marginal sampler: make transition kernel re-compute log_pdf of current state
if recompute_log_pdf:
current_log_pdf = None
# generate proposal and acceptance probability
proposals[stage, prop_idx], prop_target_logpdf[stage, prop_idx], current_log_pdf, prop_prob_logpdf[stage, prop_idx], backw_logpdf, current_kwargs = transition_kernel.proposal(prev[prop_idx], prev_logp[prop_idx], **{})
logweights[stage, prop_idx] = prop_target_logpdf[stage, prop_idx] - prop_prob_logpdf[stage, prop_idx]
if rao_blackwell_generation:
try:
all_prop_logpdfs = np.array([transition_kernel.proposal_log_pdf(prev[it - start_it], proposals[stage, :]) for it in range_it])
prop_prob_logpdf[stage, :] = logsumexp(all_prop_logpdfs, 0)
except:
assert()
else:
# print('norb')
# assert()
np.all(logweights[stage, :] == prop_target_logpdf[stage, :] - prop_prob_logpdf[stage, :])
logweights[stage, :] = prop_target_logpdf[stage, :] - prop_prob_logpdf[stage, :]
res_idx = system_res(range(pop_size), logweights[stage, :],)
samples[range_it] = proposals[stage, res_idx]
log_pdf[range_it] = prop_target_logpdf[stage, res_idx]
prev = samples[range_it]
prev_logp = log_pdf[range_it]
# update transition kernel, might do nothing
transition_kernel.next_iteration()
if weighted_update:
transition_kernel.update(np.vstack(proposals[:stage + 1, :]), pop_size, np.hstack(logweights[:stage + 1, :]))
else:
transition_kernel.update(samples[:range_it[-1] + 1], pop_size)
if resample_at_end:
#FIXME: the last stage might not have drawn the full set of samples
all_lweights = np.hstack(logweights[:stage+1, :])
res_idx = system_res(range(it+1), all_lweights)
(samples, log_pdf) = (np.vstack(proposals[:stage+1, :])[res_idx], prop_target_logpdf[res_idx])
else:
samples, log_pdf = samples[:it], log_pdf[:it]
# recall it might be less than last iterations due to time budget
return samples, log_pdf, times[:it]
def mini_rb_pmc(transition_kernel, start, num_iter, pop_size, D, recompute_log_pdf=False, time_budget=None, ):
# PMC results
print(pop_size)
assert(num_iter % pop_size == 0)
# following not implemented yet
assert(recompute_log_pdf == False)
proposals = np.zeros((num_iter // pop_size, pop_size, D)) + np.nan
logweights = np.zeros((num_iter // pop_size, pop_size)) + np.nan
prop_target_logpdf = np.zeros((num_iter // pop_size, pop_size)) + np.nan
prop_prob_logpdf = np.zeros((num_iter // pop_size, pop_size)) + np.nan
samples = np.zeros((num_iter, D)) + np.nan
log_pdf = np.zeros(num_iter) + np.nan
# timings for output and time limit
times = np.zeros(num_iter)
# for adaptive transition kernels
avg_accept = 0.
current = np.array([start]*pop_size)
current_log_pdf = np.zeros(pop_size)+np.nan
acc_prob = np.zeros(pop_size) + np.nan
logger.info("Starting PMC using %s in D=%d dimensions" % \
(transition_kernel.__class__.__name__, D,))
it = 0
for stage in range(num_iter // pop_size):
start_it = stage * pop_size
# stop sampling if time budget exceeded
if time_budget is not None and not np.isnan(times[start_it]):
if times[start_it] > times[0] + time_budget:
logger.info("Time limit of %ds exceeded. Stopping MCMC at iteration %d." % (time_budget, it))
break
# print progress
if False and stage > 1:
log_str = "PMC iteration %d/%d, current log_pdf: %.6f, avg acceptance: %.3f" % (it + 1, num_iter,
np.nan if log_pdf[it - 1] is None else log_pdf[it - 1],
avg_accept)
logger.info(log_str)
range_it = range(start_it, start_it + pop_size)
for it in range_it:
prop_idx = it - start_it
if np.isnan(current_log_pdf[prop_idx]):
cur_lpdf = None
else:
cur_lpdf = current_log_pdf[prop_idx]
times[it] = time.time()
# marginal sampler: make transition kernel re-compute log_pdf of current state
if recompute_log_pdf:
current_log_pdf = None
# generate proposal and acceptance probability
logger.debug("Performing GRIS sample %d" % it)
proposals[stage, prop_idx], prop_target_logpdf[stage, prop_idx], current_log_pdf[prop_idx], prop_prob_logpdf[stage, prop_idx], backw_logpdf, current_kwargs = transition_kernel.proposal(current[prop_idx], cur_lpdf, **{})
#logweights[stage, prop_idx] = prop_target_logpdf[stage, prop_idx] - prop_prob_logpdf[stage, prop_idx]
#Rao-Blackwellize over all used proposals
all_prop_logpdfs = np.array([transition_kernel.proposal_log_pdf(current[it - start_it], proposals[stage, :]) for it in range_it])
prop_prob_logpdf[stage, :] = logsumexp(all_prop_logpdfs, 0)
logweights[stage, :] = prop_target_logpdf[stage, :] - prop_prob_logpdf[stage, :]
res_idx = system_res(range(pop_size), logweights[stage, :])
samples[range_it], log_pdf[range_it] = proposals[stage, res_idx], prop_prob_logpdf[stage, res_idx]
ess = compute_ess(logweights[stage, :], normalize=True)
if ess/float(pop_size) > 0.5:
current = proposals[stage, :]
current_log_pdf = prop_target_logpdf[stage, :]
else:
current = proposals[stage, res_idx]
current_log_pdf = prop_prob_logpdf[stage, res_idx]
#print(ess, float(pop_size)/ess)
transition_kernel.next_iteration()
transition_kernel.update(np.vstack(proposals[:stage+1, :]), pop_size, np.hstack(logweights[:stage+1, :]))
res_idx = system_res(range(pop_size*stage)*10, weights=logweights[:stage, :].flatten())
unw_samp = np.vstack(proposals[:pop_size*stage])
unw_logtarg = np.hstack(prop_target_logpdf[:pop_size*stage])
# recall it might be less than last iterations due to time budget
return samples[:it], log_pdf[:it], unw_samp, unw_logtarg, np.hstack(logweights[:pop_size*stage]), times[:it]
def rvs(self, num_samps):
res_idx = system_res(range(len(self.lw)), self.lw)
rval = np.zeros((num_samps, self.centers.shape[1]))
for i in range(num_samps):
rval[i], _, _, _, _, _ = self.kern.proposal(self.centers[res_idx[i]], -4, **{})
return rval
def mini_smc(num_samples, # will give size of sample in final iteration
population_size,
prior, # some distribution object that our algorithm will have no problem with
log_targ, # actual target
proposal_obj,
targ_ef_bridge=0.5,
targ_ef_stop=0.9,
ef_tolerance=0.02,
reweight=False,
across=False,
estim_evid=False,
ess=False):
# ToDO: - adaptive resampling (only use importance resampling if ESS < threshold)
# - reweight earlier iterations for actual target
# - use weighted approximation instead of approx after resampling for final iteration
"""
Sample from a geometric sequence of target distributions between prior and target,
reweighting samples from early iterations for estimating the actual target.
Uses a geometric bridge for now
= Parameters =
num_samples - size of sample in final iteration
population_size - size of particle system except for final iteration
prior - some easy target distribution, prior will likely do well
log_target - actual target
proposal_obj - object with which proposals are generated
targ_ef_bridge - target efficiency factor for bridge, i.e. what should eff/num_particles be in a bridge step
targ_ef_stop - target efficiency factor with respect to final target
ef_tolerance - efficiency factor tolerance
reweight - False (only use last iteration), True (reweight for actual target and weight iterations by ESS)
across - Resampling across iterations after reweighting? True (resample. only if reweight = True)
estim_evid - return estimate of evidence/normalizing constant of log_target
ess - Return ESS of last iteration? Defaults to False.
"""
logger.info("Starting SMC using %s" % \
(proposal_obj.get_name()))
if not reweight:
assert(not across)
log_target = lambda x: np.apply_along_axis(lambda y: np.atleast_1d(log_targ(y)), 1, x)
# will be returned
step_sizes = [proposal_obj.step_size]
acceptance_rates = []
initial_guesses = prior.rvs(population_size)
population_size = initial_guesses.shape[0]
dim = initial_guesses.shape[1]
lprior = np.empty(num_samples + population_size * 3)
lprior[:population_size] = prior.logpdf(initial_guesses)
lpost = np.empty(num_samples + population_size * 3)
lpost[:population_size] = log_target(initial_guesses).flatten()
rval = np.r_[initial_guesses, np.zeros((num_samples + population_size * 3, dim))]
def ensure(size):
old = len(lprior)
if old < size:
lprior.resize(old * 2)
lpost.resize(old * 2)
rval.resize((old * 2, rval.shape[1]))
def seq_value(br, prior_value, posterior_value):
"""
Compute logprobability from lprior and lpost according to
bridge parameter/temperature br
= Parameters =
br - bridge parameter/temperature
prior_value - value according to br = 0
posterior_value - value according to br = 1
"""
return prior_value * (1. - br) + posterior_value * br
def incr_weight(idx_from, idx_to, br_old, br_new, return_ef=False):
inc_w = (seq_value(br_new, lprior[idx_from:idx_to], lpost[idx_from:idx_to]) # use lpost[idx_beg:idx_mid] here for approximating actual target
- seq_value(br_old, lprior[idx_from:idx_to], lpost[idx_from:idx_to]))
assert(not np.any(np.isnan(rval)))
if return_ef:
norm_inc_w = inc_w - logsumexp(inc_w)
ESS = exp(2 * logsumexp(norm_inc_w) - logsumexp(2 * norm_inc_w))
EF = ESS / (idx_to - idx_from)
return (inc_w, EF)
return inc_w
def mcmc_rejuvenate(cur_sample, cur_lprior, cur_lpost, br):
"""
Make an MCMC move using proposal_obj, overwriting input (except br)
Returns the acceptance probabilities.
= Parameters =
cur_sample - the particles to be moved (will be overwritten with new state after move)
cur_lpost - the logposteriors according to final target distribution in distribution sequence (will be overwritten)
cur_lprior - the logpriors according to first target distribution in distribution sequence (will be overwritten)
br - the bridge parameter/temperature which determines how posterior and prior are mixed for current target distribution in the sequence
= Return =
Acceptance probabilites
"""
# set the target to be the intermediary distribution
save_target_logpdf = proposal_obj.target_log_pdf
proposal_obj.target_log_pdf = lambda x:seq_value(br, prior.logpdf(x), save_target_logpdf(x))
if proposal_obj.__dict__.has_key('target_grad'):
save_target_grad = proposal_obj.target_grad
proposal_obj.target_grad = lambda x:seq_value(br, prior.logpdf_grad(x), save_target_grad(x))
tmp = [proposal_obj.proposal(cur_sample[idx], seq_value(br, cur_lprior[idx], cur_lpost[idx])) for idx in range(len(cur_sample))]
# reset the target to be the actual posterior
proposal_obj.target_log_pdf = save_target_logpdf
if proposal_obj.__dict__.has_key('target_grad'):
proposal_obj.target_grad = save_target_grad
# (prop, lprob_move_forw, lprob_move_back) = [np.array(l) for l in
# zip(*tmp)]
(prop, lprob_bridge_forw, _, lprob_move_forw, lprob_move_back, current_kwargs) = [np.array(l) for l in
zip(*tmp)]
# compute log_target for proposals
lprior_forw = prior.logpdf(prop).flatten()
lpost_forw = (lprob_bridge_forw.flatten() - (1 - br) * lprior_forw.flatten()) / br
assert(np.allclose(lpost_forw, log_target(prop).flatten()))
# compute all acceptance probabilites
assert(not (np.any(np.isinf(lprob_move_forw)) or np.any(np.isnan(lprob_move_forw))))
assert(not (np.any(np.isnan(lprior_forw))))
assert(not (np.any(np.isnan(lpost_forw))))
assert(not (np.any(np.isinf(lprob_move_back)) or np.any(np.isnan(lprob_move_back))))
assert(not (np.any(np.isinf(cur_lprior)) or np.any(np.isnan(cur_lprior))))
assert(not (np.any(np.isinf(cur_lpost)) or np.any(np.isnan(cur_lpost))))
mh_ratio = (lprob_move_back + seq_value(br, lprior_forw, lpost_forw)
- lprob_move_forw - seq_value(br, cur_lprior.flatten(), cur_lpost.flatten()))
assert(mh_ratio.shape == lpost_forw.shape)
acc = exp(np.min(np.c_[np.zeros_like(mh_ratio), mh_ratio], 1))
assert(not(np.any(np.isnan(acc)) or np.any(np.isinf(acc))))
move = np.random.rand(len(acc)) < acc
assert(np.mean(acc) != 0)
cur_sample[:] = prop * np.atleast_2d(move).T + cur_sample * (1 - np.atleast_2d(move).T)
cur_lpost[:] = lpost_forw * move + cur_lpost * (1 - move)
cur_lprior[:] = prior.logpdf(cur_sample) # lprior_forw*move + cur_lprior*(1-move)
return acc
def search_bridge_param(target_ef_fact, idx_beg, idx_end, br_old, eps=ef_tolerance):
high = 1.0
low = br_old
max_eval = 9
old_EF = 0
logger.debug('Start bridge search')
for i in range(max_eval + 1):
mid = low + (high - low) / 2
(inc_w, EF) = incr_weight(idx_beg, idx_end, br_old, mid, True)
logger.debug("EF: %.4f" % EF)
d = EF - target_ef_fact
if i == max_eval or np.abs(EF - old_EF) < eps:
return (mid, inc_w)
old_EF = EF
if d < -eps:
high = mid
elif d > eps:
low = mid
else:
return (mid, inc_w)
def smc_iteration(beg, mid, end, br_old, target_ef_fact):
(br_new, inc_w) = search_bridge_param(target_ef_fact, beg, mid, br_old)
samps_idx = (np.array(system_res(range(population_size), resampled_size=end - mid, weights=inc_w))
+ beg)
rval[mid:end] = rval[samps_idx]
lpost[mid:end] = lpost[samps_idx]
lprior[mid:end] = lprior[samps_idx]
proposal_obj.set_batch(rval[samps_idx])
# pre = (rval[mid:end].copy(), lprior[mid:end].copy(), lpost[mid:end].copy())
acc = mcmc_rejuvenate(rval[mid:end], lprior[mid:end], lpost[mid:end], br_new)
mean_acc = np.mean(acc)
acceptance_rates.append(mean_acc)
proposal_obj.next_iteration()
proposal_obj.update_step_size([mean_acc])
return (br_new, inc_w, mean_acc)
br = [0.0]
evid = 0
old_EF = 0
j = 1
while True:
step_sizes += [proposal_obj.step_size]
idx_beg = (j - 1) * population_size
idx_mid = idx_beg + population_size
idx_end = idx_mid + population_size
ensure(idx_end)
(br_new, inc_w, mean_acc) = smc_iteration(idx_beg, idx_mid, idx_end, br[j - 1], targ_ef_bridge)
br.append(br_new)
evid = evid + logsumexp(inc_w) - log(inc_w.size)
norm_inc_w = inc_w - logsumexp(inc_w)
j = j + 1
ESS = exp(2 * logsumexp(norm_inc_w) - logsumexp(2 * norm_inc_w))
logger.debug("At bridge distribution #%d, ESS: %.2f, mean acc: %.4f, step_size: %.4e" %
(j, ESS, mean_acc, proposal_obj.step_size))
# test how good we are with respect to the actual distribution of interest
(inc_w_final, EF_final) = incr_weight(idx_mid, idx_end, br[j - 1], 1, True)
if (np.abs(EF_final - old_EF) < ef_tolerance or # we're not improving much
np.abs(EF_final - targ_ef_stop) < ef_tolerance): # we reached our desired efficiency factor
break
old_EF = EF_final
idx_beg = (len(br) - 1) * population_size
idx_mid = idx_beg + population_size
idx_end = idx_mid + num_samples
ensure(idx_end)
(br_new, inc_w, mean_acc) = smc_iteration(idx_beg, idx_mid, idx_end, br[-1], 1)
br.append(br_new)
evid = evid + logsumexp(inc_w) - log(inc_w.size)
norm_inc_w = inc_w - logsumexp(inc_w)
ESS = exp(2 * logsumexp(norm_inc_w) - logsumexp(2 * norm_inc_w))
logger.debug("Final approx, ESS: %.2f, mean acc: %.4f" %
(ESS, mean_acc))
if reweight == False:
(rval, lpost) = (rval[idx_mid:idx_end], lpost[idx_mid:idx_end])
else:
# reweight for actual target
power = 1. - np.repeat(br, population_size)
rew = (lpost[:idx_mid] - lprior[:idx_mid]) * power
# weight each iteration by ESS wrt actual target
rew_resh = rew.reshape((len(br), population_size))
logess = ((2 * logsumexp(rew_resh, 1) - logsumexp(2 * rew_resh, 1)))
print(exp(logess))
logess_sampsize_ratio = logess - log(population_size)
rew = rew + np.repeat(logess_sampsize_ratio.flatten(), population_size)
smp_idx = np.array(system_res(range(idx_mid), resampled_size=population_size, weights=rew))
if across:
smp_idx_acr = np.array(system_res(range(idx_end), resampled_size=idx_end, weights=np.r_[rew, np.ones(idx_end - idx_mid)]))
(rval_acr, lpost_acr) = (rval[smp_idx_acr], lpost[smp_idx_acr])
(rval, lpost) = (np.r_[rval[smp_idx], rval[idx_mid:idx_end]], np.r_[lpost[smp_idx], lpost[idx_mid:idx_end]])
all_rval = [rval, lpost, np.array(step_sizes), np.array(acceptance_rates)]
if ess:
all_rval.append(ESS)
if across:
all_rval.extend((rval_acr, lpost_acr))
if evid:
all_rval.append(evid)
return all_rval