Health Investment Model with EGM^n + ENGINE

This notebook applies EGM $^n$ and ENGINE to the health investment model from White (2015). EGM $^n$ decomposes multi-decision problems into sequential stages, and ENGINE interpolates over the resulting curvilinear grids.

from __future__ import annotations

import sys
sys.path.insert(0, "..")

from time import time
import matplotlib.pyplot as plt
import numpy as np

from ConsHealthSeparableModel import (
    SequentialHealthConsumerType, 
    BasicHealthConsumerType as OurSimulType
)
from utilities import Engine2DInterp
from engine2d import warmup_jit
from HARK.interpolation import Curvilinear2DInterp
from HARK.utilities import plot_funcs

1The Health Investment Model¶

The agent chooses consumption $c_t$ and health investment $n_t$ each period:

v_t(m_t, h_t) = \max_{c_t, n_t} \frac{c_t^{1-\rho}}{1-\rho} + \beta S_t \mathbb{E}[v_{t+1}(m_{t+1}, h_{t+1})]

(1)

where:

$H_t = h_t + g(n_t)$ is post-investment health
$S_t = 1 - \phi_t/(1 + H_t)$ is survival probability
Health provides both survival benefits and higher wage income

2EGM $^n$ Decomposition¶

The sequential approach breaks the problem into stages:

Post-decision stage: compute continuation value $\mathfrak{v}_t(a_t, H_t)$
Consumption stage: invert FOC $u'(c) = \mathfrak{v}_t^a$ to get $c_t$
Health investment stage: invert FOC $g'(n) \cdot \mathfrak{v}_t^H = \mathfrak{v}_t^a$ to get $n_t$

Each stage uses the Endogenous Grid Method, avoiding numerical optimization.

# Warm up JIT compilation for accurate timing
print("Warming up ENGINE JIT compilation...")
warmup_jit()
print("Done!")

Warming up ENGINE JIT compilation...
Done!

3Single-Period Solution¶

Let’s examine the endogenous grid structure from solving one period:

OnePeriodExample = SequentialHealthConsumerType(cycles=1, interpolator=Engine2DInterp)
OnePeriodExample.solve()

# Plot the endogenous grid
X = OnePeriodExample.solution[0].fx
Y = OnePeriodExample.solution[0].fy

plt.figure(figsize=(8, 6))
plt.plot(X.flatten(), Y.flatten(), ".k", ms=2)
plt.xlim(0.0, 100.0)
plt.ylim(0.0, 50.0)
plt.xlabel(r"Market resources $m_t$")
plt.ylabel(r"Health capital $h_t$")
plt.title(r"Endogenous Grid from EGM$^n$ + ENGINE")
plt.show()

4Timing Comparison¶

Let’s solve 100 periods and compare with HARK’s original solver:

# Sequential EGM + ENGINE
SequentialEngineExample = SequentialHealthConsumerType(
    cycles=99, interpolator=Engine2DInterp
)

t0 = time()
SequentialEngineExample.solve()
t1 = time()
seq_engine_time = t1 - t0

print(f"EGM^n + ENGINE solve time: {seq_engine_time:.2f} seconds")

EGM^n + ENGINE solve time: 3.32 seconds

# HARK's original solver
from HARK.ConsumptionSaving.ConsHealthModel import BasicHealthConsumerType as HARKHealthType

HARKExample = HARKHealthType(cycles=99)

t0 = time()
HARKExample.solve()
t1 = time()
hark_time = t1 - t0

print(f"HARK original solve time: {hark_time:.2f} seconds")
print(f"Speedup: {hark_time / seq_engine_time:.1f}×")

HARK original solve time: 8.13 seconds
Speedup: 2.4×

5Solution Accuracy¶

Let’s verify that the policies are equivalent:

# Plot grid from 100-period solution
X = SequentialEngineExample.solution[0].fx
Y = SequentialEngineExample.solution[0].fy

