Causal Inference with Multilevel Models: The Electric Company Example

Prepare Notebook

We begin by importing the necessary libraries.

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import polars as pl
import preliz as pz
import pymc as pm
import pytensor.tensor as pt
import seaborn as sns
from marginaleffects import datagrid
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OrdinalEncoder, StandardScaler

az.style.use("arviz-darkgrid")
plt.rcParams["figure.figsize"] = [10, 6]
plt.rcParams["figure.dpi"] = 100
plt.rcParams["figure.facecolor"] = "white"

%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = "retina"

seed: int = sum(map(ord, "bayes"))
rng: np.random.Generator = np.random.default_rng(seed=seed)

Read Data

We load the Electric Company dataset (from Aki Vehtari’s repository). The data contains pre-test and post-test reading scores for students in paired classrooms across four grade levels. Each pair consists of one treatment and one control classroom.

data_path = "https://raw.githubusercontent.com/avehtari/ROS-Examples/master/ElectricCompany/data/electric.csv"
raw_df = pl.read_csv(data_path).drop(["supp", ""]).sort(["grade", "pair_id"])

raw_df.head()

The dataset includes: grade (grade level), pair_id (identifier for matched classroom pairs), treatment (binary indicator), pre_test (baseline reading score), and post_test (outcome reading score).

Exploratory Data Analysis

Before modeling, it is essential to understand the structure and variability in the data. Since the design relies on paired classrooms, we should inspect the distribution of test scores across grades and pairs to identify the sources of variation that our multilevel model must capture.

g = sns.pairplot(
    raw_df.to_pandas(),
    vars=["pre_test", "grade", "post_test"],
    hue="treatment",
    diag_kind="hist",
)
g.figure.suptitle("Pairwise Relationships", fontsize=18, fontweight="bold", y=1.02);

The naive treatment effect is the difference between the mean post-test scores of the treatment and control groups:

(
    raw_df.group_by(["grade", "treatment"])
    .agg(pl.col("post_test").mean())
    .pivot(index="grade", on="treatment", values="post_test")
    .with_columns(pl.col("1").sub(pl.col("0")).alias("treatment_effect"))
    .sort(by="grade")
)

Data Preprocessing

We standardize the numeric variables (test scores) to facilitate prior specification and improve sampling efficiency. Grade and pair identifiers are encoded as integers for use as group indices in the hierarchical model.

numeric_features = ["pre_test", "post_test"]
ordinal_features = ["grade", "pair_id"]

preprocessor = ColumnTransformer(
    [
        ("num", StandardScaler(), numeric_features),
        ("ord", OrdinalEncoder(dtype=int), ordinal_features),
    ],
    remainder="passthrough",
    verbose_feature_names_out=False,
).set_output(transform="polars")


df = preprocessor.fit_transform(raw_df)

df.head()

We extract the covariate matrix containing only the pre-test scores, which will be used to control for baseline differences when estimating the treatment effect.

x_columns = ["pre_test"]
x_df = df[x_columns]

x_df.head()

Next, we set up the coordinate system for our PyMC model. This includes dimensions for covariates, grades, pairs, and observations, which will allow us to structure the hierarchical relationships in the data.

n_grades = len(preprocessor["ord"].categories_[ordinal_features.index("grade")])
n_pairs = len(preprocessor["ord"].categories_[ordinal_features.index("pair_id")])

coords = {
    # covariates
    "covariates": x_df.columns,
    # grade
    "grade": preprocessor["ord"].categories_[ordinal_features.index("grade")],
    # object categories (groups)
    "pair_id": preprocessor["ord"].categories_[ordinal_features.index("pair_id")],
    # index
    "obs_idx": np.arange(len(df)),
}

Model Specification

Motivation: Why Multilevel?

The data is clustered into pairs of classrooms. Observations within the same pair (same school/grade) are likely correlated due to shared unobserved factors like school quality, neighborhood demographics, or teacher characteristics. Ignoring this structure would violate the independence assumption of standard regression.

To estimate the causal effect, we have two main strategies to handle this clustering:

Alternative 1: Fixed Effects (FE)

We could include a dummy variable (intercept) for every pair.

• Pros: This controls for ALL time-invariant pair-level unobserved confounders. It effectively “closes back doors” related to the school or neighborhood (see The Effect, Chapter 16).

• Cons: It consumes a massive number of degrees of freedom (\(n/2\) parameters just for intercepts!). With many small groups (pairs), this can lead to noisy and inefficient estimates. Of course we could use the Frisch-Waugh-Lovell theorem to demean the data as implemented in the PyFixest package.

