1"""Heterogeneous-agent household block (Bewley–Huggett–Aiyagari with skill and discount-factor types).23State: skill s in {L, H} (permanent), discount type j (permanent), productivity e (Markov, Rouwenhorst),4assets a >= 0. Each period the household receives labor income w_{s,t} gf_t e + transfer_t, chooses5expenditure x (goods + housing, Cobb-Douglas inside the period, so the period utility is6(x/P_t)^{1-gamma}/(1-gamma) with P_t the housing-inclusive price index) and end-of-period savings a'7which earn r_{t+1}. Budget: x + a' = a + y. Euler: x^{-gamma} P^{gamma-1} = beta (1+r') E[x'^{-gamma} P'^{gamma-1}].89All quantities are detrended by gf_t = (1+g_Z)^t (hatted variables) so that the grid is stationary in the10frozen-technology reference. Solution by the endogenous grid method (Carroll 2006); distribution by11non-stochastic simulation on the asset grid (Young 2010). numba-accelerated.12"""13from __future__ import annotations1415import numpy as np16from numba import njit171819# ------------------------------------------------------------------ income process20def rouwenhorst(n, rho, sigma):21 """Rouwenhorst discretization of an AR(1) in logs with persistence rho and innovation s.d. sigma."""22 q = (1 + rho) / 223 P = np.array([[q, 1 - q], [1 - q, q]])24 for i in range(2, n):25 Pn = np.zeros((i + 1, i + 1))26 Pn[:i, :i] += q * P27 Pn[:i, 1:] += (1 - q) * P28 Pn[1:, :i] += (1 - q) * P29 Pn[1:, 1:] += q * P30 Pn[1:-1] /= 231 P = Pn32 psi = sigma * np.sqrt((n - 1) / (1 - rho ** 2))33 z = np.linspace(-psi, psi, n)34 pi = np.linalg.matrix_power(P.T, 2000)[:, 0]35 e = np.exp(z)36 e = e / np.sum(pi * e) # mean one37 return e, P, pi383940def asset_grid(a_max, n=300, curv=4.0):41 u = np.linspace(0, 1, n)42 return a_max * u ** curv434445# ------------------------------------------------------------------ EGM step (numba)46@njit(cache=True)47def _egm_step(agrid, e, Pe, y_next_e, rnext, beta_eff, gamma, Pnow, Pnext, gnext, c_next_pol, a_grid_next):48 """One backward EGM step for a given (skill, type). Inputs are detrended:49 c_next_pol[e_idx, a_idx] = next-period expenditure policy on a_grid_next (assets at start of next period,50 detrended by gf_{t+1}); returns this period's policy x(a) on agrid and savings a'(a).51 beta_eff = beta * (1+g)^{-gamma} handles detrending; gnext = 1+g."""52 ne, na = c_next_pol.shape53 x_pol = np.empty((ne, na))54 ap_pol = np.empty((ne, na))55 # marginal utility next period at each (e', a'): x'^{-gamma} P'^{gamma-1}56 for ie in range(ne):57 # for each savings level a' on agrid (detrended by gf_t), next-period assets = (1+r') a' / gnext58 rhs = np.empty(na)59 for ia in range(na):60 an = (1.0 + rnext) * agrid[ia] / gnext61 s = 0.062 for je in range(ne):63 # interpolate next-period policy at an64 # find bracket65 if an <= a_grid_next[0]:66 xn = c_next_pol[je, 0]67 elif an >= a_grid_next[na - 1]:68 # extrapolate linearly69 xn = c_next_pol[je, na - 1] + (c_next_pol[je, na - 1] - c_next_pol[je, na - 2]) / (a_grid_next[na - 1] - a_grid_next[na - 2]) * (an - a_grid_next[na - 1])70 else:71 lo, hi = 0, na - 172 while hi - lo > 1:73 mid = (lo + hi) // 274 if a_grid_next[mid] <= an:75 lo = mid76 else:77 hi = mid78 wgt = (an - a_grid_next[lo]) / (a_grid_next[hi] - a_grid_next[lo])79 xn = (1 - wgt) * c_next_pol[je, lo] + wgt * c_next_pol[je, hi]80 if xn < 1e-12:81 xn = 1e-1282 s += Pe[ie, je] * xn ** (-gamma) * Pnext ** (gamma - 1.0)83 rhs[ia] = beta_eff * (1.0 + rnext) * s84 # endogenous grid: x = (rhs / P^{gamma-1})^{-1/gamma}; cash on hand m = x + a'85 m_end = np.empty(na)86 x_end = np.empty(na)87 for ia in range(na):88 x_end[ia] = (rhs[ia] / Pnow ** (gamma - 1.0)) ** (-1.0 / gamma)89 m_end[ia] = x_end[ia] + agrid[ia]90 # map back to the exogenous asset grid: m = a + y_e91 for ia in range(na):92 m = agrid[ia] + y_next_e[ie]93 if m <= m_end[0]:94 # borrowing constraint binds: a' = 0, x = m95 x_pol[ie, ia] = m96 ap_pol[ie, ia] = 0.097 else:98 lo, hi = 0, na - 199 if m >= m_end[na - 1]:100 slope = (x_end[na - 1] - x_end[na - 2]) / (m_end[na - 1] - m_end[na - 2])101 xx = x_end[na - 1] + slope * (m - m_end[na - 1])102 else:103 while hi - lo > 1:104 mid = (lo + hi) // 2105 if m_end[mid] <= m:106 lo = mid107 else:108 hi = mid109 wgt = (m - m_end[lo]) / (m_end[hi] - m_end[lo])110 xx = (1 - wgt) * x_end[lo] + wgt * x_end[hi]111 if xx > m:112 xx = m113 if xx < 1e-12:114 xx = 1e-12115 x_pol[ie, ia] = xx116 ap_pol[ie, ia] = m - xx117 return x_pol, ap_pol118119120@njit(cache=True)121def _forward(dist, ap_pol, agrid, Pe, rnext, gnext):122 """Push the distribution over (e, a) forward one period. dist[e, a] at start of t -> at start of t+1.123 Savings a' earn r' and are deflated by gnext before being placed on the grid."""124 ne, na = dist.shape125 new = np.zeros((ne, na))126 for ie in range(ne):127 for ia in range(na):128 mass = dist[ie, ia]129 if mass == 0.0:130 continue131 an = (1.0 + rnext) * ap_pol[ie, ia] / gnext132 # lottery on the grid133 if an <= agrid[0]:134 lo, w = 0, 1.0135 elif an >= agrid[na - 1]:136 lo, w = na - 2, 0.0137 else:138 lo, hi = 0, na - 1139 while hi - lo > 1:140 mid = (lo + hi) // 2141 if agrid[mid] <= an:142 lo = mid143 else:144 hi = mid145 w = 1.0 - (an - agrid[lo]) / (agrid[lo + 1] - agrid[lo])146 for je in range(ne):147 pm = mass * Pe[ie, je]148 new[je, lo] += pm * w149 new[je, lo + 1] += pm * (1.0 - w)150 return new151152153# ------------------------------------------------------------------ household block154class HouseholdBlock:155 def __init__(self, p, betas, beta_masses, rho_e=0.95, sigma_e=0.20, ne=5, na=300, a_max=400.0):156 self.p = p157 self.betas = np.asarray(betas, float) # discount types158 self.beta_masses = np.asarray(beta_masses, float) # masses within each skill (sum to one)159 self.e, self.Pe, self.pi_e = rouwenhorst(ne, rho_e, sigma_e)160 self.agrid = asset_grid(a_max, na)161 self.ne, self.na = ne, na162 self.H = np.asarray(p.H, float)163 self.gamma = p.gamma_u164 self.g = p.g_Z165166 # ----- stationary policy for constant prices (reference BGP)167 def solve_stationary(self, r, w_by_skill, transfer, P, tol=1e-10, maxit=5000):168 pols = {}169 for s in range(2):170 y = w_by_skill[s] * self.e + transfer171 for j, b in enumerate(self.betas):172 beta_eff = b * (1 + self.g) ** (-self.gamma)173 x = np.tile(y[:, None] + 0.05 * self.agrid[None, :], (1, 1)) # initial policy: consume income + 5% assets174 for it in range(maxit):175 x_new, ap = _egm_step(self.agrid, self.e, self.Pe, y, r, beta_eff, self.gamma, P, P, 1 + self.g, x, self.agrid)176 err = np.max(np.abs(x_new - x))177 x = x_new178 if err < tol:179 break180 pols[(s, j)] = (x, ap)181 return pols182183 def stationary_distribution(self, pols, r, tol=1e-12, maxit=20000):184 dists = {}185 for key, (x, ap) in pols.items():186 d = np.zeros((self.ne, self.na))187 d[:, 0] = self.pi_e188 for it in range(maxit):189 dn = _forward(d, ap, self.agrid, self.Pe, r, 1 + self.g)190 err = np.max(np.abs(dn - d))191 d = dn192 if err < tol:193 break194 dists[key] = d195 return dists196197 def aggregate(self, pols, dists):198 """Aggregate end-of-period savings and expenditure (detrended), weighting types by masses."""199 A, X = 0.0, 0.0200 for s in range(2):201 for j in range(len(self.betas)):202 m = self.H[s] * self.beta_masses[j]203 x, ap = pols[(s, j)]204 d = dists[(s, j)]205 A += m * np.sum(d * ap)206 X += m * np.sum(d * x)207 return A, X208209 # ----- transition given price paths (arrays of length T+1 for t = 0..T; r[t+1] is the return on a'_t)210 def solve_transition(self, r, wL, wH, transfer, P, x_terminal):211 """Backward induction from terminal policies x_terminal[(s,j)] (stationary at T); returns policies for212 t = 0..T-1. Inputs are detrended wages per efficiency unit (so income = w * e); r[t] indexes the return213 between t-1 and t."""214 T = len(r) - 1215 pols = {key: [None] * T for key in x_terminal}216 for s in range(2):217 w_path = wL if s == 0 else wH218 for j, b in enumerate(self.betas):219 beta_eff = b * (1 + self.g) ** (-self.gamma)220 x_next = x_terminal[(s, j)][0]221 for t in range(T - 1, -1, -1):222 y = w_path[t] * self.e + transfer[t]223 x, ap = _egm_step(self.agrid, self.e, self.Pe, y, r[t + 1], beta_eff, self.gamma, P[t], P[t + 1], 1 + self.g, x_next, self.agrid)224 pols[(s, j)][t] = (x, ap)225 x_next = x226 return pols227228 def simulate_transition(self, pols, dists0, r):229 """Forward simulation; returns aggregate savings A_t (detrended, t=0..T-1) and the distributions."""230 T = len(next(iter(pols.values())))231 A = np.zeros(T)232 X = np.zeros(T)233 dists = {k: d.copy() for k, d in dists0.items()}234 hist = []235 for t in range(T):236 snap = {}237 for s in range(2):238 for j in range(len(self.betas)):239 key = (s, j)240 m = self.H[s] * self.beta_masses[j]241 x, ap = pols[key][t]242 d = dists[key]243 A[t] += m * np.sum(d * ap)244 X[t] += m * np.sum(d * x)245 snap[key] = d246 dists[key] = _forward(d, ap, self.agrid, self.Pe, r[t + 1], 1 + self.g)247 hist.append(snap)248 return A, X, hist249250 # ----- distribution statistics251 def wealth_stats(self, dists, pols=None):252 """Gini and top shares of end-of-period savings (or start-of-period assets if pols is None)."""253 vals, wts = [], []254 for s in range(2):255 for j in range(len(self.betas)):256 key = (s, j)257 m = self.H[s] * self.beta_masses[j]258 d = dists[key]259 a = pols[key][1] if pols is not None else np.tile(self.agrid, (self.ne, 1))260 vals.append(a.ravel())261 wts.append((m * d).ravel())262 v = np.concatenate(vals)263 w = np.concatenate(wts)264 w = w / w.sum()265 o = np.argsort(v)266 v, w = v[o], w[o]267 cw = np.cumsum(w)268 cv = np.cumsum(v * w)269 tot = cv[-1]270 gini = 1 - 2 * np.sum(w * (cv - 0.5 * v * w)) / tot if tot > 0 else np.nan271 def top(share):272 idx = np.searchsorted(cw, 1 - share)273 return 1 - (cv[idx - 1] / tot if idx > 0 else 0.0)274 return dict(gini=float(gini), top1=float(top(0.01)), top10=float(top(0.10)), bottom50=float(cv[np.searchsorted(cw, 0.5)] / tot), mean=float(tot))275