1"""r3k_03_analysis.py — core asset-pricing tests on the Russell 3000 panel (WP 0002).2Signals: AI_ADOPT, AI_INFRA, AI_RISK (section-weighted cosine similarity), AI_COMPOS = mean z(ADOPT), z(INFRA);3SPID_TECH = log(1 + technology-project mentions), SPID_AI = any AI-initiative mention (SPID database).4Blocks: summary, quintile sorts (EW/VW/market-adjusted), monthly long-short + CAPM/FF6 alphas, two-way FE panel,5Fama-MacBeth, post-ChatGPT break, sub-period spreads, size groups (Russell 1000 / 2000 proxies), S&P 500 vs non-S&P.6Output: results/r3k/results.json, data/processed_r3k/analysis_panel.parquet, ls_ts_*.parquet7"""8import json, warnings, numpy as np, pandas as pd, statsmodels.api as sm, statsmodels.formula.api as smf9from statsmodels.stats.sandwich_covariance import cov_hac10from linearmodels.panel import PanelOLS11warnings.filterwarnings("ignore")12D = "data/processed_r3k"; T = "results/r3k"13m = pd.read_parquet(f"{D}/merged_panel.parquet"); ret = pd.read_parquet(f"{D}/monthly_returns.parquet")14ff = pd.read_parquet(f"{D}/ff5_monthly.parquet").merge(pd.read_parquet(f"{D}/mom_monthly.parquet"), on="date", how="left")15RES = {}16SIGNALS = {"AI_ADOPT": "ai_adoption_wmean", "AI_INFRA": "ai_infrastructure_wmean", "AI_RISK": "ai_disruption_risk_wmean"}17PRE, POST = range(2010, 2022), range(2022, 2026) # formation cohorts; 2025 cohort holds Jul-2026.. (partial) -> dropped below18df = m[m.fwd_ret.notna() & (m.fiscal_year <= 2024)].copy() # 2025 cohort has < 6 months of forward returns1920def wins(s, lo=.01, hi=.99): return s.clip(s.quantile(lo), s.quantile(hi))21for c, w in [("fwd_ret", "fwd_w"), ("fwd_ret_mktadj", "fwd_madj_w"), ("mom_12m", "mom_w"), ("vol_12m", "vol_w")]: df[w] = wins(df[c])22def zscore(g): return (g - g.mean()) / g.std(ddof=0)23for k, col in SIGNALS.items(): df[f"z_{k}"] = df.groupby("fiscal_year")[col].transform(zscore)24df["z_AI_COMPOS"] = df[["z_AI_ADOPT", "z_AI_INFRA"]].mean(axis=1)25df["spid_tech_log"] = np.log1p(df["spid_mentions_tech"]); df["spid_ai_any"] = (df["spid_mentions_ai"] + df["spid_mentions_ai_tech_tag"] > 0).astype(float)26df["z_SPID_TECH"] = df.groupby("fiscal_year")["spid_tech_log"].transform(zscore)27df["z_size"] = df.groupby("fiscal_year")["log_mktcap"].transform(zscore); df["z_mom"] = df.groupby("fiscal_year")["mom_w"].transform(zscore); df["z_vol"] = df.groupby("fiscal_year")["vol_w"].transform(zscore)28df["post"] = (df.fiscal_year >= 2022).astype(int)29RES["n_obs"] = int(len(df)); RES["n_firms"] = int(df.ticker.nunique()); RES["yrs"] = f"{int(df.fiscal_year.min())}-{int(df.fiscal_year.max())}"30RES["n_by_year"] = df.groupby("fiscal_year").size().to_dict(); RES["n_sp500_firm_years"] = int(df.in_sp500.sum()); RES["n_large"] = int((df.size_group == "large").sum())31print("N", len(df), "firms", df.ticker.nunique(), RES["yrs"])32summ_cols = list(SIGNALS.values()) + ["spid_mentions_tech", "spid_ai_any", "fwd_ret", "fwd_ret_mktadj", "mom_12m", "vol_12m", "log_mktcap", "total_chunks"]33df[summ_cols].describe(percentiles=[.05, .25, .5, .75, .95]).T.to_csv(f"{T}/t1_summary_stats.csv")34df[["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_SPID_TECH", "z_size", "z_mom", "fwd_ret"]].corr().to_csv(f"{T}/t1b_correlations.csv")35RES["corr"] = df[["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_SPID_TECH"]].corr().round(3).to_dict()3637def nw_t(x, lags=3):38 x = pd.Series(x).dropna().values39 if len(x) < 3: return (float(x.mean()) if len(x) else np.nan), np.nan40 r = sm.OLS(x, np.ones((len(x), 1))).fit(); return float(x.mean()), float(x.mean() / np.sqrt(cov_hac(r, nlags=lags)[0, 0]))4142def sort_table(d, z, weight=None, oc="fwd_ret", minn=40):43 ann = {q: [] for q in range(1, 6)}; ls = []44 for yr, g in d.groupby("fiscal_year"):45 g = g.dropna(subset=[z, oc]).copy()46 if len(g) < minn: continue47 g["q"] = pd.qcut(g[z].rank(method="first"), 5, labels=False) + 148 for q in range(1, 6):49 s = g[g.q == q]50 if s.empty: continue51 if weight is None: ann[q].append(s[oc].mean())52 else:53 w = s[weight].clip(lower=0).fillna(0); ann[q].append(np.average(s[oc], weights=w) if w.sum() > 0 else s[oc].mean())54 if ann[5] and ann[1]: ls.append(ann[5][-1] - ann[1][-1])55 out = {f"Q{q}": nw_t(ann[q])[0] for q in range(1, 6)}56 mu, t = nw_t(ls); out["Q5_minus_Q1"] = mu; out["t_stat"] = t; out["n_years"] = len(ls); return out5758sorts = {}59for k in ["AI_ADOPT", "AI_INFRA", "AI_RISK", "AI_COMPOS", "SPID_TECH"]:60 sorts[f"{k}_EW"] = sort_table(df, f"z_{k}"); sorts[f"{k}_VW"] = sort_table(df, f"z_{k}", "mktcap"); sorts[f"{k}_EW_madj"] = sort_table(df, f"z_{k}", None, "fwd_ret_mktadj")61sd = pd.DataFrame(sorts).T; sd.to_csv(f"{T}/t2_portfolio_sorts.csv"); RES["sorts"] = sd.round(4).to_dict("index")62print(sd[["Q1", "Q5", "Q5_minus_Q1", "t_stat", "n_years"]].round(4).to_string())6364# ---- monthly long-short + alphas ----65ff["date"] = pd.to_datetime(ff["date"]).dt.to_period("M").dt.to_timestamp("M"); ff_ix = ff.set_index("date")66ret["date"] = pd.to_datetime(ret["date"]).dt.to_period("M").dt.to_timestamp("M")67ret_by = {t: g.set_index("date")["ret"] for t, g in ret.groupby("ticker")}68def ls_monthly(d, z, n=5, minn=40):69 monthly = {}70 for yr, g in d.groupby("fiscal_year"):71 g = g.dropna(subset=[z]).copy()72 if len(g) < minn: continue73 g["q"] = pd.qcut(g[z].rank(method="first"), n, labels=False) + 174 hold = pd.date_range(pd.Timestamp(yr + 1, 7, 31), pd.Timestamp(yr + 2, 6, 30), freq="ME")75 for q in (1, n):76 tks = [t for t in g[g.q == q].ticker if t in ret_by]77 for dt in hold:78 r = [ret_by[t].get(dt, np.nan) for t in tks]; r = [x for x in r if pd.notna(x)]79 if len(r) >= 10: monthly.setdefault(dt, {}).setdefault(q, []).append(np.mean(r))80 recs = [{"date": dt, "lo": np.mean(v[1]), "hi": np.mean(v[n])} for dt, v in sorted(monthly.items()) if 1 in v and n in v]81 ts = pd.DataFrame(recs); ts["ls"] = ts["hi"] - ts["lo"]; return ts82def factor_alpha(ts):83 s = ts.merge(ff_ix.reset_index(), on="date", how="inner").dropna(subset=["ls", "MktRF", "SMB", "HML", "RMW", "CMA", "MOM"]); y = s["ls"]; out = {}84 r1 = sm.OLS(y, sm.add_constant(s[["MktRF"]])).fit(cov_type="HAC", cov_kwds={"maxlags": 3})85 r2 = sm.OLS(y, sm.add_constant(s[["MktRF", "SMB", "HML", "RMW", "CMA", "MOM"]])).fit(cov_type="HAC", cov_kwds={"maxlags": 3})86 mu, t = nw_t(y.values)87 out.update(mean_ls_ann_pct=round(mu * 1200, 2), hac_t=round(t, 2), sharpe=round(mu / y.std() * np.sqrt(12), 2),88 capm_alpha_ann_pct=round(r1.params["const"] * 1200, 2), capm_alpha_t=round(r1.tvalues["const"], 2),89 ff6_alpha_ann_pct=round(r2.params["const"] * 1200, 2), ff6_alpha_t=round(r2.tvalues["const"], 2), n_months=int(len(s)),90 betas={k: round(r2.params[k], 3) for k in ["MktRF", "SMB", "HML", "RMW", "CMA", "MOM"]})91 return out92RES["ls_alpha"] = {}93POST_START = pd.Timestamp("2023-07-31")94for k in ["AI_ADOPT", "AI_INFRA", "AI_COMPOS", "AI_RISK", "SPID_TECH"]:95 ts = ls_monthly(df, f"z_{k}"); ts.to_parquet(f"{D}/ls_ts_{k}.parquet", index=False)96 RES["ls_alpha"][k] = {"full": factor_alpha(ts), "pre": factor_alpha(ts[ts.date < POST_START]), "post": factor_alpha(ts[ts.date >= POST_START])}97 print(k, "LS full", RES["ls_alpha"][k]["full"]["mean_ls_ann_pct"], "ff6a", RES["ls_alpha"][k]["full"]["ff6_alpha_ann_pct"], "(t", RES["ls_alpha"][k]["full"]["ff6_alpha_t"], ") | post", RES["ls_alpha"][k]["post"]["mean_ls_ann_pct"], "sharpe", RES["ls_alpha"][k]["post"]["sharpe"])9899# ---- two-way FE panel ----100pdf = df.dropna(subset=["fwd_w", "z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_size", "z_mom", "z_vol"]).set_index(["ticker", "fiscal_year"])101specs = {"M1_adopt": "fwd_w ~ z_AI_ADOPT", "M2_adopt_ctrl": "fwd_w ~ z_AI_ADOPT + z_size + z_mom + z_vol", "M3_infra_ctrl": "fwd_w ~ z_AI_INFRA + z_size + z_mom + z_vol",102 "M4_risk_ctrl": "fwd_w ~ z_AI_RISK + z_size + z_mom + z_vol", "M5_all": "fwd_w ~ z_AI_ADOPT + z_AI_INFRA + z_AI_RISK + z_size + z_mom + z_vol",103 "M6_spid": "fwd_w ~ z_SPID_TECH + z_size + z_mom + z_vol", "M7_all_spid": "fwd_w ~ z_AI_COMPOS + z_SPID_TECH + z_size + z_mom + z_vol"}104RES["panel_fe"] = {}105for lab, f in specs.items():106 r = PanelOLS.from_formula(f + " + EntityEffects + TimeEffects", data=pdf, drop_absorbed=True).fit(cov_type="clustered", cluster_entity=True, cluster_time=True)107 RES["panel_fe"][lab] = {"coefs": {v: {"coef": round(r.params[v], 4), "t": round(r.tstats[v], 2)} for v in r.params.index}, "n": int(r.nobs), "r2_within": round(r.rsquared_within, 4)}108 print(lab, {v: (round(r.params[v], 4), round(r.tstats[v], 2)) for v in r.params.index if v.startswith("z_")})109110# ---- Fama-MacBeth ----111def fama_macbeth(d, xvars, minn=40, oc="fwd_ret"):112 coefs = {v: [] for v in ["const"] + xvars}113 for yr, g in d.groupby("fiscal_year"):114 g = g.dropna(subset=[oc] + xvars)115 if len(g) < minn: continue116 r = sm.OLS(g[oc], sm.add_constant(g[xvars])).fit()117 for v in coefs: coefs[v].append(r.params.get(v, np.nan))118 out = {v: dict(zip(["mean_coef", "t"], [round(x, 4) for x in nw_t(coefs[v])])) for v in coefs}; out["n_years"] = len(coefs["const"]); return out119RES["fama_macbeth"] = {"adopt": fama_macbeth(df, ["z_AI_ADOPT", "z_size", "z_mom", "z_vol"]), "infra": fama_macbeth(df, ["z_AI_INFRA", "z_size", "z_mom", "z_vol"]),120 "all": fama_macbeth(df, ["z_AI_ADOPT", "z_AI_INFRA", "z_AI_RISK", "z_size", "z_mom", "z_vol"]), "compos": fama_macbeth(df, ["z_AI_COMPOS", "z_size", "z_mom", "z_vol"]),121 "spid": fama_macbeth(df, ["z_SPID_TECH", "z_size", "z_mom", "z_vol"]), "compos_spid": fama_macbeth(df, ["z_AI_COMPOS", "z_SPID_TECH", "z_size", "z_mom", "z_vol"])}122for k, v in RES["fama_macbeth"].items(): print("FM", k, {x: v[x] for x in v if x.startswith("z_AI") or x.startswith("z_SPID")})123124# ---- structural break + sub-period spreads ----125RES["structural_break"] = {}126for k in ["AI_ADOPT", "AI_INFRA", "AI_COMPOS", "SPID_TECH"]:127 d2 = df.dropna(subset=["fwd_w", f"z_{k}", "z_size", "z_mom"])128 r = smf.ols(f"fwd_w ~ z_{k} * post + z_size + z_mom", data=d2).fit(cov_type="cluster", cov_kwds={"groups": d2["fiscal_year"]})129 rf = smf.ols(f"fwd_w ~ z_{k} * post + z_size + z_mom", data=d2).fit(cov_type="cluster", cov_kwds={"groups": d2["ticker"]})130 RES["structural_break"][k] = {"pre_slope": round(r.params[f"z_{k}"], 4), "interaction": round(r.params[f"z_{k}:post"], 4), "interaction_t_year": round(r.tvalues[f"z_{k}:post"], 2),131 "interaction_t_firm": round(rf.tvalues[f"z_{k}:post"], 2), "post_slope": round(r.params[f"z_{k}"] + r.params[f"z_{k}:post"], 4)}132def spread_sub(d, z, yrs, oc="fwd_ret", minn=40):133 sp = {}134 for yr, g in d[d.fiscal_year.isin(list(yrs))].groupby("fiscal_year"):135 g = g.dropna(subset=[z, oc]).copy()136 if len(g) < minn: continue137 g["q"] = pd.qcut(g[z].rank(method="first"), 5, labels=False) + 1; sp[int(yr)] = float(g[g.q == 5][oc].mean() - g[g.q == 1][oc].mean())138 mu, t = nw_t(list(sp.values())); return {"mean": round(mu, 4), "t": round(t, 2) if t == t else None, "n": len(sp), "by_cohort": {k: round(v, 4) for k, v in sp.items()}}139RES["subperiod_spread"] = {k: {"pre": spread_sub(df, f"z_{k}", PRE), "post": spread_sub(df, f"z_{k}", POST)} for k in ["AI_ADOPT", "AI_INFRA", "AI_COMPOS", "AI_RISK", "SPID_TECH"]}140for k, v in RES["subperiod_spread"].items(): print(k, "pre", v["pre"]["mean"], v["pre"]["t"], "| post", v["post"]["mean"], v["post"]["t"], v["post"]["by_cohort"])141142# ---- size groups and index membership ----143RES["groups"] = {}144for gname, mask in [("large_top1000", df.size_group == "large"), ("small_rest", df.size_group == "small"), ("sp500", df.in_sp500 == True), ("non_sp500", df.in_sp500 == False)]:145 d = df[mask]146 RES["groups"][gname] = {"n": int(len(d)), "firms": int(d.ticker.nunique()), "sort_EW": sort_table(d, "z_AI_COMPOS"), "pre": spread_sub(d, "z_AI_COMPOS", PRE), "post": spread_sub(d, "z_AI_COMPOS", POST),147 "fm_compos": fama_macbeth(d, ["z_AI_COMPOS", "z_size", "z_mom", "z_vol"])}148 ts = ls_monthly(d, "z_AI_COMPOS"); ts.to_parquet(f"{D}/ls_ts_AI_COMPOS_{gname}.parquet", index=False)149 RES["groups"][gname]["ls"] = {"full": factor_alpha(ts), "pre": factor_alpha(ts[ts.date < POST_START]), "post": factor_alpha(ts[ts.date >= POST_START])}150 print(gname, "n", len(d), "spread full", RES["groups"][gname]["sort_EW"]["Q5_minus_Q1"], "pre", RES["groups"][gname]["pre"]["mean"], "post", RES["groups"][gname]["post"]["mean"], "ff6 post", RES["groups"][gname]["ls"]["post"]["ff6_alpha_ann_pct"])151# within-year sort re-done within size group (size-neutral AI quintiles)152df["z_AI_COMPOS_sizeneu"] = df.groupby(["fiscal_year", "size_group"])["z_AI_COMPOS"].transform(lambda g: zscore(g) if g.std(ddof=0) > 0 else g * 0)153RES["size_neutral"] = {"full": sort_table(df, "z_AI_COMPOS_sizeneu"), "pre": spread_sub(df, "z_AI_COMPOS_sizeneu", PRE), "post": spread_sub(df, "z_AI_COMPOS_sizeneu", POST)}154155# ---- side-by-side samples: Russell 3000 (all) vs S&P 500 members only ----156RES["samples"] = {}157for sname, d in [("r3k", df), ("sp500", df[df.in_sp500 == True].copy())]:158 for k, col in SIGNALS.items(): d[f"zz_{k}"] = d.groupby("fiscal_year")[col].transform(zscore)159 d["zz_AI_COMPOS"] = d[["zz_AI_ADOPT", "zz_AI_INFRA"]].mean(axis=1); d["zz_SPID_TECH"] = d.groupby("fiscal_year")["spid_tech_log"].transform(zscore)160 for c in ["size", "mom", "vol"]: d[f"zz_{c}"] = d.groupby("fiscal_year")[{"size": "log_mktcap", "mom": "mom_w", "vol": "vol_w"}[c]].transform(zscore)161 blk = {"n": int(len(d)), "firms": int(d.ticker.nunique()), "sorts": {}, "ls": {}, "fm": {}, "break": {}, "sub": {}}162 for k in ["AI_ADOPT", "AI_INFRA", "AI_RISK", "AI_COMPOS", "SPID_TECH"]:163 blk["sorts"][f"{k}_EW"] = sort_table(d, f"zz_{k}"); blk["sorts"][f"{k}_VW"] = sort_table(d, f"zz_{k}", "mktcap")164 ts = ls_monthly(d, f"zz_{k}"); ts.to_parquet(f"{D}/ls_ts_{k}_{sname}.parquet", index=False)165 blk["ls"][k] = {"full": factor_alpha(ts), "pre": factor_alpha(ts[ts.date < POST_START]), "post": factor_alpha(ts[ts.date >= POST_START])}166 blk["sub"][k] = {"pre": spread_sub(d, f"zz_{k}", PRE), "post": spread_sub(d, f"zz_{k}", POST)}167 if k != "AI_RISK":168 d2 = d.dropna(subset=["fwd_w", f"zz_{k}", "zz_size", "zz_mom"])169 r = smf.ols(f"fwd_w ~ zz_{k} * post + zz_size + zz_mom", data=d2).fit(cov_type="cluster", cov_kwds={"groups": d2["ticker"]})170 blk["break"][k] = {"pre_slope": round(r.params[f"zz_{k}"], 4), "interaction": round(r.params[f"zz_{k}:post"], 4), "interaction_t_firm": round(r.tvalues[f"zz_{k}:post"], 2)}171 blk["fm"] = {"compos": fama_macbeth(d, ["zz_AI_COMPOS", "zz_size", "zz_mom", "zz_vol"]), "all": fama_macbeth(d, ["zz_AI_ADOPT", "zz_AI_INFRA", "zz_AI_RISK", "zz_size", "zz_mom", "zz_vol"]), "spid": fama_macbeth(d, ["zz_SPID_TECH", "zz_size", "zz_mom", "zz_vol"])}172 # cohort spreads + pooled firm-clustered post test173 cs = spread_sub(d, "zz_AI_COMPOS", range(2010, 2025))["by_cohort"]; blk["cohorts_pct"] = {k_: round(v * 100, 1) for k_, v in cs.items()}174 dd = d[d.fiscal_year.isin(list(POST))].dropna(subset=["zz_AI_COMPOS", "fwd_ret"]).copy(); dd["q"] = dd.groupby("fiscal_year")["zz_AI_COMPOS"].transform(lambda s_: pd.qcut(s_.rank(method="first"), 5, labels=False) + 1)175 dd = dd[dd.q.isin([1, 5])].copy(); dd["hi"] = (dd.q == 5).astype(float); dd["yr"] = dd.fiscal_year.astype(str)176 r = smf.ols("fwd_ret ~ hi + C(yr)", data=dd).fit(cov_type="cluster", cov_kwds={"groups": dd["ticker"]}); rw = smf.ols("fwd_w ~ hi + C(yr)", data=dd).fit(cov_type="cluster", cov_kwds={"groups": dd["ticker"]})177 blk["pooled_post"] = {"coef_pct": round(r.params["hi"] * 100, 2), "t": round(r.tvalues["hi"], 2), "n": int(r.nobs), "wins_pct": round(rw.params["hi"] * 100, 2), "wins_t": round(rw.tvalues["hi"], 2)}178 RES["samples"][sname] = blk179 print("SAMPLE", sname, "n", blk["n"], "compos EW", round(blk["sorts"]["AI_COMPOS_EW"]["Q5_minus_Q1"], 4), "pre", blk["sub"]["AI_COMPOS"]["pre"]["mean"], "post", blk["sub"]["AI_COMPOS"]["post"]["mean"], "LS post sharpe", blk["ls"]["AI_COMPOS"]["post"]["sharpe"], "cohorts", blk["cohorts_pct"])180181json.dump(RES, open(f"{T}/results.json", "w"), indent=2, default=str)182df.to_parquet(f"{D}/analysis_panel.parquet", index=False)183print("saved results.json")184