1"""General equilibrium with heterogeneous households: relaxation on the interest-rate path.23Given a solved two-class scenario (paths P, model m, unknowns u), the HA block replaces the capitalists:4households of both skills (three discount types each) hold all claims. The algorithm:5 1. calibrate the discount types so that stationary HA savings at the reference prices equal the two-class6 wealth W_0 (so that the HA economy starts on the same balance sheet);7 2. iterate: (a) given r, solve the production side with r exogenous (4 capital Eulers, banded Newton) ->8 supply of claims W^s_t and wages/transfers/price index; (b) solve the HA block backward and simulate ->9 demand for claims W^d_t; (c) update r_{t+1} <- r_{t+1} - lam * log(W^d_t / W^s_t) (excess demand for10 assets lowers the equilibrium return), with a secant-type step adaptation;11 3. stop when max_t |log(W^d/W^s)| < tol.12Outputs: equilibrium paths, distributional statistics along the transition, and welfare by type.13"""14from __future__ import annotations1516import time1718import numpy as np19from scipy.optimize import brentq2021from .households import HouseholdBlock22from .steady import bgp_interest_rate23from .transition import KINDS, TwoClassModel242526def calibrate_households(p, W0, wL, wH, transfer=0.0, P=1.0, masses=(0.5, 0.3, 0.2), betas_low=(0.93, 0.97),27 verbose=True, **hb_kw):28 """Find the highest discount factor so that stationary HA savings equal W0 at the reference prices."""29 r = bgp_interest_rate(p)30 bmax = (1 + p.g_Z) ** p.gamma_u / (1 + r) - 1e-4 # no-explosion bound3132 def excess(bh):33 hb = HouseholdBlock(p, (*betas_low, bh), masses, **hb_kw)34 pols = hb.solve_stationary(r, (wL, wH), transfer, P)35 dists = hb.stationary_distribution(pols, r)36 A, X = hb.aggregate(pols, dists)37 return np.log(A / W0), hb, pols, dists3839 lo, hi = betas_low[1] + 0.002, bmax40 f_lo, f_hi = excess(lo)[0], excess(hi)[0]41 if f_lo > 0:42 bh = lo43 elif f_hi < 0:44 bh = hi45 else:46 bh = brentq(lambda b: excess(b)[0], lo, hi, xtol=1e-6)47 ex, hb, pols, dists = excess(bh)48 st = hb.wealth_stats(dists, pols)49 if verbose:50 print(f" HA calibration: beta_high = {bh:.5f}, savings/W0 = {np.exp(ex):.4f}, stats = {st}")51 return hb, pols, dists, dict(beta_high=bh, stats=st)525354def ha_inputs_from_paths(P, p):55 """Detrended wage and transfer paths for the HA block (per efficiency unit), price index, r."""56 T = len(P["Y"]) - 157 wL = P["wL"][:T + 1]58 wH = P["wH"][:T + 1]59 gf = P["gf"][:T + 1]60 mass_workers = float(np.sum(p.H))61 transfer = P["transfers"][:T + 1] / mass_workers / gf # per worker, detrended62 Pidx = P["P_idx"][:T + 1] / P["P_idx"][0]63 r = P["r"][:T + 2]64 return wL, wH, transfer, Pidx, r, gf656667def _ha_demand(hb, r_path, wL, wH, transfer, Pidx, gf, dists0, T, x_term=None):68 if x_term is None:69 x_term = hb.solve_stationary(r_path[T], (wL[T], wH[T]), transfer[T], Pidx[T])70 pols = hb.solve_transition(r_path[:T + 1], wL, wH, transfer, Pidx, x_term)71 A, X, dist_hist = hb.simulate_transition(pols, dists0, r_path)72 return A * gf[:T], X, pols, dist_hist737475def solve_ha_equilibrium(p, m: TwoClassModel, u, P, lam=1.0, tol=1e-4, maxit=60, verbose=True, hb_kw=None,76 masses=(0.5, 0.3, 0.2), betas_low=(0.93, 0.97), r_bounds=(-0.02, 0.8), jac_every=6,77 step_max=0.2, checkpoint=None, r_init=None):78 """Quasi-Newton on the return path: the Jacobian of the household block's log asset demand with respect to79 the return path is computed by finite differences (T household solves, cheap); the supply side's response is80 not differentiated (it has the opposite sign, so using the demand Jacobian alone over-estimates the required81 move, which the damping lam corrects). Steps are bounded by r_bounds."""82 hb_kw = hb_kw or {}83 T = m.T84 wL, wH, transfer, Pidx, r, gf = ha_inputs_from_paths(P, p)85 W0 = P["W"][0]86 hb, pols0, dists0, hcal = calibrate_households(p, W0, wL[0], wH[0], transfer[0], Pidx[0], masses, betas_low, verbose, **hb_kw)87 U6 = u.reshape(T, 6)88 u4 = U6[:, :5].ravel().copy() # fixed-r mode: 4 log stocks + log X_K89 r_path = r.copy() if r_init is None else np.asarray(r_init, float).copy()90 hist = []91 t0 = time.time()92 J = None93 for it in range(maxit):94 m.set_fixed_r(r_path)95 u4, ok, h = m.solve(u4, tol=1e-9, verbose=False, maxit=30)96 Pp = m.evaluate(u4, full=True)97 wL, wH, transfer, Pidx, _, gf = ha_inputs_from_paths(Pp, p)98 Ws = Pp["W"][:T]99 x_term = hb.solve_stationary(r_path[T], (wL[T], wH[T]), transfer[T], Pidx[T])100 Wd, X, pols, dist_hist = _ha_demand(hb, r_path, wL, wH, transfer, Pidx, gf, dists0, T, x_term)101 gap = np.log(Wd / Ws)102 err = float(np.max(np.abs(gap)))103 hist.append(err)104 if verbose:105 print(f" HA-GE it {it:2d}: max|log Wd/Ws| = {err:.3e}, production converged={ok}, ({time.time() - t0:.0f}s)", flush=True)106 if err < tol:107 break108 if J is None or it % jac_every == 0:109 # Jacobian of log Wd (T) w.r.t. r_{t+1} (t = 0..T-1) at the current prices110 h_fd = 1e-4111 J = np.zeros((T, T))112 for s in range(T):113 rp = r_path.copy()114 rp[s + 1] += h_fd115 Wd_p, _, _, _ = _ha_demand(hb, rp, wL, wH, transfer, Pidx, gf, dists0, T, x_term)116 J[:, s] = (np.log(Wd_p) - np.log(Wd)) / h_fd117 if verbose:118 print(f" HA-GE: household Jacobian computed ({time.time() - t0:.0f}s)", flush=True)119 # supply responds with the opposite sign; approximate its elasticity by a diagonal term120 Jtot = J + np.diag(np.full(T, 5.0)) # d log Ws / d r ~ -duration (~5 years): crude, corrected by damping121 try:122 step = np.linalg.solve(Jtot, gap)123 except np.linalg.LinAlgError:124 step = gap / 20.0125 step = np.clip(lam * step, -step_max, step_max)126 r_new = np.clip(r_path[1:T + 1] - step, r_bounds[0], r_bounds[1])127 r_path[1:T + 1] = r_new128 r_path[T + 1] = r_path[T]129 if checkpoint:130 np.save(checkpoint, r_path)131 m.set_free_r()132 stats = []133 for t in range(T):134 st = hb.wealth_stats(dist_hist[t], {k: pols[k][t] for k in pols})135 st["year"] = int(P["years"][t])136 stats.append(st)137 out = dict(r=r_path, Wd=Wd, Ws=Ws, gap=gap, hist=hist, stats=stats, hcal=hcal, X=X, paths=Pp, u4=u4,138 converged=err < tol)139 return out, hb, pols, dist_hist140