Chapter 16: Bayesian statistics

Robert Johansson

Source code listings for Numerical Python - Scientific Computing and Data Science Applications with Numpy, SciPy and Matplotlib (ISBN 979-8-8688-0412-0).

import pymc as mc

import numpy as np

%matplotlib inline
%config InlineBackend.figure_format='retina'
import matplotlib.pyplot as plt

import seaborn as sns

sns.set()

from scipy import stats

import statsmodels.api as sm

import statsmodels.formula.api as smf

import pandas as pd

Changelog:¶

160828: The keyword argument vars to the functions mc.traceplot and mc.forestplot was changed to varnames.
231125: The keyword varnames was changed to var_names

Simple example: Normal distributed random variable¶

# try this
# dir(mc.distributions)

np.random.seed(100)

mu = 4.0

sigma = 2.0

model = mc.Model()

with model:
    mc.Normal("X", mu, tau=1 / sigma**2)

model.continuous_value_vars

start = dict(X=2)

with model:
    step = mc.Metropolis()
    trace = mc.sample(10000, step=step, initvals=start)

x = np.linspace(-4, 12, 1000)

y = stats.norm(mu, sigma).pdf(x)

# import arviz as az

# az.extract(trace)

def get_values(trace, variable):
    return trace.posterior.stack(sample=["chain", "draw"])[variable].values

X = get_values(trace, "X")

fig, ax = plt.subplots(figsize=(8, 3))

ax.plot(x, y, "r", lw=2)
sns.histplot(X, ax=ax, kde=True, stat="density")
ax.set_xlim(-4, 12)
ax.set_xlabel("x")
ax.set_ylabel("Probability distribution")
fig.tight_layout()
fig.savefig("ch16-normal-distribution-sampled.pdf")
fig.savefig("ch16-normal-distribution-sampled.png", dpi=600)

fig, axes = plt.subplots(1, 2, figsize=(12, 4), squeeze=False)
mc.plot_trace(trace, axes=axes)
axes[0, 0].plot(x, y, "r", lw=0.5)
fig.tight_layout()
fig.savefig("ch16-normal-sampling-trace.png", dpi=600)
fig.savefig("ch16-normal-sampling-trace.pdf")

Dependent random variables¶

model = mc.Model()

# mc.HalfNormal??

with model:
    mean = mc.Normal("mean", 3.0)
    sigma = mc.HalfNormal("sigma", sigma=1.0)
    X = mc.Normal("X", mean, tau=sigma)

model.continuous_value_vars

with model:
    start = mc.find_MAP()

start

with model:
    step = mc.Metropolis()
    trace = mc.sample(10000, initvals=start, step=step)

get_values(trace, "sigma").mean()

X = get_values(trace, "X")

X.mean()

X.std()

fig, axes = plt.subplots(3, 2, figsize=(8, 6), squeeze=False)
mc.plot_trace(trace, var_names=["mean", "sigma", "X"], axes=axes)
fig.tight_layout()
fig.savefig("ch16-dependent-rv-sample-trace.png", dpi=600)
fig.savefig("ch16-dependent-rv-sample-trace.pdf")

Posterior distributions¶

mu = 2.5

s = 1.5

data = stats.norm(mu, s).rvs(100)

with mc.Model() as model:
    mean = mc.Normal("mean", 4.0, tau=1.0)  # true 2.5
    sigma = mc.HalfNormal("sigma", tau=3.0 * np.sqrt(np.pi / 2))  # true 1.5

    X = mc.Normal("X", mean, tau=1 / sigma**2, observed=data)

model.continuous_value_vars

with model:
    start = mc.find_MAP()
    step = mc.Metropolis()
    trace = mc.sample(10000, initvals=start, step=step)
    # step = mc.NUTS()
    # trace = mc.sample(10000, initvals=start, step=step)

start

model.continuous_value_vars

fig, axes = plt.subplots(2, 2, figsize=(8, 4), squeeze=False)
mc.plot_trace(trace, var_names=["mean", "sigma"], axes=axes)
fig.tight_layout()
fig.savefig("ch16-posterior-sample-trace.png", dpi=600)
fig.savefig("ch16-posterior-sample-trace.pdf")

(
    mu,
    get_values(trace, "mean").mean(),
)  # trace.posterior.stack(sample=['chain', 'draw'])["mean"]

(
    s,
    get_values(trace, "sigma").mean(),
)  # trace.posterior.stack(sample=['chain', 'draw'])["sigma"].values.mean()

gs = mc.plot_forest(trace, var_names=["mean", "sigma"])
plt.savefig("ch16-forestplot.pdf")
plt.savefig("ch16-forestplot.png", dpi=600)

# help(mc.summary)

mc.summary(trace, var_names=["mean", "sigma"])

Linear regression¶

dataset = sm.datasets.get_rdataset("Davis", "carData", cache=True)

data = dataset.data[dataset.data.sex == "M"]

data = data[data.weight < 110]

data.head(3)

model = smf.ols("height ~ weight", data=data)

result = model.fit()

print(result.summary())

x = np.linspace(50, 110, 25)

y = result.predict({"weight": x})

fig, ax = plt.subplots(1, 1, figsize=(8, 3))
ax.plot(data.weight, data.height, "o")
ax.plot(x, y, color="blue")
ax.set_xlabel("weight")
ax.set_ylabel("height")
fig.tight_layout()
fig.savefig("ch16-linear-ols-fit.pdf")
fig.savefig("ch16-linear-ols-fit.png", dpi=600)

with mc.Model() as model:
    sigma = mc.Uniform("sigma", 0, 10)
    intercept = mc.Normal("intercept", 125, sigma=30)
    beta = mc.Normal("beta", 0, sigma=5)

    height_mu = intercept + beta * data.weight

    # likelihood function
    mc.Normal("height", mu=height_mu, sigma=sigma, observed=data.height)

    # predict
    predict_height = mc.Normal(
        "predict_height", mu=intercept + beta * x, sigma=sigma, shape=len(x)
    )

model.continuous_value_vars

# help(mc.NUTS)

with model:
    # start = mc.find_MAP()
    step = mc.NUTS()
    trace = mc.sample(10000, step=step)  # , initvals=start)

model.continuous_value_vars

fig, axes = plt.subplots(2, 2, figsize=(8, 4), squeeze=False)
mc.plot_trace(trace, var_names=["intercept", "beta"], axes=axes)
fig.tight_layout()
fig.savefig("ch16-linear-model-sample-trace.pdf")
fig.savefig("ch16-linear-model-sample-trace.png", dpi=600)

intercept = get_values(trace, "intercept").mean()
intercept

beta = get_values(
    trace, "beta"
).mean()  # trace.posterior.stack(sample=['chain', 'draw'])["beta"].values.mean()
beta

result.params

result.predict({"weight": 90})

weight_index = np.where(x == 90)[0][0]

# trace.posterior.stack(sample=['chain', 'draw'])["predict_height"].values[weight_index, :].mean()
get_values(trace, "predict_height")[weight_index, :].mean()

# trace.get_values("predict_height")[:, weight_index].mean()

# trace.posterior.stack(sample=['chain', 'draw'])["predict_height"].values.shape

fig, ax = plt.subplots(figsize=(8, 3))

sns.histplot(
    get_values(trace, "predict_height")[weight_index, :],
    ax=ax,
    kde=True,
    stat="density",
)
ax.set_xlim(150, 210)
ax.set_xlabel("height")
ax.set_ylabel("Probability distribution")
fig.tight_layout()
fig.savefig("ch16-linear-model-predict-cut.pdf")
fig.savefig("ch16-linear-model-predict-cut.png", dpi=600)

# trace.posterior.stack(sample=['chain', 'draw'])["intercept"].values

fig, ax = plt.subplots(1, 1, figsize=(12, 4))

for n in range(500, 2000, 1):
    intercept = get_values(trace, "intercept")[n]
    beta = get_values(trace, "beta")[n]
    ax.plot(x, intercept + beta * x, color="red", lw=0.25, alpha=0.05)

intercept = get_values(trace, "intercept").mean()
beta = get_values(trace, "beta").mean()
ax.plot(x, intercept + beta * x, color="k", label="Mean Bayesian prediction")

