Hi! I’m working on a dataset where observations are best described as zero-inflated, Poisson distributed, with an upper bound (that’s a relatively low number). I was hoping that I could simply apply pm.Truncated
on pm.ZeroInflatedPoisson
, along the lines of:
zeroinfl_dist = pm.ZeroInflatedPoisson.dist()
pm.Truncated("truncated_zipoisson", zeroinfl_dist, lower=0, upper=my_upper_limit)
However, it seems that truncation is not implemented for this distribution (I get an error NotImplementedError: Truncation not implemented for SymbolicRandomVariable MarginalMixtureRV{inline=True}
). What would be the best way to go around this limitation? Any hints would be much appreciated!
I think I figured it out myself - the ZeroInflatedPoisson distribution code was pretty easy to modify:
from pymc.distributions.discrete import _zero_inflated_mixture, Poisson
class ZeroInflatedPoissonTrunc:
def __new__(cls, name, psi, mu, upper_limit, **kwargs):
pois = Poisson.dist(mu=mu)
trunc_pois = pm.Truncated.dist(pois, lower=0, upper=upper_limit)
return _zero_inflated_mixture(
name=name, nonzero_p=psi, nonzero_dist=trunc_pois, **kwargs
)
@classmethod
def dist(cls, psi, mu, upper_limit, **kwargs):
pois = Poisson.dist(mu=mu)
trunc_pois = pm.Truncated.dist(pois, lower=0, upper=upper_limit)
return _zero_inflated_mixture(
name=None, nonzero_p=psi, nonzero_dist=trunc_pois, **kwargs
)
Works like a charm!
1 Like
That’s a neat solution, you can do it without the class as well
truncated_poisson = pm.Truncated.dist(pm.Poisson.dist(), lower=0, upper=upper_limit)
zero_inflation = pm.DiracDelta(0)
truncated_zip = pm.Mixture("truncated_zip", w=..., comp_dists([zero_inflation, truncated_poisson]))
3 Likes