Lecture 16, Notebook 03: Stochastic CDICE-DEQN with AR(1) productivity shocks

RUN_MODE = "smoke"  # one of: "smoke", "teaching", "production"
SEED = 0

import os, time, math
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt

np.random.seed(SEED); tf.random.set_seed(SEED)
tf.keras.backend.set_floatx('float32')
plt.rcParams['font.size'] = 13
plt.rcParams['figure.dpi'] = 110

print('TF', tf.__version__)

class P:
    Tstep = 1.0
    vartheta = 0.015
    # Preferences / production
    rho = 0.015          # discount rate
    psi = 0.68965517     # IES (CRRA exponent = 1/psi)
    alpha = 0.30
    delta = 0.10
    # Population
    L0 = 7403.0;     Linfty = 11500.0;     deltaL = 0.0268
    # TFP trend
    A0hat = 0.010295; gA0hat = 0.0217;    deltaA = 0.005
    # Carbon intensity
    sigma0 = 0.0000955592; gSigma0 = -0.0152; deltaSigma = 0.001
    # Mitigation cost
    theta2 = 2.6;    pback = 0.55;        gback = 0.005;       c2co2 = 3.666
    # Land emissions / non-CO2 forcing
    ELand0 = 0.00070922; deltaLand = 0.023
    fex0 = 0.5;      fex1 = 1.0;           Tyears = 85.0
    # Damage function:  Omega(T) = pi1*T^pow1 + pi2*T^pow2
    pi1 = 0.0; pi2 = 0.00236; pow1 = 1.0; pow2 = 2.0
    # Carbon mass cycle (mmm calibration)
    b12_ = 0.054;    b23_ = 0.0082
    MATeq = 0.607;   MUOeq = 0.489;        MLOeq = 1.281
    # Temperature
    c1_ = 0.137;     c3_ = 0.73;           c4_ = 0.00689
    f2xco2 = 3.45;   t2xco2 = 3.25;        MATbase = 0.607
    # Initial conditions (year 2015)
    k0 = 2.926; MAT0 = 0.851; MUO0 = 0.628; MLO0 = 1.323
    TAT0 = 1.1; TOC0 = 0.27;  tau0 = 0.0
    z0  = 0.0
    # Cai–Lontzek (2019) §V.A productivity-shock calibration
    rho_z   = 0.95
    sigma_z = 0.0125
    # Derived
    b12 = Tstep * b12_;    b23 = Tstep * b23_
    b21 = MATeq / MUOeq * b12_ * Tstep
    b32 = MUOeq / MLOeq * b23_ * Tstep
    c1 = Tstep * c1_;      c1c3 = Tstep * c1_ * c3_
    c1f = Tstep * c1_ * f2xco2 / t2xco2
    c4 = Tstep * c4_

for _a in list(vars(P)):
    if not _a.startswith('_'):
        v = getattr(P, _a)
        if isinstance(v, (int, float)):
            setattr(P, _a, np.float32(v))

LOG2 = np.float32(math.log(2.0))
ONE_MINUS_DELTA_TSTEP = np.float32((1.0 - float(P.delta)) ** float(P.Tstep))

# Stationary std of z (useful below for sampling realistic initial conditions):
print(f'rho_z = {float(P.rho_z):.3f},  sigma_z = {float(P.sigma_z):.4f}')
print(f'stationary std(z) = sigma_z / sqrt(1 - rho_z^2) = '
      f'{float(P.sigma_z) / np.sqrt(1 - float(P.rho_z)**2):.4f}')

GH_Q = 5
_gh_x, _gh_w = np.polynomial.hermite.hermgauss(GH_Q)
GH_NODES   = tf.constant((np.sqrt(2.0) * _gh_x).astype('float32'))  # eps quadrature points (keep float32)
GH_WEIGHTS = tf.constant((_gh_w / np.sqrt(np.pi)).astype('float32'))  # weights, sum to 1
print('GH nodes (eps): ', GH_NODES.numpy())
print('GH weights:     ', GH_WEIGHTS.numpy(), ' (sum =', float(tf.reduce_sum(GH_WEIGHTS)), ')')

def tau2t(tau):
    return -tf.math.log(1.0 - tau) / P.vartheta
def t2tau(t):
    return 1.0 - tf.math.exp(-P.vartheta * t)
def tau2tauplus(tau):
    # one annual step in tau-space
    t = tau2t(tau)
    return t2tau(t + P.Tstep)

def population(t):
    return P.L0 + (P.Linfty - P.L0) * (1.0 - tf.exp(-P.Tstep * P.deltaL * t))

def tfp_trend(t):
    return P.A0hat * tf.exp((P.Tstep * P.gA0hat) *
        (1.0 - tf.exp(-P.Tstep * P.deltaA * t)) / (P.Tstep * P.deltaA))

