Skip to content

Instantly share code, notes, and snippets.

@martinctc
Last active September 11, 2025 09:33
Show Gist options
  • Select an option

  • Save martinctc/f3d6134a32c919893556b4cbbb017555 to your computer and use it in GitHub Desktop.

Select an option

Save martinctc/f3d6134a32c919893556b4cbbb017555 to your computer and use it in GitHub Desktop.
[Python and R code for Causal Inference plots] #Python #R

Forest plot

import matplotlib.pyplot as plt
import numpy as np

interventions = ['Leadership Digest', 'Peer Nudges', 'Training Email']
ATEs = [0.045, 0.032, 0.018]
lower_CIs = [0.015, 0.005, -0.004]
upper_CIs = [0.075, 0.059, 0.040]

error_lower = [ate - lower for ate, lower in zip(ATEs, lower_CIs)]
error_upper = [upper - ate for ate, upper in zip(ATEs, upper_CIs)]

fig, ax = plt.subplots(figsize=(16, 9))
ax.errorbar(ATEs, interventions, xerr=[error_lower, error_upper],
            fmt='o', color='blue', ecolor='gray', capsize=5)
ax.axvline(x=0, color='black', linestyle='--')
ax.set_xlabel('Average Treatment Effect (ATE)')
ax.set_title('Forest Plot of Interventions with 95% Confidence Intervals')
plt.tight_layout()
plt.show()
df_forest <- data.frame(
  Intervention = c("Leadership Digest", "Peer Nudges", "Training Email"),
  ATE = c(0.045, 0.032, 0.018),
  Lower = c(0.015, 0.005, -0.004),
  Upper = c(0.075, 0.059, 0.040)
)

ggplot(df_forest, aes(x = ATE, y = Intervention)) +
  geom_point(size = 3, colour = "blue") +
  geom_errorbarh(aes(xmin = Lower, xmax = Upper), height = 0.2, colour = "grey50") +
  geom_vline(xintercept = 0, linetype = "dashed") +
  labs(title = "Forest Plot of Interventions with 95% CI",
       x = "Average Treatment Effect (ATE)", y = "") +
  theme_minimal(base_size = 16)

Event Study DiD

weeks_event = np.arange(-4, 9)
coefficients = np.array([0.01, -0.005, 0.002, -0.003, 0.0,
                          0.015, 0.025, 0.035, 0.045, 0.055,
                          0.065, 0.075, 0.085])
ci_lower = coefficients - 0.01
ci_upper = coefficients + 0.01

fig, ax = plt.subplots(figsize=(16, 9))
ax.plot(weeks_event, coefficients, marker='o', label='Estimated Effect')
ax.fill_between(weeks_event, ci_lower, ci_upper, color='blue', alpha=0.2, label='95% CI')
ax.axvline(x=0, color='red', linestyle='--', label='Intervention Point')
ax.axhline(y=0, color='gray', linestyle='--')
ax.set_title('Event Study (DiD): Estimated Effect Over Relative Weeks')
ax.set_xlabel('Relative Week')
ax.set_ylabel('Estimated Effect')
ax.legend()
ax.grid(True)
plt.tight_layout()
plt.show()
weeks <- -4:8
coeff <- c(0.01, -0.005, 0.002, -0.003, 0.0,
           0.015, 0.025, 0.035, 0.045, 0.055,
           0.065, 0.075, 0.085)
df_event <- data.frame(Week = weeks, Coef = coeff)
df_event <- df_event %>%
  mutate(Lower = Coef - 0.01, Upper = Coef + 0.01)

