1"""Heterogeneous-agent general equilibrium for a solved scenario.2Usage: python3 scripts/30_ha_equilibrium.py <scenario> [maxit] [lam]3Outputs: results/ha/<scenario>_ha.npz (r path, wealth stats by year, consumption stats), results/ha/<scenario>_ha.json"""4import json5import os6import sys7import time89import numpy as np1011sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))12from agi_ge.ha_equilibrium import solve_ha_equilibrium # noqa: E40213from agi_ge.params import Params # noqa: E40214from agi_ge.transition import TwoClassModel # noqa: E4021516ROOT = os.path.join(os.path.dirname(__file__), "..")17name = sys.argv[1]18maxit = int(sys.argv[2]) if len(sys.argv) > 2 else 4019lam = float(sys.argv[3]) if len(sys.argv) > 3 and not sys.argv[3].startswith('--') else 1.020SC = os.path.join(ROOT, "results", "scenarios")21p = Params.from_json(os.path.join(SC, f"{name}_params.json"))22u = np.load(os.path.join(SC, f"{name}_u.npy"))23m = TwoClassModel(p)24P = m.evaluate(u, full=True)25t0 = time.time()26ckpt = os.path.join(ROOT, "results", "ha", f"{name}_r_checkpoint.npy")27os.makedirs(os.path.join(ROOT, "results", "ha"), exist_ok=True)28r_init = np.load(ckpt) if (os.path.exists(ckpt) and "--resume" in sys.argv) else None29out, hb, pols, dist_hist = solve_ha_equilibrium(p, m, u, P, lam=lam, maxit=maxit, verbose=True, checkpoint=ckpt, r_init=r_init)30os.makedirs(os.path.join(ROOT, "results", "ha"), exist_ok=True)31stats = out["stats"]32years = np.array([s["year"] for s in stats])33Pp = out["paths"]34np.savez_compressed(os.path.join(ROOT, "results", "ha", f"{name}_ha.npz"), r=out["r"], Wd=out["Wd"], Ws=out["Ws"], years=years,35 gini=np.array([s["gini"] for s in stats]), top1=np.array([s["top1"] for s in stats]),36 top10=np.array([s["top10"] for s in stats]), bottom50=np.array([s["bottom50"] for s in stats]),37 X=out["X"], r_twoclass=P["r"], W_twoclass=P["W"], GDP=Pp["GDP"], labor_share=Pp["labor_share"], wL=Pp["wL"], wH=Pp["wH"])38meta = dict(name=name, converged=bool(out["converged"]), hist=[float(h) for h in out["hist"]], beta_high=out["hcal"]["beta_high"],39 stats_2026=stats[0], seconds=time.time() - t0,40 stats_by_year={int(y): {k: float(v) for k, v in s.items() if k != "year"} for y, s in zip(years, stats) if y in (2026, 2030, 2040, 2050, 2075, 2100, 2150, 2200)})41with open(os.path.join(ROOT, "results", "ha", f"{name}_ha.json"), "w") as f:42 json.dump(meta, f, indent=1, default=float)43print(json.dumps({k: meta[k] for k in ("name", "converged", "beta_high", "seconds")}, default=float))44print("r (two-class vs HA) 2026/2050/2100:", [(float(P["r"][t + 1]), float(out["r"][t + 1])) for t in (0, 24, 74)])45for y in (2026, 2050, 2100, 2200):46 if y in meta["stats_by_year"]:47 print(y, meta["stats_by_year"][y])48