Last few questions on this topic if I may (here with actual regression data)!
I am curious about a few things:
-
Is stacking our InferenceData items recommended and if so, is this a recommended way? .posterior.stack(draws=("chain", "draw")) e.g., chain, draw or draw, chain?
-
I am curious why my expected value of y is not plotting a straight line akin to official example. The results I have also seem very counterintuitive to a VERY certain HDI result (see image below). (Prior and Posterior Predictive Checks — PyMC3 3.10.0 documentation).
Thanks again for all your time. I have spent most the weekend playing with pymc3; its super addicting and very useful.
MWE:
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as pm
import arviz as az
import theano
import xarray as xr
from sklearn.datasets import make_regression
X, target, coef = make_regression(n_samples=100,
bias=1.0,
n_informative=3,
n_features=3,
noise=2.5,
random_state=1,
coef=True)
target = target
treatment = X[:, 0]
covariates = X[:, 1:]
# coords
coords = {
# dim 1: length of data set
'obs_id': np.arange(len(target)),
# dim 2: treatment var
'treatment': ['treatment'],
# dim 3: covariates
'covariates': ['covariate_1', 'covariate_2']
}
# specify model
with pm.Model(coords=coords) as sk_lm:
# data
treatment_data = pm.Data('treatment_data', treatment, dims=('obs_id'))
covariates_data = pm.Data('covariates_data', covariates, dims=('obs_id', 'covariates'))
# priors
# intentionally off the mark
alpha = pm.Normal('alpha', mu=0, sigma=1)
treatment_beta = pm.Normal("treatment_beta", mu=30, sigma=2)
covariates_betas = pm.Normal('covariates_betas', mu=[66, 33], sigma=2, dims=('covariates'))
# model error
sigma = pm.HalfNormal("sigma", sigma=1)
# matrix-dot products
m1 = pm.math.matrix_dot(covariates, covariates_betas)
# expected value of y
mu = alpha + (treatment_beta * treatment) + m1
# Likelihood: Normal
y = pm.Normal("y",
mu=mu,
sigma=sigma,
observed=target,
dims='obs_id')
# set sampler
step = pm.NUTS([alpha, treatment_beta, covariates_betas, sigma], target_accept=0.9)
# prior analysis
with sk_lm:
prior_pc = pm.sample_prior_predictive()
# Inference button (TM)!
with sk_lm:
lm_trace = pm.sample(draws=1500,
step=step,
init='jitter+adapt_diag',
cores=4,
tune=500, # burn in
return_inferencedata=False)
# posterior analysis
with sk_lm:
ppc = pm.fast_sample_posterior_predictive(trace=lm_trace,
random_seed=1,
)
# generate InferenceData
with sk_lm:
lm_idata = az.from_pymc3(
trace=lm_trace,
prior=prior_pc,
posterior_predictive=ppc,
coords=coords
)
# Posterior Analysis
# extract posterior data for all parameters; stack chains + draws
post = lm_idata.posterior.stack(draws=("chain", "draw"))
# extract the data used to build the model
const = lm_idata.constant_data
# check data
post['treatment_beta'].shape # (6000, )
# matrix multiply covariate matrix by posterior covariates betas matrix
post["covariate_mm"] = xr.apply_ufunc(np.einsum, "...i,...i",
post["covariates_betas"],
const["covariates_data"],
input_core_dims=[[],
["covariates"],
["covariates"]])
Area of Concern
# roughly following http://docs.pymc.io/notebooks/posterior_predictive.html
# generate expected value of y; alpha + beta1 + covariates
mu = post["alpha"] + (post['treatment_beta'] * const['treatment_data']) + post["covariate_mm"]
# mean outcome plot
_, ax = plt.subplots(figsize=(8,8))
# why is this not a straight line? also seems counterintuitive to the HDI results
ax.plot(np.sort(treatment), np.sort(mu.mean('draws')), label="Mean outcome", alpha=0.6)
# plot data
ax.plot(treatment, target, "o", ms=4, alpha=0.4, label="Data")
# plot 94% HPD for E(y)
# arviz.plot_hdi has superceded arvis.plot_hpd?
az.plot_hdi(x=lm_idata.constant_data.treatment_data, # observed treatment
y=mu, # expected value of y
ax=ax,
smooth=False,
hdi_prob=0.94,
fill_kwargs={"alpha": 0.8, "label": "Mean outcome 94% HPD"})
# plot 94% HPD for posterior possible unobserved y values
az.plot_hdi(x=lm_idata.constant_data.treatment_data, # observed treatment data
y=stacked_posterior_predictive['y'].T, # set of possible unobserved y
ax=ax,
smooth=False,
hdi_prob=0.94,
fill_kwargs={"alpha": 0.8, "color": "#a1dab4", "label": "Outcome 94% HPD"})
ax.set_xlabel("Predictor")
ax.set_ylabel("Outcome")
ax.set_title("Posterior predictive checks")
ax.legend(ncol=2, fontsize=10);
