Multi-level models + pymc3.glm

Question

I need to fit a multi-level linear model using PyMC3 and I really like the glm api, because of the conciseness it provides. I would like to ask if and how this can be done. This blog post I found mentions that:

glm() does not play nice with hierarchical models yet

so I am a bit sceptical this can actually be done, but it was written a couple years ago, so this might have changed. To give an example, the following is the model I would like to rewrite using the glm api

import numpy as np
import pymc3 as pm


def generate_data():
    n, beta_0, beta_1, sd_eps = 100, 1.2, 0.6, 0.2
    b_group = np.array([0.05, 0.14, -0.23])
    x = np.random.randn(n)
    group_index = np.random.choice([0, 1, 2], n)
    y = 1.2 + 0.6 * x + sd_eps * np.random.randn(n) + b_group[group_index]
    return x, group_index, y


if __name__ == '__main__':
    x, group_index, y = generate_data()

    with pm.Model() as multi_level_model:
        sd_b_group = pm.HalfNormal("sd_b_group", sigma=100)

        b_group = pm.Normal("b_group", mu=0, sigma=sd_b_group, shape=3)

        beta_0 = pm.Normal("beta_0", mu=0, sigma=100)
        beta_1 = pm.Normal("beta_1", mu=0, sigma=100)

        sd_eps = pm.HalfNormal("sd_eps", sigma=100)
        pm.Normal("y", beta_0 + beta_1 * x + b_group[group_index],
                  sigma=sd_eps, observed=y)

I assume it would look something like this:

with pm.Model():
    mu_b_group = pm.Normal("mu_b_group", mu=0, sigma=100)
    sd_b_group = pm.HalfNormal("sd_b_group", sigma=100)

    b_group = pm.Normal("b_group", mu=mu_b_group, sigma=sd_b_group, shape=3)

    pm.glm.GLM.from_formula(formula="y ~ 1 + x",
                            vars={"Intercept": b_group[group_index]},
                            data={"y": y, "x": x})

which however produces an error here when trying to stack the coefficients

TypeError: Join() can only join tensors with the same number of dimensions.

This can be done with Bambi https://github.com/bambinos/bambi. You may want to have a look that the hierarchical model example here https://bambinos.github.io/bambi/notebooks/multi-level_regression.html — Tomas Capretto, Apr 09 '21 at 14:25

Josef Ondrej · Answer 1 · 2021-02-03T09:32:56.037

After some experimenting I came up with this (not ideal, but at least a bit useful) solution:

with pm.Model():
    sd_b_group = pm.HalfNormal("sd_b_group", sigma=100)
    b_group = pm.Normal("b_group", mu=0, sigma=sd_b_group, shape=3)

    lm = pm.glm.LinearComponent.from_formula(formula="y ~ 1 + x",
                                             data={"y": y, "x": x})

    sd_eps = pm.HalfNormal("sd_eps", sigma=100)

    likelihood = pm.Normal("y", mu=lm.y_est + b_group[group_index], 
                           sigma=sd_eps, observed=y)

Also created an example on github for future reference.

Multi-level models + pymc3.glm

1 Answers1