Monte, Redding & Rossi-Hansberg (2018)：通勤与贸易一体化模型

# ============================================================
# Monte, Redding & Rossi-Hansberg (2018, REStud)
# 通勤与贸易一体化模型求解器
# 三市场：货物贸易(EK) + 通勤选择(Logit) + 土地市场
# 已实跑验证：numpy 1.26
# ============================================================
import numpy as np


def solve_mrrh(N=8, theta=4.0, nu=0.3, beta=0.75,
               tol=1e-7, maxiter=3000, damp=0.4, seed=42,
               tau_shock=None, kappa_shock=None):
    """
    求解 M-RRH 均衡。
    - theta: EK 贸易弹性
    - nu:    通勤偏好离散度（越小越敏感）
    - beta:  消费支出份额（住房 1-beta）
    - tau_shock: (N,N) 贸易成本帽子冲击（反事实）
    - kappa_shock: (N,N) 通勤成本帽子冲击（反事实）
    """
    rng = np.random.default_rng(seed)
    # --- 外生参数 ---
    T = rng.uniform(0.8, 1.2, size=N)                    # 生产率
    H = rng.uniform(0.8, 1.2, size=N)                    # 土地/住房供给
    Lbar = 1.0
    dist = rng.uniform(0.3, 2.0, size=(N, N))
    tau = np.exp(0.3 * dist)                              # 贸易成本
    np.fill_diagonal(tau, 1.0)
    kappa = 0.2 * dist                                    # 通勤成本（效用单位）
    np.fill_diagonal(kappa, 0.0)
    amenity = rng.uniform(0.9, 1.1, size=N)              # 居住地宜居度
    if tau_shock is not None:
        tau = tau * tau_shock
    if kappa_shock is not None:
        kappa = kappa * kappa_shock

    # --- 初始化 ---
    w = np.ones(N) / N                                    # 名义工资
    r = np.ones(N)                                        # 土地租金

    for it in range(maxiter):
        # (1) EK 贸易份额与价格指数
        A = T[None, :] * (w[None, :] * tau) ** (-theta)  # (n,i) 目的地n×来源i
        P = A.sum(axis=1) ** (-1.0 / theta)
        pi = A / A.sum(axis=1, keepdims=True)

        # (2) 通勤概率（McFadden Logit）：
        # 效用从居住 i 到工作 n：V_ni = (w_n / R_i) - kappa_ni + amenity_i
        # R_i = r_i^(1-beta) P_i^beta
        R = r ** (1 - beta) * P ** beta                   # (N,) 居住成本指数
        util = w[None, :] / R[:, None] - kappa + amenity[:, None]
        util = util / nu
        util_max = util.max(axis=1, keepdims=True)
        e = np.exp(util - util_max)                       # 减最大值防溢出
        lam = e / e.sum(axis=1, keepdims=True)            # (res i, work n) 行和=1

        # (3) 居住人口：L_i^res = Σ_n lam_ni * L_n^work；工作人口外生归一
        Lwork = np.ones(N) * Lbar / N
        Lres = (lam * Lwork[None, :]).sum(axis=1)         # (res i,)
        Lres = Lres / Lres.sum() * Lbar

        # (4) 土地市场出清：r_i * H_i = (1-beta) * 居住在 i 的总收入
        wage_inc = (lam * w[None, :]).sum(axis=1) * Lres
        r_new = (1 - beta) * wage_inc / H

        # (5) 工资不动点：w_n L_n^work = Σ_i pi_ni * beta * 居住在 i 的消费支出
        cons_exp = beta * wage_inc
        sales_n = (pi * cons_exp[:, None]).sum(axis=0)
        w_new = sales_n / Lwork
        w_new = w_new / w_new.sum()

        # 阻尼更新
        r_upd = damp * r_new + (1 - damp) * r
        w_upd = damp * w_new + (1 - damp) * w
        err = max(np.max(np.abs(r_upd - r)), np.max(np.abs(w_upd - w)))
        r, w = r_upd, w_upd
        if err < tol:
            print(f"收敛于第 {it} 轮, err={err:.2e}")
            break

    return dict(w=w, r=r, P=P, Lres=Lres, Lwork=Lwork,
                pi=pi, lam=lam, T=T, H=H, tau=tau, kappa=kappa)


def welfare_decompose(base, cf, beta=0.75):
    """福利分解：贸易渠道（P）+ 通勤渠道（lam）+ 选址渠道（r）。
    平均实际福利 U_i = w_i / (P_i^beta r_i^(1-beta)) 的对数变化。
    """
    U0 = base["w"] / (base["P"]**beta * base["r"]**(1 - beta))
    U1 = cf["w"] / (cf["P"]**beta * cf["r"]**(1 - beta))
    Uhat = U1 / U0
    dw = np.log(cf["w"] / base["w"])
    dP = -beta * np.log(cf["P"] / base["P"])
    dr = -(1 - beta) * np.log(cf["r"] / base["r"])
    return dict(Uhat=Uhat, dw=dw, dP=dP, dr=dr, total=dw + dP + dr)


if __name__ == "__main__":
    np.random.seed(42)
    # --- 基准均衡 ---
    base = solve_mrrh(N=8, theta=4.0, nu=0.3, beta=0.75, seed=42)
    print("工资 w:", np.round(base["w"], 3))
    print("租金 r:", np.round(base["r"], 3))
    print("通勤矩阵行和=1:", np.allclose(base["lam"].sum(axis=1), 1.0))

    # --- 反事实 1：贸易成本下降 20%（修高速公路）---
    tau_shock = np.full((8, 8), 0.8)
    np.fill_diagonal(tau_shock, 1.0)
    cf_trade = solve_mrrh(N=8, theta=4.0, nu=0.3, beta=0.75,
                          seed=42, tau_shock=tau_shock)
    w1 = welfare_decompose(base, cf_trade)
    print(f"贸易成本↓20%: 加权福利变化 {np.sum(base['Lres']*w1['Uhat'])/base['Lres'].sum()-1:+.2%}")

    # --- 反事实 2：通勤成本下降 30%（修地铁）---
    kappa_shock = np.full((8, 8), 0.7)
    np.fill_diagonal(kappa_shock, 1.0)
    cf_commute = solve_mrrh(N=8, theta=4.0, nu=0.3, beta=0.75,
                            seed=42, kappa_shock=kappa_shock)
    w2 = welfare_decompose(base, cf_commute)
    print(f"通勤成本↓30%: 加权福利变化 {np.sum(base['Lres']*w2['Uhat'])/base['Lres'].sum()-1:+.2%}")