PyMC3 posterior prediction example; how to incorporate a matrix of betas?

Questions
13

Last few questions on this topic if I may (here with actual regression data)!

I am curious about a few things:

  1. 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?

  2. 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);

pred1
pred1760×742 108 KB