@@ -668,7 +668,7 @@ def NegBinom(a, m, x):
668
668
669
669
@classmethod
670
670
def dist (cls , mu = None , alpha = None , p = None , n = None , * args , ** kwargs ):
671
- n , p = cls .get_n_p (mu , alpha , p , n )
671
+ n , p = cls .get_n_p (mu = mu , alpha = alpha , p = p , n = n )
672
672
n = at .as_tensor_variable (floatX (n ))
673
673
p = at .as_tensor_variable (floatX (p ))
674
674
return super ().dist ([n , p ], * args , ** kwargs )
@@ -1482,6 +1482,21 @@ def logcdf(value, psi, n, p):
1482
1482
)
1483
1483
1484
1484
1485
+ class ZeroInflatedNegBinomialRV (RandomVariable ):
1486
+ name = "zero_inflated_neg_binomial"
1487
+ ndim_supp = 0
1488
+ ndims_params = [0 , 0 , 0 ]
1489
+ dtype = "int64"
1490
+ _print_name = ("ZeroInflatedNegBinom" , "\\ operatorname{ZeroInflatedNegBinom}" )
1491
+
1492
+ @classmethod
1493
+ def rng_fn (cls , rng , psi , n , p , size ):
1494
+ return rng .negative_binomial (n = n , p = p , size = size ) * (rng .random (size = size ) < psi )
1495
+
1496
+
1497
+ zero_inflated_neg_binomial = ZeroInflatedNegBinomialRV ()
1498
+
1499
+
1485
1500
class ZeroInflatedNegativeBinomial (Discrete ):
1486
1501
R"""
1487
1502
Zero-Inflated Negative binomial log-likelihood.
@@ -1551,50 +1566,17 @@ def ZeroInfNegBinom(a, m, psi, x):
1551
1566
1552
1567
"""
1553
1568
1554
- def __init__ (self , psi , mu , alpha , * args , ** kwargs ):
1555
- super ().__init__ (* args , ** kwargs )
1556
- self .mu = mu = at .as_tensor_variable (floatX (mu ))
1557
- self .alpha = alpha = at .as_tensor_variable (floatX (alpha ))
1558
- self .psi = psi = at .as_tensor_variable (floatX (psi ))
1559
- self .nb = NegativeBinomial .dist (mu , alpha )
1560
- self .mode = self .nb .mode
1569
+ rv_op = zero_inflated_neg_binomial
1561
1570
1562
- def random (self , point = None , size = None ):
1563
- r"""
1564
- Draw random values from ZeroInflatedNegativeBinomial distribution.
1565
-
1566
- Parameters
1567
- ----------
1568
- point: dict, optional
1569
- Dict of variable values on which random values are to be
1570
- conditioned (uses default point if not specified).
1571
- size: int, optional
1572
- Desired size of random sample (returns one sample if not
1573
- specified).
1574
-
1575
- Returns
1576
- -------
1577
- array
1578
- """
1579
- # mu, alpha, psi = draw_values([self.mu, self.alpha, self.psi], point=point, size=size)
1580
- # g = generate_samples(self._random, mu=mu, alpha=alpha, dist_shape=self.shape, size=size)
1581
- # g[g == 0] = np.finfo(float).eps # Just in case
1582
- # g, psi = broadcast_distribution_samples([g, psi], size=size)
1583
- # return stats.poisson.rvs(g) * (np.random.random(g.shape) < psi)
1584
-
1585
- def _random (self , mu , alpha , size ):
1586
- r"""Wrapper around stats.gamma.rvs that converts NegativeBinomial's
1587
- parametrization to scipy.gamma. All parameter arrays should have
1588
- been broadcasted properly by generate_samples at this point and size is
1589
- the scipy.rvs representation.
1590
- """
1591
- return stats .gamma .rvs (
1592
- a = alpha ,
1593
- scale = mu / alpha ,
1594
- size = size ,
1595
- )
1571
+ @classmethod
1572
+ def dist (cls , psi , mu , alpha , * args , ** kwargs ):
1573
+ psi = at .as_tensor_variable (floatX (psi ))
1574
+ n , p = NegativeBinomial .get_n_p (mu = mu , alpha = alpha )
1575
+ n = at .as_tensor_variable (floatX (n ))
1576
+ p = at .as_tensor_variable (floatX (p ))
1577
+ return super ().dist ([psi , n , p ], * args , ** kwargs )
1596
1578
1597
- def logp (self , value ):
1579
+ def logp (value , psi , n , p ):
1598
1580
r"""
1599
1581
Calculate log-probability of ZeroInflatedNegativeBinomial distribution at specified value.
1600
1582
@@ -1608,20 +1590,22 @@ def logp(self, value):
1608
1590
-------
1609
1591
TensorVariable
1610
1592
"""
1611
- alpha = self .alpha
1612
- mu = self .mu
1613
- psi = self .psi
1614
1593
1615
- logp_other = at .log (psi ) + self .nb .logp (value )
1616
- logp_0 = logaddexp (
1617
- at .log1p (- psi ), at .log (psi ) + alpha * (at .log (alpha ) - at .log (alpha + mu ))
1594
+ return bound (
1595
+ at .switch (
1596
+ at .gt (value , 0 ),
1597
+ at .log (psi ) + NegativeBinomial .logp (value , n , p ),
1598
+ logaddexp (at .log1p (- psi ), at .log (psi ) + n * at .log (p )),
1599
+ ),
1600
+ 0 <= value ,
1601
+ 0 <= psi ,
1602
+ psi <= 1 ,
1603
+ 0 < n ,
1604
+ 0 <= p ,
1605
+ p <= 1 ,
1618
1606
)
1619
1607
1620
- logp_val = at .switch (at .gt (value , 0 ), logp_other , logp_0 )
1621
-
1622
- return bound (logp_val , 0 <= value , 0 <= psi , psi <= 1 , mu > 0 , alpha > 0 )
1623
-
1624
- def logcdf (self , value ):
1608
+ def logcdf (value , psi , n , p ):
1625
1609
"""
1626
1610
Compute the log of the cumulative distribution function for ZeroInflatedNegativeBinomial distribution
1627
1611
at the specified value.
@@ -1640,13 +1624,14 @@ def logcdf(self, value):
1640
1624
raise TypeError (
1641
1625
f"ZeroInflatedNegativeBinomial.logcdf expects a scalar value but received a { np .ndim (value )} -dimensional object."
1642
1626
)
1643
- psi = self .psi
1644
1627
1645
1628
return bound (
1646
- logaddexp (at .log1p (- psi ), at .log (psi ) + self . nb . logcdf (value )),
1629
+ logaddexp (at .log1p (- psi ), at .log (psi ) + NegativeBinomial . logcdf (value , n , p )),
1647
1630
0 <= value ,
1648
1631
0 <= psi ,
1649
1632
psi <= 1 ,
1633
+ 0 < p ,
1634
+ p <= 1 ,
1650
1635
)
1651
1636
1652
1637
0 commit comments