1"""Mapping the task-based stochastic growth economy to the LQ delegation game (§1.9 of the specification).23Steps4-----51. Deterministic steady state for given compute stock K_Q and delegation share D.62. Linearised household problem solved in sequence space (horizon T): consumption and wealth paths for given price7 paths; with beta*R = 1 wealth has a unit root, so the terminal condition is a_{T+1} = a_T (new steady wealth level).83. kappa0 = response of log savings a_1 to a known news shock theta = eps_1 with the aggregate capital path fixed;9 kappa_full = symmetric rational-expectations response (aggregate = individual); alpha = 1 - kappa0/kappa_full.104. |U_xx| in log units = abar^2 * [ -u''(c) - beta u''(c) R MPC ] with the MPC from the same linear solve.115. Delegation share from the threshold rule (A.10) and compute-market clearing (A.4) in the short run.12"""13from __future__ import annotations1415from dataclasses import dataclass, field1617import numpy as np18from scipy.optimize import brentq19from scipy.stats import lognorm2021from age.economy.tasks import TaskTech22from age.equilibrium.lq_game import LQParams, private_value232425@dataclass(frozen=True)26class EconParams:27 beta: float = 0.9628 gamma: float = 2.029 alpha: float = 0.3330 delta: float = 0.0631 rho_z: float = 0.9032 sigma_eps: float = 0.0233 tasks: TaskTech = field(default_factory=lambda: TaskTech(alpha=0.33, iota_bar=0.5, psi0=1.0, psi1=3.0))34 # compute / delegation technology35 delta_Q: float = 0.2036 pK_Q: float = 1.037 qbar: float = 0.00238 ebar: float = 0.00139 pE: float = 1.040 tau_s: float = 0.0141 # information about the TFP innovation (savings mapping; precisions in units of log TFP)42 tau_H: float = 500.043 tau_A: float = 5000.044 M: int = 145 # information about the predictable return component (portfolio mapping; s.d. in return units, ILLUSTRATIVE)46 sigma_R: float = 0.16 # total risky-return s.d. (calibrated: Shiller 1984-2025 real equity s.d. 0.157)47 sig_theta_R: float = 0.06 # s.d. of the predictable return component theta48 sig_H_R: float = 0.10 # human forecast noise s.d.49 sig_A_R: float = 0.04 # foundation-model forecast noise s.d.50 # delegation disutility, lognormal in goods units (fraction of steady-state consumption per year)51 kappa_mu: float = -6.052 kappa_sigma: float = 1.05354 @property55 def tau_theta(self):56 return 1.0 / self.sigma_eps ** 25758 def u2(self, c): # u''(c)59 return -self.gamma * c ** (-self.gamma - 1.0)6061 def u1(self, c): # u'(c)62 return c ** (-self.gamma)636465# --------------------------------------------------------------------------- 1. steady state66def steady_state(p: EconParams, KQ: float = 0.0, D: float = 0.0, L_endowment: float = 1.0) -> dict:67 """Deterministic steady state (z = 0) for given compute stock and delegation share.6869 Labour is inelastic; supervision time tau_s*D is subtracted from hours. Capital from R = 1/beta.70 The compute price in the steady state is the production compute rental; the compute-capital price is held at pK_Q.71 """72 L = L_endowment - p.tau_s * D73 QP = max(KQ - p.qbar * D, 0.0)74 RK = 1.0 / p.beta - 1.0 + p.delta75 T = p.tasks.bundle(QP, L)76 K = T * (p.alpha / RK) ** (1.0 / (1.0 - p.alpha))77 pr = p.tasks.prices(0.0, K, QP, L)78 Y = pr['Y']79 IQ = p.delta_Q * KQ * p.pK_Q80 fee = pr['pQ'] * p.qbar + p.pE * p.ebar81 C = Y - p.delta * K - IQ - p.pE * p.ebar * D82 assert C > 0, "non-positive steady-state consumption"83 return dict(K=K, L=L, QP=QP, KQ=KQ, D=D, Y=Y, C=C, RK=RK, R=1.0 - p.delta + RK, w=pr['w'], pQ=pr['pQ'],84 iota=pr['iota'], share_L=pr['share_L'], share_Q=pr['share_Q'], share_K=pr['share_K'],85 fee=fee, IQ=IQ, I=p.delta * K, user_cost_Q=(1.0 / p.beta - 1.0 + p.delta_Q) * p.pK_Q)868788# --------------------------------------------------------------------------- 2. linearised household (sequence space)89def _solve_household(p: EconParams, ss: dict, what: np.ndarray, Rhat: np.ndarray, a0: float = 0.0,90 windfall: float = 0.0, ge: bool = False, z: np.ndarray | None = None):91 """Solve the linearised household problem for given log-deviation price paths (length T+1).9293 Unknowns: c_hat[0..T], a_hat[1..T+1]. Equations: budget (T+1), Euler (T), terminal a_{T+1} = a_T (1).94 `windfall` adds a lump-sum (in units of abar) to beginning-of-period wealth at t=0 (for the MPC).95 With ge=True the aggregate capital path equals the household's own path (symmetric equilibrium) and prices respond:96 Y_hat = z + alpha a_hat, w_hat = Y_hat, R_hat = (RK/R)(z - (1-alpha) a_hat); `what`/`Rhat` are then ignored.97 Returns (c_hat, a_hat) with a_hat indexed 0..T+1.98 """99 T = len(what) - 1100 cbar, abar, wbar, Rbar, L = ss['C'], ss['K'], ss['w'], ss['R'], ss['L']101 rr = ss['RK'] / ss['R']102 n = 2 * T + 2103 A = np.zeros((n, n))104 b = np.zeros(n)105 ic = lambda t: t # c_hat[t], t = 0..T106 ia = lambda t: T + 1 + (t - 1) # a_hat[t], t = 1..T+1107 if ge:108 what = z.copy()109 Rhat = rr * z110 row = 0111 for t in range(T + 1): # budget: cbar c_t + abar a_{t+1} = wbar L w_t + Rbar abar (R_t + a_t)112 A[row, ic(t)] = cbar113 A[row, ia(t + 1)] = abar114 rhs = wbar * L * what[t] + Rbar * abar * Rhat[t]115 if t == 0:116 rhs += Rbar * abar * a0 + abar * windfall117 else:118 A[row, ia(t)] -= Rbar * abar119 if ge: # price responses to own (= aggregate) capital120 A[row, ia(t)] -= wbar * L * p.alpha121 A[row, ia(t)] -= Rbar * abar * rr * (-(1.0 - p.alpha))122 b[row] = rhs123 row += 1124 for t in range(T): # Euler: c_t = c_{t+1} - (1/gamma) R_{t+1}125 A[row, ic(t)] = 1.0126 A[row, ic(t + 1)] = -1.0127 b[row] = -(1.0 / p.gamma) * Rhat[t + 1]128 if ge:129 A[row, ia(t + 1)] -= (1.0 / p.gamma) * rr * (-(1.0 - p.alpha))130 row += 1131 A[row, ia(T + 1)] = 1.0 # terminal: a_{T+1} = a_T132 A[row, ia(T)] = -1.0133 x = np.linalg.solve(A, b)134 c_hat = x[: T + 1]135 a_hat = np.concatenate([[a0], x[T + 1:]])136 return c_hat, a_hat137138139def _price_paths_pe(p: EconParams, ss: dict, theta: float, T: int):140 """Price paths when the aggregate capital path is fixed and z_s = rho^(s-1) theta for s>=1 (news at t=0)."""141 z = np.zeros(T + 1)142 for s in range(1, T + 1):143 z[s] = p.rho_z ** (s - 1) * theta144 what = z.copy() # w = (1-a)(1-iota) Y/L, Y ∝ e^z with K, QP, L fixed145 Rhat = (ss['RK'] / ss['R']) * z # R = 1 - delta + RK, RK ∝ e^z146 return what, Rhat, z147148149def _solve_ge(p: EconParams, ss: dict, theta: float, T: int):150 """Symmetric rational-expectations response (aggregate k_hat = a_hat), one linear solve."""151 what, Rhat, z = _price_paths_pe(p, ss, theta, T)152 return _solve_household(p, ss, what, Rhat, ge=True, z=z)153154155def linearized_household(p: EconParams, ss: dict, T: int = 400) -> dict:156 theta = 1.0157 what, Rhat, _ = _price_paths_pe(p, ss, theta, T)158 c_pe, a_pe = _solve_household(p, ss, what, Rhat)159 kappa0 = a_pe[1] / theta160 c_ge, a_ge = _solve_ge(p, ss, theta, T)161 kappa_full = a_ge[1] / theta162 alpha = 1.0 - kappa0 / kappa_full163 # MPC out of a beginning-of-period windfall (units: consumption per unit of wealth)164 c_w, a_w = _solve_household(p, ss, np.zeros(T + 1), np.zeros(T + 1), windfall=1.0)165 MPC = c_w[0] * ss['C'] / ss['K'] # dc/da in levels166 cbar = ss['C']167 Uxx_level = p.u2(cbar) + p.beta * p.u2(cbar) * ss['R'] * MPC168 Uxx_log = ss['K'] ** 2 * Uxx_level169 return dict(kappa0=kappa0, kappa_full=kappa_full, alpha=alpha, MPC=MPC, Uxx_level=Uxx_level, Uxx_log=Uxx_log,170 a_pe=a_pe, a_ge=a_ge, c_pe=c_pe, c_ge=c_ge)171172173def lq_params_from_economy(p: EconParams, ss: dict, lin: dict | None = None) -> tuple[LQParams, dict]:174 lin = lin or linearized_household(p, ss)175 lqp = LQParams(kappa0=lin['kappa0'], alpha=lin['alpha'], tau_theta=p.tau_theta, tau_H=p.tau_H,176 tau_A=p.tau_A, M=p.M)177 return lqp, lin178179180# --------------------------------------------------------------------------- 5. delegation share and compute market181def private_value_goods(p: EconParams, ss: dict, lqp: LQParams, lin: dict, D: float) -> float:182 """Private value of delegation in goods per year.183184 Savings mapping: PV (units |U_xx|) * |U_xx|_log / u'(c). Portfolio mapping: PV * |U_xx|_goods (already in goods).185 """186 if 'Uxx_goods' in lin:187 return private_value(lqp, D) * lin['Uxx_goods']188 return private_value(lqp, D) * abs(lin['Uxx_log']) / p.u1(ss['C'])189190191def delegation_share(p: EconParams, ss: dict, lqp: LQParams, lin: dict, pQ: float, D_guess: float = 0.0) -> float:192 """D = F_kappa( PV_goods(D) - fee(pQ) - tau_s w ) solved as a fixed point (bisection on D)."""193 fee = pQ * p.qbar + p.pE * p.ebar194 kap = lognorm(s=p.kappa_sigma, scale=np.exp(p.kappa_mu) * ss['C'])195196 def f(D):197 net = private_value_goods(p, ss, lqp, lin, D) - fee - p.tau_s * ss['w']198 return D - (kap.cdf(net) if net > 0 else 0.0)199200 if f(0.0) >= 0:201 return 0.0202 if f(1.0) <= 0:203 return 1.0204 return float(brentq(f, 0.0, 1.0, xtol=1e-10))205206207def compute_market_short_run(p: EconParams, ss: dict, lqp: LQParams, lin: dict, KQ: float, z: float = 0.0) -> dict:208 """Short-run equilibrium with fixed K, L, K_Q: find Q_P such that Q_P + qbar D(pQ(Q_P)) = K_Q (A.4).209210 pQ(Q_P) is the production compute rental when the firm uses Q_P (decreasing in Q_P); D(pQ) is the delegation211 share at that rental (increasing in Q_P). Excess demand is increasing in Q_P -> unique root.212 """213 K, L = ss['K'], ss['L']214 if KQ <= 0:215 return dict(QP=0.0, D=0.0, pQ=0.0, **p.tasks.prices(z, K, 0.0, L))216217 def excess(QP):218 pr = p.tasks.prices(z, K, QP, L)219 D = delegation_share(p, ss, lqp, lin, pr['pQ'])220 return QP + p.qbar * D - KQ221222 if excess(KQ) <= 0: # all compute to production, delegation demand nil223 QP = KQ224 else:225 QP = brentq(excess, 1e-10, KQ, xtol=1e-12)226 pr = p.tasks.prices(z, K, QP, L)227 D = delegation_share(p, ss, lqp, lin, pr['pQ'])228 fee = pr['pQ'] * p.qbar + p.pE * p.ebar229 return dict(QP=QP, D=D, fee=fee, KQ=KQ, **pr)230231232# --------------------------------------------------------------------------- portfolio (allocation) mapping233def portfolio_lq_params(p: EconParams, ss: dict, investors_earn_wages: bool = True, M: int | None = None) -> dict:234 """Mean-variance allocation of wealth between risky capital claims and a safe store of value at fixed R_f.235236 k_i = (E_i[R^K'] - R_f)/(gamma_a sigma_R^2), gamma_a = gamma / W (W = steady-state wealth = K).237 Next-period expected return: E[R^K'] = Rbar^K [1 - (1-alpha) X] + theta, theta ~ N(0, sig_theta_R^2) the predictable238 return component (news), X the aggregate log capital deviation. Hence, in log holdings,239 x_i = kappa0 E_i[theta] + alpha_s E_i[X], kappa0 = 1/(gamma sigma_R^2), alpha_s = -(1-alpha) Rbar^K kappa0.240 Certainty-equivalent loss of a holding error x (log units) is (1/2) gamma sigma_R^2 K x^2 -> |U_xx| in goods.241 CRS restriction (Lemma 2) holds iff investors also receive wage income (U_XX = -U_xX); with separate investors and242 workers U_XX = 0 (upsilon = 0) and Proposition 4(a) applies. Signals in return units: tau = 1/sigma^2.243 """244 kappa0 = 1.0 / (p.gamma * p.sigma_R ** 2)245 alpha_s = -(1.0 - p.alpha) * ss['RK'] * kappa0246 Uxx_goods = p.gamma * p.sigma_R ** 2 * ss['K']247 ups = 1.0 if investors_earn_wages else 0.0248 lqp = LQParams(kappa0=kappa0, alpha=alpha_s, tau_theta=p.sig_theta_R ** -2, tau_H=p.sig_H_R ** -2,249 tau_A=p.sig_A_R ** -2, M=M or p.M, upsilon=ups)250 return dict(kappa0=kappa0, alpha=alpha_s, Uxx_goods=Uxx_goods, upsilon=ups, lqp=lqp)251