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-431 m = TwoClassModel(p)32 u0 = m.initial_guess(r0=bgp_interest_rate(p), g=p.g_Z)33 f0 = m.evaluate(u0)34 # the BGP guess solves the dynamic system up to the adjustment-cost term omitted from the base-point35 # compute rental (phi_Q g (delta_Q + g/2) ~ 4e-4; see docs/assumptions_register.md, C13)36 assert np.abs(f0).max() < 1e-3373839def test_banded_jacobian_matches_dense(frozen):40 p, info = frozen41 m = TwoClassModel(p.replace(T=6))42 u0 = m.smart_guess()43 f0 = m.evaluate(u0)44 n = len(u0)45 Jd = np.zeros((n, n))46 h = 1e-547 for i in range(n):48 up = u0.copy(); up[i] += h49 Jd[:, i] = (m.evaluate(up) - f0) / h50 Jb = m.jacobian(u0, f0, h=h).toarray()51 assert np.abs(Jd - Jb).max() < 1e-6 * max(1.0, np.abs(Jd).max())525354def test_newton_converges_on_small_shock(frozen):55 p, info = calibrate(overrides=dict(g_pQ=0.03, g_eps=0.02, g_alpha=0.03))56 m = TwoClassModel(p.replace(T=12))57 u, ok, hist = m.solve(m.smart_guess(), tol=1e-9, verbose=False, maxit=30)58 assert ok59 P = m.evaluate(u, full=True)60 assert np.abs(P["income_residual"]).max() < 1e-861 assert np.abs(P["goods_gap"][:m.T]).max() < 1e-9 # t = T is the extrapolated terminal period62 # no-arbitrage on the aggregate claim: (1+r_{t+1}) W_t = payout_{t+1} + W_{t+1}63 T = m.T64 payout = sum(P["D"][k] for k in KINDS) + P["RTp"] + P["RTHp"]65 gap = (1 + P["r"][1:T + 1]) * P["W"][:T] - payout[1:T + 1] - P["W"][1:T + 1]66 assert np.abs(gap / P["W"][:T]).max() < 1e-6676869saved = sorted(glob.glob(os.path.join(SC, "*_paths.npz")))707172@pytest.mark.parametrize("path", saved, ids=[os.path.basename(s) for s in saved])73def test_saved_scenario_is_an_equilibrium(path):74 P = load_paths(path)75 T = len(P["Y"]) - 176 assert np.abs(P["residuals"]).max() < 1e-777 assert np.abs(P["income_residual"] / P["GDP"]).max() < 1e-678 assert np.abs(P["goods_gap"][:T]).max() < 1e-7 # t = T is the extrapolated terminal period79 assert np.all(P["X_K"] > 0) and np.all(P["C_W"] > 0)80 meta = json.load(open(path.replace("_paths.npz", "_meta.json")))81 assert meta["converged"]82 prm = json.load(open(path.replace("_paths.npz", "_params.json")))83 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:84 # no-arbitrage on the aggregate claim: (1+r) W_t = payout_{t+1} + W_{t+1}85 T = len(P["Y"]) - 186 payout = sum(P[f"D_{k}"] for k in KINDS) + P["RTp"] + P["RTHp"]87 gap = (1 + P["r"][1:T + 1]) * P["W"][:T] - payout[1:T + 1] - P["W"][1:T + 1]88 assert np.abs(gap / P["W"][:T]).max() < 1e-689