cass_koopmans.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "<a id='cass-koopmans'></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",
   "metadata": {},
   "source": [
    "# Cass-Koopmans Optimal Growth Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contents\n",
    "\n",
    "- [Cass-Koopmans Optimal Growth Model](#Cass-Koopmans-Optimal-Growth-Model)  \n",
    "  - [Overview](#Overview)  \n",
    "  - [The Growth Model](#The-Growth-Model)  \n",
    "  - [Competitive Equilibrium](#Competitive-Equilibrium)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Coauthor: Brandon Kaplowitz**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "This lecture describes a model that Tjalling Koopmans [[Koo65]](https://lectures.quantecon.org/zreferences.html#koopmans)\n",
    "and David Cass [[Cas65]](https://lectures.quantecon.org/zreferences.html#cass) used to analyze optimal growth.\n",
    "\n",
    "The model can be viewed as an extension of the model of Robert Solow\n",
    "described in [an earlier lecture](https://lectures.quantecon.org/py/python_oop.html)\n",
    "but adapted to make the savings rate the outcome of an optimal choice.\n",
    "\n",
    "(Solow assumed a constant saving rate determined outside the model).\n",
    "\n",
    "We describe two versions of the model to illustrate what is, in fact, a\n",
    "more general connection between a **planned economy** and an economy\n",
    "organized as a **competitive equilibrium**.\n",
    "\n",
    "The lecture uses important ideas including\n",
    "\n",
    "- Hicks-Arrow prices named after John R. Hicks and Kenneth Arrow.  \n",
    "- A min-max problem for solving a planning problem.  \n",
    "- A **shooting algorithm** for solving difference equations subject\n",
    "  to initial and terminal conditions.  \n",
    "- A connection between some Lagrange multipliers in the min-max\n",
    "  problem and the Hicks-Arrow prices.  \n",
    "- A **Big** $ K $ **, little** $ k $ trick widely used in\n",
    "  macroeconomic dynamics.  \n",
    "\n",
    "\n",
    "> - We shall encounter this trick in [this lecture](https://lectures.quantecon.org/py/rational_expectations.html#)\n",
    "  and also in [this lecture](https://lectures.quantecon.org/py/dyn_stack.html#).  \n",
    "\n",
    "\n",
    "\n",
    "- An application of a **guess and verify** method for solving a\n",
    "  system of difference equations.  \n",
    "- The intimate connection between the cases for the optimality of two\n",
    "  competing visions of good ways to organize an economy, namely:  \n",
    "\n",
    "\n",
    "> - **socialism** in which a central planner commands the\n",
    "  allocation of resources, and  \n",
    "- **capitalism** (also known as **a free markets economy**) in\n",
    "  which competitive equilibrium **prices** induce individual\n",
    "  consumers and producers to choose a socially optimal allocation\n",
    "  as an unintended consequence of their completely selfish\n",
    "  decisions  \n",
    "\n",
    "\n",
    "\n",
    "- A **turnpike** property that describes optimal paths for\n",
    "  long-but-finite horizon economies.  \n",
    "- A non-stochastic version of a theory of the **term structure of\n",
    "  interest rates**.  \n",
    "\n",
    "\n",
    "Let’s start with some standard imports:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numba import njit\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Growth Model\n",
    "\n",
    "Time is discrete and takes values $ t = 0, 1 , \\ldots, T $.\n",
    "\n",
    "(We leave open the possibility that $ T = + \\infty $, but that\n",
    "will require special care in interpreting and using a **terminal\n",
    "condition** on $ K_t $ at $ t = T+1 $ to be described\n",
    "below).\n",
    "\n",
    "A single good can either be consumed or invested in physical capital.\n",
    "\n",
    "The consumption good is not durable and depreciates completely if not\n",
    "consumed immediately.\n",
    "\n",
    "The capital good is durable but depreciates each period at rate\n",
    "$ \\delta \\in (0,1) $.\n",
    "\n",
    "We let $ C_t $ be a nondurable consumption good at time t.\n",
    "\n",
    "Let $ K_t $ be the stock of physical capital at time t.\n",
    "\n",
    "Let $ \\vec{C} $ = $ \\{C_0,\\dots, C_T\\} $ and\n",
    "$ \\vec{K} $ = $ \\{K_1,\\dots,K_{T+1}\\} $.\n",
    "\n",
    "A representative household is endowed with one unit of labor at each\n",
    "$ t $ and likes the consumption good at each $ t $.\n",
    "\n",
    "The representative household inelastically supplies a single unit of\n",
    "labor $ N_t $ at each $ t $, so that\n",
    "$ N_t =1 \\text{ for all } t \\in [0,T] $.\n",
    "\n",
    "The representative household has preferences over consumption bundles\n",
    "ordered by the utility functional:\n",
    "\n",
    "\n",
    "<a id='equation-utility-functional'></a>\n",
    "$$\n",
    "U(\\vec{C}) = \\sum_{t=0}^{T} \\beta^t \\frac{C_t^{1-\\gamma}}{1-\\gamma} \\tag{1}\n",
    "$$\n",
    "\n",
    "where $ \\beta \\in (0,1) $ is a discount factor and $ \\gamma >0 $\n",
    "governs the curvature of the one-period utility function.\n",
    "\n",
    "Note that\n",
    "\n",
    "\n",
    "<a id='equation-utility-oneperiod'></a>\n",
    "$$\n",
    "u(C_t) = \\frac{C_t^{1-\\gamma}}{1-\\gamma} \\tag{2}\n",
    "$$\n",
    "\n",
    "satisfies $ u'>0,u''<0 $.\n",
    "\n",
    "$ u' > 0 $ asserts the consumer prefers more to less.\n",
    "\n",
    "$ u''< 0 $ asserts that marginal utility declines with increases\n",
    "in $ C_t $.\n",
    "\n",
    "We assume that $ K_0 > 0 $ is a given exogenous level of initial\n",
    "capital.\n",
    "\n",
    "There is an economy-wide production function\n",
    "\n",
    "\n",
    "<a id='equation-production-function'></a>\n",
    "$$\n",
    "F(K_t,N_t) = A K_t^{\\alpha}N_t^{1-\\alpha} \\tag{3}\n",
    "$$\n",
    "\n",
    "with $ 0 < \\alpha<1 $, $ A > 0 $.\n",
    "\n",
    "A feasible allocation $ \\vec C, \\vec K $ satisfies\n",
    "\n",
    "\n",
    "<a id='equation-allocation'></a>\n",
    "$$\n",
    "C_t + K_{t+1} \\leq F(K_t,N_t) + (1-\\delta) K_t, \\quad \\text{for all } t \\in [0, T] \\tag{4}\n",
    "$$\n",
    "\n",
    "where $ \\delta \\in (0,1) $ is a depreciation rate of capital."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Planning Problem\n",
    "\n",
    "A planner chooses an allocation $ \\{\\vec{C},\\vec{K}\\} $ to\n",
    "maximize [(1)](#equation-utility-functional) subject to [(4)](#equation-allocation).\n",
    "\n",
    "Let $ \\vec{\\mu}=\\{\\mu_0,\\dots,\\mu_T\\} $ be a sequence of\n",
    "nonnegative **Lagrange multipliers**.\n",
    "\n",
    "To find an optimal allocation, we form a Lagrangian\n",
    "\n",
    "$$\n",
    "\\mathcal{L}(\\vec{C},\\vec{K},\\vec{\\mu}) =\n",
    "\\sum_{t=0}^T \\beta^t\\left\\{ u(C_t)+ \\mu_t\n",
    "\\left(F(K_t,1) + (1-\\delta) K_t- C_t - K_{t+1} \\right)\\right\\}\n",
    "$$\n",
    "\n",
    "and then solve the following min-max problem:\n",
    "\n",
    "\n",
    "<a id='equation-min-max-prob'></a>\n",
    "$$\n",
    "\\min_{\\vec{\\mu}} \\max_{\\vec{C},\\vec{K}} \\mathcal{L}(\\vec{C},\\vec{K},\\vec{\\mu}) \\tag{5}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Useful Properties of Linearly Homogeneous Production Function\n",
    "\n",
    "The following technicalities will help us.\n",
    "\n",
    "Notice that\n",
    "\n",
    "$$\n",
    "F(K_t,N_t) = A K_t^\\alpha N_t^{1-\\alpha} = N_t A\\left(\\frac{K_t}{N_t}\\right)^\\alpha\n",
    "$$\n",
    "\n",
    "Define the **output per-capita production function**\n",
    "\n",
    "$$\n",
    "f\\left(\\frac{K_t}{N_t}\\right) = A\\left(\\frac{K_t}{N_t}\\right)^\\alpha\n",
    "$$\n",
    "\n",
    "whose argument is **capital per-capita**.\n",
    "\n",
    "Evidently,\n",
    "\n",
    "$$\n",
    "F(K_t,N_t)=N_t f \\left( \\frac{K_t}{N_t} \\right)\n",
    "$$\n",
    "\n",
    "Now for some useful calculations.\n",
    "\n",
    "First\n",
    "\n",
    "\n",
    "<a id='equation-useful-calc1'></a>\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\frac{\\partial F}{\\partial K}\n",
    "&=\n",
    "\\frac{\\partial N_t f\\left( \\frac{K_t}{N_t}\\right)}{\\partial K_t}\n",
    "\\\\ &=\n",
    "N_t f'\\left(\\frac{K_t}{N_t}\\right)\\frac{1}{N_t} \\quad \\text{(Chain rule)}\n",
    "\\\\ &=\n",
    "f'\\left.\\left(\\frac{K_t}{N_t}\\right)\\right|_{N_t=1}\n",
    "\\\\ &= f'(K_t)\n",
    "\\end{aligned} \\tag{6}\n",
    "$$\n",
    "\n",
    "Also\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "\\frac{\\partial F}{\\partial N}\n",
    "&=\n",
    "\\frac{\\partial N_t f\\left( \\frac{K_t}{N_t}\\right)}{\\partial N_t} \\quad \\text{(Product rule)}\n",
    "\\\\ &=\n",
    "f\\left(\\frac{K_t}{N_t}\\right){+} N_t f'\\left(\\frac{K_t}{N_t}\\right) \\frac{-K_t}{N_t^2} \\quad \\text{(Chain rule)}\n",
    "\\\\ &=\n",
    "f\\left(\\frac{K_t}{N_t}\\right){-}\\frac{K_t}{N_t}f'\\left.\\left(\\frac{K_t}{N_t}\\right)\\right|_{N_t=1}\n",
    "\\\\ &=\n",
    "f(K_t) - f'(K_t) K_t\n",
    "\\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Back to Solving the Problem\n",
    "\n",
    "To solve the Lagrangian extremization problem, we compute first\n",
    "derivatives of the Lagrangian and set them equal to 0.\n",
    "\n",
    "- **Note:** Our objective function and constraints satisfy\n",
    "  conditions that work to assure that required second-order\n",
    "  conditions are satisfied at an allocation that satisfies the\n",
    "  first-order conditions that we are about to compute.  \n",
    "\n",
    "\n",
    "Here are the **first order necessary conditions** for extremization\n",
    "(i.e., maximization with respect to $ \\vec C, \\vec K $,\n",
    "minimization with respect to $ \\vec \\mu $):\n",
    "\n",
    "\n",
    "<a id='equation-constraint1'></a>\n",
    "$$\n",
    "C_t: \\qquad u'(C_t)-\\mu_t=0 \\qquad \\text{for all} \\quad t= 0,1,\\dots,T \\tag{7}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-constraint2'></a>\n",
    "$$\n",
    "K_t: \\qquad \\beta \\mu_t\\left[(1-\\delta)+f'(K_t)\\right] - \\mu_{t-1}=0 \\qquad \\text{for all } \\quad t=1,2,\\dots,T \\tag{8}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-constraint3'></a>\n",
    "$$\n",
    "\\mu_t:\\qquad F(K_t,1)+ (1-\\delta) K_t  - C_t - K_{t+1}=0 \\qquad \\text{for all } \\quad t=0,1,\\dots,T \\tag{9}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-constraint4'></a>\n",
    "$$\n",
    "K_{T+1}: \\qquad -\\mu_T \\leq 0, \\ \\leq 0 \\text{ if } K_{T+1}=0; \\ =0 \\text{ if } K_{T+1}>0 \\tag{10}\n",
    "$$\n",
    "\n",
    "Note that in [(8)](#equation-constraint2) we plugged in for\n",
    "$ \\frac{\\partial F}{\\partial K} $ using our formula [(6)](#equation-useful-calc1)\n",
    "above.\n",
    "\n",
    "Because $ N_t = 1 $ for $ t = 1, \\ldots, T $, need not\n",
    "differentiate with respect to those arguments.\n",
    "\n",
    "Note that [(9)](#equation-constraint3) comes from the occurrence\n",
    "of $ K_t $ in both the period $ t $ and period $ t-1 $\n",
    "feasibility constraints.\n",
    "\n",
    "[(10)](#equation-constraint4) comes from differentiating with respect\n",
    "to $ K_{T+1} $ in the last period and applying the following\n",
    "condition called a **Karush-Kuhn-Tucker condition** (KKT):\n",
    "\n",
    "\n",
    "<a id='equation-kkt'></a>\n",
    "$$\n",
    "\\mu_T K_{T+1}=0 \\tag{11}\n",
    "$$\n",
    "\n",
    "See [Karush-Kuhn-Tucker conditions](https://en.wikipedia.org/wiki/Karush-Kuhn-Tucker_conditions).\n",
    "\n",
    "Combining [(7)](#equation-constraint1) and [(8)](#equation-constraint2) gives\n",
    "\n",
    "$$\n",
    "u'\\left(C_t\\right)\\left[(1-\\delta)+f'\\left(K_t\\right)\\right]-u'\\left(C_{t-1}\\right)=0\n",
    "\\quad \\text{ for all } t=1,2,\\dots, T+1\n",
    "$$\n",
    "\n",
    "Rewriting gives\n",
    "\n",
    "\n",
    "<a id='equation-l12'></a>\n",
    "$$\n",
    "u'\\left(C_{t+1}\\right)\\left[(1-\\delta)+f'\\left(K_{t+1}\\right)\\right]=\n",
    "u'\\left(C_{t}\\right) \\quad \\text{ for all } t=0,1,\\dots, T \\tag{12}\n",
    "$$\n",
    "\n",
    "Taking the inverse of the utility function on both sides of the above\n",
    "equation gives\n",
    "\n",
    "$$\n",
    "C_{t+1} =u'^{-1}\\left(\\left(\\frac{\\beta}{u'(C_t)}[f'(K_{t+1}) +(1-\\delta)]\\right)^{-1}\\right)\n",
    "$$\n",
    "\n",
    "or using our utility function [(2)](#equation-utility-oneperiod)\n",
    "\n",
    "$$\n",
    "\\begin{aligned} C_{t+1} =\\left(\\beta C_t^{\\gamma}[f'(K_{t+1}) +\n",
    "(1-\\delta)]\\right)^{1/\\gamma} \\notag\\\\= C_t\\left(\\beta [f'(K_{t+1}) +\n",
    "(1-\\delta)]\\right)^{1/\\gamma} \\end{aligned}\n",
    "$$\n",
    "\n",
    "The above first-order condition for consumption is called an **Euler\n",
    "equation**.\n",
    "\n",
    "It tells us how consumption in adjacent periods are optimally related\n",
    "to each other and to capital next period.\n",
    "\n",
    "We now use some of the equations above to calculate some\n",
    "variables and functions that we’ll soon use to solve the planning\n",
    "problem with Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def u(c, γ):\n",
    "    '''\n",
    "    Utility function\n",
    "    ASIDE: If you have a utility function that is hard to solve by hand\n",
    "    you can use automatic or symbolic  differentiation\n",
    "    See https://github.com/HIPS/autograd\n",
    "    '''\n",
    "    if γ == 1:\n",
    "        # If γ = 1 we can show via L'hopital's Rule that the utility\n",
    "        # becomes log\n",
    "        return np.log(c)\n",
    "    else:\n",
    "        return c**(1 - γ) / (1 - γ)\n",
    "\n",
    "@njit\n",
    "def u_prime(c, γ):\n",
    "    '''Derivative of utility'''\n",
    "    if γ == 1:\n",
    "        return 1 / c\n",
    "    else:\n",
    "        return c**(-γ)\n",
    "\n",
    "@njit\n",
    "def u_prime_inv(c, γ):\n",
    "    '''Inverse utility'''\n",
    "    if γ == 1:\n",
    "        return c\n",
    "    else:\n",
    "        return c**(-1 / γ)\n",
    "\n",
    "@njit\n",
    "def f(A, k, α):\n",
    "    '''Production function'''\n",
    "    return A * k**α\n",
    "\n",
    "@njit\n",
    "def f_prime(A, k, α):\n",
    "    '''Derivative of production function'''\n",
    "    return α * A * k**(α - 1)\n",
    "\n",
    "@njit\n",
    "def f_prime_inv(A, k, α):\n",
    "    return (k / (A * α))**(1 / (α - 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Shooting Method\n",
    "\n",
    "We shall use a **shooting method** to compute an optimal allocation\n",
    "$ \\vec C, \\vec K $ and an associated Lagrange multiplier sequence\n",
    "$ \\vec \\mu $.\n",
    "\n",
    "The first-order necessary conditions for the planning problem,\n",
    "namely, equations [(7)](#equation-constraint1), [(8)](#equation-constraint2), and\n",
    "[(9)](#equation-constraint3), form a system of **difference equations** with\n",
    "two boundary conditions:\n",
    "\n",
    "- $ K_0 $ is a given **initial condition** for capital  \n",
    "- $ K_{T+1} =0 $ is a **terminal condition** for capital that we\n",
    "  deduced from the first-order necessary condition for $ K_{T+1} $\n",
    "  the KKT condition [(11)](#equation-kkt)  \n",
    "\n",
    "\n",
    "We have no initial condition for the Lagrange multiplier\n",
    "$ \\mu_0 $.\n",
    "\n",
    "If we did, solving for the allocation would be simple:\n",
    "\n",
    "- Given $ \\mu_0 $ and $ k_0 $, we could compute $ c_0 $ from\n",
    "  equation [(7)](#equation-constraint1) and then $ k_1 $ from equation\n",
    "  [(9)](#equation-constraint3) and $ \\mu_1 $ from equation\n",
    "  [(8)](#equation-constraint2).  \n",
    "- We could then iterate on to compute the remaining elements of\n",
    "  $ \\vec C, \\vec K, \\vec \\mu $.  \n",
    "\n",
    "\n",
    "But we don’t have an initial condition for $ \\mu_0 $, so this\n",
    "won’t work.\n",
    "\n",
    "But a simple modification called the **shooting algorithm** will\n",
    "work.\n",
    "\n",
    "The **shooting algorithm** is an instance of a **guess and verify**\n",
    "algorithm.\n",
    "\n",
    "It proceeds as follows:\n",
    "\n",
    "- Guess a value for the initial Lagrange multiplier $ \\mu_0 $.  \n",
    "- Apply the **simple algorithm** described above.  \n",
    "- Compute the implied value of $ k_{T+1} $ and check whether it\n",
    "  equals zero.  \n",
    "- If the implied $ K_{T+1} =0 $, we have solved the problem.  \n",
    "- If $ K_{T+1} > 0 $, lower $ \\mu_0 $ and try again.  \n",
    "- If $ K_{T+1} < 0 $, raise $ \\mu_0 $ and try again.  \n",
    "\n",
    "\n",
    "The following Python code implements the shooting algorithm for the\n",
    "planning problem.\n",
    "\n",
    "We make a slight modification starting with a guess of\n",
    "$ c_0 $ but since $ c_0 $ is a function of $ \\mu_0 $\n",
    "there is no difference to the procedure above.\n",
    "\n",
    "We’ll apply it with an initial guess that will turn out not to be\n",
    "perfect, as we’ll soon see."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Parameters\n",
    "γ = 2\n",
    "δ = 0.02\n",
    "β = 0.95\n",
    "α = 0.33\n",
    "A = 1\n",
    "\n",
    "# Initial guesses\n",
    "T = 10\n",
    "c = np.zeros(T+1)  # T periods of consumption initialized to 0\n",
    "# T periods of capital initialized to 0 (T+2 to include t+1 variable as well)\n",
    "k = np.zeros(T+2)\n",
    "k[0] = 0.3  # Initial k\n",
    "c[0] = 0.2  # Guess of c_0\n",
    "\n",
    "@njit\n",
    "def shooting_method(c, # Initial consumption\n",
    "                    k,   # Initial capital\n",
    "                    γ,   # Coefficient of relative risk aversion\n",
    "                    δ,   # Depreciation rate on capital# Depreciation rate\n",
    "                    β,   # Discount factor\n",
    "                    α,   # Return to capital per capita\n",
    "                    A):  # Technology\n",
    "\n",
    "    T = len(c) - 1\n",
    "\n",
    "    for t in range(T):\n",
    "        # Equation 1 with inequality\n",
    "        k[t+1] = f(A=A, k=k[t], α=α) + (1 - δ) * k[t] - c[t]\n",
    "        if k[t+1] < 0:   # Ensure nonnegativity\n",
    "            k[t+1] = 0\n",
    "\n",
    "      # Equation 2: We keep in the general form to show how we would\n",
    "      # solve if we didn't want to do any simplification\n",
    "\n",
    "        if β * (f_prime(A=A, k=k[t+1], α=α) + (1 - δ)) == np.inf:\n",
    "            # This only occurs if k[t+1] is 0, in which case, we won't\n",
    "            # produce anything next period, so consumption will have to be 0\n",
    "            c[t+1] = 0\n",
    "        else:\n",
    "            c[t+1] = u_prime_inv(u_prime(c=c[t], γ=γ) \\\n",
    "            / (β * (f_prime(A=A, k=k[t+1], α=α) + (1 - δ))), γ=γ)\n",
    "\n",
    "    # Terminal condition calculation\n",
    "    k[T+1] = f(A=A, k=k[T], α=α) + (1 - δ) * k[T] - c[T]\n",
    "\n",
    "    return c, k\n",
    "\n",
    "paths = shooting_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
    "colors = ['blue', 'red']\n",
    "titles = ['Consumption', 'Capital']\n",
    "ylabels = ['$c_t$', '$k_t$']\n",
    "\n",
    "for path, color, title, y, ax in zip(paths, colors, titles, ylabels, axes):\n",
    "    ax.plot(path, c=color, alpha=0.7)\n",
    "    ax.set(title=title, ylabel=y, xlabel='t')\n",
    "\n",
    "ax.scatter(T+1, 0, s=80)\n",
    "ax.axvline(T+1, color='k', ls='--', lw=1)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evidently, our initial guess for $ \\mu_0 $ is too high and makes\n",
    "initial consumption is too low.\n",
    "\n",
    "We know this because we miss our $ K_{T+1}=0 $ target on the high\n",
    "side.\n",
    "\n",
    "Now we automate things with a search-for-a-good $ \\mu_0 $\n",
    "algorithm that stops when we hit the target $ K_{t+1} = 0 $.\n",
    "\n",
    "The search procedure is to use a **bisection method**.\n",
    "\n",
    "Here is how we apply the bisection method.\n",
    "\n",
    "We take an initial guess for $ C_0 $ (we can eliminate\n",
    "$ \\mu_0 $ because $ C_0 $ is an exact function of\n",
    "$ \\mu_0 $).\n",
    "\n",
    "We know that the lowest $ C_0 $ can ever be is $ 0 $ and the\n",
    "largest it can be is initial output $ f(K_0) $.\n",
    "\n",
    "We take a $ C_0 $ guess and shoot forward to $ T+1 $.\n",
    "\n",
    "If the $ K_{T+1}>0 $, let it be our new **lower** bound\n",
    "on $ C_0 $.\n",
    "\n",
    "If $ K_{T+1}<0 $, let it be our new **upper** bound.\n",
    "\n",
    "Make a new guess for $ C_0 $ exactly halfway between our new\n",
    "upper and lower bounds.\n",
    "\n",
    "Shoot forward again and iterate the procedure.\n",
    "\n",
    "When $ K_{T+1} $ gets close enough to 0 (within some error\n",
    "tolerance bounds), stop and declare victory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def bisection_method(c,\n",
    "                     k,\n",
    "                     γ,              # Coefficient of relative risk aversion\n",
    "                     δ,              # Depreciation rate\n",
    "                     β,              # Discount factor\n",
    "                     α,              # Return to capital per capita\n",
    "                     A,              # Technology\n",
    "                     tol=1e-4,\n",
    "                     max_iter=1e4,\n",
    "                     terminal=0):    # Value we are shooting towards\n",
    "\n",
    "    T = len(c) - 1\n",
    "    i = 1                            # Initial iteration\n",
    "    c_high = f(k=k[0], α=α, A=A)     # Initial high value of c\n",
    "    c_low = 0                        # Initial low value of c\n",
    "\n",
    "    path_c, path_k = shooting_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "    while (np.abs((path_k[T+1] - terminal)) > tol or path_k[T] == terminal) \\\n",
    "        and i < max_iter:\n",
    "\n",
    "        if path_k[T+1] - terminal > tol:\n",
    "            # If assets are too high the c[0] we chose is now a lower bound\n",
    "            # on possible values of c[0]\n",
    "            c_low = c[0]\n",
    "        elif path_k[T+1] - terminal < -tol:\n",
    "            # If assets fell too quickly, the c[0] we chose is now an upper\n",
    "            # bound on possible values of c[0]\n",
    "            c_high=c[0]\n",
    "        elif path_k[T] == terminal:\n",
    "            # If assets fell  too quickly, the c[0] we chose is now an upper\n",
    "            # bound on possible values of c[0]\n",
    "            c_high=c[0]\n",
    "\n",
    "        c[0] = (c_high + c_low) / 2  # This is the bisection part\n",
    "        path_c, path_k = shooting_method(c, k, γ, δ, β, α, A)\n",
    "        i += 1\n",
    "\n",
    "    if np.abs(path_k[T+1] - terminal) < tol and path_k[T] != terminal:\n",
    "        print('Converged successfully on iteration', i-1)\n",
    "    else:\n",
    "        print('Failed to converge and hit maximum iteration')\n",
    "\n",
    "    μ = u_prime(c=path_c, γ=γ)\n",
    "    return path_c, path_k, μ"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 10\n",
    "c = np.zeros(T+1) # T periods of consumption initialized to 0\n",
    "# T periods of capital initialized to 0. T+2 to include t+1 variable as well\n",
    "k = np.zeros(T+2)\n",
    "\n",
    "k[0] = 0.3 # initial k\n",
    "c[0] = 0.3 # our guess of c_0\n",
    "\n",
    "paths = bisection_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "def plot_paths(paths, axes=None, ss=None):\n",
    "\n",
    "    T = len(paths[0])\n",
    "\n",
    "    if axes is None:\n",
    "        fix, axes = plt.subplots(1, 3, figsize=(13, 3))\n",
    "\n",
    "    ylabels = ['$c_t$', '$k_t$', '$\\mu_t$']\n",
    "    titles = ['Consumption', 'Capital', 'Lagrange Multiplier']\n",
    "\n",
    "    for path, y, title, ax in zip(paths, ylabels, titles, axes):\n",
    "        ax.plot(path)\n",
    "        ax.set(ylabel=y, title=title, xlabel='t')\n",
    "\n",
    "    # Plot steady state value of capital\n",
    "    if ss is not None:\n",
    "        axes[1].axhline(ss, c='k', ls='--', lw=1)\n",
    "\n",
    "    axes[1].axvline(T, c='k', ls='--', lw=1)\n",
    "    axes[1].scatter(T, paths[1][-1], s=80)\n",
    "    plt.tight_layout()\n",
    "\n",
    "plot_paths(paths)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting Intial Capital to the Steady State\n",
    "\n",
    "If $ T \\rightarrow +\\infty $, the optimal allocation converges to\n",
    "steady state values of $ C_t $ and $ K_t $.\n",
    "\n",
    "It is instructive to compute these and then to set $ K_0 $ equal\n",
    "to its steady state value.\n",
    "\n",
    "In a steady state $ K_{t+1} = K_t=\\bar{K} $ for all very\n",
    "large $ t $ the feasibility constraint [(4)](#equation-allocation) is\n",
    "\n",
    "\n",
    "<a id='equation-feasibility-constraint'></a>\n",
    "$$\n",
    "f(\\bar{K})-\\delta \\bar{K} = \\bar{C} \\tag{13}\n",
    "$$\n",
    "\n",
    "Substituting $ K_t = \\bar K $ and $ C_t=\\bar C $ for\n",
    "all $ t $ into [(12)](#equation-l12) gives\n",
    "\n",
    "$$\n",
    "1=\\beta \\frac{u'(\\bar{C})}{u'(\\bar{C})}[f'(\\bar{K})+(1-\\delta)]\n",
    "$$\n",
    "\n",
    "Defining $ \\beta = \\frac{1}{1+\\rho} $, and cancelling gives\n",
    "\n",
    "$$\n",
    "1+\\rho = 1[f'(\\bar{K}) + (1-\\delta)]\n",
    "$$\n",
    "\n",
    "Simplifying gives\n",
    "\n",
    "$$\n",
    "f'(\\bar{K}) = \\rho +\\delta\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "\\bar{K} = f'^{-1}(\\rho+\\delta)\n",
    "$$\n",
    "\n",
    "Using our production function [(3)](#equation-production-function) gives\n",
    "\n",
    "$$\n",
    "\\alpha \\bar{K}^{\\alpha-1} = \\rho + \\delta\n",
    "$$\n",
    "\n",
    "Finally, using $ \\alpha= .33 $,\n",
    "$ \\rho = 1/\\beta-1 =1/(19/20)-1 = 20/19-19/19 = 1/19 $, $ \\delta = 1/50 $,\n",
    "we get\n",
    "\n",
    "$$\n",
    "\\bar{K} = \\left(\\frac{\\frac{33}{100}}{\\frac{1}{50}+\\frac{1}{19}}\\right)^{\\frac{67}{100}} \\approx 9.57583\n",
    "$$\n",
    "\n",
    "Let’s verify this with Python and then use this steady state\n",
    "$ \\bar K $ as our initial capital stock $ K_0 $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "ρ = 1 / β - 1\n",
    "k_ss = f_prime_inv(k=ρ+δ, A=A, α=α)\n",
    "\n",
    "print(f'steady state for capital is: {k_ss}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 150\n",
    "c = np.zeros(T+1)\n",
    "k = np.zeros(T+2)\n",
    "c[0] = 0.3\n",
    "k[0] = k_ss  # Start at steady state\n",
    "paths = bisection_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "plot_paths(paths, ss=k_ss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evidently, in this economy with a large value of\n",
    "$ T $, $ K_t $ stays near its initial value at the until the\n",
    "end of time approaches closely.\n",
    "\n",
    "Evidently, the planner likes the steady state capital stock and wants\n",
    "to stay near there for a long time.\n",
    "\n",
    "Let’s see what happens when we push the initial\n",
    "$ K_0 $ below $ \\bar K $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "k_init = k_ss / 3   # Below our steady state\n",
    "T = 150\n",
    "c = np.zeros(T+1)\n",
    "k = np.zeros(T+2)\n",
    "c[0] = 0.3\n",
    "k[0] = k_init\n",
    "paths = bisection_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "plot_paths(paths, ss=k_ss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how the planner pushes capital toward the steady state, stays\n",
    "near there for a while, then pushes $ K_t $ toward the terminal\n",
    "value $ K_{T+1} =0 $ as $ t $ gets close to $ T $.\n",
    "\n",
    "The following graphs compare outcomes as we vary $ T $."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T_list = (150, 75, 50, 25)\n",
    "\n",
    "fix, axes = plt.subplots(1, 3, figsize=(13, 3))\n",
    "\n",
    "for T in T_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_init\n",
    "    paths = bisection_method(c, k, γ, δ, β, α, A)\n",
    "    plot_paths(paths, ss=k_ss, axes=axes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following calculation shows that when we set $ T $ very large\n",
    "the planner makes the capital stock spend most of its time close to\n",
    "its steady state value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T_list = (250, 150, 50, 25)\n",
    "\n",
    "fix, axes = plt.subplots(1, 3, figsize=(13, 3))\n",
    "\n",
    "for T in T_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_init\n",
    "    paths = bisection_method(c, k, γ, δ, β, α, A)\n",
    "    plot_paths(paths, ss=k_ss, axes=axes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The different colors in the above graphs are tied to outcomes with\n",
    "different horizons $ T $.\n",
    "\n",
    "Notice that as the horizon increases, the planner puts $ K_t $\n",
    "closer to the steady state value $ \\bar K $ for longer.\n",
    "\n",
    "This pattern reflects a **turnpike** property of the steady state.\n",
    "\n",
    "A rule of thumb for the planner is\n",
    "\n",
    "- for whatever $ K_0 $ you start with, push $ K_t $ toward\n",
    "  the steady state and stay there for as long as you can  \n",
    "\n",
    "\n",
    "In loose language: head for the turnpike and stay near it for as long as you can.\n",
    "\n",
    "As we drive $ T $ toward $ +\\infty $, the planner\n",
    "keeps $ K_t $ very close to its steady state for all dates after\n",
    "some transition toward the steady state.\n",
    "\n",
    "The planner makes the saving rate $ \\frac{f(K_t) - C_t}{f(K_t)} $\n",
    "vary over time.\n",
    "\n",
    "Let’s calculate it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def S(K):\n",
    "    '''Aggregate savings'''\n",
    "    T = len(K) - 2\n",
    "    S = np.zeros(T+1)\n",
    "    for t in range(T+1):\n",
    "        S[t] = K[t+1] - (1 - δ) * K[t]\n",
    "    return S\n",
    "\n",
    "@njit\n",
    "def s(K):\n",
    "    '''Savings rate'''\n",
    "    T = len(K) - 2\n",
    "    Y = f(A, K, α)\n",
    "    Y = Y[0:T+1]\n",
    "    s = S(K) / Y\n",
    "    return s\n",
    "\n",
    "def plot_savings(paths, c_ss=None, k_ss=None, s_ss=None, axes=None):\n",
    "\n",
    "    T = len(paths[0])\n",
    "    k_star = paths[1]\n",
    "    savings_path = s(k_star)\n",
    "    new_paths = (paths[0], paths[1], savings_path)\n",
    "\n",
    "    if axes is None:\n",
    "        fix, axes = plt.subplots(1, 3, figsize=(13, 3))\n",
    "\n",
    "    ylabels = ['$c_t$', '$k_t$', '$s_t$']\n",
    "    titles = ['Consumption', 'Capital', 'Savings Rate']\n",
    "\n",
    "    for path, y, title, ax in zip(new_paths, ylabels, titles, axes):\n",
    "        ax.plot(path)\n",
    "        ax.set(ylabel=y, title=title, xlabel='t')\n",
    "\n",
    "    # Plot steady state value of consumption\n",
    "    if c_ss is not None:\n",
    "        axes[0].axhline(c_ss, c='k', ls='--', lw=1)\n",
    "\n",
    "    # Plot steady state value of capital\n",
    "    if k_ss is not None:\n",
    "        axes[1].axhline(k_ss, c='k', ls='--', lw=1)\n",
    "\n",
    "    # Plot steady state value of savings\n",
    "    if s_ss is not None:\n",
    "        axes[2].axhline(s_ss, c='k', ls='--', lw=1)\n",
    "\n",
    "    axes[1].axvline(T, c='k', ls='--', lw=1)\n",
    "    axes[1].scatter(T, k_star[-1], s=80)\n",
    "    plt.tight_layout()\n",
    "\n",
    "T_list = (250, 150, 75, 50)\n",
    "\n",
    "fix, axes = plt.subplots(1, 3, figsize=(13, 3))\n",
    "\n",
    "for T in T_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_init\n",
    "    paths = bisection_method(c, k, γ, δ, β, α, A)\n",
    "    plot_savings(paths, k_ss=k_ss, axes=axes)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The Limiting Economy\n",
    "\n",
    "We now consider an economy in which $ T = +\\infty $.\n",
    "\n",
    "The appropriate thing to do is to replace terminal condition\n",
    "[(10)](#equation-constraint4) with\n",
    "\n",
    "$$\n",
    "\\lim_{T \\rightarrow +\\infty} \\beta^T u'(C_T) K_{T+1} = 0\n",
    "$$\n",
    "\n",
    "which is sometimes called a **transversality condition**.\n",
    "\n",
    "This condition will be satisfied by a path that converges to an\n",
    "optimal steady state.\n",
    "\n",
    "We can approximate the optimal path from an arbitrary initial\n",
    "$ K_0 $ and shooting towards the optimal steady state\n",
    "$ K $ at a large but finite $ T+1 $.\n",
    "\n",
    "In the following code, we do this for a large $ T $; we shoot\n",
    "towards the **steady state** and plot consumption, capital and the\n",
    "savings rate.\n",
    "\n",
    "We know that in the steady state that the saving rate must be fixed\n",
    "and that $ \\bar s= \\frac{f(\\bar K)-\\bar C}{f(\\bar K)} $.\n",
    "\n",
    "From [(13)](#equation-feasibility-constraint) the steady state saving rate equals\n",
    "\n",
    "$$\n",
    "\\bar s =\\frac{ \\delta \\bar{K}}{f(\\bar K)}\n",
    "$$\n",
    "\n",
    "The steady state savings level $ \\bar S = \\bar s f(\\bar K) $ is\n",
    "the amount required to offset capital depreciation each period.\n",
    "\n",
    "We first study optimal capital paths that start below the steady\n",
    "state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 130\n",
    "\n",
    "# Steady states\n",
    "S_ss = δ * k_ss\n",
    "c_ss = f(A, k_ss, α) - S_ss\n",
    "s_ss = S_ss / f(A, k_ss, α)\n",
    "\n",
    "c = np.zeros(T+1)\n",
    "k = np.zeros(T+2)\n",
    "c[0] = 0.3\n",
    "k[0] = k_ss / 3         # Start below steady state\n",
    "paths = bisection_method(c, k, γ, δ, β, α, A, terminal=k_ss)\n",
    "plot_savings(paths, k_ss=k_ss, s_ss=s_ss, c_ss=c_ss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since $ K_0<\\bar K $, $ f'(K_0)>\\rho +\\delta $.\n",
    "\n",
    "The planner chooses a positive saving rate above the steady state\n",
    "level offsetting depreciation that enables us to increase our capital\n",
    "stock.\n",
    "\n",
    "Note, $ f''(K)<0 $, so as $ K $ rises, $ f'(K) $ declines.\n",
    "\n",
    "The planner slowly lowers the savings rate until reaching a steady\n",
    "state where $ f'(K)=\\rho +\\delta $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "\n",
    "- Plot the optimal consumption, capital, and savings paths when the\n",
    "  initial capital level begins at 1.5 times the steady state level\n",
    "  as we shoot towards the steady state at $ T=130 $.  \n",
    "- Why does the savings rate respond like it does?  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T = 130\n",
    "\n",
    "c = np.zeros(T+1)\n",
    "k = np.zeros(T+2)\n",
    "c[0] = 0.3\n",
    "k[0] = k_ss * 1.5   # Start above steady state\n",
    "paths = bisection_method(c, k, γ, δ, β, α, A, terminal=k_ss)\n",
    "plot_savings(paths, k_ss=k_ss, s_ss=s_ss, c_ss=c_ss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Competitive Equilibrium\n",
    "\n",
    "Next, we study a decentralized version of an economy with the same\n",
    "technology and preference structure as our planned economy.\n",
    "\n",
    "But now there is no planner.\n",
    "\n",
    "Market prices adjust to reconcile distinct decisions that are made\n",
    "separately by a representative household and a representative firm.\n",
    "\n",
    "The technology for producing goods and accumulating capital via\n",
    "physical investment remains as in our planned economy.\n",
    "\n",
    "There is a representative consumer who has the same preferences over\n",
    "consumption plans as did the consumer in the planned economy.\n",
    "\n",
    "Instead of being told what to consume and save by a planner, the\n",
    "household chooses for itself subject to a budget constraint\n",
    "\n",
    "- At each time $ t $, the household receives wages and rentals\n",
    "  of capital from a firm – these comprise its **income** at\n",
    "  time $ t $.  \n",
    "- The consumer decides how much income to allocate to consumption or\n",
    "  to savings.  \n",
    "- The household can save either by acquiring additional physical\n",
    "  capital (it trades one for one with time $ t $ consumption)\n",
    "  or by acquiring claims on consumption at dates other\n",
    "  than $ t $.  \n",
    "- A utility-maximizing household owns all physical capital and labor\n",
    "  and rents them to the firm.  \n",
    "- The household consumes, supplies labor, and invests in physical\n",
    "  capital.  \n",
    "- A profit-maximizing representative firm operates the production\n",
    "  technology.  \n",
    "- The firm rents labor and capital each period from the\n",
    "  representative household and sells its output each period to the\n",
    "  household.  \n",
    "- The representative household and the representative firm are both\n",
    "  **price takers:**  \n",
    "  - they (correctly) believe that prices are not affected by their\n",
    "    choices  \n",
    "\n",
    "\n",
    "**Note:** We are free to think of there being a large number\n",
    "$ M $ of identical representative consumers and $ M $\n",
    "identical representative firms."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Firm Problem\n",
    "\n",
    "At time $ t $ the representative firm hires labor\n",
    "$ \\tilde n_t $ and capital $ \\tilde k_t $.\n",
    "\n",
    "The firm’s profits at time $ t $ are\n",
    "\n",
    "$$\n",
    "F(\\tilde k_t, \\tilde n_t)-w_t \\tilde n_t -\\eta_t \\tilde k_t\n",
    "$$\n",
    "\n",
    "where $ w_t $ is a wage rate at $ t $\n",
    "and $ \\eta_t $ is the rental rate on capital at $ t $.\n",
    "\n",
    "As in the planned economy model\n",
    "\n",
    "$$\n",
    "F(\\tilde k_t, \\tilde n_t) = A \\tilde k_t^\\alpha \\tilde n_t^{1-\\alpha}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Zero Profit Conditions\n",
    "\n",
    "Zero-profits condition for capital and labor are\n",
    "\n",
    "$$\n",
    "F_k(\\tilde k_t, \\tilde n_t) =\\eta_t\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "\n",
    "<a id='equation-zero-profits'></a>\n",
    "$$\n",
    "F_n(\\tilde k_t, \\tilde n_t) =w_t \\tag{14}\n",
    "$$\n",
    "\n",
    "These conditions emerge from a no-arbitrage requirement.\n",
    "\n",
    "To describe this line of reasoning, we begin by applying a theorem of\n",
    "Euler about linearly homogenous functions.\n",
    "\n",
    "The theorem applies to the Cobb-Douglas production function because\n",
    "it assumed displays constant returns to scale:\n",
    "\n",
    "$$\n",
    "\\alpha F(\\tilde k_t, \\tilde n_t) =  F(\\alpha  \\tilde k_t, \\alpha \\tilde n_t)\n",
    "$$\n",
    "\n",
    "for $ \\alpha \\in (0,1) $.\n",
    "\n",
    "Taking the partial derivative\n",
    "$ \\frac{\\partial F }{\\partial \\alpha} $ on both sides of the\n",
    "above equation gives\n",
    "\n",
    "$$\n",
    "F(\\tilde k_t,\\tilde n_t) =_\\text{chain rule} \\frac{\\partial F}{\\partial \\tilde k_t}\n",
    "\\tilde k_t + \\frac{\\partial F}{\\partial \\tilde  n_t} \\tilde n_t\n",
    "$$\n",
    "\n",
    "Rewrite the firm’s profits as\n",
    "\n",
    "$$\n",
    "\\frac{\\partial F}{\\partial \\tilde k_t} \\tilde k_t +\n",
    "\\frac{\\partial F}{\\partial \\tilde  n_t} \\tilde n_t-w_t \\tilde n_t -\\eta_t k_t\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "$$\n",
    "\\left(\\frac{\\partial F}{\\partial \\tilde k_t}-\\eta_t\\right) \\tilde k_t +\n",
    "\\left(\\frac{\\partial F}{\\partial \\tilde  n_t}-w_t\\right) \\tilde n_t\n",
    "$$\n",
    "\n",
    "Because $ F $ is homogeneous of degree $ 1 $, it follows\n",
    "that $ \\frac{\\partial F}{\\partial \\tilde k_t} $ and\n",
    "$ \\frac{\\partial F}{\\partial \\tilde n_t} $ are homogeneous of\n",
    "degree $ 0 $ and therefore fixed with respect to\n",
    "$ \\tilde k_t $ and $ \\tilde n_t $.\n",
    "\n",
    "If $ \\frac{\\partial F}{\\partial \\tilde k_t}> \\eta_t $, then the\n",
    "firm makes positive profits on each additional unit of\n",
    "$ \\tilde k_t $, so it will want to make $ \\tilde k_t $\n",
    "arbitrarily large.\n",
    "\n",
    "But setting $ \\tilde k_t = + \\infty $ is not physically feasible,\n",
    "so presumably **equilibrium** prices will assume values that present\n",
    "the firm with no such arbitrage opportunity.\n",
    "\n",
    "A related argument applies if\n",
    "$ \\frac{\\partial F}{\\partial \\tilde n_t}> w_t $.\n",
    "\n",
    "If $ \\frac{\\partial \\tilde k_t}{\\partial \\tilde k_t}< \\eta_t $,\n",
    "the firm will set $ \\tilde k_t $ to zero.\n",
    "\n",
    "Again, **equilibrium** prices won’t incentive the firm to do that.\n",
    "\n",
    "And so on…\n",
    "\n",
    "It is convenient to define\n",
    "$ \\vec w_t =\\{w_0, \\dots,w_T\\} $and $ \\vec \\eta_t = \\{\\eta_0, \\dots, \\eta_T\\} $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Household Problem\n",
    "\n",
    "A representative household lives at $ t=0,1,\\dots, T $.\n",
    "\n",
    "At $ t $, the household rents $ 1 $ unit of labor\n",
    "and $ k_t $ units of capital to a firm and receives income\n",
    "\n",
    "$$\n",
    "w_t 1+ \\eta_t k_t\n",
    "$$\n",
    "\n",
    "At $ t $ the household allocates its income to the following\n",
    "purchases\n",
    "\n",
    "$$\n",
    "\\left(c_t + (k_{t+1} -(1-\\delta)k_t)\\right)\n",
    "$$\n",
    "\n",
    "Here $ \\left(k_{t+1} -(1-\\delta)k_t\\right) $ is the household’s\n",
    "net investment in physical capital and $ \\delta \\in (0,1) $ is\n",
    "again a depreciation rate of capital.\n",
    "\n",
    "In period $ t $ is free to purchase more goods to be consumed and\n",
    "invested in physical capital than its income from supplying capital\n",
    "and labor to the firm, provided that in some other periods its income\n",
    "exceeds its purchases.\n",
    "\n",
    "A household’s net excess demand for time $ t $ consumption goods\n",
    "is the gap\n",
    "\n",
    "$$\n",
    "e_t \\equiv \\left(c_t + (k_{t+1} -(1-\\delta)k_t)\\right)-(w_t 1 + \\eta_t k_t)\n",
    "$$\n",
    "\n",
    "Let $ \\vec c = \\{c_0,\\dots,c_T\\} $ and let $ \\vec k = \\{k_1,\\dots,k_{T+1}\\} $.\n",
    "\n",
    "$ k_0 $ is given to the household."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Market Structure for Intertemporal Trades\n",
    "\n",
    "There is a **single** grand competitive market in which a\n",
    "representative household can trade date $ 0 $ goods for goods at\n",
    "all other dates $ t=1, 2, \\ldots, T $.\n",
    "\n",
    "What matters are not **bilateral** trades of the good at one date\n",
    "$ t $ for the good at another date $ \\tilde t \\neq t $.\n",
    "\n",
    "Instead, think of there being **multilateral** and **multitemporal**\n",
    "trades in which bundles of goods at some dates can be traded for\n",
    "bundles of goods at some other dates.\n",
    "\n",
    "There exist **complete markets** in such bundles with associated\n",
    "market prices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Market Prices\n",
    "\n",
    "Let $ q^0_t $ be the price of a good at date $ t $ relative\n",
    "to a good at date $ 0 $.\n",
    "\n",
    "$ \\{q^0_t\\}_{t=0}^T $ is a vector of **Hicks-Arrow prices**,\n",
    "named after the 1972 joint economics Nobel prize winners who used\n",
    "such prices in some of their important work.\n",
    "\n",
    "Evidently,\n",
    "\n",
    "$$\n",
    "q^0_t=\\frac{\\text{number of time 0 goods}}{\\text{number of time t goods}}\n",
    "$$\n",
    "\n",
    "Because $ q^0_t $ is a **relative price**, the units in terms of\n",
    "which prices are quoted are arbitrary – we can normalize them\n",
    "without substantial consequence.\n",
    "\n",
    "If we use the price vector $ \\{q^0_t\\}_{t=0}^T $ to evaluate a\n",
    "stream of excess demands $ \\{e_t\\}_{t=0}^T $ we compute the\n",
    "**present value** of $ \\{e_t\\}_{t=0}^T $ to be\n",
    "$ \\sum_{t=0}^T q^0_t e_t $.\n",
    "\n",
    "That the market is **multitemporal** is reflected in the situation\n",
    "that the household faces a **single** budget constraint.\n",
    "\n",
    "It states that the present value of the household’s net excess\n",
    "demands must be zero:\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^T q^0_t e_t  \\leq 0\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^T q^0_t  \\left(c_t + (k_{t+1} -(1-\\delta)k_t)-(w_t 1 + \\eta_t k_t) \\right) \\leq 0\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Household Problem\n",
    "\n",
    "The household faces the constrained optimization problem:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}& \\max_{\\vec c, \\vec k}  \\sum_{t=0}^T \\beta^t u(c_t) \\\\ \\text{subject to} \\ \\   & \\sum_{t=0}^T q_t^0\\left(c_t +\\left(k_{t+1}-(1-\\delta) k_t\\right) -w_t -\\eta_t k_t\\right) \\leq 0  \\notag \\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Definitions\n",
    "\n",
    "- A **price system** is a sequence\n",
    "  $ \\{q_t^0,\\eta_t,w_t\\}_{t=0}^T= \\{\\vec q, \\vec \\eta, \\vec w\\} $.  \n",
    "- An **allocation** is a sequence\n",
    "  $ \\{c_t,k_{t+1},n_t=1\\}_{t=0}^T = \\{\\vec c, \\vec k, \\vec n =1\\} $.  \n",
    "- A **competitive equilibrium** is a price system and an allocation\n",
    "  for which  \n",
    "  - Given the price system, the allocation solves the household’s\n",
    "    problem.  \n",
    "  - Given the price system, the allocation solves the firm’s\n",
    "    problem.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing a Competitive Equilibrium\n",
    "\n",
    "We shall compute a competitive equilibrium using a **guess and\n",
    "verify** approach.\n",
    "\n",
    "- We shall **guess** equilibrium price sequences\n",
    "  $ \\{\\vec q, \\vec \\eta, \\vec w\\} $.  \n",
    "- We shall then **verify** that at those prices, the household and\n",
    "  the firm choose the same allocation.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Guess for Price System\n",
    "\n",
    "We have computed an allocation $ \\{\\vec C, \\vec K, \\vec 1\\} $\n",
    "that solves the planning problem.\n",
    "\n",
    "We use that allocation to construct our guess for the equilibrium\n",
    "price system.\n",
    "\n",
    "In particular, we guess that for $ t=0,\\dots,T $:\n",
    "\n",
    "\n",
    "<a id='equation-eq-price'></a>\n",
    "$$\n",
    "\\lambda q_t^0 = \\beta^t u'(K_t) =\\beta^t \\mu_t \\tag{15}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-eq-price2'></a>\n",
    "$$\n",
    "w_t = f(K_t) -K_t f'(K_t) \\tag{16}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-eq-price3'></a>\n",
    "$$\n",
    "\\eta_t = f'(K_t) \\tag{17}\n",
    "$$\n",
    "\n",
    "At these prices, let the capital chosen by the household be\n",
    "\n",
    "\n",
    "<a id='equation-eq-pr4'></a>\n",
    "$$\n",
    "k^*_t(\\vec q, \\vec w, \\vec \\eta) , \\quad t \\geq 0 \\tag{18}\n",
    "$$\n",
    "\n",
    "and let the allocation chosen by the firm be\n",
    "\n",
    "$$\n",
    "\\tilde k^*_t(\\vec q, \\vec  w, \\vec \\eta), \\quad t \\geq 0\n",
    "$$\n",
    "\n",
    "and so on.\n",
    "\n",
    "If our guess for the equilibrium price system is correct, then it\n",
    "must occur that\n",
    "\n",
    "\n",
    "<a id='equation-ge1'></a>\n",
    "$$\n",
    "k_t^*  = \\tilde k_t^* \\tag{19}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-ge2'></a>\n",
    "$$\n",
    "1   = \\tilde n_t^* \\tag{20}\n",
    "$$\n",
    "\n",
    "$$\n",
    "c_t^* + k_{t+1}^* - (1-\\delta) k_t^*  = F(\\tilde k_t^*, \\tilde n_t^*)\n",
    "$$\n",
    "\n",
    "We shall verify that for $ t=0,\\dots,T $ the allocations chosen\n",
    "by the household and the firm both equal the allocation that solves\n",
    "the planning problem:\n",
    "\n",
    "\n",
    "<a id='equation-eq-pl'></a>\n",
    "$$\n",
    "k^*_t = \\tilde k^*_t=K_t, \\tilde n_t=1, c^*_t=C_t \\tag{21}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Verification Procedure\n",
    "\n",
    "Our approach is to stare at first-order necessary conditions for the\n",
    "optimization problems of the household and the firm.\n",
    "\n",
    "At the price system we have guessed, both sets of first-order\n",
    "conditions are satisfied at the allocation that solves the planning\n",
    "problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Household’s Lagrangian\n",
    "\n",
    "To solve the household’s problem, we formulate the appropriate\n",
    "Lagrangian and pose the min-max problem:\n",
    "\n",
    "$$\n",
    "\\min_{\\lambda} \\max_{\\vec{c},\\vec{k}}\\mathcal{L}(\\vec{c},\\vec{k},\\lambda)=\n",
    "\\sum_{t=0}^T \\beta^t u(c_t)+ \\lambda \\left(\\sum_{t=0}^T q_t^0\\left(\\left((1-\\delta)\n",
    "k_t -w_t\\right) +\\eta_t k_t -c_t  - k_{t+1}\\right)\\right)\n",
    "$$\n",
    "\n",
    "First-order conditions are\n",
    "\n",
    "\n",
    "<a id='equation-cond1'></a>\n",
    "$$\n",
    "c_t: \\quad \\beta^t u'(c_t)-\\lambda q_t^0=0 \\quad  t=0,1,\\dots,T \\tag{22}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-cond2'></a>\n",
    "$$\n",
    "k_t: \\quad -\\lambda q_t^0 \\left[(1-\\delta)+\\eta_t \\right]+\\lambda q^0_{t-1}=0 \\quad  t=1,2,\\dots,T+1 \\tag{23}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-cond3'></a>\n",
    "$$\n",
    "\\lambda:  \\quad \\left(\\sum_{t=0}^T q_t^0\\left(c_t + \\left(k_{t+1}-(1-\\delta) k_t\\right) -w_t -\\eta_t k_t\\right)\\right) \\leq 0 \\tag{24}\n",
    "$$\n",
    "\n",
    "\n",
    "<a id='equation-cond4'></a>\n",
    "$$\n",
    "k_{T+1}: \\quad -\\lambda q_0^{T+1} \\leq 0, \\ \\leq 0 \\text{ if } k_{T+1}=0; \\ =0 \\text{ if } k_{T+1}>0 \\tag{25}\n",
    "$$\n",
    "\n",
    "Now we plug in for our guesses of prices and derive all the FONC of\n",
    "the planner problem [(7)](#equation-constraint1)-[(10)](#equation-constraint4):\n",
    "\n",
    "Combining [(22)](#equation-cond1) and [(15)](#equation-eq-price), we get:\n",
    "\n",
    "$$\n",
    "u'(C_t) = \\mu_t\n",
    "$$\n",
    "\n",
    "which is [(7)](#equation-constraint1).\n",
    "\n",
    "Combining [(23)](#equation-cond2), [(15)](#equation-eq-price), and\n",
    "[(17)](#equation-eq-price3) we get:\n",
    "\n",
    "\n",
    "<a id='equation-co-re'></a>\n",
    "$$\n",
    "-\\lambda \\beta^t \\mu_t\\left[(1-\\delta) +f'(K_t)\\right] +\\lambda \\beta^{t-1}\\mu_{t-1}=0 \\tag{26}\n",
    "$$\n",
    "\n",
    "Rewriting [(26)](#equation-co-re) by dividing by $ \\lambda $ on\n",
    "both sides (which is nonzero due to u’>0) we get:\n",
    "\n",
    "$$\n",
    "\\beta^t \\mu_t [(1-\\delta+f'(K_t)] = \\beta^{t-1} \\mu_{t-1}\n",
    "$$\n",
    "\n",
    "or\n",
    "\n",
    "$$\n",
    "\\beta \\mu_t [(1-\\delta+f'(K_t)] = \\mu_{t-1}\n",
    "$$\n",
    "\n",
    "which is [(8)](#equation-constraint2).\n",
    "\n",
    "Combining [(24)](#equation-cond3), [(15)](#equation-eq-price), [(16)](#equation-eq-price2)\n",
    "and [(17)](#equation-eq-price3) after multiplying both sides of\n",
    "[(24)](#equation-cond3) by $ \\lambda $, we get:\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^T \\beta^t \\mu_{t} \\left(C_t+ (K_{t+1} -(1-\\delta)K_t)-f(K_t)+K_t f'(K_t)-f'(K_t)K_t\\right) \\leq 0\n",
    "$$\n",
    "\n",
    "Cancelling,\n",
    "\n",
    "$$\n",
    "\\sum_{t=0}^T  \\beta^t \\mu_{t} \\left(C_t +K_{t+1} -(1-\\delta)K_t - F(K_t,1)\\right) \\leq 0\n",
    "$$\n",
    "\n",
    "Since $ \\beta^t $ and $ \\mu_t $ are always positive here,\n",
    "(excepting perhaps the T+1 period) we get:\n",
    "\n",
    "$$\n",
    "C_t+K_{t+1}-(1-\\delta)K_t -F(K_t,1)=0 \\quad  \\text{ for all }t \\text{ in } 0,\\dots,T\n",
    "$$\n",
    "\n",
    "which is [(9)](#equation-constraint3).\n",
    "\n",
    "Combining [(25)](#equation-cond4) and [(15)](#equation-eq-price), we get:\n",
    "\n",
    "$$\n",
    "- \\beta^{T+1} \\mu_{T+1} \\leq 0\n",
    "$$\n",
    "\n",
    "Dividing both sides by $ \\beta^{T+1} $ which will be strictly\n",
    "positive here, we get:\n",
    "\n",
    "$$\n",
    "-\\mu_{T+1} \\leq 0\n",
    "$$\n",
    "\n",
    "which is the [(10)](#equation-constraint4) of our planning problem.\n",
    "\n",
    "**Thus, at our guess of the equilibrium price system, the allocation\n",
    "that solves the planning problem also solves the problem faced by a\n",
    "representative household living in a competitive equilibrium**.\n",
    "\n",
    "We now consider the problem faced by a firm in a competitive\n",
    "equilibrium:\n",
    "\n",
    "If we plug in [(21)](#equation-eq-pl) into [(14)](#equation-zero-profits) for all t, we\n",
    "get\n",
    "\n",
    "$$\n",
    "\\frac{\\partial F(K_t, 1)}{\\partial K_t} = f'(K_t) = \\eta_t\n",
    "$$\n",
    "\n",
    "which is [(17)](#equation-eq-price3).\n",
    "\n",
    "If we now plug [(21)](#equation-eq-pl) into [(14)](#equation-zero-profits) for all t, we\n",
    "get:\n",
    "\n",
    "$$\n",
    "\\frac{\\partial F(\\tilde K_t, 1)}{\\partial \\tilde L} = f(K_t)-f'(K_t)K_t=w_t\n",
    "$$\n",
    "\n",
    "which is exactly [(18)](#equation-eq-pr4).\n",
    "\n",
    "**Thus, at our guess of the equilibrium price system, the allocation\n",
    "that solves the planning problem also solves the problem faced by a firm\n",
    "within a competitive equilibrium**.\n",
    "\n",
    "By [(19)](#equation-ge1) and [(20)](#equation-ge2) this allocation is\n",
    "identical to the one that solves the consumer’s problem.\n",
    "\n",
    "**Note:** Because budget sets are affected only by relative prices,\n",
    "$ \\{q_0^t\\} $ is determined only up to multiplication by a\n",
    "positive constant.\n",
    "\n",
    "**Normalization:** We are free to choose a $ \\{q_0^t\\} $ that\n",
    "makes $ \\lambda=1 $, thereby making $ q_0^t $ be measured in\n",
    "units of the marginal utility of time $ 0 $ goods.\n",
    "\n",
    "We will also plot q, w and $ \\eta $ below to show the prices that\n",
    "induce the same aggregate movements we saw earlier in the planning\n",
    "problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def q_func(β, c, γ):\n",
    "    # Here we choose numeraire to be u'(c_0) -- this is q^(t_0)_t\n",
    "    T = len(c) - 2\n",
    "    q = np.zeros(T+1)\n",
    "    q[0] = 1\n",
    "    for t in range(1, T+2):\n",
    "        q[t] = β**t * u_prime(c[t], γ)\n",
    "    return q\n",
    "\n",
    "@njit\n",
    "def w_func(A, k, α):\n",
    "    w = f(A, k, α) - k * f_prime(A, k, α)\n",
    "    return w\n",
    "\n",
    "@njit\n",
    "def η_func(A, k, α):\n",
    "    η = f_prime(A, k, α)\n",
    "    return η"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we calculate and plot for each $ T $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "T_list = (250, 150, 75, 50)\n",
    "\n",
    "fix, axes = plt.subplots(2, 3, figsize=(13, 6))\n",
    "titles = ['Arrow-Hicks Prices', 'Labor Rental Rate', 'Capital Rental Rate',\n",
    "          'Consumption', 'Capital', 'Lagrange Multiplier']\n",
    "ylabels = ['$q_t^0$', '$w_t$', '$\\eta_t$', '$c_t$', '$k_t$', '$\\mu_t$']\n",
    "\n",
    "for T in T_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_ss / 3\n",
    "    c, k, μ = bisection_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "    q = q_func(β, c, γ)\n",
    "    w = w_func(β, k, α)[:-1]\n",
    "    η = η_func(A, k, α)[:-1]\n",
    "    plots = [q, w, η, c, k, μ]\n",
    "\n",
    "    for ax, plot, title, y in zip(axes.flatten(), plots, titles, ylabels):\n",
    "        ax.plot(plot)\n",
    "        ax.set(title=title, ylabel=y, xlabel='t')\n",
    "        if title is 'Capital':\n",
    "            ax.axhline(k_ss, lw=1, ls='--', c='k')\n",
    "        if title is 'Consumption':\n",
    "            ax.axhline(c_ss, lw=1, ls='--', c='k')\n",
    "\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Varying Curvature\n",
    "\n",
    "Now we see how our results change if we keep $ T $ constant, but allow\n",
    "the curvature parameter, $ \\gamma $ to vary, starting\n",
    "with $ K_0 $ below the steady state.\n",
    "\n",
    "We plot the results for $ T=150 $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "γ_list = (1.1, 4, 6, 8)\n",
    "T = 150\n",
    "\n",
    "fix, axes = plt.subplots(2, 3, figsize=(13, 6))\n",
    "\n",
    "for γ in γ_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_ss / 3\n",
    "    c, k, μ = bisection_method(c, k, γ, δ, β, α, A)\n",
    "\n",
    "    q = q_func(β, c, γ)\n",
    "    w = w_func(β, k, α)[:-1]\n",
    "    η = η_func(A, k, α)[:-1]\n",
    "    plots = [q, w, η, c, k, μ]\n",
    "\n",
    "    for ax, plot, title, y in zip(axes.flatten(), plots, titles, ylabels):\n",
    "        ax.plot(plot, label=f'$\\gamma = {γ}$')\n",
    "        ax.set(title=title, ylabel=y, xlabel='t')\n",
    "        if title is 'Capital':\n",
    "            ax.axhline(k_ss, lw=1, ls='--', c='k')\n",
    "        if title is 'Consumption':\n",
    "            ax.axhline(c_ss, lw=1, ls='--', c='k')\n",
    "\n",
    "axes[0, 0].legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adjusting $ \\gamma $ means adjusting how much individuals prefer\n",
    "to smooth consumption.\n",
    "\n",
    "Higher $ \\gamma $ means individuals prefer to smooth more\n",
    "resulting in slower adjustments to the steady state allocations.\n",
    "\n",
    "Vice-versa for lower $ \\gamma $."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Yield Curves and Hicks-Arrow Prices Again\n",
    "\n",
    "Now, we compute Hicks-Arrow prices again, but also calculate the\n",
    "implied yields to maturity.\n",
    "\n",
    "This will let us plot a **yield curve**.\n",
    "\n",
    "The key formulas are:\n",
    "\n",
    "The **yield to maturity**\n",
    "\n",
    "$$\n",
    "r_{t_0,t}= -\\frac{\\log q^{t_0}_t}{t - t_0}\n",
    "$$\n",
    "\n",
    "A generic Hicks-Arrow price for any base-year $ t_0\\leq t $\n",
    "\n",
    "$$\n",
    "q^{t_0}_t = \\beta^{t-t_0} \\frac{u'(c_t)}{u'(c_{t_0})}= \\beta^{t-t_0}\n",
    "\\frac{c_t^{-\\gamma}}{c_{t_0}^{-\\gamma}}\n",
    "$$\n",
    "\n",
    "We redefine our function for $ q $ to allow arbitrary base\n",
    "years, and define a new function for $ r $, then plot both.\n",
    "\n",
    "First, we plot when $ t_0=0 $ as before, for different values of\n",
    "$ T $, with $ K_0 $ below the steady state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "@njit\n",
    "def q_func(t_0, β, c, γ):\n",
    "    # Here we choose numeraire to be u'(c_0) -- this is q^(t_0)_t\n",
    "    T = len(c)\n",
    "    q = np.zeros(T+1-t_0)\n",
    "    q[0] = 1\n",
    "    for t in range(t_0+1, T):\n",
    "        q[t-t_0] = β**(t - t_0) * u_prime(c[t], γ) / u_prime(c[t_0], γ)\n",
    "    return q\n",
    "\n",
    "@njit\n",
    "def r_func(t_0, β, c, γ):\n",
    "    '''Yield to maturity'''\n",
    "    T = len(c) - 1\n",
    "    r = np.zeros(T+1-t_0)\n",
    "    for t in range(t_0+1, T+1):\n",
    "        r[t-t_0]= -np.log(q_func(t_0, β, c, γ)[t-t_0]) / (t - t_0)\n",
    "    return r\n",
    "\n",
    "t_0 = 0\n",
    "T_list = [150, 75, 50]\n",
    "γ = 2\n",
    "titles = ['Hicks-Arrow Prices', 'Yields']\n",
    "ylabels = ['$q_t^0$', '$r_t^0$']\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "\n",
    "for T in T_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_ss / 3\n",
    "    c, k, μ = bisection_method(c, k, γ, δ, β, α, A)\n",
    "    q = q_func(t_0, β, c, γ)\n",
    "    r = r_func(t_0, β, c, γ)\n",
    "\n",
    "    for ax, plot, title, y in zip(axes, (q, r), titles, ylabels):\n",
    "        ax.plot(plot)\n",
    "        ax.set(title=title, ylabel=y, xlabel='t')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we plot when $ t_0=20 $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "t_0 = 20\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(10, 5))\n",
    "\n",
    "for T in T_list:\n",
    "    c = np.zeros(T+1)\n",
    "    k = np.zeros(T+2)\n",
    "    c[0] = 0.3\n",
    "    k[0] = k_ss / 3\n",
    "    c, k, μ = bisection_method(c, k, γ, δ, β, α, A)\n",
    "    q = q_func(t_0, β, c, γ)\n",
    "    r = r_func(t_0, β, c, γ)\n",
    "\n",
    "    for ax, plot, title, y in zip(axes, (q, r), titles, ylabels):\n",
    "        ax.plot(plot)\n",
    "        ax.set(title=title, ylabel=y, xlabel='t')\n",
    "\n",
    "axes[1].set_title(f'Yields at $t_0 = {t_0}$')\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We shall have more to say about the term structure of interest rates\n",
    "in a later lecture on the topic."
   ]
  }
 ],
 "metadata": {
  "filename": "cass_koopmans.rst",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "Cass-Koopmans Optimal Growth Model"
 },
 "nbformat": 4,
 "nbformat_minor": 2
}