xjtulyc/causal-inference

Causal Inference with DoWhy and CausalML

Causal inference goes beyond statistical association to estimate the effect of interventions. This skill covers the complete causal inference workflow: DAG construction, identification via backdoor/frontdoor criteria, estimation with multiple methods, refutation testing, propensity score matching, and uplift modeling.

---

1. Core Concepts

The potential outcomes framework defines causal effects:

• ATE: Average Treatment Effect = E[Y(1) - Y(0)]
• ATT: Average Treatment Effect on the Treated
• CATE: Conditional ATE (heterogeneous treatment effects)
• Backdoor criterion: A set Z satisfies it if Z blocks all backdoor paths from T to Y
• Frontdoor criterion: Applicable when all paths go through mediator M
• Confounding: Variables that affect both treatment and outcome

import numpy as np
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import GradientBoostingRegressor


def generate_confounded_data(n=2000, seed=42):
    """
    Generate synthetic observational data with confounding.

    Data-generating process:
        U1, U2 ~ N(0, 1)   (unobserved confounders — for illustration)
        X1 = U1 + noise     (confounder 1: ability/IQ proxy)
        X2 = U2 + noise     (confounder 2: family income)
        T = sigmoid(0.8*X1 + 0.6*X2 - 1.0) > Bernoulli  (treatment: college)
        Y = 2.5*T + 1.2*X1 + 0.9*X2 + noise             (outcome: income)

    True ATE = 2.5 (units of income per treatment unit)
    """
    rng = np.random.default_rng(seed)

    X1 = rng.standard_normal(n)  # ability proxy
    X2 = rng.standard_normal(n)  # family income proxy
    X3 = rng.standard_normal(n)  # irrelevant covariate

    # Propensity score model
    log_odds = 0.8 * X1 + 0.6 * X2 - 1.0
    propensity = 1 / (1 + np.exp(-log_odds))
    T = (rng.uniform(size=n) < propensity).astype(int)

    # Outcome model (true ATE = 2.5)
    noise = 0.5 * rng.standard_normal(n)
    Y = 2.5 * T + 1.2 * X1 + 0.9 * X2 + noise

    df = pd.DataFrame({
        "treatment": T,
        "outcome": Y,
        "ability": X1,
        "family_income": X2,
        "noise_var": X3,
        "propensity_true": propensity,
    })
    return df


print("True ATE = 2.5")
df = generate_confounded_data(n=3000)
print(f"Dataset shape: {df.shape}")
print(f"Treatment prevalence: {df.treatment.mean():.3f}")
print(f"Naive difference in means: {df.groupby('treatment')['outcome'].mean().diff().iloc[-1]:.4f}")

---

2. DAG Construction

import networkx as nx
import matplotlib.pyplot as plt


def build_causal_dag(nodes, edges, latent_nodes=None):
    """
    Build a causal DAG from node/edge lists.

    Parameters
    ----------
    nodes : list of str
        Variable names in the DAG.
    edges : list of (str, str)
        Directed edges (cause, effect).
    latent_nodes : list of str, optional
        Nodes representing unobserved variables.

    Returns
    -------
    G : nx.DiGraph
        Directed acyclic graph with node attributes.
    """
    G = nx.DiGraph()
    G.add_nodes_from(nodes)
    G.add_edges_from(edges)

    # Validate acyclicity
    if not nx.is_directed_acyclic_graph(G):
        raise ValueError("Graph contains cycles — not a valid DAG.")

    # Mark node types
    if latent_nodes:
        for node in latent_nodes:
            G.nodes[node]["latent"] = True

    return G


