{
"cells": [
{
"cell_type": "markdown",
"id": "5cc960e7",
"metadata": {},
"source": [
"\n",
"<a id='optgrowth'></a>\n",
"<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
" <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
" <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
" </a>\n",
"</div>"
]
},
{
"cell_type": "markdown",
"id": "d54ff0be",
"metadata": {},
"source": [
"# Optimal Growth II: Accelerating the Code with Numba"
]
},
{
"cell_type": "markdown",
"id": "672453eb",
"metadata": {},
"source": [
"## Contents\n",
"\n",
"- [Optimal Growth II: Accelerating the Code with Numba](#Optimal-Growth-II:-Accelerating-the-Code-with-Numba) \n",
" - [Overview](#Overview) \n",
" - [The Model](#The-Model) \n",
" - [Computation](#Computation) \n",
" - [Exercises](#Exercises) "
]
},
{
"cell_type": "markdown",
"id": "940bfbb4",
"metadata": {},
"source": [
"In addition to what’s in Anaconda, this lecture will need the following libraries:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "13118123",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"!pip install quantecon"
]
},
{
"cell_type": "markdown",
"id": "5378a46f",
"metadata": {},
"source": [
"## Overview\n",
"\n",
"[Previously](https://python.quantecon.org/optgrowth.html), we studied a stochastic optimal\n",
"growth model with one representative agent.\n",
"\n",
"We solved the model using dynamic programming.\n",
"\n",
"In writing our code, we focused on clarity and flexibility.\n",
"\n",
"These are important, but there’s often a trade-off between flexibility and\n",
"speed.\n",
"\n",
"The reason is that, when code is less flexible, we can exploit structure more\n",
"easily.\n",
"\n",
"(This is true about algorithms and mathematical problems more generally:\n",
"more specific problems have more structure, which, with some thought, can be\n",
"exploited for better results.)\n",
"\n",
"So, in this lecture, we are going to accept less flexibility while gaining\n",
"speed, using just-in-time (JIT) compilation to\n",
"accelerate our code.\n",
"\n",
"Let’s start with some imports:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a767dd72",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"import matplotlib.pyplot as plt\n",
"import numpy as np\n",
"from numba import jit, jit\n",
"from quantecon.optimize.scalar_maximization import brent_max"
]
},
{
"cell_type": "markdown",
"id": "acf13381",
"metadata": {},
"source": [
"The function `brent_max` is also designed for embedding in JIT-compiled code.\n",
"\n",
"These are alternatives to similar functions in SciPy (which, unfortunately, are not JIT-aware)."
]
},
{
"cell_type": "markdown",
"id": "9c015d0e",
"metadata": {},
"source": [
"## The Model\n",
"\n",
"\n",
"<a id='index-1'></a>\n",
"The model is the same as discussed in our [previous lecture](https://python.quantecon.org/optgrowth.html)\n",
"on optimal growth.\n",
"\n",
"We will start with log utility:\n",
"\n",
"$$\n",
"u(c) = \\ln(c)\n",
"$$\n",
"\n",
"We continue to assume that\n",
"\n",
"- $ f(k) = k^{\\alpha} $ \n",
"- $ \\phi $ is the distribution of $ \\xi := \\exp(\\mu + s \\zeta) $ when $ \\zeta $ is standard normal \n",
"\n",
"\n",
"We will once again use value function iteration to solve the model.\n",
"\n",
"In particular, the algorithm is unchanged, and the only difference is in the implementation itself.\n",
"\n",
"As before, we will be able to compare with the true solutions"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f23c7193",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"\n",
"def v_star(y, α, β, μ):\n",
" \"\"\"\n",
" True value function\n",
" \"\"\"\n",
" c1 = np.log(1 - α * β) / (1 - β)\n",
" c2 = (μ + α * np.log(α * β)) / (1 - α)\n",
" c3 = 1 / (1 - β)\n",
" c4 = 1 / (1 - α * β)\n",
" return c1 + c2 * (c3 - c4) + c4 * np.log(y)\n",
"\n",
"def σ_star(y, α, β):\n",
" \"\"\"\n",
" True optimal policy\n",
" \"\"\"\n",
" return (1 - α * β) * y"
]
},
{
"cell_type": "markdown",
"id": "d9e2c877",
"metadata": {},
"source": [
"## Computation\n",
"\n",
"\n",
"<a id='index-2'></a>\n",
"We will again store the primitives of the optimal growth model in a class.\n",
"\n",
"But now we are going to use [Numba’s](https://python-programming.quantecon.org/numba.html) `@jitclass` decorator to target our class for JIT compilation.\n",
"\n",
"Because we are going to use Numba to compile our class, we need to specify the data types.\n",
"\n",
"You will see this as a list called `opt_growth_data` above our class.\n",
"\n",
"Unlike in the [previous lecture](https://python.quantecon.org/optgrowth.html), we\n",
"hardwire the production and utility specifications into the\n",
"class.\n",
"\n",
"This is where we sacrifice flexibility in order to gain more speed."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0d451c14",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"from numba import float64\n",
"from numba.experimental import jitclass\n",
"\n",
"opt_growth_data = [\n",
" ('α', float64), # Production parameter\n",
" ('β', float64), # Discount factor\n",
" ('μ', float64), # Shock location parameter\n",
" ('s', float64), # Shock scale parameter\n",
" ('grid', float64[:]), # Grid (array)\n",
" ('shocks', float64[:]) # Shock draws (array)\n",
"]\n",
"\n",
"@jitclass(opt_growth_data)\n",
"class OptimalGrowthModel:\n",
"\n",
" def __init__(self,\n",
" α=0.4,\n",
" β=0.96,\n",
" μ=0,\n",
" s=0.1,\n",
" grid_max=4,\n",
" grid_size=120,\n",
" shock_size=250,\n",
" seed=1234):\n",
"\n",
" self.α, self.β, self.μ, self.s = α, β, μ, s\n",
"\n",
" # Set up grid\n",
" self.grid = np.linspace(1e-5, grid_max, grid_size)\n",
"\n",
" # Store shocks (with a seed, so results are reproducible)\n",
" np.random.seed(seed)\n",
" self.shocks = np.exp(μ + s * np.random.randn(shock_size))\n",
"\n",
"\n",
" def f(self, k):\n",
" \"The production function\"\n",
" return k**self.α\n",
"\n",
"\n",
" def u(self, c):\n",
" \"The utility function\"\n",
" return np.log(c)\n",
"\n",
" def f_prime(self, k):\n",
" \"Derivative of f\"\n",
" return self.α * (k**(self.α - 1))\n",
"\n",
"\n",
" def u_prime(self, c):\n",
" \"Derivative of u\"\n",
" return 1/c\n",
"\n",
" def u_prime_inv(self, c):\n",
" \"Inverse of u'\"\n",
" return 1/c"
]
},
{
"cell_type": "markdown",
"id": "e4e282c0",
"metadata": {},
"source": [
"The class includes some methods such as `u_prime` that we do not need now\n",
"but will use in later lectures."
]
},
{
"cell_type": "markdown",
"id": "3bc3f1cc",
"metadata": {},
"source": [
"### The Bellman Operator\n",
"\n",
"We will use JIT compilation to accelerate the Bellman operator.\n",
"\n",
"First, here’s a function that returns the value of a particular consumption choice `c`, given state `y`, as per the Bellman equation [(40.9)](https://python.quantecon.org/optgrowth.html#equation-fpb30)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d7356b3",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"@jit\n",
"def state_action_value(c, y, v_array, og):\n",
" \"\"\"\n",
" Right hand side of the Bellman equation.\n",
"\n",
" * c is consumption\n",
" * y is income\n",
" * og is an instance of OptimalGrowthModel\n",
" * v_array represents a guess of the value function on the grid\n",
"\n",
" \"\"\"\n",
"\n",
" u, f, β, shocks = og.u, og.f, og.β, og.shocks\n",
"\n",
" v = lambda x: np.interp(x, og.grid, v_array)\n",
"\n",
" return u(c) + β * np.mean(v(f(y - c) * shocks))"
]
},
{
"cell_type": "markdown",
"id": "839d119e",
"metadata": {},
"source": [
"Now we can implement the Bellman operator, which maximizes the right hand side\n",
"of the Bellman equation:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5fdb5184",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"@jit\n",
"def T(v, og):\n",
" \"\"\"\n",
" The Bellman operator.\n",
"\n",
" * og is an instance of OptimalGrowthModel\n",
" * v is an array representing a guess of the value function\n",
"\n",
" \"\"\"\n",
"\n",
" v_new = np.empty_like(v)\n",
" v_greedy = np.empty_like(v)\n",
"\n",
" for i in range(len(og.grid)):\n",
" y = og.grid[i]\n",
"\n",
" # Maximize RHS of Bellman equation at state y\n",
" result = brent_max(state_action_value, 1e-10, y, args=(y, v, og))\n",
" v_greedy[i], v_new[i] = result[0], result[1]\n",
"\n",
" return v_greedy, v_new"
]
},
{
"cell_type": "markdown",
"id": "c8861af2",
"metadata": {},
"source": [
"We use the `solve_model` function to perform iteration until convergence."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b8ee5c57",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"def solve_model(og,\n",
" tol=1e-4,\n",
" max_iter=1000,\n",
" verbose=True,\n",
" print_skip=25):\n",
" \"\"\"\n",
" Solve model by iterating with the Bellman operator.\n",
"\n",
" \"\"\"\n",
"\n",
" # Set up loop\n",
" v = og.u(og.grid) # Initial condition\n",
" i = 0\n",
" error = tol + 1\n",
"\n",
" while i < max_iter and error > tol:\n",
" v_greedy, v_new = T(v, og)\n",
" error = np.max(np.abs(v - v_new))\n",
" i += 1\n",
" if verbose and i % print_skip == 0:\n",
" print(f\"Error at iteration {i} is {error}.\")\n",
" v = v_new\n",
"\n",
" if error > tol:\n",
" print(\"Failed to converge!\")\n",
" elif verbose:\n",
" print(f\"\\nConverged in {i} iterations.\")\n",
"\n",
" return v_greedy, v_new"
]
},
{
"cell_type": "markdown",
"id": "6ab74baf",
"metadata": {},
"source": [
"Let’s compute the approximate solution at the default parameters.\n",
"\n",
"First we create an instance:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "955a7341",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"og = OptimalGrowthModel()"
]
},
{
"cell_type": "markdown",
"id": "6a19b35b",
"metadata": {},
"source": [
"Now we call `solve_model`, using the `%%time` magic to check how long it\n",
"takes."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7ed32aad",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"%%time\n",
"v_greedy, v_solution = solve_model(og)"
]
},
{
"cell_type": "markdown",
"id": "67b21522",
"metadata": {},
"source": [
"You will notice that this is *much* faster than our [original implementation](https://python.quantecon.org/optgrowth.html).\n",
"\n",
"Here is a plot of the resulting policy, compared with the true policy:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "29748319",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"fig, ax = plt.subplots()\n",
"\n",
"ax.plot(og.grid, v_greedy, lw=2,\n",
" alpha=0.8, label='approximate policy function')\n",
"\n",
"ax.plot(og.grid, σ_star(og.grid, og.α, og.β), 'k--',\n",
" lw=2, alpha=0.8, label='true policy function')\n",
"\n",
"ax.legend()\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "71c597f3",
"metadata": {},
"source": [
"Again, the fit is excellent — this is as expected since we have not changed\n",
"the algorithm.\n",
"\n",
"The maximal absolute deviation between the two policies is"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fe8031f3",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"np.max(np.abs(v_greedy - σ_star(og.grid, og.α, og.β)))"
]
},
{
"cell_type": "markdown",
"id": "a50b8c91",
"metadata": {},
"source": [
"## Exercises"
]
},
{
"cell_type": "markdown",
"id": "a965e2f4",
"metadata": {},
"source": [
"## Exercise 41.1\n",
"\n",
"Time how long it takes to iterate with the Bellman operator\n",
"20 times, starting from initial condition $ v(y) = u(y) $.\n",
"\n",
"Use the default parameterization."
]
},