def tfp(t, z):
    # effective stochastic TFP: trend * exp(z)
    return tfp_trend(t) * tf.exp(z)

def gr_tfp(t):
    return P.gA0hat * tf.exp(-P.Tstep * P.deltaA * t)
def gr_lab(t):
    return P.deltaL / ((P.Linfty / (P.Linfty - P.L0)) *
                        tf.exp(P.Tstep * P.deltaL * t) - 1.0)
def sigma(t):
    log_factor = np.float32(math.log(1.0 + P.Tstep * P.deltaSigma))
    return P.sigma0 * tf.exp(P.Tstep * P.gSigma0 / log_factor *
                    ((1.0 + P.Tstep * P.deltaSigma) ** t - 1.0))
def theta1(t):
    return P.pback * (1000.0 * P.c2co2 * sigma(t)) * tf.exp(-P.Tstep * P.gback * t) / P.theta2
def land_emissions(t):
    return P.ELand0 * tf.exp(-P.Tstep * P.deltaLand * t)
def ext_forcing(t):
    Yr = P.Tyears / P.Tstep
    return P.fex0 + (1.0 / Yr) * (P.fex1 - P.fex0) * tf.minimum(t, Yr)
def beta_hat(t):
    # effective per-period discount factor (Friedl et al. eq. 68)
    return tf.exp((-P.rho + (1.0 - 1.0/P.psi) * gr_tfp(t) + gr_lab(t)) * P.Tstep)

def Omega(TAT):
    return P.pi1 * tf.pow(TAT, P.pow1) + P.pi2 * tf.pow(TAT, P.pow2)
def Omega_prime(TAT):
    return P.pow1 * P.pi1 * tf.pow(TAT, P.pow1 - 1.0) + P.pow2 * P.pi2 * tf.pow(TAT, P.pow2 - 1.0)

def carbon_next(MAT, MUO, MLO, mu, k, t, z):
    A_ = tfp(t, z); L_ = population(t); sig_ = sigma(t); EL_ = land_emissions(t)
    E_ind = (1.0 - mu) * sig_ * A_ * L_ * tf.pow(k, P.alpha)
    MAT_n = (1.0 - P.b12) * MAT + P.b21 * MUO + P.Tstep * E_ind + P.Tstep * EL_
    MUO_n = P.b12 * MAT + (1.0 - P.b21 - P.b23) * MUO + P.b32 * MLO
    MLO_n = P.b23 * MUO + (1.0 - P.b32) * MLO
    return MAT_n, MUO_n, MLO_n

def temperature_next(TAT, TOC, MAT, t):
    Forcing = P.f2xco2 * tf.math.log(tf.maximum(MAT, 1e-6) / P.MATbase) / LOG2 + ext_forcing(t)
    TAT_n = (1.0 - P.c1c3 - P.c1f) * TAT + P.c1c3 * TOC + P.c1 * Forcing
    TOC_n = P.c4 * TAT + (1.0 - P.c4) * TOC
    return TAT_n, TOC_n

S_MIN = np.array([0.5,  0.5,  0.2,  1.0,  0.0,  0.0, 0.00, -0.5], dtype='float32')
S_MAX = np.array([60.0, 3.0,  3.0,  4.0, 10.0,  4.0, 0.99,  0.5], dtype='float32')
N_STATE = 8

def normalize(x):
    return (x - S_MIN) / (S_MAX - S_MIN + 1e-8)

def build_net(n_hidden=512):
    return tf.keras.Sequential([
        tf.keras.layers.Input(shape=(N_STATE,)),
        tf.keras.layers.Dense(n_hidden, activation='relu', kernel_initializer='glorot_uniform'),
        tf.keras.layers.Dense(n_hidden, activation='relu', kernel_initializer='glorot_uniform'),
        tf.keras.layers.Dense(8),
    ])

def split_policy(raw):
    k_plus     = tf.nn.softplus(raw[:, 0:1])
    lambd_hat  = tf.nn.softplus(raw[:, 1:2])
    mu         = tf.nn.softplus(raw[:, 2:3])
    nu_AT_hat  = tf.nn.softplus(raw[:, 3:4])
    nu_UO_hat  = raw[:, 4:5]
    nu_LO_hat  = raw[:, 5:6]
    eta_AT_hat = raw[:, 6:7]
    eta_OC_hat = raw[:, 7:8]
    return k_plus, lambd_hat, mu, nu_AT_hat, nu_UO_hat, nu_LO_hat, eta_AT_hat, eta_OC_hat

policy_net = build_net(n_hidden=512)
policy_net.summary()

SIGMA_Z_TF = tf.constant(float(P.sigma_z), dtype=tf.float32)