ax.plot(data.weight, data.height, "o")
ax.plot(x, y, "--", color="blue", label="OLS prediction")
ax.set_xlabel("weight")
ax.set_ylabel("height")
ax.legend(loc=0)
fig.tight_layout()
fig.savefig("ch16-linear-model-fit.pdf")
fig.savefig("ch16-linear-model-fit.png", dpi=600)

import bambi

model = bambi.Model("height ~ weight", data)

trace = model.fit(draws=2000)  #  chains=4

fig, axes = plt.subplots(3, 2, figsize=(8, 6), squeeze=False)
mc.plot_trace(trace, var_names=["Intercept", "weight", "height_sigma"], axes=axes)
fig.tight_layout()
fig.savefig("ch16-glm-sample-trace.pdf")
fig.savefig("ch16-glm-sample-trace.png", dpi=600)

Multilevel model¶

dataset = sm.datasets.get_rdataset("Davis", "carData", cache=True)

data = dataset.data.copy()
data = data[data.weight < 110]

data["sex"] = data["sex"].apply(lambda x: 1 if x == "F" else 0)

data.head()

with mc.Model() as model:
    # heirarchical model: hyper priors
    # intercept_mu = mc.Normal("intercept_mu", 125)
    # intercept_sigma = 30.0 #mc.Uniform('intercept_sigma', lower=0, upper=50)
    # beta_mu = mc.Normal("beta_mu", 0.0)
    # beta_sigma = 5.0 #mc.Uniform('beta_sigma', lower=0, upper=10)

    # multilevel model: prior parameters
    intercept_mu, intercept_sigma = 125, 30
    beta_mu, beta_sigma = 0.0, 5.0

    # priors
    intercept = mc.Normal("intercept", intercept_mu, sigma=intercept_sigma, shape=2)
    beta = mc.Normal("beta", beta_mu, sigma=beta_sigma, shape=2)
    error = mc.Uniform("error", 0, 10)

    # model equation
    sex_idx = data.sex.values
    height_mu = intercept[sex_idx] + beta[sex_idx] * data.weight

    mc.Normal("height", mu=height_mu, sigma=error, observed=data.height)

model.continuous_value_vars

with model:
    start = mc.find_MAP()
    # hessian = mc.find_hessian(start)
    step = mc.NUTS()
    trace = mc.sample(5000, step=step, initvals=start)

fig, axes = plt.subplots(3, 2, figsize=(8, 6), squeeze=False)
mc.plot_trace(trace, var_names=["intercept", "beta", "error"], axes=axes)
fig.tight_layout()
fig.savefig("ch16-multilevel-sample-trace.pdf")
fig.savefig("ch16-multilevel-sample-trace.png", dpi=600)

intercept_m, intercept_f = get_values(trace, "intercept").mean(axis=1)

intercept = get_values(trace, "intercept").mean()

get_values(trace, "beta")

beta_m, beta_f = get_values(trace, "beta").mean(axis=1)

beta = get_values(trace, "beta").mean()

fig, ax = plt.subplots(1, 1, figsize=(8, 3))

mask_m = data.sex == 0
mask_f = data.sex == 1

ax.plot(
    data.weight[mask_m],
    data.height[mask_m],
    "o",
    color="steelblue",
    label="male",
    alpha=0.5,
)
ax.plot(
    data.weight[mask_f],
    data.height[mask_f],
    "o",
    color="green",
    label="female",
    alpha=0.5,
)

x = np.linspace(35, 110, 50)
ax.plot(x, intercept_m + x * beta_m, color="steelblue", label="model male group")
ax.plot(x, intercept_f + x * beta_f, color="green", label="model female group")
ax.plot(x, intercept + x * beta, color="black", label="model both groups")

ax.set_xlabel("weight")
ax.set_ylabel("height")
ax.legend(loc=0)
fig.tight_layout()
fig.savefig("ch16-multilevel-linear-model-fit.pdf")
fig.savefig("ch16-multilevel-linear-model-fit.png", dpi=600)

get_values(trace, "error").mean()

References¶

Johansson, R. (2024). Numerical Python: Scientific Computing and Data Science Applications with Numpy, SciPy and Matplotlib. Apress. 10.1007/979-8-8688-0413-7