From 269dd75326b3a768ca0d693fc0ae543416bf2317 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Sat, 29 Mar 2025 10:36:58 -0600
Subject: [PATCH 01/13] dist and rv init commit

---
 pymc_extras/distributions/__init__.py |   1 +
 pymc_extras/distributions/discrete.py | 100 ++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)

diff --git a/pymc_extras/distributions/__init__.py b/pymc_extras/distributions/__init__.py
index 046e3b601..404298b55 100644
--- a/pymc_extras/distributions/__init__.py
+++ b/pymc_extras/distributions/__init__.py
@@ -37,4 +37,5 @@
     "R2D2M2CP",
     "Skellam",
     "histogram_approximation",
+    "GrassiaIIGeometric",
 ]
diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 0bfe8e7c4..bd73f2dcf 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -397,3 +397,103 @@ def dist(cls, mu1, mu2, **kwargs):
             class_name="Skellam",
             **kwargs,
         )
+
+
+class GrassiaIIGeometricRV(RandomVariable):
+    name = "g2g"
+    signature = "(),()->()"
+
+    dtype = "int64"
+    _print_name = ("GrassiaIIGeometric", "\\operatorname{GrassiaIIGeometric}")
+
+    def __call__(self, r, alpha, size=None, **kwargs):
+        return super().__call__(r, alpha, size=size, **kwargs)
+
+    @classmethod
+    def rng_fn(cls, rng, r, alpha, size):
+        if size is None:
+            size = np.broadcast_shapes(r.shape, alpha.shape)
+
+        r = np.asarray(r)
+        alpha = np.asarray(alpha)
+
+        r = np.broadcast_to(r, size)
+        alpha = np.broadcast_to(alpha, size)
+
+        output = np.zeros(shape=size + (1,))  # noqa:RUF005
+
+        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)
+
+            t = rng.geometric(p)
+
+            return np.array([t])
+
+        for index in np.ndindex(*size):
+            output[index] = sim_data(lam[index])
+
+        return output
+
+
+g2g = GrassiaIIGeometricRV()
+
+
+class GrassiaIIGeometric(UnitContinuous):
+    r"""Grassia(II)-Geometric distribution for a discrete-time, contractual customer population.
+
+    Described by Hardie and Fader in [1]_, this distribution is comprised by the following PMF and survival functions:
+
+    .. math::
+        \mathbb{P}T=t|r,\alpha,\beta;Z(t)) = (\frac{\alpha}{\alpha+C(t-1)})^{r} - (\frac{\alpha}{\alpha+C(t)})^{r}  \\
+        \begin{align}
+        \mathbb{S}(t|r,\alpha,\beta;Z(t)) = (\frac{\alpha}{\alpha+C(t)})^{r} \\
+        \end{align}
+    ========  ===============================================
+    Support   :math:`0 < t <= T` for :math: `t = 1, 2, \dots, T`
+    ========  ===============================================
+
+    Parameters
+    ----------
+    r : tensor_like of float
+        Shape parameter of Gamma distribution describing customer heterogeneity. (r > 0)
+    alpha : tensor_like of float
+        Scale parameter of Gamma distribution describing customer heterogeneity. (alpha > 0)
+
+    References
+    ----------
+    .. [1] Fader, Peter & G. S. Hardie, Bruce (2020).
+       "Incorporating Time-Varying Covariates in a Simple Mixture Model for Discrete-Time Duration Data."
+       https://www.brucehardie.com/notes/037/time-varying_covariates_in_BG.pdf
+    """
+
+    rv_op = g2g
+
+    @classmethod
+    def dist(cls, r, alpha, *args, **kwargs):
+        r = pt.as_tensor_variable(r)
+        alpha = pt.as_tensor_variable(alpha)
+        return super().dist([r, alpha], *args, **kwargs)
+
+    def logp(value, r, alpha):
+        logp = -r * (pt.log(alpha + value - 1) + pt.log(alpha + value))
+
+        return check_parameters(
+            logp,
+            r > 0,
+            alpha > 0,
+            msg="s > 0, alpha > 0",
+        )
+
+    def logcdf(value, r, alpha):
+        # TODO: Math may not be correct here
+        logcdf = r * (pt.log(value) - pt.log(alpha + value))
+
+        return check_parameters(
+            logcdf,
+            r > 0,
+            alpha > 0,
+            msg="s > 0, alpha > 0",
+        )
\ No newline at end of file

From d734c68393ec9dabbed44fa05537f8f5e0c0e85d Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Tue, 15 Apr 2025 10:20:31 -0600
Subject: [PATCH 02/13] docstrings

---
 pymc_extras/distributions/discrete.py | 34 +++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 5 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index bd73f2dcf..0331aeba9 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -442,25 +442,49 @@ def sim_data(lam):
 
 
 class GrassiaIIGeometric(UnitContinuous):
-    r"""Grassia(II)-Geometric distribution for a discrete-time, contractual customer population.
+    r"""Grassia(II)-Geometric distribution.
 
-    Described by Hardie and Fader in [1]_, this distribution is comprised by the following PMF and survival functions:
+    This distribution is a flexible alternative to the Geometric distribution for the
+    number of trials until a discrete event, and can be easily extended to support both static
+    and time-varying covariates.
+
+    Hardie and Fader describe this distribution with the following PMF and survival functions in [1]_:
 
     .. math::
         \mathbb{P}T=t|r,\alpha,\beta;Z(t)) = (\frac{\alpha}{\alpha+C(t-1)})^{r} - (\frac{\alpha}{\alpha+C(t)})^{r}  \\
         \begin{align}
         \mathbb{S}(t|r,\alpha,\beta;Z(t)) = (\frac{\alpha}{\alpha+C(t)})^{r} \\
         \end{align}
+
+    .. plot::
+        :context: close-figs
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import scipy.stats as st
+        import arviz as az
+        plt.style.use('arviz-darkgrid')
+        t = np.arange(1, 11)
+        alpha_vals = [1., 1., 2., 2.]
+        r_vals = [.1, .25, .5, 1.]
+        for alpha, r in zip(alpha_vals, r_vals):
+            pmf = (alpha/(alpha + t - 1))**r - (alpha/(alpha+t))**r
+            plt.plot(t, pmf, '-o', label=r'$\alpha$ = {}, $r$ = {}'.format(alpha, r))
+        plt.xlabel('t', fontsize=12)
+        plt.ylabel('p(t)', fontsize=12)
+        plt.legend(loc=1)
+        plt.show()
+
     ========  ===============================================
-    Support   :math:`0 < t <= T` for :math: `t = 1, 2, \dots, T`
+    Support   :math:`t \in \mathbb{N}_{>0}`
     ========  ===============================================
 
     Parameters
     ----------
     r : tensor_like of float
-        Shape parameter of Gamma distribution describing customer heterogeneity. (r > 0)
+        Shape parameter (r > 0).
     alpha : tensor_like of float
-        Scale parameter of Gamma distribution describing customer heterogeneity. (alpha > 0)
+        Scale parameter (alpha > 0).
 
     References
     ----------

From 93c4a6029a2ae5f42e5177cd524a52a4315a5a52 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Sun, 20 Apr 2025 14:09:55 -0500
Subject: [PATCH 03/13] unit tests

---
 pymc_extras/distributions/__init__.py |   1 +
 pymc_extras/distributions/discrete.py |  28 ++++--
 tests/distributions/test_discrete.py  | 117 ++++++++++++++++++++++++++
 3 files changed, 141 insertions(+), 5 deletions(-)

diff --git a/pymc_extras/distributions/__init__.py b/pymc_extras/distributions/__init__.py
index 404298b55..956c9e244 100644
--- a/pymc_extras/distributions/__init__.py
+++ b/pymc_extras/distributions/__init__.py
@@ -22,6 +22,7 @@
     BetaNegativeBinomial,
     GeneralizedPoisson,
     Skellam,
+    GrassiaIIGeometric,
 )
 from pymc_extras.distributions.histogram_utils import histogram_approximation
 from pymc_extras.distributions.multivariate import R2D2M2CP
diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 0331aeba9..9ba8cec22 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -15,6 +15,7 @@
 import numpy as np
 import pymc as pm
 
+from pymc.distributions.distribution import Discrete
 from pymc.distributions.dist_math import betaln, check_parameters, factln, logpow
 from pymc.distributions.shape_utils import rv_size_is_none
 from pytensor import tensor as pt
@@ -441,12 +442,11 @@ def sim_data(lam):
 g2g = GrassiaIIGeometricRV()
 
 
