Weights for Model in Bambi

Questions
11

I think I made it work. See I’m using pm.CustomDist instead of pm.DensityDist. I had a look at the PyMC codebase and it looks like DensityDist will be deprecated at some point and CustomDist is the new thing.

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)
    }
)

def logp(value, mu, sigma, weights):
    # This is the same as the logp of a normal, but we have the weights multiplying
    return weights * pm.Normal.logp(value, mu, sigma)

def random(mu, sigma, rng=None, size=None):
    # Weights don't make sense when generating new observations. 
    # They are a property of actual observations
    return rng.normal(loc=mu, scale=sigma, size=size)

class WeightedNormal:
    def __new__(self, name, mu, sigma, weights, **kwargs):
        return pm.CustomDist(name, mu, sigma, weights, logp=logp, **kwargs)
    
    @classmethod
    def dist(cls, mu, sigma, **kwargs):
        return pm.CustomDist.dist(
            mu, sigma, class_name="WeightedNormal", random=random, **kwargs
        )

# 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):
        if "observed" in 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()

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

idata_2 = model.predict(idata, data = data, kind="pps", inplace = False)

See it’s a lot of code. Hopefully, we can have this inside Bambi in the future.