Alternative 2: Random Effects (RE) / Hierarchical Intercepts

We model the pair intercepts as coming from a common distribution, e.g., \(\alpha_{j} \sim \text{Normal}(\mu, \sigma)\).

• Pros: This uses partial pooling. The model learns the variance \(\sigma\) and “shrinks” noisy pair estimates toward the global mean. It is far more efficient than FE.

• Cons: For this we need to expand a bit more …

Addressing Random Effects: A Note of Caution

A fundamental concern with Random Effects models, discussed extensively in econometrics and causal inference, is the assumption that group effects are uncorrelated with the predictors (\(\text{Cov}(\alpha_j, X) = 0\)). When this assumption is violated—for example, if high-performing schools both score higher at baseline and implement treatments differently—standard Random Effects estimates can be biased. Fixed Effects models avoid this problem by differencing out all time-invariant group characteristics, making no assumptions about their correlation with predictors. This is why Fixed Effects are often preferred in observational studies where such correlations are likely (Huntington-Klein, 2021, The Effect, Chapter 16).

Why Random Effects are Valid in this Randomized Design:

However, the concern about correlation between group effects and predictors does not apply uniformly to all variables. Crucially, in this experiment, treatment was randomized WITHIN pairs. - Because of randomization, the treatment assignment is uncorrelated with the pair’s baseline characteristics (the “random effect”) by design. - The randomization mechanism ensures \(\text{Cov}(\alpha_j, T_i) = 0\) for the treatment variable, even if \(\text{Cov}(\alpha_j, X_i) \neq 0\) for other covariates. - Therefore, the standard critique of Random Effects does not threaten the validity of our causal effect estimate for treatment, though we may still need to control for other confounders like pre-test scores.

We can thus use a Hierarchical Intercept Model to efficiently control for pair-level heterogeneity and obtain correct standard errors without the heavy penalty of Fixed Effects.

\[\begin{align*} \text{post_test}_i &\sim \text{Normal}(\mu_i, \sigma_y) \\ \mu_i &= \alpha_{\text{pair}[i]} + \theta \cdot T_i + \beta_x \cdot \text{pre_test}_i\\ \alpha_j &\sim \text{Normal}(\mu_\alpha, \sigma_\alpha) \end{align*}\]

Here, \(i\) indexes observations, \(j\) indexes pairs, \(T_i\) is the treatment indicator, and \(\alpha_j\) are the pair-specific intercepts. The hierarchical structure on \(\alpha_j\) implements partial pooling: pairs with little data are shrunk toward the grade-level mean \(\mu_\alpha\), while pairs with more data are allowed to deviate. We implement this using a non-centered parametrization for improved sampling efficiency.

Regarding the hierarchical structure: we use a non-centered parametrization for the random intercepts to improve sampling efficiency (for details, see A Primer on Bayesian Methods for Multilevel Modeling).

with pm.Model(coords=coords) as model:
    # --- Data Containers ---
    # covariates
    x_data = pm.Data("x_data", x_df, dims=("obs_idx", "covariates"))
    # grade
    grade_idx_data = pm.Data("grade_idx_data", df["grade"].to_numpy(), dims="obs_idx")
    # object categories
    pair_idx_data = pm.Data("pair_idx_data", df["pair_id"].to_numpy(), dims="obs_idx")
    # treatment
    treatment_data = pm.Data(
        "treatment_data", df["treatment"].to_numpy(), dims=("obs_idx")
    )
    # outcome
    post_test_data = pm.Data(
        "post_test_data", df["post_test"].to_numpy(), dims="obs_idx"
    )

    # --- Priors ---
    mu_alpha = pm.Normal("mu_alpha", mu=0, sigma=1, dims=("grade"))
    sigma_alpha = pm.HalfNormal("sigma_alpha", sigma=1, dims=("grade"))
    # auxiliary variable for the non-centered parametrization
    z_alpha = pm.Normal("z_alpha", mu=0, sigma=1, dims=("pair_id", "grade"))

    theta = pm.Normal("theta", mu=0, sigma=1, dims=("grade"))
    beta_x = pm.Normal("beta_x", mu=0, sigma=1, dims=("grade", "covariates"))

    sigma_outcome = pm.HalfNormal("sigma_outcome", sigma=1, dims=("grade"))

    # --- Parametrization ---
    # Non-centered parametrization for the random intercepts
    alpha = pm.Deterministic(
        "alpha", mu_alpha + z_alpha * sigma_alpha, dims=("pair_id", "grade")
    )

    mu_outcome = pm.Deterministic(
        "mu_outcome",
        alpha[pair_idx_data, grade_idx_data]
        + theta[grade_idx_data] * treatment_data
        + (beta_x[grade_idx_data] * x_data).sum(axis=-1),
        dims=("obs_idx"),
    )

    # --- Likelihood ---
    pm.Normal(
        "post_test_obs",
        mu=mu_outcome,
        sigma=sigma_outcome[grade_idx_data],
        observed=post_test_data,
        dims="obs_idx",
    )


pm.model_to_graphviz(model)

The graphical representation above shows the dependency structure of our model: how the observed data (post_test_obs) depends on the hierarchical parameters through the mean structure (mu_outcome).

Prior Predictive Check

Before fitting the model to data, we simulate from the prior distribution to ensure our priors are reasonable and produce plausible outcomes. This step helps detect specification errors and poorly calibrated priors.

with model:
    idata = pm.sample_prior_predictive(random_seed=rng)

fig, ax = plt.subplots()
az.plot_ppc(idata, group="prior", ax=ax)
az.plot_dist(df["post_test"].to_numpy(), color="black", ax=ax)
ax.set_title("Prior Predictive Check", fontsize=18, fontweight="bold", y=1.02);

The prior predictive distribution (blue) is compared against the observed data (black). The priors allow for a wide range of plausible outcomes, which is appropriate given our standardized data. The overlap indicates that the observed data are not surprising under our prior assumptions.

Posterior Inference

We now fit the model using Hamiltonian Monte Carlo (HMC) via the NumPyro backend.

with model:
    idata.extend(
        pm.sample(
            tune=1_500,
            draws=1_000,
            target_accept=0.9,
            chains=4,
            nuts_sampler="numpyro",
            random_seed=rng,
        )
    )

    idata.extend(pm.sample_posterior_predictive(idata, random_seed=rng))

After sampling, we generate posterior predictive samples to assess model fit. These predictions will be used to evaluate how well the model captures the observed data patterns.

Model Diagnostics

We check for sampling pathologies that would indicate problems with the posterior geometry or sampler configuration. The primary diagnostic is the number of divergent transitions, which should ideally be zero.

idata["sample_stats"]["diverging"].sum().item()

0

No divergent transitions indicate that the sampler successfully explored the posterior without encountering problematic curvature. We also examine the summary statistics and trace plots to assess convergence.

az.summary(
    idata,
    var_names=[
        "beta_x",
        "mu_alpha",
        "sigma_alpha",
        "sigma_outcome",
        "theta",
    ],
)

The summary statistics show effective sample sizes (ESS) and \(\hat{R}\) diagnostics for key parameters. Values of \(\hat{R} \approx 1\) indicate convergence across chains. The treatment effect parameter (theta) and variance components are our main interest.

axes = az.plot_trace(
    idata,
    var_names=[
        "beta_x",
        "mu_alpha",
        "sigma_alpha",
        "sigma_outcome",
        "theta",
    ],
    figsize=(10, 8),
)


plt.gcf().suptitle("Model Trace", fontsize=18, fontweight="bold", y=1.02);

The trace plots show good mixing across chains (overlapping colors) and stable posterior distributions (smooth histograms on the left). The chains have converged to stationary distributions, confirming the reliability of our posterior estimates.

Next, we check the posterior predictive distribution.

fig, ax = plt.subplots()
az.plot_ppc(idata, group="posterior", num_pp_samples=500, ax=ax)
ax.set_title("Posterior Predictive Check", fontsize=18, fontweight="bold", y=1.02);

The posterior predictive distribution closely matches the observed data, indicating that the model captures the essential features of the data-generating process.

Treatment Effect Estimates

We now examine the estimated treatment effects. The parameter theta represents the average causal effect of watching “The Electric Company” on standardized test scores, estimated separately for each grade level.

ax, *_ = az.plot_forest(
    idata,
    combined=True,
    var_names=["theta"],
    figsize=(10, 6),
)
ax.set_title(
    r"Treatment Effect Estimates $(94\% HDI)$",
    fontsize=18,
    fontweight="bold",
    y=1.02,
);

The forest plot shows the posterior distributions of treatment effects for each grade. The point estimates and credible intervals provide evidence about the magnitude and uncertainty of the causal effect. To interpret these effects in the original scale, we transform them back from standardized units.

ax, *_ = az.plot_forest(
    idata["posterior"]["theta"]
    * preprocessor["num"].scale_[numeric_features.index("post_test")],
    combined=True,
    var_names=["theta"],
    figsize=(10, 6),
)
ax.set_title(
    r"Treatment Effect Estimates $(94\% HDI)$ (Original Scale)",
    fontsize=18,
    fontweight="bold",
    y=1.02,
);

