add logcdf

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -489,6 +489,8 @@ def dist(cls, alpha, beta, *args, **kwargs):
         return super().dist([alpha, beta], *args, **kwargs)
 
     def logp(value, alpha, beta):
+        """From Expression (5) on p.6 of Fader & Hardie (2007)"""
+
         logp = betaln(alpha + 1, beta + value - 1) - betaln(alpha, beta)
 
         logp = pt.switch(
@@ -510,6 +512,18 @@ def logp(value, alpha, beta):
             msg="alpha > 0, beta > 0",
         )
 
+    def logcdf(value, alpha, beta):
+        """Adapted from Expression (6) on p.6 of Fader & Hardie (2007)"""
+        # survival function from paper
+        logS = (
+            pt.gammaln(beta + value)
+            - pt.gammaln(beta)
+            + pt.gammaln(alpha + beta)
+            - pt.gammaln(alpha + beta + value)
+        )
+        # log(1-exp())
+        return pt.log1mexp(logS)
+
     def support_point(rv, size, alpha, beta):
         """Calculate a reasonable starting point for sampling.
 

tests/distributions/test_discrete.py

@@ -24,9 +24,11 @@
     BaseTestDistributionRandom,
     Domain,
     I,
+    Nat,
     Rplus,
     assert_support_point_is_expected,
     check_logp,
+    check_selfconsistency_discrete_logcdf,
     discrete_random_tester,
 )
 from pytensor import config
@@ -241,8 +243,8 @@ def test_random_basic_properties(self):
 
         def test_random_edge_cases(self):
             """Test with very small and large beta and alpha values"""
-            beta_vals = [20.0, 0.07, 18.0, 0.05]
-            alpha_vals = [20.0, 14.0, 0.06, 0.05]
+            beta_vals = [20.0, 0.08, 18.0, 0.06]
+            alpha_vals = [20.0, 14.0, 0.07, 0.06]
 
             for beta in beta_vals:
                 for alpha in alpha_vals:
@@ -299,6 +301,13 @@ def test_logp_matches_paper(self):
         logp_fn = pytensor.function([value, alpha_, beta_], logp)
         np.testing.assert_allclose(logp_fn(t_vec, alpha, beta), expected, rtol=1e-2)
 
+    def test_logcdf(self):
+        check_selfconsistency_discrete_logcdf(
+            distribution=ShiftedBetaGeometric,
+            domain=Nat,
+            paramdomains={"alpha": Rplus, "beta": Rplus},
+        )
+
     @pytest.mark.parametrize(
         "alpha, beta, size, expected_shape",
         [