def binomial_dist(total_count, probs, *, overdispersion=0.0):
"""
Returns a Beta-Binomial distribution that is an overdispersed version of a
Binomial distribution, according to a parameter ``overdispersion``,
typically set in the range 0.1 to 0.5.
This is useful for (1) fitting real data that is overdispersed relative to
a Binomial distribution, and (2) relaxing models of large populations to
improve inference. In particular the ``overdispersion`` parameter lower
bounds the relative uncertainty in stochastic models such that increasing
population leads to a limiting scale-free dynamical system with bounded
stochasticity, in contrast to Binomial-based SDEs that converge to
deterministic ODEs in the large population limit.
This parameterization satisfies the following properties:
1. Variance increases monotonically in ``overdispersion``.
2. ``overdispersion = 0`` results in a Binomial distribution.
3. ``overdispersion`` lower bounds the relative uncertainty ``std_dev /
(total_count * p * q)``, where ``probs = p = 1 - q``, and serves as an
asymptote for relative uncertainty as ``total_count → ∞``. This
contrasts the Binomial whose relative uncertainty tends to zero.
4. If ``X ~ binomial_dist(n, p, overdispersion=σ)`` then in the large
population limit ``n → ∞``, the scaled random variable ``X / n``
converges in distribution to ``LogitNormal(log(p/(1-p)), σ)``.
To achieve these properties we set ``p = probs``, ``q = 1 - p``, and::
concentration = 1 / (p * q * overdispersion**2) - 1
:param total_count: Number of Bernoulli trials.
:type total_count: int or torch.Tensor
:param probs: Event probabilities.
:type probs: float or torch.Tensor
:param overdispersion: Amount of overdispersion, in the half open interval
[0,2). Defaults to zero.
:type overdispersion: float or torch.tensor
"""
_validate_overdispersion(overdispersion)
if _is_zero(overdispersion):
if _RELAX:
return _relaxed_binomial(total_count, probs)
return dist.ExtendedBinomial(total_count, probs)
p = probs
q = 1 - p
od2 = (overdispersion + 1e-8)**2
concentration1 = 1 / (q * od2 + 1e-8) - p
concentration0 = 1 / (p * od2 + 1e-8) - q
# At this point we have
# concentration1 + concentration0 == 1 / (p + q + od2 + 1e-8) - 1
if _RELAX:
return _relaxed_beta_binomial(concentration1, concentration0,
total_count)
return dist.ExtendedBetaBinomial(concentration1, concentration0,
total_count)
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 beta_binomial_dist(concentration1,
concentration0,
total_count,
*,
overdispersion=0.0):
"""
Returns a Beta-Binomial distribution that is an overdispersed version of a
the usual Beta-Binomial distribution, according to an extra parameter
``overdispersion``, typically set in the range 0.1 to 0.5.
:param concentration1: 1st concentration parameter (alpha) for the
Beta distribution.
:type concentration1: float or torch.Tensor
:param concentration0: 2nd concentration parameter (beta) for the
Beta distribution.
:type concentration0: float or torch.Tensor
:param total_count: Number of Bernoulli trials.
:type total_count: float or torch.Tensor
:param overdispersion: Amount of overdispersion, in the half open interval
[0,2). Defaults to zero.
:type overdispersion: float or torch.tensor
"""
_validate_overdispersion(overdispersion)
if not _is_zero(overdispersion):
# Compute harmonic sum of two sources of concentration resulting in
# final concentration c = 1 / (1 / c_1 + 1 / c_2)
od2 = (overdispersion + 1e-8)**2
c_1 = concentration1 + concentration0
c_2 = c_1**2 / (concentration1 * concentration0 * od2 + 1e-8) - 1
factor = 1 + c_1 / c_2
concentration1 = concentration1 / factor
concentration0 = concentration0 / factor
if _RELAX:
return _relaxed_beta_binomial(concentration1, concentration0,
total_count)
return dist.ExtendedBetaBinomial(concentration1, concentration0,
total_count)