small lam value tests

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -426,11 +426,17 @@ def rng_fn(cls, rng, r, alpha, size):
         lam = rng.gamma(shape=r, scale=1/alpha, size=size)
 
         def sim_data(lam):
-            # TODO: To support time-varying covariates, covariate vector may need to be added
-            p = 1 - np.exp(-lam)
-
+            # Handle numerical stability for very small lambda values
+            p = np.where(
+                lam < 0.001,
+                lam,  # For small lambda, p ≈ lambda
+                1 - np.exp(-lam)  # Standard formula for larger lambda
+            )
+            
+            # Ensure p is in valid range for geometric distribution
+            p = np.clip(p, 1e-5, 1.)
+            
             t = rng.geometric(p)
-
             return np.array([t])
 
         for index in np.ndindex(*size):

tests/distributions/test_discrete.py

@@ -222,13 +222,29 @@ class TestRandomVariable(BaseTestDistributionRandom):
         ]
 
         def test_random_basic_properties(self):
+            # Test standard parameter values
             discrete_random_tester(
                 dist=self.pymc_dist,
-                paramdomains={"r": Rplus, "alpha": Rplus},
+                paramdomains={
+                    "r": Domain([0.5, 1.0, 2.0], edges=(None, None)),  # Standard values
+                    "alpha": Domain([0.5, 1.0, 2.0], edges=(None, None)),  # Standard values
+                },
                 ref_rand=lambda r, alpha, size: np.random.geometric(
                     1 - np.exp(-np.random.gamma(r, 1/alpha, size=size)), size=size
                 ),
             )
+            
+            # Test small parameter values that could generate small lambda values
+            discrete_random_tester(
+                dist=self.pymc_dist,
+                paramdomains={
+                    "r": Domain([0.01, 0.1], edges=(None, None)),  # Small r values
+                    "alpha": Domain([10.0, 100.0], edges=(None, None)),  # Large alpha values
+                },
+                ref_rand=lambda r, alpha, size: np.random.geometric(
+                    np.clip(np.random.gamma(r, 1/alpha, size=size), 1e-5, 1.0), size=size
+                ),
+            )
 
         @pytest.mark.parametrize("r,alpha", [
             (0.5, 1.0),