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
Example #2
0
    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)
Example #3
0
    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]
Example #6
0
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
Example #8
0
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
Example #11
0
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