wehub-resource-sync
537 lines
17 KiB
Python
537 lines
17 KiB
Python
"""Shared causal inference utilities for Ch15 notebooks and case study DML.
|
||
|
||
Provides:
|
||
- block_permute(): Block permutation preserving autocorrelation
|
||
- manual_dml_timeseries(): Walk-forward DML with embargo
|
||
- run_dml_analysis(): Full DML pipeline (naive + DML + refutation)
|
||
|
||
Used by teaching notebooks (02-04, 07) and case study 09_causal_dml.py.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
|
||
# Pin OpenMP threading before sklearn imports — HistGradientBoostingRegressor
|
||
# uses OMP-parallel histogram construction whose floating-point reduction order
|
||
# is non-deterministic across threads. With OMP_NUM_THREADS=1 the placebo loop
|
||
# is bit-reproducible across runs at the same seed/spec/data.
|
||
os.environ.setdefault("OMP_NUM_THREADS", "1")
|
||
|
||
import time
|
||
from datetime import UTC, datetime
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import statsmodels.api as sm
|
||
from scipy import stats
|
||
from sklearn.ensemble import HistGradientBoostingRegressor
|
||
from statsmodels.regression.linear_model import OLS
|
||
|
||
from utils.modeling import RANDOM_SEED, seed_everything
|
||
|
||
|
||
def embargo_from_buffer(label_buffer: str) -> int:
|
||
"""Convert a label buffer string to an integer embargo period count.
|
||
|
||
Supports all pandas duration units (D, H/h, M, T/min).
|
||
Returns the number of observation periods to skip between train and test.
|
||
|
||
Bar-frequency assumptions baked into the conversion:
|
||
- D: one period per `value` days (e.g. "5D" → 5)
|
||
- H/h: the buffer is interpreted as the bar cadence; the result is the
|
||
number of `value`-hour bars in one trading day (24 // value), so "8H"
|
||
→ 3 bars (a one-day embargo on 8-hour bars)
|
||
- M: `value` months × 21 trading days
|
||
- T/min: the buffer is the cadence; the result is the bars in 15 minutes
|
||
(`value` // 15), assuming 15-minute base bars
|
||
"""
|
||
import re
|
||
|
||
match = re.match(r"(\d+)(D|H|h|M|T|min)", label_buffer.strip())
|
||
if not match:
|
||
raise ValueError(f"Cannot parse label_buffer: {label_buffer}")
|
||
value, unit = int(match.group(1)), match.group(2)
|
||
return {
|
||
"D": value,
|
||
"H": max(1, 24 // value),
|
||
"h": max(1, 24 // value),
|
||
"M": value * 21,
|
||
"T": max(1, value // 15),
|
||
"min": max(1, value // 15),
|
||
}[unit]
|
||
|
||
|
||
def block_permute(
|
||
arr: np.ndarray, block_size: int, rng: np.random.Generator | None = None
|
||
) -> np.ndarray:
|
||
"""Permute array in blocks to preserve autocorrelation structure.
|
||
|
||
Essential for refutation tests on time series data. Random permutation
|
||
destroys autocorrelation, making placebo tests too easy to pass.
|
||
|
||
Parameters
|
||
----------
|
||
arr : array-like
|
||
Array to permute.
|
||
block_size : int
|
||
Size of blocks to preserve.
|
||
rng : np.random.Generator, optional
|
||
Random number generator for reproducibility.
|
||
|
||
Returns
|
||
-------
|
||
np.ndarray
|
||
Block-permuted array.
|
||
"""
|
||
arr = np.asarray(arr)
|
||
n = len(arr)
|
||
if rng is None:
|
||
rng = np.random.default_rng()
|
||
|
||
n_blocks = n // block_size
|
||
if n_blocks < 2:
|
||
return rng.permutation(arr)
|
||
|
||
block_indices = rng.permutation(n_blocks)
|
||
|
||
result = []
|
||
for idx in block_indices:
|
||
start = idx * block_size
|
||
result.append(arr[start : start + block_size])
|
||
|
||
# Handle remainder
|
||
remainder_start = n_blocks * block_size
|
||
if remainder_start < n:
|
||
result.append(arr[remainder_start:])
|
||
|
||
return np.concatenate(result)
|
||
|
||
|
||
def manual_dml_timeseries(
|
||
Y: np.ndarray,
|
||
T: np.ndarray,
|
||
X: np.ndarray,
|
||
n_folds: int = 5,
|
||
embargo: int = 21,
|
||
model_y=None,
|
||
model_t=None,
|
||
return_residuals: bool = False,
|
||
hac_maxlags: int | None = None,
|
||
) -> dict:
|
||
"""Walk-forward DML with embargo for temporal data.
|
||
|
||
Follows Chernozhukov et al. (2017) and de Prado (2018):
|
||
1. Split data into K temporal folds (not random)
|
||
2. For each fold, train on earlier data, predict on later
|
||
3. Embargo gap between train and test prevents autocorrelation leakage
|
||
4. HAC standard errors account for residual autocorrelation
|
||
|
||
Parameters
|
||
----------
|
||
Y : array
|
||
Outcome variable.
|
||
T : array
|
||
Treatment variable.
|
||
X : array
|
||
Confounder matrix.
|
||
n_folds : int
|
||
Number of temporal folds.
|
||
embargo : int
|
||
Gap periods between train and test sets.
|
||
model_y, model_t : sklearn estimator, optional
|
||
Nuisance models for E[Y|X] and E[T|X].
|
||
return_residuals : bool
|
||
If True, include residual arrays in result dict.
|
||
hac_maxlags : int or None
|
||
HAC (Newey-West) bandwidth. If None, uses cube-root rule: max(1, int(n**(1/3))).
|
||
|
||
Returns
|
||
-------
|
||
dict
|
||
Keys: theta, se_iid, se_hac, t_stat_iid, t_stat_hac, p_value_hac,
|
||
n_obs, hac_maxlags. If return_residuals: also Y_res, T_res.
|
||
"""
|
||
seed_everything(RANDOM_SEED)
|
||
|
||
n = len(Y)
|
||
|
||
# Initialize residual arrays
|
||
Y_res = np.full(n, np.nan)
|
||
T_res = np.full(n, np.nan)
|
||
|
||
fold_size = n // (n_folds + 1)
|
||
|
||
for fold in range(n_folds):
|
||
train_end = (fold + 1) * fold_size
|
||
test_start = train_end + embargo
|
||
test_end = min(test_start + fold_size, n)
|
||
|
||
if test_end <= test_start:
|
||
continue
|
||
|
||
train_idx = np.arange(0, train_end)
|
||
test_idx = np.arange(test_start, test_end)
|
||
|
||
if len(train_idx) < 50 or len(test_idx) < 10:
|
||
continue
|
||
|
||
# Fit nuisance models on training data (clone to avoid mutation)
|
||
from sklearn.base import clone
|
||
|
||
_default_y = HistGradientBoostingRegressor(max_iter=50, max_depth=3, random_state=42)
|
||
_default_t = HistGradientBoostingRegressor(max_iter=50, max_depth=3, random_state=42)
|
||
my = clone(model_y) if model_y is not None else _default_y
|
||
mt = clone(model_t) if model_t is not None else _default_t
|
||
|
||
my.fit(X[train_idx], Y[train_idx])
|
||
mt.fit(X[train_idx], T[train_idx])
|
||
|
||
Y_res[test_idx] = Y[test_idx] - my.predict(X[test_idx])
|
||
T_res[test_idx] = T[test_idx] - mt.predict(X[test_idx])
|
||
|
||
# Drop observations without residuals
|
||
valid = ~np.isnan(Y_res) & ~np.isnan(T_res)
|
||
Y_v = Y_res[valid]
|
||
T_v = T_res[valid]
|
||
n_valid = len(Y_v)
|
||
|
||
empty = {
|
||
"theta": np.nan,
|
||
"se_iid": np.nan,
|
||
"se_hac": np.nan,
|
||
"t_stat_iid": np.nan,
|
||
"t_stat_hac": np.nan,
|
||
"p_value_hac": np.nan,
|
||
"n_obs": n_valid,
|
||
"hac_maxlags": 0,
|
||
}
|
||
if n_valid < 50:
|
||
if return_residuals:
|
||
empty["Y_res"] = Y_res
|
||
empty["T_res"] = T_res
|
||
return empty
|
||
|
||
# Final stage: Y_res = alpha + theta * T_res + epsilon
|
||
# Must include intercept: cross-fitting residuals may have non-zero mean
|
||
# when training data varies across folds (expanding window).
|
||
if hac_maxlags is None:
|
||
hac_maxlags = max(1, int(n_valid ** (1 / 3)))
|
||
|
||
T_const = sm.add_constant(T_v)
|
||
ols_iid = OLS(Y_v, T_const).fit()
|
||
theta = ols_iid.params[1]
|
||
|
||
# HC0 standard error
|
||
se_iid = np.sqrt(ols_iid.cov_HC0[1, 1])
|
||
|
||
# HAC (Newey-West) standard error with frequency-adaptive bandwidth
|
||
se_hac = se_iid
|
||
try:
|
||
ols_hac = OLS(Y_v, T_const).fit(cov_type="HAC", cov_kwds={"maxlags": hac_maxlags})
|
||
cov = ols_hac.cov_params()
|
||
se_hac = np.sqrt(cov.iloc[1, 1] if hasattr(cov, "iloc") else cov[1, 1])
|
||
except Exception:
|
||
pass # Fall back to HC0 standard errors on numerical failure
|
||
|
||
t_stat_hac = theta / se_hac if se_hac > 0 else np.nan
|
||
p_value_hac = (
|
||
2 * (1 - stats.t.cdf(abs(t_stat_hac), df=n_valid - 2))
|
||
if not np.isnan(t_stat_hac)
|
||
else np.nan
|
||
)
|
||
|
||
result = {
|
||
"theta": theta,
|
||
"se_iid": se_iid,
|
||
"se_hac": se_hac,
|
||
"t_stat_iid": theta / se_iid if se_iid > 0 else np.nan,
|
||
"t_stat_hac": t_stat_hac,
|
||
"p_value_hac": p_value_hac,
|
||
"n_obs": n_valid,
|
||
"hac_maxlags": hac_maxlags,
|
||
}
|
||
|
||
if return_residuals:
|
||
result["Y_res"] = Y_res
|
||
result["T_res"] = T_res
|
||
|
||
return result
|
||
|
||
|
||
REFUTATION_ALPHA = 0.05
|
||
|
||
|
||
def classify_refutation(empirical_p: float) -> str:
|
||
"""Binary pass/fail of the block-permutation refutation test at 5 %.
|
||
|
||
Returns "Passes" if the empirical placebo p-value is below 5 %
|
||
(the observed effect cannot be reproduced by permutation in
|
||
most placebo runs); "Fails" otherwise. Always read the raw
|
||
`empirical_p` alongside the label.
|
||
"""
|
||
return "Passes" if empirical_p < REFUTATION_ALPHA else "Fails"
|
||
|
||
|
||
def run_dml_analysis(
|
||
df: pd.DataFrame,
|
||
treatment_col: str,
|
||
outcome_col: str,
|
||
confounder_cols: list[str],
|
||
n_folds: int = 5,
|
||
embargo: int = 21,
|
||
n_placebo: int = 100,
|
||
block_size: int = 21,
|
||
seed: int = 42,
|
||
hac_maxlags: int | None = None,
|
||
) -> dict:
|
||
"""Full DML analysis pipeline: naive OLS, DML, and refutation tests.
|
||
|
||
Parameters
|
||
----------
|
||
df : pd.DataFrame
|
||
Analysis dataset sorted by time.
|
||
treatment_col : str
|
||
Treatment variable column name.
|
||
outcome_col : str
|
||
Outcome variable column name.
|
||
confounder_cols : list[str]
|
||
Confounder column names.
|
||
n_folds : int
|
||
Number of walk-forward CV folds.
|
||
embargo : int
|
||
Gap periods between train and test.
|
||
n_placebo : int
|
||
Number of block permutation replications.
|
||
block_size : int
|
||
Block size for permutation test.
|
||
seed : int
|
||
Random seed.
|
||
hac_maxlags : int or None
|
||
HAC bandwidth. If None, uses cube-root rule.
|
||
|
||
Returns
|
||
-------
|
||
dict
|
||
Comprehensive results with keys: naive_effect, dml_result,
|
||
confounding_bias, confounding_bias_pct, refutation (z_score,
|
||
empirical_p, placebo_mean, placebo_std, placebo_effects,
|
||
refutation_class), p_value_hac, hac_maxlags, and n_obs.
|
||
"""
|
||
# Input validation
|
||
n = len(df)
|
||
min_rows = (n_folds + 1) * 50 + n_folds * embargo
|
||
if n < min_rows:
|
||
raise ValueError(
|
||
f"Need at least {min_rows} rows for {n_folds}-fold CV with embargo={embargo}, got {n}"
|
||
)
|
||
if df[treatment_col].std() < 1e-10:
|
||
raise ValueError(f"Treatment '{treatment_col}' has near-zero variance")
|
||
if df[outcome_col].std() < 1e-10:
|
||
raise ValueError(f"Outcome '{outcome_col}' has near-zero variance")
|
||
|
||
_dml_started_at = datetime.now(UTC).isoformat()
|
||
_dml_t0 = time.perf_counter()
|
||
|
||
rng = np.random.default_rng(seed)
|
||
|
||
T = df[treatment_col].values
|
||
Y = df[outcome_col].values
|
||
X = df[confounder_cols].values
|
||
|
||
# Naive OLS
|
||
T_const = np.column_stack([np.ones(len(T)), T])
|
||
naive_coef = np.linalg.lstsq(T_const, Y, rcond=None)[0]
|
||
naive_effect = naive_coef[1]
|
||
|
||
# DML estimate
|
||
dml = manual_dml_timeseries(
|
||
Y,
|
||
T,
|
||
X,
|
||
n_folds=n_folds,
|
||
embargo=embargo,
|
||
return_residuals=True,
|
||
hac_maxlags=hac_maxlags,
|
||
)
|
||
|
||
# Confounding bias
|
||
dml_effect = dml["theta"]
|
||
bias = naive_effect - dml_effect
|
||
bias_pct = bias / abs(dml_effect) * 100 if dml_effect != 0 else 0.0
|
||
|
||
# Block permutation refutation
|
||
placebo_effects = []
|
||
for _ in range(n_placebo):
|
||
T_perm = block_permute(T, block_size, rng=rng)
|
||
perm_result = manual_dml_timeseries(
|
||
Y,
|
||
T_perm,
|
||
X,
|
||
n_folds=min(3, n_folds),
|
||
embargo=embargo,
|
||
)
|
||
if not np.isnan(perm_result["theta"]):
|
||
placebo_effects.append(perm_result["theta"])
|
||
|
||
refutation = {}
|
||
if len(placebo_effects) >= 10:
|
||
placebo_arr = np.array(placebo_effects)
|
||
p_mean = np.mean(placebo_arr)
|
||
p_std = np.std(placebo_arr)
|
||
z = (dml_effect - p_mean) / p_std if p_std > 0 else np.inf
|
||
emp_p = np.mean(np.abs(placebo_arr) >= np.abs(dml_effect))
|
||
ref_class = classify_refutation(emp_p)
|
||
refutation = {
|
||
"z_score": z,
|
||
"empirical_p": emp_p,
|
||
"placebo_mean": p_mean,
|
||
"placebo_std": p_std,
|
||
"n_successful": len(placebo_effects),
|
||
"placebo_effects": placebo_effects,
|
||
"refutation_class": ref_class,
|
||
}
|
||
|
||
return {
|
||
"naive_effect": naive_effect,
|
||
"dml_result": dml,
|
||
"confounding_bias": bias,
|
||
"confounding_bias_pct": bias_pct,
|
||
"refutation": refutation,
|
||
"p_value_hac": dml.get("p_value_hac", np.nan),
|
||
"hac_maxlags": dml.get("hac_maxlags", 0),
|
||
"n_obs": len(df),
|
||
"started_at": _dml_started_at,
|
||
"elapsed_s": time.perf_counter() - _dml_t0,
|
||
}
|
||
|
||
|
||
def format_dml_summary(results: dict) -> str:
|
||
"""Format DML analysis results for display."""
|
||
dml = results["dml_result"]
|
||
p_hac = results.get("p_value_hac", dml.get("p_value_hac", np.nan))
|
||
hac_lags = results.get("hac_maxlags", dml.get("hac_maxlags", "?"))
|
||
lines = [
|
||
"=" * 60,
|
||
"DML ANALYSIS SUMMARY",
|
||
"=" * 60,
|
||
f"Observations: {results['n_obs']:,}",
|
||
f"HAC bandwidth: {hac_lags} lags (cube-root rule)",
|
||
"",
|
||
f"Naive OLS effect: {results['naive_effect']:.6f}",
|
||
f"DML effect: {dml['theta']:.6f}",
|
||
f" SE (IID): {dml['se_iid']:.6f}",
|
||
f" SE (HAC): {dml['se_hac']:.6f}",
|
||
f" t-stat (HAC): {dml['t_stat_hac']:.2f}",
|
||
f" p-value (HAC): {p_hac:.4f}",
|
||
"",
|
||
f"Confounding bias: {results['confounding_bias']:.6f} ({results['confounding_bias_pct']:+.1f}%)",
|
||
]
|
||
|
||
ref = results.get("refutation", {})
|
||
if ref:
|
||
ref_class = ref.get("refutation_class", classify_refutation(ref["empirical_p"]))
|
||
lines += [
|
||
"",
|
||
"Refutation (block permutation):",
|
||
f" Z-score: {ref['z_score']:.2f}",
|
||
f" Empirical p: {ref['empirical_p']:.4f}",
|
||
f" Classification: {ref_class}",
|
||
f" Placebos: {ref['n_successful']}",
|
||
]
|
||
|
||
lines.append("=" * 60)
|
||
return "\n".join(lines)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Registry integration
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def register_causal_run(
|
||
case_study_id: str,
|
||
label: str,
|
||
results: dict,
|
||
predictions=None,
|
||
*,
|
||
treatment_col: str = "",
|
||
confounder_cols: list[str] | None = None,
|
||
n_folds: int = 5,
|
||
embargo: int = 0,
|
||
notebook: str = "causal_dml",
|
||
case_dir=None,
|
||
started_at: str | None = None,
|
||
elapsed_s: float | None = None,
|
||
) -> str:
|
||
"""Register a causal DML run in the dedicated `causal_runs` table.
|
||
|
||
Causal DML estimates a treatment effect rather than a cross-sectional
|
||
score, so it lives in its own table — distinct from `training_runs`,
|
||
`prediction_sets`, and `prediction_metrics` which serve predictive
|
||
families. The `predictions` argument (per-row residuals + ATE) is
|
||
accepted for backward compatibility but no longer persisted: it has
|
||
no downstream readers, and re-running the case study notebook is the
|
||
canonical way to regenerate diagnostics.
|
||
"""
|
||
import json
|
||
|
||
# Alias the registration helper to avoid shadowing this wrapper's own name —
|
||
# a future refactor that hoists this import to module level would otherwise
|
||
# turn the call below into infinite recursion.
|
||
from case_studies.utils.registry.registration import (
|
||
register_causal_run as _register_causal_run,
|
||
)
|
||
from case_studies.utils.registry.specs import (
|
||
build_training_spec,
|
||
canonical_json,
|
||
training_hash_from_spec,
|
||
)
|
||
|
||
dml_result = results.get("dml_result", {})
|
||
ref = results.get("refutation", {})
|
||
|
||
causal_params = {"treatment": treatment_col, "embargo": embargo}
|
||
if confounder_cols:
|
||
causal_params["confounders"] = confounder_cols
|
||
|
||
spec = build_training_spec(
|
||
"causal_dml",
|
||
"dml",
|
||
label,
|
||
n_folds=n_folds,
|
||
causal_params=causal_params,
|
||
)
|
||
causal_hash = training_hash_from_spec(spec)
|
||
|
||
# Preserve NULLs for unknown p-values rather than silently coercing them
|
||
# to 1.0 — a HAC p-value that underflows to exactly 0.0 is a strongly
|
||
# significant result, and ``or 1.0`` would flip its meaning.
|
||
p_value_hac = results.get("p_value_hac")
|
||
refutation_p = ref.get("empirical_p")
|
||
|
||
_register_causal_run(
|
||
case_study_id,
|
||
causal_hash,
|
||
label=label,
|
||
treatment=treatment_col,
|
||
confounders_json=json.dumps(confounder_cols or []),
|
||
embargo=embargo,
|
||
n_folds=n_folds,
|
||
n_obs=int(dml_result.get("n_obs", 0)),
|
||
dml_effect=float(dml_result.get("theta", 0.0)),
|
||
dml_se_hac=float(dml_result.get("se_hac", 0.0)),
|
||
p_value_hac=float(p_value_hac) if p_value_hac is not None else None,
|
||
naive_effect=float(results.get("naive_effect", 0.0)),
|
||
confounding_bias_pct=float(results.get("confounding_bias_pct", 0.0)),
|
||
refutation_p=float(refutation_p) if refutation_p is not None else None,
|
||
spec_json=canonical_json(spec),
|
||
notebook=notebook,
|
||
started_at=started_at or results.get("started_at"),
|
||
elapsed_s=elapsed_s if elapsed_s is not None else results.get("elapsed_s"),
|
||
)
|
||
|
||
p_hac_display = f"{float(p_value_hac):.4f}" if p_value_hac is not None else "n/a"
|
||
print(f" -> registered causal_dml (causal_hash={causal_hash}, p_hac={p_hac_display})")
|
||
return causal_hash
|