1"""Static task block: output maximization over labor and compute allocations (Restrepo's work equation).23The competitive equilibrium of the static block is the solution of the concave program45 max_{L_i, Q_i >= 0} Y = Z * ( G(V(X_B({X_i}), K_N), T, E_Y) + sum_a gamma_a X_a )6 s.t. X_i = L_i + Q_i / alpha_i, sum_{i in Omega(s)} L_i = H_s, sum_i Q_i = Q,78and factor prices are the multipliers (envelope theorem): w_s = max_{i in Omega(s)} F_i,9A = max_i F_i / alpha_i, where F_i = dY/dX_i. Tasks with alpha_i = inf cannot use compute.1011Nests (zero weights switch a nest off, so the pure Restrepo benchmark is the case nu_N = nu_T = nu_E = 0):12 X_B = (sum_i nu_i X_i^rho_B)^(1/rho_B)13 V = (nu_N K_N^rho_N + (1-nu_N) X_B^rho_N)^(1/rho_N)14 G = (nu_T T^rho_T + nu_E E_Y^rho_T + nu_V V^rho_T)^(1/rho_T), nu_V = 1 - nu_T - nu_E15"""16from __future__ import annotations1718from dataclasses import dataclass19from typing import Optional2021import numpy as np22from scipy.optimize import minimize, fsolve2324from .ces import ces, ces_grad2526_XFLOOR = 1e-10272829@dataclass30class Technology:31 nu: np.ndarray # task weights (bottleneck), zeros for accessory32 gamma: np.ndarray # linear accessory coefficients, zeros for bottleneck33 skill: np.ndarray # skill index of each task34 rho_B: float35 nu_N: float36 rho_N: float37 nu_T: float38 nu_E: float39 rho_T: float4041 @property42 def n(self):43 return len(self.nu)4445 @property46 def nskill(self):47 return int(self.skill.max()) + 14849 # ----- nested aggregator and its marginal products -----50 def aggregate(self, X, KN, T, EY, Z):51 """Return dict with Y, XB, V, G and marginal values F_i = dY/dX_i, R_N, R_T, P_E."""52 X = np.maximum(np.asarray(X, dtype=float), _XFLOOR)53 XB = ces(self.nu, X, self.rho_B)54 if self.nu_N > 0:55 V = ces([self.nu_N, 1 - self.nu_N], [KN, XB], self.rho_N)56 gV = ces_grad([self.nu_N, 1 - self.nu_N], [KN, XB], self.rho_N)57 dV_dKN, dV_dXB = gV[0], gV[1]58 else:59 V, dV_dKN, dV_dXB = XB, 0.0, 1.060 nu_V = 1.0 - self.nu_T - self.nu_E61 if self.nu_T > 0 or self.nu_E > 0:62 wts = [self.nu_T, self.nu_E, nu_V]63 G = ces(wts, [T, EY, V], self.rho_T)64 gG = ces_grad(wts, [T, EY, V], self.rho_T)65 dG_dT, dG_dE, dG_dV = gG[0], gG[1], gG[2]66 else:67 G, dG_dT, dG_dE, dG_dV = V, 0.0, 0.0, 1.068 acc = float(np.sum(self.gamma * X))69 Y = Z * (G + acc)70 dXB = ces_grad(self.nu, X, self.rho_B)71 F = Z * (dG_dV * dV_dXB * dXB + self.gamma)72 return dict(Y=Y, XB=XB, V=V, G=G, F=F,73 RN=Z * dG_dV * dV_dKN, RT=Z * dG_dT, PE=Z * dG_dE)7475 # ----- the static program -----76 def solve(self, Q, H, KN, T, EY, Z, alpha, x0: Optional[np.ndarray] = None, tol=1e-14):77 """Solve the static allocation. Variables are scaled: l_i = L_i/H_{s(i)}, q_i = Q_i/Q in [0,1]."""78 n = self.n79 alpha = np.asarray(alpha, dtype=float)80 fin = np.isfinite(alpha)81 inv_alpha = np.where(fin, 1.0 / np.where(fin, alpha, 1.0), 0.0)82 H = np.asarray(H, dtype=float)83 Hs = H[self.skill] # supply of the skill owning each task84 nS = self.nskill85 use_Q = bool(fin.any()) and Q > 086 Qs = float(Q) if use_Q else 0.08788 def X_of(x):89 L = x[:n] * Hs90 Qa = x[n:] * Qs if use_Q else np.zeros(n)91 return L, Qa, L + Qa * inv_alpha9293 # objective scale: output at the interior starting point94 def _x0():95 l0 = np.array([1.0 / np.sum(self.skill == s) for s in self.skill])96 if use_Q:97 q0 = np.where(fin, 1.0 / fin.sum(), 0.0)98 return np.concatenate([l0, q0])99 return l0100101 if x0 is None or len(x0) != (2 * n if use_Q else n):102 x0 = _x0()103 scale = max(self.aggregate(X_of(_x0())[2], KN, T, EY, Z)["Y"], 1e-12)104105 def obj(x):106 return -self.aggregate(X_of(x)[2], KN, T, EY, Z)["Y"] / scale107108 def jac(x):109 F = self.aggregate(X_of(x)[2], KN, T, EY, Z)["F"]110 g = np.concatenate([F * Hs, F * inv_alpha * Qs]) if use_Q else F * Hs111 return -g / scale112113 cons = []114 for s in range(nS):115 m = (self.skill == s).astype(float)116 mm = np.concatenate([m, np.zeros(n)]) if use_Q else m117 cons.append({"type": "eq", "fun": (lambda x, mm=mm: 1.0 - mm @ x), "jac": (lambda x, mm=mm: -mm)})118 if use_Q:119 mq = np.concatenate([np.zeros(n), fin.astype(float)])120 cons.append({"type": "eq", "fun": (lambda x, mq=mq: 1.0 - mq @ x), "jac": (lambda x, mq=mq: -mq)})121122 bounds = [(0.0, 1.0)] * n123 if use_Q:124 bounds += [((0.0, 1.0) if fin[i] else (0.0, 0.0)) for i in range(n)]125 x0 = np.clip(x0, 0.0, 1.0)126 res = minimize(obj, x0, jac=jac, bounds=bounds, constraints=cons, method="SLSQP",127 options={"ftol": tol, "maxiter": 1000})128 x, nit = res.x, res.nit129 res2 = minimize(obj, x, jac=jac, bounds=bounds, constraints=cons, method="SLSQP",130 options={"ftol": tol, "maxiter": 1000})131 if res2.fun <= res.fun:132 x, res = res2.x, res2133 x = self._polish(x, n, use_Q, fin, inv_alpha, Hs, Qs, X_of, KN, T, EY, Z, nS)134 L, Qa, X = X_of(x)135 agg = self.aggregate(X, KN, T, EY, Z)136 F = agg["F"]137 w = np.array([F[self.skill == s].max() for s in range(nS)])138 A = float(np.max((F * inv_alpha)[fin])) if use_Q else 0.0139 out = dict(L=L, Q=Qa, X=X, w=w, A=A, Y=agg["Y"], XB=agg["XB"], V=agg["V"], G=agg["G"], F=F,140 RN=agg["RN"], RT=agg["RT"], PE=agg["PE"], x=x, success=bool(res.success), nit=nit)141 return out142143 def _polish(self, x, n, use_Q, fin, inv_alpha, Hs, Qs, X_of, KN, T, EY, Z, nS, tol_active=1e-7):144 """Newton polish of the KKT equalities on the active set identified by SLSQP.145146 Free variables: l_i with l_i > tol, q_i with q_i > tol; unknown prices w_s, A. Equations: F_i = w_s for147 free labor, F_i/alpha_i = A for free compute, adding-up constraints. The result does not depend on the148 SLSQP warm start (to machine precision), which keeps finite-difference Jacobians of the dynamic model clean.149 """150 freeL = np.where(x[:n] > tol_active)[0]151 freeQ = np.where(x[n:] > tol_active)[0] if use_Q else np.array([], dtype=int)152 nL, nQ = len(freeL), len(freeQ)153 if nL == 0:154 return x155156 def eqs(z):157 xx = x.copy()158 xx[freeL] = z[:nL]159 if use_Q:160 xx[n + freeQ] = z[nL:nL + nQ]161 w = z[nL + nQ:nL + nQ + nS]162 A = z[-1] if use_Q else 0.0163 F = self.aggregate(X_of(xx)[2], KN, T, EY, Z)["F"]164 r = [F[i] - w[self.skill[i]] for i in freeL]165 r += [F[i] * inv_alpha[i] - A for i in freeQ]166 r += [np.sum(xx[:n][self.skill == s]) - 1.0 for s in range(nS)]167 if use_Q:168 r.append(np.sum(xx[n:][fin]) - 1.0)169 return np.array(r)170171 F0 = self.aggregate(X_of(x)[2], KN, T, EY, Z)["F"]172 w0 = np.array([F0[self.skill == s].max() for s in range(nS)])173 z0 = np.concatenate([x[freeL], x[n + freeQ] if use_Q else [], w0, [np.max((F0 * inv_alpha)[fin])] if use_Q else []])174 # scale: prices can be tiny; solve in logs of prices175 try:176 z, info, ier, msg = fsolve(eqs, z0, full_output=True, xtol=1e-14)177 except Exception:178 return x179 # accept on the residual, not on fsolve's status flag (it complains about xtol at machine precision)180 if np.max(np.abs(eqs(z))) > 1e-11 * max(1.0, np.max(np.abs(z))) or np.any(z[:nL + nQ] < -1e-12) or np.any(z[:nL + nQ] > 1 + 1e-12):181 return x182 xx = x.copy()183 xx[freeL] = np.clip(z[:nL], 0.0, 1.0)184 if use_Q:185 xx[n + freeQ] = np.clip(z[nL:nL + nQ], 0.0, 1.0)186 # accept only if the polished point is at least as good and KKT-consistent (no inactive variable wants in)187 F = self.aggregate(X_of(xx)[2], KN, T, EY, Z)["F"]188 w = z[nL + nQ:nL + nQ + nS]189 A = z[-1] if use_Q else 0.0190 if np.any(F - w[self.skill] > 1e-9 * np.abs(w[self.skill]) + 1e-14):191 return x192 if use_Q and np.any((F * inv_alpha)[fin] - A > 1e-9 * abs(A) + 1e-14):193 return x194 return xx195196 def check_kkt(self, sol, H, Q, alpha, tol=1e-6):197 """Complementary slackness of the task allocation; returns max violation."""198 F, L, Qa, w, A = sol["F"], sol["L"], sol["Q"], sol["w"], sol["A"]199 alpha = np.asarray(alpha, dtype=float)200 viol = 0.0201 for i in range(self.n):202 s = self.skill[i]203 # labor used only where F_i = w_s204 if L[i] > tol * max(H[s], 1.0):205 viol = max(viol, abs(F[i] - w[s]) / max(w[s], 1e-12))206 else:207 viol = max(viol, max(F[i] - w[s], 0.0) / max(w[s], 1e-12))208 if np.isfinite(alpha[i]) and Q > 0:209 if Qa[i] > tol * max(Q, 1.0):210 viol = max(viol, abs(F[i] / alpha[i] - A) / max(A, 1e-12))211 else:212 viol = max(viol, max(F[i] / alpha[i] - A, 0.0) / max(A, 1e-12))213 return viol214215216def technology_from_params(p) -> Technology:217 ts = p.tasks218 return Technology(nu=np.array(ts.nu, dtype=float), gamma=np.array(ts.gamma, dtype=float),219 skill=np.array(ts.skill, dtype=int), rho_B=p.rho_B, nu_N=p.nu_N, rho_N=p.rho_N,220 nu_T=p.nu_T, nu_E=p.nu_E, rho_T=p.rho_T)221222223def restrepo_benchmark(nu, gamma, skill, rho_B, alpha, H, Q_grid, Z=1.0):224 """Pure endowment economy of Restrepo (2025): Y = Z (X_B + sum gamma_a X_a), compute an endowment Q.225226 Returns a list of static solutions along Q_grid (used for the replication of Propositions 1-5)."""227 tech = Technology(nu=np.asarray(nu, float), gamma=np.asarray(gamma, float), skill=np.asarray(skill, int),228 rho_B=rho_B, nu_N=0.0, rho_N=0.0, nu_T=0.0, nu_E=0.0, rho_T=0.0)229 out, x0 = [], None230 for Q in Q_grid:231 sol = tech.solve(Q, H, KN=0.0, T=0.0, EY=0.0, Z=Z, alpha=alpha, x0=x0)232 x0 = sol["x"]233 H_arr = np.asarray(H, float)234 wages_bill = float(np.sum(sol["w"] * H_arr))235 sol["labor_share"] = wages_bill / sol["Y"]236 sol["compute_share"] = sol["A"] * Q / sol["Y"]237 sol["Q_total"] = Q238 out.append(sol)239 return out240