def propose_assignments2(model, data, state, update=False):
    N, K = state.Z.shape
    state = state.copy()
    cache = ibp.IBPCache.from_state(model, data, state, np.zeros(N, dtype=bool))

    proposal_prob = 0.

    for i in range(N):
        obs = data.mask[i, :]
        x = state.X[i, :]
        
        evidence = np.zeros(4)
        prior_odds = np.zeros(4)
        for c, (z1, z2) in enumerate(CHOICES):
            z = np.array([z1, z2])
            mu = cache.fpost.predictive_mu(z)
            ssq = cache.fpost.predictive_ssq(z) + state.sigma_sq_n
            evidence[c] = ibp.gauss_loglik_vec_C2(x[obs], mu[obs], ssq)

            for k in [0, 1]:
                if cache.counts[k] > 0:
                    prior_odds[c] += np.log(cache.counts[k]) - np.log(cache.num_included - cache.counts[k] + 1)
                else:
                    prior_odds[c] += np.log(model.alpha) - np.log(cache.num_included + 1)

        odds = evidence + prior_odds
        dist = distributions.MultinomialDistribution.from_odds(odds)
        if update:
            state.Z[i, :] = CHOICES[dist.sample().argmax()]
        proposal_prob += dist.loglik(CHOICES.index(tuple(state.Z[i, :])))
        cache.add(i, state.Z[i, :], state.X[i, :])

        assert np.isfinite(proposal_prob)

    return state, proposal_prob
def next_assignment_proposal(model, data, state, cache, Sigma_info, i, k):
    assert not cache.rows_included[i]
    x = state.X[i, :]

    evidence = np.zeros(2)
    for assignment in [0, 1]:
        mu = Sigma_info.mu_for(k, assignment)
        ssq = Sigma_info.sigma_sq_for(k, assignment) + state.sigma_sq_n
        evidence[assignment] = ibp.gauss_loglik_vec_C2(x, mu, ssq)
    data_odds = evidence[1] - evidence[0]

    if cache.counts[k] > 0:
        prior_odds = np.log(cache.counts[k]) - np.log(cache.num_included - cache.counts[k] + 1)
    else:
        #prior_odds = poisson(1, 0.5 * model.alpha / (i+1)) - poisson(0, 0.5 * model.alpha / (i+1))
        prior_odds = np.log(model.alpha) - np.log(cache.num_included + 1)

    return distributions.BernoulliDistribution.from_odds(data_odds + prior_odds)
def next_assignment_proposal(model, data, state, cache, Sigma_info, i, k):
    assert not cache.rows_included[i]
    x = state.X[i, :]

    evidence = np.zeros(2)
    for assignment in [0, 1]:
        mu = Sigma_info.mu_for(k, assignment)
        ssq = Sigma_info.sigma_sq_for(k, assignment) + state.sigma_sq_n
        evidence[assignment] = ibp.gauss_loglik_vec_C2(x, mu, ssq)
    data_odds = evidence[1] - evidence[0]

    if cache.counts[k] > 0:
        prior_odds = np.log(
            cache.counts[k]) - np.log(cache.num_included - cache.counts[k] + 1)
    else:
        #prior_odds = poisson(1, 0.5 * model.alpha / (i+1)) - poisson(0, 0.5 * model.alpha / (i+1))
        prior_odds = np.log(model.alpha) - np.log(cache.num_included + 1)

    return distributions.BernoulliDistribution.from_odds(data_odds +
                                                         prior_odds)
def propose_assignments2(model, data, state, update=False):
    N, K = state.Z.shape
    state = state.copy()
    cache = ibp.IBPCache.from_state(model, data, state, np.zeros(N,
                                                                 dtype=bool))

    proposal_prob = 0.

    for i in range(N):
        obs = data.mask[i, :]
        x = state.X[i, :]

        evidence = np.zeros(4)
        prior_odds = np.zeros(4)
        for c, (z1, z2) in enumerate(CHOICES):
            z = np.array([z1, z2])
            mu = cache.fpost.predictive_mu(z)
            ssq = cache.fpost.predictive_ssq(z) + state.sigma_sq_n
            evidence[c] = ibp.gauss_loglik_vec_C2(x[obs], mu[obs], ssq)

            for k in [0, 1]:
                if cache.counts[k] > 0:
                    prior_odds[c] += np.log(
                        cache.counts[k]) - np.log(cache.num_included -
                                                  cache.counts[k] + 1)
                else:
                    prior_odds[c] += np.log(
                        model.alpha) - np.log(cache.num_included + 1)

        odds = evidence + prior_odds
        dist = distributions.MultinomialDistribution.from_odds(odds)
        if update:
            state.Z[i, :] = CHOICES[dist.sample().argmax()]
        proposal_prob += dist.loglik(CHOICES.index(tuple(state.Z[i, :])))
        cache.add(i, state.Z[i, :], state.X[i, :])

        assert np.isfinite(proposal_prob)

    return state, proposal_prob