
{
"cell_type": "markdown",
"id": "fdb873ea",
"metadata": {},
"source": [
"## Solution to[ Exercise 41.1](https://python.quantecon.org/#ogfast_ex1)\n",
"\n",
"Let’s set up the initial condition."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5676c612",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"v = og.u(og.grid)"
]
},
{
"cell_type": "markdown",
"id": "e189c7aa",
"metadata": {},
"source": [
"Here’s the timing:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4b364141",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"%%time\n",
"\n",
"for i in range(20):\n",
" v_greedy, v_new = T(v, og)\n",
" v = v_new"
]
},
{
"cell_type": "markdown",
"id": "6a418599",
"metadata": {},
"source": [
"Compared with our [timing](https://python.quantecon.org/optgrowth.html#og_ex2) for the non-compiled version of\n",
"value function iteration, the JIT-compiled code is usually an order of magnitude faster."
]
},
{
"cell_type": "markdown",
"id": "358b3861",
"metadata": {},
"source": [
"## Exercise 41.2\n",
"\n",
"Modify the optimal growth model to use the CRRA utility specification.\n",
"\n",
"$$\n",
"u(c) = \\frac{c^{1 - \\gamma} } {1 - \\gamma}\n",
"$$\n",
"\n",
"Set `γ = 1.5` as the default value and maintaining other specifications.\n",
"\n",
"(Note that `jitclass` currently does not support inheritance, so you will\n",
"have to copy the class and change the relevant parameters and methods.)\n",
"\n",
"Compute an estimate of the optimal policy, plot it and compare visually with\n",
"the same plot from the [analogous exercise](https://python.quantecon.org/optgrowth.html#og_ex1) in the first optimal\n",
"growth lecture.\n",
"\n",
"Compare execution time as well."
]
},
{
"cell_type": "markdown",
"id": "1c8bbe95",
"metadata": {},
"source": [
"## Solution to[ Exercise 41.2](https://python.quantecon.org/#ogfast_ex2)\n",
"\n",
"Here’s our CRRA version of `OptimalGrowthModel`:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4cf24e3f",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"from numba import float64\n",
"from numba.experimental import jitclass\n",
"\n",
"opt_growth_data = [\n",
" ('α', float64), # Production parameter\n",
" ('β', float64), # Discount factor\n",
" ('μ', float64), # Shock location parameter\n",
" ('γ', float64), # Preference parameter\n",
" ('s', float64), # Shock scale parameter\n",
" ('grid', float64[:]), # Grid (array)\n",
" ('shocks', float64[:]) # Shock draws (array)\n",
"]\n",
"\n",
"@jitclass(opt_growth_data)\n",
"class OptimalGrowthModel_CRRA:\n",
"\n",
" def __init__(self,\n",
" α=0.4,\n",
" β=0.96,\n",
" μ=0,\n",
" s=0.1,\n",
" γ=1.5,\n",
" grid_max=4,\n",
" grid_size=120,\n",
" shock_size=250,\n",
" seed=1234):\n",
"\n",
" self.α, self.β, self.γ, self.μ, self.s = α, β, γ, μ, s\n",
"\n",
" # Set up grid\n",
" self.grid = np.linspace(1e-5, grid_max, grid_size)\n",
"\n",
" # Store shocks (with a seed, so results are reproducible)\n",
" np.random.seed(seed)\n",
" self.shocks = np.exp(μ + s * np.random.randn(shock_size))\n",
"\n",
"\n",
" def f(self, k):\n",
" \"The production function.\"\n",
" return k**self.α\n",
"\n",
" def u(self, c):\n",
" \"The utility function.\"\n",
" return c**(1 - self.γ) / (1 - self.γ)\n",
"\n",
" def f_prime(self, k):\n",
" \"Derivative of f.\"\n",
" return self.α * (k**(self.α - 1))\n",
"\n",
" def u_prime(self, c):\n",
" \"Derivative of u.\"\n",
" return c**(-self.γ)\n",
"\n",
" def u_prime_inv(c):\n",
" return c**(-1 / self.γ)"
]
},
{
"cell_type": "markdown",
"id": "8c170fa1",
"metadata": {},
"source": [
"Let’s create an instance:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "979bd13a",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"og_crra = OptimalGrowthModel_CRRA()"
]
},
{
"cell_type": "markdown",
"id": "20a0e5e8",
"metadata": {},
"source": [
"Now we call `solve_model`, using the `%%time` magic to check how long it\n",
"takes."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "d2ecd45b",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"%%time\n",
"v_greedy, v_solution = solve_model(og_crra)"
]
},
{
"cell_type": "markdown",
"id": "ff985c82",
"metadata": {},
"source": [
"Here is a plot of the resulting policy:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0392cdbc",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"fig, ax = plt.subplots()\n",
"\n",
"ax.plot(og.grid, v_greedy, lw=2,\n",
" alpha=0.6, label='Approximate value function')\n",
"\n",
"ax.legend(loc='lower right')\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "674ba864",
"metadata": {},
"source": [
"This matches the solution that we obtained in our non-jitted code,\n",
"[in the exercises](https://python.quantecon.org/optgrowth.html#og_ex1).\n",
"\n",
"Execution time is an order of magnitude faster."
]
},
{
"cell_type": "markdown",
"id": "de485549",
"metadata": {},
"source": [
"## Exercise 41.3\n",
"\n",
"In this exercise we return to the original log utility specification.\n",
"\n",
"Once an optimal consumption policy $ \\sigma $ is given, income follows\n",
"\n",
"$$\n",
"y_{t+1} = f(y_t - \\sigma(y_t)) \\xi_{t+1}\n",
"$$\n",
"\n",
"The next figure shows a simulation of 100 elements of this sequence for three\n",
"different discount factors (and hence three different policies).\n",
"\n",
"\n",
"\n",
" \n",
"In each sequence, the initial condition is $ y_0 = 0.1 $.\n",
"\n",
"The discount factors are `discount_factors = (0.8, 0.9, 0.98)`.\n",
"\n",
"We have also dialed down the shocks a bit with `s = 0.05`.\n",
"\n",
"Otherwise, the parameters and primitives are the same as the log-linear model discussed earlier in the lecture.\n",
"\n",
"Notice that more patient agents typically have higher wealth.\n",
"\n",
"Replicate the figure modulo randomness."
]
},
{
"cell_type": "markdown",
"id": "65d4b597",
"metadata": {},
"source": [
"## Solution to[ Exercise 41.3](https://python.quantecon.org/#ogfast_ex3)\n",
"\n",
"Here’s one solution:"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "55b99752",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"def simulate_og(σ_func, og, y0=0.1, ts_length=100):\n",
" '''\n",
" Compute a time series given consumption policy σ.\n",
" '''\n",
" y = np.empty(ts_length)\n",
" ξ = np.random.randn(ts_length-1)\n",
" y[0] = y0\n",
" for t in range(ts_length-1):\n",
" y[t+1] = (y[t] - σ_func(y[t]))**og.α * np.exp(og.μ + og.s * ξ[t])\n",
" return y"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f0019390",
"metadata": {
"hide-output": false
},
"outputs": [],
"source": [
"fig, ax = plt.subplots()\n",
"\n",
"for β in (0.8, 0.9, 0.98):\n",
"\n",
" og = OptimalGrowthModel(β=β, s=0.05)\n",
"\n",
" v_greedy, v_solution = solve_model(og, verbose=False)\n",
"\n",
" # Define an optimal policy function\n",
" σ_func = lambda x: np.interp(x, og.grid, v_greedy)\n",
" y = simulate_og(σ_func, og)\n",
" ax.plot(y, lw=2, alpha=0.6, label=rf'$\\beta = {β}$')\n",
"\n",
"ax.legend(loc='lower right')\n",
"plt.show()"
]
}
],
"metadata": {
"date": 1750302297.2459168,
"filename": "optgrowth_fast.md",
"kernelspec": {
"display_name": "Python",
"language": "python3",
"name": "python3"
},
"title": "Optimal Growth II: Accelerating the Code with Numba"
},
"nbformat": 4,
"nbformat_minor": 5
}