Causal Inference in Python - Part 3: Effect Heterogeneity and Personalization

import numpy as np
import pandas as pd

np.random.seed(42)
n = 2000

struggling = np.random.binomial(1, 0.4, n)
prior_grade = np.clip(np.random.normal(70 - 12*struggling, 10, n), 0, 100)
grade_level = np.random.choice([7, 8, 9], n)
tutoring = np.random.binomial(1, 0.5, n)

grade_level_effect = dict(zip([7,8,9], [0, 2, 4]))
baseline = 0.8*prior_grade + 20 - 14*struggling + np.vectorize(grade_level_effect.get)(grade_level)

tent = np.maximum(0, 1 - np.abs(prior_grade - 65) / 12)
cate_s = 6 + 10 * tent

quad = np.maximum(0, 1 - ((prior_grade - 72) / 15)**2)
cate_ns = 2 + 4 * (grade_level == 9) * quad

true_cate = np.where(struggling, cate_s, cate_ns)
final_grade = np.clip(baseline + true_cate*tutoring + np.random.normal(0, 5, n), 0, 100)

df = pd.DataFrame(dict(struggling=struggling, prior_grade=prior_grade.round(1),
    grade_level=grade_level, tutoring=tutoring,
    true_cate=true_cate.round(3), final_grade=final_grade.round(1)))
df.head()

import matplotlib.pyplot as plt
from itertools import product

fig, ax = plt.subplots(figsize=(9, 4))
colors = {(0,7):'#aec6e8', (0,8):'#4a90d9', (0,9):'#1a4f8a', (1,7):'#f4a582', (1,8):'#d6604d', (1,9):'#8b1a1a'}
markers = {7:'o', 8:'s', 9:'^'}
for s, g in product([0,1], [7,8,9]):
    mask = (df.struggling==s) & (df.grade_level==g)
    label = f"{'Struggling' if s else 'Not struggling'}, G{g}"
    ax.scatter(df.loc[mask,'prior_grade'], df.loc[mask,'true_cate'],
               alpha=0.4, s=12, color=colors[(s,g)], marker=markers[g], label=label)
ax.set(xlabel='Prior Grade', ylabel='True CATE', title='True CATE vs Prior Grade')
ax.legend(fontsize=7, ncol=2)
plt.tight_layout()
plt.show()

print(df.final_grade.describe())
print(df.groupby('struggling').final_grade.mean())

from sklearn.ensemble import RandomForestRegressor

X,y,T =['struggling', 'prior_grade', 'grade_level'], 'final_grade', 'tutoring'

rf = RandomForestRegressor(random_state=42).fit(df[X + [T]], df[y])
df['ml_pred'] = rf.predict(df[X + [T]])

m0 = RandomForestRegressor(random_state=42).fit(df.query(f"{T}==0")[X], df.query(f"{T}==0")[y])
m1 = RandomForestRegressor(random_state=42).fit(df.query(f"{T}==1")[X], df.query(f"{T}==1")[y])
df['t_learner_pred'] = m1.predict(df[X]) - m0.predict(df[X])

df['random_pred'] = np.random.default_rng(42).uniform(-10, 10, len(df))

df.head()

from fklearn.causal.validation.curves import effect_by_segment

fig, axs = plt.subplots(1, 3, figsize=(12, 4), sharey=True)
for ax, m in zip(axs,['random_pred', 'ml_pred', 't_learner_pred']): 
    eff = effect_by_segment(df, 'tutoring', 'final_grade', m)
    ax.bar(range(len(eff)), eff.values)
    ax.set(title=m, xlabel='Quantile')
axs[0].set_ylabel('Estimated Effect')
plt.tight_layout()

from fklearn.causal.validation.curves import cumulative_gain_curve, relative_cumulative_gain_curve

fig, axs = plt.subplots(1, 2, figsize=(12, 4))
for m in['random_pred', 'ml_pred', 't_learner_pred']:
    cg = cumulative_gain_curve(df, 'tutoring', 'final_grade', m)
    rcg = relative_cumulative_gain_curve(df, 'tutoring', 'final_grade', m)
    x_pct = np.linspace(0, 100, len(cg))
    axs[0].plot(x_pct, cg, label=m)
    axs[1].plot(x_pct, rcg, label=m)
axs[0].plot([0, 100],[0, cumulative_gain_curve(df, 'tutoring', 'final_grade', 't_learner_pred')[-1]], 'k--', label='baseline')
axs[1].axhline(0, color='k', linestyle='--', label='baseline')
axs[0].set(xlabel='Top %', ylabel='Cumulative Gain'); axs[1].set(xlabel='Top %', ylabel='Relative Cumulative Gain')
for ax in axs: ax.legend()
plt.tight_layout()

def auc(df, t, y, m): return relative_cumulative_gain_curve(df, t, y, m).mean()
for m in ['random_pred', 'ml_pred', 't_learner_pred']: print(f"{m} AUC: {auc(df, 'tutoring', 'final_grade', m):.2f}")

import statsmodels.formula.api as smf
from sklearn.metrics import mean_squared_error

# 1. Isolate the residuals for outcome and treatment

X_cols = "struggling + prior_grade + C(grade_level)"
y_res = smf.ols(f"final_grade ~ {X_cols}", data=df).fit().resid
t_res = smf.ols(f"tutoring ~ {X_cols}", data=df).fit().resid

# 2. Construct the Y* target and the stabilizing weights

y_star = y_res / t_res
weights = t_res**2

# 3. Calculate weighted MSE for each model

for m in ['random_pred', 'ml_pred', 't_learner_pred']:
    wmse = mean_squared_error(y_star, df[m], sample_weight=weights)
    print(f"{m} Weighted MSE: {wmse:.2f}")