def test_neg_binomial(total_count: float, logit: float) -> None:
"""
Test to check that maximizing the likelihood recovers the parameters
"""
# generate samples
total_counts = torch.zeros((NUM_SAMPLES,)) + total_count
logits = torch.zeros((NUM_SAMPLES,)) + logit
neg_bin_distr = NegativeBinomial(total_count=total_counts, logits=logits)
samples = neg_bin_distr.sample()
init_biases = [
inv_softplus(total_count - START_TOL_MULTIPLE * TOL * total_count),
logit - START_TOL_MULTIPLE * TOL * logit,
]
total_count_hat, logit_hat = maximum_likelihood_estimate_sgd(
NegativeBinomialOutput(),
samples,
init_biases=init_biases,
num_epochs=15,
)
assert (
np.abs(total_count_hat - total_count) < TOL * total_count
), f"total_count did not match: total_count = {total_count}, total_count_hat = {total_count_hat}"
assert (
np.abs(logit_hat - logit) < TOL * logit
), f"logit did not match: logit = {logit}, logit_hat = {logit_hat}"
def test_neg_binomial(mu_alpha: Tuple[float, float]) -> None:
"""
Test to check that maximizing the likelihood recovers the parameters
"""
# test instance
mu, alpha = mu_alpha
# generate samples
mus = torch.zeros((NUM_SAMPLES, )) + mu
alphas = torch.zeros((NUM_SAMPLES, )) + alpha
neg_bin_distr = NegativeBinomial(total_count=1.0 / alphas,
probs=mus * alphas / (1.0 + mus * alphas))
samples = neg_bin_distr.sample()
init_biases = [
inv_softplus(mu - START_TOL_MULTIPLE * TOL * mu),
inv_softplus(alpha + START_TOL_MULTIPLE * TOL * alpha),
]
mu_hat, alpha_hat = maximum_likelihood_estimate_sgd(
NegativeBinomialOutput(),
samples,
init_biases=init_biases,
num_epochs=15,
)
assert (np.abs(mu_hat - mu) <
TOL * mu), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
assert (np.abs(alpha_hat - alpha) < TOL * alpha
), f"alpha did not match: alpha = {alpha}, alpha_hat = {alpha_hat}"
def test_zinb_0_gate(total_count, probs):
# if gate is 0 ZINB is NegativeBinomial
zinb_ = ZeroInflatedNegativeBinomial(torch.zeros(1),
total_count=torch.tensor(total_count),
probs=torch.tensor(probs))
neg_bin = NegativeBinomial(torch.tensor(total_count),
probs=torch.tensor(probs))
s = neg_bin.sample((20, ))
zinb_prob = zinb_.log_prob(s)
neg_bin_prob = neg_bin.log_prob(s)
assert_close(zinb_prob, neg_bin_prob, atol=1e-06)