These are the treatment effects expressed in the original units of the test scores. Positive values indicate that the treatment group (who watched the show) scored higher on average than the control group. These values match the results from the book.

Counterfactual Prediction Grids

Here we want to show another way to compute the average treatment effect (ATE) using the marginaleffects package and PyMC do operator to generate counterfactual predictions.

We construct balanced prediction grids that span the range of pre-test scores for both treatment and control conditions. This allows us to compute counterfactual predictions: what would have happened to the same students under both treatment and control.

Remark: We can use the real data instead of the prediction grid. However, this is unnecessary and potentially computationally expensive if the data is large.

# Construct prediction grids (control)
raw_df_control_grid = datagrid(
    newdata=raw_df,
    treatment=0,
    grade=preprocessor["ord"].categories_[ordinal_features.index("grade")],
    pair_id=preprocessor["ord"].categories_[ordinal_features.index("pair_id")],
    pre_test=np.linspace(raw_df["pre_test"].min(), raw_df["pre_test"].max(), 10),
)
# Preprocess the prediction grid
df_control_grid = preprocessor.transform(raw_df_control_grid)
df_control_grid.columns = [col.split("__")[-1] for col in df_control_grid.columns]
x_df_control_grid = df_control_grid[x_columns]

# Construct prediction grids (treatment)
raw_df_treatment_grid = datagrid(
    newdata=raw_df,
    treatment=1,
    grade=preprocessor["ord"].categories_[ordinal_features.index("grade")],
    pair_id=preprocessor["ord"].categories_[ordinal_features.index("pair_id")],
    pre_test=np.linspace(raw_df["pre_test"].min(), raw_df["pre_test"].max(), 10),
)
# Preprocess the prediction grid
df_treatment_grid = preprocessor.transform(raw_df_treatment_grid)
df_treatment_grid.columns = [col.split("__")[-1] for col in df_treatment_grid.columns]
x_df_treatment_grid = df_treatment_grid[x_columns]

We have created two grids: one with treatment set to 0 (control) and one with treatment set to 1 (treated), both spanning the observed range of pre-test scores and all grade-pair combinations.

Counterfactual Estimates

We now generate posterior predictive samples under each counterfactual scenario. For the control grid, we predict outcomes as if all students were in the control condition. For the treatment grid, we predict outcomes as if all students watched the show.

with model:
    pm.set_data(
        new_data={
            "x_data": x_df_control_grid,
            "grade_idx_data": df_control_grid["grade"].to_numpy(),
            "pair_idx_data": df_control_grid["pair_id"].to_numpy(),
            "treatment_data": df_control_grid["treatment"].to_numpy(),
            "post_test_data": df_control_grid["post_test"].to_numpy(),
        },
        coords={
            "covariates": x_df_control_grid.columns,
            "grade": preprocessor["ord"].categories_[ordinal_features.index("grade")],
            "pair_id": preprocessor["ord"].categories_[
                ordinal_features.index("pair_id")
            ],
            "obs_idx": np.arange(len(df_control_grid)),
        },
    )

    posterior_predictive_control = pm.sample_posterior_predictive(
        idata, var_names=["post_test_obs", "mu_outcome"], random_seed=rng
    )

The control counterfactual predictions represent the expected outcomes in the absence of treatment, conditioning on the observed covariates and the estimated model parameters.

with model:
    pm.set_data(
        new_data={
            "x_data": x_df_treatment_grid,
            "grade_idx_data": df_treatment_grid["grade"].to_numpy(),
            "pair_idx_data": df_treatment_grid["pair_id"].to_numpy(),
            "treatment_data": df_treatment_grid["treatment"].to_numpy(),
            "post_test_data": df_treatment_grid["post_test"].to_numpy(),
        },
        coords={
            "covariates": x_df_treatment_grid.columns,
            "grade": preprocessor["ord"].categories_[ordinal_features.index("grade")],
            "pair_id": preprocessor["ord"].categories_[
                ordinal_features.index("pair_id")
            ],
            "obs_idx": np.arange(len(df_treatment_grid)),
        },
    )

    posterior_predictive_treatment = pm.sample_posterior_predictive(
        idata, var_names=["post_test_obs", "mu_outcome"], random_seed=rng
    )

Similarly, the treatment counterfactual predictions represent expected outcomes if all students had watched the show. The difference between these two counterfactuals yields the causal effect.

Results: Treatment Effect Estimates

We can now compute the average treatment effect (ATE) by comparing the posterior predictive distributions under the counterfactual scenarios (Treatment vs. Control). The difference between these predictions, averaged across all observations in the grid, provides a distribution for the causal effect.

