Lecture 09, Notebook 12: Krusell–Smith with all-in-one deep learning

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

# Core scientific Python
import numpy as np
import matplotlib.pyplot as plt

# TensorFlow 2.x with Keras
import tensorflow as tf
from tensorflow import keras

# Reproducibility -- reseed from the cell-1 SEED constant.
np.random.seed(SEED)
tf.random.set_seed(SEED)

print("TensorFlow version:", tf.__version__)
print("Seed:", SEED)
plt.rcParams["font.size"] = 12

# --- Preferences ---
beta   = tf.constant(0.99, dtype=tf.float32)
sigma  = tf.constant(1.0,  dtype=tf.float32)    # log utility

# --- Technology ---
alpha  = tf.constant(0.36, dtype=tf.float32)
delta  = tf.constant(0.025, dtype=tf.float32)

# --- Aggregate and idiosyncratic shock levels ---
z_levels = tf.constant([0.99, 1.01], dtype=tf.float32)   # (bad, good)
u_rates  = tf.constant([0.10, 0.04], dtype=tf.float32)   # (bad, good)

# Steady-state average unemployment (used for labor normalisation)
u_ss   = 0.07
l_bar  = tf.constant(1.0 / (1.0 - u_ss), dtype=tf.float32)

# --- Panel size, dispatched on RUN_MODE ---
_PANEL = {"smoke": 100, "teaching": 300, "production": 1000}[RUN_MODE]
N_AGENTS = _PANEL

# --- Deterministic steady-state k (for panel initialisation) ---
r_ss_det = 1.0/float(beta.numpy()) - 1.0 + float(delta.numpy())
K_ss_det = ((float(alpha.numpy()) / r_ss_det) ** (1.0/(1.0-float(alpha.numpy()))))
# Adjust for L_ss (avg employed):
L_ss = 1.0 - u_ss
K_SS = K_ss_det * L_ss          # pulls the SS down a bit via the production function
print(f"log-utility KS calibration: β={float(beta.numpy())}, α={float(alpha.numpy())}, δ={float(delta.numpy())}")
print(f"deterministic steady-state K ≈ {K_SS:.2f}")

# --- Capital domain ---
# Used for state-space sampling in Phase A.  Covers the ergodic range + margin.
K_MIN = 0.5 * K_SS
K_MAX = 1.5 * K_SS
k_min = tf.constant(1e-3, dtype=tf.float32)   # avoid log(0)
k_ind_min = 0.2 * K_SS                         # individual k sampled in Phase A
k_ind_max = 4.0 * K_SS                         # (wide enough to cover realistic holdings)
print(f"training domain: K̄ ∈ [{K_MIN:.2f}, {K_MAX:.2f}],  k ∈ [{k_ind_min:.2f}, {k_ind_max:.2f}]")

# Aggregate-state transition
P_agg = np.array([
    [0.875, 0.125],    # bad  -> (bad, good)
    [0.125, 0.875],    # good -> (bad, good)
], dtype=np.float32)

# Conditional idiosyncratic transition P(eps' | a, a', eps)
# Shape (a, a', eps, eps'); rows over eps' sum to 1.
P_eps = np.zeros((2, 2, 2, 2), dtype=np.float32)

P_eps[0, 0] = [[0.525, 0.475],
               [0.350, 0.650]]            # bad  -> bad
P_eps[0, 1] = [[0.093, 0.907],
               [0.038, 0.962]]            # bad  -> good
P_eps[1, 0] = [[0.840, 0.160],
               [0.200, 0.800]]            # good -> bad
P_eps[1, 1] = [[0.292, 0.708],
               [0.042, 0.958]]            # good -> good

assert np.allclose(P_eps.sum(axis=-1), 1.0, atol=1e-6)

P_agg_tf = tf.constant(P_agg)
P_eps_tf = tf.constant(P_eps)
print("Aggregate transition matrix:\n", P_agg)

@tf.function
def prices(K_agg, a_idx):
    '''Aggregate prices w, r at (K̄, a).

    K_agg: () tf.float32 scalar
    a_idx: () int32 scalar; 0 = bad, 1 = good.
    '''
    z = tf.gather(z_levels, a_idx)
    u = tf.gather(u_rates,  a_idx)
    L = 1.0 - u
    w = (1.0 - alpha) * z * tf.pow(K_agg, alpha)        * tf.pow(L, -alpha)
    r = alpha         * z * tf.pow(K_agg, alpha - 1.0)  * tf.pow(L, 1.0 - alpha)
    return w, r

