Custom Zero-Inflated Binomial in Bambi

Questions version agnostic
1

@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
2

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):
    SUPPORTED_LINKS = {
        "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 = ZeroInflatedBinomial("zinb", likelihood, links)
zinb_family

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

1 Like
3

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

output
output965×603 194 KB

2 Likes
4

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

Z.

5

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):
    SUPPORTED_LINKS = {
        "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"}
zinb_family = ZeroInflatedBinomial("zinb", likelihood, links)
zinb_family

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"].

But again, if you don’t really need this, don’t do it because it’s kind of a Frankenstein and only works because I know how you can inject things in the right places to make things work. This will be much better when we support it natively.

1 Like
6

Thanks. Although it is convoluted, it’s pretty cool that you can use it it in this setup.

1 Like