<jupyter_start><jupyter_text>GLM: Mini-batch ADVI on hierarchical regression model:::{post} Sept 23, 2021:tags: generalized linear model, hierarchical model, pymc3.Minibatch, pymc3.Model, pymc3.NUTS, pymc3.Normal, pymc3.Uniform, variational inference:category: intermediate::: Unlike Gaussian mixture models, (hierarchical) regression models have independent variables. These variables affect the likelihood function, but are not random variables. When using mini-batch, we should take care of that.<jupyter_code>%env THEANO_FLAGS=device=cpu, floatX=float32, warn_float64=ignore

import os

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as pm
import seaborn as sns
import theano
import theano.tensor as tt

from scipy import stats

print(f"Running on PyMC3 v{pm.__version__}")
%config InlineBackend.figure_format = 'retina'
RANDOM_SEED = 8927
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-darkgrid")
try:
    data = pd.read_csv(os.path.join("..", "data", "radon.csv"))
except FileNotFoundError:
    data = pd.read_csv(pm.get_data("radon.csv"))

data
county_idx = data["county_code"].values
floor_idx = data["floor"].values
log_radon_idx = data["log_radon"].values

coords = {"counties": data.county.unique()}<jupyter_output><empty_output><jupyter_text>Here, `log_radon_idx_t` is a dependent variable, while `floor_idx_t` and `county_idx_t` determine the independent variable.<jupyter_code>log_radon_idx_t = pm.Minibatch(log_radon_idx, 100)
floor_idx_t = pm.Minibatch(floor_idx, 100)
county_idx_t = pm.Minibatch(county_idx, 100)
with pm.Model(coords=coords) as hierarchical_model:
    # Hyperpriors for group nodes
    mu_a = pm.Normal("mu_alpha", mu=0.0, sigma=100 ** 2)
    sigma_a = pm.Uniform("sigma_alpha", lower=0, upper=100)
    mu_b = pm.Normal("mu_beta", mu=0.0, sigma=100 ** 2)
    sigma_b = pm.Uniform("sigma_beta", lower=0, upper=100)<jupyter_output><empty_output><jupyter_text>Intercept for each county, distributed around group mean `mu_a`. Above we just set `mu` and `sd` to a fixed value while here we plug in a common group distribution for all `a` and `b` (which are vectors with the same length as the number of unique counties in our example).<jupyter_code>with hierarchical_model:

    a = pm.Normal("alpha", mu=mu_a, sigma=sigma_a, dims="counties")
    # Intercept for each county, distributed around group mean mu_a
    b = pm.Normal("beta", mu=mu_b, sigma=sigma_b, dims="counties")<jupyter_output><empty_output><jupyter_text>Model prediction of radon level `a[county_idx]` translates to `a[0, 0, 0, 1, 1, ...]`, we thus link multiple household measures of a county to its coefficients.<jupyter_code>with hierarchical_model:

    radon_est = a[county_idx_t] + b[county_idx_t] * floor_idx_t<jupyter_output><empty_output><jupyter_text>Finally, we specify the likelihood:<jupyter_code>with hierarchical_model:

    # Model error
    eps = pm.Uniform("eps", lower=0, upper=100)

    # Data likelihood
    radon_like = pm.Normal(
        "radon_like", mu=radon_est, sigma=eps, observed=log_radon_idx_t, total_size=len(data)
    )<jupyter_output><empty_output><jupyter_text>Random variables `radon_like`, associated with `log_radon_idx_t`, should be given to the function for ADVI to denote that as observations in the likelihood term.  On the other hand, `minibatches` should include the three variables above.  Then, run ADVI with mini-batch.<jupyter_code>with hierarchical_model:
    approx = pm.fit(100000, callbacks=[pm.callbacks.CheckParametersConvergence(tolerance=1e-4)])
    idata_advi = az.from_pymc3(approx.sample(500))<jupyter_output>/home/ada/.local/lib/python3.8/site-packages/theano/gpuarray/dnn.py:192: UserWarning: Your cuDNN version is more recent than Theano. If you encounter problems, try updating Theano or downgrading cuDNN to a version >= v5 and <= v7.
  warnings.warn(<jupyter_text>Check the trace of ELBO and compare the result with MCMC.<jupyter_code>plt.plot(approx.hist);
# Inference button (TM)!
with pm.Model(coords=coords):

    mu_a = pm.Normal("mu_alpha", mu=0.0, sigma=100 ** 2)
    sigma_a = pm.Uniform("sigma_alpha", lower=0, upper=100)
    mu_b = pm.Normal("mu_beta", mu=0.0, sigma=100 ** 2)
    sigma_b = pm.Uniform("sigma_beta", lower=0, upper=100)

    a = pm.Normal("alpha", mu=mu_a, sigma=sigma_a, dims="counties")
    b = pm.Normal("beta", mu=mu_b, sigma=sigma_b, dims="counties")

    # Model error
    eps = pm.Uniform("eps", lower=0, upper=100)

    radon_est = a[county_idx] + b[county_idx] * floor_idx

    radon_like = pm.Normal("radon_like", mu=radon_est, sigma=eps, observed=log_radon_idx)

    # essentially, this is what init='advi' does
    step = pm.NUTS(scaling=approx.cov.eval(), is_cov=True)
    hierarchical_trace = pm.sample(
        2000, step, start=approx.sample()[0], progressbar=True, return_inferencedata=True
    )
az.plot_density(
    [idata_advi, hierarchical_trace], var_names=["~alpha", "~beta"], data_labels=["ADVI", "NUTS"]
);<jupyter_output><empty_output><jupyter_text>Watermark<jupyter_code>%load_ext watermark
%watermark -n -u -v -iv -w -p xarray<jupyter_output>Last updated: Thu Sep 23 2021

Python implementation: CPython
Python version       : 3.8.10
IPython version      : 7.25.0

xarray: 0.17.0

numpy     : 1.21.0
seaborn   : 0.11.1
pandas    : 1.2.1
matplotlib: 3.3.4
theano    : 1.1.2
pymc3     : 3.11.2
scipy     : 1.7.1
arviz     : 0.11.2

Watermark: 2.2.0
{'fraction_non_alphanumeric': 0.12060394034247837, 'fraction_numerical': 0.028908120051555884, 'mean_word_length': 5.151755379388448, 'pattern_counts': {'":': 1, '<': 32, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '463130', 'license': 'MIT', 'n_tokens_mistral': 2084, 'n_tokens_neox': 1944, 'n_words': 484, 'path': 'examples/variational_inference/GLM-hierarchical-advi-minibatch.ipynb', 'repo_name': 'mjhajharia/pymc-examples', 'size': '5431'}