@tf.function
def cash_on_hand_panel(k, eps_float, w, r):
    '''Cash on hand for a panel of agents.  Shapes: k,eps_float (N,); w,r (). Returns (N,).'''
    return w * l_bar * eps_float + (1.0 + r - delta) * k

class LogNormaliseLayer(keras.layers.Layer):
    '''Compute log(max(x, eps) / K_ss) as a Keras-native layer.
    Keras Functional models require ops in the namespace keras.ops or keras.layers.
    '''
    def __init__(self, K_ss, eps=1e-3, **kw):
        super().__init__(**kw)
        self.K_ss = float(K_ss)
        self.eps  = float(eps)
    def call(self, x):
        import tensorflow as tf
        return tf.math.log(tf.maximum(x, self.eps) / self.K_ss)

def build_policy_network():
    k_in   = keras.Input(shape=(1,), name="k")
    eps_in = keras.Input(shape=(2,), name="eps")
    K_in   = keras.Input(shape=(1,), name="K")
    a_in   = keras.Input(shape=(2,), name="a")

    # Log-normalise the continuous inputs around the deterministic SS
    k_feat = LogNormaliseLayer(K_SS, name="log_k_over_Kss")(k_in)
    K_feat = LogNormaliseLayer(K_SS, name="log_K_over_Kss")(K_in)

    h = keras.layers.concatenate([k_feat, eps_in, K_feat, a_in])
    h = keras.layers.Dense(64, activation="swish")(h)
    h = keras.layers.Dense(64, activation="swish")(h)
    out = keras.layers.Dense(1, activation="sigmoid", name="phi")(h)
    return keras.Model([k_in, eps_in, K_in, a_in], out)

policy_net = build_policy_network()
policy_net.summary()

def onehot2(idx_tensor, n=2):
    '''Create one-hot tensor; works on either Python int or tf tensor of ints.'''
    return tf.one_hot(idx_tensor, depth=n, dtype=tf.float32)

def apply_policy(k, eps_int, K_agg, a_idx):
    '''Apply the policy network to a batch of individual states + single aggregate state.

    k:       (N,)   tf.float32, individual capital
    eps_int: (N,)   tf.int32, employment indicator (0,1)
    K_agg:   ()     tf.float32, aggregate capital
    a_idx:   ()     tf.int32,  aggregate state index

    Returns phi, c, k_next, y, w, r
    '''
    w, r = prices(K_agg, a_idx)
    eps_float = tf.cast(eps_int, tf.float32)
    y = cash_on_hand_panel(k, eps_float, w, r)

    N = tf.shape(k)[0]
    k_col    = tf.reshape(k, (-1, 1))
    eps_1hot = onehot2(eps_int)
    K_col    = tf.fill([N, 1], K_agg)
    a_1hot   = tf.tile(tf.reshape(onehot2(a_idx), (1, 2)), [N, 1])

    phi = policy_net([k_col, eps_1hot, K_col, a_1hot], training=True)
    phi = tf.squeeze(phi, axis=-1)
    # strict interior to avoid log(0) / div-by-0 in the Euler residual
    phi = tf.clip_by_value(phi, 1e-3, 1.0 - 1e-3)

    c      = phi * y
    k_next = (1.0 - phi) * y
    return phi, c, k_next, y, w, r

