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