-class GrassiaIIGeometric(UnitContinuous):
+class GrassiaIIGeometric(Discrete):
     r"""Grassia(II)-Geometric distribution.
 
-    This distribution is a flexible alternative to the Geometric distribution for the
-    number of trials until a discrete event, and can be easily extended to support both static
-    and time-varying covariates.
+    This distribution is a flexible alternative to the Geometric distribution for the number of trials until a
+    discrete event, and can be extended to support both static and time-varying covariates.
 
     Hardie and Fader describe this distribution with the following PMF and survival functions in [1]_:
 
@@ -520,4 +520,22 @@ def logcdf(value, r, alpha):
             r > 0,
             alpha > 0,
             msg="s > 0, alpha > 0",
-        )
\ No newline at end of file
+        )
+
+    def support_point(rv, size, r, alpha):
+        """Calculate a reasonable starting point for sampling.
+        
+        For the GrassiaIIGeometric distribution, we use a point estimate based on
+        the expected value of the mixing distribution. Since the mixing distribution
+        is Gamma(r, 1/alpha), its mean is r/alpha. We then transform this through
+        the geometric link function and round to ensure an integer value.
+        """
+        # E[lambda] = r/alpha for Gamma(r, 1/alpha)
+        # p = 1 - exp(-lambda) for geometric
+        # E[T] = 1/p for geometric
+        mean = pt.ceil(pt.exp(alpha/r))  # Conservative upper bound
+        
+        if not rv_size_is_none(size):
+            mean = pt.full(size, mean)
+        
+        return mean
\ No newline at end of file
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index 8e803a31b..69071cd54 100644
--- a/tests/distributions/test_discrete.py
+++ b/tests/distributions/test_discrete.py
@@ -34,6 +34,7 @@
     BetaNegativeBinomial,
     GeneralizedPoisson,
     Skellam,
+    GrassiaIIGeometric,
 )
 
 
@@ -208,3 +209,119 @@ def test_logp(self):
                 {"mu1": Rplus_small, "mu2": Rplus_small},
                 lambda value, mu1, mu2: scipy.stats.skellam.logpmf(value, mu1, mu2),
             )
+
+
+class TestGrassiaIIGeometric:
+    class TestRandomVariable(BaseTestDistributionRandom):
+        pymc_dist = GrassiaIIGeometric
+        pymc_dist_params = {"r": 1.0, "alpha": 2.0}
+        expected_rv_op_params = {"r": 1.0, "alpha": 2.0}
+        tests_to_run = [
+            "check_pymc_params_match_rv_op",
+            "check_rv_size",
+        ]
+
+        def test_random_basic_properties(self):
+            discrete_random_tester(
+                dist=self.pymc_dist,
+                paramdomains={"r": Rplus, "alpha": Rplus},
+                ref_rand=lambda r, alpha, size: np.random.geometric(
+                    1 - np.exp(-np.random.gamma(r, 1/alpha, size=size)), size=size
+                ),
+            )
+
+        @pytest.mark.parametrize("r,alpha", [
+            (0.5, 1.0),
+            (1.0, 2.0),
+            (2.0, 0.5),
+            (5.0, 1.0),
+        ])
+        def test_random_moments(self, r, alpha):
+            dist = self.pymc_dist.dist(r=r, alpha=alpha, size=10_000)
+            draws = dist.eval()
+
+            # Check that all values are positive integers
+            assert np.all(draws > 0)
+            assert np.all(draws.astype(int) == draws)
+
+            # Check that values are reasonably distributed
+            # Note: Exact moments are complex for this distribution
+            # so we just check basic properties
+            assert np.mean(draws) > 0
+            assert np.var(draws) > 0
+
+    def test_logp_basic(self):
+        r = pt.scalar("r")
+        alpha = pt.scalar("alpha")
+        value = pt.vector("value", dtype="int64")
+
+        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha), value)
+        logp_fn = pytensor.function([value, r, alpha], logp)
+
+        # Test basic properties of logp
+        test_value = np.array([1, 2, 3, 4, 5])
+        test_r = 1.0
+        test_alpha = 1.0
+        
+        logp_vals = logp_fn(test_value, test_r, test_alpha)
+        assert not np.any(np.isnan(logp_vals))
+        assert np.all(np.isfinite(logp_vals))
+
+        # Test invalid values
+        assert logp_fn(np.array([0]), test_r, test_alpha) == np.inf  # Value must be > 0
+        
+        with pytest.raises(TypeError):
+            logp_fn(np.array([1.5]), test_r, test_alpha) == -np.inf  # Value must be integer
+
+        # Test parameter restrictions
+        with pytest.raises(ParameterValueError):
+            logp_fn(np.array([1]), -1.0, test_alpha)  # r must be > 0
+
+        with pytest.raises(ParameterValueError):
+            logp_fn(np.array([1]), test_r, -1.0)  # alpha must be > 0
+
+    def test_sampling_consistency(self):
+        """Test that sampling from the distribution produces reasonable results"""
+        r = 2.0
+        alpha = 1.0
+        with pm.Model():
+            x = GrassiaIIGeometric("x", r=r, alpha=alpha)
+            trace = pm.sample(chains=1, draws=1000, random_seed=42).posterior
+
+        samples = trace["x"].values.flatten()
+        
+        # Check basic properties of samples
+        assert np.all(samples > 0)  # All values should be positive
+        assert np.all(samples.astype(int) == samples)  # All values should be integers
+        
+        # Check mean and variance are reasonable
+        # (exact values depend on the parameterization)
+        assert 0 < np.mean(samples) < np.inf
+        assert 0 < np.var(samples) < np.inf
+
+    @pytest.mark.parametrize(
+        "r, alpha, size, expected_shape",
+        [
+            (1.0, 1.0, None, ()),  # Scalar output
+            ([1.0, 2.0], 1.0, None, (2,)),  # Vector output from r
+            (1.0, [1.0, 2.0], None, (2,)),  # Vector output from alpha
+            (1.0, 1.0, (3, 2), (3, 2)),  # Explicit size
+        ],
+    )
+    def test_support_point(self, r, alpha, size, expected_shape):
+        """Test that support_point returns reasonable values with correct shapes"""
+        with pm.Model() as model:
+            GrassiaIIGeometric("x", r=r, alpha=alpha, size=size)
+        
+        init_point = model.initial_point()["x"]
+        
+        # Check shape
+        assert init_point.shape == expected_shape
+        
+        # Check values are positive integers
+        assert np.all(init_point > 0)
+        assert np.all(init_point.astype(int) == init_point)
+        
+        # Check values are finite and reasonable
+        assert np.all(np.isfinite(init_point))
+        assert np.all(init_point < 1e6)  # Should not be extremely large

From d2e72b59c2f13c7490d2f4647d9b31522c327871 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Sun, 20 Apr 2025 18:33:44 -0500
Subject: [PATCH 04/13] alpha min value

---
 pymc_extras/distributions/discrete.py | 10 +++-------
 tests/distributions/test_discrete.py  |  8 ++++----
 2 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 9ba8cec22..b4b9eefd4 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -512,14 +512,13 @@ def logp(value, r, alpha):
         )
 
     def logcdf(value, r, alpha):
-        # TODO: Math may not be correct here
         logcdf = r * (pt.log(value) - pt.log(alpha + value))
 
         return check_parameters(
             logcdf,
             r > 0,
-            alpha > 0,
-            msg="s > 0, alpha > 0",
+            alpha > 0.6181, # alpha must be greater than 0.6181 for convergence
+            msg="r > 0, alpha > 0",
         )
 
     def support_point(rv, size, r, alpha):
@@ -530,10 +529,7 @@ def support_point(rv, size, r, alpha):
         is Gamma(r, 1/alpha), its mean is r/alpha. We then transform this through
         the geometric link function and round to ensure an integer value.
         """
-        # E[lambda] = r/alpha for Gamma(r, 1/alpha)
-        # p = 1 - exp(-lambda) for geometric
-        # E[T] = 1/p for geometric
-        mean = pt.ceil(pt.exp(alpha/r))  # Conservative upper bound
+        mean = pt.ceil(pt.exp(alpha/r))
         
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index 69071cd54..0bb5787d8 100644
--- a/tests/distributions/test_discrete.py
+++ b/tests/distributions/test_discrete.py
@@ -214,8 +214,8 @@ def test_logp(self):
 class TestGrassiaIIGeometric:
     class TestRandomVariable(BaseTestDistributionRandom):
         pymc_dist = GrassiaIIGeometric
-        pymc_dist_params = {"r": 1.0, "alpha": 2.0}
-        expected_rv_op_params = {"r": 1.0, "alpha": 2.0}
+        pymc_dist_params = {"r": .5, "alpha": 2.0}
+        expected_rv_op_params = {"r": .5, "alpha": 2.0}
         tests_to_run = [
             "check_pymc_params_match_rv_op",
             "check_rv_size",
@@ -289,11 +289,11 @@ def test_sampling_consistency(self):
             trace = pm.sample(chains=1, draws=1000, random_seed=42).posterior
 
         samples = trace["x"].values.flatten()
-        
+
         # Check basic properties of samples
         assert np.all(samples > 0)  # All values should be positive
         assert np.all(samples.astype(int) == samples)  # All values should be integers
-        
+
         # Check mean and variance are reasonable
         # (exact values depend on the parameterization)
         assert 0 < np.mean(samples) < np.inf

From 8685005bbf72b51a948f03c6e76436d1e95456d2 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Mon, 21 Apr 2025 15:03:54 -0500
Subject: [PATCH 05/13] revert alpha lim

---
 pymc_extras/distributions/discrete.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index b4b9eefd4..50238f5d6 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -423,7 +423,7 @@ def rng_fn(cls, rng, r, alpha, size):
 
         output = np.zeros(shape=size + (1,))  # noqa:RUF005
 
-        lam = rng.gamma(shape=r, scale=1 / alpha, size=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
@@ -517,7 +517,7 @@ def logcdf(value, r, alpha):
         return check_parameters(
             logcdf,
             r > 0,
-            alpha > 0.6181, # alpha must be greater than 0.6181 for convergence
+            alpha > 0, # alpha must be greater than 0.6181 for convergence
             msg="r > 0, alpha > 0",
         )
 

From 026f18272f0aae2b99421b7a078e548d4f07a79b Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Tue, 22 Apr 2025 00:52:00 -0500
Subject: [PATCH 06/13] small lam value tests

---
 pymc_extras/distributions/discrete.py | 14 ++++++++++----
 tests/distributions/test_discrete.py  | 18 +++++++++++++++++-
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 50238f5d6..26951c003 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/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):
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index 0bb5787d8..55dcfdf13 100644
--- a/tests/distributions/test_discrete.py
+++ b/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),

From d12dd0b1cece7db595abd7bd75eb73f12bff3311 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Tue, 22 Apr 2025 00:56:48 -0500
Subject: [PATCH 07/13] ruff formatting

---
 pymc_extras/distributions/discrete.py | 14 +++++++-------
 tests/distributions/test_discrete.py  | 16 ++++++++--------
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 26951c003..f99b9f268 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -15,8 +15,8 @@
 import numpy as np
 import pymc as pm
 
-from pymc.distributions.distribution import Discrete
 from pymc.distributions.dist_math import betaln, check_parameters, factln, logpow
+from pymc.distributions.distribution import Discrete
 from pymc.distributions.shape_utils import rv_size_is_none
 from pytensor import tensor as pt
 from pytensor.tensor.random.op import RandomVariable
@@ -432,10 +432,10 @@ def sim_data(lam):
                 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])
 
@@ -529,15 +529,15 @@ def logcdf(value, r, alpha):
 
     def support_point(rv, size, r, alpha):
         """Calculate a reasonable starting point for sampling.
-        
+
         For the GrassiaIIGeometric distribution, we use a point estimate based on
         the expected value of the mixing distribution. Since the mixing distribution
         is Gamma(r, 1/alpha), its mean is r/alpha. We then transform this through
         the geometric link function and round to ensure an integer value.
         """
         mean = pt.ceil(pt.exp(alpha/r))