def plot_dag(G, treatment=None, outcome=None, confounders=None, title="Causal DAG"):
    """
    Visualize a causal DAG with color-coded nodes.

    Color convention:
        Red   = treatment
        Green = outcome
        Orange = confounder/covariate
        Gray  = other
    """
    node_colors = []
    for node in G.nodes():
        if treatment and node == treatment:
            node_colors.append("tomato")
        elif outcome and node == outcome:
            node_colors.append("mediumseagreen")
        elif confounders and node in confounders:
            node_colors.append("darkorange")
        else:
            node_colors.append("lightgray")

    pos = nx.spring_layout(G, seed=7, k=2.0)
    fig, ax = plt.subplots(figsize=(9, 6))
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=1800, ax=ax, alpha=0.9)
    nx.draw_networkx_labels(G, pos, font_size=10, font_weight="bold", ax=ax)
    nx.draw_networkx_edges(G, pos, ax=ax, arrows=True,
                           arrowstyle="-|>", arrowsize=20,
                           edge_color="steelblue", width=2)
    ax.set_title(title, fontsize=13, fontweight="bold")
    ax.axis("off")

    # Legend
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(color="tomato", label="Treatment"),
        Patch(color="mediumseagreen", label="Outcome"),
        Patch(color="darkorange", label="Confounder"),
    ]
    ax.legend(handles=legend_elements, loc="upper left")
    plt.tight_layout()
    plt.show()
    return fig


# Education-income DAG
nodes = ["ability", "family_income", "treatment", "outcome", "noise_var"]
edges = [
    ("ability", "treatment"),
    ("ability", "outcome"),
    ("family_income", "treatment"),
    ("family_income", "outcome"),
    ("treatment", "outcome"),
]

dag = build_causal_dag(nodes, edges)
plot_dag(dag, treatment="treatment", outcome="outcome",
         confounders=["ability", "family_income"],
         title="Education → Income DAG")

# Check backdoor paths
def find_backdoor_paths(G, treatment, outcome):
    """Find all backdoor paths (paths with arrow into treatment)."""
    G_undirected = G.to_undirected()
    all_paths = list(nx.all_simple_paths(G_undirected, treatment, outcome))
    backdoor = []
    for path in all_paths:
        if len(path) > 2:
            # First edge must go INTO treatment
            if G.has_edge(path[1], path[0]):
                backdoor.append(path)
    return backdoor

backdoor_paths = find_backdoor_paths(dag, "treatment", "outcome")
print(f"\nBackdoor paths from 'treatment' to 'outcome':")
for p in backdoor_paths:
    print(f"  {' → '.join(p)}")

---

3. Identification and Estimation

import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import GradientBoostingRegressor, RandomForestClassifier
from sklearn.preprocessing import StandardScaler
import warnings
warnings.filterwarnings("ignore")


def identify_effect(dag, treatment, outcome, covariates):
    """
    Check if the backdoor criterion is satisfied and return adjustment set.

    The backdoor criterion holds if covariates:
    1. Block all backdoor paths from treatment to outcome.
    2. Do not include any descendants of treatment.

    Returns
    -------
    dict with identification status and adjustment set.
    """
    # Check for descendant confounders
    try:
        descendants = nx.descendants(dag, treatment)
        adjustment_set = [c for c in covariates if c not in descendants]
        has_backdoor = len([c for c in covariates if c in adjustment_set]) > 0
        return {
            "identified": True,
            "method": "backdoor",
            "adjustment_set": adjustment_set,
            "invalid_covariates": [c for c in covariates if c in descendants],
        }
    except Exception as e:
        return {"identified": False, "error": str(e)}


