Using “potentials” in Bambi to complicate a regression model

I have had my eye on a python package called Bambi for a while now, because I often need a regression model that is a little more complicated that sklearn.linear_model.LinearRegression but not complicated enough to make a whole new PyMC model.

Here is a minimal example (adapted from Awesome Open Source):

import bambi as bmb
import numpy as np, pandas as pd, arviz as az

data = pd.DataFrame({
    "y": np.random.normal(size=50),
    "x1": np.random.normal(size=50),
    "x2": np.random.normal(size=50)
})

model = bmb.Model("y ~ x1 + x2", data)
results = model.fit()
az.summary(results)

One cool thing about Bambi is that while it is simpler that writing a whole new PyMC model, it is a lot like writing a PyMC model. For example, if I need to add an informative prior, that is pretty easy:

priors = {'x1': bmb.Prior("Uniform", lower=0, upper=.05)}
model = bmb.Model("y ~ x1 + x2", data,
                 priors=priors)
results = model.fit()
az.summary(results)

And if I need a more complex distribution on that prior, Bambi exposes a “potential” parameter that puts additional terms in the posterior distribution, just like PyMC:

potentials = [
 (('x1', 'x2'),
 lambda x1, x2: bmb.math.switch((x1+x2 <= 0.0), 0, -999)),
]

model = bmb.Model("y ~ x1 + x2", data,
                 potentials=potentials)
results = model.fit()
az.summary(results)

I’m guessing that the syntax will continue to evolve, which is just one more reason Bambi is a python package that I am going to continue to watch.