def model2():
c1 = pyro.param("c1",
torch.tensor(0.5),
constraint=constraints.positive)
c0 = pyro.param("c0",
torch.tensor(1.5),
constraint=constraints.positive)
pyro.sample("obs", dist.BetaBinomial(c1, c0, total_count), obs=data)
def test_beta_binomial_approx_sample(concentration1, concentration0,
total_count):
sample_shape = (10000, )
d = dist.BetaBinomial(concentration1, concentration0, total_count)
expected = d.sample(sample_shape)
with set_approx_sample_thresh(200):
actual = d.sample(sample_shape)
assert_close(expected.mean(), actual.mean(), rtol=0.1)
assert_close(expected.std(), actual.std(), rtol=0.1)
def test_beta_binomial(concentration1, concentration0, total_count):
# For small overdispersion, beta_binomial_dist is close to BetaBinomial.
d1 = dist.BetaBinomial(concentration1, concentration0, total_count)
d2 = beta_binomial_dist(concentration1, concentration0, total_count,
overdispersion=0.01)
# CRPS is equivalent to the Cramer-von Mises test.
# https://en.wikipedia.org/wiki/Cram%C3%A9r%E2%80%93von_Mises_criterion
k = torch.arange(0., total_count + 1.)
cdf1 = d1.log_prob(k).exp().cumsum(-1)
cdf2 = d2.log_prob(k).exp().cumsum(-1)
crps = (cdf1 - cdf2).pow(2).mean()
assert crps < 0.01
def test_extended_beta_binomial(tol):
with set_approx_log_prob_tol(tol):
concentration1 = torch.tensor([0.2, 1.0, 2.0, 1.0]).requires_grad_()
concentration0 = torch.tensor([0.2, 0.5, 1.0, 2.0]).requires_grad_()
total_count = torch.tensor([0.0, 1.0, 2.0, 10.0])
d1 = dist.BetaBinomial(concentration1, concentration0, total_count)
d2 = dist.ExtendedBetaBinomial(concentration1, concentration0,
total_count)
# Check on good data.
data = d1.sample((100, ))
assert_equal(d1.log_prob(data), d2.log_prob(data))
# Check on extended data.
data = torch.arange(-10.0, 20.0).unsqueeze(-1)
with pytest.raises(ValueError):
d1.log_prob(data)
log_prob = d2.log_prob(data)
valid = d1.support.check(data)
assert ((log_prob > -math.inf) == valid).all()
check_grad(log_prob, concentration1, concentration0)
# Check on shape error.
with pytest.raises(ValueError):
d2.log_prob(torch.tensor([0.0, 0.0]))
# Check on value error.
with pytest.raises(ValueError):
d2.log_prob(torch.tensor(0.5))
# Check on negative total_count.
concentration1 = torch.tensor(1.5).requires_grad_()
concentration0 = torch.tensor(1.5).requires_grad_()
total_count = torch.arange(-10, 0.0)
d = dist.ExtendedBetaBinomial(concentration1, concentration0,
total_count)
log_prob = d.log_prob(data)
assert (log_prob == -math.inf).all()
check_grad(log_prob, concentration1, concentration0)
def model(
s,
m,
y=None,
gamma_hyper=1.0,
pi0=1.0,
rho0=1.0,
epsilon0=0.01,
alpha0=1000.0,
dtype=torch.float32,
device="cpu",
):
# Cast inputs and set device
m, gamma_hyper, pi0, rho0, epsilon0, alpha0 = [
torch.tensor(v, dtype=dtype, device=device)
for v in [m, gamma_hyper, pi0, rho0, epsilon0, alpha0]
]
if y is not None:
y = torch.tensor(y)
n, g = m.shape
with pyro.plate("position", g, dim=-1):
with pyro.plate("strain", s, dim=-2):
gamma = pyro.sample(
"gamma",
dist.RelaxedBernoulli(temperature=gamma_hyper, logits=0.0),
)
# gamma.shape == (s, g)
rho_hyper = pyro.sample("rho_hyper", dist.Gamma(rho0, 1.0))
rho = pyro.sample(
"rho",
dist.RelaxedOneHotCategorical(
temperature=rho_hyper,
logits=torch.zeros(s, dtype=dtype, device=device),
),
)
epsilon_hyper = pyro.sample("epsilon_hyper", dist.Beta(1.0, 1 / epsilon0))
alpha_hyper = pyro.sample("alpha_hyper", dist.Gamma(alpha0, 1.0))
pi_hyper = pyro.sample("pi_hyper", dist.Gamma(pi0, 1.0))
with pyro.plate("sample", n, dim=-1):
pi = pyro.sample(
"pi",
dist.RelaxedOneHotCategorical(temperature=pi_hyper, probs=rho),
)
alpha = pyro.sample("alpha", dist.Gamma(alpha_hyper, 1.0)).unsqueeze(
-1
)
epsilon = pyro.sample(
"epsilon", dist.Beta(1.0, 1 / epsilon_hyper)
).unsqueeze(-1)
# pi.shape == (n, s)
# alpha.shape == epsilon.shape == (n,)
p_noerr = pyro.deterministic("p_noerr", pi @ gamma)
p = pyro.deterministic(
"p", (1 - epsilon / 2) * (p_noerr) + (epsilon / 2) * (1 - p_noerr)
)
# p.shape == (n, g)
y = pyro.sample(
"y",
dist.BetaBinomial(
concentration1=alpha * p,
concentration0=alpha * (1 - p),
total_count=m,
),
obs=y,
)
# y.shape == (n, g)
return y