def estimate_ate(df, treatment, outcome, covariates, method="regression"):
    """
    Estimate the Average Treatment Effect using various methods.

    Parameters
    ----------
    df : pd.DataFrame
    treatment : str
        Binary treatment column name.
    outcome : str
        Continuous outcome column name.
    covariates : list of str
        Adjustment variables (confounders).
    method : str
        One of: 'regression', 'ipw', 'matching', 'doubly_robust', 'dml'

    Returns
    -------
    dict with 'ate', 'std_error', 'ci_lower', 'ci_upper'
    """
    T = df[treatment].values
    Y = df[outcome].values
    X = df[covariates].values

    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    if method == "regression":
        # OLS with treatment indicator and confounders
        X_full = np.column_stack([T, X_scaled])
        reg = LinearRegression().fit(X_full, Y)
        ate = reg.coef_[0]

        # Bootstrap standard error
        rng = np.random.default_rng(0)
        boot_ates = []
        for _ in range(500):
            idx = rng.integers(0, len(df), len(df))
            X_b = X_full[idx]
            Y_b = Y[idx]
            boot_ates.append(LinearRegression().fit(X_b, Y_b).coef_[0])
        std_error = np.std(boot_ates)

    elif method == "ipw":
        # Inverse Probability Weighting
        ps_model = LogisticRegression(C=1.0, max_iter=1000)
        ps_model.fit(X_scaled, T)
        ps = ps_model.predict_proba(X_scaled)[:, 1]
        ps = np.clip(ps, 0.01, 0.99)  # trim extreme weights

        weights = T / ps + (1 - T) / (1 - ps)
        ate = np.mean(weights * Y * T) / np.mean(weights * T) - \
              np.mean(weights * Y * (1 - T)) / np.mean(weights * (1 - T))

        # Bootstrap SE
        rng = np.random.default_rng(0)
        boot_ates = []
        for _ in range(500):
            idx = rng.integers(0, len(df), len(df))
            ps_b = ps[idx]
            T_b, Y_b = T[idx], Y[idx]
            w_b = T_b / ps_b + (1 - T_b) / (1 - ps_b)
            boot_ates.append(
                np.mean(w_b * Y_b * T_b) / np.mean(w_b * T_b) -
                np.mean(w_b * Y_b * (1 - T_b)) / np.mean(w_b * (1 - T_b))
            )
        std_error = np.std(boot_ates)

    elif method == "doubly_robust":
        # Doubly Robust (AIPW) estimator
        # Outcome model
        X_treated = np.column_stack([np.ones(X_scaled.shape[0]), X_scaled])
        treated_mask = T == 1
        control_mask = T == 0

        mu1 = GradientBoostingRegressor(n_estimators=100, random_state=0)
        mu1.fit(X_scaled[treated_mask], Y[treated_mask])
        mu0 = GradientBoostingRegressor(n_estimators=100, random_state=0)
        mu0.fit(X_scaled[control_mask], Y[control_mask])

        mu1_hat = mu1.predict(X_scaled)
        mu0_hat = mu0.predict(X_scaled)

        # Propensity model
        ps_model = LogisticRegression(C=1.0, max_iter=1000)
        ps_model.fit(X_scaled, T)
        ps = np.clip(ps_model.predict_proba(X_scaled)[:, 1], 0.01, 0.99)

        # AIPW scores
        tau_hat = (
            mu1_hat - mu0_hat
            + T * (Y - mu1_hat) / ps
            - (1 - T) * (Y - mu0_hat) / (1 - ps)
        )
        ate = np.mean(tau_hat)
        std_error = np.std(tau_hat) / np.sqrt(len(tau_hat))

    else:
        raise ValueError(f"Unknown method: {method}. Choose from: regression, ipw, doubly_robust")

    ci_lower = ate - 1.96 * std_error
    ci_upper = ate + 1.96 * std_error

    return {
        "method": method,
        "ate": ate,
        "std_error": std_error,
        "ci_lower": ci_lower,
        "ci_upper": ci_upper,
    }


# Run all methods on synthetic data
df = generate_confounded_data(n=3000)
covariates = ["ability", "family_income"]
true_ate = 2.5

results = {}
for method in ["regression", "ipw", "doubly_robust"]:
    res = estimate_ate(df, "treatment", "outcome", covariates, method=method)
    results[method] = res
    print(
        f"{method:15s}: ATE = {res['ate']:.4f} ± {res['std_error']:.4f}  "
        f"95% CI = [{res['ci_lower']:.3f}, {res['ci_upper']:.3f}]"
    )

print(f"\nTrue ATE = {true_ate}")

---

4. Propensity Score Matching

import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import NearestNeighbors
import matplotlib.pyplot as plt


def estimate_propensity_scores(df, treatment, covariates, model="logistic"):
    """
    Estimate propensity scores P(T=1 | X).

    Parameters
    ----------
    model : str
        'logistic' or 'gradient_boosting'
    """
    X = df[covariates].values
    T = df[treatment].values
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    if model == "logistic":
        clf = LogisticRegression(C=1.0, max_iter=1000, solver="lbfgs")
    elif model == "gradient_boosting":
        from sklearn.ensemble import GradientBoostingClassifier
        clf = GradientBoostingClassifier(n_estimators=100, max_depth=3, random_state=42)
    else:
        raise ValueError(f"Unknown model: {model}")

    clf.fit(X_scaled, T)
    ps = clf.predict_proba(X_scaled)[:, 1]
    return ps, scaler, clf


