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, scales=None):68 if x_term is None:69 x_term = hb.solve_stationary(r_path[T], (wL[T], wH[T]), transfer[T], Pidx[T], scale=1.0 if scales is None else scales[T])70 pols = hb.solve_transition(r_path[:T + 1], wL, wH, transfer, Pidx, x_term, scales=scales)71 A, X, dist_hist = hb.simulate_transition(pols, dists0, r_path, scales=scales)72 return A * gf[:T], X, pols, dist_hist737475def grid_scales(Ws, gf, W0, T):76 """Grid scale per period: detrended aggregate wealth relative to 2026, never below one, smoothed."""77 sc = np.ones(T + 2)78 w = np.maximum(Ws[:T] / gf[:T] / W0, 1.0)79 sc[:T] = w80 sc[T:] = w[-1]81 return sc828384def solve_ha_equilibrium(p, m: TwoClassModel, u, P, lam=1.0, tol=1e-4, maxit=60, verbose=True, hb_kw=None,85 masses=(0.5, 0.3, 0.2), betas_low=(0.93, 0.97), r_bounds=(-0.02, 0.8), jac_every=6,86 step_max=0.2, checkpoint=None, r_init=None):87 """Quasi-Newton on the return path: the Jacobian of the household block's log asset demand with respect to88 the return path is computed by finite differences (T household solves, cheap); the supply side's response is89 not differentiated (it has the opposite sign, so using the demand Jacobian alone over-estimates the required90 move, which the damping lam corrects). Steps are bounded by r_bounds."""91 hb_kw = hb_kw or {}92 T = m.T93 wL, wH, transfer, Pidx, r, gf = ha_inputs_from_paths(P, p)94 W0 = (1.0 - p.theta_F) * P["W"][0]95 hb, pols0, dists0, hcal = calibrate_households(p, W0, wL[0], wH[0], transfer[0], Pidx[0], masses, betas_low, verbose, **hb_kw)96 U6 = u.reshape(T, 6)97 u4 = U6[:, :5].ravel().copy() # fixed-r mode: 4 log stocks + log X_K98 r_path = r.copy() if r_init is None else np.asarray(r_init, float).copy()99 hist = []100 t0 = time.time()101 J = None102 for it in range(maxit):103 m.set_fixed_r(r_path)104 u4, ok, h = m.solve(u4, tol=1e-9, verbose=False, maxit=30)105 Pp = m.evaluate(u4, full=True)106 wL, wH, transfer, Pidx, _, gf = ha_inputs_from_paths(Pp, p)107 Ws = (1.0 - p.theta_F) * Pp["W"][:T] # households hold the claims not held by the citizen fund108 scales = grid_scales((1.0 - p.theta_F) * Pp["W"], gf, W0, T)109 x_term = hb.solve_stationary(r_path[T], (wL[T], wH[T]), transfer[T], Pidx[T], scale=scales[T])110 Wd, X, pols, dist_hist = _ha_demand(hb, r_path, wL, wH, transfer, Pidx, gf, dists0, T, x_term, scales)111 gap = np.log(Wd / Ws)112 err = float(np.max(np.abs(gap)))113 hist.append(err)114 if verbose:115 print(f" HA-GE it {it:2d}: max|log Wd/Ws| = {err:.3e}, production converged={ok}, ({time.time() - t0:.0f}s)", flush=True)116 if err < tol:117 break118 if J is None or it % jac_every == 0:119 # Jacobian of log Wd (T) w.r.t. r_{t+1} (t = 0..T-1) at the current prices120 h_fd = 1e-4121 J = np.zeros((T, T))122 for s in range(T):123 rp = r_path.copy()124 rp[s + 1] += h_fd125 Wd_p, _, _, _ = _ha_demand(hb, rp, wL, wH, transfer, Pidx, gf, dists0, T, x_term, scales)126 J[:, s] = (np.log(Wd_p) - np.log(Wd)) / h_fd127 if verbose:128 print(f" HA-GE: household Jacobian computed ({time.time() - t0:.0f}s)", flush=True)129 # supply responds with the opposite sign; approximate its elasticity by a diagonal term130 Jtot = J + np.diag(np.full(T, 5.0)) # d log Ws / d r ~ -duration (~5 years): crude, corrected by damping131 try:132 step = np.linalg.solve(Jtot, gap)133 except np.linalg.LinAlgError:134 step = gap / 20.0135 step = np.clip(lam * step, -step_max, step_max)136 r_new = np.clip(r_path[1:T + 1] - step, r_bounds[0], r_bounds[1])137 r_path[1:T + 1] = r_new138 r_path[T + 1] = r_path[T]139 if checkpoint:140 np.save(checkpoint, r_path)141 m.set_free_r()142 stats = []143 for t in range(T):144 st = hb.wealth_stats(dist_hist[t], {k: pols[k][t] for k in pols})145 st["year"] = int(P["years"][t])146 stats.append(st)147 out = dict(r=r_path, Wd=Wd, Ws=Ws, gap=gap, hist=hist, stats=stats, hcal=hcal, X=X, paths=Pp, u4=u4,148 converged=err < tol)149 return out, hb, pols, dist_hist150