TeX 53.2%
Python 46.6%
1"""WP18 — equilibrium computation.23* Stationary recursive competitive equilibrium (SRCE, spec B.10): unknowns (log p, log p~_r,4 estate lottery, transfer rebate) solved by damped Broyden with a finite-difference Jacobian5 for the price block and an identity block for the lottery/rebate fixed points. Residuals:6 relative excess demand on the owner-occupied and rental markets, lottery change, rebate7 change. Every result carries `equilibrium = "GE"`; `solve_stationary(..., fixed_prices=True)`8 only iterates the lottery at given prices and is labelled "PE (prices fixed)".9* Supply normalisation (`calibrate_supply`): the supply scales H_bar, Hr_bar are set so that10 the baseline clears at p = p_r = 1 (KMV-style normalisation); counterfactuals keep the scales.11* Perfect-foresight transition after an unanticipated (MIT) shock: sequence-space shooting on12 (log p_t, log p~_r,t)_{t<T} with a finite-difference Jacobian computed in parallel and13 Broyden updates; durable stocks H_t = (1 - delta_H) H_{t-1} + delta_H S(p_t).14* Static-expectations variant (Assumption A6 of the appendix): households use the stationary15 continuation values, prices clear period by period.16"""17from __future__ import annotations1819import json20import os21import time22from concurrent.futures import ProcessPoolExecutor23from dataclasses import dataclass, field24from multiprocessing import get_context25from pathlib import Path26from typing import Any2728import numpy as np2930from . import model as M31from .distributions import (Panel, PeriodRecord, estate_lottery, pool_records, simulate_period,32 simulate_stationary, stack_policies)33from .household_solver import AgePolicy, bellman_residuals, ownership_threshold, solve_lifecycle34from .model import Params, Prices, stationary_prices353637# --------------------------------------------------------------------------------------38# Logging39# --------------------------------------------------------------------------------------40class RunLog:41 """Structured JSON-lines log in results/logs/<name>.jsonl plus a console echo."""4243 def __init__(self, name: str, echo: bool = True):44 self.path = M.RESULTS_DIR / "logs" / f"{name}.jsonl"45 self.path.parent.mkdir(parents=True, exist_ok=True)46 self.echo = echo47 self.t0 = time.time()48 self.fh = open(self.path, "a")4950 def __call__(self, event: str, **kw):51 rec = dict(t=round(time.time() - self.t0, 2), event=event, **{k: _jsonable(v) for k, v in kw.items()})52 self.fh.write(json.dumps(rec) + "\n")53 self.fh.flush()54 if self.echo:55 print(f"[{rec['t']:8.1f}s] {event}: " + ", ".join(f"{k}={_fmt(v)}" for k, v in kw.items()), flush=True)5657 def close(self):58 self.fh.close()596061def _jsonable(v):62 if isinstance(v, np.ndarray):63 return v.tolist()64 if isinstance(v, (np.floating, np.integer)):65 return v.item()66 return v676869def _fmt(v):70 if isinstance(v, float):71 return f"{v:.4g}"72 if isinstance(v, np.ndarray):73 return np.array2string(v, precision=4, suppress_small=True)74 return str(v)757677# --------------------------------------------------------------------------------------78# Supply helpers79# --------------------------------------------------------------------------------------80def rental_price(P: Params, p: float, ptil_r: float) -> float:81 """p_r = sigma p~_r + (1 - sigma) p (spec B.6)."""82 return P.sigma_seg * ptil_r + (1.0 - P.sigma_seg) * p838485def excess_demands(P: Params, H_d: float, S_d: float, p: float, ptil_r: float,86 H_stock: float | None = None, Hr_stock: float | None = None) -> np.ndarray:87 """Relative excess demands (owner-occupied, rental) given demands and prices.8889 With stocks given (transition), the stocks replace the stationary supply curves. For90 sigma < 1 a share (1 - sigma) of rental demand is served by converted owner-occupied units.91 """92 sg = P.sigma_seg93 S_own = M.supply_stock(P, p, P.H_bar, P.eps) if H_stock is None else H_stock94 S_rent = M.supply_stock(P, p, P.Hr_bar, P.eps_r) if Hr_stock is None else Hr_stock # converted margin, priced at p95 if sg >= 1.0:96 S_seg = M.supply_stock(P, ptil_r, P.Hr_bar, P.eps_r) if Hr_stock is None else Hr_stock97 return np.array([(H_d - S_own) / S_own, (S_d - S_seg) / S_seg])98 if sg <= 0.0:99 tot = S_own + S_rent100 return np.array([(H_d + S_d - tot) / tot, np.log(ptil_r / p)])101 S_seg = M.supply_stock(P, ptil_r, sg * P.Hr_bar, P.eps_r) if Hr_stock is None else sg * Hr_stock102 conv = S_own + (1 - sg) * S_rent103 return np.array([(H_d + (1 - sg) * S_d - conv) / conv, (sg * S_d - S_seg) / S_seg])104105106# --------------------------------------------------------------------------------------107# Stationary evaluation108# --------------------------------------------------------------------------------------109@dataclass110class SSEval:111 pols: list[AgePolicy]112 rec: PeriodRecord113 panel: Panel114 prices: Prices115 H_d: float116 S_d: float117 lottery_new: tuple[np.ndarray, np.ndarray]118 rebate_new: float119 forced_share: float120 clip_share: float121122123def evaluate_ss(P: Params, p: float, ptil_r: float, lottery, rebate: float, panel0: Panel | None = None,124 burn: int | None = None) -> SSEval:125 p_r = rental_price(P, p, ptil_r)126 pr = stationary_prices(P, p, p_r, rebate)127 pols = solve_lifecycle(P, pr, lottery)128 recs, panel = simulate_stationary(P, pols, pr, panel0=panel0, burn=burn)129 rec = pool_records(recs)130 lot_new = estate_lottery(P, rec)131 # transfer tax revenue per age-1 household: tau_b * (mean gross gift at j_gift + mean gross estate at J)132 g_mean = float(rec.gift[rec.age == P.j_gift].mean()) if (rec.age == P.j_gift).any() else 0.0133 e_mean = float(rec.estate[rec.age == P.J].mean()) if (rec.age == P.J).any() else 0.0134 rebate_new = P.tau_b * (g_mean + e_mean)135 return SSEval(pols=pols, rec=rec, panel=panel, prices=pr, H_d=rec.housing_demand(P), S_d=rec.rental_demand(P),136 lottery_new=lot_new, rebate_new=rebate_new,137 forced_share=float((rec.flag == 1).mean()), clip_share=float((rec.flag == 2).mean()))138139140@dataclass141class SSResult:142 P: Params143 p: float144 ptil_r: float145 p_r: float146 rho: float147 lottery: tuple[np.ndarray, np.ndarray]148 rebate: float149 ev: SSEval150 residuals: np.ndarray151 converged: bool152 iterations: int153 history: list[dict]154 equilibrium: str155 time_sec: float156 bellman: dict = field(default_factory=dict)157 thresholds: dict = field(default_factory=dict)158159 @property160 def H(self) -> float:161 return self.ev.H_d162163 @property164 def Hr(self) -> float:165 return self.ev.S_d166167168@dataclass169class SSCore:170 """Light-weight stationary-equilibrium data needed by the transition solver (picklable)."""171 p: float172 ptil_r: float173 p_r: float174 rebate: float175 lottery: tuple[np.ndarray, np.ndarray]176 V_ss: list[tuple[np.ndarray, np.ndarray]]177 panel: Panel178 H: float179 Hr: float180 equilibrium: str181182183def ss_core(res: "SSResult") -> SSCore:184 P = res.P185 return SSCore(p=res.p, ptil_r=res.ptil_r, p_r=res.p_r, rebate=res.rebate, lottery=res.lottery,186 V_ss=[(pol.V_own, pol.V_rent) for pol in res.ev.pols], panel=res.ev.panel,187 H=M.supply_stock(P, res.p, P.H_bar, P.eps),188 Hr=M.supply_stock(P, res.ptil_r if P.sigma_seg > 0 else res.p, P.Hr_bar, P.eps_r),189 equilibrium=res.equilibrium)190191192def _unpack(P: Params, u: np.ndarray, K: int):193 p, ptil = float(np.exp(u[0])), float(np.exp(u[1]))194 pts = np.maximum(u[2:2 + K], 0.0)195 rebate = float(u[2 + K]) if P.tau_b > 0 else 0.0196 return p, ptil, (pts, np.full(K, 1.0 / K)), rebate197198199def _residual(P: Params, ev: SSEval, p: float, ptil: float, lottery, rebate: float, fixed_prices: bool) -> np.ndarray:200 ex = excess_demands(P, ev.H_d, ev.S_d, p, ptil)201 if fixed_prices:202 ex = np.zeros(2)203 lot_res = ev.lottery_new[0] - lottery[0]204 parts = [ex, lot_res]205 if P.tau_b > 0:206 parts.append(np.array([ev.rebate_new - rebate]))207 return np.concatenate(parts)208209210def illinois(fun, x0: float, slope: float, tol: float, max_iter: int = 14, x_step: float = 0.04,211 log=None, label: str = "") -> tuple[float, float, Any, int]:212 """Root of a decreasing scalar function in log-price space.213214 Secant step from the initial slope guess, then bracketing and Illinois regula falsi once a215 sign change is found. Returns (x, f(x), payload, n_evals). `fun` returns (f, payload).216 """217 xa, (fa, pa) = x0, fun(x0)218 n = 1219 if abs(fa) < tol:220 return xa, fa, pa, n221 # second point in the direction that should reduce f222 xb = xa - np.clip(fa / slope, -x_step, x_step)223 fb, pb = fun(xb); n += 1224 lo, flo, plo = (xa, fa, pa) if fa > 0 else (xb, fb, pb)225 hi, fhi, phi = (xb, fb, pb) if fa > 0 else (xa, fa, pa)226 # expand until bracketed (f(lo) > 0 > f(hi) with lo < hi for a decreasing function)227 k = 0228 while not (flo > 0 > fhi) and k < 8:229 k += 1230 if abs(fb) < tol:231 return xb, fb, pb, n232 slope_est = (fb - fa) / (xb - xa) if abs(xb - xa) > 1e-10 and (fb - fa) / (xb - xa) < 0 else slope233 xn = xb - np.clip(fb / slope_est, -2 * x_step, 2 * x_step)234 xa, fa, pa = xb, fb, pb235 xb, (fb, pb) = xn, fun(xn); n += 1236 if fb > 0:237 lo, flo, plo = xb, fb, pb238 else:239 hi, fhi, phi = xb, fb, pb240 best = (xb, fb, pb) if abs(fb) < abs(fa) else (xa, fa, pa)241 if not (flo > 0 > fhi):242 return best[0], best[1], best[2], n243 side = 0244 for _ in range(max_iter):245 if abs(best[1]) < tol:246 break247 xm = (lo * fhi - hi * flo) / (fhi - flo)248 xm = min(max(xm, lo + 0.05 * (hi - lo)), hi - 0.05 * (hi - lo))249 fm, pm = fun(xm); n += 1250 if abs(fm) < abs(best[1]):251 best = (xm, fm, pm)252 if fm > 0:253 lo, flo, plo = xm, fm, pm254 if side == 1:255 fhi *= 0.5256 side = 1257 else:258 hi, fhi, phi = xm, fm, pm259 if side == -1:260 flo *= 0.5261 side = -1262 if hi - lo < 1e-6:263 break264 return best[0], best[1], best[2], n265266267def clear_prices(P: Params, p0: float, ptil0: float, lottery, rebate: float, tol: float, log: RunLog | None,268 panel0: Panel | None = None, x_step: float = 0.03) -> tuple[float, float, SSEval, np.ndarray, int]:269 """Nested market clearing: rental market (inner, in p~_r) within the owner market (outer, in p)."""270 n_evals = [0]271 ptil_state = [ptil0]272273 def rental_cleared(logp):274 p = float(np.exp(logp))275 if P.sigma_seg <= 0.0:276 ev = evaluate_ss(P, p, p, lottery, rebate, panel0=panel0); n_evals[0] += 1277 ex = excess_demands(P, ev.H_d, ev.S_d, p, p)278 return ex[0], (p, ev, ex)279280 def inner(logpt):281 pt = float(np.exp(logpt))282 ev = evaluate_ss(P, p, pt, lottery, rebate, panel0=panel0); n_evals[0] += 1283 ex = excess_demands(P, ev.H_d, ev.S_d, p, pt)284 return ex[1], (pt, ev, ex)285 # slope guess: rental excess falls with ptil through supply (eps_r) and demand (~ -10)286 x, fx, (pt, ev, ex), _ = illinois(inner, np.log(ptil_state[0]), -(15.0 + P.eps_r), tol * 0.6, x_step=x_step)287 ptil_state[0] = pt288 return ex[0], (pt, ev, ex)289290 x, fx, (pt, ev, ex), _ = illinois(rental_cleared, np.log(p0), -(5.0 + P.eps), tol, x_step=x_step)291 p = float(np.exp(x))292 if log:293 log("clear_prices", p=p, ptil_r=pt, excess=ex, n_evals=n_evals[0])294 return p, pt, ev, ex, n_evals[0]295296297def clear_prices_broyden(P: Params, p0: float, ptil0: float, lottery, rebate: float, tol: float, log: RunLog | None,298 J0: np.ndarray | None = None, panel0: Panel | None = None, max_iter: int = 8):299 """2-D damped Broyden on (log p, log p~_r) from a good starting point (later lottery rounds).300 Returns (p, ptil, ev, ex, n_evals, J, success)."""301 n_evals = 0302303 def F(u):304 nonlocal n_evals305 p, pt = float(np.exp(u[0])), float(np.exp(u[1]))306 if P.sigma_seg <= 0.0:307 pt = p308 ev = evaluate_ss(P, p, pt, lottery, rebate, panel0=panel0); n_evals += 1309 ex = excess_demands(P, ev.H_d, ev.S_d, p, pt)310 if P.sigma_seg <= 0.0:311 ex = np.array([ex[0], 0.0])312 return ex, ev313314 u = np.array([np.log(p0), np.log(ptil0)])315 f, ev = F(u)316 if np.max(np.abs(f)) < tol:317 return float(np.exp(u[0])), float(np.exp(u[1])), ev, f, n_evals, J0, True318 if J0 is None:319 J0 = np.eye(2)320 for k in range(2 if P.sigma_seg > 0 else 1):321 du = np.zeros(2); du[k] = 0.01322 fk, _ = F(u + du)323 J0[:, k] = (fk - f) / 0.01324 if P.sigma_seg <= 0.0:325 J0[1, 1] = 1.0; J0[0, 1] = 0.0; J0[1, 0] = 0.0326 Jm = J0.copy()327 best = (float(np.max(np.abs(f))), u.copy(), f.copy(), ev)328 fails = 0329 for _ in range(max_iter):330 try:331 step = -np.linalg.solve(Jm, f)332 except np.linalg.LinAlgError:333 break334 step = np.clip(step, -0.05, 0.05)335 if P.sigma_seg <= 0.0:336 step[1] = step[0]337 lam = 1.0338 accepted = False339 for _ in range(3):340 u_new = u + lam * step341 f_new, ev_new = F(u_new)342 if np.linalg.norm(f_new) < np.linalg.norm(f):343 accepted = True344 break345 lam *= 0.5346 if not accepted:347 fails += 1348 if fails >= 2:349 break350 continue351 sv = u_new - u; yv = f_new - f352 if float(sv @ sv) > 1e-16:353 Jm = Jm + np.outer(yv - Jm @ sv, sv) / float(sv @ sv)354 u, f, ev = u_new, f_new, ev_new355 if float(np.max(np.abs(f))) < best[0]:356 best = (float(np.max(np.abs(f))), u.copy(), f.copy(), ev)357 if np.max(np.abs(f)) < tol:358 break359 _, u, f, ev = best360 ok = bool(np.max(np.abs(f)) < tol)361 if log:362 log("clear_prices_broyden", p=float(np.exp(u[0])), ptil_r=float(np.exp(u[1])), excess=f, n_evals=n_evals, success=ok)363 return float(np.exp(u[0])), float(np.exp(u[1])), ev, f, n_evals, Jm, ok364365366def solve_stationary(P: Params, p0: float = 1.0, ptil0: float = 1.0, lottery0=None, rebate0: float = 0.0,367 fixed_prices: bool = False, log: RunLog | None = None, name: str = "ss",368 diagnostics: bool = True, tol: float | None = None, max_iter: int | None = None,369 panel0: Panel | None = None) -> SSResult:370 """Stationary equilibrium: nested price clearing (Illinois regula falsi on each market) inside371 a damped fixed-point iteration on the estate lottery and the transfer rebate.372373 fixed_prices=True keeps (p0, ptil0) and only iterates the lottery/rebate fixed point: the374 result is labelled "PE (prices fixed)".375 """376 t0 = time.time()377 cfg_eq = P.cfg["equilibrium"]378 tol = float(cfg_eq["tol_excess"]) if tol is None else tol379 tol_lot = float(cfg_eq["tol_lottery"])380 max_outer = int(cfg_eq.get("max_outer", 10)) if max_iter is None else max_iter381 damp = float(cfg_eq["damping"])382 K = P.K_estate383 if lottery0 is None:384 lottery0 = (np.array([0.0, 0.0, 0.05, 0.4, 1.0])[:K] if K == 5 else np.linspace(0, 1, K), np.full(K, 1.0 / K))385 log = log or RunLog(name)386 lottery, rebate = (np.array(lottery0[0], dtype=float), lottery0[1]), rebate0387 p, ptil = p0, ptil0388 history = []389 converged = False390 total_evals = 0391 it = 0392 ev = None393 Jm = None394 for it in range(1, max_outer + 1):395 if fixed_prices:396 ev = evaluate_ss(P, p, ptil, lottery, rebate, panel0=panel0); total_evals += 1397 ex = np.zeros(2)398 else:399 # nested regula falsi (robust to the Monte-Carlo noise in excess-demand differences);400 # later rounds start from the previous prices with small initial steps401 p, ptil, ev, ex, n_ev = clear_prices(P, p, ptil, lottery, rebate, tol, log, panel0=panel0,402 x_step=0.03 if it == 1 else 0.008)403 total_evals += n_ev404 lot_new = ev.lottery_new[0]405 lot_change = float(np.max(np.abs(lot_new - lottery[0])) / max(1.0, float(np.max(lot_new))))406 reb_change = float(abs(ev.rebate_new - rebate))407 history.append(dict(iter=it, excess_owner=float(ex[0]), excess_rental=float(ex[1]), lottery_change=lot_change,408 rebate_change=reb_change, p=p, ptil_r=ptil, evals=total_evals))409 log("ss_outer", iter=it, p=p, ptil_r=ptil, excess=ex, lottery=lot_new, lottery_change=lot_change,410 rebate=ev.rebate_new, rebate_change=reb_change, evals=total_evals)411 if lot_change < tol_lot and reb_change < tol_lot and float(np.max(np.abs(ex))) < tol:412 converged = True413 break414 lottery = ((1.0 - damp) * lottery[0] + damp * lot_new, lottery[1])415 rebate = (1.0 - damp) * rebate + damp * ev.rebate_new416 f = np.concatenate([ex, lot_new - lottery[0], [ev.rebate_new - rebate] if P.tau_b > 0 else []])417 lot = (lot_new, lottery[1]) if converged else lottery418 res = SSResult(P=P, p=p, ptil_r=ptil, p_r=rental_price(P, p, ptil), rho=ev.prices.rho, lottery=lot, rebate=rebate,419 ev=ev, residuals=f, converged=converged, iterations=it, history=history,420 equilibrium="PE (prices fixed)" if fixed_prices else "GE", time_sec=time.time() - t0)421 if diagnostics:422 res.bellman = bellman_residuals(P, ev.pols, ev.prices, lot, n_points=400)423 res.thresholds = {j: ownership_threshold(P, ev.pols[j - 1], ev.prices, j) for j in (1, 2, 3)}424 log("ss_diagnostics", converged=converged, bellman=res.bellman, forced=ev.forced_share, clipped=ev.clip_share)425 log("ss_done", p=p, p_r=res.p_r, rho=res.rho, converged=converged, outer_iterations=it, evals=total_evals,426 resid=f, sec=res.time_sec)427 return res428429430def lottery_pts(u, K):431 return u[2:2 + K]432433434def calibrate_supply(P: Params, log: RunLog | None = None) -> tuple[float, float, SSResult]:435 """Set H_bar, Hr_bar so that the economy clears at p = p~_r = p_ref (lottery converged)."""436 P.H_bar, P.Hr_bar = 1.0, 1.0 # placeholders, prices fixed437 res = solve_stationary(P, p0=P.p_ref, ptil0=P.p_ref, fixed_prices=True, log=log, name="calibrate_supply",438 diagnostics=False)439 P.H_bar, P.Hr_bar = res.ev.H_d, res.ev.S_d440 return P.H_bar, P.Hr_bar, res441442443def price_scan(P: Params, p_values: np.ndarray, lottery, log: RunLog | None = None) -> list[dict]:444 """Owner-market excess demand along a grid of p, with p~_r cleared (1-D secant) and the445 lottery held at the equilibrium one: used to count equilibria (multiplicity search)."""446 out = []447 for p in p_values:448 ptil = p449 f_prev = None; x_prev = None450 for it in range(12):451 ev = evaluate_ss(P, float(p), float(ptil), lottery, 0.0)452 ex = excess_demands(P, ev.H_d, ev.S_d, float(p), float(ptil))453 if abs(ex[1]) < 2e-4 or P.sigma_seg <= 0:454 break455 x = np.log(ptil)456 if f_prev is None:457 x_new = x - ex[1] / (-(1.0 + P.eps_r)) # elasticity guess458 else:459 slope = (ex[1] - f_prev) / (x - x_prev) if abs(x - x_prev) > 1e-9 else -(1 + P.eps_r)460 x_new = x - ex[1] / slope if abs(slope) > 1e-6 else x461 x_prev, f_prev = x, ex[1]462 ptil = float(np.exp(np.clip(x_new, x - 0.3, x + 0.3)))463 own_share = float(ev.rec.owner1().mean())464 out.append(dict(p=float(p), ptil_r=float(ptil), excess_owner=float(ex[0]), excess_rental=float(ex[1]),465 H_d=ev.H_d, S_d=ev.S_d, ownership=own_share))466 if log:467 log("price_scan", p=float(p), ptil_r=float(ptil), excess_owner=float(ex[0]), ownership=own_share)468 return out469470471# --------------------------------------------------------------------------------------472# Transition after an MIT shock473# --------------------------------------------------------------------------------------474@dataclass475class Shock:476 kind: str = "inflation" # inflation | real_rate | none477 pi_u: float = 0.10478 persistence: int = 1 # number of periods with elevated inflation / rate479 contract: str = "frm"480 phi_r: float = 0.0 # real-rate response r_t = r + phi_r * pi_u in shock periods481 dr: float = 0.0 # real-rate shock (period units) for kind == real_rate482 eta_rent: float | None = None483 eta_wage: float | None = None484485486def shock_paths(P: Params, sh: Shock, T: int) -> dict:487 """Exogenous paths: inflation pi_t (between t and t+1 for the deflation of balances, i.e.488 pi_next at t), real gross rate R_t, nominal rate i_t, surprise at t = 0."""489 pibar, R = P.pibar, P.R490 pi_real = np.full(T + 2, pibar) # realised inflation INTO period t (deflates balances from t-1 to t)491 Rt = np.full(T + 2, R)492 if sh.kind == "inflation":493 for k in range(sh.persistence):494 pi_real[k] = (1.0 + pibar) * (1.0 + sh.pi_u) - 1.0495 if sh.phi_r != 0.0:496 for k in range(sh.persistence):497 Rt[k] = R * (1.0 + sh.phi_r * sh.pi_u) if k == 0 else R * (1.0 + sh.phi_r * sh.pi_u)498 elif sh.kind == "real_rate":499 for k in range(sh.persistence):500 Rt[k] = R * (1.0 + sh.dr)501 pi_u = sh.pi_u if sh.kind == "inflation" else 0.0502 # nominal rate quoted at t for (t, t+1): i_t = R_t (1 + pi_{t+1}) - 1 (+ spread)503 i_new = Rt[:T + 1] * (1.0 + pi_real[1:T + 2]) - 1.0 + P.spread504 # realised real deposit return at t: (1 + i_{t-1}) / (1 + pi_t); at t = 0 the pre-shock i* applies505 R_a = np.empty(T + 1)506 R_a[0] = (1.0 + P.i_nom) / (1.0 + pi_real[0]) * (1.0 + P.dep_premium)507 R_a[1:] = (1.0 + i_new[:T]) / (1.0 + pi_real[1:T + 1]) * (1.0 + P.dep_premium)508 reset, within = M.contract_flags(sh.contract)509 balance_factor = (1.0 + pibar) / (1.0 + pi_real[0]) if not within else 1.0510 return dict(pi_real=pi_real, Rt=Rt, i_new=i_new, R_a=R_a, pi_u=pi_u, reset=reset, within=within,511 balance_factor=balance_factor)512513514def frm_rate_by_n(P: Params, paths: dict, t: int) -> np.ndarray:515 """Contract rate of an FRM with n remaining payments at date t: the rate at origination516 t0 = t - (N - n), equal to the stationary rate if t0 < 0."""517 out = np.empty(P.N)518 for n in range(1, P.N + 1):519 t0 = t - (P.N - n)520 out[n - 1] = P.i_nom if t0 < 0 else float(paths["i_new"][t0])521 return out522523524def build_price_path(P: Params, paths: dict, sh: Shock, T: int, logp: np.ndarray, logptil: np.ndarray,525 ss: SSCore) -> list[Prices]:526 """Prices objects for t = 0..T (t = T uses stationary prices as the terminal condition)."""527 p = np.concatenate([np.exp(logp), [ss.p]])528 ptil = np.concatenate([np.exp(logptil), [ss.ptil_r]])529 p_r = np.array([rental_price(P, p[t], ptil[t]) for t in range(T + 1)])530 eta_rent = P.eta_rent if sh.eta_rent is None else sh.eta_rent531 eta_wage = P.eta_wage if sh.eta_wage is None else sh.eta_wage532 out = []533 for t in range(T + 1):534 p_r_next = p_r[t + 1] if t + 1 <= T else ss.p_r535 rho = M.rent_from_price(P, p_r[t], p_r_next, float(paths["Rt"][t]))536 pi_u = float(paths["pi_u"]) if (t == 0 and sh.kind == "inflation") else 0.0537 out.append(Prices(p=float(p[t]), p_r=float(p_r[t]), rho=rho, R_a=float(paths["R_a"][t]), R_next=float(paths["Rt"][t]),538 i_new=float(paths["i_new"][t]), i_q=float(paths["i_new"][t]) + P.stress,539 pi_next=float(paths["pi_real"][t + 1]), pi_u=pi_u,540 i_c_by_n=frm_rate_by_n(P, paths, t), reset=bool(paths["reset"]),541 wage_factor=(1.0 + pi_u) ** (-(1.0 - eta_wage)) if pi_u else 1.0,542 rent_rigid_factor=(1.0 + pi_u) ** (-(1.0 - eta_rent)) if pi_u else 1.0,543 p_next=float(p[t + 1]) if t + 1 <= T else ss.p, transfer_rebate=ss.rebate))544 return out545546547def solve_households_path(P: Params, prices: list[Prices], ss: SSCore, T: int,548 static_expectations: bool = False) -> list[list[AgePolicy]]:549 """Backward solve over dates t = T-1..0; pols_by_t[t][j-1]. Terminal values at T = SS."""550 V_next = list(ss.V_ss) # V_{j, T} for all j551 pols_by_t: list[list[AgePolicy] | None] = [None] * T552 for t in range(T - 1, -1, -1):553 # continuation for age j at t is V_{j+1, t+1}554 V_next_by_age = [None] * P.J555 for j in range(1, P.J):556 V_next_by_age[j - 1] = V_next[j] if not static_expectations else ss.V_ss[j]557 pols = solve_lifecycle(P, [prices[t]] * P.J, ss.lottery, V_next_by_age=V_next_by_age)558 pols_by_t[t] = pols559 V_next = [(pol.V_own, pol.V_rent) for pol in pols]560 return pols_by_t # type: ignore[return-value]561562563def simulate_path(P: Params, pols_by_t, prices: list[Prices], panel0: Panel, T: int, balance_factor: float,564 seed_offset: int = 100000) -> list[PeriodRecord]:565 panel = panel0.copy()566 recs = []567 for t in range(T):568 S = stack_policies(pols_by_t[t])569 rec = simulate_period(P, S, prices[t], panel, (P.seed, seed_offset + t),570 balance_factor=balance_factor if t == 0 else 1.0)571 recs.append(rec)572 return recs573574575def stocks_along_path(P: Params, p: np.ndarray, ptil: np.ndarray, H0: float, Hr0: float) -> tuple[np.ndarray, np.ndarray]:576 T = len(p)577 H = np.empty(T); Hr = np.empty(T)578 h, hr = H0, Hr0579 for t in range(T):580 h = M.stock_next(P, h, p[t], P.H_bar, P.eps)581 pr_seg = ptil[t] if P.sigma_seg > 0 else p[t]582 hr = M.stock_next(P, hr, pr_seg, P.Hr_bar, P.eps_r) if P.sigma_seg >= 1.0 else \583 M.stock_next(P, hr, pr_seg, P.Hr_bar, P.eps_r)584 H[t], Hr[t] = h, hr585 return H, Hr586587588def transition_residual(P: Params, ss: SSCore, sh: Shock, paths: dict, T: int, u: np.ndarray,589 static_expectations: bool = False, return_all: bool = False):590 logp, logptil = u[:T], u[T:2 * T]591 prices = build_price_path(P, paths, sh, T, logp, logptil, ss)592 pols_by_t = solve_households_path(P, prices, ss, T, static_expectations=static_expectations)593 recs = simulate_path(P, pols_by_t, prices, ss.panel, T, float(paths["balance_factor"]))594 p = np.exp(logp); ptil = np.exp(logptil)595 H, Hr = stocks_along_path(P, p, ptil, ss.H, ss.Hr)596 res = np.empty(2 * T)597 for t in range(T):598 ex = excess_demands(P, recs[t].housing_demand(P), recs[t].rental_demand(P), float(p[t]), float(ptil[t]),599 H_stock=float(H[t]), Hr_stock=float(Hr[t]))600 res[t], res[T + t] = ex601 if return_all:602 return res, dict(prices=prices, pols_by_t=pols_by_t, recs=recs, H=H, Hr=Hr)603 return res604605606# ---- parallel Jacobian workers -------------------------------------------------------607_CTX: dict[str, Any] = {}608609610def _init_worker(payload, n_threads: int):611 import numba612 numba.set_num_threads(max(1, n_threads))613 _CTX.update(payload)614615616def _worker_column(args):617 k, h, u = args618 P, ss, sh, paths, T, static = (_CTX[n] for n in ("P", "ss", "sh", "paths", "T", "static"))619 du = np.zeros_like(u); du[k] = h620 return k, transition_residual(P, ss, sh, paths, T, u + du, static_expectations=static)621622623def _fd_jacobian_parallel(P, ss, sh, paths, T, u, f0, h, static, n_workers: int, log: RunLog | None):624 n = len(u)625 J = np.empty((n, n))626 payload = dict(P=P, ss=ss, sh=sh, paths=paths, T=T, static=static)627 n_workers = max(1, min(n_workers, n))628 threads = max(1, (os.cpu_count() or 4) // n_workers)629 if n_workers == 1:630 for k in range(n):631 _, fk = _worker_column_local(P, ss, sh, paths, T, static, k, h, u)632 J[:, k] = (fk - f0) / h633 return J634 ctx = get_context("spawn")635 with ProcessPoolExecutor(max_workers=n_workers, mp_context=ctx, initializer=_init_worker,636 initargs=(payload, threads)) as ex:637 for k, fk in ex.map(_worker_column, [(k, h, u) for k in range(n)]):638 J[:, k] = (fk - f0) / h639 if log and k % 8 == 0:640 log("jacobian_column", k=k)641 return J642643644def _worker_column_local(P, ss, sh, paths, T, static, k, h, u):645 du = np.zeros_like(u); du[k] = h646 return k, transition_residual(P, ss, sh, paths, T, u + du, static_expectations=static)647648649@dataclass650class TransitionResult:651 P: Params652 shock: Shock653 T: int654 p: np.ndarray655 ptil_r: np.ndarray656 p_r: np.ndarray657 rho: np.ndarray658 H: np.ndarray659 Hr: np.ndarray660 residuals: np.ndarray661 converged: bool662 iterations: int663 recs: list[PeriodRecord]664 prices: list[Prices]665 paths: dict666 equilibrium: str667 time_sec: float668 history: list[dict]669670671def solve_transition(P: Params, ss_in: "SSResult | SSCore", sh: Shock, T: int | None = None, mode: str = "GE",672 log: RunLog | None = None, name: str = "transition", n_workers: int | None = None,673 u0: np.ndarray | None = None) -> TransitionResult:674 """mode: "GE" (perfect foresight, prices clear), "PE" (prices fixed at the SRCE — labelled675 partial equilibrium), "static" (A6: stationary continuation values, prices clear).676 u0: optional warm start, stacked (log p_t, log p~_r,t) for t < T (e.g. the path of a677 converged neighbouring experiment); the Jacobian is then computed at u0."""678 t0 = time.time()679 ss = ss_in if isinstance(ss_in, SSCore) else ss_core(ss_in)680 cfg_tr = P.cfg["transition"]681 T = int(cfg_tr["T"]) if T is None else T682 tol = float(cfg_tr["tol"]); max_iter = int(cfg_tr["max_iter"]); h = float(cfg_tr["jac_step"])683 log = log or RunLog(name)684 paths = shock_paths(P, sh, T)685 static = mode == "static"686 u = np.concatenate([np.full(T, np.log(ss.p)), np.full(T, np.log(ss.ptil_r))])687 if u0 is not None and mode != "PE" and len(u0) == 2 * T:688 u = np.asarray(u0, dtype=float).copy()689 log("transition_warm_start", p0=float(np.exp(u[0])))690 history = []691 f, extra = transition_residual(P, ss, sh, paths, T, u, static_expectations=static, return_all=True)692 log("transition_eval", iter=0, max_resid=float(np.max(np.abs(f))), mode=mode)693 converged = False694 it = 0695 if mode == "PE":696 converged = True697 else:698 # --- damped tatonnement pre-phase for large shocks: move each period's log prices in the699 # direction of its excess demand until the path is within the Jacobian's linear range700 kappa = 0.15701 n_tat = 0702 while float(np.max(np.abs(f))) > 0.01 and n_tat < 80:703 u_try = u + kappa * f704 f_try, extra_try = transition_residual(P, ss, sh, paths, T, u_try, static_expectations=static, return_all=True)705 if np.linalg.norm(f_try) < np.linalg.norm(f):706 u, f, extra = u_try, f_try, extra_try707 kappa = min(kappa * 1.2, 0.5)708 else:709 kappa *= 0.5710 if kappa < 1e-3:711 break712 n_tat += 1713 log("transition_tatonnement", iter=n_tat, max_resid=float(np.max(np.abs(f))), kappa=kappa, p0=float(np.exp(u[0])))714 n_workers = n_workers or max(1, min(16, (os.cpu_count() or 4) - 2))715 Jm = _fd_jacobian_parallel(P, ss, sh, paths, T, u, f, h, static, n_workers, log)716 log("jacobian_done", n=len(u), sec=time.time() - t0)717 fails = 0718 refreshes = 0719 best = (float(np.max(np.abs(f))), u.copy(), f.copy(), extra)720 for it in range(1, max_iter + 1):721 history.append(dict(iter=it - 1, max_resid=float(np.max(np.abs(f)))))722 if np.max(np.abs(f)) < tol:723 converged = True724 break725 try:726 step = -np.linalg.solve(Jm, f)727 except np.linalg.LinAlgError:728 step = -0.5 * f729 step = np.clip(step, -0.2, 0.2)730 lam = 1.0731 accepted = False732 for _ in range(3):733 u_new = u + lam * step734 f_new, extra_new = transition_residual(P, ss, sh, paths, T, u_new, static_expectations=static, return_all=True)735 if np.linalg.norm(f_new) < np.linalg.norm(f):736 accepted = True737 break738 lam *= 0.5739 if accepted:740 s_ = u_new - u; yv = f_new - f741 if float(s_ @ s_) > 1e-16:742 Jm = Jm + np.outer(yv - Jm @ s_, s_) / float(s_ @ s_)743 u, f, extra = u_new, f_new, extra_new744 fails = 0745 else:746 fails += 1747 if fails >= 2:748 if refreshes >= 1:749 log("transition_stall", msg="no improvement after a Jacobian refresh: stopping at the best iterate (noise floor)")750 break751 u, f, extra = best[1].copy(), best[2].copy(), best[3]752 Jm = _fd_jacobian_parallel(P, ss, sh, paths, T, u, f, 1.5 * h, static, n_workers, log)753 log("jacobian_refresh", iter=it)754 refreshes += 1755 fails = 0756 if float(np.max(np.abs(f))) < best[0]:757 best = (float(np.max(np.abs(f))), u.copy(), f.copy(), extra)758 log("transition_eval", iter=it, max_resid=float(np.max(np.abs(f))), lam=lam, accepted=accepted, p0=float(np.exp(u[0])))759 u, f, extra = best[1], best[2], best[3]760 if not converged and np.max(np.abs(f)) < tol:761 converged = True762 p = np.exp(u[:T]); ptil = np.exp(u[T:])763 p_r = np.array([pr.p_r for pr in extra["prices"][:T]]); rho = np.array([pr.rho for pr in extra["prices"][:T]])764 label = {"GE": "GE", "PE": "PE (prices fixed)", "static": "static expectations (A6)"}[mode]765 if mode == "GE" and not converged:766 label = "GE NOT CONVERGED (max residual %.2e)" % float(np.max(np.abs(f)))767 res = TransitionResult(P=P, shock=sh, T=T, p=p, ptil_r=ptil, p_r=p_r, rho=rho, H=extra["H"], Hr=extra["Hr"],768 residuals=f, converged=converged, iterations=it, recs=extra["recs"], prices=extra["prices"],769 paths=paths, equilibrium=label, time_sec=time.time() - t0, history=history)770 log("transition_done", mode=mode, converged=converged, max_resid=float(np.max(np.abs(f))), iterations=it,771 sec=res.time_sec, p_path=p[:6], rho_path=rho[:6])772 return res773