Allen-Arkolakis (2014)：空间经济统一框架与网络结构

"""
Allen & Arkolakis (2014) 空间经济统一框架
- 地点网络, 一般化引力 pi_ni, 迁移选址 L_i ∝ V_i^eps, 土地出清
- 网络不动点迭代: 工资 + 人口
"""
import numpy as np
np.random.seed(20240912)


def solve_spatial_network(N, A, B, tau, theta=5.0, eps=1.5,
                          beta=0.75, Lbar=1.0, tol=1e-10, damp=0.3):
    L = np.ones(N) / N
    w = np.ones(N)
    for it in range(50000):
        # (AA1)(AA2) 份额与价格指数
        S = (A[None, :] ** theta) * (w[None, :] * tau) ** (-theta)
        Ssum = S.sum(axis=1, keepdims=True)
        pi = S / Ssum
        P = Ssum.ravel() ** (-1.0 / theta)
        # (AA3) 贸易平衡 -> 隐含工资
        w_implied = pi.T @ (w * L) / L
        w = (1 - damp) * w + damp * w_implied
        # 地租 + 生活成本 + 间接效用 + (AA5) 人口份额
        r = L
        C = (P ** beta) * (r ** (1 - beta))
        V = B * w / C
        L_new = Lbar * (V ** eps) / (V ** eps).sum()
        L = (1 - damp) * L + damp * L_new
        if it > 200 and np.max(np.abs(w_implied / w - 1)) < tol \
                   and np.max(np.abs(L_new / L - 1)) < tol:
            break
    S = (A[None, :] ** theta) * (w[None, :] * tau) ** (-theta)
    Ssum = S.sum(axis=1, keepdims=True)
    pi = S / Ssum; P = Ssum.ravel() ** (-1.0 / theta)
    r = L; C = P ** beta * r ** (1 - beta); V = B * w / C
    return dict(w=w, L=L, P=P, r=r, pi=pi, V=V)


def welfare(base, cf):
    """(AA6) 人口加权总福利变化."""
    W0 = (base["L"] * np.log(base["V"])).sum()
    W1 = (cf["L"] * np.log(cf["V"])).sum()
    return np.exp(W1 - W0) - 1.0


# ---------------- N=6 网络示例 ----------------
N = 6
pos = np.linspace(0, 5, N)
dist = np.abs(pos[None, :] - pos[:, None])
tau = np.exp(0.6 * dist); np.fill_diagonal(tau, 1.0)
A = np.array([0.8, 1.0, 1.6, 1.5, 0.9, 0.7])   # 核心 2,3
B = np.array([1.3, 1.0, 0.8, 0.85, 1.0, 1.2])
theta, eps, beta = 5.0, 1.5, 0.75

base = solve_spatial_network(N, A, B, tau, theta, eps, beta)
print("工资 w_i   :", np.round(base["w"], 3))
print("人口 L_i   :", np.round(base["L"], 3))
print("贸易行和=1 :", np.round(base["pi"].sum(axis=1), 4))

# 反事实: 核心 2,3 对外连通成本降 40%
tau1 = tau.copy()
for c in (2, 3):
    tau1[c, :] = np.where(np.arange(N) != c, tau[c, :] * 0.6, 1.0)
    tau1[:, c] = np.where(np.arange(N) != c, tau[:, c] * 0.6, 1.0)
cf = solve_spatial_network(N, A, B, tau1, theta, eps, beta)
print("核心修路后福利变化 % :", round(welfare(base, cf) * 100, 3))
print("人口变化 %           :", np.round((cf["L"]/base["L"]-1)*100, 2))
# ACR 充分统计量核对
acr = (np.diag(cf["pi"]).mean()/np.diag(base["pi"]).mean()) ** (-1/theta) - 1
print("ACR 近似 %          :", round(acr * 100, 3))