For illustration, we focus on Grade 1 students. We first extract and transform the predictions back to the original scale of test scores.

control_mask = (
    raw_df_control_grid.select(pl.col("grade").eq(pl.lit(1))).to_numpy().flatten()
)
control_posterior_grade = posterior_predictive_control["posterior_predictive"][
    "mu_outcome"
][:, :, control_mask]

original_scale_control_posterior_grade = (
    control_posterior_grade
    * preprocessor["num"].scale_[numeric_features.index("post_test")]
    + preprocessor["num"].mean_[numeric_features.index("post_test")]
)

treatment_mask = (
    raw_df_treatment_grid.select(pl.col("grade").eq(pl.lit(1))).to_numpy().flatten()
)
treatment_posterior_grade = posterior_predictive_treatment["posterior_predictive"][
    "mu_outcome"
][:, :, treatment_mask]

original_scale_treatment_posterior_grade = (
    treatment_posterior_grade
    * preprocessor["num"].scale_[numeric_features.index("post_test")]
    + preprocessor["num"].mean_[numeric_features.index("post_test")]
)

We now visualize the posterior distribution of the average treatment effect for Grade 1. The first plot shows the effect computed from the counterfactual predictions, while the second shows the direct coefficient estimate. These should be consistent, providing a validation of our approach.

fig, ax = plt.subplots(
    nrows=2,
    ncols=1,
    sharex=True,
    sharey=True,
    figsize=(10, 8),
    layout="constrained",
)


az.plot_posterior(
    (
        original_scale_treatment_posterior_grade
        - original_scale_control_posterior_grade
    ).mean(dim=("obs_idx")),
    ref_val=0,
    ax=ax[0],
)
ax[0].set(title="ATE from Counterfactual Predictions")

az.plot_posterior(
    (
        idata["posterior"]["theta"]
        * preprocessor["num"].scale_[numeric_features.index("post_test")]
    ).sel(grade=1),
    ref_val=0,
    ax=ax[1],
)
ax[1].set(
    title="ATE from Parameter Estimates (beta coefficient)",
    xlabel="ATE (Original Scale)",
);

Both approaches yield similar posterior distributions, confirming the robustness of our causal effect estimate. The positive treatment effect indicates that watching “The Electric Company” improved reading scores for Grade 1 students. The credible intervals quantify our uncertainty about the magnitude of this effect.

Motivation: Modeling Treatment Effect Heterogeneity

The hierarchical intercept model in Part 1 assumes that the treatment effect is constant across all pairs (after accounting for grade differences). While this model efficiently controls for pair-level confounding, it does not allow us to explore whether the treatment effect varies across pairs. In practice, some schools or classrooms may respond more strongly to the intervention than others due to unmeasured factors such as teacher implementation fidelity, student engagement, or local context.

To capture this treatment effect heterogeneity, we extend our model to allow both the baseline intercept \(\alpha_j\) and the treatment effect \(\theta_j\) to vary across pairs. Moreover, we model the joint distribution of these pair-specific parameters using a bivariate normal distribution with a covariance structure:

\[ \begin{pmatrix} \alpha_j \\ \theta_j \end{pmatrix} \sim \text{MvNormal}\left(\begin{pmatrix} \mu_\alpha \\ \mu_\theta \end{pmatrix}, \Sigma\right) \]

where \(\Sigma\) is a \(2 \times 2\) covariance matrix. The off-diagonal elements of \(\Sigma\) capture the correlation between baseline performance and treatment response. For example, a positive correlation would indicate that pairs with higher baseline scores also tend to experience larger treatment effects, while a negative correlation would suggest compensatory effects (treatment helps struggling pairs more).

def vectorized_correlation_matrices(corr_values, size=2):
    n_matrices = corr_values.shape[0]

    # Reshape for broadcasting
    # Use reshape or expand_dims instead of dimshuffle
    corr_expanded = pt.reshape(corr_values, (n_matrices, 1, 1))

    # Create base: all elements are correlation values
    base = corr_expanded * pt.ones((n_matrices, size, size))

    # Create diagonal mask
    diag_mask = pt.eye(size, dtype="bool")

    # Set diagonal to 1
    return pt.where(diag_mask, 1.0, base)


def vectorized_diagonal_matrices_v4(values):
    k = values.shape[1]  # 2

    # Create identity matrix (2, 2)
    identity_matrix = pt.eye(k)

    # Reshape values for broadcasting: (4, 2) -> (4, 2, 1)
    values_expanded = values[:, :, None]

    # Multiply: (4, 2, 1) * (2, 2) -> (4, 2, 2)
    # This puts values[i, j] at position [i, j, j]
    return values_expanded * identity_matrix