def make_loss(sigma_z_value):
    sig_z = tf.constant(float(sigma_z_value), dtype=tf.float32)

    @tf.function
    def loss_fn(states):
        states = tf.cast(states, tf.float32)
        k   = states[:, 0:1]; MAT = states[:, 1:2]
        MUO = states[:, 2:3]; MLO = states[:, 3:4]
        TAT = states[:, 4:5]; TOC = states[:, 5:6]
        tau = states[:, 6:7]; z   = states[:, 7:8]
        t   = tau2t(tau)

        raw = policy_net(normalize(states))
        k_plus, lambd_hat, mu, nu_AT, nu_UO, nu_LO, eta_AT, eta_OC = split_policy(raw)

        A_ = tfp(t, z); L_ = population(t); sig_ = sigma(t)
        th1_ = theta1(t); _bhat = beta_hat(t)
        growth_factor = tf.exp(P.Tstep * (gr_tfp(t) + gr_lab(t)))

        Omega_ = Omega(TAT)
        Theta_ = th1_ * tf.pow(mu, P.theta2)
        Theta_prime = th1_ * P.theta2 * tf.pow(mu, P.theta2 - 1.0)
        con = tf.pow(lambd_hat, -P.psi)

        MAT_p, MUO_p, MLO_p = carbon_next(MAT, MUO, MLO, mu, k, t, z)
        TAT_p, TOC_p = temperature_next(TAT, TOC, MAT, t)
        tau_p = tau2tauplus(tau)
        t_n = tau2t(tau_p)

        # next-period exogenous quantities that don't depend on z':
        L_n_const   = population(t_n)
        sig_n_const = sigma(t_n)
        th1_n_const = theta1(t_n)

        # Gauss-Hermite expectation: 5 weighted next-period network evaluations,
        # one per quadrature node z'_q = rho_z*z + sigma_z*sqrt(2)*x_q.
        E_eq1 = tf.zeros_like(k); E_eq4 = tf.zeros_like(k)
        E_eq5 = tf.zeros_like(k); E_eq6 = tf.zeros_like(k)
        E_eq7 = tf.zeros_like(k); E_eq8 = tf.zeros_like(k)
        for q in range(GH_Q):
            w_q   = GH_WEIGHTS[q]
            eps_q = GH_NODES[q]
            z_n   = P.rho_z * z + sig_z * eps_q
            st_n  = tf.concat([k_plus, MAT_p, MUO_p, MLO_p, TAT_p, TOC_p, tau_p, z_n], axis=1)
            raw_n = policy_net(normalize(st_n))
            _, lh, mu_n, nuAT, nuUO, nuLO, etAT, etOC = split_policy(raw_n)
            A_nq = tfp(t_n, z_n)
            Theta_n_q = th1_n_const * tf.pow(mu_n, P.theta2)

            E_eq1 += w_q * (lh * (P.Tstep * (1.0 - Theta_n_q - Omega(TAT_p))
                                  * P.alpha * tf.pow(k_plus, P.alpha - 1.0)
                                  + ONE_MINUS_DELTA_TSTEP)
                            + (-nuAT) * (1.0 - mu_n) * P.Tstep * sig_n_const * A_nq * L_n_const
                              * P.alpha * tf.pow(k_plus, P.alpha - 1.0))
            E_eq4 += w_q * (lh * (-P.Tstep * Omega_prime(TAT_p)) * tf.pow(k_plus, P.alpha)
                            + etAT * (1.0 - P.c1c3 - P.c1f) + etOC * P.c4)
            E_eq5 += w_q * ((-nuAT) * (1.0 - P.b12) + nuUO * P.b12
                            + etAT * P.c1 * P.f2xco2 / (LOG2 * tf.maximum(MAT_p, 1e-6)))
            E_eq6 += w_q * ((-nuAT) * P.b21 + nuUO * (1.0 - P.b21 - P.b23) + nuLO * P.b23)
            E_eq7 += w_q * (nuUO * P.b32 + nuLO * (1.0 - P.b32))
            E_eq8 += w_q * (etAT * P.c1c3 + etOC * (1.0 - P.c4))

        # Equilibrium residuals
        eq1 = growth_factor * lambd_hat - _bhat * E_eq1
        eq2 = (P.Tstep * (1.0 - Theta_ - Omega_) * tf.pow(k, P.alpha)
               - P.Tstep * con + ONE_MINUS_DELTA_TSTEP * k - growth_factor * k_plus)
        lambdMU = (-lambd_hat * P.Tstep * Theta_prime * tf.pow(k, P.alpha)
                   - (-nu_AT) * P.Tstep * sig_ * A_ * L_ * tf.pow(k, P.alpha))
        eq3 = lambdMU + (1.0 - mu) - tf.sqrt(lambdMU * lambdMU + (1.0 - mu) * (1.0 - mu) + 1e-12)
        eq4 = eta_AT  - _bhat * E_eq4
        eq5 = (-nu_AT) - _bhat * E_eq5
        eq6 = nu_UO   - _bhat * E_eq6
        eq7 = nu_LO   - _bhat * E_eq7
        eq8 = eta_OC  - _bhat * E_eq8

        pen_mu_upper = 1e-2 * tf.reduce_mean(tf.maximum(mu - 1.0, 0.0) ** 2)

        L_eq = (tf.reduce_mean(eq1**2) + tf.reduce_mean(eq2**2)
                + tf.reduce_mean(eq3**2) + tf.reduce_mean(eq4**2)
                + tf.reduce_mean(eq5**2) + tf.reduce_mean(eq6**2)
                + tf.reduce_mean(eq7**2) + tf.reduce_mean(eq8**2)) / 8.0
        return L_eq + pen_mu_upper

    return loss_fn