-        
+
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)
-        
-        return mean
\ No newline at end of file
+
+        return mean
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index 55dcfdf13..e49b4a200 100644
--- a/tests/distributions/test_discrete.py
+++ b/tests/distributions/test_discrete.py
@@ -33,8 +33,8 @@
 from pymc_extras.distributions import (
     BetaNegativeBinomial,
     GeneralizedPoisson,
-    Skellam,
     GrassiaIIGeometric,
+    Skellam,
 )
 
 
@@ -233,7 +233,7 @@ def test_random_basic_properties(self):
                     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,
@@ -278,14 +278,14 @@ def test_logp_basic(self):
         test_value = np.array([1, 2, 3, 4, 5])
         test_r = 1.0
         test_alpha = 1.0
-        
+
         logp_vals = logp_fn(test_value, test_r, test_alpha)
         assert not np.any(np.isnan(logp_vals))
         assert np.all(np.isfinite(logp_vals))
 
         # Test invalid values
         assert logp_fn(np.array([0]), test_r, test_alpha) == np.inf  # Value must be > 0
-        
+
         with pytest.raises(TypeError):
             logp_fn(np.array([1.5]), test_r, test_alpha) == -np.inf  # Value must be integer
 
@@ -328,16 +328,16 @@ def test_support_point(self, r, alpha, size, expected_shape):
         """Test that support_point returns reasonable values with correct shapes"""
         with pm.Model() as model:
             GrassiaIIGeometric("x", r=r, alpha=alpha, size=size)
-        
+
         init_point = model.initial_point()["x"]
-        
+
         # Check shape
         assert init_point.shape == expected_shape
-        
+
         # Check values are positive integers
         assert np.all(init_point > 0)
         assert np.all(init_point.astype(int) == init_point)
-        
+
         # Check values are finite and reasonable
         assert np.all(np.isfinite(init_point))
         assert np.all(init_point < 1e6)  # Should not be extremely large

From bcd9cac4ed6d41add023e0191a1de198e004fce2 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Tue, 22 Apr 2025 00:04:48 -0600
Subject: [PATCH 08/13] TODOs

---
 pymc_extras/distributions/discrete.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index f99b9f268..02a509855 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -410,6 +410,7 @@ class GrassiaIIGeometricRV(RandomVariable):
     def __call__(self, r, alpha, size=None, **kwargs):
         return super().__call__(r, alpha, size=size, **kwargs)
 
+    # TODO: Param will need to be added for dot product of time-varying covariates
     @classmethod
     def rng_fn(cls, rng, r, alpha, size):
         if size is None:
@@ -430,7 +431,7 @@ def sim_data(lam):
             p = np.where(
                 lam < 0.001,
                 lam,  # For small lambda, p ≈ lambda
-                1 - np.exp(-lam)  # Standard formula for larger lambda
+                1 - np.exp(-lam)  # TODO: covariate param added here as 1 - np.exp(-lam * np.expcovar_dot)
             )
 
             # Ensure p is in valid range for geometric distribution
@@ -447,7 +448,7 @@ def sim_data(lam):
 
 g2g = GrassiaIIGeometricRV()
 
-
+# TODO: Add time-varying covariates. May simply replace the t-value , but is a continuous parameter
 class GrassiaIIGeometric(Discrete):
     r"""Grassia(II)-Geometric distribution.
 

From 78be1074a8639ef7e183ba19d258c8366c941064 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Tue, 22 Apr 2025 11:33:27 -0600
Subject: [PATCH 09/13] WIP add covar support to RV

---
 pymc_extras/distributions/discrete.py | 38 ++++++++++++++-------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 02a509855..bed3965bf 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -402,46 +402,48 @@ def dist(cls, mu1, mu2, **kwargs):
 
 class GrassiaIIGeometricRV(RandomVariable):
     name = "g2g"
-    signature = "(),()->()"
+    signature = "(),(),()->()"
 
     dtype = "int64"
     _print_name = ("GrassiaIIGeometric", "\\operatorname{GrassiaIIGeometric}")
 
-    def __call__(self, r, alpha, size=None, **kwargs):
-        return super().__call__(r, alpha, size=size, **kwargs)
+    def __call__(self, r, alpha, time_covar_dot=None,size=None, **kwargs):
+        return super().__call__(r, alpha, time_covar_dot=time_covar_dot, size=size, **kwargs)
 
-    # TODO: Param will need to be added for dot product of time-varying covariates
     @classmethod
-    def rng_fn(cls, rng, r, alpha, size):
+    def rng_fn(cls, rng, r, alpha, time_covar_dot,size):
+        if time_covar_dot is None:
+            time_covar_dot = np.array(0)
         if size is None:
-            size = np.broadcast_shapes(r.shape, alpha.shape)
-
-        r = np.asarray(r)
-        alpha = np.asarray(alpha)
+            size = np.broadcast_shapes(r.shape, alpha.shape, time_covar_dot.shape)
 
         r = np.broadcast_to(r, size)
         alpha = np.broadcast_to(alpha, size)
+        time_covar_dot = np.broadcast_to(time_covar_dot,size)
 
         output = np.zeros(shape=size + (1,))  # noqa:RUF005
 
         lam = rng.gamma(shape=r, scale=1/alpha, size=size)
 
-        def sim_data(lam):
+        exp_time_covar_dot = np.exp(time_covar_dot)
+
+        def sim_data(lam, exp_time_covar_dot):
             # Handle numerical stability for very small lambda values
-            p = np.where(
-                lam < 0.001,
-                lam,  # For small lambda, p ≈ lambda
-                1 - np.exp(-lam)  # TODO: covariate param added here as 1 - np.exp(-lam * np.expcovar_dot)
-            )
+            # p = np.where(
+            #     lam < 0.0001,
+            #     lam,  # For small lambda, p ≈ lambda
+            #     1 - np.exp(-lam * exp_time_covar_dot)
+            # )
 
-            # Ensure p is in valid range for geometric distribution
-            p = np.clip(p, 1e-5, 1.)
+            # Ensure lam is in valid range for geometric distribution
+            lam = np.clip(lam, np.finfo(float).tiny, 1.)
+            p = 1 - np.exp(-lam * exp_time_covar_dot)
 
             t = rng.geometric(p)
             return np.array([t])
 
         for index in np.ndindex(*size):
-            output[index] = sim_data(lam[index])
+            output[index] = sim_data(lam[index], exp_time_covar_dot[index])
 
         return output
 

From 8a304599f0feccb646a3f09d46d24e093b5944e3 Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Fri, 20 Jun 2025 08:37:10 -0600
Subject: [PATCH 10/13] WIP time indexing

---
 pymc_extras/distributions/discrete.py | 149 ++++++++++++++++++--------
 tests/distributions/test_discrete.py  | 109 +++++++++++++------
 2 files changed, 178 insertions(+), 80 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index bed3965bf..e7bdec46d 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -399,7 +399,7 @@ def dist(cls, mu1, mu2, **kwargs):
             **kwargs,
         )
 
-
+# TODO: C expressions are not correct. Both value and covariate broadcasting must be handled.
 class GrassiaIIGeometricRV(RandomVariable):
     name = "g2g"
     signature = "(),(),()->()"
@@ -407,51 +407,62 @@ class GrassiaIIGeometricRV(RandomVariable):
     dtype = "int64"
     _print_name = ("GrassiaIIGeometric", "\\operatorname{GrassiaIIGeometric}")
 
-    def __call__(self, r, alpha, time_covar_dot=None,size=None, **kwargs):
-        return super().__call__(r, alpha, time_covar_dot=time_covar_dot, size=size, **kwargs)
+    def __call__(self, r, alpha, time_covariates_sum=None, size=None, **kwargs):
+        return super().__call__(r, alpha, time_covariates_sum, size=size, **kwargs)
 
     @classmethod
-    def rng_fn(cls, rng, r, alpha, time_covar_dot,size):
-        if time_covar_dot is None:
-            time_covar_dot = np.array(0)
+    def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
+        if time_covariates_sum is None:
+            time_covariates_sum = np.array(0)
         if size is None:
-            size = np.broadcast_shapes(r.shape, alpha.shape, time_covar_dot.shape)
+            size = np.broadcast_shapes(r.shape, alpha.shape, time_covariates_sum.shape)
 
         r = np.broadcast_to(r, size)
         alpha = np.broadcast_to(alpha, size)
-        time_covar_dot = np.broadcast_to(time_covar_dot,size)
-
-        output = np.zeros(shape=size + (1,))  # noqa:RUF005
-
-        lam = rng.gamma(shape=r, scale=1/alpha, size=size)
-
-        exp_time_covar_dot = np.exp(time_covar_dot)
-
-        def sim_data(lam, exp_time_covar_dot):
-            # Handle numerical stability for very small lambda values
-            # p = np.where(
-            #     lam < 0.0001,
-            #     lam,  # For small lambda, p ≈ lambda
-            #     1 - np.exp(-lam * exp_time_covar_dot)
-            # )
-
-            # Ensure lam is in valid range for geometric distribution
-            lam = np.clip(lam, np.finfo(float).tiny, 1.)
-            p = 1 - np.exp(-lam * exp_time_covar_dot)
-
-            t = rng.geometric(p)
-            return np.array([t])
-
-        for index in np.ndindex(*size):
-            output[index] = sim_data(lam[index], exp_time_covar_dot[index])
-
+        time_covariates_sum = np.broadcast_to(time_covariates_sum, size)
+
+        # Calculate exp(time_covariates_sum) for all samples
+        exp_time_covar_sum = np.exp(time_covariates_sum)
+        
+        # Initialize output array
+        output = np.zeros(size, dtype=np.int64)
+        
+        # For each sample, generate a value from the distribution
+        for idx in np.ndindex(*size):
+            # Calculate survival probabilities for each possible value
+            t = 1
+            while True:
+                C_t = t + exp_time_covar_sum[idx]
+                C_tm1 = (t - 1) + exp_time_covar_sum[idx]
+                
+                # Calculate PMF for current t
+                pmf = (
+                    (alpha[idx] / (alpha[idx] + C_tm1)) ** r[idx] - 
+                    (alpha[idx] / (alpha[idx] + C_t)) ** r[idx]
+                )
+                
+                # If PMF is negative or NaN, we've gone too far
+                if pmf <= 0 or np.isnan(pmf):
+                    break
+                    
+                # Accept this value with probability proportional to PMF
+                if rng.random() < pmf:
+                    output[idx] = t
+                    break
+                    
+                t += 1
+                
+                # Safety check to prevent infinite loops
+                if t > 1000:  # Arbitrary large number
+                    output[idx] = t
+                    break
+        
         return output
 
 
 g2g = GrassiaIIGeometricRV()
 
-# TODO: Add time-varying covariates. May simply replace the t-value , but is a continuous parameter
-class GrassiaIIGeometric(Discrete):
+# TODO: C expressions are not correct. Both value and covariate broadcasting must be handled.
     r"""Grassia(II)-Geometric distribution.
 
     This distribution is a flexible alternative to the Geometric distribution for the number of trials until a
