WIP time indexing

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -423,40 +423,40 @@ 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
             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_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,

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"]