ggplot(df_event, aes(x = Week, y = Coef)) +
  geom_ribbon(aes(ymin = Lower, ymax = Upper), fill = "blue", alpha = 0.2) +
  geom_line(colour = "blue") +
  geom_point() +
  geom_vline(xintercept = 0, linetype = "dashed", colour = "red") +
  geom_hline(yintercept = 0, linetype = "dashed", colour = "grey50") +
  labs(title = "Event Study (DiD): Estimated Effect Over Relative Weeks",
       x = "Relative Week", y = "Estimated Effect") +
  theme_minimal(base_size = 

ITSA

weeks_its = np.arange(1, 21)
actual = np.concatenate([np.linspace(10, 15, 10), np.linspace(17, 27, 10)])
counterfactual = np.concatenate([np.linspace(10, 15, 10), np.linspace(15.5, 20.5, 10)])

fig, ax = plt.subplots(figsize=(16, 9))
ax.plot(weeks_its, actual, label='Actual Outcome', marker='o')
ax.plot(weeks_its, counterfactual, label='Counterfactual Outcome', linestyle='--', marker='x')
ax.axvline(x=10.5, color='red', linestyle='--', label='Intervention Point')
ax.set_title('Interrupted Time Series: Actual vs Counterfactual Outcome')
ax.set_xlabel('Week')
ax.set_ylabel('Outcome Metric')
ax.legend()
ax.grid(True)
plt.tight_layout()
plt.show()
weeks_its <- 1:20
actual <- c(seq(10, 15, length.out = 10), seq(17, 27, length.out = 10))
counterfactual <- c(seq(10, 15, length.out = 10), seq(15.5, 20.5, length.out = 10))
df_its <- data.frame(Week = weeks_its,
                     Actual = actual,
                     Counterfactual = counterfactual)

ggplot(df_its, aes(x = Week)) +
  geom_line(aes(y = Actual, colour = "Actual"), size = 1) +
  geom_point(aes(y = Actual), colour = "blue") +
  geom_line(aes(y = Counterfactual, colour = "Counterfactual"), linetype = "dashed", size = 1) +
  geom_point(aes(y = Counterfactual), colour = "grey50") +
  geom_vline(xintercept = 10.5, linetype = "dashed", colour = "red") +
  labs(title = "Interrupted Time Series: Actual vs Counterfactual Outcome",
       x = "Week", y = "Outcome Metric", colour = "") +
  theme_minimal(base_size = 16)

CATE heatmap

# After creating the heatmap and annotations as before:
for i, seg in enumerate(pivot.index):
    for j, intv in enumerate(pivot.columns):
        if df[(df['Segment']==seg) & (df['Intervention']==intv)]['Significant'].values[0]:
            # Position marker slightly to the right of the text
            ax.scatter(j + 0.9, i + 0.5, s=90, facecolors='white',
                       edgecolors='black', linewidths=0.8)

# Add a custom legend for significance
from matplotlib.lines import Line2D
legend_elements = [Line2D([0], [0], marker='o', color='w', label='Significant (95% CI)',
                          markerfacecolor='white', markeredgecolor='black', markersize=10)]
ax.legend(handles=legend_elements, loc='upper right', frameon=True)
# Python 3.x
# Dependencies: pandas, numpy, seaborn, matplotlib

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import pandas as pd
from matplotlib.lines import Line2D

# -------------------------------
# 1) Simulated CATE inputs
#    (Δ weekly Copilot actions per user)
# -------------------------------
segments = [
    "Managers",
    "Individual Contributors",
    "New Hires (≤6 mo)",
    "Tenured (>2 yrs)",
    "Customer-facing",
    "Internal Platform",
    "EMEA",
    "AMER",
    "APAC",
]
interventions = ["Leadership Digest", "Peer Nudges", "Training Email"]

# Fabricated CATE matrix (rows = segments, cols = interventions)
CATE = np.array([
    [0.060, 0.050, 0.020],  # Managers
    [0.020, 0.030, 0.010],  # ICs
    [0.050, 0.070, 0.030],  # New Hires
    [0.020, 0.020, 0.000],  # Tenured
    [0.050, 0.060, 0.020],  # Customer-facing
    [0.010, 0.015, 0.000],  # Internal Platform
    [0.030, 0.040, 0.015],  # EMEA
    [0.040, 0.050, 0.020],  # AMER
    [0.025, 0.035, 0.010],  # APAC
])

# Significance mask (True if 95% CI excludes 0) — simulated
sig = np.array([
    [ True,  True, False],
    [False,  True, False],
    [ True,  True,  True],
    [False, False, False],
    [ True,  True, False],
    [False, False, False],
    [False,  True, False],
    [ True,  True, False],
    [False,  True, False],
])

# -------------------------------
# 2) Build plotting DataFrame
# -------------------------------
rows = []
for i, seg in enumerate(segments):
    for j, intv in enumerate(interventions):
        rows.append({
            "Segment": seg,
            "Intervention": intv,
            "CATE": CATE[i, j],
            "Significant": sig[i, j],
        })
df = pd.DataFrame(rows)

# Matrix for heatmap, and annotation strings
pivot = df.pivot(index="Segment", columns="Intervention", values="CATE")
annot = np.array([[f"{v:+.3f}" for v in row] for row in pivot.values])

# -------------------------------
# 3) Plot — 16×9; marker at right edge
# -------------------------------
plt.figure(figsize=(16, 9))
ax = sns.heatmap(
    pivot,
    annot=annot, fmt="",
    cmap=sns.diverging_palette(220, 20, as_cmap=True),
    center=0,
    linewidths=0.5, linecolor="white",
    cbar_kws={"label": "CATE (Δ weekly Copilot actions per user)"},
)

# Move significance markers to the RIGHT side of each cell
for i, seg in enumerate(pivot.index):
    for j, intv in enumerate(pivot.columns):
        is_sig = df[(df["Segment"] == seg) & (df["Intervention"] == intv)]["Significant"].values[0]
        if is_sig:
            # x: j + 0.90 => right side; y: i + 0.50 => vertical center
            ax.scatter(j + 0.90, i + 0.50,
                       s=90, facecolors="white", edgecolors="black", linewidths=0.8)

# Legend for significance
legend_elements = [Line2D([0], [0], marker="o", color="w",
                          label="Significant (95% CI)",
                          markerfacecolor="white", markeredgecolor="black",
                          markersize=10)]
ax.legend(handles=legend_elements, loc="upper right", frameon=True)

# Titles & labels
plt.suptitle("Where do interventions work best? (CATE heatmap)", fontsize=20, y=0.95)
plt.title("Estimated heterogeneous treatment effects by segment and intervention\n"
          "Circles mark effects whose 95% CI excludes zero (simulated data)", fontsize=12)
ax.set_xlabel("Intervention")
ax.set_ylabel("Segment")

# Layout & save
plt.tight_layout(rect=[0, 0, 1, 0.93])
plt.savefig("cate_heatmap_16x9_side_marker.png", dpi=200)
plt.show()

# -------------------------------
# Notes for real data:
# - Replace CATE with your estimates.
# - If you have standard errors (se), significance can be set via:
#     sig = (np.abs(CATE) - 1.96 * se) > 0
# - You can push the marker slightly outside with j + 1.02 if you prefer.
# -------------------------------

Dose-Response curve (GPS)

import numpy as np, pandas as pd, matplotlib.pyplot as plt
from scipy.stats import poisson
import statsmodels.api as sm

# 1) Simulate synthetic data
def simulate_data(n=2500, max_emails=8, seed=123):
    rng = np.random.default_rng(seed)
    prior = np.clip(rng.normal(15, 5, n), 0, None)
    tenure = rng.uniform(0, 10, n)
    culture = rng.binomial(1, 0.4, n)
    lam = np.exp(-0.5 + 0.04*prior + 0.05*tenure + 0.4*culture)
    T = np.clip(rng.poisson(lam), 0, max_emails)
    gT = 1.6*T - 0.13*(T**2)
    y = 8 + gT + 0.6*prior + 0.25*tenure + 1.2*culture + rng.normal(0, 2.0, n)
    return pd.DataFrame(dict(y=y, T=T, prior=prior, tenure=tenure, culture=culture))

# 2) Poisson treatment model → GPS
def fit_gps_poisson(df):
    X = sm.add_constant(df[['prior','tenure','culture']])
    m = sm.GLM(df['T'], X, family=sm.families.Poisson()).fit()
    lam_hat = np.exp(m.predict(X))
    return m, lam_hat

# 3) Outcome stage
def dose_response(df, lam_hat, t_grid):
    r_obs = poisson.pmf(df['T'].values, lam_hat)
    Z = pd.DataFrame({'const':1,'T':df['T'],'T2':df['T']**2,'r':r_obs,'r2':r_obs**2,'Tr':df['T']*r_obs})
    ols = sm.OLS(df['y'], Z).fit()
    mu = []
    for t in t_grid:
        r_t = poisson.pmf(t, lam_hat)
        Zt = pd.DataFrame({'const':1,'T':t,'T2':t**2,'r':r_t,'r2':r_t**2,'Tr':t*r_t})
        mu.append(ols.predict(Zt).mean())
    return np.array(mu)

# 4) Bootstrap CIs
def bootstrap_mu(df, t_grid, B=150, seed=42):
    rng = np.random.default_rng(seed)
    n, out = len(df), np.empty((B, len(t_grid)))
    for b in range(B):
        idx = rng.integers(0, n, n)
        m_b, lam_b = fit_gps_poisson(df.iloc[idx])
        out[b] = dose_response(df.iloc[idx], lam_b, t_grid)
    mu = out.mean(0)
    l, u = np.percentile(out, [2.5, 97.5], axis=0)
    return mu, l, u

# Run
df = simulate_data()
m, lam = fit_gps_poisson(df)
t_grid = np.arange(0, 9)
mu, lo, hi = bootstrap_mu(df, t_grid)

# Plot (16:9)
plt.figure(figsize=(12.8, 7.2))
plt.fill_between(t_grid, lo, hi, color="#4285F4", alpha=.18, label="95% CI (GPS)")
plt.plot(t_grid, mu, color="#1a5fb4", lw=3, label="Estimated dose–response (GPS)")
plt.title("Training Emails → Copilot Usage: Dose–Response")
plt.xlabel("Training emails received (count)")
plt.ylabel("Copilot usage (e.g., actions/week)")
plt.xticks(t_grid)
plt.grid(alpha=.25)
plt.legend()
plt.tight_layout()
set.seed(123)
n <- 2500; max_emails <- 8

# Simulate
prior <- pmax(rnorm(n, 15, 5), 0)
tenure <- runif(n, 0, 10)
culture <- rbinom(n, 1, 0.4)
lambda <- exp(-0.5 + 0.04*prior + 0.05*tenure + 0.4*culture)
Tcount <- pmin(rpois(n, lambda), max_emails)
gT <- 1.6*Tcount - 0.13*Tcount^2
y <- 8 + gT + 0.6*prior + 0.25*tenure + 1.2*culture + rnorm(n,0,2)
df <- data.frame(y, T=Tcount, prior, tenure, culture)

# Poisson treatment model
m <- glm(T ~ prior + tenure + culture, data=df, family=poisson())
lam_hat <- fitted(m)

# Outcome stage
pmf <- function(t, lam) dpois(t, lam)
r_obs <- pmf(df$T, lam_hat)
Z <- data.frame(T=df$T, T2=df$T^2, r=r_obs, r2=r_obs^2, Tr=df$T*r_obs)
ols <- lm(y ~ T + T2 + r + r2 + Tr, data=cbind(df, Z))

mu_fun <- function(t){
  r_t <- pmf(t, lam_hat)
  Zt <- data.frame(T=t, T2=t^2, r=r_t, r2=r_t^2, Tr=t*r_t)
  mean(predict(ols, newdata=Zt))
}

t_grid <- 0:max_emails
mu_hat <- sapply(t_grid, mu_fun)

# Bootstrap CIs
B <- 150
boot_mat <- matrix(NA_real_, nrow=B, ncol=length(t_grid))
for(b in 1:B){
  idx <- sample.int(n, n, replace=TRUE)
  dFb <- df[idx,]
  mb <- glm(T ~ prior + tenure + culture, data=dFb, family=poisson())
  lam_b <- fitted(mb)
  r_b <- dpois(dFb$T, lam_b)
  Zb <- data.frame(T=dFb$T, T2=dFb$T^2, r=r_b, r2=r_b^2, Tr=dFb$T*r_b)
  ols_b <- lm(y ~ T + T2 + r + r2 + Tr, data=cbind(dFb, Zb))
  mu_b <- sapply(t_grid, function(t){
    r_t <- dpois(t, lam_b)
    mean(predict(ols_b, newdata=data.frame(T=t, T2=t^2, r=r_t, r2=r_t^2, Tr=t*r_t)))
  })
  boot_mat[b,] <- mu_b
}
ci_l <- apply(boot_mat, 2, quantile, 0.025)
ci_u <- apply(boot_mat, 2, quantile, 0.975)

# Plot
library(ggplot2)
ggplot() +
  geom_ribbon(aes(x=t_grid, ymin=ci_l, ymax=ci_u), fill="#4285F4", alpha=.18) +
  geom_line(aes(x=t_grid, y=mu_hat), color="#1a5fb4", linewidth=1.2) +
  scale_x_continuous(breaks=t_grid) +
  labs(title="Training Emails → Copilot Usage: Dose–Response",
       x="Training emails received (count)", y="Copilot usage (e.g., actions/week)") +
  theme_minimal(base_size = 12)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment