Custom Zero-Inflated Binomial in Bambi

@tcapretto I tried to setup a custom ZINB distribution. I got 90% there but I don’t know how to set “n” (trials).

I get the following error: "TypeError: ZeroInflatedBinomial.new() missing 1 required positional argument: ‘n’

My test notebook is here.

2 Likes

Turns out this is a little more complex than I anticipated because one of the parameters in the distribution is usually an observed variable, the `n`.

This is like the `Binomial` family. See the implementation here bambi/univariate.py at 9ba92e1b8fa2833370468de46d55bffac3fda101 · bambinos/bambi · GitHub where we need to implement a custom `posterior_predictive` and `transform_backend_kwargs` methods.

This will make your example work

``````from bambi.families.univariate import UnivariateFamily

class ZeroInflatedBinomial(UnivariateFamily):
"p": ["identity", "logit", "probit", "cloglog"],
"psi": ["logit", "probit", "cloglog"]
}

@staticmethod
def transform_backend_kwargs(kwargs):
observed = kwargs.pop("observed")
kwargs["observed"] = observed[:, 0].squeeze()
kwargs["n"] = observed[:, 1].squeeze()
return kwargs

likelihood = bmb.Likelihood("ZeroInflatedBinomial", params=["p", "psi"], parent="p")
links = {"p": "logit", "psi": "logit"}
zinb_family
``````

Notice I’m not implementing the `posterior_predictive` method yet. This would take more work.

1 Like

It works, thanks. The true values were found easily. I appreciate it.

2 Likes

@tcapretto How would one use pm.Censored() with this setup? Or is it not possible?

Z.

You could if you created also a custom likelihood function. But only if you really needed it. This will be easier once we support censored responses natively in Bambi (which should happen soon).

See the example

``````from functools import partial

from bambi.families.univariate import UnivariateFamily

def CensoredZeroInflatedBinomial(name, psi, n, p, lower, upper, observed):
dist = pm.ZeroInflatedBinomial.dist(psi=psi, n=n, p=p)
return pm.Censored(name, dist, lower=lower, upper=upper, observed=observed)

dist_fn = partial(CensoredZeroInflatedBinomial, lower=0, upper=10)

class ZeroInflatedBinomial(UnivariateFamily):
"p": ["identity", "logit", "probit", "cloglog"],
"psi": ["logit", "probit", "cloglog"]
}

@staticmethod
def transform_backend_kwargs(kwargs):
observed = kwargs.pop("observed")
kwargs["observed"] = observed[:, 0].squeeze()
kwargs["n"] = observed[:, 1].squeeze()
return kwargs

likelihood = bmb.Likelihood("CensoredZeroInflatedBinomial", params=["p", "psi"], parent="p", dist=dist_fn)
links = {"p": "logit", "psi": "logit"}
See the usage of the custom likelihood function, which is passed to the `dist` argument in `bmb.Likelihood`. The limitation is that lower and upper can’t be taken from the dataset, they have to be fixed to some values. If you want to make them equal to arrays instead of scalars, you can try fixing them to `data["lower"]` and `data["upper"]`.