1"""Two-class perfect-foresight general-equilibrium transition.23Unknowns (5 per period t = 0..T-1): log K^N_{t+1}, log K^Q_{t+1}, log K^E_{t+1}, log K^H_{t+1}, r_{t+1}.4Residuals: four capital Euler equations (Tobin's q with quadratic adjustment costs) and the capitalists'5consumption Euler equation. Capitalists' consumption is the goods-market residual (Walras' law), so goods-6market clearing holds identically and is verified ex post together with the income identity and the7asset-market identity (capitalists' wealth = market value of all claims).89Terminal conditions: investment rates, the interest rate and the growth rate of capitalists' expenditure are10held at their period-(T-1) values beyond T; land prices at T are the Gordon value of the terminal rent.11Truncation error is checked by `horizon_check` (re-solving with a longer horizon).1213The Jacobian is banded (each residual depends on unknowns at t-1, t, t+1) and is computed by colored14finite differences; Newton steps use a sparse LU solve and a backtracking line search.15"""16from __future__ import annotations1718import time19from dataclasses import dataclass2021import numpy as np22import scipy.sparse as sp23import scipy.sparse.linalg as spla2425from .production import housing_rent, income_identity, static_period26from .tasks import technology_from_params2728KINDS = ("N", "Q", "E", "H")293031@dataclass32class Paths:33 """Container of equilibrium paths (arrays of length T, except stocks of length T+1)."""34 data: dict3536 def __getitem__(self, k):37 return self.data[k]3839 def keys(self):40 return self.data.keys()414243class TwoClassModel:44 def __init__(self, p):45 self.p = p46 self.tech = technology_from_params(p)47 self.T = p.T48 self.alpha = p.alpha_path()49 self.Z = p.Z_path()50 self.gf = (1.0 + p.g_Z) ** np.arange(p.T + 1) # efficiency units of labor, land, resource51 self.pQ = p.pQ_path()52 self.eps = p.eps_path()53 # extend exogenous paths one period beyond T for terminal evaluations54 self.alpha = np.vstack([self.alpha, self.alpha[-1:]])55 self.Z = np.append(self.Z, self.Z[-1])56 self.pQ = np.append(self.pQ, self.pQ[-1] * max(np.exp(-p.g_pQ), p.pQ_floor_ratio * p.pQ0 / self.pQ[-1]))57 self.eps = np.append(self.eps, self.eps[-1] * max(np.exp(-p.g_eps), p.eps_floor_ratio * p.eps0 / self.eps[-1]))58 self.price = {"N": np.ones(self.T + 2), "Q": np.append(self.pQ, self.pQ[-1]), "E": np.ones(self.T + 2), "H": np.ones(self.T + 2)}59 self._warm = None60 self.n_static = 061 self.nv = 6 # unknowns per period: 4 log stocks, log X_K (capitalists' expenditure), r62 self.r_given = None # exogenous r path (length T+2) when nv == 5 (fixed-r mode)63 self.XK_residual = False # if True, X_K is computed from the goods market (used to build guesses)64 self.W_target = None # optional wealth path (HA block demand) used in evaluate(full) diagnostics6566 # ------------------------------------------------------------------ helpers67 def set_fixed_r(self, r_path):68 """Switch to the 4-unknown mode with an exogenous interest-rate path r[1..T] (array of length T+1 or T+2)."""69 r = np.asarray(r_path, float)70 full = np.empty(self.T + 2)71 full[:len(r)] = r[:self.T + 2]72 if len(r) < self.T + 2:73 full[len(r):] = r[-1]74 self.r_given = full75 self.nv = 57677 def set_free_r(self):78 self.r_given = None79 self.nv = 68081 @property82 def ir(self):83 """column index of r (None in fixed-r mode)"""84 return 5 if self.nv == 6 else None8586 def unpack(self, u):87 T = self.T88 U = u.reshape(T, self.nv)89 K = {k: np.empty(T + 2) for k in KINDS}90 K0 = {"N": self.p.KN0, "Q": self.p.KQ0, "E": self.p.KE0, "H": self.p.KH0}91 for j, k in enumerate(KINDS):92 K[k][0] = K0[k]93 K[k][1:T + 1] = np.exp(np.clip(U[:, j], np.log(K0[k]) - 12.0, np.log(K0[k]) + 90.0))94 d = self.p.delta[k]95 x_last = K[k][T] / K[k][T - 1] - (1 - d)96 K[k][T + 1] = K[k][T] * (1 - d + x_last) # terminal: constant investment rate97 if self.nv == 6:98 r = np.empty(T + 2)99 r[1:T + 1] = U[:, 5]100 r[0] = np.nan101 r[T + 1] = r[T]102 else:103 r = self.r_given.copy()104 XK = np.empty(T + 1)105 XK[:T] = np.exp(np.clip(U[:, 4], -30, 30))106 XK[T] = XK[T - 1] ** 2 / XK[T - 2] # terminal: constant growth of expenditure107 return K, r, XK108109 def initial_guess(self, r0=0.04, g=None):110 p = self.p111 g = p.g_Z if g is None else g112 T = self.T113 U = np.empty((T, 6))114 K0 = {"N": p.KN0, "Q": p.KQ0, "E": p.KE0, "H": p.KH0}115 for j, k in enumerate(KINDS):116 U[:, j] = np.log(K0[k] * (1 + g) ** np.arange(1, T + 1))117 U[:, 5] = r0118 U[:, 4] = 0.0119 u = U.ravel()120 return self.fill_XK(u)121122 def fill_XK(self, u):123 """Set the log X_K column from the goods-market residual at the current stocks (guess builder)."""124 self.XK_residual = True125 P = self.evaluate(u, full=True)126 self.XK_residual = False127 U = u.reshape(self.T, self.nv)128 U[:, 4] = np.log(np.maximum(P["X_K"][:self.T], 1e-3))129 return U.ravel()130131 def smart_guess(self, r0=None, verbose=False, max_growth=0.35):132 """Partial-equilibrium guess: other stocks on the frozen BGP path, compute stock chosen period by133 period so that its net rental equals the user cost (1+r)p_{t-1} - (1-delta)p_t (no adjustment cost)."""134 from scipy.optimize import brentq135 from .steady import bgp_interest_rate136 p, T = self.p, self.T137 r0 = bgp_interest_rate(p) if r0 is None else r0138 u = self.initial_guess(r0=r0, g=p.g_Z)139 U = u.reshape(T, 6)140 K0 = {"N": p.KN0, "E": p.KE0, "H": p.KH0}141 x0 = None142 d = p.delta["Q"]143 for t in range(1, T + 1):144 KN, KE, KH = (K0[k] * (1 + p.g_Z) ** t for k in ("N", "E", "H"))145 uc = (1 + r0) * self.pQ[t - 1] - (1 - d) * self.pQ[t]146 def stat(lk):147 return static_period(p, self.tech, KN, np.exp(lk), KE, KH, self.alpha[t], self.Z[t], self.eps[t], x0=x0, gf=self.gf[t])148149 def f(lk):150 s = stat(lk)151 return np.log(max(s["RQ"], 1e-300)) - np.log(uc)152 lo, hi = np.log(p.KQ0) - 5, np.log(p.KQ0) + 40153 try:154 if f(lo) > 0 and f(hi) < 0:155 lk = brentq(f, lo, hi, xtol=1e-3)156 elif f(hi) >= 0:157 lk = hi158 else:159 lk = lo160 # do not guess idle capacity: cap at the energy-feasible stock (utilization = 1)161 if stat(lk)["util"] < 0.999:162 a, b = lo, lk163 for _ in range(25):164 mid = 0.5 * (a + b)165 if stat(mid)["util"] < 0.999:166 b = mid167 else:168 a = mid169 lk = a170 except Exception:171 lk = U[t - 1, 1]172 prev = np.log(p.KQ0) if t == 1 else U[t - 2, 1]173 lk = max(lk, prev + np.log(1 - d)) # firms cannot shed capital faster than depreciation174 lk = min(lk, prev + np.log(1 + max_growth)) # and adjustment costs make very fast growth unattractive175 U[t - 1, 1] = lk176 if verbose:177 print(" smart guess K_Q path (first/last):", np.exp(U[0, 1]), np.exp(U[-1, 1]))178 return self.fill_XK(U.ravel())179180 # ------------------------------------------------------------------ core evaluation181 def evaluate(self, u, full=False):182 """Static blocks for t = 0..T, flows, prices; returns residual vector (and paths if full)."""183 p, T = self.p, self.T184 K, r, XK_u = self.unpack(u)185 per = []186 x0 = None187 for t in range(T + 1):188 s = static_period(p, self.tech, K["N"][t], K["Q"][t], K["E"][t], K["H"][t],189 self.alpha[t], self.Z[t], self.eps[t], x0=x0, gf=self.gf[t])190 x0 = s["x"]191 per.append(s)192 self.n_static += 1193 # investment rates, q, adjustment costs194 x = {k: K[k][1:T + 2] / K[k][:T + 1] - (1 - p.delta[k]) for k in KINDS} # t = 0..T195 I = {k: x[k] * K[k][:T + 1] for k in KINDS}196 q = {k: self.price[k][:T + 1] * (1 + p.phi[k] * (x[k] - p.delta[k])) for k in KINDS}197 adj = {k: self.price[k][:T + 1] * 0.5 * p.phi[k] * (x[k] - p.delta[k]) ** 2 * K[k][:T + 1] for k in KINDS}198 inv_goods = sum(self.price[k][:T + 1] * I[k] + adj[k] for k in KINDS)199 R = {"N": np.array([s["RN"] for s in per]), "Q": np.array([s["RQ"] for s in per]),200 "E": np.array([s["RE"] for s in per]), "H": np.zeros(T + 1)}201 Y = np.array([s["Y"] for s in per])202 wage = np.array([s["wage_bill"] for s in per])203 wL = np.array([s["w"][0] for s in per])204 wH = np.array([s["w"][1] for s in per])205 A = np.array([s["A"] for s in per])206 RT = np.array([s["RT"] for s in per])207 rentR = np.array([s["rentR"] for s in per])208 # taxes and transfers209 T_eff = p.T_land * self.gf210 TH_eff = p.TH_land * self.gf211 cap_income_gross = sum(R[k] * K[k][:T + 1] for k in ("N", "Q", "E")) + RT * T_eff + rentR212 # foreign ownership of compute: net outflow = share * (payout of compute sector)213 payout_Q = R["Q"] * K["Q"][:T + 1] - self.price["Q"][:T + 1] * I["Q"] - adj["Q"]214 foreign_flow = p.foreign_share_Q * payout_Q215 # workers hand-to-mouth: expenditure = wages + transfers216 H = np.asarray(p.H)217 # goods consumption total = Y - investment - foreign net outflow218 C_total = Y - inv_goods - foreign_flow219 for t in range(T + 1):220 housing_rent(p, per[t], C_total[t])221 RH = np.array([s["RH"] for s in per])222 R["H"] = np.array([s["RHK"] for s in per])223 RTH = np.array([s["RTH"] for s in per])224 cap_income_gross = cap_income_gross + R["H"] * K["H"][:T + 1] + RTH * TH_eff225 # policy: land tax, capital income tax, compute tax -> lump-sum transfers to workers; citizen fund226 tax = p.tau_T * (RT * T_eff + RTH * TH_eff) + p.tau_K * cap_income_gross + p.tau_Q * R["Q"] * K["Q"][:T + 1]227 payout_all = sum(R[k] * K[k][:T + 1] - self.price[k][:T + 1] * I[k] - adj[k] for k in KINDS) + RT * T_eff + RTH * TH_eff228 fund = p.theta_F * payout_all229 mass_workers = H.sum()230 transfers_workers = tax + fund * mass_workers / (mass_workers + 1.0) # capitalists have unit mass231 X_W = wage + transfers_workers # workers' total expenditure (goods + housing)232 C_W = (1 - p.theta_H) * X_W233 if self.XK_residual:234 C_K = C_total - C_W # capitalists' goods consumption from the goods market235 X_K = C_K / (1 - p.theta_H)236 else:237 X_K = XK_u238 C_K = (1 - p.theta_H) * X_K239 goods_gap = (C_total - C_W - C_K) / Y # goods-market residual (zero in equilibrium)240 P_idx = RH ** p.theta_H # ideal price index (up to a constant)241 # ---- residuals, t = 0..T-1242 res = np.empty((T, self.nv))243 for j, k in enumerate(KINDS):244 d, ph = p.delta[k], p.phi[k]245 pk = self.price[k]246 rhs = (R[k][1:T + 1] + pk[1:T + 1] * (ph * (x[k][1:T + 1] - d) * x[k][1:T + 1] - 0.5 * ph * (x[k][1:T + 1] - d) ** 2)247 + (1 - d) * q[k][1:T + 1])248 # AI-provider markup: investment uses marginal revenue A/mu instead of A249 if k == "Q" and p.markup_Q != 1.0:250 util = np.array([s["util"] for s in per])251 PE = np.array([s["PE"] for s in per])252 rhs = rhs - util[1:T + 1] * A[1:T + 1] * (1 - 1 / p.markup_Q)253 if k == "Q" and p.tau_Q > 0:254 rhs = rhs - p.tau_Q * R["Q"][1:T + 1]255 if p.tau_K > 0:256 rhs = rhs - p.tau_K * R[k][1:T + 1]257 res[:, j] = (q[k][:T] * (1 + r[1:T + 1]) - rhs) / pk[:T]258 gX = X_K[1:T + 1] / X_K[:T]259 gP = P_idx[1:T + 1] / P_idx[:T]260 res[:, 4] = goods_gap[:T]261 if self.nv == 6:262 res[:, 5] = p.gamma_u * np.log(gX) + (1 - p.gamma_u) * np.log(gP) - np.log(p.beta * (1 + r[1:T + 1]))263 if not full:264 out = res.ravel()265 return np.where(np.isfinite(out), out, 1e6)266 # ---- asset prices and wealth267 # land prices per physical unit; rents per physical unit = rent per efficiency unit * gf268 RTp, RTHp = RT * self.gf, RTH * self.gf269 PT = np.empty(T + 1)270 PTH = np.empty(T + 1)271 gR = RTp[T] / RTp[T - 1] - 1272 gRH = RTHp[T] / RTHp[T - 1] - 1273 rT = r[T]274 PT[T] = RTp[T] * (1 - p.tau_T) * (1 + gR) / max(rT - gR, 1e-4)275 PTH[T] = RTHp[T] * (1 - p.tau_T) * (1 + gRH) / max(rT - gRH, 1e-4)276 for t in range(T - 1, -1, -1):277 PT[t] = (RTp[t + 1] * (1 - p.tau_T) + PT[t + 1]) / (1 + r[t + 1])278 PTH[t] = (RTHp[t + 1] * (1 - p.tau_T) + PTH[t + 1]) / (1 + r[t + 1])279 V = {k: q[k] * K[k][1:T + 2] for k in KINDS} # ex-dividend equity values280 D = {k: R[k] * K[k][:T + 1] - self.price[k][:T + 1] * I[k] - adj[k] for k in KINDS} # net payouts281 W = sum(V[k] for k in KINDS) + PT * p.T_land + PTH * p.TH_land282 GDP = np.array([s["GDP"] for s in per])283 out = dict(years=np.arange(p.year0, p.year0 + T + 1), K=K, r=r, x=x, I=I, q=q, adj=adj, R=R, Y=Y, GDP=GDP,284 wage=wage, wL=wL, wH=wH, A=A, RT=RT, RTH=RTH, RTp=RTp, RTHp=RTHp, RH=RH, rentR=rentR, PE=np.array([s["PE"] for s in per]),285 util=np.array([s["util"] for s in per]), Qtot=np.array([s["Qtot"] for s in per]),286 E=np.array([s["E"] for s in per]), EQ=np.array([s["EQ"] for s in per]),287 C_total=C_total, C_K=C_K, C_W=C_W, X_K=X_K, X_W=X_W, P_idx=P_idx, transfers=transfers_workers,288 tax=tax, fund=fund, PT=PT, PTH=PTH, V=V, D=D, W=W, inv_goods=inv_goods, foreign_flow=foreign_flow,289 labor_share=wage / GDP, per=per, residuals=res, gf=self.gf, T_eff=T_eff, TH_eff=TH_eff, goods_gap=goods_gap,290 Li=np.array([s["L"] for s in per]), Qi=np.array([s["Qi"] for s in per]),291 Xi=np.array([s["X"] for s in per]), HS=np.array([s["HS"] for s in per]),292 alpha=self.alpha[:T + 1], pQ=self.pQ[:T + 1], eps=self.eps[:T + 1], Z=self.Z[:T + 1])293 out["income_residual"] = np.array([income_identity(p, s) for s in per])294 return Paths(out)295296 # ------------------------------------------------------------------ Jacobian (colored FD, banded)297 def jacobian(self, u, f0, h=1e-5, band=2):298 T = self.T299 nv = self.nv300 n = nv * T301 ncol = nv * (2 * band + 1)302 J = sp.lil_matrix((n, n))303 for c in range(ncol):304 idx = np.arange(c, n, ncol)305 up = u.copy()306 up[idx] += h307 fp = self.evaluate(up)308 df = (fp - f0) / h309 df = np.where(np.isfinite(df), df, 0.0)310 for i in idx:311 t = i // nv312 lo, hi = max(0, t - band) * nv, min(T, t + band + 1) * nv313 J[lo:hi, i] = df[lo:hi].reshape(-1, 1)314 return J.tocsc()315316 def solve(self, u0=None, tol=1e-9, maxit=60, verbose=True, h=1e-5, band=2):317 u = self.initial_guess() if u0 is None else u0.copy()318 f = self.evaluate(u)319 nf = np.max(np.abs(f))320 t0 = time.time()321 hist = []322 J = None323 for it in range(maxit):324 if verbose:325 im = int(np.argmax(np.abs(f)))326 print(f" newton it {it:2d} |res|_inf = {nf:.3e} at t={im // self.nv}, eq={('N','Q','E','H','goods','euler')[im % self.nv]} ({time.time() - t0:.0f}s)", flush=True)327 hist.append(nf)328 if nf < tol:329 break330 J = self.jacobian(u, f, h=h, band=band)331 lu = spla.splu(J)332 du = -lu.solve(f)333 lam = 1.0334 best = None335 for _ in range(16):336 un = u + lam * du337 fn = self.evaluate(un)338 nfn = np.max(np.abs(fn))339 if np.isfinite(nfn) and (best is None or nfn < best[2]):340 best = (un, fn, nfn)341 if np.isfinite(nfn) and nfn < nf:342 break343 lam *= 0.5344 else:345 # no decrease along the Newton direction: take the best finite trial point anyway346 if best is None:347 raise RuntimeError("Newton: no finite trial point found")348 un, fn, nfn = best349 u, f, nf = un, fn, nfn350 converged = nf < tol351 if verbose:352 print(f" done: |res|_inf = {nf:.3e}, converged={converged}, {time.time() - t0:.0f}s, static solves={self.n_static}")353 return u, converged, hist354355356def solve_continuation(p, horizons=(20, 40, 80, 140, 200), verbose=True, tol=1e-9, **kw):357 """Horizon continuation: solve on short horizons first, extend the solution (terminal growth extrapolated)358 and re-solve, up to p.T. Returns (model, u, converged, history)."""359 hs = [h for h in horizons if h < p.T] + [p.T]360 u_prev, T_prev = None, None361 hist_all = []362 for Th in hs:363 pT = p.replace(T=Th)364 m = TwoClassModel(pT)365 if u_prev is None:366 u0 = m.smart_guess()367 else:368 U = u_prev.reshape(T_prev, 6)369 extra = Th - T_prev370 Ul = np.vstack([U, np.tile(U[-1], (extra, 1))])371 g = U[-1, :5] - U[-2, :5]372 for j in range(extra):373 Ul[T_prev + j, :5] = U[-1, :5] + g * (j + 1)374 u0 = Ul.ravel()375 if verbose:376 print(f"[continuation] horizon {Th}")377 u, ok, hist = m.solve(u0, tol=tol, verbose=verbose, **kw)378 hist_all.append((Th, ok, hist[-1]))379 if not ok and verbose:380 print(f" warning: horizon {Th} not converged (|res| = {hist[-1]:.2e})")381 u_prev, T_prev = u, Th382 return m, u_prev, ok, hist_all383384385def sparsity_pattern(T, band=2):386 n = 5 * T387 S = sp.lil_matrix((n, n), dtype=np.int8)388 for t in range(T):389 lo, hi = max(0, t - band) * 5, min(T, t + band + 1) * 5390 S[t * 5:(t + 1) * 5, lo:hi] = 1391 return S.tocsr()392393394def solve_trf(model, u0, tol=1e-10, max_nfev=400, verbose=True, band=2):395 """Trust-region least-squares (scipy TRF) with the banded sparsity pattern; robust to the flat directions396 created by the task tie (A = w/alpha) and to kinks of the static block."""397 from scipy.optimize import least_squares398 t0 = time.time()399 S = sparsity_pattern(model.T, band)400 res = least_squares(model.evaluate, u0, jac="2-point", jac_sparsity=S, method="trf", x_scale="jac",401 tr_solver="lsmr", ftol=1e-15, xtol=1e-15, gtol=1e-15, max_nfev=max_nfev,402 verbose=1 if verbose else 0, diff_step=1e-6)403 f = model.evaluate(res.x)404 nf = float(np.max(np.abs(f)))405 if verbose:406 print(f" trf: |res|_inf = {nf:.3e}, nfev={res.nfev}, {time.time() - t0:.0f}s, status={res.status}")407 return res.x, nf < tol, nf408409410def solve_homotopy(p, lambdas=(0.0, 0.15, 0.3, 0.5, 0.7, 0.85, 1.0), verbose=True, tol=1e-9, u0=None, **kw):411 """Continuation in the size of the technology shock: g_pQ, g_eps and g_alpha are scaled by lambda412 (alpha_shock_year scenarios scale the log-distance to alpha_inf instead)."""413 base = dict(g_pQ=p.g_pQ, g_eps=p.g_eps, g_alpha=p.g_alpha)414 u, ok, hist_all = u0, False, []415 for lam in lambdas:416 pl = p.replace(g_pQ=lam * base["g_pQ"], g_eps=lam * base["g_eps"], g_alpha=lam * base["g_alpha"])417 m = TwoClassModel(pl)418 if u is None:419 u = m.smart_guess()420 if verbose:421 print(f"[homotopy] lambda = {lam}")422 u, ok, hist = m.solve(u, tol=tol, verbose=verbose, **kw)423 hist_all.append((lam, ok, hist[-1]))424 return m, u, ok, hist_all425426427def horizon_check(p, u_short, extra=50, **kw):428 """Re-solve with a longer horizon from the short solution and report the max relative change of the429 first p.T-extra periods of the capital and rate paths (truncation error)."""430 pl = p.replace(T=p.T + extra)431 ml = TwoClassModel(pl)432 U = u_short.reshape(p.T, 6)433 Ul = np.vstack([U, np.tile(U[-1], (extra, 1))])434 # extend capital logs along the terminal growth rate435 g = U[-1, :4] - U[-2, :4]436 for j in range(extra):437 Ul[p.T + j, :4] = U[-1, :4] + g * (j + 1)438 ul, ok, _ = ml.solve(Ul.ravel(), **kw)439 Ks, rs, _ = ml.unpack(ul)440 Ks0, rs0, _ = TwoClassModel(p).unpack(u_short)441 n = p.T - extra442 dev = max(np.max(np.abs(np.log(Ks[k][1:n + 1] / Ks0[k][1:n + 1]))) for k in KINDS)443 devr = np.max(np.abs(rs[1:n + 1] - rs0[1:n + 1]))444 return dict(ok=ok, max_log_dev_K=dev, max_dev_r=devr, u_long=ul, model_long=ml)445