Prior Specification for Correlation

Before specifying the full model, we visualize our prior distribution for the correlation parameter. We use a \(\text{Beta}(20, 5)\) distribution (scaled to \([-1, 1]\)) to encode a prior belief that the correlation between intercepts and treatment effects is likely positive but with substantial uncertainty. This prior is weakly informative and allows the data to dominate the posterior inference.

pz.Beta(alpha=20, beta=5).plot_pdf(pointinterval=True);

The Beta distribution shown above is transformed to the correlation scale via the mapping \(\psi = 2 \times \text{Beta}(20, 5) - 1\). This induces a prior that favors positive correlations but remains flexible enough to accommodate a range of values.

Model Specification

We now specify the full covariance model. The key difference from Part 1 is that both the intercept and the treatment effect vary by pair. The model structure is:

\[ \begin{align} \text{post_test}_i &\sim \text{Normal}(\mu_i, \sigma_y) \\ \mu_i &= \alpha_{\text{pair}[i]} + \theta_{\text{pair}[i]} \cdot T_i + \beta_x \cdot \text{pre_test}_i \\ \begin{pmatrix} \alpha_j \\ \theta_j \end{pmatrix} &\sim \text{Normal}\left(\begin{pmatrix} \mu_\alpha \\ \mu_\theta \end{pmatrix}, \Sigma\right) \end{align} \]

where \(\Sigma\) is decomposed as \(\Sigma = D \Omega D\), with \(D\) being a diagonal matrix of standard deviations and \(\Omega\) being the correlation matrix. This decomposition allows us to place separate priors on the marginal variances and the correlation structure, which improves interpretability and sampling efficiency.

As in the model in Part 1, we use a non-centered parametrization to improve sampling efficiency.

coords.update({"effect": ["intercept", "slope"], "effect_copy": ["intercept", "slope"]})
coords.update({"corr_dim": ["corr_dim_1"]})

with pm.Model(coords=coords) as cov_model:
    # --- Data Containers ---
    # covariates
    x_data = pm.Data("x_data", x_df, dims=("obs_idx", "covariates"))
    # grade
    grade_idx_data = pm.Data("grade_idx_data", df["grade"].to_numpy(), dims="obs_idx")
    # object categories
    pair_idx_data = pm.Data("pair_idx_data", df["pair_id"].to_numpy(), dims="obs_idx")
    # treatment
    treatment_data = pm.Data(
        "treatment_data", df["treatment"].to_numpy(), dims=("obs_idx")
    )
    # outcome
    post_test_data = pm.Data(
        "post_test_data", df["post_test"].to_numpy(), dims="obs_idx"
    )

    # --- Priors ---

    beta_x = pm.Normal("beta_x", mu=0, sigma=1, dims=("grade", "covariates"))
    sigma_outcome = pm.HalfNormal("sigma_outcome", sigma=1, dims=("grade"))

    mu_alpha = pm.Normal("mu_alpha", mu=0, sigma=0.5, dims=("grade"))
    mu_theta = pm.Normal("mu_theta", mu=0, sigma=0.5, dims=("grade"))

    # Group-level standard deviations
    sigma_u = pm.HalfNormal(
        "sigma_u", sigma=np.array([0.2, 0.2]), dims=("grade", "effect")
    )

    # Triangular upper part of the correlation matrix
    # omega_triu = pm.LKJCorr("omega_triu", eta=1, n=2, dims=("grade", "corr_dim")) <- Not supported yet! # noqa: E501
    omega_triu = pm.Beta("omega_triu", alpha=20, beta=5, dims=("grade", "corr_dim"))
    omega_triu_scaled = pm.Deterministic(
        "omega_triu_scaled", omega_triu * 2 - 1, dims=("grade", "corr_dim")
    )

    # Construct correlation matrix
    omega = pm.Deterministic(
        "omega",
        vectorized_correlation_matrices(omega_triu_scaled),
        dims=("grade", "effect", "effect_copy"),
    )

    # Construct diagonal matrix of standard deviation
    sigma_diagonal = pm.Deterministic(
        "sigma_diagonal",
        vectorized_diagonal_matrices_v4(sigma_u),
        dims=("grade", "effect", "effect_copy"),
    )

    # Compute covariance matrix
    cov = pm.Deterministic(
        "cov",
        pt.einsum("bij,bjk,bkl->bil", sigma_diagonal, omega, sigma_diagonal),
        dims=("grade", "effect", "effect_copy"),
    )

    # Cholesky decomposition of covariance matrix
    cholesky_cov = pm.Deterministic(
        "cholesky_cov",
        pt.slinalg.cholesky(cov),
        dims=("grade", "effect", "effect_copy"),
    )

    # And finally get group-specific coefficients
    u_raw = pm.Normal("u_raw", mu=0, sigma=1, dims=("grade", "effect", "pair_id"))
    u = pm.Deterministic(
        "u",
        pt.einsum("bik,bkj->bji", cholesky_cov, u_raw),
        dims=("grade", "pair_id", "effect"),
    )

    # Extract the intercept and slope components deviations from the population means
    u0 = pm.Deterministic("u0", u[:, :, 0], dims=("grade", "pair_id"))
    u1 = pm.Deterministic("u1", u[:, :, 1], dims=("grade", "pair_id"))

    # Extract the intercept and slope components
    alpha = pm.Deterministic("alpha", mu_alpha + u0.T, dims=("pair_id", "grade"))
    theta = pm.Deterministic("theta", mu_theta + u1.T, dims=("pair_id", "grade"))

    mu_outcome = pm.Deterministic(
        "mu_outcome",
        alpha[pair_idx_data, grade_idx_data]
        + theta[pair_idx_data, grade_idx_data] * treatment_data
        + (beta_x[grade_idx_data] * x_data).sum(axis=-1),
        dims=("obs_idx"),
    )

    # --- Likelihood ---
    pm.Normal(
        "post_test_obs",
        mu=mu_outcome,
        sigma=sigma_outcome[grade_idx_data],
        observed=post_test_data,
        dims="obs_idx",
    )

