Week 7: Physics-Informed Neural Networks

Deep Learning for Macroeconomics — Honours, The University of Edinburgh
Instructor: Juan Zurita · juan.zurita@ed.ac.uk

---

Learning objectives

By the end of this notebook you will be able to:

1. Explain the PINN principle: minimise the PDE/ODE residual over collocation points
2. Solve a simple ODE boundary-value problem with a PINN
3. Apply PINNs to the cake-eating HJB equation — a continuous-time consumption problem
4. Discuss soft vs hard boundary conditions and their impact on training

---

# Setup
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

torch.manual_seed(42)
np.random.seed(42)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"PyTorch {torch.__version__} on {device}")

1. From discrete to continuous time

So far we’ve worked in discrete time: \(k_{t+1} = f(k_t) - c_t\). Many models in macro and finance are naturally formulated in continuous time:

\[\frac{dk}{dt} = f(k) - c, \quad \rho V(k) = \max_c \Big\{ u(c) + V'(k) [f(k) - c] \Big\}\]

The second equation is the Hamilton–Jacobi–Bellman (HJB) equation — a PDE that the value function must satisfy.

The PINN idea

A Physics-Informed Neural Network parameterises the solution \(V(k; \theta)\) as a neural network and minimises:

\[\mathcal{L}(\theta) = \frac{1}{N} \sum_{i=1}^{N} \Big| \rho V(k_i; \theta) - \max_c \big\{ u(c) + V'(k_i; \theta) [f(k_i) - c] \big\} \Big|^2\]

The key: \(V'(k; \theta)\) is computed via automatic differentiation — differentiating the network output with respect to its input \(k\).

2. Warm-up: solving an ODE with a PINN

Before tackling the HJB equation, let’s solve a simple ODE:

\[y''(x) + y(x) = 0, \quad y(0) = 0, \quad y(\pi/2) = 1\]

The analytical solution is \(y(x) = \sin(x)\).

# PINN for a simple ODE: y'' + y = 0

class PINN_ODE(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(1, 32), nn.Tanh(),
            nn.Linear(32, 32), nn.Tanh(),
            nn.Linear(32, 1)
        )
    
    def forward(self, x):
        return self.net(x)

model = PINN_ODE()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Collocation points (where we enforce the ODE)
x_col = torch.linspace(0, np.pi/2, 50).reshape(-1, 1).requires_grad_(True)

# Boundary points
x_bc0 = torch.tensor([[0.0]])
x_bc1 = torch.tensor([[np.pi/2]])

for epoch in range(5000):
    # ODE residual: y'' + y = 0
    y = model(x_col)
    dy = torch.autograd.grad(y, x_col, torch.ones_like(y), create_graph=True)[0]
    d2y = torch.autograd.grad(dy, x_col, torch.ones_like(dy), create_graph=True)[0]
    ode_residual = d2y + y
    
    # Boundary conditions
    bc0_loss = (model(x_bc0) - 0.0)**2
    bc1_loss = (model(x_bc1) - 1.0)**2
    
    # Total loss
    loss = torch.mean(ode_residual**2) + 100 * (bc0_loss + bc1_loss)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Plot
model.eval()
x_test = torch.linspace(0, np.pi/2, 200).reshape(-1, 1)
with torch.no_grad():
    y_pinn = model(x_test).numpy().flatten()

fig, ax = plt.subplots(figsize=(8, 4))
ax.plot(x_test.numpy(), np.sin(x_test.numpy()), 'k-', linewidth=2, label='Analytical: sin(x)')
ax.plot(x_test.numpy(), y_pinn, '--', color='#7a2318', linewidth=2, label='PINN solution')
ax.set_xlabel('x'); ax.set_ylabel('y(x)')
ax.set_title('PINN Solution to y\'\' + y = 0'); ax.legend(); ax.grid(True, alpha=0.3)
plt.tight_layout(); plt.show()
print(f"Max error: {np.max(np.abs(y_pinn - np.sin(x_test.numpy().flatten()))):.6f}")

3. The cake-eating HJB equation

A consumer with a cake of size \(k\) solves:

\[\rho V(k) = \max_c \Big\{ \ln(c) + V'(k) (-c) \Big\}\]

The FOC gives \(c^*(k) = 1/V'(k)\), substituting back:

\[\rho V(k) = \ln\left(\frac{1}{V'(k)}\right) - 1\]

This is a nonlinear ODE for \(V(k)\) — perfect for a PINN.

4. Soft vs hard boundary conditions

Two approaches:

Soft BCs (penalty method): add \(\lambda \cdot (\text{BC violation})^2\) to the loss. Simple but requires tuning \(\lambda\).

Hard BCs (construction): write the network output so BCs are satisfied by construction: \[\hat{y}(x) = A(x) + B(x) \cdot \text{NN}(x)\] where \(A(x)\) satisfies the BCs and \(B(x) = 0\) at the boundaries.

Hard BCs are preferred when possible — they guarantee exact boundary satisfaction and make training more stable.

---

Exercises

Exercise 1: Solve the cake-eating HJB with a PINN. Verify against the analytical solution \(V(k) = \frac{1}{\rho}[\ln(\rho k) - 1]\).

Exercise 2: Solve the Black–Scholes PDE: \(\frac{\partial V}{\partial t} + \frac{1}{2}\sigma^2 S^2 \frac{\partial^2 V}{\partial S^2} + rS\frac{\partial V}{\partial S} - rV = 0\)

---

Next week: Heterogeneous Agents — the frontier where deep learning really shines.