def _length_log_probs_with_rates(self, log_rates):
n_classes = log_rates.size(-1)
max_length = self.max_k
# max_length x n_classes
time_steps = torch.arange(max_length, device=log_rates.device).unsqueeze(-1).expand(max_length,
n_classes).float()
if max_length == 1:
return torch.FloatTensor([0, -1000]).unsqueeze(-1).expand(2, n_classes).to(log_rates.device)
# return torch.zeros(max_length, n_classes).to(log_rates.device)
poissons = Poisson(torch.exp(log_rates))
if log_rates.dim() == 2:
time_steps = time_steps.unsqueeze(1).expand(max_length, log_rates.size(0), n_classes)
return poissons.log_prob(time_steps).transpose(0, 1)
else:
assert log_rates.dim() == 1
return poissons.log_prob(time_steps)
def test_zip_0_gate(rate):
# if gate is 0 ZIP is Poisson
zip_ = ZeroInflatedPoisson(torch.zeros(1), torch.tensor(rate))
pois = Poisson(torch.tensor(rate))
s = pois.sample((20, ))
zip_prob = zip_.log_prob(s)
pois_prob = pois.log_prob(s)
assert_close(zip_prob, pois_prob, atol=1e-06)
def tile_map_prior(prior: ImagePrior, tile_map):
# Source probabilities
dist_sources = Poisson(torch.tensor(prior.mean_sources))
log_prob_no_source = dist_sources.log_prob(torch.tensor(0))
log_prob_one_source = dist_sources.log_prob(torch.tensor(1))
log_prob_source = (tile_map["n_sources"] == 0) * log_prob_no_source + (
tile_map["n_sources"] == 1) * log_prob_one_source
# Binary probabilities
galaxy_log_prob = torch.tensor(0.7).log()
star_log_prob = torch.tensor(0.3).log()
log_prob_binary = (galaxy_log_prob * tile_map["galaxy_bools"] +
star_log_prob * tile_map["star_bools"])
# Galaxy probabiltiies
gal_dist = Normal(0.0, 1.0)
galaxy_probs = gal_dist.log_prob(
tile_map["galaxy_params"]) * tile_map["galaxy_bools"]
# prob_normalized =
return log_prob_source.sum() + log_prob_binary.sum() + galaxy_probs.sum()
def k_new(self, X, Z, A, i, truncation):
'''
i: The loop calling this function is asking this function
"how many new features (k_new) should data point i draw?"
truncation: When computing the un-normalized posterior for k_new|X,Z,A, we cannot
compute the posterior for the infinite amount of values k_new could take on. So instead
we compute from 0 up to some high number, truncation, and then normalize. In practice,
the posterior probability for k_new is so low that it underflows past truncation=20.
'''
log_likelihood = torch.zeros(truncation)
log_poisson_probs = torch.zeros(truncation)
N, K = Z.size()
D = X.size()[1]
p_k_new = Pois(torch.tensor([self.alpha / N]))
cur_X_minus_ZA = X - Z @ A
for j in range(truncation):
# Compute the log likelihood of k_new equaling j
log_likelihood[j] = self.log_likelihood_given_k_new(
cur_X_minus_ZA, Z, D, i, j)
# Compute the prior probability of k_new equaling j
log_poisson_probs[j] = p_k_new.log_prob(j)
# Add new column to Z for next feature
zeros = torch.zeros(N)
Z = torch.cat((Z, torch.zeros(N, 1)), 1)
Z[i][-1] = 1
# Compute log posterior of k_new and exp/normalize
log_sample_probs = log_likelihood + log_poisson_probs
sample_probs = self.renormalize_log_probs(log_sample_probs)
# Important: we changed Z for calculating p(k_new|...)
# so we must take off the extra rows
Z = Z[:, :-truncation]
assert Z.size()[1] == K
posterior_k_new = Categorical(sample_probs)
return posterior_k_new.sample()
def log_prob(self, likelihood: Poisson, x):
# [B, V]
counts = self.make_counts(x)
# [B, V] -> [B]
return likelihood.log_prob(counts).sum(-1)
def step1_logprob(self, μs):
N = len(μs)
poisson = Poisson(self.λ)
return poisson.log_prob(N)