pm.model_to_graphviz(cov_model)

Prior Predictive Check

As in Part 1, we sample from the prior distribution to verify that our priors produce reasonable predictions before observing the data. This is especially important for the covariance model, where the additional flexibility could lead to implausible outcomes if priors are poorly calibrated.

with cov_model:
    cov_idata = pm.sample_prior_predictive(random_seed=rng)

fig, ax = plt.subplots()
az.plot_ppc(cov_idata, group="prior", ax=ax)
az.plot_dist(df["post_test"].to_numpy(), color="black", ax=ax)
ax.set_title(
    "Prior Predictive Check - Covariance Model", fontsize=18, fontweight="bold", y=1.02
);

This prior predictive distribution looks quite reasonable.

Posterior Inference

We now fit the covariance model using HMC. Due to the increased complexity of the model (additional parameters and correlation structure), we use a longer tuning phase (2,000 iterations) and a higher target acceptance rate (0.95) to ensure thorough exploration of the posterior geometry.

with cov_model:
    cov_idata.extend(
        pm.sample(
            tune=2_000,
            draws=1_000,
            chains=4,
            nuts_sampler="numpyro",
            target_accept=0.95,
            random_seed=rng,
        )
    )

    cov_idata.extend(pm.sample_posterior_predictive(cov_idata))

Posterior sampling is complete without any issues.

Model Diagnostics

We examine diagnostic statistics to ensure the sampler converged successfully. The covariance model has a more complex posterior geometry, so careful attention to diagnostics is essential.

cov_idata["sample_stats"]["diverging"].sum().item()

0

The absence of divergent transitions indicates successful sampling. The more complex model structure did not cause problematic posterior geometry, which validates our use of the non-centered parametrization and appropriate tuning parameters.

Next, let’s check the posterior predictive distribution.

fig, ax = plt.subplots()
az.plot_ppc(cov_idata, group="posterior", num_pp_samples=500, ax=ax)
ax.set_title(
    "Posterior Predictive Check - Covariance Model",
    fontsize=18,
    fontweight="bold",
    y=1.02,
);

Overall, the posterior predictive distribution looks good.

As the key ingredient of the covariance model is the correlation parameter, let’s visualize the posterior distribution of the correlation parameter and compare it to the prior.

axes = az.plot_dist_comparison(
    cov_idata,
    var_names=["omega_triu_scaled"],
    figsize=(10, 18),
)
plt.gcf().suptitle(
    "Prior vs Posterior Comparison - Covariance Model",
    fontsize=18,
    fontweight="bold",
    y=1.02,
);

We do not see much difference between the prior and posterior distribution of the correlation parameter. Hence, this shows the covariance structure is not overly influential in the average treatment effect for this example. This is consistent with Gelman and Hill’s findings in their book.

Treatment Effect Estimates

We now examine the estimated treatment effects from the covariance model. Unlike Part 1, where we had a single treatment effect parameter per grade, here we have a population-level mean treatment effect (mu_theta) and pair-specific deviations from this mean (theta). The parameter mu_theta represents the average causal effect across all pairs within a grade.