@@ -494,7 +505,9 @@ class GrassiaIIGeometric(Discrete):
         Shape parameter (r > 0).
     alpha : tensor_like of float
         Scale parameter (alpha > 0).
-
+    time_covariates_sum : tensor_like of float, optional
+        Optional dot product of time-varying covariates and their coefficients, summed over time.
+        
     References
     ----------
     .. [1] Fader, Peter & G. S. Hardie, Bruce (2020).
@@ -505,22 +518,56 @@ class GrassiaIIGeometric(Discrete):
     rv_op = g2g
 
     @classmethod
-    def dist(cls, r, alpha, *args, **kwargs):
+    def dist(cls, r, alpha, time_covariates_sum=None, *args, **kwargs):
         r = pt.as_tensor_variable(r)
         alpha = pt.as_tensor_variable(alpha)
-        return super().dist([r, alpha], *args, **kwargs)
-
-    def logp(value, r, alpha):
-        logp = -r * (pt.log(alpha + value - 1) + pt.log(alpha + value))
+        if time_covariates_sum is None:
+            time_covariates_sum = pt.constant(0.0)
+        time_covariates_sum = pt.as_tensor_variable(time_covariates_sum)
+        return super().dist([r, alpha, time_covariates_sum], *args, **kwargs)
 
+    def logp(value, r, alpha, time_covariates_sum=None):
+        """
+        Log probability function for GrassiaIIGeometric distribution.
+        
+        The PMF is:
+        P(T=t|r,α,β;Z(t)) = (α/(α+C(t-1)))^r - (α/(α+C(t)))^r
+        
+        where C(t) = t + exp(time_covariates_sum)
+        """
+        if time_covariates_sum is None:
+            time_covariates_sum = pt.constant(0.0)
+            
+        # Calculate C(t) and C(t-1)
+        C_t = value + pt.exp(time_covariates_sum)
+        C_tm1 = (value - 1) + pt.exp(time_covariates_sum)
+        
+        # Calculate the PMF on log scale
+        logp = pt.log(
+            pt.pow(alpha / (alpha + C_tm1), r) - 
+            pt.pow(alpha / (alpha + C_t), r)
+        )
+        
+        # Handle invalid values
+        logp = pt.switch(
+            pt.or_(
+                value < 1,  # Value must be >= 1
+                pt.isnan(logp),  # Handle NaN cases
+            ),
+            -np.inf,
+            logp
+        )
+        
         return check_parameters(
             logp,
             r > 0,
             alpha > 0,
-            msg="s > 0, alpha > 0",
+            msg="r > 0, alpha > 0",
         )
 
-    def logcdf(value, r, alpha):
+    def logcdf(value, r, alpha, time_covariates_sum=None):
+        if time_covariates_sum is not None:
+            value = time_covariates_sum
         logcdf = r * (pt.log(value) - pt.log(alpha + value))
 
         return check_parameters(
@@ -530,15 +577,27 @@ def logcdf(value, r, alpha):
             msg="r > 0, alpha > 0",
         )
 
-    def support_point(rv, size, r, alpha):
+    def support_point(rv, size, r, alpha, time_covariates_sum=None):
         """Calculate a reasonable starting point for sampling.
 
         For the GrassiaIIGeometric distribution, we use a point estimate based on
         the expected value of the mixing distribution. Since the mixing distribution
         is Gamma(r, 1/alpha), its mean is r/alpha. We then transform this through
         the geometric link function and round to ensure an integer value.
+
+        When time_covariates_sum is provided, it affects the expected value through
+        the exponential link function: exp(time_covariates_sum).
         """
-        mean = pt.ceil(pt.exp(alpha/r))
+        # Base mean without covariates
+        mean = pt.exp(alpha/r)
+
+        # Apply time-varying covariates if provided
+        if time_covariates_sum is None:
+            time_covariates_sum = pt.constant(0.0)
+        mean = mean * pt.exp(time_covariates_sum)
+
+        # Round up to nearest integer
+        mean = pt.ceil(mean)
 
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index e49b4a200..9259b31a2 100644
--- a/tests/distributions/test_discrete.py
+++ b/tests/distributions/test_discrete.py
@@ -214,23 +214,24 @@ def test_logp(self):
 class TestGrassiaIIGeometric:
     class TestRandomVariable(BaseTestDistributionRandom):
         pymc_dist = GrassiaIIGeometric
-        pymc_dist_params = {"r": .5, "alpha": 2.0}
-        expected_rv_op_params = {"r": .5, "alpha": 2.0}
+        pymc_dist_params = {"r": .5, "alpha": 2.0, "time_covariates_sum": 1.0}
+        expected_rv_op_params = {"r": .5, "alpha": 2.0, "time_covariates_sum": 1.0}
         tests_to_run = [
             "check_pymc_params_match_rv_op",
             "check_rv_size",
         ]
 
         def test_random_basic_properties(self):
-            # Test standard parameter values
+            # Test standard parameter values with time covariates
             discrete_random_tester(
                 dist=self.pymc_dist,
                 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
+                    "time_covariates_sum": Domain([-1.0, 1.0, 2.0], edges=(None, None)),  # Time covariates
                 },
-                ref_rand=lambda r, alpha, size: np.random.geometric(
-                    1 - np.exp(-np.random.gamma(r, 1/alpha, size=size)), size=size
+                ref_rand=lambda r, alpha, time_covariates_sum, size: np.random.geometric(
+                    1 - np.exp(-np.random.gamma(r, 1/alpha, size=size) * np.exp(time_covariates_sum)), size=size
                 ),
             )
 
@@ -240,20 +241,21 @@ def test_random_basic_properties(self):
                 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
+                    "time_covariates_sum": Domain([0.0, 1.0], edges=(None, None)),  # Time covariates
                 },
-                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
+                ref_rand=lambda r, alpha, time_covariates_sum, size: np.random.geometric(
+                    np.clip(np.random.gamma(r, 1/alpha, size=size) * np.exp(time_covariates_sum), 1e-5, 1.0), size=size
                 ),
             )
 
-        @pytest.mark.parametrize("r,alpha", [
-            (0.5, 1.0),
-            (1.0, 2.0),
-            (2.0, 0.5),
-            (5.0, 1.0),
+        @pytest.mark.parametrize("r,alpha,time_covariates_sum", [
+            (0.5, 1.0, 0.0),
+            (1.0, 2.0, 1.0),
+            (2.0, 0.5, -1.0),
+            (5.0, 1.0, None),
         ])
-        def test_random_moments(self, r, alpha):
-            dist = self.pymc_dist.dist(r=r, alpha=alpha, size=10_000)
+        def test_random_moments(self, r, alpha, time_covariates_sum):
+            dist = self.pymc_dist.dist(r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=10_000)
             draws = dist.eval()
 
             # Check that all values are positive integers
@@ -269,65 +271,102 @@ def test_random_moments(self, r, alpha):
     def test_logp_basic(self):
         r = pt.scalar("r")
         alpha = pt.scalar("alpha")
+        time_covariates_sum = pt.scalar("time_covariates_sum")
         value = pt.vector("value", dtype="int64")
 
-        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha), value)
-        logp_fn = pytensor.function([value, r, alpha], logp)
+        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha, time_covariates_sum), value)
+        logp_fn = pytensor.function([value, r, alpha, time_covariates_sum], logp)
 
         # Test basic properties of logp
         test_value = np.array([1, 2, 3, 4, 5])
         test_r = 1.0
         test_alpha = 1.0
+        test_time_covariates_sum = 1.0
 
-        logp_vals = logp_fn(test_value, test_r, test_alpha)
+        logp_vals = logp_fn(test_value, test_r, test_alpha, test_time_covariates_sum)
         assert not np.any(np.isnan(logp_vals))
         assert np.all(np.isfinite(logp_vals))
 
         # Test invalid values
-        assert logp_fn(np.array([0]), test_r, test_alpha) == np.inf  # Value must be > 0
+        assert logp_fn(np.array([0]), test_r, test_alpha, test_time_covariates_sum) == np.inf  # Value must be > 0
 
         with pytest.raises(TypeError):
-            logp_fn(np.array([1.5]), test_r, test_alpha) == -np.inf  # Value must be integer
+            logp_fn(np.array([1.5]), test_r, test_alpha, test_time_covariates_sum)  # Value must be integer
 
         # Test parameter restrictions
         with pytest.raises(ParameterValueError):
-            logp_fn(np.array([1]), -1.0, test_alpha)  # r must be > 0
+            logp_fn(np.array([1]), -1.0, test_alpha, test_time_covariates_sum)  # r must be > 0
 
         with pytest.raises(ParameterValueError):
-            logp_fn(np.array([1]), test_r, -1.0)  # alpha must be > 0
+            logp_fn(np.array([1]), test_r, -1.0, test_time_covariates_sum)  # alpha must be > 0
 
     def test_sampling_consistency(self):
         """Test that sampling from the distribution produces reasonable results"""
         r = 2.0
         alpha = 1.0
+        time_covariates_sum = None
+        
+        # First test direct sampling from the distribution
+        dist = GrassiaIIGeometric.dist(r=r, alpha=alpha, time_covariates_sum=time_covariates_sum)
+        direct_samples = dist.eval()
+        
+        # Convert to numpy array if it's not already
+        if not isinstance(direct_samples, np.ndarray):
+            direct_samples = np.array([direct_samples])
+        
+        # Ensure we have a 1D array
+        if direct_samples.ndim == 0:
+            direct_samples = direct_samples.reshape(1)
+            
+        assert direct_samples.size > 0, "Direct sampling produced no samples"
+        assert np.all(direct_samples > 0), "Direct sampling produced non-positive values"
+        assert np.all(direct_samples.astype(int) == direct_samples), "Direct sampling produced non-integer values"
+        
+        # Then test MCMC sampling
         with pm.Model():
