WIP tests

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -433,6 +433,11 @@ def rng_fn(cls, rng, r, alpha, time_covariate_vector, size):
         lam_covar = lam * exp_time_covar
 
         p = 1 - np.exp(-lam_covar)
+        # TODO: This is a hack to ensure valid probability in (0, 1]
+        # We should find a better way to do this.
+        # Ensure valid probability in (0, 1]
+        tiny = np.finfo(p.dtype).tiny
+        p = np.clip(p, tiny, 1.0)
         samples = rng.geometric(p)
         # samples = np.ceil(np.log(1 - rng.uniform(size=size)) / (-lam_covar))
 
@@ -576,12 +581,11 @@ def support_point(rv, size, r, alpha, time_covariate_vector):
             1.0 / (1.0 - pt.exp(-base_lambda)),  # Full expression for larger lambda
         )
 
-        # Apply time covariates if provided
+        # Apply time covariates if provided: multiply by exp(sum over axis=0)
+        # This yields a scalar for 1D covariates and a time-length vector for 2D (features x time)
         tcv = pt.as_tensor_variable(time_covariate_vector)
         if tcv.ndim != 0:
-            # If 1D, treat as per-time vector; if 2D+, sum features while preserving time axis
-            cov_time = tcv if tcv.ndim == 1 else tcv.sum(axis=0)
-            mean = mean * pt.exp(cov_time)
+            mean = mean * pt.exp(tcv.sum(axis=0))
 
         # Round up to nearest integer and ensure >= 1
         mean = pt.maximum(pt.ceil(mean), 1.0)
@@ -603,8 +607,8 @@ def C_t(t: pt.TensorVariable, time_covariate_vector: pt.TensorVariable) -> pt.Te
     if time_covariate_vector.ndim == 1:
         per_time_sum = pt.exp(time_covariate_vector)
     else:
-        feature_axes = tuple(range(time_covariate_vector.ndim - 1))
-        per_time_sum = pt.sum(pt.exp(time_covariate_vector), axis=feature_axes)
+        # If axis=0 is time and axis>0 are features, sum over features (axis>0)
+        per_time_sum = pt.sum(pt.exp(time_covariate_vector), axis=0)
 
     # Build cumulative sum up to each t without advanced indexing
     time_length = pt.shape(per_time_sum)[0]
@@ -617,9 +621,5 @@ def C_t(t: pt.TensorVariable, time_covariate_vector: pt.TensorVariable) -> pt.Te
     mask = pt.lt(time_idx, pt.shape_padright(t_vec, 1))
     # Sum per-time contributions over time axis
     base_sum = pt.sum(pt.shape_padleft(per_time_sum) * mask, axis=-1)
-    # Carry-forward last per-time value for t beyond time_length
-    last_value = per_time_sum[-1]
-    excess_steps = pt.maximum(t_vec - time_length, 0)
-    carried = base_sum + excess_steps * last_value
-    # If original t was scalar, return scalar
-    return pt.squeeze(carried)
+    # If original t was scalar, return scalar (saturate at last time step)
+    return pt.squeeze(base_sum)