1"""Model B — heterogeneous-agent stationary recursive competitive equilibrium with AI delegation and compute ownership.23Households (continuum): CRRA utility, idiosyncratic earnings e (Rouwenhorst), borrowing limit a_min, cash-on-hand state4x = (R + nu) a + w e, where nu is the household's idiosyncratic return term reflecting *decision quality*:5 self-managed (d=0): nu ~ N(-m_H, sigma_rH^2) (mean cost of mistakes + dispersion)6 delegated (d=1): nu ~ N(0, sigma_rA^2) at fee phi0 + phi1 * a' (goods, paid when saving)7Mean mistakes are a resource cost; dispersion is zero-sum redistribution. Delegation d is chosen each period together8with savings a' (value-function iteration with a discrete choice, numba kernels).9Firms: task-based technology of Model A (age.economy.tasks) with compute stock K_Q. Compute supply: K_Q = KQ0 (pQ/uQ)^eps_Q10(eps_Q = 0: fixed stock). Ownership of compute rents: broad (added to every household's return) or concentrated (paid11to the wealthiest mass n_K as a proportional-to-wealth bonus return).12Equilibrium: r such that household savings A(r) = K(r) + pK_Q K_Q (capital market), compute market (A.4), L from the13earnings distribution. Solved by bisection on r with warm starts.14"""15from __future__ import annotations1617from dataclasses import dataclass, field1819import numpy as np20from numba import njit21from scipy.optimize import brentq2223from age.economy.tasks import TaskTech242526# --------------------------------------------------------------------------- parameters27@dataclass28class HAParams:29 beta: float = 0.9630 gamma: float = 2.031 alpha: float = 0.3332 delta: float = 0.0633 rho_e: float = 0.9634 sigma_e: float = 0.2035 n_e: int = 736 a_min: float = 0.037 a_max: float = 400.038 n_a: int = 20039 # decision-quality returns40 m_H: float = 0.01 # mean cost of self-managed mistakes (per unit of wealth per year)41 sigma_rH: float = 0.04 # dispersion of self-managed idiosyncratic return42 sigma_rA: float = 0.01 # dispersion under delegation43 phi0: float = 0.02 # fixed fee (goods per year)44 phi1: float = 0.002 # proportional fee (per unit saved)45 n_gh: int = 5 # Gauss-Hermite nodes for the return term46 delegation_on: bool = True47 # technology / compute48 tasks: TaskTech = field(default_factory=lambda: TaskTech(alpha=0.33, iota_bar=0.5, psi0=1.0, psi1=3.0))49 KQ0: float = 0.0 # compute stock at pQ = uQ50 eps_Q: float = 0.0 # compute supply elasticity51 pK_Q: float = 1.052 delta_Q: float = 0.2053 qbar: float = 0.002 # compute per delegated household-year54 ownership: str = 'broad' # 'broad' | 'concentrated'55 n_K: float = 0.10 # mass of owners if concentrated (top of the wealth distribution)56 # numerics57 tol_v: float = 1e-658 max_iter: int = 200059 howard: int = 20606162# --------------------------------------------------------------------------- discretisations63def rouwenhorst(n, rho, sigma):64 """Rouwenhorst (1995) discretisation of an AR(1) with unconditional s.d. sigma/sqrt(1-rho^2)."""65 p = (1 + rho) / 266 P = np.array([[p, 1 - p], [1 - p, p]])67 for i in range(2, n):68 Pn = np.zeros((i + 1, i + 1))69 Pn[:i, :i] += p * P70 Pn[:i, 1:] += (1 - p) * P71 Pn[1:, :i] += (1 - p) * P72 Pn[1:, 1:] += p * P73 Pn[1:i] /= 274 P = Pn75 sd = sigma / np.sqrt(1 - rho ** 2)76 grid = np.linspace(-sd * np.sqrt(n - 1), sd * np.sqrt(n - 1), n)77 return grid, P787980def gauss_hermite(n):81 x, w = np.polynomial.hermite_e.hermegauss(n)82 return x, w / w.sum()838485def make_grid(a_min, a_max, n):86 g = np.linspace(0, 1, n) ** 387 return a_min + (a_max - a_min) * g888990# --------------------------------------------------------------------------- numba kernels91@njit(cache=True)92def _interp(xg, yg, x):93 n = xg.shape[0]94 if x <= xg[0]:95 return yg[0] + (yg[1] - yg[0]) * (x - xg[0]) / (xg[1] - xg[0])96 if x >= xg[n - 1]:97 return yg[n - 1] + (yg[n - 1] - yg[n - 2]) * (x - xg[n - 1]) / (xg[n - 1] - xg[n - 2])98 lo, hi = 0, n - 199 while hi - lo > 1:100 mid = (lo + hi) // 2101 if xg[mid] <= x:102 lo = mid103 else:104 hi = mid105 t = (x - xg[lo]) / (xg[hi] - xg[lo])106 return yg[lo] + t * (yg[hi] - yg[lo])107108109@njit(cache=True)110def _expected_value(V, xgrid, egrid, P, agrid, w, R, nu_nodes, nu_w, ie):111 """EV[a'] for current e-index ie: sum over e', nu of V(x' = (R+nu) a' + w e', e')."""112 na = agrid.shape[0]113 ne = egrid.shape[0]114 EV = np.zeros(na)115 for ia in range(na):116 s = 0.0117 for je in range(ne):118 pe = P[ie, je]119 if pe == 0.0:120 continue121 for k in range(nu_nodes.shape[0]):122 xp = (R + nu_nodes[k]) * agrid[ia] + w * egrid[je]123 s += pe * nu_w[k] * _interp(xgrid, V[:, je], xp)124 EV[ia] = s125 return EV126127128@njit(cache=True)129def _bellman(V, xgrid, egrid, P, agrid, w, R, gamma, beta, nuH, wH, nuA, wA, phi0, phi1, deleg_on):130 nx = xgrid.shape[0]131 ne = egrid.shape[0]132 na = agrid.shape[0]133 Vn = np.empty((nx, ne))134 pol_a = np.empty((nx, ne))135 pol_d = np.zeros((nx, ne), dtype=np.int64)136 pol_c = np.empty((nx, ne))137 for ie in range(ne):138 EV0 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuH, wH, ie)139 if deleg_on:140 EV1 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuA, wA, ie)141 else:142 EV1 = EV0143 for ix in range(nx):144 x = xgrid[ix]145 best = -1e300146 ba = agrid[0]147 bd = 0148 bc = 0.0149 # d = 0150 for ia in range(na):151 c = x - agrid[ia]152 if c <= 1e-12:153 break154 u = (c ** (1.0 - gamma) - 1.0) / (1.0 - gamma) if gamma != 1.0 else np.log(c)155 v = u + beta * EV0[ia]156 if v > best:157 best = v158 ba = agrid[ia]159 bd = 0160 bc = c161 if deleg_on:162 for ia in range(na):163 c = x - agrid[ia] - phi0 - phi1 * agrid[ia]164 if c <= 1e-12:165 break166 u = (c ** (1.0 - gamma) - 1.0) / (1.0 - gamma) if gamma != 1.0 else np.log(c)167 v = u + beta * EV1[ia]168 if v > best:169 best = v170 ba = agrid[ia]171 bd = 1172 bc = c173 Vn[ix, ie] = best174 pol_a[ix, ie] = ba175 pol_d[ix, ie] = bd176 pol_c[ix, ie] = bc177 return Vn, pol_a, pol_d, pol_c178179180@njit(cache=True)181def _howard(V, xgrid, egrid, P, agrid, w, R, gamma, beta, nuH, wH, nuA, wA, pol_a, pol_d, pol_c, n_steps):182 nx = xgrid.shape[0]183 ne = egrid.shape[0]184 for _ in range(n_steps):185 Vn = np.empty((nx, ne))186 for ie in range(ne):187 EV0 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuH, wH, ie)188 EV1 = _expected_value(V, xgrid, egrid, P, agrid, w, R, nuA, wA, ie)189 for ix in range(nx):190 c = pol_c[ix, ie]191 u = (c ** (1.0 - gamma) - 1.0) / (1.0 - gamma) if gamma != 1.0 else np.log(c)192 ev = _interp(agrid, EV1 if pol_d[ix, ie] == 1 else EV0, pol_a[ix, ie])193 Vn[ix, ie] = u + beta * ev194 V = Vn195 return V196197198@njit(cache=True)199def _stationary(xgrid, egrid, P, pol_a, pol_d, w, R, nuH, wH, nuA, wA, n_iter, tol):200 """Non-stochastic simulation (Young 2010) of the distribution over (x, e)."""201 nx = xgrid.shape[0]202 ne = egrid.shape[0]203 lam = np.ones((nx, ne)) / (nx * ne)204 for it in range(n_iter):205 new = np.zeros((nx, ne))206 for ie in range(ne):207 for ix in range(nx):208 m = lam[ix, ie]209 if m == 0.0:210 continue211 ap = pol_a[ix, ie]212 if pol_d[ix, ie] == 1:213 nodes = nuA214 wts = wA215 else:216 nodes = nuH217 wts = wH218 for je in range(ne):219 pe = P[ie, je]220 if pe == 0.0:221 continue222 for k in range(nodes.shape[0]):223 xp = (R + nodes[k]) * ap + w * egrid[je]224 # lottery onto grid225 if xp <= xgrid[0]:226 new[0, je] += m * pe * wts[k]227 elif xp >= xgrid[nx - 1]:228 new[nx - 1, je] += m * pe * wts[k]229 else:230 lo, hi = 0, nx - 1231 while hi - lo > 1:232 mid = (lo + hi) // 2233 if xgrid[mid] <= xp:234 lo = mid235 else:236 hi = mid237 t = (xp - xgrid[lo]) / (xgrid[hi] - xgrid[lo])238 new[lo, je] += m * pe * wts[k] * (1.0 - t)239 new[hi, je] += m * pe * wts[k] * t240 diff = np.max(np.abs(new - lam))241 lam = new242 if diff < tol:243 break244 return lam245246247# --------------------------------------------------------------------------- household block248class HouseholdBlock:249 def __init__(self, p: HAParams):250 self.p = p251 self.egrid_log, self.P = rouwenhorst(p.n_e, p.rho_e, p.sigma_e)252 self.egrid = np.exp(self.egrid_log)253 # normalise mean labour efficiency to 1254 pi = self.stationary_e()255 self.egrid = self.egrid / (pi @ self.egrid)256 self.agrid = make_grid(p.a_min, p.a_max, p.n_a)257 self.xgrid = make_grid(p.a_min + 1e-3, p.a_max * 1.2, p.n_a)258 gh_x, gh_w = gauss_hermite(p.n_gh)259 self.nuH = -p.m_H + p.sigma_rH * gh_x260 self.wH = gh_w261 self.nuA = p.sigma_rA * gh_x262 self.wA = gh_w263 self.V = None264265 def stationary_e(self):266 w, v = np.linalg.eig(self.P.T)267 pi = np.real(v[:, np.argmin(np.abs(w - 1))])268 return pi / pi.sum()269270 def solve(self, R: float, w: float, bonus: np.ndarray | None = None):271 """Value function iteration with Howard steps. `bonus` (per x-grid point) adds a wealth-proportional return272 bonus for concentrated compute ownership: R_eff(x) = R + bonus[x]; implemented by shifting nu nodes per x is273 expensive, so bonus is applied as an extra income r_bonus*a' through the return R for the top group only274 via a separate solve (see Economy)."""275 p = self.p276 if self.V is None:277 c0 = np.maximum(self.xgrid[:, None] * 0.05, 1e-3) * np.ones((1, p.n_e))278 self.V = ((c0 ** (1 - p.gamma) - 1) / (1 - p.gamma)) / (1 - p.beta)279 V = self.V280 for it in range(p.max_iter):281 Vn, pol_a, pol_d, pol_c = _bellman(V, self.xgrid, self.egrid, self.P, self.agrid, w, R, p.gamma, p.beta,282 self.nuH, self.wH, self.nuA, self.wA, p.phi0, p.phi1, p.delegation_on)283 diff = np.max(np.abs(Vn - V))284 V = _howard(Vn, self.xgrid, self.egrid, self.P, self.agrid, w, R, p.gamma, p.beta,285 self.nuH, self.wH, self.nuA, self.wA, pol_a, pol_d, pol_c, p.howard)286 if diff < p.tol_v:287 break288 self.V = V289 self.pol_a, self.pol_d, self.pol_c = pol_a, pol_d, pol_c290 self.n_iter = it + 1291 return V292293 def distribution(self, R, w):294 self.lam = _stationary(self.xgrid, self.egrid, self.P, self.pol_a, self.pol_d, w, R,295 self.nuH, self.wH, self.nuA, self.wA, 20000, 1e-11)296 return self.lam297298 def aggregates(self):299 lam = self.lam300 A = float(np.sum(lam * self.pol_a))301 D = float(np.sum(lam * self.pol_d))302 C = float(np.sum(lam * self.pol_c))303 L = float(np.sum(lam.sum(axis=0) * self.egrid))304 fees = float(np.sum(lam * self.pol_d * (self.p.phi0 + self.p.phi1 * self.pol_a)))305 mistakes = float(np.sum(lam * (1 - self.pol_d) * self.p.m_H * self.pol_a))306 return dict(A=A, D=D, C=C, L=L, fees=fees, mistakes=mistakes)307308 # ----- distributional statistics on end-of-period wealth a'309 def wealth_stats(self, extra_mass: float = 0.0, extra_wealth: float = 0.0):310 lam = self.lam311 a = self.pol_a.ravel()312 m = lam.ravel()313 d_all = self.pol_d.ravel()314 if extra_mass > 0:315 a = np.append(a, extra_wealth)316 m = np.append(m * (1 - extra_mass), extra_mass)317 d_all = np.append(d_all, 0)318 order = np.argsort(a)319 a, m = a[order], m[order]320 cm = np.cumsum(m)321 cw = np.cumsum(a * m)322 W = cw[-1]323 lorenz_x = cm / cm[-1]324 lorenz_y = cw / W325 gini = 1 - 2 * np.trapezoid(lorenz_y, lorenz_x)326 lx = np.concatenate([[0.0], lorenz_x])327 ly = np.concatenate([[0.0], lorenz_y])328 def top(share):329 return 1 - float(np.interp(1 - share, lx, ly))330 # delegation by wealth decile, households only (the owner class is outside the household block)331 a_h = self.pol_a.ravel(); m_h = lam.ravel(); d_h = self.pol_d.ravel()332 oh = np.argsort(a_h); cm_h = np.cumsum(m_h[oh]) / m_h.sum()333 dec = np.minimum((cm_h * 10).astype(int), 9)334 deleg_by_decile = np.array([np.sum(m_h[oh][dec == k] * d_h[oh][dec == k]) / max(np.sum(m_h[oh][dec == k]), 1e-12) for k in range(10)])335 return dict(gini=float(gini), top1=float(top(0.01)), top10=float(top(0.10)), bottom50=float(np.interp(0.5, lx, ly)),336 lorenz_x=lorenz_x, lorenz_y=lorenz_y, deleg_by_decile=deleg_by_decile, mean_wealth=float(W))337338339# --------------------------------------------------------------------------- general equilibrium340class EconomyB:341 def __init__(self, p: HAParams):342 self.p = p343 self.hh = HouseholdBlock(p)344 self.uQ = (1 / p.beta - 1 + p.delta_Q) * p.pK_Q # user cost of compute at the discount rate345346 def compute_supply(self, pQ):347 p = self.p348 if p.KQ0 <= 0:349 return 0.0350 if p.eps_Q == 0:351 return p.KQ0352 return p.KQ0 * max(pQ / self.uQ, 1e-9) ** p.eps_Q353354 def firm_side(self, r, L, D):355 """Given r, L and delegation share, clear the compute market and return prices and quantities."""356 p = self.p357 RK = r + p.delta358 def solve_for_KQ(KQ):359 QP = max(KQ - p.qbar * D, 1e-12) if KQ > 0 else 0.0360 T = p.tasks.bundle(QP, L)361 K = T * (p.alpha / RK) ** (1 / (1 - p.alpha))362 pr = p.tasks.prices(0.0, K, QP, L)363 return K, QP, pr364 if p.KQ0 <= 0:365 K, QP, pr = solve_for_KQ(0.0)366 KQ = 0.0367 elif p.eps_Q == 0:368 KQ = p.KQ0369 K, QP, pr = solve_for_KQ(KQ)370 else:371 f = lambda KQ: self.compute_supply(solve_for_KQ(KQ)[2]['pQ']) - KQ372 KQ = brentq(f, p.qbar * D + 1e-6, 50.0)373 K, QP, pr = solve_for_KQ(KQ)374 return dict(K=K, KQ=KQ, QP=QP, **pr)375376 def excess_assets(self, r, D_guess=0.0, L=1.0):377 """Household savings minus K + pK_Q*KQ at interest rate r (inner fixed point on D for the compute market)."""378 p = self.p379 D = D_guess380 for _ in range(5):381 fs = self.firm_side(r, L, D)382 pQ, KQ, K = fs['pQ'], fs['KQ'], fs['K']383 rents = (pQ - self.uQ) * KQ if KQ > 0 else 0.0 # scarcity rents over user cost384 assets_total = K + p.pK_Q * KQ385 # fund return: capital earns r (= RK - delta); compute earns uQ (user cost) + rents386 if p.ownership == 'broad' or rents <= 0:387 R = 1 + r + rents / max(assets_total, 1e-9)388 else:389 R = 1 + r390 self.hh.solve(R, fs['w'])391 self.hh.distribution(R, fs['w'])392 agg = self.hh.aggregates()393 if abs(agg['D'] - D) < 1e-3:394 D = agg['D']395 break396 D = agg['D']397 self.last = dict(r=r, R=R, D=D, rents=rents, agg=agg, fs=fs, assets_total=assets_total)398 return agg['A'] - assets_total399400 def solve(self, r_lo=-0.02, r_hi=None, tol=1e-4, verbose=False):401 p = self.p402 L = 1.0 # efficiency units normalised403 r_hi = r_hi if r_hi is not None else 1 / p.beta - 1 - 2e-4404 f = lambda r: self.excess_assets(r, L=L)405 # coarse scan for a sign change (excess assets is increasing in r), then Brent406 grid = np.linspace(r_lo, r_hi, 8)407 vals = [f(r) for r in grid]408 idx = next((i for i in range(len(grid) - 1) if vals[i] < 0 <= vals[i + 1]), None)409 if idx is None:410 raise RuntimeError(f'no sign change in excess assets on [{r_lo}, {r_hi}]: {np.round(vals, 3)}')411 r = brentq(f, grid[idx], grid[idx + 1], xtol=tol)412 f(r)413 last = self.last414 agg, fs = last['agg'], last['fs']415 ws = self.hh.wealth_stats()416 if p.ownership == 'concentrated' and last['rents'] > 0:417 # owner class (mass n_K) holds the right to the scarcity rents, valued as a perpetuity at r; its members are418 # appended to the wealth distribution as a point mass (two-class economy of Proposition 6)419 ws = self.hh.wealth_stats(extra_mass=p.n_K, extra_wealth=last['rents'] / p.n_K / max(r, 1e-4))420 Y = fs['Y']421 shares = dict(labor=fs['w'] * L / Y, capital=fs['RK'] * fs['K'] / Y, compute=fs['pQ'] * fs['QP'] / Y)422 return dict(r=r, R=last['R'], w=fs['w'], K=fs['K'], KQ=fs['KQ'], QP=fs['QP'], pQ=fs['pQ'], iota=fs['iota'],423 Y=Y, C=agg['C'], D=agg['D'], fees=agg['fees'], mistakes=agg['mistakes'], rents=last['rents'],424 shares=shares, gini=ws['gini'], top1=ws['top1'], top10=ws['top10'], bottom50=ws['bottom50'],425 deleg_by_decile=ws['deleg_by_decile'], lorenz=(ws['lorenz_x'], ws['lorenz_y']),426 mean_wealth=ws['mean_wealth'], n_iter_v=self.hh.n_iter)427