-            x = GrassiaIIGeometric("x", r=r, alpha=alpha)
+            x = GrassiaIIGeometric("x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum)
             trace = pm.sample(chains=1, draws=1000, random_seed=42).posterior
 
-        samples = trace["x"].values.flatten()
+        # Extract samples and ensure they're in the correct shape
+        samples = trace["x"].values
+        assert samples is not None, "No samples were returned from MCMC"
+        assert samples.size > 0, "MCMC sampling produced empty array"
+        
+        if samples.ndim > 1:
+            samples = samples.reshape(-1)  # Flatten if needed
 
         # Check basic properties of samples
-        assert np.all(samples > 0)  # All values should be positive
-        assert np.all(samples.astype(int) == samples)  # All values should be integers
+        assert samples.size > 0, "No samples after reshaping"
+        assert np.all(samples > 0), "Found non-positive values in samples"
+        assert np.all(samples.astype(int) == samples), "Found non-integer values in samples"
 
         # Check mean and variance are reasonable
-        # (exact values depend on the parameterization)
-        assert 0 < np.mean(samples) < np.inf
-        assert 0 < np.var(samples) < np.inf
+        mean = np.mean(samples)
+        var = np.var(samples)
+        assert 0 < mean < np.inf, f"Mean {mean} is not in valid range"
+        assert 0 < var < np.inf, f"Variance {var} is not in valid range"
+
+        # Additional checks for distribution properties
+        # The mean should be greater than 1 for these parameters
+        assert mean > 1, f"Mean {mean} is not greater than 1"
+        # The variance should be positive and finite
+        assert var > 0, f"Variance {var} is not positive"
 
     @pytest.mark.parametrize(
-        "r, alpha, size, expected_shape",
+        "r, alpha, time_covariates_sum, size, expected_shape",
         [
-            (1.0, 1.0, None, ()),  # Scalar output
-            ([1.0, 2.0], 1.0, None, (2,)),  # Vector output from r
-            (1.0, [1.0, 2.0], None, (2,)),  # Vector output from alpha
-            (1.0, 1.0, (3, 2), (3, 2)),  # Explicit size
+            (1.0, 1.0, 1.0, None, ()),  # Scalar output with covariates
+            ([1.0, 2.0], 1.0, 1.0, None, (2,)),  # Vector output from r
+            (1.0, [1.0, 2.0], 1.0, None, (2,)),  # Vector output from alpha
+            (1.0, 1.0, None, None, ()),  # No time covariates
+            (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
         ],
     )
-    def test_support_point(self, r, alpha, size, expected_shape):
+    def test_support_point(self, r, alpha, time_covariates_sum, size, expected_shape):
         """Test that support_point returns reasonable values with correct shapes"""
         with pm.Model() as model:
-            GrassiaIIGeometric("x", r=r, alpha=alpha, size=size)
+            GrassiaIIGeometric("x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=size)
 
         init_point = model.initial_point()["x"]
 

From 7c7afc842d9c4c27b054f69755e970ec68107ffa Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Fri, 20 Jun 2025 08:37:49 -0600
Subject: [PATCH 11/13] WIP time indexing

---
 pymc_extras/distributions/discrete.py | 34 ++++++-------
 tests/distributions/test_discrete.py  | 72 ++++++++++++++++++---------
 2 files changed, 66 insertions(+), 40 deletions(-)

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index e7bdec46d..73df33416 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -423,10 +423,10 @@ def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
 
         # Calculate exp(time_covariates_sum) for all samples
         exp_time_covar_sum = np.exp(time_covariates_sum)
-        
+
         # Initialize output array
         output = np.zeros(size, dtype=np.int64)
-        
+
         # For each sample, generate a value from the distribution
         for idx in np.ndindex(*size):
             # Calculate survival probabilities for each possible value
@@ -434,29 +434,29 @@ def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
             while True:
                 C_t = t + exp_time_covar_sum[idx]
                 C_tm1 = (t - 1) + exp_time_covar_sum[idx]
-                
+
                 # Calculate PMF for current t
                 pmf = (
-                    (alpha[idx] / (alpha[idx] + C_tm1)) ** r[idx] - 
+                    (alpha[idx] / (alpha[idx] + C_tm1)) ** r[idx] -
                     (alpha[idx] / (alpha[idx] + C_t)) ** r[idx]
                 )
-                
+
                 # If PMF is negative or NaN, we've gone too far
                 if pmf <= 0 or np.isnan(pmf):
                     break
-                    
+
                 # Accept this value with probability proportional to PMF
                 if rng.random() < pmf:
                     output[idx] = t
                     break
-                    
+
                 t += 1
-                
+
                 # Safety check to prevent infinite loops
                 if t > 1000:  # Arbitrary large number
                     output[idx] = t
                     break
-        
+
         return output
 
 
@@ -507,7 +507,7 @@ def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
         Scale parameter (alpha > 0).
     time_covariates_sum : tensor_like of float, optional
         Optional dot product of time-varying covariates and their coefficients, summed over time.
-        
+
     References
     ----------
     .. [1] Fader, Peter & G. S. Hardie, Bruce (2020).
@@ -529,25 +529,25 @@ def dist(cls, r, alpha, time_covariates_sum=None, *args, **kwargs):
     def logp(value, r, alpha, time_covariates_sum=None):
         """
         Log probability function for GrassiaIIGeometric distribution.
-        
+
         The PMF is:
         P(T=t|r,α,β;Z(t)) = (α/(α+C(t-1)))^r - (α/(α+C(t)))^r
-        
+
         where C(t) = t + exp(time_covariates_sum)
         """
         if time_covariates_sum is None:
             time_covariates_sum = pt.constant(0.0)
-            
+
         # Calculate C(t) and C(t-1)
         C_t = value + pt.exp(time_covariates_sum)
         C_tm1 = (value - 1) + pt.exp(time_covariates_sum)
-        
+
         # Calculate the PMF on log scale
         logp = pt.log(
-            pt.pow(alpha / (alpha + C_tm1), r) - 
+            pt.pow(alpha / (alpha + C_tm1), r) -
             pt.pow(alpha / (alpha + C_t), r)
         )
-        
+
         # Handle invalid values
         logp = pt.switch(
             pt.or_(
@@ -557,7 +557,7 @@ def logp(value, r, alpha, time_covariates_sum=None):
             -np.inf,
             logp
         )
-        
+
         return check_parameters(
             logp,
             r > 0,
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index 9259b31a2..93cd80b60 100644
--- a/tests/distributions/test_discrete.py
+++ b/tests/distributions/test_discrete.py
@@ -214,8 +214,8 @@ def test_logp(self):
 class TestGrassiaIIGeometric:
     class TestRandomVariable(BaseTestDistributionRandom):
         pymc_dist = GrassiaIIGeometric
-        pymc_dist_params = {"r": .5, "alpha": 2.0, "time_covariates_sum": 1.0}
-        expected_rv_op_params = {"r": .5, "alpha": 2.0, "time_covariates_sum": 1.0}
+        pymc_dist_params = {"r": 0.5, "alpha": 2.0, "time_covariates_sum": 1.0}
+        expected_rv_op_params = {"r": 0.5, "alpha": 2.0, "time_covariates_sum": 1.0}
         tests_to_run = [
             "check_pymc_params_match_rv_op",
             "check_rv_size",
@@ -228,10 +228,16 @@ def test_random_basic_properties(self):
                 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
-                    "time_covariates_sum": Domain([-1.0, 1.0, 2.0], edges=(None, None)),  # Time covariates
+                    "time_covariates_sum": Domain(
+                        [-1.0, 1.0, 2.0], edges=(None, None)
+                    ),  # Time covariates
                 },
                 ref_rand=lambda r, alpha, time_covariates_sum, size: np.random.geometric(
-                    1 - np.exp(-np.random.gamma(r, 1/alpha, size=size) * np.exp(time_covariates_sum)), size=size
+                    1
+                    - np.exp(
+                        -np.random.gamma(r, 1 / alpha, size=size) * np.exp(time_covariates_sum)
+                    ),
+                    size=size,
                 ),
             )
 
@@ -241,21 +247,33 @@ def test_random_basic_properties(self):
                 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
-                    "time_covariates_sum": Domain([0.0, 1.0], edges=(None, None)),  # Time covariates
+                    "time_covariates_sum": Domain(
+                        [0.0, 1.0], edges=(None, None)
+                    ),  # Time covariates
                 },
                 ref_rand=lambda r, alpha, time_covariates_sum, size: np.random.geometric(
-                    np.clip(np.random.gamma(r, 1/alpha, size=size) * np.exp(time_covariates_sum), 1e-5, 1.0), size=size
+                    np.clip(
+                        np.random.gamma(r, 1 / alpha, size=size) * np.exp(time_covariates_sum),
+                        1e-5,
+                        1.0,
+                    ),
+                    size=size,
                 ),
             )
 
-        @pytest.mark.parametrize("r,alpha,time_covariates_sum", [
-            (0.5, 1.0, 0.0),
-            (1.0, 2.0, 1.0),
-            (2.0, 0.5, -1.0),
-            (5.0, 1.0, None),
-        ])
+        @pytest.mark.parametrize(
+            "r,alpha,time_covariates_sum",
+            [
+                (0.5, 1.0, 0.0),
+                (1.0, 2.0, 1.0),
+                (2.0, 0.5, -1.0),
+                (5.0, 1.0, None),
+            ],
+        )
         def test_random_moments(self, r, alpha, time_covariates_sum):
-            dist = self.pymc_dist.dist(r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=10_000)
+            dist = self.pymc_dist.dist(
+                r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=10_000
+            )
             draws = dist.eval()
 
             # Check that all values are positive integers
@@ -288,10 +306,14 @@ def test_logp_basic(self):
         assert np.all(np.isfinite(logp_vals))
 
         # Test invalid values
-        assert logp_fn(np.array([0]), test_r, test_alpha, test_time_covariates_sum) == np.inf  # Value must be > 0
+        assert (
+            logp_fn(np.array([0]), test_r, test_alpha, test_time_covariates_sum) == np.inf
+        )  # Value must be > 0
 
         with pytest.raises(TypeError):
-            logp_fn(np.array([1.5]), test_r, test_alpha, test_time_covariates_sum)  # Value must be integer
+            logp_fn(
+                np.array([1.5]), test_r, test_alpha, test_time_covariates_sum
+            )  # Value must be integer
 
         # Test parameter restrictions
         with pytest.raises(ParameterValueError):
@@ -305,23 +327,25 @@ def test_sampling_consistency(self):
         r = 2.0
         alpha = 1.0
         time_covariates_sum = None
-        
+
         # First test direct sampling from the distribution
         dist = GrassiaIIGeometric.dist(r=r, alpha=alpha, time_covariates_sum=time_covariates_sum)
         direct_samples = dist.eval()
-        
+
         # Convert to numpy array if it's not already
         if not isinstance(direct_samples, np.ndarray):
             direct_samples = np.array([direct_samples])
-        
+
         # Ensure we have a 1D array
         if direct_samples.ndim == 0:
             direct_samples = direct_samples.reshape(1)
-            
+
         assert direct_samples.size > 0, "Direct sampling produced no samples"
         assert np.all(direct_samples > 0), "Direct sampling produced non-positive values"
-        assert np.all(direct_samples.astype(int) == direct_samples), "Direct sampling produced non-integer values"
-        
+        assert np.all(
+            direct_samples.astype(int) == direct_samples
+        ), "Direct sampling produced non-integer values"
+
         # Then test MCMC sampling
         with pm.Model():
             x = GrassiaIIGeometric("x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum)
@@ -331,7 +355,7 @@ def test_sampling_consistency(self):
         samples = trace["x"].values
         assert samples is not None, "No samples were returned from MCMC"
         assert samples.size > 0, "MCMC sampling produced empty array"
-        
+
         if samples.ndim > 1:
             samples = samples.reshape(-1)  # Flatten if needed
 
@@ -366,7 +390,9 @@ def test_sampling_consistency(self):
     def test_support_point(self, r, alpha, time_covariates_sum, size, expected_shape):
         """Test that support_point returns reasonable values with correct shapes"""
         with pm.Model() as model:
-            GrassiaIIGeometric("x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=size)
+            GrassiaIIGeometric(
+                "x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=size
+            )
 
         init_point = model.initial_point()["x"]
 

From b957333c4bf0ede37fa068a17cb1a3935e8a6f4a Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Fri, 20 Jun 2025 13:57:32 -0600
Subject: [PATCH 12/13] WIP symbolic indexing

---
 pymc_extras/distributions/discrete.py | 171 ++++++++++++-----------
 test_simple.py                        |  61 +++++++++
 tests/distributions/test_discrete.py  | 189 ++++++++++++++++----------
 3 files changed, 270 insertions(+), 151 deletions(-)
 create mode 100644 test_simple.py

diff --git a/pymc_extras/distributions/discrete.py b/pymc_extras/distributions/discrete.py
index 4af28b5c8..09c6467e2 100644
--- a/pymc_extras/distributions/discrete.py
+++ b/pymc_extras/distributions/discrete.py
@@ -401,7 +401,7 @@ def dist(cls, mu1, mu2, **kwargs):
             **kwargs,
         )
 
-# TODO: C expressions are not correct. Both value and covariate broadcasting must be handled.
+
 class GrassiaIIGeometricRV(RandomVariable):
     name = "g2g"
     signature = "(),(),()->()"
@@ -409,62 +409,55 @@ class GrassiaIIGeometricRV(RandomVariable):
     dtype = "int64"
     _print_name = ("GrassiaIIGeometric", "\\operatorname{GrassiaIIGeometric}")
 
-    def __call__(self, r, alpha, time_covariates_sum=None, size=None, **kwargs):
-        return super().__call__(r, alpha, time_covariates_sum, size=size, **kwargs)
+    def __call__(self, r, alpha, time_covariate_vector=None, size=None, **kwargs):
+        return super().__call__(r, alpha, time_covariate_vector, size=size, **kwargs)
 
     @classmethod
-    def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
-        if time_covariates_sum is None:
-            time_covariates_sum = np.array(0)
+    def rng_fn(cls, rng, r, alpha, time_covariate_vector, size):
+        # Handle None case for time_covariate_vector
+        if time_covariate_vector is None:
+            time_covariate_vector = 0.0
+
+        # Convert inputs to numpy arrays - these should be concrete values
+        r = np.asarray(r, dtype=np.float64)
+        alpha = np.asarray(alpha, dtype=np.float64)
+        time_covariate_vector = np.asarray(time_covariate_vector, dtype=np.float64)
+
+        # Determine output size
         if size is None:
-            size = np.broadcast_shapes(r.shape, alpha.shape, time_covariates_sum.shape)
+            size = np.broadcast_shapes(r.shape, alpha.shape, time_covariate_vector.shape)
 
+        # Broadcast parameters to the output size
         r = np.broadcast_to(r, size)
         alpha = np.broadcast_to(alpha, size)
-        time_covariates_sum = np.broadcast_to(time_covariates_sum, size)
-
-        # Calculate exp(time_covariates_sum) for all samples
-        exp_time_covar_sum = np.exp(time_covariates_sum)
-
-        # Initialize output array
-        output = np.zeros(size, dtype=np.int64)
-
-        # For each sample, generate a value from the distribution
-        for idx in np.ndindex(*size):
-            # Calculate survival probabilities for each possible value
-            t = 1
-            while True:
-                C_t = t + exp_time_covar_sum[idx]
-                C_tm1 = (t - 1) + exp_time_covar_sum[idx]
-
-                # Calculate PMF for current t
-                pmf = (
-                    (alpha[idx] / (alpha[idx] + C_tm1)) ** r[idx] -
-                    (alpha[idx] / (alpha[idx] + C_t)) ** r[idx]
-                )
+        time_covariate_vector = np.broadcast_to(time_covariate_vector, size)
 
-                # If PMF is negative or NaN, we've gone too far
-                if pmf <= 0 or np.isnan(pmf):
-                    break
+        # Calculate exp(time_covariate_vector) for all samples
+        exp_time_covar_sum = np.exp(time_covariate_vector)
 
-                # Accept this value with probability proportional to PMF
-                if rng.random() < pmf:
-                    output[idx] = t
-                    break
+        # Use a simpler approach: generate from a geometric distribution with transformed parameters
+        # This is an approximation but should be much faster and more reliable
+        lam = rng.gamma(shape=r, scale=1 / alpha, size=size)
+        lam_covar = lam * exp_time_covar_sum
 
-                t += 1
+        # Handle numerical stability for very small lambda values
+        p = np.where(
+            lam_covar < 0.0001,
+            lam_covar,  # For small values, set this to p
+            1 - np.exp(-lam_covar),
+        )
 
-                # Safety check to prevent infinite loops
-                if t > 1000:  # Arbitrary large number
-                    output[idx] = t
-                    break
+        # Ensure p is in valid range for geometric distribution
+        p = np.clip(p, np.finfo(float).tiny, 1.0)
 
-        return output
+        # Generate geometric samples
+        return rng.geometric(p)
 
 
 g2g = GrassiaIIGeometricRV()
 
-# TODO: C expressions are not correct. Both value and covariate broadcasting must be handled.
+
+class GrassiaIIGeometric(Discrete):
     r"""Grassia(II)-Geometric distribution.
 
     This distribution is a flexible alternative to the Geometric distribution for the number of trials until a
@@ -507,8 +500,8 @@ def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
         Shape parameter (r > 0).
     alpha : tensor_like of float
         Scale parameter (alpha > 0).
-    time_covariates_sum : tensor_like of float, optional
-        Optional dot product of time-varying covariates and their coefficients, summed over time.
+    time_covariate_vector : tensor_like of float, optional
+        Optional vector of dot product of time-varying covariates and their coefficients by time period.
 
     References
     ----------
@@ -520,34 +513,36 @@ def rng_fn(cls, rng, r, alpha, time_covariates_sum, size):
     rv_op = g2g
 
     @classmethod
-    def dist(cls, r, alpha, time_covariates_sum=None, *args, **kwargs):
+    def dist(cls, r, alpha, time_covariate_vector=None, *args, **kwargs):
         r = pt.as_tensor_variable(r)
         alpha = pt.as_tensor_variable(alpha)
-        if time_covariates_sum is None:
-            time_covariates_sum = pt.constant(0.0)
-        time_covariates_sum = pt.as_tensor_variable(time_covariates_sum)
-        return super().dist([r, alpha, time_covariates_sum], *args, **kwargs)
-
-    def logp(value, r, alpha, time_covariates_sum=None):
-        """
-        Log probability function for GrassiaIIGeometric distribution.
-
-        The PMF is:
-        P(T=t|r,α,β;Z(t)) = (α/(α+C(t-1)))^r - (α/(α+C(t)))^r
-
-        where C(t) = t + exp(time_covariates_sum)
-        """
-        if time_covariates_sum is None:
-            time_covariates_sum = pt.constant(0.0)
-
-        # Calculate C(t) and C(t-1)
-        C_t = value + pt.exp(time_covariates_sum)
-        C_tm1 = (value - 1) + pt.exp(time_covariates_sum)
+        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)
+
+    def logp(value, r, alpha, time_covariate_vector=None):
+        if time_covariate_vector is None:
+            time_covariate_vector = pt.constant(0.0)
+        time_covariate_vector = pt.as_tensor_variable(time_covariate_vector)
+
+        def C_t(t):
+            # Aggregate time_covariate_vector over active time periods
+            if t == 0:
+                return pt.constant(1.0)
+            # Handle case where time_covariate_vector is a scalar
+            if time_covariate_vector.ndim == 0:
+                return t * pt.exp(time_covariate_vector)
+            else:
+                # For vector time_covariate_vector, we need to handle symbolic indexing
+                # Since we can't slice with symbolic indices, we'll use a different approach
+                # For now, we'll use the first element multiplied by t
+                # This is a simplification but should work for basic cases
+                return t * pt.exp(time_covariate_vector[:t])
 
         # Calculate the PMF on log scale
         logp = pt.log(
-            pt.pow(alpha / (alpha + C_tm1), r) -
-            pt.pow(alpha / (alpha + C_t), r)
+            pt.pow(alpha / (alpha + C_t(value - 1)), r) - pt.pow(alpha / (alpha + C_t(value)), r)
         )
 
         # Handle invalid values
@@ -557,7 +552,7 @@ def logp(value, r, alpha, time_covariates_sum=None):
                 pt.isnan(logp),  # Handle NaN cases
             ),
             -np.inf,
-            logp
+            logp,
         )
 
         return check_parameters(
@@ -567,19 +562,35 @@ def logp(value, r, alpha, time_covariates_sum=None):
             msg="r > 0, alpha > 0",
         )
 
-    def logcdf(value, r, alpha, time_covariates_sum=None):
-        if time_covariates_sum is not None:
-            value = time_covariates_sum
-        logcdf = r * (pt.log(value) - pt.log(alpha + value))
+    def logcdf(value, r, alpha, time_covariate_vector=None):
+        if time_covariate_vector is None:
+            time_covariate_vector = pt.constant(0.0)
+        time_covariate_vector = pt.as_tensor_variable(time_covariate_vector)
+
+        # Calculate CDF on log scale
+        # For the GrassiaIIGeometric, the CDF is 1 - survival function
+        # S(t) = (alpha/(alpha + C(t)))^r
+        # CDF(t) = 1 - S(t)
+
+        def C_t(t):
+            if t == 0:
+                return pt.constant(1.0)
+            if time_covariate_vector.ndim == 0:
+                return t * pt.exp(time_covariate_vector)
+            else:
+                return t * pt.exp(time_covariate_vector[:t])
+
+        survival = pt.pow(alpha / (alpha + C_t(value)), r)
+        logcdf = pt.log(1 - survival)
 
         return check_parameters(
             logcdf,
             r > 0,
-            alpha > 0, # alpha must be greater than 0.6181 for convergence
+            alpha > 0,
             msg="r > 0, alpha > 0",
         )
 
-    def support_point(rv, size, r, alpha, time_covariates_sum=None):
+    def support_point(rv, size, r, alpha, time_covariate_vector=None):
         """Calculate a reasonable starting point for sampling.
 
         For the GrassiaIIGeometric distribution, we use a point estimate based on
@@ -587,16 +598,16 @@ def support_point(rv, size, r, alpha, time_covariates_sum=None):
         is Gamma(r, 1/alpha), its mean is r/alpha. We then transform this through
         the geometric link function and round to ensure an integer value.
 
-        When time_covariates_sum is provided, it affects the expected value through
-        the exponential link function: exp(time_covariates_sum).
+        When time_covariate_vector is provided, it affects the expected value through
+        the exponential link function: exp(time_covariate_vector).
         """
         # Base mean without covariates
-        mean = pt.exp(alpha/r)
+        mean = pt.exp(alpha / r)
 
         # Apply time-varying covariates if provided
-        if time_covariates_sum is None:
-            time_covariates_sum = pt.constant(0.0)
-        mean = mean * pt.exp(time_covariates_sum)
+        if time_covariate_vector is None:
+            time_covariate_vector = pt.constant(0.0)
+        mean = mean * pt.exp(time_covariate_vector)
 
         # Round up to nearest integer
         mean = pt.ceil(mean)
diff --git a/test_simple.py b/test_simple.py
new file mode 100644
index 000000000..ada1a0680
--- /dev/null
+++ b/test_simple.py
@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+
+import pymc as pm
+import pytensor.tensor as pt
+
+from pymc_extras.distributions import GrassiaIIGeometric
+
+
+def test_basic_functionality():
+    """Test basic functionality of GrassiaIIGeometric distribution"""
+    print("Testing basic GrassiaIIGeometric functionality...")
+
+    # Test 1: Create distribution with None time_covariate_vector
+    try:
+        dist = GrassiaIIGeometric.dist(r=2.0, alpha=1.0, time_covariate_vector=None)
+        print("✓ Distribution created successfully with None time_covariate_vector")
+
+        # Test sampling
+        samples = dist.eval()
+        print(f"✓ Direct sampling successful: {samples}")
+
+    except Exception as e:
+        print(f"✗ Failed to create distribution with None time_covariate_vector: {e}")
+        return False
+
+    # Test 2: Create distribution with scalar time_covariate_vector
+    try:
+        dist = GrassiaIIGeometric.dist(r=2.0, alpha=1.0, time_covariate_vector=0.5)
+        print("✓ Distribution created successfully with scalar time_covariate_vector")
+
+        # Test sampling
+        samples = dist.eval()
+        print(f"✓ Direct sampling successful: {samples}")
+
+    except Exception as e:
+        print(f"✗ Failed to create distribution with scalar time_covariate_vector: {e}")
+        return False
+
+    # Test 3: Test logp function
+    try:
+        r = pt.scalar("r")
+        alpha = pt.scalar("alpha")
+        time_covariate_vector = pt.scalar("time_covariate_vector")
+        value = pt.scalar("value", dtype="int64")
+
+        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha, time_covariate_vector), value)
+        logp_fn = pm.compile_fn([value, r, alpha, time_covariate_vector], logp)
+
+        result = logp_fn(2, 1.0, 1.0, 0.0)
+        print(f"✓ Logp function works: {result}")
+
+    except Exception as e:
+        print(f"✗ Failed to test logp function: {e}")
+        return False
+
+    print("✓ All basic functionality tests passed!")
+    return True
+
+
+if __name__ == "__main__":
+    test_basic_functionality()
diff --git a/tests/distributions/test_discrete.py b/tests/distributions/test_discrete.py
index 93cd80b60..befc7594e 100644
--- a/tests/distributions/test_discrete.py
+++ b/tests/distributions/test_discrete.py
@@ -214,8 +214,8 @@ def test_logp(self):
 class TestGrassiaIIGeometric:
     class TestRandomVariable(BaseTestDistributionRandom):
         pymc_dist = GrassiaIIGeometric
-        pymc_dist_params = {"r": 0.5, "alpha": 2.0, "time_covariates_sum": 1.0}
-        expected_rv_op_params = {"r": 0.5, "alpha": 2.0, "time_covariates_sum": 1.0}
+        pymc_dist_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": 1.0}
+        expected_rv_op_params = {"r": 0.5, "alpha": 2.0, "time_covariate_vector": 1.0}
         tests_to_run = [
             "check_pymc_params_match_rv_op",
             "check_rv_size",
@@ -228,14 +228,14 @@ def test_random_basic_properties(self):
                 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
-                    "time_covariates_sum": Domain(
+                    "time_covariate_vector": Domain(
                         [-1.0, 1.0, 2.0], edges=(None, None)
                     ),  # Time covariates
                 },
-                ref_rand=lambda r, alpha, time_covariates_sum, size: np.random.geometric(
+                ref_rand=lambda r, alpha, time_covariate_vector, size: np.random.geometric(
                     1
                     - np.exp(
-                        -np.random.gamma(r, 1 / alpha, size=size) * np.exp(time_covariates_sum)
+                        -np.random.gamma(r, 1 / alpha, size=size) * np.exp(time_covariate_vector)
                     ),
                     size=size,
                 ),
@@ -247,13 +247,13 @@ def test_random_basic_properties(self):
                 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
-                    "time_covariates_sum": Domain(
+                    "time_covariate_vector": Domain(
                         [0.0, 1.0], edges=(None, None)
                     ),  # Time covariates
                 },
-                ref_rand=lambda r, alpha, time_covariates_sum, size: np.random.geometric(
+                ref_rand=lambda r, alpha, time_covariate_vector, size: np.random.geometric(
                     np.clip(
-                        np.random.gamma(r, 1 / alpha, size=size) * np.exp(time_covariates_sum),
+                        np.random.gamma(r, 1 / alpha, size=size) * np.exp(time_covariate_vector),
                         1e-5,
                         1.0,
                     ),
@@ -262,7 +262,7 @@ def test_random_basic_properties(self):
             )
 
         @pytest.mark.parametrize(
-            "r,alpha,time_covariates_sum",
+            "r,alpha,time_covariate_vector",
             [
                 (0.5, 1.0, 0.0),
                 (1.0, 2.0, 1.0),
@@ -270,9 +270,9 @@ def test_random_basic_properties(self):
                 (5.0, 1.0, None),
             ],
         )
-        def test_random_moments(self, r, alpha, time_covariates_sum):
+        def test_random_moments(self, r, alpha, time_covariate_vector):
             dist = self.pymc_dist.dist(
-                r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=10_000
+                r=r, alpha=alpha, time_covariate_vector=time_covariate_vector, size=10_000
             )
             draws = dist.eval()
 
@@ -289,109 +289,156 @@ def test_random_moments(self, r, alpha, time_covariates_sum):
     def test_logp_basic(self):
         r = pt.scalar("r")
         alpha = pt.scalar("alpha")
-        time_covariates_sum = pt.scalar("time_covariates_sum")
+        time_covariate_vector = pt.vector("time_covariate_vector")
         value = pt.vector("value", dtype="int64")
 
-        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha, time_covariates_sum), value)
-        logp_fn = pytensor.function([value, r, alpha, time_covariates_sum], logp)
+        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha, time_covariate_vector), value)
+        logp_fn = pytensor.function([value, r, alpha, time_covariate_vector], logp)
 
         # Test basic properties of logp
-        test_value = np.array([1, 2, 3, 4, 5])
+        test_value = np.array([1, 1, 2, 3, 4, 5])
         test_r = 1.0
         test_alpha = 1.0
-        test_time_covariates_sum = 1.0
+        test_time_covariate_vector = np.array(
+            [
+                None,
+                [1],
+                [1, 2],
+                [1, 2, 3],
+                [1, 2, 3, 4],
+                [1, 2, 3, 4, 5],
+            ]
+        )
 
-        logp_vals = logp_fn(test_value, test_r, test_alpha, test_time_covariates_sum)
+        logp_vals = logp_fn(test_value, test_r, test_alpha, test_time_covariate_vector)
         assert not np.any(np.isnan(logp_vals))
         assert np.all(np.isfinite(logp_vals))
 
         # Test invalid values
         assert (
-            logp_fn(np.array([0]), test_r, test_alpha, test_time_covariates_sum) == np.inf
+            logp_fn(np.array([0]), test_r, test_alpha, test_time_covariate_vector) == np.inf
         )  # Value must be > 0
 
         with pytest.raises(TypeError):
             logp_fn(
-                np.array([1.5]), test_r, test_alpha, test_time_covariates_sum
+                np.array([1.5]), test_r, test_alpha, test_time_covariate_vector
             )  # Value must be integer
 
         # Test parameter restrictions
         with pytest.raises(ParameterValueError):
-            logp_fn(np.array([1]), -1.0, test_alpha, test_time_covariates_sum)  # r must be > 0
+            logp_fn(np.array([1]), -1.0, test_alpha, test_time_covariate_vector)  # r must be > 0
 
         with pytest.raises(ParameterValueError):
-            logp_fn(np.array([1]), test_r, -1.0, test_time_covariates_sum)  # alpha must be > 0
+            logp_fn(np.array([1]), test_r, -1.0, test_time_covariate_vector)  # alpha must be > 0
 
     def test_sampling_consistency(self):
         """Test that sampling from the distribution produces reasonable results"""
         r = 2.0
         alpha = 1.0
-        time_covariates_sum = None
+        time_covariate_vector = None  # Start with just None case
 
         # First test direct sampling from the distribution
-        dist = GrassiaIIGeometric.dist(r=r, alpha=alpha, time_covariates_sum=time_covariates_sum)
-        direct_samples = dist.eval()
+        try:
+            dist = GrassiaIIGeometric.dist(
+                r=r, alpha=alpha, time_covariate_vector=time_covariate_vector
+            )
+
+            direct_samples = dist.eval()
 
-        # Convert to numpy array if it's not already
-        if not isinstance(direct_samples, np.ndarray):
-            direct_samples = np.array([direct_samples])
+            # Convert to numpy array if it's not already
+            if not isinstance(direct_samples, np.ndarray):
+                direct_samples = np.array([direct_samples])
 
-        # Ensure we have a 1D array
-        if direct_samples.ndim == 0:
-            direct_samples = direct_samples.reshape(1)
+            # Ensure we have a 1D array
+            if direct_samples.ndim == 0:
+                direct_samples = direct_samples.reshape(1)
 
-        assert direct_samples.size > 0, "Direct sampling produced no samples"
-        assert np.all(direct_samples > 0), "Direct sampling produced non-positive values"
-        assert np.all(
-            direct_samples.astype(int) == direct_samples
-        ), "Direct sampling produced non-integer values"
+            assert (
+                direct_samples.size > 0
+            ), f"Direct sampling produced no samples for {time_covariate_vector}"
+            assert np.all(
+                direct_samples > 0
+            ), f"Direct sampling produced non-positive values for {time_covariate_vector}"
+            assert np.all(
+                direct_samples.astype(int) == direct_samples
+            ), f"Direct sampling produced non-integer values for {time_covariate_vector}"
+
+        except Exception as e:
+            import traceback
+
+            traceback.print_exc()
+            raise
 
         # Then test MCMC sampling
-        with pm.Model():
-            x = GrassiaIIGeometric("x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum)
-            trace = pm.sample(chains=1, draws=1000, random_seed=42).posterior
-
-        # Extract samples and ensure they're in the correct shape
-        samples = trace["x"].values
-        assert samples is not None, "No samples were returned from MCMC"
-        assert samples.size > 0, "MCMC sampling produced empty array"
-
-        if samples.ndim > 1:
-            samples = samples.reshape(-1)  # Flatten if needed
-
-        # Check basic properties of samples
-        assert samples.size > 0, "No samples after reshaping"
-        assert np.all(samples > 0), "Found non-positive values in samples"
-        assert np.all(samples.astype(int) == samples), "Found non-integer values in samples"
-
-        # Check mean and variance are reasonable
-        mean = np.mean(samples)
-        var = np.var(samples)
-        assert 0 < mean < np.inf, f"Mean {mean} is not in valid range"
-        assert 0 < var < np.inf, f"Variance {var} is not in valid range"
-
-        # Additional checks for distribution properties
-        # The mean should be greater than 1 for these parameters
-        assert mean > 1, f"Mean {mean} is not greater than 1"
-        # The variance should be positive and finite
-        assert var > 0, f"Variance {var} is not positive"
+        try:
+            with pm.Model():
+                x = GrassiaIIGeometric(
+                    "x", r=r, alpha=alpha, time_covariate_vector=time_covariate_vector
+                )
+
+                trace = pm.sample(
+                    chains=1, draws=50, tune=0, random_seed=42, progressbar=False
+                ).posterior
+
+            # Extract samples and ensure they're in the correct shape
+            samples = trace["x"].values
+
+            assert (
+                samples is not None
+            ), f"No samples were returned from MCMC for {time_covariate_vector}"
+            assert (
+                samples.size > 0
+            ), f"MCMC sampling produced empty array for {time_covariate_vector}"
+
+            if samples.ndim > 1:
+                samples = samples.reshape(-1)  # Flatten if needed
+
+            # Check basic properties of samples
+            assert samples.size > 0, f"No samples after reshaping for {time_covariate_vector}"
+            assert np.all(
+                samples > 0
+            ), f"Found non-positive values in samples for {time_covariate_vector}"
+            assert np.all(
+                samples.astype(int) == samples
+            ), f"Found non-integer values in samples for {time_covariate_vector}"
+
+            # Check mean and variance are reasonable
+            mean = np.mean(samples)
+            var = np.var(samples)
+            assert (
+                0 < mean < np.inf
+            ), f"Mean {mean} is not in valid range for {time_covariate_vector}"
+            assert (
+                0 < var < np.inf
+            ), f"Variance {var} is not in valid range for {time_covariate_vector}"
+
+            # Additional checks for distribution properties
+            # The mean should be greater than 1 for these parameters
+            assert mean > 1, f"Mean {mean} is not greater than 1 for {time_covariate_vector}"
+            # The variance should be positive and finite
+            assert var > 0, f"Variance {var} is not positive for {time_covariate_vector}"
+
+        except Exception as e:
+            import traceback
+
+            traceback.print_exc()
+            raise
 
     @pytest.mark.parametrize(
-        "r, alpha, time_covariates_sum, size, expected_shape",
+        "r, alpha, time_covariate_vector, size, expected_shape",
         [
-            (1.0, 1.0, 1.0, None, ()),  # Scalar output with covariates
-            ([1.0, 2.0], 1.0, 1.0, None, (2,)),  # Vector output from r
-            (1.0, [1.0, 2.0], 1.0, None, (2,)),  # Vector output from alpha
-            (1.0, 1.0, None, None, ()),  # No time covariates
+            (1.0, 1.0, None, None, ()),  # Scalar output with no covariates
+            ([1.0, 2.0], 1.0, [1.0], None, (2,)),  # Vector output from r
+            (1.0, [1.0, 2.0], [1.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
+            (1.0, 1.0, [1.0], (3, 2), (3, 2)),  # Explicit size
         ],
     )
-    def test_support_point(self, r, alpha, time_covariates_sum, size, expected_shape):
+    def test_support_point(self, r, alpha, time_covariate_vector, size, expected_shape):
         """Test that support_point returns reasonable values with correct shapes"""
         with pm.Model() as model:
             GrassiaIIGeometric(
-                "x", r=r, alpha=alpha, time_covariates_sum=time_covariates_sum, size=size
+                "x", r=r, alpha=alpha, time_covariate_vector=time_covariate_vector, size=size
             )
 
         init_point = model.initial_point()["x"]

From d0c1d980e26687cb36739083a5f89a9b6b453d3b Mon Sep 17 00:00:00 2001
From: ColtAllen <colt.allen1@gmail.com>
Date: Fri, 20 Jun 2025 13:58:48 -0600
Subject: [PATCH 13/13] delete test_simple.py

---
 test_simple.py | 61 --------------------------------------------------
 1 file changed, 61 deletions(-)
 delete mode 100644 test_simple.py

diff --git a/test_simple.py b/test_simple.py
deleted file mode 100644
index ada1a0680..000000000
--- a/test_simple.py
+++ /dev/null
@@ -1,61 +0,0 @@
-#!/usr/bin/env python3
-
-import pymc as pm
-import pytensor.tensor as pt
-
-from pymc_extras.distributions import GrassiaIIGeometric
-
-
-def test_basic_functionality():
-    """Test basic functionality of GrassiaIIGeometric distribution"""
-    print("Testing basic GrassiaIIGeometric functionality...")
-
-    # Test 1: Create distribution with None time_covariate_vector
-    try:
-        dist = GrassiaIIGeometric.dist(r=2.0, alpha=1.0, time_covariate_vector=None)
-        print("✓ Distribution created successfully with None time_covariate_vector")
-
-        # Test sampling
-        samples = dist.eval()
-        print(f"✓ Direct sampling successful: {samples}")
-
-    except Exception as e:
-        print(f"✗ Failed to create distribution with None time_covariate_vector: {e}")
-        return False
-
-    # Test 2: Create distribution with scalar time_covariate_vector
-    try:
-        dist = GrassiaIIGeometric.dist(r=2.0, alpha=1.0, time_covariate_vector=0.5)
-        print("✓ Distribution created successfully with scalar time_covariate_vector")
-
-        # Test sampling
-        samples = dist.eval()
-        print(f"✓ Direct sampling successful: {samples}")
-
-    except Exception as e:
-        print(f"✗ Failed to create distribution with scalar time_covariate_vector: {e}")
-        return False
-
-    # Test 3: Test logp function
-    try:
-        r = pt.scalar("r")
-        alpha = pt.scalar("alpha")
-        time_covariate_vector = pt.scalar("time_covariate_vector")
-        value = pt.scalar("value", dtype="int64")
-
-        logp = pm.logp(GrassiaIIGeometric.dist(r, alpha, time_covariate_vector), value)
-        logp_fn = pm.compile_fn([value, r, alpha, time_covariate_vector], logp)
-
-        result = logp_fn(2, 1.0, 1.0, 0.0)
-        print(f"✓ Logp function works: {result}")
-
-    except Exception as e:
-        print(f"✗ Failed to test logp function: {e}")
-        return False
-
-    print("✓ All basic functionality tests passed!")
-    return True
-
-
-if __name__ == "__main__":
-    test_basic_functionality()