def propensity_score_matching(df, treatment, outcome, covariates, caliper=0.05, n_neighbors=1):
    """
    Perform 1:1 nearest-neighbor propensity score matching.

    Parameters
    ----------
    caliper : float
        Maximum allowed difference in propensity scores (in std units).
        Typically 0.2 * std(logit(ps)).
    n_neighbors : int
        Number of controls matched per treated unit.

    Returns
    -------
    matched_df : pd.DataFrame
        Dataset of matched pairs.
    att : float
        Average Treatment Effect on the Treated.
    """
    ps, _, _ = estimate_propensity_scores(df, treatment, covariates)
    df = df.copy()
    df["propensity_score"] = ps
    df["logit_ps"] = np.log(ps / (1 - ps))

    treated = df[df[treatment] == 1].reset_index(drop=True)
    control = df[df[treatment] == 0].reset_index(drop=True)

    # Fit nearest-neighbor on logit propensity scores
    nn = NearestNeighbors(n_neighbors=n_neighbors, metric="euclidean")
    nn.fit(control["logit_ps"].values.reshape(-1, 1))

    distances, indices = nn.kneighbors(treated["logit_ps"].values.reshape(-1, 1))

    # Apply caliper
    caliper_abs = caliper * np.std(df["logit_ps"])
    matched_pairs = []
    for i, (dist_row, idx_row) in enumerate(zip(distances, indices)):
        for dist, idx in zip(dist_row, idx_row):
            if dist <= caliper_abs:
                matched_pairs.append({
                    "treated_idx": i,
                    "control_idx": idx,
                    "distance": dist,
                    f"{outcome}_treated": treated[outcome].iloc[i],
                    f"{outcome}_control": control[outcome].iloc[idx],
                    "ps_treated": treated["propensity_score"].iloc[i],
                    "ps_control": control["propensity_score"].iloc[idx],
                })

    if not matched_pairs:
        raise ValueError("No matches found within caliper. Try increasing caliper value.")

    pairs_df = pd.DataFrame(matched_pairs)
    att = (pairs_df[f"{outcome}_treated"] - pairs_df[f"{outcome}_control"]).mean()
    att_se = (pairs_df[f"{outcome}_treated"] - pairs_df[f"{outcome}_control"]).std() / np.sqrt(len(pairs_df))

    print(f"\nPropensity Score Matching Results:")
    print(f"  Treated units:  {len(treated)}")
    print(f"  Control units:  {len(control)}")
    print(f"  Matched pairs:  {len(pairs_df)}")
    print(f"  Match rate:     {len(pairs_df)/len(treated)*100:.1f}%")
    print(f"  ATT = {att:.4f} ± {att_se:.4f}")
    print(f"  95% CI = [{att - 1.96*att_se:.3f}, {att + 1.96*att_se:.3f}]")

    return pairs_df, att, att_se


def check_balance(df, treatment, covariates, ps=None):
    """
    Compute standardized mean differences (SMD) before/after matching.
    SMD < 0.1 is typically considered good balance.
    """
    treated = df[df[treatment] == 1]
    control = df[df[treatment] == 0]

    print("\nCovariate Balance (Standardized Mean Differences):")
    print(f"{'Covariate':20s} {'SMD':>10}")
    print("-" * 32)

    smds = {}
    for col in covariates:
        mean_t = treated[col].mean()
        mean_c = control[col].mean()
        std_pooled = np.sqrt((treated[col].var() + control[col].var()) / 2)
        smd = (mean_t - mean_c) / (std_pooled + 1e-10)
        smds[col] = abs(smd)
        flag = " !" if abs(smd) > 0.1 else "  "
        print(f"{col:20s} {smd:>10.4f}{flag}")

    return smds


def plot_propensity_overlap(df, treatment, ps_col="propensity_score"):
    """Plot propensity score distributions to check overlap (common support)."""
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))

    treated_ps = df[df[treatment] == 1][ps_col]
    control_ps = df[df[treatment] == 0][ps_col]

    axes[0].hist(treated_ps, bins=40, alpha=0.6, color="tomato", density=True, label="Treated")
    axes[0].hist(control_ps, bins=40, alpha=0.6, color="steelblue", density=True, label="Control")
    axes[0].set_xlabel("Propensity Score")
    axes[0].set_ylabel("Density")
    axes[0].set_title("Propensity Score Overlap")
    axes[0].legend()

    # Love plot (SMD)
    smds = check_balance(df, treatment, [c for c in df.columns if c not in [treatment, ps_col, "outcome", "propensity_score", "logit_ps", "propensity_true"]])
    axes[1].barh(list(smds.keys()), list(smds.values()), color=["tomato" if v > 0.1 else "steelblue" for v in smds.values()])
    axes[1].axvline(0.1, color="red", ls="--", lw=1, label="SMD=0.1 threshold")
    axes[1].set_xlabel("Absolute SMD")
    axes[1].set_title("Covariate Balance (Love Plot)")
    axes[1].legend()

    plt.tight_layout()
    plt.show()


