"""Two-class perfect-foresight general-equilibrium transition. Unknowns (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}. Residuals: four capital Euler equations (Tobin's q with quadratic adjustment costs) and the capitalists' consumption Euler equation. Capitalists' consumption is the goods-market residual (Walras' law), so goods- market clearing holds identically and is verified ex post together with the income identity and the asset-market identity (capitalists' wealth = market value of all claims). Terminal conditions: investment rates, the interest rate and the growth rate of capitalists' expenditure are held at their period-(T-1) values beyond T; land prices at T are the Gordon value of the terminal rent. Truncation error is checked by `horizon_check` (re-solving with a longer horizon). The Jacobian is banded (each residual depends on unknowns at t-1, t, t+1) and is computed by colored finite differences; Newton steps use a sparse LU solve and a backtracking line search. """ from __future__ import annotations import time from dataclasses import dataclass import numpy as np import scipy.sparse as sp import scipy.sparse.linalg as spla from .production import housing_rent, income_identity, static_period from .tasks import technology_from_params KINDS = ("N", "Q", "E", "H") @dataclass class Paths: """Container of equilibrium paths (arrays of length T, except stocks of length T+1).""" data: dict def __getitem__(self, k): return self.data[k] def keys(self): return self.data.keys() class TwoClassModel: def __init__(self, p): self.p = p self.tech = technology_from_params(p) self.T = p.T self.alpha = p.alpha_path() self.Z = p.Z_path() self.gf = (1.0 + p.g_Z) ** np.arange(p.T + 1) # efficiency units of labor, land, resource self.pQ = p.pQ_path() self.eps = p.eps_path() # extend exogenous paths one period beyond T for terminal evaluations self.alpha = np.vstack([self.alpha, self.alpha[-1:]]) self.Z = np.append(self.Z, self.Z[-1]) self.pQ = np.append(self.pQ, self.pQ[-1] * max(np.exp(-p.g_pQ), p.pQ_floor_ratio * p.pQ0 / self.pQ[-1])) self.eps = np.append(self.eps, self.eps[-1] * max(np.exp(-p.g_eps), p.eps_floor_ratio * p.eps0 / self.eps[-1])) 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)} self._warm = None self.n_static = 0 self.nv = 6 # unknowns per period: 4 log stocks, log X_K (capitalists' expenditure), r self.r_given = None # exogenous r path (length T+2) when nv == 5 (fixed-r mode) self.XK_residual = False # if True, X_K is computed from the goods market (used to build guesses) self.W_target = None # optional wealth path (HA block demand) used in evaluate(full) diagnostics # ------------------------------------------------------------------ helpers def set_fixed_r(self, r_path): """Switch to the 4-unknown mode with an exogenous interest-rate path r[1..T] (array of length T+1 or T+2).""" r = np.asarray(r_path, float) full = np.empty(self.T + 2) full[:len(r)] = r[:self.T + 2] if len(r) < self.T + 2: full[len(r):] = r[-1] self.r_given = full self.nv = 5 def set_free_r(self): self.r_given = None self.nv = 6 @property def ir(self): """column index of r (None in fixed-r mode)""" return 5 if self.nv == 6 else None def unpack(self, u): T = self.T U = u.reshape(T, self.nv) K = {k: np.empty(T + 2) for k in KINDS} K0 = {"N": self.p.KN0, "Q": self.p.KQ0, "E": self.p.KE0, "H": self.p.KH0} for j, k in enumerate(KINDS): K[k][0] = K0[k] K[k][1:T + 1] = np.exp(np.clip(U[:, j], np.log(K0[k]) - 12.0, np.log(K0[k]) + 90.0)) d = self.p.delta[k] x_last = K[k][T] / K[k][T - 1] - (1 - d) K[k][T + 1] = K[k][T] * (1 - d + x_last) # terminal: constant investment rate if self.nv == 6: r = np.empty(T + 2) r[1:T + 1] = U[:, 5] r[0] = np.nan r[T + 1] = r[T] else: r = self.r_given.copy() XK = np.empty(T + 1) XK[:T] = np.exp(np.clip(U[:, 4], -30, 30)) XK[T] = XK[T - 1] ** 2 / XK[T - 2] # terminal: constant growth of expenditure return K, r, XK def initial_guess(self, r0=0.04, g=None): p = self.p g = p.g_Z if g is None else g T = self.T U = np.empty((T, 6)) K0 = {"N": p.KN0, "Q": p.KQ0, "E": p.KE0, "H": p.KH0} for j, k in enumerate(KINDS): U[:, j] = np.log(K0[k] * (1 + g) ** np.arange(1, T + 1)) U[:, 5] = r0 U[:, 4] = 0.0 u = U.ravel() return self.fill_XK(u) def fill_XK(self, u): """Set the log X_K column from the goods-market residual at the current stocks (guess builder).""" self.XK_residual = True P = self.evaluate(u, full=True) self.XK_residual = False U = u.reshape(self.T, self.nv) U[:, 4] = np.log(np.maximum(P["X_K"][:self.T], 1e-3)) return U.ravel() def smart_guess(self, r0=None, verbose=False, max_growth=0.35): """Partial-equilibrium guess: other stocks on the frozen BGP path, compute stock chosen period by period so that its net rental equals the user cost (1+r)p_{t-1} - (1-delta)p_t (no adjustment cost).""" from scipy.optimize import brentq from .steady import bgp_interest_rate p, T = self.p, self.T r0 = bgp_interest_rate(p) if r0 is None else r0 u = self.initial_guess(r0=r0, g=p.g_Z) U = u.reshape(T, 6) K0 = {"N": p.KN0, "E": p.KE0, "H": p.KH0} x0 = None d = p.delta["Q"] for t in range(1, T + 1): KN, KE, KH = (K0[k] * (1 + p.g_Z) ** t for k in ("N", "E", "H")) uc = (1 + r0) * self.pQ[t - 1] - (1 - d) * self.pQ[t] def stat(lk): 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]) def f(lk): s = stat(lk) return np.log(max(s["RQ"], 1e-300)) - np.log(uc) lo, hi = np.log(p.KQ0) - 5, np.log(p.KQ0) + 40 try: if f(lo) > 0 and f(hi) < 0: lk = brentq(f, lo, hi, xtol=1e-3) elif f(hi) >= 0: lk = hi else: lk = lo # do not guess idle capacity: cap at the energy-feasible stock (utilization = 1) if stat(lk)["util"] < 0.999: a, b = lo, lk for _ in range(25): mid = 0.5 * (a + b) if stat(mid)["util"] < 0.999: b = mid else: a = mid lk = a except Exception: lk = U[t - 1, 1] prev = np.log(p.KQ0) if t == 1 else U[t - 2, 1] lk = max(lk, prev + np.log(1 - d)) # firms cannot shed capital faster than depreciation lk = min(lk, prev + np.log(1 + max_growth)) # and adjustment costs make very fast growth unattractive U[t - 1, 1] = lk if verbose: print(" smart guess K_Q path (first/last):", np.exp(U[0, 1]), np.exp(U[-1, 1])) return self.fill_XK(U.ravel()) # ------------------------------------------------------------------ core evaluation def evaluate(self, u, full=False): """Static blocks for t = 0..T, flows, prices; returns residual vector (and paths if full).""" p, T = self.p, self.T K, r, XK_u = self.unpack(u) per = [] x0 = None for t in range(T + 1): s = static_period(p, self.tech, K["N"][t], K["Q"][t], K["E"][t], K["H"][t], self.alpha[t], self.Z[t], self.eps[t], x0=x0, gf=self.gf[t]) x0 = s["x"] per.append(s) self.n_static += 1 # investment rates, q, adjustment costs x = {k: K[k][1:T + 2] / K[k][:T + 1] - (1 - p.delta[k]) for k in KINDS} # t = 0..T I = {k: x[k] * K[k][:T + 1] for k in KINDS} q = {k: self.price[k][:T + 1] * (1 + p.phi[k] * (x[k] - p.delta[k])) for k in KINDS} 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} inv_goods = sum(self.price[k][:T + 1] * I[k] + adj[k] for k in KINDS) R = {"N": np.array([s["RN"] for s in per]), "Q": np.array([s["RQ"] for s in per]), "E": np.array([s["RE"] for s in per]), "H": np.zeros(T + 1)} Y = np.array([s["Y"] for s in per]) wage = np.array([s["wage_bill"] for s in per]) wL = np.array([s["w"][0] for s in per]) wH = np.array([s["w"][1] for s in per]) A = np.array([s["A"] for s in per]) RT = np.array([s["RT"] for s in per]) rentR = np.array([s["rentR"] for s in per]) # taxes and transfers T_eff = p.T_land * self.gf TH_eff = p.TH_land * self.gf cap_income_gross = sum(R[k] * K[k][:T + 1] for k in ("N", "Q", "E")) + RT * T_eff + rentR # foreign ownership of compute: net outflow = share * (payout of compute sector) payout_Q = R["Q"] * K["Q"][:T + 1] - self.price["Q"][:T + 1] * I["Q"] - adj["Q"] foreign_flow = p.foreign_share_Q * payout_Q # workers hand-to-mouth: expenditure = wages + transfers H = np.asarray(p.H) # goods consumption total = Y - investment - foreign net outflow C_total = Y - inv_goods - foreign_flow for t in range(T + 1): housing_rent(p, per[t], C_total[t]) RH = np.array([s["RH"] for s in per]) R["H"] = np.array([s["RHK"] for s in per]) RTH = np.array([s["RTH"] for s in per]) cap_income_gross = cap_income_gross + R["H"] * K["H"][:T + 1] + RTH * TH_eff # policy: land tax, capital income tax, compute tax -> lump-sum transfers to workers; citizen fund 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] 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_eff fund = p.theta_F * payout_all mass_workers = H.sum() transfers_workers = tax + fund * mass_workers / (mass_workers + 1.0) # capitalists have unit mass X_W = wage + transfers_workers # workers' total expenditure (goods + housing) C_W = (1 - p.theta_H) * X_W if self.XK_residual: C_K = C_total - C_W # capitalists' goods consumption from the goods market X_K = C_K / (1 - p.theta_H) else: X_K = XK_u C_K = (1 - p.theta_H) * X_K goods_gap = (C_total - C_W - C_K) / Y # goods-market residual (zero in equilibrium) P_idx = RH ** p.theta_H # ideal price index (up to a constant) # ---- residuals, t = 0..T-1 res = np.empty((T, self.nv)) for j, k in enumerate(KINDS): d, ph = p.delta[k], p.phi[k] pk = self.price[k] 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) + (1 - d) * q[k][1:T + 1]) # AI-provider markup: investment uses marginal revenue A/mu instead of A if k == "Q" and p.markup_Q != 1.0: util = np.array([s["util"] for s in per]) PE = np.array([s["PE"] for s in per]) rhs = rhs - util[1:T + 1] * A[1:T + 1] * (1 - 1 / p.markup_Q) if k == "Q" and p.tau_Q > 0: rhs = rhs - p.tau_Q * R["Q"][1:T + 1] if p.tau_K > 0: rhs = rhs - p.tau_K * R[k][1:T + 1] res[:, j] = (q[k][:T] * (1 + r[1:T + 1]) - rhs) / pk[:T] gX = X_K[1:T + 1] / X_K[:T] gP = P_idx[1:T + 1] / P_idx[:T] res[:, 4] = goods_gap[:T] if self.nv == 6: res[:, 5] = p.gamma_u * np.log(gX) + (1 - p.gamma_u) * np.log(gP) - np.log(p.beta * (1 + r[1:T + 1])) if not full: out = res.ravel() return np.where(np.isfinite(out), out, 1e6) # ---- asset prices and wealth # land prices per physical unit; rents per physical unit = rent per efficiency unit * gf RTp, RTHp = RT * self.gf, RTH * self.gf PT = np.empty(T + 1) PTH = np.empty(T + 1) gR = RTp[T] / RTp[T - 1] - 1 gRH = RTHp[T] / RTHp[T - 1] - 1 rT = r[T] PT[T] = RTp[T] * (1 - p.tau_T) * (1 + gR) / max(rT - gR, 1e-4) PTH[T] = RTHp[T] * (1 - p.tau_T) * (1 + gRH) / max(rT - gRH, 1e-4) for t in range(T - 1, -1, -1): PT[t] = (RTp[t + 1] * (1 - p.tau_T) + PT[t + 1]) / (1 + r[t + 1]) PTH[t] = (RTHp[t + 1] * (1 - p.tau_T) + PTH[t + 1]) / (1 + r[t + 1]) V = {k: q[k] * K[k][1:T + 2] for k in KINDS} # ex-dividend equity values D = {k: R[k] * K[k][:T + 1] - self.price[k][:T + 1] * I[k] - adj[k] for k in KINDS} # net payouts W = sum(V[k] for k in KINDS) + PT * p.T_land + PTH * p.TH_land GDP = np.array([s["GDP"] for s in per]) 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, 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]), util=np.array([s["util"] for s in per]), Qtot=np.array([s["Qtot"] for s in per]), E=np.array([s["E"] for s in per]), EQ=np.array([s["EQ"] for s in per]), 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, tax=tax, fund=fund, PT=PT, PTH=PTH, V=V, D=D, W=W, inv_goods=inv_goods, foreign_flow=foreign_flow, labor_share=wage / GDP, per=per, residuals=res, gf=self.gf, T_eff=T_eff, TH_eff=TH_eff, goods_gap=goods_gap, Li=np.array([s["L"] for s in per]), Qi=np.array([s["Qi"] for s in per]), Xi=np.array([s["X"] for s in per]), HS=np.array([s["HS"] for s in per]), alpha=self.alpha[:T + 1], pQ=self.pQ[:T + 1], eps=self.eps[:T + 1], Z=self.Z[:T + 1]) out["income_residual"] = np.array([income_identity(p, s) for s in per]) return Paths(out) # ------------------------------------------------------------------ Jacobian (colored FD, banded) def jacobian(self, u, f0, h=1e-5, band=2): T = self.T nv = self.nv n = nv * T ncol = nv * (2 * band + 1) J = sp.lil_matrix((n, n)) for c in range(ncol): idx = np.arange(c, n, ncol) up = u.copy() up[idx] += h fp = self.evaluate(up) df = (fp - f0) / h df = np.where(np.isfinite(df), df, 0.0) for i in idx: t = i // nv lo, hi = max(0, t - band) * nv, min(T, t + band + 1) * nv J[lo:hi, i] = df[lo:hi].reshape(-1, 1) return J.tocsc() def solve(self, u0=None, tol=1e-9, maxit=60, verbose=True, h=1e-5, band=2): u = self.initial_guess() if u0 is None else u0.copy() f = self.evaluate(u) nf = np.max(np.abs(f)) t0 = time.time() hist = [] J = None for it in range(maxit): if verbose: im = int(np.argmax(np.abs(f))) 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) hist.append(nf) if nf < tol: break J = self.jacobian(u, f, h=h, band=band) lu = spla.splu(J) du = -lu.solve(f) lam = 1.0 best = None for _ in range(16): un = u + lam * du fn = self.evaluate(un) nfn = np.max(np.abs(fn)) if np.isfinite(nfn) and (best is None or nfn < best[2]): best = (un, fn, nfn) if np.isfinite(nfn) and nfn < nf: break lam *= 0.5 else: # no decrease along the Newton direction: take the best finite trial point anyway if best is None: raise RuntimeError("Newton: no finite trial point found") un, fn, nfn = best u, f, nf = un, fn, nfn converged = nf < tol if verbose: print(f" done: |res|_inf = {nf:.3e}, converged={converged}, {time.time() - t0:.0f}s, static solves={self.n_static}") return u, converged, hist def solve_continuation(p, horizons=(20, 40, 80, 140, 200), verbose=True, tol=1e-9, **kw): """Horizon continuation: solve on short horizons first, extend the solution (terminal growth extrapolated) and re-solve, up to p.T. Returns (model, u, converged, history).""" hs = [h for h in horizons if h < p.T] + [p.T] u_prev, T_prev = None, None hist_all = [] for Th in hs: pT = p.replace(T=Th) m = TwoClassModel(pT) if u_prev is None: u0 = m.smart_guess() else: U = u_prev.reshape(T_prev, 6) extra = Th - T_prev Ul = np.vstack([U, np.tile(U[-1], (extra, 1))]) g = U[-1, :5] - U[-2, :5] for j in range(extra): Ul[T_prev + j, :5] = U[-1, :5] + g * (j + 1) u0 = Ul.ravel() if verbose: print(f"[continuation] horizon {Th}") u, ok, hist = m.solve(u0, tol=tol, verbose=verbose, **kw) hist_all.append((Th, ok, hist[-1])) if not ok and verbose: print(f" warning: horizon {Th} not converged (|res| = {hist[-1]:.2e})") u_prev, T_prev = u, Th return m, u_prev, ok, hist_all def sparsity_pattern(T, band=2): n = 5 * T S = sp.lil_matrix((n, n), dtype=np.int8) for t in range(T): lo, hi = max(0, t - band) * 5, min(T, t + band + 1) * 5 S[t * 5:(t + 1) * 5, lo:hi] = 1 return S.tocsr() def solve_trf(model, u0, tol=1e-10, max_nfev=400, verbose=True, band=2): """Trust-region least-squares (scipy TRF) with the banded sparsity pattern; robust to the flat directions created by the task tie (A = w/alpha) and to kinks of the static block.""" from scipy.optimize import least_squares t0 = time.time() S = sparsity_pattern(model.T, band) res = least_squares(model.evaluate, u0, jac="2-point", jac_sparsity=S, method="trf", x_scale="jac", tr_solver="lsmr", ftol=1e-15, xtol=1e-15, gtol=1e-15, max_nfev=max_nfev, verbose=1 if verbose else 0, diff_step=1e-6) f = model.evaluate(res.x) nf = float(np.max(np.abs(f))) if verbose: print(f" trf: |res|_inf = {nf:.3e}, nfev={res.nfev}, {time.time() - t0:.0f}s, status={res.status}") return res.x, nf < tol, nf def 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): """Continuation in the size of the technology shock: g_pQ, g_eps and g_alpha are scaled by lambda (alpha_shock_year scenarios scale the log-distance to alpha_inf instead).""" base = dict(g_pQ=p.g_pQ, g_eps=p.g_eps, g_alpha=p.g_alpha) u, ok, hist_all = u0, False, [] for lam in lambdas: pl = p.replace(g_pQ=lam * base["g_pQ"], g_eps=lam * base["g_eps"], g_alpha=lam * base["g_alpha"]) m = TwoClassModel(pl) if u is None: u = m.smart_guess() if verbose: print(f"[homotopy] lambda = {lam}") u, ok, hist = m.solve(u, tol=tol, verbose=verbose, **kw) hist_all.append((lam, ok, hist[-1])) return m, u, ok, hist_all def horizon_check(p, u_short, extra=50, **kw): """Re-solve with a longer horizon from the short solution and report the max relative change of the first p.T-extra periods of the capital and rate paths (truncation error).""" pl = p.replace(T=p.T + extra) ml = TwoClassModel(pl) U = u_short.reshape(p.T, 6) Ul = np.vstack([U, np.tile(U[-1], (extra, 1))]) # extend capital logs along the terminal growth rate g = U[-1, :4] - U[-2, :4] for j in range(extra): Ul[p.T + j, :4] = U[-1, :4] + g * (j + 1) ul, ok, _ = ml.solve(Ul.ravel(), **kw) Ks, rs, _ = ml.unpack(ul) Ks0, rs0, _ = TwoClassModel(p).unpack(u_short) n = p.T - extra dev = max(np.max(np.abs(np.log(Ks[k][1:n + 1] / Ks0[k][1:n + 1]))) for k in KINDS) devr = np.max(np.abs(rs[1:n + 1] - rs0[1:n + 1])) return dict(ok=ok, max_log_dev_K=dev, max_dev_r=devr, u_long=ul, model_long=ml)