Allow negative total_count in Extended distributions (#2439)

pyro-ppl · Apr 24, 2020 · e284f64 · e284f64
1 parent 560aafb
commit e284f64
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 8 deletions.
diff --git a/examples/sir_hmc.py b/examples/sir_hmc.py
@@ -174,7 +174,7 @@ def reparameterized_discrete_model(args, data):
                     dist.ExtendedBinomial(I_prev, prob_i),
                     obs=I2R)
         pyro.sample("obs_{}".format(t),
-                    dist.ExtendedBinomial(S2I.clamp(min=0), rho),
+                    dist.ExtendedBinomial(S2I, rho),
                     obs=datum)
 
 
@@ -301,7 +301,7 @@ def continuous_model(args, data):
                     dist.ExtendedBinomial(I_prev, prob_i),
                     obs=I2R)
         pyro.sample("obs_{}".format(t),
-                    dist.ExtendedBinomial(S2I.clamp(min=0), rho),
+                    dist.ExtendedBinomial(S2I, rho),
                     obs=datum)
 
 
@@ -408,7 +408,7 @@ def vectorized_model(args, data):
     # Compute probability factors.
     S2I_logp = dist.ExtendedBinomial(S_prev, -(rate_s * I_prev).expm1()).log_prob(S2I)
     I2R_logp = dist.ExtendedBinomial(I_prev, prob_i).log_prob(I2R)
-    obs_logp = dist.ExtendedBinomial(S2I.clamp(min=0), rho).log_prob(data)
+    obs_logp = dist.ExtendedBinomial(S2I, rho).log_prob(data)
 
     # Manually perform variable elimination.
     logp = S_logp + (I_logp + obs_logp) + S2I_logp + I2R_logp

diff --git a/pyro/contrib/epidemiology/sir.py b/pyro/contrib/epidemiology/sir.py
@@ -98,5 +98,5 @@ def transition_bwd(self, params, prev, curr, t):
                     dist.ExtendedBinomial(prev["I"], prob_i),
                     obs=I2R)
         pyro.sample("obs_{}".format(t),
-                    dist.ExtendedBinomial(S2I.clamp(min=0), rho),
+                    dist.ExtendedBinomial(S2I, rho),
                     obs=self.data[t])
diff --git a/pyro/distributions/extended.py b/pyro/distributions/extended.py
@@ -12,9 +12,13 @@
 class ExtendedBinomial(Binomial):
     """
     EXPERIMENTAL :class:`~pyro.distributions.Binomial` distribution extended to
-    have logical support the entire integers. Numerical support is still the
-    integer interval ``[0, total_count]``
+    have logical support the entire integers and to allow arbitrary integer
+    ``total_count``. Numerical support is still the integer interval ``[0,
+    total_count]``.
     """
+    arg_constraints = {"total_count": constraints.integer,
+                       "probs": constraints.unit_interval,
+                       "logits": constraints.real}
     support = constraints.integer
 
     def log_prob(self, value):
@@ -26,9 +30,13 @@ def log_prob(self, value):
 class ExtendedBetaBinomial(BetaBinomial):
     """
     EXPERIMENTAL :class:`~pyro.distributions.BetaBinomial` distribution
-    extended to have logical support the entire integers. Numerical support is
-    still the integer interval ``[0, total_count]``
+    extended to have logical support the entire integers and to allow arbitrary
+    integer ``total_count``. Numerical support is still the integer interval
+    ``[0, total_count]``.
     """
+    arg_constraints = {"concentration1": constraints.positive,
+                       "concentration0": constraints.positive,
+                       "total_count": constraints.integer}
     support = constraints.integer
 
     def log_prob(self, value):

diff --git a/tests/distributions/test_extended.py b/tests/distributions/test_extended.py
@@ -37,6 +37,11 @@ def test_extended_binomial():
     with pytest.raises(ValueError):
         d2.log_prob(torch.tensor(0.5))
 
+    # Check on negative total_count.
+    total_count = torch.arange(-10, 0.)
+    d = dist.ExtendedBinomial(total_count, 0.5)
+    assert (d.log_prob(data) == -math.inf).all()
+
 
 def test_extended_beta_binomial():
     concentration1 = torch.tensor([1.0, 2.0, 1.0])
@@ -65,3 +70,8 @@ def test_extended_beta_binomial():
     # Check on value error.
     with pytest.raises(ValueError):
         d2.log_prob(torch.tensor(0.5))
+
+    # Check on negative total_count.
+    total_count = torch.arange(-10, 0.)
+    d = dist.ExtendedBetaBinomial(1.5, 1.5, total_count)
+    assert (d.log_prob(data) == -math.inf).all()