fig, ax = plt.subplots(
    nrows=2,
    ncols=1,
    sharex=True,
    sharey=True,
    figsize=(10, 8),
    layout="constrained",
)

az.plot_forest(
    idata["posterior"]["theta"]
    * preprocessor["num"].scale_[numeric_features.index("post_test")],
    combined=True,
    var_names=["theta"],
    ax=ax[0],
)
ax[0].set(title="Hierarchical Intercept Model")

az.plot_forest(
    cov_idata["posterior"]["mu_theta"]
    * preprocessor["num"].scale_[numeric_features.index("post_test")],
    combined=True,
    var_names=["mu_theta"],
    ax=ax[1],
)
ax[1].set(title="Covariance Model")
fig.suptitle(
    r"Treatment Effect Estimates $(94\% HDI)$ (Original Scale)",
    fontsize=18,
    fontweight="bold",
    y=1.05,
);

The forest plot shows the posterior distribution of the average treatment effect (mu_theta) for each grade. These estimates are comparable to the theta estimates from Part 1, but they now represent the mean of a distribution of pair-specific effects rather than a fixed common effect.

Hierarchical Shrinkage: Population vs. Group-Level Effects

One of the key advantages of the covariance model is that it allows us to estimate both population-level parameters (the means \(\mu_\alpha\) and \(\mu_\theta\)) and group-level parameters (the pair-specific \(\alpha_j\) and \(\theta_j\)). Through partial pooling, the model borrows strength across pairs: estimates for pairs with sparse data are shrunk toward the population mean, while estimates for data-rich pairs are allowed to deviate more.

The visualization below compares population-level and group-level estimates for Grade 1. The top row shows the population means (\(\mu_\alpha\) and \(\mu_\theta\)), while the bottom row shows the distribution of pair-specific effects (\(\alpha_j\) and \(\theta_j\)) for pairs within Grade 1. The degree of shrinkage depends on the estimated between-pair variance: if pairs are highly similar, estimates will be heavily pooled; if pairs differ substantially, estimates will retain more individual variation.

fig, ax = plt.subplots(
    nrows=2,
    ncols=2,
    height_ratios=[0.3, 1],
    figsize=(15, 8),
    sharex=True,
    sharey=False,
    layout="constrained",
)

az.plot_forest(
    cov_idata["posterior"].sel(grade=1),
    combined=True,
    var_names=["mu_alpha"],
    colors="C0",
    ax=ax[0, 0],
)
ax[0, 0].set_title("Intercepts Population-Level Estimates (Grade 1)", fontsize=14)


az.plot_forest(
    cov_idata["posterior"]
    .sel(grade=1)
    .where(
        cov_idata["posterior"].pair_id.isin(
            raw_df.group_by("grade")
            .agg(pl.col("pair_id").unique())
            .filter(pl.col("grade").eq(pl.lit(1)))["pair_id"]
            .to_list()
        ),
        drop=True,
    ),
    var_names=["alpha"],
    combined=True,
    colors="C1",
    ax=ax[1, 0],
)
ax[1, 0].set_title("Intercepts Group-Level Estimates (Grade 1)", fontsize=14)


az.plot_forest(
    cov_idata["posterior"].sel(grade=1),
    combined=True,
    var_names=["mu_theta"],
    colors="C0",
    ax=ax[0, 1],
)
ax[0, 1].set_title("Treatment Effect Population-Level Estimates (Grade 1)", fontsize=14)

az.plot_forest(
    cov_idata["posterior"]
    .sel(grade=1)
    .where(
        cov_idata["posterior"].pair_id.isin(
            raw_df.group_by("grade")
            .agg(pl.col("pair_id").unique())
            .filter(pl.col("grade").eq(pl.lit(1)))["pair_id"]
            .to_list()
        ),
        drop=True,
    ),
    var_names=["theta"],
    combined=True,
    colors="C1",
    ax=ax[1, 1],
)
ax[1, 1].set_title("Treatment Effect Group-Level Estimates (Grade 1)", fontsize=14)
fig.suptitle(
    "Hierarchical Shrinkage: Population vs. Group-Level Effects (Grade 1)",
    fontsize=18,
    fontweight="bold",
    y=1.07,
);

References

• Gelman, A., & Hill, J. (2006). Data Analysis Using Regression and Multilevel/Hierarchical Models. Cambridge University Press. Chapter 23: Causal Inference Using Multilevel Models.
• Huntington-Klein, N. (2021). The Effect: An Introduction to Research Design and Causality. Chapman and Hall/CRC. Chapter 16: Fixed Effects. Available online at https://theeffectbook.net/ch-FixedEffects.html