Truncating a Zero Inflated Poisson distribution

Questions version agnostic
1

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!

2

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!

3

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]))
2 Likes