Multiple linear regression#

[1]:
import arviz as az
import bambi as bmb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
[2]:
az.style.use("arviz-darkgrid")
SEED = 7355608

Load and examine Eugene-Springfield community sample data#

Bambi comes with several datasets. These can be accessed via the load_data() function.

[3]:
data = bmb.load_data("ESCS")
np.round(data.describe(), 2)
[3]:
drugs n e o a c hones emoti extra agree consc openn
count 604.00 604.00 604.00 604.00 604.00 604.00 604.00 604.00 604.00 604.00 604.00 604.00
mean 2.21 80.04 106.52 113.87 124.63 124.23 3.89 3.18 3.21 3.13 3.57 3.41
std 0.65 23.21 19.88 21.12 16.67 18.69 0.45 0.46 0.53 0.47 0.44 0.52
min 1.00 23.00 42.00 51.00 63.00 44.00 2.56 1.47 1.62 1.59 2.00 1.28
25% 1.71 65.75 93.00 101.00 115.00 113.00 3.59 2.88 2.84 2.84 3.31 3.06
50% 2.14 76.00 107.00 112.00 126.00 125.00 3.88 3.19 3.22 3.16 3.56 3.44
75% 2.64 93.00 120.00 129.00 136.00 136.00 4.20 3.47 3.56 3.44 3.84 3.75
max 4.29 163.00 158.00 174.00 171.00 180.00 4.94 4.62 4.75 4.44 4.75 4.72

It’s always a good idea to start off with some basic plotting. Here’s what our outcome variable drugs (some index of self-reported illegal drug use) looks like:

[4]:
data["drugs"].hist();
../_images/notebooks_ESCS_multiple_regression_6_0.png

The five numerical predictors that we’ll use are sum-scores measuring participants’ standings on the Big Five personality dimensions. The dimensions are:

  • O = Openness to experience

  • C = Conscientiousness

  • E = Extraversion

  • A = Agreeableness

  • N = Neuroticism

Here’s what our predictors look like:

[5]:
az.plot_pair(data[["o", "c", "e", "a", "n"]].to_dict("list"), marginals=True, textsize=24);
../_images/notebooks_ESCS_multiple_regression_8_0.png

We can easily see all the predictors are more or less symmetrically distributed without outliers and the pairwise correlations between them are not strong.

Specify model and examine priors#

We’re going to fit a pretty straightforward additive multiple regression model predicting drug index from all 5 personality dimension scores. It’s simple to specify the model using a familiar formula interface. Here we also tell Bambi to run two parallel Markov Chain Monte Carlo (MCMC) chains, each one with 2000 draws. The first 1000 draws are tuning steps that we discard and the last 1000 draws are considered to be taken from the joint posterior distribution of all the parameters (to be confirmed when we analyze the convergence of the chains).

[6]:
model = bmb.Model("drugs ~ o + c + e + a + n", data)
fitted = model.fit(tune=2000, draws=2000, init="adapt_diag", random_seed=SEED)
Auto-assigning NUTS sampler...
Initializing NUTS using adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [Intercept, o, c, e, a, n, drugs_sigma]
100.00% [16000/16000 00:21<00:00 Sampling 4 chains, 0 divergences]
Sampling 4 chains for 2_000 tune and 2_000 draw iterations (8_000 + 8_000 draws total) took 22 seconds.

Great! But this is a Bayesian model, right? What about the priors? If no priors are given explicitly by the user, then Bambi chooses smart default priors for all parameters of the model based on the implied partial correlations between the outcome and the predictors. Here’s what the default priors look like in this case – the plots below show 1000 draws from each prior distribution:

[7]:
model.plot_priors();
../_images/notebooks_ESCS_multiple_regression_14_0.png
[8]:
# Normal priors on the coefficients
{x.name: x.prior.args for x in model.terms.values()}
[8]:
{'Intercept': {'mu': array(2.21014664), 'sigma': array(21.19375074)},
 'o': {'mu': array(0.), 'sigma': array(0.0768135)},
 'c': {'mu': array(0.), 'sigma': array(0.08679683)},
 'e': {'mu': array(0.), 'sigma': array(0.0815892)},
 'a': {'mu': array(0.), 'sigma': array(0.09727366)},
 'n': {'mu': array(0.), 'sigma': array(0.06987412)}}