# Run matching
df = generate_confounded_data(n=3000)
df["propensity_score"] = estimate_propensity_scores(df, "treatment", ["ability", "family_income"])[0]
plot_propensity_overlap(df, "treatment")
pairs_df, att, att_se = propensity_score_matching(
    df, "treatment", "outcome", ["ability", "family_income"], caliper=0.1
)

---

5. Refutation Tests

import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler


def refute_estimate(df, treatment, outcome, covariates, estimated_ate, method="placebo"):
    """
    Run refutation tests to validate the causal estimate.

    Refutation methods:
    - 'random_common_cause': Add a random variable as confounder; ATE should be stable.
    - 'placebo_treatment': Replace treatment with random noise; ATE should go to ~0.
    - 'data_subset': Re-estimate on 80% random subsets; ATE should be stable.
    - 'bootstrap': Bootstrap distribution of ATE.

    Returns
    -------
    dict with refutation results and p-value.
    """

    def compute_ate_regression(data, t_col, y_col, x_cols):
        X = np.column_stack([data[t_col].values, StandardScaler().fit_transform(data[x_cols].values)])
        reg = LinearRegression().fit(X, data[y_col].values)
        return reg.coef_[0]

    rng = np.random.default_rng(123)

    if method == "random_common_cause":
        n_trials = 100
        perturbed_ates = []
        for _ in range(n_trials):
            df_perturbed = df.copy()
            df_perturbed["random_confounder"] = rng.standard_normal(len(df))
            new_covs = covariates + ["random_confounder"]
            ate_new = compute_ate_regression(df_perturbed, treatment, outcome, new_covs)
            perturbed_ates.append(ate_new)

        mean_perturbed = np.mean(perturbed_ates)
        std_perturbed = np.std(perturbed_ates)
        print(f"\nRandom Common Cause Refutation:")
        print(f"  Original ATE:   {estimated_ate:.4f}")
        print(f"  Perturbed ATE:  {mean_perturbed:.4f} ± {std_perturbed:.4f}")
        print(f"  Change:         {abs(mean_perturbed - estimated_ate):.4f}")
        result = "PASS" if abs(mean_perturbed - estimated_ate) < 0.3 else "FAIL"
        print(f"  Result: {result}")
        return {"method": method, "original_ate": estimated_ate,
                "new_ate": mean_perturbed, "passed": result == "PASS"}

    elif method == "placebo_treatment":
        n_trials = 200
        placebo_ates = []
        for _ in range(n_trials):
            df_placebo = df.copy()
            # Random binary placebo with same prevalence as original treatment
            p_treat = df[treatment].mean()
            df_placebo["placebo_treatment"] = (rng.uniform(size=len(df)) < p_treat).astype(int)
            ate_placebo = compute_ate_regression(df_placebo, "placebo_treatment", outcome, covariates)
            placebo_ates.append(ate_placebo)

        mean_placebo = np.mean(placebo_ates)
        # p-value: fraction of placebo ATEs more extreme than original
        p_value = np.mean(np.abs(placebo_ates) >= np.abs(estimated_ate))
        print(f"\nPlacebo Treatment Refutation:")
        print(f"  Original ATE:  {estimated_ate:.4f}")
        print(f"  Placebo ATE:   {mean_placebo:.4f} ± {np.std(placebo_ates):.4f}")
        print(f"  p-value:       {p_value:.4f}")
        result = "PASS" if abs(mean_placebo) < 0.2 else "FAIL"
        print(f"  Result: {result}")
        return {"method": method, "original_ate": estimated_ate,
                "placebo_ate": mean_placebo, "p_value": p_value, "passed": result == "PASS"}

    elif method == "data_subset":
        n_trials = 50
        subset_ates = []
        for _ in range(n_trials):
            idx = rng.choice(len(df), size=int(0.8 * len(df)), replace=False)
            df_sub = df.iloc[idx]
            ate_sub = compute_ate_regression(df_sub, treatment, outcome, covariates)
            subset_ates.append(ate_sub)

        mean_subset = np.mean(subset_ates)
        std_subset = np.std(subset_ates)
        print(f"\nData Subset Refutation:")
        print(f"  Original ATE:  {estimated_ate:.4f}")
        print(f"  Subset ATE:    {mean_subset:.4f} ± {std_subset:.4f}")
        result = "PASS" if abs(mean_subset - estimated_ate) < 0.15 else "FAIL"
        print(f"  Result: {result}")
        return {"method": method, "original_ate": estimated_ate,
                "new_ate": mean_subset, "std": std_subset, "passed": result == "PASS"}

    else:
        raise ValueError(f"Unknown refutation method: {method}")


