WIP rng_fn testing

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -411,24 +411,27 @@ class GrassiaIIGeometricRV(RandomVariable):
 
     @classmethod
     def rng_fn(cls, rng, r, alpha, time_covariate_vector, size):
+        # Aggregate time covariates for each sample before broadcasting
+        exp_time_covar = np.exp(
+            time_covariate_vector.sum(axis=0)
+        )  # TODO: try np.exp(time_covariate_vector).sum(axis=0) instead?
+
         # Determine output size
         if size is None:
-            size = np.broadcast_shapes(r.shape, alpha.shape, time_covariate_vector.shape)
+            size = np.broadcast_shapes(r.shape, alpha.shape, exp_time_covar.shape)
 
         # Broadcast parameters to output size
         r = np.broadcast_to(r, size)
         alpha = np.broadcast_to(alpha, size)
-        time_covariate_vector = np.broadcast_to(time_covariate_vector, size)
+        exp_time_covar = np.broadcast_to(exp_time_covar, size)
 
         lam = rng.gamma(shape=r, scale=1 / alpha, size=size)
 
-        # Aggregate time covariates for each sample
-        exp_time_covar = np.exp(
-            time_covariate_vector.sum(axis=0)
-        )  # TODO: try np.exp(time_covariate_vector).sum(axis=0) instead?
-        lam_covar = lam * exp_time_covar
+        lam_covar = lam * exp_time_covar  # TODO: test summing over this in a notebook as well?
 
-        samples = np.ceil(np.log(1 - rng.uniform(size=size)) / (-lam_covar))
+        p = 1 - np.exp(-lam_covar)
+        samples = rng.geometric(p)
+        # samples = np.ceil(np.log(1 - rng.uniform(size=size)) / (-lam_covar))
 
         return samples
 
@@ -560,7 +563,7 @@ def support_point(rv, size, r, alpha, time_covariate_vector):
         )
 
         # Apply time covariates if provided
-        mean = mean * pt.exp(time_covariate_vector)
+        mean = mean * pt.exp(time_covariate_vector.sum(axis=0))
 
         # Round up to nearest integer and ensure >= 1
         mean = pt.maximum(pt.ceil(mean), 1.0)

tests/distributions/test_discrete.py

@@ -24,6 +24,7 @@
     BaseTestDistributionRandom,
     Domain,
     I,
+    NatBig,
     Rplus,
     assert_support_point_is_expected,
     check_logp,
@@ -216,8 +217,8 @@ def test_logp(self):
 class TestGrassiaIIGeometric:
     class TestRandomVariable(BaseTestDistributionRandom):
         pymc_dist = GrassiaIIGeometric
-        pymc_dist_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": 0.0}
-        expected_rv_op_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": 0.0}
+        pymc_dist_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": [0.0]}
+        expected_rv_op_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": [0.0]}
         tests_to_run = [
             "check_pymc_params_match_rv_op",
             "check_rv_size",
@@ -228,7 +229,7 @@ def test_random_basic_properties(self):
             # Test with standard parameter values
             r_vals = [0.5, 1.0, 2.0]
             alpha_vals = [0.5, 1.0, 2.0]
-            time_cov_vals = [-1.0, 1.0, 2.0]
+            time_cov_vals = [[0.0], [1.0], [2.0]]
 
             for r in r_vals:
                 for alpha in alpha_vals:
@@ -275,7 +276,7 @@ def test_random_none_covariates(self):
                     dist = self.pymc_dist.dist(
                         r=r,
                         alpha=alpha,
-                        time_covariate_vector=0.0,  # Changed from None to avoid zip issues
+                        time_covariate_vector=[0.0],  # Changed from None to avoid zip issues
                         size=1000,
                     )
                     draws = dist.eval()
@@ -289,10 +290,10 @@ def test_random_none_covariates(self):
         @pytest.mark.parametrize(
             "r,alpha,time_covariate_vector",
             [
-                (0.5, 1.0, 0.0),
-                (1.0, 2.0, 1.0),
-                (2.0, 0.5, -1.0),
-                (5.0, 1.0, 0.0),  # Changed from None to avoid zip issues
+                (0.5, 1.0, None),
+                (1.0, 2.0, [1.0]),
+                (2.0, 0.5, [[1.0], [2.0]]),
+                ([5.0], [1.0], None),
             ],
         )
         def test_random_moments(self, r, alpha, time_covariate_vector):
@@ -306,11 +307,11 @@ def test_random_moments(self, r, alpha, time_covariate_vector):
             assert np.mean(draws) > 0
             assert np.var(draws) > 0
 
-    def test_logp_basic(self):
+    def test_logp(self):
         # Create PyTensor variables with explicit values to ensure proper initialization
         r = pt.as_tensor_variable(1.0)
         alpha = pt.as_tensor_variable(2.0)
-        time_covariate_vector = pt.as_tensor_variable(0.5)
+        time_covariate_vector = pt.as_tensor_variable([0.5, 1.0])
         value = pt.vector("value", dtype="int64")
 
         # Create the distribution with the PyTensor variables
@@ -334,17 +335,17 @@ def test_logp_basic(self):
     def test_logcdf(self):
         # test logcdf matches log sums across parameter values
         check_selfconsistency_discrete_logcdf(
-            GrassiaIIGeometric, I, {"r": Rplus, "alpha": Rplus, "time_covariate_vector": I}
+            GrassiaIIGeometric, NatBig, {"r": Rplus, "alpha": Rplus, "time_covariate_vector": Rplus}
         )
 
     @pytest.mark.parametrize(
         "r, alpha, time_covariate_vector, size, expected_shape",
         [
-            (1.0, 1.0, 0.0, None, ()),  # Scalar output with no covariates (0.0 instead of None)
-            ([1.0, 2.0], 1.0, 0.0, None, (2,)),  # Vector output from r
-            (1.0, [1.0, 2.0], 0.0, None, (2,)),  # Vector output from alpha
-            (1.0, 1.0, [1.0, 2.0], None, (2,)),  # Vector output from time covariates
-            (1.0, 1.0, 1.0, (3, 2), (3, 2)),  # Explicit size with scalar time covariates
+            (1.0, 1.0, None, None, ()),  # Scalar output with no covariates
+            ([1.0, 2.0], 1.0, [0.0], None, (2,)),  # Vector output from r
+            (1.0, [1.0, 2.0], [0.0], None, (2,)),  # Vector output from alpha
+            (1.0, 1.0, [1.0, 2.0], None, ()),  # Vector output from time covariates
+            (1.0, 1.0, [1.0, 2.0], (3, 2), (3, 2)),  # Explicit size with time covariates
         ],
     )
     def test_support_point(self, r, alpha, time_covariate_vector, size, expected_shape):