1"""High-level driver: calibrate a scenario, solve its perfect-foresight transition, save results.23Strategy (docs/replication.md): (1) calibrate the 2026 base point consistently with the scenario's4anticipated technology paths; (2) homotopy in the size of the technology shock on a short horizon;5(3) horizon continuation up to the full horizon; (4) store paths, residual diagnostics and a horizon check.6"""7from __future__ import annotations89import json10import os11import time1213import numpy as np1415from .calibration import calibrate16from .transition import KINDS, TwoClassModel1718SCENARIO_RATE_KEYS = ("g_pQ", "g_eps", "g_alpha")192021def _params_for(overrides, lam, T, param_overrides=None):22 ov = dict(overrides or {})23 base = {k: ov.get(k, None) for k in SCENARIO_RATE_KEYS}24 p, info = calibrate(overrides=ov)25 # scale the technology rates by lambda (calibration stays consistent with the scaled rates)26 scaled = {k: lam * (base[k] if base[k] is not None else getattr(p, k)) for k in SCENARIO_RATE_KEYS}27 p, info = calibrate(overrides={**ov, **scaled})28 p = p.replace(T=T, **(param_overrides or {}))29 if lam < 1.0 and p.alpha_shock_year is not None:30 # for a jump scenario, homotopy scales the log distance of the jump instead31 p = p.replace(alpha_shock_year=None, g_alpha=lam * 0.5)32 return p, info333435def extend_guess(u_prev, T_prev, T_new, nv=6):36 U = u_prev.reshape(T_prev, nv)37 extra = T_new - T_prev38 Ul = np.vstack([U, np.tile(U[-1], (extra, 1))])39 g = U[-1, :5] - U[-2, :5]40 for j in range(extra):41 Ul[T_prev + j, :5] = U[-1, :5] + g * (j + 1)42 return Ul.ravel()434445def solve_scenario(name, overrides=None, param_overrides=None, T=275, T_short=None,46 lambdas=(1.0,), fallback_lambdas=(0.0, 0.25, 0.5, 0.75, 1.0), horizons=(),47 tol=1e-9, verbose=True, outdir=None, horizon_check_years=50):48 t0 = time.time()49 log = []50 u, T_prev = None, None51 T_short = T if T_short is None else T_short # homotopy on the full horizon by default (terminal near BGP)52 # --- homotopy53 for lam in lambdas:54 p, info = _params_for(overrides, lam, T_short, param_overrides)55 m = TwoClassModel(p)56 if u is None:57 u = m.smart_guess()58 if verbose:59 print(f"[{name}] homotopy lambda={lam:.2f} T={T_short}", flush=True)60 u, ok, hist = m.solve(u, tol=tol, verbose=verbose, maxit=40)61 log.append(dict(stage="homotopy", lam=lam, T=T_short, converged=bool(ok), res=float(hist[-1])))62 if not ok and fallback_lambdas and tuple(lambdas) != tuple(fallback_lambdas):63 # direct solve failed: homotopy in the shock size from the frozen economy64 if verbose:65 print(f"[{name}] direct solve not converged (|res|={hist[-1]:.2e}); switching to homotopy", flush=True)66 u = None67 for lam in fallback_lambdas:68 p, info = _params_for(overrides, lam, T_short, param_overrides)69 m = TwoClassModel(p)70 if u is None:71 u = m.smart_guess()72 if verbose:73 print(f"[{name}] homotopy lambda={lam:.2f} T={T_short}", flush=True)74 u, ok, hist = m.solve(u, tol=tol, verbose=verbose, maxit=40)75 log.append(dict(stage="homotopy-fallback", lam=lam, T=T_short, converged=bool(ok), res=float(hist[-1])))76 T_prev = T_short77 # --- horizon continuation78 for Th in [h for h in horizons if T_short < h < T] + [T]:79 p, info = _params_for(overrides, 1.0, Th, param_overrides)80 m = TwoClassModel(p)81 u = extend_guess(u, T_prev, Th)82 if verbose:83 print(f"[{name}] horizon {Th}", flush=True)84 u, ok, hist = m.solve(u, tol=tol, verbose=verbose, maxit=40)85 log.append(dict(stage="horizon", lam=1.0, T=Th, converged=bool(ok), res=float(hist[-1])))86 T_prev = Th87 P = m.evaluate(u, full=True)88 # --- horizon (truncation) check89 hc = None90 if horizon_check_years:91 pl, _ = _params_for(overrides, 1.0, T + horizon_check_years, param_overrides)92 ml = TwoClassModel(pl)93 ul, okl, histl = ml.solve(extend_guess(u, T, T + horizon_check_years), tol=tol, verbose=False, maxit=40)94 Kl, rl, _ = ml.unpack(ul)95 K0, r0, _ = m.unpack(u)96 n = T - horizon_check_years97 hc = dict(converged=bool(okl), res=float(histl[-1]),98 max_log_dev_K={k: float(np.max(np.abs(np.log(Kl[k][1:n + 1] / K0[k][1:n + 1])))) for k in KINDS},99 max_dev_r=float(np.max(np.abs(rl[1:n + 1] - r0[1:n + 1]))), years_compared=int(n))100 out = dict(name=name, overrides=overrides or {}, param_overrides=param_overrides or {}, T=T,101 log=log, horizon_check=hc, seconds=time.time() - t0, converged=bool(ok), final_res=float(hist[-1]),102 calibration_info={k: (v.tolist() if hasattr(v, "tolist") else v) for k, v in info.items()})103 if outdir:104 os.makedirs(outdir, exist_ok=True)105 np.save(os.path.join(outdir, f"{name}_u.npy"), u)106 p.to_json(os.path.join(outdir, f"{name}_params.json"))107 save_paths(P, os.path.join(outdir, f"{name}_paths.npz"))108 with open(os.path.join(outdir, f"{name}_meta.json"), "w") as f:109 json.dump(out, f, indent=2, default=float)110 return p, m, u, P, out111112113PATH_KEYS = ("years", "r", "Y", "GDP", "wage", "wL", "wH", "A", "RT", "RTH", "RTp", "RTHp", "RH", "rentR", "PE", "util",114 "Qtot", "E", "EQ", "C_total", "C_K", "C_W", "X_K", "X_W", "P_idx", "transfers", "tax", "fund", "PT", "PTH",115 "W", "inv_goods", "foreign_flow", "labor_share", "income_residual", "Li", "Qi", "Xi", "HS", "alpha", "pQ",116 "eps", "Z", "gf", "T_eff", "TH_eff", "residuals", "goods_gap")117118119def save_paths(P, path):120 d = {}121 for k in PATH_KEYS:122 d[k] = np.asarray(P[k])123 for k in KINDS:124 d[f"K_{k}"] = P["K"][k]125 d[f"x_{k}"] = P["x"][k]126 d[f"I_{k}"] = P["I"][k]127 d[f"q_{k}"] = P["q"][k]128 d[f"R_{k}"] = P["R"][k]129 d[f"V_{k}"] = P["V"][k]130 d[f"D_{k}"] = P["D"][k]131 d[f"adj_{k}"] = P["adj"][k]132 np.savez_compressed(path, **d)133134135def load_paths(path):136 z = np.load(path)137 return {k: z[k] for k in z.files}138