spb/wp3_uqo

1#!/usr/bin/env python32# Author: Simon-Pierre Boucher — contact@spboucher.ai3#4"""Inter-quantile Wald tests with a joint bootstrap covariance (October 20265revision).67The published test (script 02) treats the tau = 0.10 and tau = 0.90 estimates8as independent, i.e. it omits the covariance term of Eq. (iqr_test). Here the9two quantile regressions are re-estimated on the same 150,000-observation10subsample (seed 42) and the difference beta(0.10) - beta(0.90) is bootstrapped11jointly: (i) a paired observation bootstrap and (ii) a paired ZIP3-cluster12bootstrap that also allows for spatial dependence within three-digit ZIP areas.13Each bootstrap draw re-fits both quantile regressions (about 15 s per draw);14draws run in parallel. Default: 200 draws, roughly 8-10 minutes on 16 workers.1516Produces results/revision/iqr_wald_bootstrap.csv.1718Usage:  python scripts/07_revision_qr_bootstrap.py [--draws 200] [--workers 8]19"""2021import argparse22import sys23import time24from concurrent.futures import ProcessPoolExecutor25from pathlib import Path2627import numpy as np28import pandas as pd29import statsmodels.api as sm30from scipy import stats3132sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))33from wp3 import config, data  # noqa: E4023435OUT = config.RESULTS_DIR / "revision"36_X = _Y = _CODES = None  # populated in each worker via initializer373839def _init(X, y, codes):40    global _X, _Y, _CODES41    _X, _Y, _CODES = X, y, codes424344def _fit_pair(idx):45    r10 = sm.QuantReg(_Y[idx], _X[idx]).fit(q=0.10, max_iter=2000)46    r90 = sm.QuantReg(_Y[idx], _X[idx]).fit(q=0.90, max_iter=2000)47    return r10.params, r90.params484950def _draw(args):51    seed, scheme = args52    rng = np.random.RandomState(seed)53    n = len(_Y)54    if scheme == "obs":55        idx = rng.randint(0, n, n)56    else:57        uniq = np.unique(_CODES)58        members = {g: np.where(_CODES == g)[0] for g in uniq}59        draw = rng.choice(uniq, size=len(uniq), replace=True)60        idx = np.concatenate([members[g] for g in draw])61    b10, b90 = _fit_pair(idx)62    return b10 - b90636465def main():66    parser = argparse.ArgumentParser(description=__doc__)67    parser.add_argument("--draws", type=int, default=200)68    parser.add_argument("--workers", type=int, default=8)69    args = parser.parse_args()70    OUT.mkdir(parents=True, exist_ok=True)71    t0 = time.time()7273    md = data.load_model_data()74    df = md["df_clean"]75    X_const, y = md["X_const"], md["y_clean"]76    rng = np.random.RandomState(config.RANDOM_STATE)77    idx = rng.choice(len(X_const), size=config.QR_SUBSAMPLE_SIZE, replace=False)78    X = X_const.iloc[idx].values.astype(float)79    yy = y.iloc[idx].values.astype(float)80    codes = pd.factorize(df["address_zipcode"].fillna("000").astype(str).str[:3].iloc[idx])[0]81    names = list(X_const.columns)82    kpos = [names.index(v) for v in config.KEY_VARS]8384    print("Point estimates on the seed-42 subsample...")85    _init(X, yy, codes)86    r10 = sm.QuantReg(yy, X).fit(q=0.10, max_iter=2000)87    r90 = sm.QuantReg(yy, X).fit(q=0.90, max_iter=2000)88    diff = r10.params - r90.params89    se_indep = np.sqrt(r10.bse ** 2 + r90.bse ** 2)9091    results = {}92    for scheme in ["obs", "zip3"]:93        print(f"Bootstrap ({scheme}), {args.draws} draws on {args.workers} workers...")94        t = time.time()95        with ProcessPoolExecutor(max_workers=args.workers, initializer=_init,96                                 initargs=(X, yy, codes)) as ex:97            draws = list(ex.map(_draw, [(1000 + i, scheme) for i in range(args.draws)]))98        results[scheme] = np.vstack(draws)99        print(f"   done in {time.time() - t:.0f}s")100101    rows = []102    for v, j in zip(config.KEY_VARS, kpos):103        d_obs = results["obs"][:, j]104        d_z3 = results["zip3"][:, j]105        row = {"variable": v, "beta_010": r10.params[j], "beta_090": r90.params[j],106               "difference": diff[j],107               "se_independent": se_indep[j], "z_independent": diff[j] / se_indep[j],108               "se_boot_obs": d_obs.std(ddof=1), "z_boot_obs": diff[j] / d_obs.std(ddof=1),109               "ci_low_obs": np.percentile(d_obs, 2.5), "ci_high_obs": np.percentile(d_obs, 97.5),110               "se_boot_zip3": d_z3.std(ddof=1), "z_boot_zip3": diff[j] / d_z3.std(ddof=1),111               "ci_low_zip3": np.percentile(d_z3, 2.5), "ci_high_zip3": np.percentile(d_z3, 97.5)}112        row["p_boot_zip3"] = 2 * (1 - stats.norm.cdf(abs(row["z_boot_zip3"])))113        row["significant_5pct_zip3"] = bool(abs(row["z_boot_zip3"]) > 1.96)114        rows.append(row)115    out = pd.DataFrame(rows)116    out["n_draws"] = args.draws117    out["n_zip3_clusters"] = len(np.unique(codes))118    out.to_csv(OUT / "iqr_wald_bootstrap.csv", index=False)119    print(out[["variable", "difference", "z_independent", "z_boot_obs", "z_boot_zip3",120               "significant_5pct_zip3"]].round(3).to_string(index=False))121    print(f"\nDone in {time.time() - t0:.0f}s")122123124if __name__ == "__main__":125    main()126