# Run all refutation tests
df = generate_confounded_data(n=3000)
res = estimate_ate(df, "treatment", "outcome", ["ability", "family_income"], method="doubly_robust")
estimated_ate = res["ate"]

for refute_method in ["random_common_cause", "placebo_treatment", "data_subset"]:
    refute_estimate(df, "treatment", "outcome", ["ability", "family_income"],
                    estimated_ate, method=refute_method)

---

6. Complete Example A — Education Effect on Income

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import GradientBoostingRegressor


def education_income_study(n=5000, seed=0):
    """
    Estimate the causal effect of college education on income
    using backdoor adjustment on synthetic observational data.

    True causal effect = 3.0 (income units per year of education).
    """
    rng = np.random.default_rng(seed)

    # Observed confounders
    parents_edu = rng.integers(8, 20, size=n).astype(float)
    iq = 100 + 15 * rng.standard_normal(n)
    region = rng.integers(0, 3, size=n)  # 0=rural, 1=suburban, 2=urban

    # Treatment: years of college (0-4), influenced by confounders
    log_odds = (0.05 * (parents_edu - 12) + 0.02 * (iq - 100)
                + 0.3 * (region == 2).astype(float) - 0.5)
    p_college = 1 / (1 + np.exp(-log_odds))
    college_years = rng.binomial(4, p_college)

    # Outcome: income ($k/yr), TRUE causal effect = 3.0 per year
    noise = 5 * rng.standard_normal(n)
    income = (20 + 3.0 * college_years
              + 0.4 * parents_edu
              + 0.08 * (iq - 100)
              + 4 * (region == 2).astype(float)
              + noise)

    df = pd.DataFrame({
        "college_years": college_years,
        "income": income,
        "parents_edu": parents_edu,
        "iq": iq,
        "region": region,
    })

    covariates = ["parents_edu", "iq", "region"]
    true_ate = 3.0

    print("=" * 55)
    print("Education → Income Causal Effect Study")
    print("=" * 55)
    print(f"  N = {n}, True ATE = {true_ate}")
    print(f"  Naive OLS (no adjustment): {_naive_ols(df):.4f}")

    # Method comparison
    methods_results = {}
    for method in ["regression", "ipw", "doubly_robust"]:
        res = estimate_ate(df, "college_years", "income", covariates, method=method)
        methods_results[method] = res
        print(f"  {method:15s}: ATE = {res['ate']:.4f} (bias={res['ate']-true_ate:+.4f})")

    # Visualization
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))

    # 1. Scatter: college years vs income
    for region_val, color, label in [(0, "steelblue", "Rural"), (1, "darkorange", "Suburban"), (2, "green", "Urban")]:
        mask = df["region"] == region_val
        axes[0].scatter(df.loc[mask, "college_years"] + 0.05 * rng.standard_normal(mask.sum()),
                        df.loc[mask, "income"], alpha=0.15, s=5, color=color, label=label)
    axes[0].set_xlabel("College Years")
    axes[0].set_ylabel("Income ($k/yr)")
    axes[0].set_title("Raw Data by Region")
    axes[0].legend(markerscale=3)

    # 2. ATE comparison
    method_names = list(methods_results.keys())
    ates = [methods_results[m]["ate"] for m in method_names]
    cis = [(methods_results[m]["ci_lower"], methods_results[m]["ci_upper"]) for m in method_names]
    y_pos = range(len(method_names))
    axes[1].barh(y_pos, ates, xerr=[(a - ci[0]) for a, ci in zip(ates, cis)], color="steelblue", alpha=0.7)
    axes[1].axvline(true_ate, color="red", ls="--", lw=2, label=f"True ATE={true_ate}")
    axes[1].set_yticks(list(y_pos))
    axes[1].set_yticklabels(method_names)
    axes[1].set_xlabel("Estimated ATE")
    axes[1].set_title("ATE Estimates by Method")
    axes[1].legend()

    # 3. Propensity score distribution
    X_cov = StandardScaler().fit_transform(df[covariates].values)
    # Use college >=2 as binary treatment proxy
    T_binary = (df["college_years"] >= 2).astype(int)
    ps_model = LogisticRegression(C=1.0, max_iter=1000)
    ps_model.fit(X_cov, T_binary)
    ps = ps_model.predict_proba(X_cov)[:, 1]
    axes[2].hist(ps[T_binary == 1], bins=40, alpha=0.6, color="tomato", density=True, label="College ≥2yr")
    axes[2].hist(ps[T_binary == 0], bins=40, alpha=0.6, color="steelblue", density=True, label="College <2yr")
    axes[2].set_xlabel("Propensity Score")
    axes[2].set_ylabel("Density")
    axes[2].set_title("Propensity Score Overlap")
    axes[2].legend()

    plt.suptitle("Education → Income Causal Study", fontsize=13, fontweight="bold")
    plt.tight_layout()
    plt.show()

    return df, methods_results


