WIP time covars required param

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -412,9 +412,7 @@ 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?
+        exp_time_covar = np.exp(time_covariate_vector).sum(axis=0)
 
         # Determine output size
         if size is None:
@@ -427,7 +425,7 @@ def rng_fn(cls, rng, r, alpha, time_covariate_vector, size):
 
         lam = rng.gamma(shape=r, scale=1 / alpha, size=size)
 
-        lam_covar = lam * exp_time_covar  # TODO: test summing over this in a notebook as well?
+        lam_covar = lam * exp_time_covar
 
         p = 1 - np.exp(-lam_covar)
         samples = rng.geometric(p)
@@ -483,8 +481,8 @@ class GrassiaIIGeometric(Discrete):
         Shape parameter (r > 0).
     alpha : tensor_like of float
         Scale parameter (alpha > 0).
-    time_covariate_vector : tensor_like of float, optional
-        Optional vector containing dot products of time-varying covariates and coefficients.
+    time_covariate_vector : tensor_like of float
+        Vector containing dot products of time-varying covariates and coefficients.
 
     References
     ----------
@@ -496,11 +494,9 @@ class GrassiaIIGeometric(Discrete):
     rv_op = g2g
 
     @classmethod
-    def dist(cls, r, alpha, time_covariate_vector=None, *args, **kwargs):
+    def dist(cls, r, alpha, time_covariate_vector, *args, **kwargs):
         r = pt.as_tensor_variable(r)
         alpha = pt.as_tensor_variable(alpha)
-        if time_covariate_vector is None:
-            time_covariate_vector = pt.constant(0.0)
         time_covariate_vector = pt.as_tensor_variable(time_covariate_vector)
         return super().dist([r, alpha, time_covariate_vector], *args, **kwargs)
 
@@ -537,7 +533,6 @@ def logcdf(value, r, alpha, time_covariate_vector):
             logcdf,
             r > 0,
             alpha > 0,
-            time_covariate_vector >= 0,
             msg="r > 0, alpha > 0",
         )
 
@@ -575,16 +570,12 @@ def support_point(rv, size, r, alpha, time_covariate_vector):
         return mean
 
 
-# TODO: can this be moved into logp? Indexing not required for logcdf
 def C_t(t: pt.TensorVariable, time_covariate_vector: pt.TensorVariable) -> pt.TensorVariable:
     """Utility for processing time-varying covariates in GrassiaIIGeometric distribution."""
-    if time_covariate_vector.ndim == 0:
-        return t
-    else:
-        # Ensure t is a valid index
-        t_idx = pt.maximum(0, t - 1)  # Convert to 0-based indexing
-        # If t_idx exceeds length of time_covariate_vector, use last value
-        max_idx = pt.shape(time_covariate_vector)[0] - 1
-        safe_idx = pt.minimum(t_idx, max_idx)
-        covariate_value = time_covariate_vector[..., safe_idx]
-        return pt.exp(covariate_value).sum()
+    # Ensure t is a valid index
+    t_idx = pt.maximum(0, t - 1)  # Convert to 0-based indexing
+    # If t_idx exceeds length of time_covariate_vector, use last value
+    max_idx = pt.shape(time_covariate_vector)[0] - 1
+    safe_idx = pt.minimum(t_idx, max_idx)
+    covariate_value = time_covariate_vector[..., safe_idx]
+    return pt.exp(covariate_value).sum(axis=0)

tests/distributions/test_discrete.py

@@ -217,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": [1.0, 2.0, 3.0]}
+        expected_rv_op_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": [1.0, 2.0, 3.0]}
         tests_to_run = [
             "check_pymc_params_match_rv_op",
             "check_rv_size",
@@ -250,7 +250,7 @@ def test_random_edge_cases(self):
             # Test with small r and large alpha values
             r_vals = [0.1, 0.5]
             alpha_vals = [5.0, 10.0]
-            time_cov_vals = [[0.0], [1.0]]
+            time_cov_vals = [[0.0, 1.0, 2.0], [5.0, 10.0, 15.0]]
 
             for r in r_vals:
                 for alpha in alpha_vals:
@@ -266,34 +266,13 @@ def test_random_edge_cases(self):
                         assert np.mean(draws) > 0
                         assert np.var(draws) > 0
 
-        def test_random_none_covariates(self):
-            """Test random sampling with None time_covariate_vector"""
-            r_vals = [0.5, 1.0, 2.0]
-            alpha_vals = [0.5, 1.0, 2.0]
-
-            for r in r_vals:
-                for alpha in alpha_vals:
-                    dist = self.pymc_dist.dist(
-                        r=r,
-                        alpha=alpha,
-                        time_covariate_vector=[0.0],  # Changed from None to avoid zip issues
-                        size=1000,
-                    )
-                    draws = dist.eval()
-
-                    # Check basic properties
-                    assert np.all(draws > 0)
-                    assert np.all(draws.astype(int) == draws)
-                    assert np.mean(draws) > 0
-                    assert np.var(draws) > 0
-
         @pytest.mark.parametrize(
             "r,alpha,time_covariate_vector",
             [
-                (0.5, 1.0, None),
+                (0.5, 1.0, [[0.0], [0.0], [0.0]]),
                 (1.0, 2.0, [1.0]),
                 (2.0, 0.5, [[1.0], [2.0]]),
-                ([5.0], [1.0], None),
+                ([5.0], [1.0], [0.0, 0.0, 0.0]),
             ],
         )
         def test_random_moments(self, r, alpha, time_covariate_vector):
@@ -311,7 +290,7 @@ 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, 1.0])
+        time_covariate_vector = pt.as_tensor_variable([[0.5, 1.0, 1.5], [0.0, 0.0, 0.0]])
         value = pt.vector("value", dtype="int64")
 
         # Create the distribution with the PyTensor variables
@@ -335,16 +314,16 @@ def test_logp(self):
     def test_logcdf(self):
         # test logcdf matches log sums across parameter values
         check_selfconsistency_discrete_logcdf(
-            GrassiaIIGeometric, NatBig, {"r": Rplus, "alpha": Rplus, "time_covariate_vector": Rplus}
+            GrassiaIIGeometric, NatBig, {"r": Rplus, "alpha": Rplus, "time_covariate_vector": I}
         )
 
     @pytest.mark.parametrize(
         "r, alpha, time_covariate_vector, size, expected_shape",
         [
-            (1.0, 1.0, None, None, ()),  # Scalar output with no covariates
+            (1.0, 1.0, [0.0, 0.0, 0.0], None, ()),  # Scalar output
             ([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.0, 4.0]], None, (2,)),  # Vector output from time covariates
             (1.0, 1.0, [1.0, 2.0], (3, 2), (3, 2)),  # Explicit size with time covariates
         ],
     )