def euler_residual(k_panel, eps_int, K_agg, a_idx):
    '''Return (loss, ree, c, k_next, K_next) for a batch / panel.'''
    # --- today ---
    phi, c, k_next, y, w, r = apply_policy(k_panel, eps_int, K_agg, a_idx)

    # Next-period aggregate capital: average of the network's own savings
    K_next = tf.reduce_mean(k_next)

    # --- tomorrow: exact sum over (a', eps') pairs ---
    # Joint conditional prob of (a', eps') given (a, eps_i), for each agent i
    # P_eps_tf[a_idx, :, eps_i, :]  has shape (2 a', 2 eps'); we stack over agents.
    pe_stack = tf.gather(P_eps_tf[a_idx], eps_int, axis=1)       # (2 a', N, 2 eps')
    pe_stack = tf.transpose(pe_stack, [1, 0, 2])                 # (N, 2 a', 2 eps')
    P_ap     = P_agg_tf[a_idx]                                   # (2,) over a'
    prob     = pe_stack * tf.reshape(P_ap, (1, 2, 1))            # (N, 2 a', 2 eps')

    # Accumulate E[(1 + r' - delta) / c'] across (a', eps')
    EMU = tf.zeros_like(c)
    for a_next in range(2):
        w_n, r_n = prices(K_next, a_next)
        for eps_next_val in range(2):
            eps_next_vec = tf.fill(tf.shape(eps_int), eps_next_val)
            eps_n_float  = tf.cast(eps_next_vec, tf.float32)
            y_n = cash_on_hand_panel(k_next, eps_n_float, w_n, r_n)

            N = tf.shape(k_next)[0]
            k_col    = tf.reshape(k_next, (-1, 1))
            eps_1hot = onehot2(eps_next_vec)
            K_col    = tf.fill([N, 1], K_next)
            a_1hot   = tf.tile(tf.reshape(onehot2(a_next), (1, 2)), [N, 1])
            phi_n = tf.squeeze(
                policy_net([k_col, eps_1hot, K_col, a_1hot], training=True), axis=-1)
            phi_n = tf.clip_by_value(phi_n, 1e-3, 1.0 - 1e-3)
            c_n = phi_n * y_n

            contrib = (1.0 + r_n - delta) / c_n                            # (N,)
            weight  = prob[:, a_next, eps_next_val]                         # (N,)
            EMU += weight * contrib

    # Relative Euler error (unit-free)
    ree = beta * c * EMU - 1.0
    loss = tf.reduce_mean(ree ** 2)
    return loss, ree, c, k_next, K_next

def step_panel(k, eps_int, a_idx):
    '''One-period simulation: new (k', eps', a').  Inputs are tf tensors, output too.'''
    K_agg = tf.reduce_mean(k)
    _, _, k_next, _, _, _ = apply_policy(k, eps_int, K_agg, a_idx)
    k_next_np = k_next.numpy()

    a_next = int(np.random.choice(2, p=P_agg[a_idx]))

    eps_int_np = eps_int.numpy() if isinstance(eps_int, tf.Tensor) else np.asarray(eps_int)
    eps_next_np = np.empty_like(eps_int_np)
    for i in range(len(eps_int_np)):
        eps_next_np[i] = np.random.choice(2, p=P_eps[a_idx, a_next, eps_int_np[i]])

    return (
        tf.constant(k_next_np,   dtype=tf.float32),
        tf.constant(eps_next_np, dtype=tf.int32),
        a_next,
    )

def sample_state_batch(N):
    '''Draw a random training batch uniformly over the bounded training domain.

    Returns (k, eps_int, K_agg, a_idx) as tf tensors / Python ints.
    '''
    k       = np.random.uniform(k_ind_min, k_ind_max, size=N).astype(np.float32)
    eps_int = np.random.binomial(1, 1.0 - u_ss, size=N).astype(np.int32)
    K_agg   = float(np.random.uniform(K_MIN, K_MAX))
    a_idx   = int(np.random.randint(0, 2))
    return (tf.constant(k), tf.constant(eps_int),
            tf.constant(K_agg, dtype=tf.float32), tf.constant(a_idx, dtype=tf.int32))

# ---- Hyperparameters ----
N_PHASE_A = 2000
N_PHASE_B = 3000
N_TOTAL   = N_PHASE_A + N_PHASE_B
LR_MAX, LR_MIN = 1e-3, 1e-5

sched = keras.optimizers.schedules.CosineDecay(
    initial_learning_rate=LR_MAX,
    decay_steps=N_TOTAL,
    alpha=LR_MIN / LR_MAX,
)
optimizer = keras.optimizers.Adam(learning_rate=sched)

# Storage
loss_history    = np.zeros(N_TOTAL, dtype=np.float32)
ree_max_history = np.zeros(N_TOTAL, dtype=np.float32)
K_history       = np.zeros(N_TOTAL, dtype=np.float32)
a_history       = np.zeros(N_TOTAL, dtype=np.int32)
phase_history   = np.zeros(N_TOTAL, dtype=np.int32)  # 0 = A, 1 = B

# --- Phase A: exogenous sampling ---
print(f"Phase A: {N_PHASE_A} iterations of state-space sampling...")
for i in range(N_PHASE_A):
    k_b, eps_b, K_b, a_b = sample_state_batch(N_AGENTS)
    with tf.GradientTape() as tape:
        loss, ree, c, k_next, K_next = euler_residual(k_b, eps_b, K_b, a_b)
    grads = tape.gradient(loss, policy_net.trainable_variables)
    optimizer.apply_gradients(zip(grads, policy_net.trainable_variables))

    loss_history[i]    = loss.numpy()
    ree_max_history[i] = tf.reduce_max(tf.abs(ree)).numpy()
    K_history[i]       = float(K_b.numpy())
    a_history[i]       = int(a_b.numpy())
    phase_history[i]   = 0

    if i % 200 == 0:
        print(f"  A-iter {i:4d}  loss={loss.numpy():.3e}  "
              f"max|REE|={ree_max_history[i]:.3e}")

