Weights for Model in Bambi

Questions
7

Thanks for the example!

I see now what you want (i.e. to weigh the contribution of each observation to the likelihood). I
can think of a complex solution that I think solves your problem but it’s not scalable to other likelihoods.

This solution uses a custom family with a custom distribution which is created with DensityDist, that’s why it’s very complex.

import arviz as az
import bambi as bmb
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt

from bambi.families.univariate import UnivariateFamily

rng = np.random.default_rng(1234)
data = pd.DataFrame(
    {
        "score": np.concatenate([rng.normal(size=1000), rng.normal(loc=1, size=1000)]),
        "weights": np.repeat([1, 2], 1000)
    }
)

# Create custom "distribution" with DensityDist
def logp(value, mu, sigma, weights):
    # This is the same as the logp of a normal, but we have the weights multiplying
    res = weights * (-0.5 * pt.pow((value - mu) / sigma, 2) - pt.log(pt.sqrt(2.0 * np.pi)) - pt.log(sigma))
    return res

def WeightedNormal(name, mu, sigma, weights, *args, **kwargs):
    return pm.DensityDist(name, mu, sigma, weights, *args, **kwargs, logp=logp)

# Create custom family. We need to use the weights in a different way, so we need 'transform_backend_kwargs'.
class WeightedNormalFamily(UnivariateFamily):
    SUPPORTED_LINKS = {"mu": ["identity", "log"], "sigma": ["log"]}
    @staticmethod
    def transform_backend_kwargs(kwargs):
        observed = kwargs["observed"]
        kwargs["observed"] = observed[:, 0]
        kwargs["weights"] = observed[:, 1]
        return kwargs

# Create the custom family instance
likelihood = bmb.Likelihood("WeightedNormal", params=["mu", "sigma"], parent="mu", dist=WeightedNormal)
links = {"mu": "identity", "sigma": "log"}
wn_family = WeightedNormalFamily("weighted-normal", likelihood, links)

# Create the model in Bambi using the custom family
priors = {
    "mu": bmb.Prior("Normal", mu=0, sigma=1),
    "sigma": bmb.Prior("HalfNormal", sigma=1),
}
model = bmb.Model("c(score, weights) ~ 1", data, family=wn_family, priors=priors)
model.build()
model.graph()

image

idata = model.fit()
az.summary(idata)

image
image934×192 17.1 KB

So, this seems to work but it requires to implement a bunch of custom code, which depends on the likelihood. This should be way less painful when we implement it in Bambi.

I imagine something like

bmb.Model("weighted(y, weights) ~ ...")

and that should be it… But for now it’s an expression of desire

1 Like