1"""Tests of the dynamic blocks: exact frozen BGP, Jacobian accuracy, and equilibrium checks of saved scenarios."""2import glob3import json4import os5import sys67import numpy as np8import pytest910sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))11from agi_ge.calibration import calibrate # noqa: E40212from agi_ge.solve import load_paths # noqa: E40213from agi_ge.steady import bgp_interest_rate, solve_bgp # noqa: E40214from agi_ge.transition import KINDS, TwoClassModel # noqa: E4021516ROOT = os.path.join(os.path.dirname(__file__), "..")17SC = os.path.join(ROOT, "results", "scenarios")181920@pytest.fixture(scope="module")21def frozen():22 p, info = calibrate(overrides=dict(g_pQ=0.0, g_eps=0.0, g_alpha=0.0))23 return p.replace(T=12), info242526def test_frozen_bgp_is_exact(frozen):27 p, info = frozen28 b = solve_bgp(p.replace(alpha_freeze=True))29 for k in ("N", "E", "H"):30 assert abs(np.log(b["K"][k] / {"N": p.KN0, "E": p.KE0, "H": p.KH0}[k])) < 1e-631 m = TwoClassModel(p)32 u0 = m.initial_guess(r0=bgp_interest_rate(p), g=p.g_Z)33 f0 = m.evaluate(u0)34 assert np.abs(f0).max() < 1e-6 # the BGP guess solves the dynamic system353637def test_banded_jacobian_matches_dense(frozen):38 p, info = frozen39 m = TwoClassModel(p.replace(T=6))40 u0 = m.smart_guess()41 f0 = m.evaluate(u0)42 n = len(u0)43 Jd = np.zeros((n, n))44 h = 1e-545 for i in range(n):46 up = u0.copy(); up[i] += h47 Jd[:, i] = (m.evaluate(up) - f0) / h48 Jb = m.jacobian(u0, f0, h=h).toarray()49 assert np.abs(Jd - Jb).max() < 1e-6 * max(1.0, np.abs(Jd).max())505152def test_newton_converges_on_small_shock(frozen):53 p, info = calibrate(overrides=dict(g_pQ=0.03, g_eps=0.02, g_alpha=0.03))54 m = TwoClassModel(p.replace(T=12))55 u, ok, hist = m.solve(m.smart_guess(), tol=1e-9, verbose=False, maxit=30)56 assert ok57 P = m.evaluate(u, full=True)58 assert np.abs(P["income_residual"]).max() < 1e-859 assert np.abs(P["goods_gap"]).max() < 1e-960 # no-arbitrage on the aggregate claim: (1+r_{t+1}) W_t = payout_{t+1} + W_{t+1}61 T = m.T62 payout = sum(P["D"][k] for k in KINDS) + P["RTp"] + P["RTHp"]63 gap = (1 + P["r"][1:T + 1]) * P["W"][:T] - payout[1:T + 1] - P["W"][1:T + 1]64 assert np.abs(gap / P["W"][:T]).max() < 1e-6656667saved = sorted(glob.glob(os.path.join(SC, "*_paths.npz")))686970@pytest.mark.parametrize("path", saved, ids=[os.path.basename(s) for s in saved])71def test_saved_scenario_is_an_equilibrium(path):72 P = load_paths(path)73 assert np.abs(P["residuals"]).max() < 1e-774 assert np.abs(P["income_residual"] / P["GDP"]).max() < 1e-775 assert np.abs(P["goods_gap"]).max() < 1e-776 assert np.all(P["X_K"] > 0) and np.all(P["C_W"] > 0)77 meta = json.load(open(path.replace("_paths.npz", "_meta.json")))78 assert meta["converged"]79 prm = json.load(open(path.replace("_paths.npz", "_params.json")))80 if prm["tau_T"] == 0 and prm["tau_K"] == 0 and prm["tau_Q"] == 0 and prm["markup_Q"] == 1 and prm["foreign_share_Q"] == 0:81 # no-arbitrage on the aggregate claim: (1+r) W_t = payout_{t+1} + W_{t+1}82 T = len(P["Y"]) - 183 payout = sum(P[f"D_{k}"] for k in KINDS) + P["RTp"] + P["RTHp"]84 gap = (1 + P["r"][1:T + 1]) * P["W"][:T] - payout[1:T + 1] - P["W"][1:T + 1]85 assert np.abs(gap / P["W"][:T]).max() < 1e-686