plt.figure(figsize=(8, 6))
plt.plot(X.flatten(), Y.flatten(), ".k", ms=2)
plt.xlim(0.0, 100.0)
plt.ylim(0.0, 50.0)
plt.xlabel(r"Market resources $m_t$")
plt.ylabel(r"Health capital $h_t$")
plt.title(r"Endogenous Grid at $t=0$ (100-period model)")
plt.show()

Policy Comparison¶

# Define policy cross-sections
t = 0
hLvl = 20.0

def C_engine(mLvl):
    return SequentialEngineExample.solution[t](mLvl, hLvl * np.ones_like(mLvl))[1]

def N_engine(mLvl):
    return SequentialEngineExample.solution[t](mLvl, hLvl * np.ones_like(mLvl))[2]

def C_hark(mLvl):
    return HARKExample.solution[t](mLvl, hLvl * np.ones_like(mLvl))[1]

def N_hark(mLvl):
    return HARKExample.solution[t](mLvl, hLvl * np.ones_like(mLvl))[2]

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

mLvl = np.linspace(0.1, 20.0, 100)

axes[0].plot(mLvl, C_engine(mLvl), "b-", lw=2, label=r"EGM$^n$ + ENGINE")
axes[0].plot(mLvl, C_hark(mLvl), "r--", lw=2, label="HARK original")
axes[0].set_xlabel(r"Market resources $m_t$")
axes[0].set_ylabel(r"Consumption $c_t$")
axes[0].set_title("Consumption Policy")
axes[0].legend()

axes[1].plot(mLvl, N_engine(mLvl), "b-", lw=2, label=r"EGM$^n$ + ENGINE")
axes[1].plot(mLvl, N_hark(mLvl), "r--", lw=2, label="HARK original")
axes[1].set_xlabel(r"Market resources $m_t$")
axes[1].set_ylabel(r"Health investment $n_t$")
axes[1].set_title("Health Investment Policy")
axes[1].legend()

plt.tight_layout()
plt.show()

Numerical Differences¶

# Compute differences on a grid
mLvl_test = np.linspace(1.0, 50.0, 200)
hLvl_test = np.linspace(1.0, 40.0, 200)
mLvl_grid, hLvl_grid = np.meshgrid(mLvl_test, hLvl_test, indexing="ij")

_, c_engine, n_engine = SequentialEngineExample.solution[0](mLvl_grid, hLvl_grid)
_, c_hark, n_hark = HARKExample.solution[0](mLvl_grid, hLvl_grid)

c_diff = np.abs(c_engine - c_hark)
n_diff = np.abs(n_engine - n_hark)

print("=== Policy Differences ===")
print(f"Consumption: max={np.max(c_diff):.2e}, mean={np.mean(c_diff):.2e}")
print(f"Health inv:  max={np.max(n_diff):.2e}, mean={np.mean(n_diff):.2e}")

=== Policy Differences ===
Consumption: max=2.60e-03, mean=1.41e-03
Health inv:  max=4.73e-03, mean=1.61e-03

Verifying Solver Equivalence¶

# Our simultaneous solver with HARK's interpolator should match exactly
OurCurvi = OurSimulType(cycles=99, interpolator=Curvilinear2DInterp)
OurCurvi.solve()
_, c_ours, n_ours = OurCurvi.solution[0](mLvl_grid, hLvl_grid)

print("\n=== Our Simul+Curvilinear vs HARK original ===")
print(f"Consumption max diff: {np.max(np.abs(c_ours - c_hark)):.2e}")
print(f"Health inv max diff:  {np.max(np.abs(n_ours - n_hark)):.2e}")
print("\n→ Solvers are mathematically equivalent!")
print("→ Any differences with ENGINE are purely interpolation method.")


=== Our Simul+Curvilinear vs HARK original ===
Consumption max diff: 1.30e-04
Health inv max diff:  1.44e-03

→ Solvers are mathematically equivalent!
→ Any differences with ENGINE are purely interpolation method.

EGM $^n$ produces the same policies as simultaneous optimization. The method applies to any problem with separable structure; see the paper for the mathematical details.