<jupyter_start><jupyter_text>Bayesian moderation analysis in PyMC3**Author:** [<NAME>](https://github.com/drbenvincent)This notebook provides an example of how to conduct a Bayesian [moderation analysis](https://en.wikipedia.org/wiki/Moderation_(statistics)). Note that this is _not_ mediation analysis, which is different.**NOTE:** Priors in the PyMC3 model need to be updated to fit your particular situation. This is especially the case for the moderation term $\beta_2$.  The moderation modelA 'full' moderation model defines an outcome variable $y$ as:$$y = \mathrm{Normal}(\beta_0 + (\beta_1 \cdot x) + (\beta_2 \cdot x \cdot moderator) + (\beta_3 \cdot moderator), \sigma^2)$$Where $x$ and $moderator$ are your observed data, and the following are the model parameters:- $\beta_0$ is the intercept- $\beta_1$ is the rate at which $y$ increases per unit of $x$. If $\beta_1 \neq 0$ then we have a main effect of $x$. - $\beta_2$ is the coefficient for the moderation term $x \cdot moderator$.- $\beta_3$ is the way in which y increases as $moderator$ increases. If $\beta_3 \neq 0$ then we have a main effect of $moderator$. - $\sigma$ is the standard deviation of the observation noise Optional: If you want to use autoformatter, then     pip install nb-black<jupyter_code>%load_ext lab_black
import numpy as np
import pandas as pd

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

plt.rcParams.update({"font.size": 14})

import pymc3 as pm
import arviz as az

print(f"Running on PyMC3 v{pm.__version__}")<jupyter_output>Running on PyMC3 v3.8<jupyter_text>Define options<jupyter_code>SEED = 123

N = 500

# define parameters (only used in the context of simulated data)
β0, β1, β2, β3, σ = 0.5, 1, 0.7, 3, 1
params = {"β0": β0, "β1": β1, "β2": β2, "β3": β3, "σ": σ}

n_samples = 1_000

sample_options = {
    "tune": 1_000,
    "draws": n_samples,
    "chains": 2,
    "cores": 2,
    "nuts_kwargs": {"target_accept": 0.95},
    "random_seed": SEED,
}

np.random.seed(SEED)

cmap = plt.cm.get_cmap("Spectral")<jupyter_output><empty_output><jupyter_text>Create synthetic dataset<jupyter_code>def make_data():
    x = np.random.randn(N)
    moderator = np.random.randn(N)
    μ = β0 + β1 * x + β2 * x * moderator + β3 * moderator
    ϵ = np.random.randn(N) * σ
    y = μ + ϵ
    return x, moderator, y


x, moderator, y = make_data()<jupyter_output><empty_output><jupyter_text>Plot data and true model Set up `ScalarMappable` so we can access colormap across both scatter and linesWe can use:- `scalarMap.cmap` to use as a colormap, and- `scalarMap.to_rgba(moderator_value)` to grab a colour for a given moderator value<jupyter_code>from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

scalarMap = ScalarMappable(
    norm=Normalize(vmin=np.min(moderator), vmax=np.max(moderator)), cmap="viridis"
)
def plot_data(x, moderator, y, ax=None):
    if ax is None:
        fig, ax = plt.subplots(1, 1)
    else:
        fig = plt.gcf()

    h = ax.scatter(x, y, c=moderator, cmap=scalarMap.cmap)
    ax.set(xlabel="x", ylabel="y")
    # colourbar for moderator
    cbar = fig.colorbar(h)
    cbar.ax.set_ylabel("moderator")
    return ax


def plot_true(ax):
    _x = np.linspace(np.min(x), np.max(x), 10)
    moderator_percentiles = [2.5, 25, 50, 75, 97.5]
    m_levels = np.percentile(moderator, moderator_percentiles)
    for p, m in zip(moderator_percentiles, m_levels):
        μ = β0 + β1 * _x + β2 * _x * m + β3 * m
        ax.plot(
            _x, μ, lw=3, label=f"{p}th percentile of moderator", c=scalarMap.to_rgba(m)
        )
    ax.legend(fontsize=9)


fig, ax = plt.subplots(figsize=(10, 6))
ax = plot_data(x, moderator, y, ax=ax)
plot_true(ax)
ax.set_title("Data and true model")<jupyter_output><empty_output><jupyter_text>Make model<jupyter_code>def bayesian_moderation(y, x, m):
    """Conduct simple Bayesian moderation. 
    Returns the model and the trace."""

    print("WARNING: Priors must be updated to suit your speicific data context")

    with pm.Model() as model:

        β0 = pm.Normal("β0", mu=0, sd=1)
        β1 = pm.Normal("β1", mu=0, sd=1)
        β2 = pm.Normal("β2", mu=0, sd=1)
        β3 = pm.Normal("β3", mu=0, sd=1)
        σ = pm.HalfCauchy("σ", 1)

        y = pm.Normal("y", mu=β0 + (β1 * x) + (β2 * x * m) + (β3 * m), sd=σ, observed=y)

    with model:
        trace = pm.sample(**sample_options)

    return model, trace
model, trace = bayesian_moderation(y, x, moderator)<jupyter_output>WARNING: Priors must be updated to suit your speicific data context<jupyter_text>Check the trace if you like, using:    pm.traceplot(trace)  Plot posterior over parameters.<jupyter_code>def plot_posteriors():
    fig, ax = plt.subplots(1, 5, figsize=(14, 5))
    for i, (key, val) in enumerate(params.items()):
        az.plot_posterior(trace, var_names=key, ax=ax[i], ref_val=val)


plot_posteriors()<jupyter_output><empty_output><jupyter_text>Posterior predictionFirst plot posterior prediction in data space<jupyter_code>def posterior_prediction_plot(
    trace, y, x, moderator, ax=None, percentile_list=[2.5, 25, 50, 75, 97.5]
):
    """Plot posterior predicted `y` for the defined moderator percentiles"""
    if ax is None:
        fig, ax = plt.subplots(1, 1)

    # PLOT POSTERIOR PREDICTED `y` FOR DEFINED MODERATOR LEVELS
    xi = np.linspace(np.min(x), np.max(x), 20)
    m_levels = np.percentile(moderator, percentile_list)

    for p, m in zip(percentile_list, m_levels):
        β0 = np.expand_dims(trace["β0"], axis=1)
        β1 = np.expand_dims(trace["β1"], axis=1)
        β2 = np.expand_dims(trace["β2"], axis=1)
        β3 = np.expand_dims(trace["β3"], axis=1)
        _y = β0 + β1 * xi + β2 * xi * m + β3 * m
        region = np.percentile(_y, [2.5, 50, 95], axis=0)

        ax.fill_between(
            xi,
            region[0, :],
            region[2, :],
            alpha=0.2,
            color=scalarMap.to_rgba(m),
            edgecolor="w",
        )
        ax.plot(
            xi,
            region[1, :],
            color=scalarMap.to_rgba(m),
            linewidth=2,
            label=f"{p}th percentile of moderator",
        )

    ax.legend(fontsize=9)

    ax.set(xlabel="x", ylabel="y")
    return ax
fig, ax = plt.subplots(figsize=(10, 6))
plot_data(x, moderator, y, ax=ax)
posterior_prediction_plot(trace, y, x, moderator, ax=ax)
ax.set_title("Data and posterior prediction")<jupyter_output><empty_output><jupyter_text>Now examine the moderation effect in more detail.The moderation model $$y = \mathrm{Normal}(\beta_0 + (\beta_1 \cdot x) + (\beta_2 \cdot x \cdot moderator) + (\beta_3 \cdot moderator), \sigma^2)$$can be re-written as$$y = \mathrm{Normal}(\beta_0 + x (\beta_1 + \beta_2 \cdot moderator) + (\beta_3 \cdot moderator), \sigma^2)$$where $\beta_1 + \beta_2 \cdot moderator$ describes how the moderator influences the linear relationship between $x$ and $y$.<jupyter_code>def plot_moderation_effect(
    ax=None,
    true_β1=None,
    true_β2=None,
    percentile_list=[0, 2.5, 25, 50, 75, 97.5, 100],
    samples_to_plot=100,
):

    if ax is None:
        fig, ax = plt.subplots(1, 1)

    m = np.linspace(np.min(moderator), np.max(moderator), 2)

    # sample from posterior
    if samples_to_plot < n_samples:
        ind = np.random.choice(range(n_samples), samples_to_plot)

    β1 = np.expand_dims(trace["β1"], axis=1)[ind, :]
    β2 = np.expand_dims(trace["β2"], axis=1)[ind, :]
    rate = β1 + β2 * m

    # spaghetti plot
    for n in range(samples_to_plot):
        ax.plot(m, rate[n, :], "k", alpha=0.1)

    # plot true model, if known
    if (true_β1 is not None) and (true_β2 is not None):
        true = true_β1 + true_β2 * m
        ax.plot(m, true, "r", lw=3, label="true")

    # plot points at each percentile of m
    m_levels = np.percentile(moderator, percentile_list)
    for m in m_levels:
        ax.plot(
            m, np.mean(β1) + np.mean(β2) * m, "o", c=scalarMap.to_rgba(m), markersize=10
        )

    ax.legend()

    ax.set(
        title="moderation effect",
        xlabel="$moderator$",
        ylabel=r"$\beta_1 + \beta_2 \cdot moderator$",
    )


fig, ax = plt.subplots(1, 2, figsize=(10, 6))

plot_moderation_effect(ax[0], true_β1=β1, true_β2=β2)
az.plot_posterior(trace, var_names="β2", ax=ax[1], ref_val=β2)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13779055835130602, 'fraction_numerical': 0.027797747423915647, 'mean_word_length': 3.4328199681359535, 'pattern_counts': {'":': 13, '<': 33, '<?xml version=': 0, '>': 32, 'https://': 2, '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': '476878', 'license': 'MIT', 'n_tokens_mistral': 3261, 'n_tokens_neox': 3049, 'n_words': 829, 'path': 'moderation analysis.ipynb', 'repo_name': 'drbenvincent/pymc3-demo-code', 'size': '8351'}