def _naive_ols(df):
    """Simple bivariate regression (confounded estimate)."""
    from sklearn.linear_model import LinearRegression
    X = df["college_years"].values.reshape(-1, 1)
    Y = df["income"].values
    return LinearRegression().fit(X, Y).coef_[0]


if __name__ == "__main__":
    df_study, results = education_income_study()

---

7. Complete Example B — Propensity Score Matching for Observational Study

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression


def observational_matching_study(n=2000, seed=42):
    """
    Simulate a clinical observational study and apply propensity score matching.

    Scenario: Effect of a new drug on blood pressure reduction.
    True ATT = 8.0 mmHg reduction.
    """
    rng = np.random.default_rng(seed)

    # Baseline covariates
    age = rng.integers(40, 75, size=n).astype(float)
    bmi = 25 + 5 * rng.standard_normal(n)
    baseline_bp = 140 + 10 * rng.standard_normal(n)
    smoker = (rng.uniform(size=n) < 0.3).astype(int)
    diabetes = (rng.uniform(size=n) < 0.2).astype(int)

    # Treatment assignment (sicker patients more likely to receive drug)
    log_odds = (0.04 * (age - 55) + 0.05 * (bmi - 25) +
                0.01 * (baseline_bp - 140) + 0.5 * smoker + 0.4 * diabetes - 0.5)
    ps_true = 1 / (1 + np.exp(-log_odds))
    drug = (rng.uniform(size=n) < ps_true).astype(int)

    # Outcome: BP reduction (true ATT = 8.0)
    noise = 3.0 * rng.standard_normal(n)
    bp_reduction = (8.0 * drug
                    + 0.15 * (age - 55)
                    + 0.2 * (bmi - 25)
                    + 0.05 * (baseline_bp - 140)
                    + 2.0 * smoker
                    + 1.5 * diabetes
                    + noise)

    df = pd.DataFrame({
        "drug": drug,
        "bp_reduction": bp_reduction,
        "age": age,
        "bmi": bmi,
        "baseline_bp": baseline_bp,
        "smoker": smoker,
        "diabetes": diabetes,
        "ps_true": ps_true,
    })

    covariates = ["age", "bmi", "baseline_bp", "smoker", "diabetes"]
    true_att = 8.0

    print("=" * 55)
    print("Clinical Drug Study — Propensity Score Matching")
    print("=" * 55)
    naive_diff = df.groupby("drug")["bp_reduction"].mean().diff().iloc[-1]
    print(f"  Naive difference in means:  {naive_diff:.4f} (true ATT = {true_att})")

    # Estimate PS and run matching
    X = StandardScaler().fit_transform(df[covariates].values)
    T = df["drug"].values
    ps_model = LogisticRegression(C=1.0, max_iter=1000)
    ps_model.fit(X, T)
    df["ps_estimated"] = ps_model.predict_proba(X)[:, 1]

    # Check overlap before matching
    from sklearn.calibration import calibration_curve
    fraction_of_positives, mean_predicted_value = calibration_curve(
        T, df["ps_estimated"], n_bins=10
    )

    pairs_df, att, att_se = propensity_score_matching(
        df, "drug", "bp_reduction", covariates, caliper=0.1
    )

    print(f"\n  Matched ATT:               {att:.4f} ± {att_se:.4f}")
    print(f"  Bias (matched):            {att - true_att:+.4f}")

    # Also run doubly-robust for comparison
    dr_result = estimate_ate(df, "drug", "bp_reduction", covariates, method="doubly_robust")
    print(f"  Doubly-robust ATT:         {dr_result['ate']:.4f} ± {dr_result['std_error']:.4f}")

    # Balance check after matching
    # Build matched dataset
    treated_matched_idx = pairs_df["treated_idx"].values
    control_matched_idx = pairs_df["control_idx"].values
    treated_df = df[df["drug"] == 1].reset_index(drop=True).iloc[treated_matched_idx]
    control_df = df[df["drug"] == 0].reset_index(drop=True).iloc[control_matched_idx]
    df_matched = pd.concat([treated_df, control_df], ignore_index=True)

    print("\nBalance before matching:")
    check_balance(df, "drug", covariates)
    print("\nBalance after matching:")
    check_balance(df_matched, "drug", covariates)

    # Final plot
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))

    axes[0].hist(df[df["drug"] == 1]["ps_estimated"], bins=30, alpha=0.6,
                 color="tomato", density=True, label="Drug=1")
    axes[0].hist(df[df["drug"] == 0]["ps_estimated"], bins=30, alpha=0.6,
                 color="steelblue", density=True, label="Drug=0")
    axes[0].set_title("Propensity Score Distribution (Before Matching)")
    axes[0].set_xlabel("Estimated PS")
    axes[0].legend()

    diff_pairs = pairs_df["bp_reduction_treated"] - pairs_df["bp_reduction_control"]
    axes[1].hist(diff_pairs, bins=40, color="steelblue", alpha=0.7, edgecolor="white")
    axes[1].axvline(att, color="red", lw=2, label=f"ATT={att:.2f}")
    axes[1].axvline(true_att, color="green", lw=2, ls="--", label=f"True ATT={true_att}")
    axes[1].set_xlabel("Individual Treatment Effect (matched pairs)")
    axes[1].set_ylabel("Count")
    axes[1].set_title("Matched Pair Differences")
    axes[1].legend()

    plt.tight_layout()
    plt.show()

    return df, pairs_df, att