[9]:
# HalfStudentT prior on the residual standard deviation
model.family.likelihood.priors
[9]:
{'sigma': HalfStudentT(nu: 4, sigma: 0.6482)}

You could also just print the model and see it also contains the same information about the priors

[10]:
model
[10]:
Formula: drugs ~ o + c + e + a + n
Family name: Gaussian
Link: identity
Observations: 604
Priors:
  Common-level effects
    Intercept ~ Normal(mu: 2.2101, sigma: 21.1938)
    o ~ Normal(mu: 0.0, sigma: 0.0768)
    c ~ Normal(mu: 0.0, sigma: 0.0868)
    e ~ Normal(mu: 0.0, sigma: 0.0816)
    a ~ Normal(mu: 0.0, sigma: 0.0973)
    n ~ Normal(mu: 0.0, sigma: 0.0699)

  Auxiliary parameters
    sigma ~ HalfStudentT(nu: 4, sigma: 0.6482)
------
* To see a plot of the priors call the .plot_priors() method.
* To see a summary or plot of the posterior pass the object returned by .fit() to az.summary() or az.plot_trace()

Some more info about the default prior distributions can be found in this technical paper.

Notice the apparently small SDs of the slope priors. This is due to the relative scales of the outcome and the predictors: remember from the plots above that the outcome, drugs, ranges from 1 to about 4, while the predictors all range from about 20 to 180 or so. A one-unit change in any of the predictors – which is a trivial increase on the scale of the predictors – is likely to lead to a very small absolute change in the outcome. Believe it or not, these priors are actually quite wide on the partial correlation scale!

Examine the model results#

Let’s start with a pretty picture of the parameter estimates!

[11]:
az.plot_trace(fitted);
../_images/notebooks_ESCS_multiple_regression_23_0.png

The left panels show the marginal posterior distributions for all of the model’s parameters, which summarize the most plausible values of the regression coefficients, given the data we have now observed. These posterior density plots show two overlaid distributions because we ran two MCMC chains. The panels on the right are “trace plots” showing the sampling paths of the two MCMC chains as they wander through the parameter space. If any of these paths exhibited a pattern other than white noise we would be concerned about the convergence of the chains.

A much more succinct (non-graphical) summary of the parameter estimates can be found like so:

[12]:
az.summary(fitted)
[12]:
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
Intercept 3.300 0.361 2.644 4.005 0.004 0.003 7177.0 5823.0 1.0
o 0.006 0.001 0.004 0.008 0.000 0.000 8429.0 6970.0 1.0
c -0.004 0.001 -0.007 -0.001 0.000 0.000 7935.0 6729.0 1.0
e 0.003 0.001 0.001 0.006 0.000 0.000 9255.0 7278.0 1.0
a -0.012 0.002 -0.015 -0.010 0.000 0.000 9516.0 6394.0 1.0
n -0.001 0.001 -0.004 0.001 0.000 0.000 8076.0 6415.0 1.0
drugs_sigma 0.592 0.017 0.561 0.624 0.000 0.000 9487.0 6676.0 1.0

When there are multiple MCMC chains, the default summary output includes some basic convergence diagnostic info (the effective MCMC sample sizes and the Gelman-Rubin “R-hat” statistics), although in this case it’s pretty clear from the trace plots above that the chains have converged just fine.

Summarize effects on partial correlation scale#

[13]:
samples = fitted.posterior

It turns out that we can convert each regression coefficient into a partial correlation by multiplying it by a constant that depends on (1) the SD of the predictor, (2) the SD of the outcome, and (3) the degree of multicollinearity with the set of other predictors. Two of these statistics are actually already computed and stored in the fitted model object, in a dictionary called dm_statistics (for design matrix statistics), because they are used internally. We will compute the others manually.

Some information about the relationship between linear regression parameters and partial correlation can be found here.

[14]:
# the names of the predictors
varnames = ['o', 'c', 'e', 'a', 'n']

# compute the needed statistics like R-squared when each predictor is response and all the
# other predictors are the predictor