# --- Phase B: simulated-panel sampling ---
# Initialise the panel close to the deterministic steady state
k_panel  = tf.constant(np.full(N_AGENTS, K_SS, dtype=np.float32))
eps_panel = tf.constant(
    np.random.binomial(1, 1.0 - u_ss, size=N_AGENTS).astype(np.int32))
a_state   = 1  # start in the good state

print(f"Phase B: {N_PHASE_B} iterations of simulated-panel sampling...")
for j in range(N_PHASE_B):
    K_agg = tf.reduce_mean(k_panel)
    a_tf  = tf.constant(a_state, dtype=tf.int32)
    with tf.GradientTape() as tape:
        loss, ree, c, k_next, K_next = euler_residual(k_panel, eps_panel, K_agg, a_tf)
    grads = tape.gradient(loss, policy_net.trainable_variables)
    optimizer.apply_gradients(zip(grads, policy_net.trainable_variables))

    idx = N_PHASE_A + j
    loss_history[idx]    = loss.numpy()
    ree_max_history[idx] = tf.reduce_max(tf.abs(ree)).numpy()
    K_history[idx]       = float(K_agg.numpy())
    a_history[idx]       = a_state
    phase_history[idx]   = 1

    # Advance the panel one period
    k_panel, eps_panel, a_state = step_panel(k_panel, eps_panel, a_state)

    if j % 200 == 0:
        print(f"  B-iter {j:4d}  loss={loss.numpy():.3e}  "
              f"max|REE|={ree_max_history[idx]:.3e}  "
              f"K̄={K_history[idx]:.2f}")

print("training done.")

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

x = np.arange(N_TOTAL)
for ax, (y, title) in zip(
    axes,
    [(loss_history,    "Training loss (MSE of relative Euler error)"),
     (ree_max_history, "Worst-agent |relative Euler error|")],
):
    ax.semilogy(x[phase_history == 0], y[phase_history == 0], ".", markersize=2,
                color="#5555cc", label="Phase A (exogenous sampling)")
    ax.semilogy(x[phase_history == 1], y[phase_history == 1], ".", markersize=2,
                color="#cc3333", label="Phase B (simulated panel)")
    ax.axvline(N_PHASE_A, color="gray", ls="--", lw=1)
    ax.set_xlabel("iteration")
    ax.set_ylabel(title.split(" ")[0])
    ax.set_title(title)
    ax.grid(True, which="both", alpha=0.3)
    ax.legend(loc="upper right", fontsize=9)

plt.tight_layout()
plt.show()

# Report post-training statistics (last 200 iterations of Phase B)
tail = slice(-200, None)
print(f"\nFinal (last 200 iterations of Phase B):")
print(f"  mean loss       = {loss_history[tail].mean():.3e}")
print(f"  mean max|REE|   = {ree_max_history[tail].mean():.3e}")
print(f"  median max|REE| = {np.median(ree_max_history[tail]):.3e}")
print(f"\nSuccess criterion: median max|REE| < 5e-2 and mean loss < 1e-3.")
pass_crit = (np.median(ree_max_history[tail]) < 5e-2) and (loss_history[tail].mean() < 1e-3)
print(f"  --> {'PASS' if pass_crit else 'did not fully converge; see notes below'}")

# Restrict to Phase B for ergodic analysis (the only phase with a real panel)
mask_B = (phase_history == 1)
K_B = K_history[mask_B]
a_B = a_history[mask_B]
t_B = np.where(mask_B)[0]

fig, axes = plt.subplots(1, 2, figsize=(12, 4))

axes[0].scatter(t_B[a_B == 1], K_B[a_B == 1], s=3, color="#cc3333", label="good")
axes[0].scatter(t_B[a_B == 0], K_B[a_B == 0], s=3, color="#3333cc", label="bad")
axes[0].axhline(K_SS, color="k", ls="--", lw=1, label="det. steady state")
axes[0].set_xlabel("iteration")
axes[0].set_ylabel(r"aggregate capital $\bar K$")
axes[0].set_title(r"$\bar K$ along the simulated panel (Phase B)")
axes[0].legend()
axes[0].grid(alpha=0.3)

axes[1].hist(K_B[a_B == 1], bins=30, color="#cc3333", alpha=0.6, label="good")
axes[1].hist(K_B[a_B == 0], bins=30, color="#3333cc", alpha=0.6, label="bad")
axes[1].set_xlabel(r"$\bar K$")
axes[1].set_ylabel("frequency")
axes[1].set_title("Ergodic distribution of aggregate capital")
axes[1].legend()
axes[1].grid(alpha=0.3)

plt.tight_layout()
plt.show()

print(f"Good-state (Phase B):  mean K̄ = {K_B[a_B==1].mean():.2f},  std = {K_B[a_B==1].std():.3f}")
print(f"Bad-state  (Phase B):  mean K̄ = {K_B[a_B==0].mean():.2f},  std = {K_B[a_B==0].std():.3f}")
print(f"Deterministic steady state K* = {K_SS:.2f}")

# Phase-B-only trajectory
burn = 500
mask_late = (phase_history == 1)
idx_B = np.where(mask_late)[0]
K_seq = K_history[idx_B][burn:]
a_seq = a_history[idx_B][burn:]

logK = np.log(K_seq)
logK_t  = logK[:-1]
logK_tp = logK[1:]
a_t_seq = a_seq[:-1]

print("Law-of-motion regressions  log K̄_t+1 = A(a) + B(a) log K̄_t:")
print("-" * 72)
for a_val, lbl in [(0, "bad state "), (1, "good state")]:
    m = (a_t_seq == a_val)
    if m.sum() < 20:
        continue
    x, y = logK_t[m], logK_tp[m]
    # polyfit returns (slope, intercept)
    slope, intercept = np.polyfit(x, y, deg=1)
    y_hat = slope * x + intercept
    ss_res = float(np.sum((y - y_hat) ** 2))
    ss_tot = float(np.sum((y - y.mean()) ** 2) + 1e-30)
    r2 = 1.0 - ss_res / ss_tot
    print(f"  {lbl}:   slope B = {slope:+.5f},  intercept A = {intercept:+.5f},"
          f"   R² = {r2:.6f}   (n = {m.sum()})")
print("-" * 72)
print("Krusell--Smith approximate-aggregation target:  R² >= 0.99")

K_bar_good = K_B[a_B == 1].mean()
K_bar_bad  = K_B[a_B == 0].mean()

def eval_policy(k_grid, eps_int, K_agg, a_idx):
    N = len(k_grid)
    k_col    = k_grid.reshape(-1, 1)
    eps_1hot = np.tile(np.eye(2, dtype=np.float32)[eps_int], (N, 1))
    K_col    = np.full((N, 1), K_agg, dtype=np.float32)
    a_1hot   = np.tile(np.eye(2, dtype=np.float32)[a_idx], (N, 1))
    phi = policy_net([k_col, eps_1hot, K_col, a_1hot], training=False).numpy().flatten()
    z_val = z_levels[a_idx].numpy()
    u_val = u_rates[a_idx].numpy()
    L_val = 1.0 - u_val
    w = (1.0 - alpha.numpy()) * z_val * K_agg**alpha.numpy()     * L_val**(-alpha.numpy())
    r = alpha.numpy()         * z_val * K_agg**(alpha.numpy()-1) * L_val**(1.0 - alpha.numpy())
    y = w * float(l_bar.numpy()) * eps_int + (1.0 + r - float(delta.numpy())) * k_grid
    k_next = (1.0 - phi) * y
    c      = phi * y
    return phi, c, k_next, w, r

k_grid = np.linspace(0.05 * K_SS, 5.0 * K_SS, 200).astype(np.float32)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))

for ax, (eps_int, lbl) in zip(axes, [(1, "employed"), (0, "unemployed")]):
    _, _, kp_g, _, _ = eval_policy(k_grid, eps_int, K_bar_good, 1)
    _, _, kp_b, _, _ = eval_policy(k_grid, eps_int, K_bar_bad,  0)
    ax.plot(k_grid, kp_g, color="#cc3333", lw=2, label=rf"good state, $\bar K = {K_bar_good:.1f}$")
    ax.plot(k_grid, kp_b, color="#3333cc", lw=2, label=rf"bad state, $\bar K = {K_bar_bad:.1f}$")
    ax.plot(k_grid, k_grid, color="gray", ls="--", lw=1, label=r"$45°$")
    ax.set_xlabel("individual capital $k$")
    ax.set_ylabel("savings $k'$")
    ax.set_title(f"Learned savings policy — {lbl}")
    ax.legend(fontsize=9)
    ax.grid(alpha=0.3)

plt.tight_layout()
plt.show()