@@ -404,7 +404,7 @@ def dist(cls, mu1, mu2, **kwargs):
404
404
405
405
class GrassiaIIGeometricRV (RandomVariable ):
406
406
name = "g2g"
407
- signature = "(),(),()->()"
407
+ signature = "(),(),(t )->()"
408
408
409
409
dtype = "int64"
410
410
_print_name = ("GrassiaIIGeometric" , "\\ operatorname{GrassiaIIGeometric}" )
@@ -422,15 +422,13 @@ def rng_fn(cls, rng, r, alpha, time_covariate_vector, size):
422
422
423
423
lam = rng .gamma (shape = r , scale = 1 / alpha , size = size )
424
424
425
- # Calculate exp(time_covariate_vector) for all samples
425
+ # Aggregate time covariates for each sample
426
426
exp_time_covar = np .exp (
427
- time_covariate_vector
428
- ). mean () # Approximation required to return a t-scalar from a covariate vector
427
+ time_covariate_vector . sum ( axis = 0 )
428
+ ) # TODO: try np.exp(time_covariate_vector).sum(axis=0) instead?
429
429
lam_covar = lam * exp_time_covar
430
430
431
- # Take uniform draws from the inverse CDF
432
- u = rng .uniform (size = size )
433
- samples = np .ceil (np .log (1 - u ) / (- lam_covar ))
431
+ samples = np .ceil (np .log (1 - rng .uniform (size = size )) / (- lam_covar ))
434
432
435
433
return samples
436
434
@@ -536,6 +534,7 @@ def logcdf(value, r, alpha, time_covariate_vector):
536
534
logcdf ,
537
535
r > 0 ,
538
536
alpha > 0 ,
537
+ time_covariate_vector >= 0 ,
539
538
msg = "r > 0, alpha > 0" ,
540
539
)
541
540
@@ -573,6 +572,7 @@ def support_point(rv, size, r, alpha, time_covariate_vector):
573
572
return mean
574
573
575
574
575
+ # TODO: can this be moved into logp? Indexing not required for logcdf
576
576
def C_t (t : pt .TensorVariable , time_covariate_vector : pt .TensorVariable ) -> pt .TensorVariable :
577
577
"""Utility for processing time-varying covariates in GrassiaIIGeometric distribution."""
578
578
if time_covariate_vector .ndim == 0 :
0 commit comments