# x_matrix = common effects design matrix (excluding intercept/constant term)
terms = [t for t in model.common_terms.values() if t.name != "Intercept"]
x_matrix = [pd.DataFrame(x.data, columns=x.levels) for x in terms]
x_matrix = pd.concat(x_matrix, axis=1)

dm_statistics = {
    'r2_x': pd.Series(
        {
            x: sm.OLS(
                endog=x_matrix[x],
                exog=sm.add_constant(x_matrix.drop(x, axis=1))
                if "Intercept" in model.term_names
                else x_matrix.drop(x, axis=1),
            )
            .fit()
            .rsquared
            for x in list(x_matrix.columns)
        }
    ),
    'sigma_x': x_matrix.std(),
    'mean_x': x_matrix.mean(axis=0),
}

r2_x = dm_statistics['r2_x']
sd_x = dm_statistics['sigma_x']
r2_y = pd.Series([sm.OLS(endog=data['drugs'],
                         exog=sm.add_constant(data[[p for p in varnames if p != x]])).fit().rsquared
                  for x in varnames], index=varnames)
sd_y = data['drugs'].std()

# compute the products to multiply each slope with to produce the partial correlations
slope_constant = (sd_x[varnames] / sd_y) * ((1 - r2_x[varnames]) / (1 - r2_y)) ** 0.5
slope_constant
[14]:
o    32.392557
c    27.674284
e    30.305117
a    26.113299
n    34.130431
dtype: float64

Now we just multiply each sampled regression coefficient by its corresponding slope_constant to transform it into a sample partial correlation coefficient.

[15]:
# We also stack draws and chains into samples.
# This will be useful for the following plots
pcorr_samples = (samples[varnames] * slope_constant).stack(samples=("draw", "chain"))

And voilà! We now have a joint posterior distribution for the partial correlation coefficients. Let’s plot the marginal posterior distributions:

[16]:
# Pass the same axes to az.plot_kde to have all the densities in the same plot
_, ax = plt.subplots()
for idx, (k, v) in enumerate(pcorr_samples.items()):
    az.plot_kde(v.values, label=k, plot_kwargs={'color':f'C{idx}'}, ax=ax)
ax.axvline(x=0, color='k', linestyle='--');
../_images/notebooks_ESCS_multiple_regression_35_0.png

The means of these distributions serve as good point estimates of the partial correlations:

[17]:
pcorr_samples.mean()
[17]:
<xarray.Dataset>
Dimensions:  ()
Data variables:
    o        float64 0.1958
    c        float64 -0.1048
    e        float64 0.1027
    a        float64 -0.3244
    n        float64 -0.05116