def gen_traj(n_traj=64, n_steps=300, sigma_z_realized=None):
    if sigma_z_realized is None:
        sigma_z_realized = float(P.sigma_z)
    rng = np.random.default_rng()
    k0   = (float(P.k0)   * rng.uniform(0.8, 1.2, (n_traj, 1))).astype('float32')
    MAT0 = (float(P.MAT0) * rng.uniform(0.9, 1.1, (n_traj, 1))).astype('float32')
    MUO0 = np.full((n_traj, 1), float(P.MUO0), dtype='float32')
    MLO0 = np.full((n_traj, 1), float(P.MLO0), dtype='float32')
    TAT0 = (float(P.TAT0) * rng.uniform(0.8, 1.2, (n_traj, 1))).astype('float32')
    TOC0 = np.full((n_traj, 1), float(P.TOC0), dtype='float32')
    tau0 = np.full((n_traj, 1), float(P.tau0), dtype='float32')
    if sigma_z_realized > 0:
        z0 = rng.normal(0.0, sigma_z_realized / np.sqrt(max(1 - float(P.rho_z)**2, 1e-6)),
                        (n_traj, 1)).astype('float32')
    else:
        z0 = np.zeros((n_traj, 1), dtype='float32')
    state = np.concatenate([k0, MAT0, MUO0, MLO0, TAT0, TOC0, tau0, z0], axis=1)
    out = [state.copy()]
    for _ in range(n_steps - 1):
        s_norm = (state - S_MIN) / (S_MAX - S_MIN + 1e-8)
        raw = policy_net(s_norm, training=False).numpy()
        k_plus = np.log1p(np.exp(raw[:, 0:1]))
        mu     = np.log1p(np.exp(raw[:, 2:3]))
        k = np.maximum(state[:, 0:1], 1e-6)
        MAT = state[:, 1:2]; MUO = state[:, 2:3]; MLO = state[:, 3:4]
        TAT = state[:, 4:5]; TOC = state[:, 5:6]; tau = state[:, 6:7]; z = state[:, 7:8]
        t = (-np.log(np.maximum(1.0 - tau, 1e-6)) / float(P.vartheta)).astype('float32')
        t_tf = tf.constant(t); z_tf = tf.constant(z)
        A_ = tfp(t_tf, z_tf).numpy(); L_ = population(t_tf).numpy()
        sig_ = sigma(t_tf).numpy(); EL_ = land_emissions(t_tf).numpy()
        E_ind = (1.0 - mu) * sig_ * A_ * L_ * np.power(k, float(P.alpha))
        MAT_n = (1-float(P.b12))*MAT + float(P.b21)*MUO + float(P.Tstep)*E_ind + float(P.Tstep)*EL_
        MUO_n = float(P.b12)*MAT + (1-float(P.b21)-float(P.b23))*MUO + float(P.b32)*MLO
        MLO_n = float(P.b23)*MUO + (1-float(P.b32))*MLO
        Forcing = float(P.f2xco2)*np.log(np.maximum(MAT, 1e-6)/float(P.MATbase))/np.log(2.0) + ext_forcing(t_tf).numpy()
        TAT_n = (1-float(P.c1c3)-float(P.c1f))*TAT + float(P.c1c3)*TOC + float(P.c1)*Forcing
        TOC_n = float(P.c4)*TAT + (1-float(P.c4))*TOC
        tau_n = (1.0 - np.exp(-float(P.vartheta) * (t + float(P.Tstep)))).astype('float32')
        eps_n = rng.normal(0.0, 1.0, (n_traj, 1)).astype('float32')
        z_n   = (float(P.rho_z) * z + sigma_z_realized * eps_n).astype('float32')
        state = np.concatenate(
            [np.maximum(k_plus, 0.01), MAT_n, MUO_n, MLO_n, TAT_n, TOC_n, tau_n, z_n],
            axis=1).astype('float32')
        out.append(state.copy())
    return np.concatenate(out, axis=0)

optimizer = tf.keras.optimizers.Adam(learning_rate=5e-5, clipvalue=1.0)
LR_SCHED = {0: 5e-5, 80: 2e-5, 160: 1e-5}

SIGMA_FOR_LOSS = float(P.sigma_z)
print(f'Training with sigma_z = {SIGMA_FOR_LOSS:.4f}')
loss_fn = make_loss(SIGMA_FOR_LOSS)

@tf.function
def train_step(batch):
    with tf.GradientTape() as tape:
        L = loss_fn(batch)
    grads = tape.gradient(L, policy_net.trainable_variables)
    optimizer.apply_gradients(zip(grads, policy_net.trainable_variables))
    return L

N_EPISODES = {"smoke": 12, "teaching": 200, "production": 1000}[RUN_MODE]
N_INNER    = {"smoke": 30, "teaching": 100, "production": 100}[RUN_MODE]
N_TRAJ     = 64
N_STEPS    = 300
BATCH      = 512

losses = []
t0 = time.time()
for ep in range(N_EPISODES):
    if ep in LR_SCHED:
        optimizer.learning_rate.assign(LR_SCHED[ep])
    traj = gen_traj(n_traj=N_TRAJ, n_steps=N_STEPS, sigma_z_realized=SIGMA_FOR_LOSS)
    for _ in range(N_INNER):
        idx = np.random.choice(len(traj), BATCH, replace=False)
        L = train_step(traj[idx])
        losses.append(float(L.numpy()))
    if ep % 10 == 0:
        print(f'  ep {ep:3d}  L={float(L):.3e}  ({time.time()-t0:.0f}s)', flush=True)
print(f'done {(time.time()-t0)/60:.1f} min')

fig, ax = plt.subplots(figsize=(8.5, 3.6))
ax.plot(np.arange(len(losses)) / N_INNER, losses, lw=0.6, alpha=0.6)
# moving average for clarity
window = max(1, len(losses) // 200)
ma = np.convolve(losses, np.ones(window)/window, mode='valid')
ax.plot(np.arange(len(ma)) / N_INNER, ma, color='C3', lw=1.6, label=f'moving avg ({window} steps)')
ax.set_xlabel('episode'); ax.set_ylabel('mean-squared residual')
ax.set_yscale('log'); ax.grid(alpha=0.3); ax.legend()
ax.set_title('DEQN training loss — 8 FOCs + Gauss-Hermite expectations')
plt.tight_layout(); plt.show()

def simulate_one(n=300, sigma_z_realized=0.0, seed=None):
    rng = np.random.default_rng(seed)
    state = [float(P.k0), float(P.MAT0), float(P.MUO0), float(P.MLO0),
             float(P.TAT0), float(P.TOC0), float(P.tau0), float(P.z0)]
    out = {ky: [] for ky in ['year', 'k', 'MAT', 'TAT', 'mu', 'con', 'scc', 'z']}
    for s in range(n):
        year = 2015.0 + s
        k_, MAT_, MUO_, MLO_, TAT_, TOC_, tau_, z_ = state
        t_val = -np.log(max(1 - tau_, 1e-6)) / float(P.vartheta)
        x = np.array([state], dtype='float32')
        raw = policy_net((x - S_MIN) / (S_MAX - S_MIN + 1e-8), training=False).numpy()
        kp = np.log1p(np.exp(raw[0, 0]))
        lh = np.log1p(np.exp(raw[0, 1]))
        mu_t = np.log1p(np.exp(raw[0, 2]))
        nuAT = np.log1p(np.exp(raw[0, 3]))
        nuUO = raw[0, 4]; etAT = raw[0, 6]
        t_tf = tf.constant([[t_val]], dtype=tf.float32)
        z_tf = tf.constant([[z_]], dtype=tf.float32)
        sig_ = sigma(t_tf).numpy().item()
        A_ = tfp(t_tf, z_tf).numpy().item(); L_ = population(t_tf).numpy().item()
        EL_ = land_emissions(t_tf).numpy().item()
        Fex_ = ext_forcing(t_tf).numpy().item()
        th1_ = (float(P.pback) * (1000.0 * float(P.c2co2) * sig_)
                * np.exp(-float(P.Tstep) * float(P.gback) * t_val) / float(P.theta2))
        Theta_ = th1_ * mu_t**float(P.theta2)
        Omega_ = float(P.pi2) * TAT_**2
        con = lh ** (-float(P.psi))
        k_safe = max(k_, 1e-6)
        dvdk = (lh * (float(P.Tstep) * (1 - Theta_ - Omega_) * float(P.alpha) * k_safe**(float(P.alpha)-1)
                      + (1-float(P.delta))**float(P.Tstep))
                + (-nuAT) * (1 - mu_t) * float(P.Tstep) * sig_ * A_ * L_ * float(P.alpha) * k_safe**(float(P.alpha)-1))
        dvdMAT = ((-nuAT) * (1 - float(P.b12)) + nuUO * float(P.b12)
                  + etAT * float(P.c1) * float(P.f2xco2) / (np.log(2.0) * max(MAT_, 1e-6)))
        scc = -dvdMAT / (dvdk + 1e-15) * A_ * L_ / float(P.c2co2)
        out['year'].append(year); out['k'].append(k_ * A_ * L_); out['MAT'].append(MAT_ * 1000)
        out['TAT'].append(TAT_); out['mu'].append(mu_t)
        out['con'].append(con * A_ * L_); out['scc'].append(scc); out['z'].append(z_)
        E_ind = (1 - mu_t) * sig_ * A_ * L_ * k_safe**float(P.alpha)
        MAT_n = (1-float(P.b12))*MAT_ + float(P.b21)*MUO_ + float(P.Tstep)*E_ind + float(P.Tstep)*EL_
        MUO_n = float(P.b12)*MAT_ + (1-float(P.b21)-float(P.b23))*MUO_ + float(P.b32)*MLO_
        MLO_n = float(P.b23)*MUO_ + (1-float(P.b32))*MLO_
        Forcing = float(P.f2xco2)*np.log(max(MAT_, 1e-6)/float(P.MATbase))/np.log(2.0) + Fex_
        TAT_n = (1-float(P.c1c3)-float(P.c1f))*TAT_ + float(P.c1c3)*TOC_ + float(P.c1)*Forcing
        TOC_n = float(P.c4)*TAT_ + (1-float(P.c4))*TOC_
        tau_n = 1.0 - np.exp(-float(P.vartheta) * (t_val + float(P.Tstep)))
        eps_n = rng.normal(0.0, 1.0)
        z_n   = float(P.rho_z) * z_ + sigma_z_realized * eps_n
        state = [max(kp, 0.01), MAT_n, MUO_n, MLO_n, TAT_n, TOC_n, tau_n, z_n]
    return {k: np.array(v) for k, v in out.items()}

res_det = simulate_one(n=300, sigma_z_realized=0.0, seed=0)
ref = {(0, 'TAT'): 1.10, (0, 'MAT'): 851.0, (0, 'mu'): 0.144, (0, 'scc'): 24.82,
       (85, 'TAT'): 2.92, (85, 'MAT'): 1222.9, (85, 'mu'): 0.673, (85, 'scc'): 186.4,
       (285, 'TAT'): 3.10, (285, 'mu'): 0.999}
print(f'{"year":<6}{"var":<6}{"got":>10}{"ref":>10}{"diff":>10}  status')
n_ok = n_close = n_fail = 0
for (idx, var), v in ref.items():
    g = res_det[var][idx]
    err = abs(g - v) / max(abs(v), 1e-6) * 100
    flag = 'OK' if err < 5 else ('CLOSE' if err < 15 else 'FAIL')
    n_ok += (flag == 'OK'); n_close += (flag == 'CLOSE'); n_fail += (flag == 'FAIL')
    print(f'{2015+idx:<6}{var:<6}{g:>10.3f}{v:>10.3f}{err:>9.1f}%   [{flag}]')
print(f'\nGate A: {n_ok} OK, {n_close} CLOSE, {n_fail} FAIL  out of {len(ref)}')

det = res_det
fig, axes = plt.subplots(2, 2, figsize=(10, 6.5), sharex=True)
yr = det['year']
axes[0, 0].plot(yr, det['TAT'], color='C3', lw=2)
axes[0, 0].set_ylabel('atm. temp $T_{AT}$ [°C]')
axes[0, 0].axhline(2.0, color='k', ls=':', alpha=0.5, label='2°C')
axes[0, 0].legend(loc='lower right'); axes[0, 0].grid(alpha=0.3)

axes[0, 1].plot(yr, det['MAT'], color='C0', lw=2)
axes[0, 1].set_ylabel('atm. carbon $M_{AT}$ [GtC]')
axes[0, 1].grid(alpha=0.3)

axes[1, 0].plot(yr, det['mu'], color='C2', lw=2)
axes[1, 0].set_ylabel(r'abatement rate $\mu$')
axes[1, 0].set_xlabel('year'); axes[1, 0].grid(alpha=0.3)
axes[1, 0].set_ylim(-0.05, 1.05)

axes[1, 1].plot(yr, det['con'], color='C4', lw=2)
axes[1, 1].set_ylabel('consumption $c \cdot A \cdot L$ [\\$\\,trill.]')
axes[1, 1].set_xlabel('year'); axes[1, 1].grid(alpha=0.3)

fig.suptitle('Deterministic optimal path ($z \\equiv 0$)')
plt.tight_layout(); plt.show()

scc = det['scc']; yr = det['year']

fig, axes = plt.subplots(1, 2, figsize=(11, 3.8))

axes[0].plot(yr, scc, color='C1', lw=2)
axes[0].set_xlabel('year'); axes[0].set_ylabel('SCC [\\$/tCO$_2$]')
axes[0].set_title('SCC, linear scale (convex shape)')
axes[0].grid(alpha=0.3)
# Annotate three reference years
for y_ref, label in [(2015, '2015'), (2050, '2050'), (2100, '2100')]:
    i = int(y_ref - 2015); axes[0].plot(yr[i], scc[i], 'o', color='C3')
    axes[0].annotate(f'{label}\n\\${scc[i]:.0f}', (yr[i], scc[i]),
                     xytext=(5, -25), textcoords='offset points', fontsize=10)

axes[1].semilogy(yr, scc, color='C1', lw=2)
axes[1].set_xlabel('year'); axes[1].set_ylabel('SCC [\\$/tCO$_2$, log scale]')
axes[1].set_title('SCC, log scale (linearity ↔ exponential growth)')
axes[1].grid(alpha=0.3, which='both')

plt.tight_layout(); plt.show()

# Numerical convexity / growth-rate check
print(f'SCC(2015) = {scc[0]:.2f}')
print(f'SCC(2050) = {scc[35]:.2f}    avg growth 2015-2050: {(np.log(scc[35]/scc[0])/35*100):.2f} %/yr')
print(f'SCC(2100) = {scc[85]:.2f}    avg growth 2050-2100: {(np.log(scc[85]/scc[35])/50*100):.2f} %/yr')
slopes = np.gradient(scc, yr)
n_increasing = int(np.all(np.diff(scc[:200]) > 0))
print(f'SCC strictly increasing through year 2215: {bool(n_increasing)}')
print(f'd^2 SCC / dt^2 averaged over 2015-2100: {np.mean(np.diff(slopes[:85])):.4f} (>0 = convex)')

N_MC = {"smoke": 20, "teaching": 200, "production": 500}[RUN_MODE]
N_YEARS = 285
ts_scc = np.zeros((N_MC, N_YEARS))
ts_TAT = np.zeros((N_MC, N_YEARS))
ts_mu  = np.zeros((N_MC, N_YEARS))
ts_z   = np.zeros((N_MC, N_YEARS))
print('Running 200 Monte-Carlo paths ...', flush=True)
t0 = time.time()
for i in range(N_MC):
    r = simulate_one(n=N_YEARS, sigma_z_realized=float(P.sigma_z), seed=1000 + i)
    ts_scc[i] = r['scc']; ts_TAT[i] = r['TAT']
    ts_mu[i]  = r['mu'];  ts_z[i]   = r['z']
    if (i + 1) % 50 == 0:
        print(f'  {i+1}/{N_MC}  ({time.time()-t0:.0f}s)', flush=True)
years = np.arange(2015, 2015 + N_YEARS)
print(f'done in {time.time()-t0:.0f}s')

def fan(ax, years, ts, color, label):
    p05 = np.percentile(ts, 5,  axis=0); p25 = np.percentile(ts, 25, axis=0)
    p50 = np.percentile(ts, 50, axis=0); p75 = np.percentile(ts, 75, axis=0)
    p95 = np.percentile(ts, 95, axis=0)
    ax.fill_between(years, p05, p95, color=color, alpha=0.18, label=f'{label} 5-95%')
    ax.fill_between(years, p25, p75, color=color, alpha=0.35, label=f'{label} 25-75%')
    ax.plot(years, p50, color=color, lw=1.6, label=f'{label} median')

fig, ax = plt.subplots(figsize=(10, 4.8))
fan(ax, years, ts_scc, 'C0', 'stochastic')
ax.plot(years, det['scc'][:N_YEARS], color='C1', lw=2.0, ls='--', label='deterministic ($z\\equiv 0$)')
ax.set_xlabel('year'); ax.set_ylabel('SCC [\\$/tCO$_2$]')
ax.set_title(f'SCC distribution over time — {N_MC} Monte Carlo paths, '
             rf'$\rho_z = {float(P.rho_z):.2f}$, $\sigma_z = {float(P.sigma_z):.4f}$')
ax.legend(loc='upper left', fontsize=10); ax.grid(alpha=0.3)
ax.set_xlim(2015, 2200)
plt.tight_layout(); plt.show()

fig, axes = plt.subplots(1, 2, figsize=(11.5, 4))
fan(axes[0], years, ts_TAT, 'C3', 'stochastic')
axes[0].plot(years, det['TAT'][:N_YEARS], color='k', lw=1.6, ls='--', label='deterministic')
axes[0].axhline(2.0, color='gray', ls=':', alpha=0.6, label='2°C')
axes[0].set_xlabel('year'); axes[0].set_ylabel('$T_{AT}$ [°C]')
axes[0].set_title('Atmospheric temperature distribution'); axes[0].legend(fontsize=9)
axes[0].grid(alpha=0.3); axes[0].set_xlim(2015, 2200)

fan(axes[1], years, ts_mu, 'C2', 'stochastic')
axes[1].plot(years, det['mu'][:N_YEARS], color='k', lw=1.6, ls='--', label='deterministic')
axes[1].set_xlabel('year'); axes[1].set_ylabel(r'$\mu$ (abatement rate)')
axes[1].set_title('Optimal abatement rate distribution'); axes[1].legend(fontsize=9, loc='lower right')
axes[1].grid(alpha=0.3); axes[1].set_xlim(2015, 2200); axes[1].set_ylim(-0.02, 1.05)

plt.tight_layout(); plt.show()

fig, axes = plt.subplots(1, 2, figsize=(11.5, 3.8))
for ax, year_target in zip(axes, [2050, 2100]):
    i = year_target - 2015
    samp = ts_scc[:, i]
    det_v = det['scc'][i]
    ax.hist(samp, bins=25, color='C0', alpha=0.7, edgecolor='white')
    ax.axvline(det_v, color='C1', lw=2.0, ls='--', label=f'deterministic = \\${det_v:.1f}')
    ax.axvline(samp.mean(), color='C3', lw=1.6, label=f'stoch mean = \\${samp.mean():.1f}')
    ax.set_xlabel(f'SCC at {year_target} [\\$/tCO$_2$]'); ax.set_ylabel('count')
    ax.set_title(f'{year_target}:  std = \\${samp.std():.2f},  '
                 f'5-95% = [\\${np.percentile(samp,5):.0f}, \\${np.percentile(samp,95):.0f}]')
    ax.legend(fontsize=10); ax.grid(alpha=0.3)
plt.tight_layout(); plt.show()

print('Cai-Lontzek qualitative checks:')
i50, i100 = 35, 85
for year_target, i in [(2050, i50), (2100, i100)]:
    samp = ts_scc[:, i]; det_v = det['scc'][i]
    spread_ok = samp.std() > 0.05 * det_v
    mean_ok   = abs(samp.mean() - det_v) / det_v < 0.25
    print(f'  {year_target}: std/det = {samp.std()/det_v:.3f}  '
          f'(distribution non-degenerate: {spread_ok});  '
          f'|stoch mean - det|/det = {abs(samp.mean()-det_v)/det_v:.3f}  '
          f'(mean within 25%: {mean_ok})')

rng = np.random.default_rng(SEED)
sel = rng.choice(N_MC, 10, replace=False)
fig, axes = plt.subplots(1, 3, figsize=(14, 3.6))

for j in sel:
    axes[0].plot(years, ts_TAT[j], color='C3', alpha=0.35, lw=1.0)
    axes[1].plot(years, ts_mu[j],  color='C2', alpha=0.35, lw=1.0)
    axes[2].plot(years, ts_scc[j], color='C0', alpha=0.35, lw=1.0)

axes[0].plot(years, det['TAT'][:N_YEARS], color='k', lw=2.0, label='deterministic')
axes[1].plot(years, det['mu'][:N_YEARS],  color='k', lw=2.0, label='deterministic')
axes[2].plot(years, det['scc'][:N_YEARS], color='k', lw=2.0, label='deterministic')

for ax, ylab, ttl in zip(axes,
                         ['$T_{AT}$ [°C]', r'$\mu$', 'SCC [\\$/tCO$_2$]'],
                         ['Temperature', 'Abatement rate', 'SCC']):
    ax.set_xlabel('year'); ax.set_ylabel(ylab); ax.set_title(ttl)
    ax.set_xlim(2015, 2200); ax.grid(alpha=0.3); ax.legend(fontsize=9)

plt.tight_layout(); plt.show()

# Validation: in smoke mode the short run will not match the reference; just confirm the run is sane.
if RUN_MODE == "smoke":
    assert np.isfinite(losses[-1]) and losses[-1] < losses[0], f"Training loss did not decrease: {losses[0]:.2e} -> {losses[-1]:.2e}"
    assert np.isfinite(res_det["scc"]).all() and res_det["scc"][0] > 0.0, "Deterministic SCC not finite-positive"
    assert np.isfinite(ts_scc).all() and (ts_scc[:, 85] > 0.0).all(), "Monte-Carlo SCC at 2100 not finite-positive"
else:
    assert n_fail == 0, f"Gate A: {n_fail} reference checks failed (out of {len(ref)})"