if __name__ == "__main__":
    observational_matching_study()

---

Quick Reference

| Task | Method / Function | |---------------------------------|-------------------------------------------------------| | Build causal DAG | build_causal_dag(nodes, edges) | | Visualize DAG | plot_dag(G, treatment, outcome, confounders) | | Find backdoor paths | find_backdoor_paths(G, treatment, outcome) | | Regression adjustment ATE | estimate_ate(df, T, Y, X, method='regression') | | IPW estimation | estimate_ate(df, T, Y, X, method='ipw') | | Doubly-robust (AIPW) | estimate_ate(df, T, Y, X, method='doubly_robust') | | Propensity score estimation | estimate_propensity_scores(df, T, X) | | PSM with caliper | propensity_score_matching(df, T, Y, X, caliper) | | Check covariate balance | check_balance(df, T, covariates) | | Placebo refutation | refute_estimate(..., method='placebo_treatment') | | Random confounder refutation | refute_estimate(..., method='random_common_cause') | | Data subset refutation | refute_estimate(..., method='data_subset') |

Key Assumptions to Check

1. Unconfoundedness (ignorability): All confounders are measured. Test with sensitivity analysis.
2. Overlap (positivity): 0 < P(T=1|X) < 1 for all X. Check propensity score histograms.
3. SUTVA: No interference between units; no hidden treatment versions.
4. Correct model specification: Outcome/propensity models must be well-specified for doubly-robust to work.