Relative importance: Which predictors have the strongest effects (defined in terms of squared partial correlation?#

We just take the square of the partial correlation coefficients, so it’s easy to get posteriors on that scale too:

[18]:
_, ax = plt.subplots()
for idx, (k, v) in enumerate(pcorr_samples.items()):
    az.plot_kde(v.values ** 2, label=k, plot_kwargs={'color':f'C{idx}'}, ax=ax)
ax.set_ylim(0, 80);
../_images/notebooks_ESCS_multiple_regression_40_0.png

With these posteriors we can ask: What is the probability that the squared partial correlation for Openness (blue) is greater than the squared partial correlation for Conscientiousness (orange)?

[19]:
(pcorr_samples['o'] ** 2 > pcorr_samples['c'] ** 2).mean()
[19]:
<xarray.DataArray ()>
array(0.927875)

If we contrast this result with the plot we’ve just shown, we may think the probability is too high when looking at the overlap between the blue and orange curves. However, the previous plot is only showing marginal posteriors, which don’t account for correlations between the coefficients. In our Bayesian world, our model parameters’ are random variables (and consequently, any combination of them are too). As such, squared partial correlation have a joint distribution. When computing probabilities involving at least two of these parameters, one has to use the joint distribution. Otherwise, if we choose to work only with marginals, we are implicitly assuming independence.

Let’s check the joint distribution of the squared partial correlation for Openness and Conscientiousness. We highlight with a blue color the draws where the coefficient for Openness is greater than the coefficient for Conscientiousness.

[20]:
sq_partial_c = pcorr_samples['c'].values ** 2
sq_partial_o = pcorr_samples['o'].values ** 2
[21]:
colors = list(np.where(sq_partial_c > sq_partial_o, "C1", "C0"))

plt.scatter(sq_partial_o, sq_partial_c, c=colors)
plt.xlabel("Openness to experience")
plt.ylabel("Conscientiousness");
../_images/notebooks_ESCS_multiple_regression_45_0.png

We can see that in the great majority of the draws (92.8%) the squared partial correlation for Openness is greater than the one for Conscientiousness. In fact, we can check the correlation between them is

[22]:
np.corrcoef(sq_partial_c, sq_partial_o)[0, 1]
[22]:
-1.874138083091745e-01

which explains why ony looking at the marginal posteriors (i.e. assuming independence) is not the best approach here.

For each predictor, what is the probability that it has the largest squared partial correlation?

[23]:
pc_df = pcorr_samples.to_dataframe()
(pc_df**2).idxmax(axis=1).value_counts() / len(pc_df.index)
[23]:
a    0.990625
o    0.009250
c    0.000125
dtype: float64

Agreeableness is clearly the strongest predictor of drug use among the Big Five personality traits in terms of partial correlation, but it’s still not a particularly strong predictor in an absolute sense. Walter Mischel famously claimed that it is rare to see correlations between personality measure and relevant behavioral outcomes exceed 0.3. In this case, the probability that the agreeableness partial correlation exceeds 0.3 is:

[24]:
(np.abs(pcorr_samples['a']) > 0.3).mean()
[24]:
<xarray.DataArray 'a' ()>
array(0.729125)

Posterior Predictive#

Once we have computed the posterior distribution, we can use it to compute the posterior predictive distribution. As the name implies, these are predictions assuming the model’s parameter are distributed as the posterior. Thus, the posterior predictive includes the uncertainty about the parameters.

With bambi we can use the model’s predict() method with the fitted az.InferenceData to generate a posterior predictive samples, which are then automatically added to the az.InferenceData object

[25]:
posterior_predictive = model.predict(fitted, kind="pps")
fitted
[25]:
arviz.InferenceData
    • <xarray.Dataset>
      Dimensions:      (chain: 4, draw: 2000)
      Coordinates:
        * chain        (chain) int64 0 1 2 3
        * draw         (draw) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
      Data variables:
          Intercept    (chain, draw) float64 3.405 3.692 3.085 ... 3.546 3.538 3.121
          o            (chain, draw) float64 0.006116 0.006835 ... 0.006975 0.007124
          c            (chain, draw) float64 -0.00307 -0.004805 ... -0.001573
          e            (chain, draw) float64 0.004453 0.002762 ... 0.0006999 0.001612
          a            (chain, draw) float64 -0.01442 -0.01332 ... -0.01356 -0.0127
          n            (chain, draw) float64 -0.002132 -0.003611 ... -0.001556
          drugs_sigma  (chain, draw) float64 0.5764 0.6029 0.5865 ... 0.6052 0.5884
      Attributes:
          created_at:                  2022-06-01T18:53:40.956618
          arviz_version:               0.12.1
          inference_library:           pymc
          inference_library_version:   4.0.0b6
          sampling_time:               21.887134313583374
          tuning_steps:                2000
          modeling_interface:          bambi
          modeling_interface_version:  0.9.0

    • <xarray.Dataset>
      Dimensions:      (chain: 4, draw: 500, drugs_dim_0: 604)
      Coordinates:
        * chain        (chain) int64 0 1 2 3
        * draw         (draw) int64 0 1 2 3 4 5 6 7 ... 493 494 495 496 497 498 499
        * drugs_dim_0  (drugs_dim_0) int64 0 1 2 3 4 5 6 ... 598 599 600 601 602 603
      Data variables:
          drugs        (chain, draw, drugs_dim_0) float64 1.893 1.686 ... 2.172 3.063
      Attributes:
          created_at:                  2022-06-01T18:53:45.790414
          arviz_version:               0.12.1
          modeling_interface:          bambi
          modeling_interface_version:  0.9.0

    • <xarray.Dataset>
      Dimensions:      (chain: 4, draw: 2000, drugs_dim_0: 604)
      Coordinates:
        * chain        (chain) int64 0 1 2 3
        * draw         (draw) int64 0 1 2 3 4 5 6 ... 1994 1995 1996 1997 1998 1999
        * drugs_dim_0  (drugs_dim_0) int64 0 1 2 3 4 5 6 ... 598 599 600 601 602 603
      Data variables:
          drugs        (chain, draw, drugs_dim_0) float64 -0.8509 -1.756 ... -2.754
      Attributes:
          created_at:                  2022-06-01T18:53:41.397576
          arviz_version:               0.12.1
          inference_library:           pymc
          inference_library_version:   4.0.0b6
          modeling_interface:          bambi
          modeling_interface_version:  0.9.0

    • <xarray.Dataset>
      Dimensions:             (chain: 4, draw: 2000)
      Coordinates:
        * chain               (chain) int64 0 1 2 3
        * draw                (draw) int64 0 1 2 3 4 5 ... 1995 1996 1997 1998 1999
      Data variables: (12/13)
          energy              (chain, draw) float64 540.0 539.1 539.8 ... 541.4 538.1
          tree_depth          (chain, draw) int64 2 3 2 2 3 2 2 2 ... 2 3 2 3 3 2 3 2
          perf_counter_diff   (chain, draw) float64 0.001072 0.009217 ... 0.0008147
          diverging           (chain, draw) bool False False False ... False False
          acceptance_rate     (chain, draw) float64 0.9234 0.8429 ... 0.8391 0.9759
          energy_error        (chain, draw) float64 -0.4879 -0.1736 ... 0.5002 -1.064
          ...                  ...
          perf_counter_start  (chain, draw) float64 2.503e+04 2.503e+04 ... 2.504e+04
          step_size_bar       (chain, draw) float64 0.8968 0.8968 ... 0.8924 0.8924
          n_steps             (chain, draw) float64 3.0 7.0 3.0 3.0 ... 3.0 7.0 3.0
          lp                  (chain, draw) float64 -536.0 -535.8 ... -537.9 -535.7
          max_energy_error    (chain, draw) float64 0.6073 0.509 ... 0.6089 -1.064
          step_size           (chain, draw) float64 0.807 0.807 ... 0.8926 0.8926
      Attributes:
          created_at:                  2022-06-01T18:53:40.961202
          arviz_version:               0.12.1
          inference_library:           pymc
          inference_library_version:   4.0.0b6
          sampling_time:               21.887134313583374
          tuning_steps:                2000
          modeling_interface:          bambi
          modeling_interface_version:  0.9.0

    • <xarray.Dataset>
      Dimensions:      (drugs_dim_0: 604)
      Coordinates:
        * drugs_dim_0  (drugs_dim_0) int64 0 1 2 3 4 5 6 ... 598 599 600 601 602 603
      Data variables:
          drugs        (drugs_dim_0) float64 1.857 3.071 1.571 2.214 ... 1.5 2.5 3.357
      Attributes:
          created_at:                  2022-06-01T18:53:41.398388
          arviz_version:               0.12.1
          inference_library:           pymc
          inference_library_version:   4.0.0b6
          modeling_interface:          bambi
          modeling_interface_version:  0.9.0

One use of the posterior predictive is as a diagnostic tool, shown below using az.plot_ppc().The blue lines represent the posterior predictive distribution estimates, and the black line represents the observed data. Our posterior predictions seems perform an adequately good job representing the observed data in all regions except near the value of 1, where the observed data and posterior estimates diverge.

[26]:
az.plot_ppc(fitted);
/home/tomas/anaconda3/envs/bambi/lib/python3.9/site-packages/IPython/core/pylabtools.py:151: UserWarning: Creating legend with loc="best" can be slow with large amounts of data.
  fig.canvas.print_figure(bytes_io, **kw)
../_images/notebooks_ESCS_multiple_regression_57_1.png
[27]:
%load_ext watermark
%watermark -n -u -v -iv -w
Last updated: Wed Jun 01 2022

Python implementation: CPython
Python version       : 3.9.7
IPython version      : 8.3.0

numpy      : 1.21.5
matplotlib : 3.5.1
pandas     : 1.4.2
bambi      : 0.7.1
statsmodels: 0.13.2
sys        : 3.9.7 (default, Sep 16 2021, 13:09:58)
[GCC 7.5.0]
arviz      : 0.12.1

Watermark: 2.3.0