{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 7 Solution Methods to Solve the Growth Model with Julia\n",
    "\n",
    "This notebook is part of a computational appendix that accompanies the paper\n",
    "\n",
    "> MATLAB, Python, Julia: What to Choose in Economics?\n",
    "> > Coleman, Lyon, Maliar, and Maliar (2017)\n",
    "\n",
    "In order to run the codes in this notebook you will need to install and configure a few Julia packages. We recommend downloading [JuliaPro](https://juliacomputing.com/products/juliapro.html) and/or following the instructions on [quantecon.org](https://lectures.quantecon.org/jl/getting_started.html). Once your Julia installation is up and running, there are a few additional packages you will need in order to run the code here. To do this uncomment the lines in the cell below (by deleting the `#` and space at the beginning of each line) and run the cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Pkg.add(\"BasisMatrices\")\n",
    "# Pkg.add(\"Optim\")\n",
    "# Pkg.add(\"Parameters\")\n",
    "# Pkg.add(\"QuantEcon\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "using BasisMatrices, Optim, QuantEcon, Parameters\n",
    "using BasisMatrices: Degree, Derivative"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "This section gives a short description of the commonly used stochastic Neoclassical growth model.\n",
    "\n",
    "There is a single infinitely-lived representative agent who consumes and saves using capital. The consumer discounts the future with factor $\\beta$ and derives utility from only consumption. Additionally, saved capital will depreciate at $\\delta$.\n",
    "\n",
    "The consumer has access to a Cobb-Douglas technology which uses capital saved from the previous period to produce and is subject to stochastic productivity shocks.\n",
    "\n",
    "Productivity shocks follow an AR(1) in logs.\n",
    "\n",
    "The agent's problem can be written recursively using the following Bellman equation\n",
    "\n",
    "\\begin{align}\n",
    "  V(k_t, z_t) &= \\max_{k_{t+1}} u(c_t) + \\beta E \\left[ V(k_{t+1}, z_{t+1}) \\right] \\\\\n",
    "  &\\text{subject to } \\\\\n",
    "  c_t &= z_t f(k_t) + (1 - \\delta) k_t - k_{t+1} \\\\\n",
    "  \\log z_{t+1} &= \\rho \\log z_t + \\sigma \\varepsilon\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Julia Code\n",
    "\n",
    "We begin by defining a type that describes our model. It will hold the three things\n",
    "\n",
    "1. Parameters of the growth model\n",
    "2. Grids used for approximating the solution\n",
    "3. Nodes and weights used to approximate integration\n",
    "\n",
    "Note the `@with_kw` comes from the `Parameters` package -- It allows one to specify default arguments for the parameters when building a type (for more information refer to their [documentation](http://parametersjl.readthedocs.io/en/latest/)). One of the benefits of using the `Parameters` package is their code allows us to do things like, `@unpack a, b, c = Params` which takes elements from inside the type `Params` and \"unpacks\" them... i.e. it automates code of the form `a, b, c = Params.a, Params.b, Params.c`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NeoclassicalGrowth"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "The stochastic Neoclassical growth model type contains parameters\n",
    "which define the model\n",
    "\n",
    "* α: Capital share in output\n",
    "* β: Discount factor\n",
    "* δ: Depreciation rate\n",
    "* γ: Risk aversion\n",
    "* ρ: Persistence of the log of the productivity level\n",
    "* σ: Standard deviation of shocks to log productivity level\n",
    "* A: Coefficient on C-D production function\n",
    "\n",
    "* kgrid: Grid over capital\n",
    "* zgrid: Grid over productivity\n",
    "* grid: Grid of (k, z) pairs\n",
    "* eps_nodes: Nodes used to integrate\n",
    "* weights: Weights used to integrate\n",
    "* z1: A grid of the possible z1s tomorrow given eps_nodes and zgrid\n",
    "\"\"\"\n",
    "@with_kw struct NeoclassicalGrowth\n",
    "    # Parameters\n",
    "    α::Float64 = 0.36\n",
    "    β::Float64 = 0.99\n",
    "    δ::Float64 = 0.02\n",
    "    γ::Float64 = 2.0\n",
    "    ρ::Float64 = 0.95\n",
    "    σ::Float64 = 0.01\n",
    "    A::Float64 = (1.0/β - (1 - δ)) / α\n",
    "\n",
    "    # Grids\n",
    "    kgrid::Vector{Float64} = collect(linspace(0.9, 1.1, 10))\n",
    "    zgrid::Vector{Float64} = collect(linspace(0.9, 1.1, 10))\n",
    "    grid::Matrix{Float64} = gridmake(kgrid, zgrid)\n",
    "    eps_nodes::Vector{Float64} = qnwnorm(5, 0.0, σ^2)[1]\n",
    "    weights::Vector{Float64} = qnwnorm(5, 0.0, σ^2)[2]\n",
    "    z1::Matrix{Float64} = (zgrid.^(ρ))' .* exp.(eps_nodes)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also define some useful functions so that we [don't repeat ourselves](https://lectures.quantecon.org/py/writing_good_code.html#don-t-repeat-yourself) later in the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "expendables_t (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Helper functions\n",
    "f(ncgm::NeoclassicalGrowth, k, z) = z .* (ncgm.A*k.^ncgm.α)\n",
    "df(ncgm::NeoclassicalGrowth, k, z) = ncgm.α * z .* (ncgm.A * k.^(ncgm.α-1.0))\n",
    "\n",
    "u(ncgm::NeoclassicalGrowth, c) = c > 1e-10 ? (c^(1-ncgm.γ)-1)/(1-ncgm.γ) : -1e10\n",
    "du(ncgm::NeoclassicalGrowth, c) = c > 1e-10 ? c^(-ncgm.γ) : 1e10\n",
    "duinv(ncgm::NeoclassicalGrowth, u) = u .^ (-1 / ncgm.γ)\n",
    "\n",
    "expendables_t(ncgm::NeoclassicalGrowth, k, z) = (1-ncgm.δ)*k + f(ncgm, k, z)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solution Methods\n",
    "\n",
    "In this notebook, we describe the following solution methods:\n",
    "\n",
    "* Conventional Value Function Iteration\n",
    "* Envelope Condition Value Function Iteration\n",
    "* Envelope Condition Derivative Value Function Iteration\n",
    "* Endogenous Grid Value Function Iteration\n",
    "* Conventional Policy Function Iteration\n",
    "* Envelope Condition Policy Function Iteration\n",
    "* Euler Equation Method\n",
    "\n",
    "Each of these solution methods will have a very similar structure that follows a few basic steps:\n",
    "\n",
    "1. Guess a function (either value function or policy function).\n",
    "2. Using this function, update our guess of both the value and policy functions.\n",
    "3. Check whether the function we guessed and what it was updated to are similar enough. If so, proceed. If not, return to step 2 using the updated functions.\n",
    "4. Output the policy and value functions.\n",
    "\n",
    "In order to reduce the amount of repeated code and keep the exposition as clean as possible (the notebook is plenty long as is...), we will define multiple solution types that will have a more general (abstract) type called `SolutionMethod`. A solution can then be characterized by a concrete type `ValueCoeffs` (a special case for each solution method) which consists of an approximation degree, coefficients for the value function, and coefficients for the policy function. The rest of the functions below that are just more helper methods. We will then define a general solve method that applies steps 1, 3, and 4 from the algorithm above. Finally, we will implement a special method to do step 2 for each of the algorithms.\n",
    "\n",
    "These implementation may seem a bit confusing at first (though hopefully the idea itself feels intuitive) -- The implementation takes advantage of a powerful type system in Julia."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Types for solution methods\n",
    "abstract type SolutionMethod end\n",
    "\n",
    "struct IterateOnPolicy <: SolutionMethod end\n",
    "struct VFI_ECM <: SolutionMethod end\n",
    "struct VFI_EGM <: SolutionMethod end\n",
    "struct VFI <: SolutionMethod end\n",
    "struct PFI_ECM <: SolutionMethod end\n",
    "struct PFI <: SolutionMethod end\n",
    "struct dVFI_ECM <: SolutionMethod end\n",
    "struct EulEq <: SolutionMethod end\n",
    "\n",
    "#\n",
    "# Type for Approximating Value and Policy\n",
    "#\n",
    "mutable struct ValueCoeffs{T<:SolutionMethod,D<:Degree}\n",
    "    d::D\n",
    "    v_coeffs::Vector{Float64}\n",
    "    k_coeffs::Vector{Float64}\n",
    "end\n",
    "\n",
    "function ValueCoeffs{T<:SolutionMethod,d}(::Type{Val{d}}, method::T)\n",
    "    # Initialize two vectors of zeros\n",
    "    deg = Degree{d}()\n",
    "    n = n_complete(2, deg)\n",
    "    v_coeffs = zeros(n)\n",
    "    k_coeffs = zeros(n)\n",
    "\n",
    "    return ValueCoeffs{T,Degree{d}}(deg, v_coeffs, k_coeffs)\n",
    "end\n",
    "\n",
    "function ValueCoeffs{T<:SolutionMethod,d}(\n",
    "        ncgm::NeoclassicalGrowth, ::Type{Val{d}}, method::T\n",
    "    )\n",
    "    # Initialize with vector of zeros\n",
    "    deg = Degree{d}()\n",
    "    n = n_complete(2, deg)\n",
    "    v_coeffs = zeros(n)\n",
    "\n",
    "    # Policy guesses based on k and z\n",
    "    k, z = ncgm.grid[:, 1], ncgm.grid[:, 2]\n",
    "    css = ncgm.A - ncgm.δ\n",
    "    yss = ncgm.A\n",
    "    c_pol = f(ncgm, k, z) * (css/yss)\n",
    "\n",
    "    # Figure out what kp is\n",
    "    k_pol = expendables_t(ncgm, k, z) - c_pol\n",
    "    k_coeffs = complete_polynomial(ncgm.grid, d) \\ k_pol\n",
    "\n",
    "    return ValueCoeffs{T,Degree{d}}(deg, v_coeffs, k_coeffs)\n",
    "end\n",
    "\n",
    "solutionmethod{T<:SolutionMethod}(::ValueCoeffs{T}) = T\n",
    "\n",
    "# A few copy methods to make life easier\n",
    "Base.copy{T,D}(vp::ValueCoeffs{T,D}) =\n",
    "    ValueCoeffs{T,D}(vp.d, vp.v_coeffs, vp.k_coeffs)\n",
    "\n",
    "Base.copy{T1,D,T2<:SolutionMethod}(vp::ValueCoeffs{T1,D}, ::T2) =\n",
    "    ValueCoeffs{T2,D}(vp.d, vp.v_coeffs, vp.k_coeffs)\n",
    "\n",
    "function Base.copy{T,new_degree}(\n",
    "        ncgm::NeoclassicalGrowth, vp::ValueCoeffs{T}, ::Type{Val{new_degree}}\n",
    "    )\n",
    "    # Build Value and policy matrix\n",
    "    deg = Degree{new_degree}()\n",
    "    V = build_V(ncgm, vp)\n",
    "    k = build_k(ncgm, vp)\n",
    "\n",
    "    # Build new Phi\n",
    "    Phi = complete_polynomial(ncgm.grid, deg)\n",
    "    v_coeffs = Phi \\ V\n",
    "    k_coeffs = Phi \\ k\n",
    "\n",
    "    return ValueCoeffs{T,Degree{new_degree}}(deg, v_coeffs, k_coeffs)\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will need to repeatedly update coefficients, build $V$ (or $dV$ depending on the solution method), and be able to compute expected values, so we define some additional helper functions below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "build_k (generic function with 1 method)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Updates the coefficients for the value function inplace in `vp`\n",
    "\"\"\"\n",
    "function update_v!(vp::ValueCoeffs, new_coeffs::Vector{Float64}, dampen::Float64)\n",
    "    vp.v_coeffs = (1-dampen)*vp.v_coeffs + dampen*new_coeffs\n",
    "end\n",
    "\n",
    "\"\"\"\n",
    "Updates the coefficients for the policy function inplace in `vp`\n",
    "\"\"\"\n",
    "function update_k!(vp::ValueCoeffs, new_coeffs::Vector{Float64}, dampen::Float64)\n",
    "    vp.k_coeffs = (1-dampen)*vp.k_coeffs + dampen*new_coeffs\n",
    "end\n",
    "\n",
    "\"\"\"\n",
    "Builds either V or dV depending on the solution method that is given. If it\n",
    "is a solution method that iterates on the derivative of the value function\n",
    "then it will return derivative of the value function, otherwise the\n",
    "value function itself\n",
    "\"\"\"\n",
    "build_V_or_dV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs) =\n",
    "    build_V_or_dV(ncgm, vp, solutionmethod(vp)())\n",
    "\n",
    "build_V_or_dV(ncgm, vp::ValueCoeffs, ::SolutionMethod) = build_V(ncgm, vp)\n",
    "build_V_or_dV(ncgm, vp::ValueCoeffs, T::dVFI_ECM) = build_dV(ncgm, vp)\n",
    "\n",
    "function build_dV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)\n",
    "    Φ = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())\n",
    "    Φ*vp.v_coeffs\n",
    "end\n",
    "\n",
    "function build_V(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)\n",
    "    Φ = complete_polynomial(ncgm.grid, vp.d)\n",
    "    Φ*vp.v_coeffs\n",
    "end\n",
    "\n",
    "function build_k(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)\n",
    "    Φ = complete_polynomial(ncgm.grid, vp.d)\n",
    "    Φ*vp.k_coeffs\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, in order to evaluate the value function, we will need to be able to take expectations.\n",
    "\n",
    "These functions evaluates expectations by taking the policy $k_{t+1}$ and the current productivity state $z_t$ as inputs. They then integrates over the possible $z_{t+1}$s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "compute_dEV (generic function with 1 method)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function compute_EV!(cp_kpzp::Vector{Float64}, ncgm::NeoclassicalGrowth,\n",
    "                     vp::ValueCoeffs, kp, iz)\n",
    "    # Pull out information from types\n",
    "    z1, weightsz = ncgm.z1, ncgm.weights\n",
    "\n",
    "    # Get number nodes\n",
    "    nzp = length(weightsz)\n",
    "\n",
    "    EV = 0.0\n",
    "    for izp in 1:nzp\n",
    "        zp = z1[izp, iz]\n",
    "        complete_polynomial!(cp_kpzp, [kp, zp], vp.d)\n",
    "        EV += weightsz[izp] * dot(vp.v_coeffs, cp_kpzp)\n",
    "    end\n",
    "\n",
    "    return EV\n",
    "end\n",
    "\n",
    "function compute_EV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs, kp, iz)\n",
    "    cp_kpzp = Array{Float64}(n_complete(2, vp.d))\n",
    "\n",
    "    return compute_EV!(cp_kpzp, ncgm, vp, kp, iz)\n",
    "end\n",
    "\n",
    "function compute_EV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)\n",
    "    # Get length of k and z grids\n",
    "    kgrid, zgrid = ncgm.kgrid, ncgm.zgrid\n",
    "    nk, nz = length(kgrid), length(zgrid)\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "\n",
    "    # Allocate space to store EV\n",
    "    EV = Array{Float64}(nk*nz)\n",
    "\n",
    "    for ik in 1:nk, iz in 1:nz\n",
    "        # Pull out states\n",
    "        k = kgrid[ik]\n",
    "        z = zgrid[iz]\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "\n",
    "        # Pass to scalar EV\n",
    "        complete_polynomial!(temp, [k, z], vp.d)\n",
    "        kp = dot(vp.k_coeffs, temp)\n",
    "        EV[ikiz_index] = compute_EV!(temp, ncgm, vp, kp, iz)\n",
    "    end\n",
    "\n",
    "    return EV\n",
    "end\n",
    "\n",
    "\n",
    "function compute_dEV!(cp_dkpzp::Vector,  ncgm::NeoclassicalGrowth,\n",
    "                      vp::ValueCoeffs, kp, iz)\n",
    "    # Pull out information from types\n",
    "    z1, weightsz = ncgm.z1, ncgm.weights\n",
    "\n",
    "    # Get number nodes\n",
    "    nzp = length(weightsz)\n",
    "\n",
    "    dEV = 0.0\n",
    "    for izp in 1:nzp\n",
    "        zp = z1[izp, iz]\n",
    "        complete_polynomial!(cp_dkpzp, [kp, zp], vp.d, Derivative{1}())\n",
    "        dEV += weightsz[izp] * dot(vp.v_coeffs, cp_dkpzp)\n",
    "    end\n",
    "\n",
    "    return dEV\n",
    "end\n",
    "\n",
    "function compute_dEV(ncgm::NeoclassicalGrowth, vp::ValueCoeffs, kp, iz)\n",
    "    compute_dEV!(Array{Float64}(n_complete(2, vp.d)), ncgm, vp, kp, iz)\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  General Solution Method\n",
    "\n",
    "As promised, below is some code that \"generally\" applies the algorithm that we described -- Notice that it is implemented for a type `ValueCoeffs{SolutionMethod}` which is our abstract type. We will define a special version of `update` for each solution method and then we will only need this one `solve` method and won't repeat the more tedious portions of our code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "solve (generic function with 1 method)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function solve{T<:SolutionMethod}(ncgm::NeoclassicalGrowth, vp::ValueCoeffs{T};\n",
    "                                  tol::Float64=1e-6, maxiter::Int=5000, dampen::Float64=1.0,\n",
    "                                  nskipprint::Int=1, verbose::Bool=true)\n",
    "    # Get number of k and z on grid\n",
    "    nk, nz = length(ncgm.kgrid), length(ncgm.zgrid)\n",
    "\n",
    "    # Build basis matrix and value function\n",
    "    dPhi = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())\n",
    "    Phi = complete_polynomial(ncgm.grid, vp.d)\n",
    "    V = build_V_or_dV(ncgm, vp)\n",
    "    k = build_k(ncgm, vp)\n",
    "    Vnew = copy(V)\n",
    "    knew = copy(k)\n",
    "\n",
    "    # Print column names\n",
    "    if verbose\n",
    "        @printf(\"| Iteration | Distance V | Distance K |\\n\")\n",
    "    end\n",
    "\n",
    "    # Iterate to convergence\n",
    "    dist, iter = 10.0, 0\n",
    "    while (tol < dist) & (iter < maxiter)\n",
    "        # Update the value function using appropriate update methods\n",
    "        update!(Vnew, knew, ncgm, vp, Phi, dPhi)\n",
    "\n",
    "        # Compute distance and update all relevant elements\n",
    "        iter += 1\n",
    "        dist_v = maximum(abs, 1.0 .- Vnew./V)\n",
    "        dist_k = maximum(abs, 1.0 .- knew./k)\n",
    "        copy!(V, Vnew)\n",
    "        copy!(k, knew)\n",
    "\n",
    "        # If we are iterating on a policy, use the difference of values\n",
    "        # otherwise use the distance on policy\n",
    "        dist = ifelse(solutionmethod(vp) == IterateOnPolicy, dist_v, dist_k)\n",
    "\n",
    "        # Print status update\n",
    "        if verbose && (iter%nskipprint == 0)\n",
    "            @printf(\"|%-11d|%-12e|%-12e|\\n\", iter, dist_v, dist_k)\n",
    "        end\n",
    "    end\n",
    "\n",
    "    # Update value and policy functions one last time as long as the\n",
    "    # solution method isn't IterateOnPolicy\n",
    "    if ~(solutionmethod(vp) == IterateOnPolicy)\n",
    "        # Update capital policy after finished\n",
    "        kp = env_condition_kp(ncgm, vp)\n",
    "        update_k!(vp, complete_polynomial(ncgm.grid, vp.d) \\ kp, 1.0)\n",
    "\n",
    "        # Update value function according to specified policy\n",
    "        vp_igp = copy(vp, IterateOnPolicy())\n",
    "        solve(ncgm, vp_igp; tol=1e-10, maxiter=5000, verbose=false)\n",
    "        update_v!(vp, vp_igp.v_coeffs, 1.0)\n",
    "\n",
    "    end\n",
    "\n",
    "    return vp\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterating to Convergence (given policy)\n",
    "\n",
    "This isn't one of the methods described above, but it is used as an element of a few of our methods (and also as a way to get a first guess at the value function). This method takes an initial policy function, $\\bar{k}(k_t, z_t)$, as given, and then, without changing the policy, iterates until the value function has converged.\n",
    "\n",
    "Thus the \"update section\" of the algorithm in this instance would be:\n",
    "\n",
    "* Leave policy function unchanged\n",
    "* At each point of grid, $(k_t, z_t)$, compute $\\hat{V}(k_t, z_t) = u(c(\\bar{k}(k_t, z_t))) + \\beta E \\left[ V(\\bar{k}(k_t, z_t), z_{t+1}) \\right]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 1 method)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{IterateOnPolicy},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid;\n",
    "    nk, nz = length(kgrid), length(zgrid)\n",
    "\n",
    "    # Iterate over all states\n",
    "    for ik in 1:nk, iz in 1:nz\n",
    "        # Pull out states\n",
    "        k = kgrid[ik]\n",
    "        z = zgrid[iz]\n",
    "\n",
    "        # Pull out policy and evaluate consumption\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        k1 = kpol[ikiz_index]\n",
    "        c = expendables_t(ncgm, k, z) - k1\n",
    "\n",
    "        # New value\n",
    "        EV = compute_EV(ncgm, vp, k1, iz)\n",
    "        V[ikiz_index] = u(ncgm, c) + ncgm.β*EV\n",
    "    end\n",
    "\n",
    "    # Update coefficients\n",
    "    update_v!(vp, Φ \\ V, 1.0)\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "\n",
    "    return V\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conventional Value Function Iteration\n",
    "\n",
    "This is one of the first solution methods for macroeconomics a graduate student in economics typically learns.\n",
    "\n",
    "In this solution method, one takes as given a value function, $V(k_t, z_t)$, and then solves for the optimal policy given the value function.\n",
    "\n",
    "The update section takes the form:\n",
    "\n",
    "* For each point, $(k_t, z_t)$, numerically solve for $c^*(k_t, z_t)$ to satisfy the first order condition $u'(c^*) = \\beta E \\left[ V_1((1 - \\delta) k_t + z_t f(k_t) - c^*, z_{t+1}) \\right]$\n",
    "* Define $k^*(k_t, z_t) = (1 - \\delta) k_t + z_t f(k_t) - c^*(k_t, z_t)$\n",
    "* Update value function according to $\\hat{V}(k_t, z_t) = u(c^*(k_t, z_t)) + \\beta E \\left[ V(k^*(k_t, z_t), z_{t+1}) \\right]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 2 methods)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{VFI},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid\n",
    "    nk, nz = length(kgrid), length(zgrid)\n",
    "\n",
    "    # Iterate over all states\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for iz=1:nz, ik=1:nk\n",
    "        k = kgrid[ik]; z = zgrid[iz]\n",
    "\n",
    "        # Define an objective function (negative for minimization)\n",
    "        y = expendables_t(ncgm, k, z)\n",
    "        solme(kp) = du(ncgm, y - kp) - ncgm.β*compute_dEV!(temp, ncgm, vp, kp, iz)\n",
    "\n",
    "        # Find sol to foc\n",
    "        kp = brent(solme, 1e-8, y-1e-8; rtol=1e-12)\n",
    "        c = expendables_t(ncgm, k, z) - kp\n",
    "\n",
    "        # New value\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        EV = compute_EV!(temp, ncgm, vp, kp, iz)\n",
    "        V[ikiz_index] = u(ncgm, c) + ncgm.β*EV\n",
    "        kpol[ikiz_index] = kp\n",
    "    end\n",
    "\n",
    "    # Update coefficients\n",
    "    update_v!(vp, Φ \\ V, 1.0)\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "\n",
    "    return V\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Endogenous Grid Value Function Iteration\n",
    "\n",
    "Method introduced by Chris Carroll. The key to this method is that the grid of points being used to approximate is over $(k_{t+1}, z_{t})$ instead of $(k_t, z_t)$. The insightful piece of this algorithm is that the transformation allows one to write a closed form for the consumption function, $c^*(k_{t+1}, z_t) = u'^{-1} \\left( V_1(k_{t+1}, z_{t+1}) \\right]$.\n",
    "\n",
    "Then for a given $(k_{t+1}, z_{t})$ the update section would be\n",
    "\n",
    "* Define $c^*(k_{t+1}, z_t) = u'^{-1} \\left( V_1(k_{t+1}, z_{t+1}) \\right]$\n",
    "* Find $k_t$ such that $k_t = (1 - \\delta) k_t + z_t f(k_t) - k_{t+1}$\n",
    "* Update value function according to $\\hat{V}(k_t, z_t) = u(c^*(k_{t+1}, z_t)) + \\beta E \\left[ V(k_{t+1}, z_{t+1}) \\right]$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 3 methods)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{VFI_EGM},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;\n",
    "    nk, nz = length(kgrid), length(zgrid)\n",
    "\n",
    "    # Iterate\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for iz=1:nz, ik=1:nk\n",
    "\n",
    "        # In EGM we use the grid points as if they were our\n",
    "        # policy for yesterday and find implied kt\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        k1 = kgrid[ik];z = zgrid[iz];\n",
    "\n",
    "        # Compute the derivative of expected values\n",
    "        dEV = compute_dEV!(temp, ncgm, vp, k1, iz)\n",
    "\n",
    "        # Compute optimal consumption\n",
    "        c = duinv(ncgm, ncgm.β*dEV)\n",
    "\n",
    "        # Need to find corresponding kt for optimal c\n",
    "        obj(kt) = expendables_t(ncgm, kt, z) - c - k1\n",
    "        kt_star = brent(obj, 0.0, 2.0, xtol=1e-10)\n",
    "\n",
    "        # New value\n",
    "        EV = compute_EV!(temp, ncgm, vp, k1, iz)\n",
    "        V[ikiz_index] = u(ncgm, c) + ncgm.β*EV\n",
    "        kpol[ikiz_index] = kt_star\n",
    "    end\n",
    "\n",
    "    # New Φ (has our new \"kt_star\" and z points)\n",
    "    Φ_egm = complete_polynomial([kpol grid[:, 2]], vp.d)\n",
    "\n",
    "    # Update coefficients\n",
    "    update_v!(vp, Φ_egm \\ V, 1.0)\n",
    "    update_k!(vp, Φ_egm \\ grid[:, 1], 1.0)\n",
    "\n",
    "    # Update V and kpol to be value and policy corresponding\n",
    "    # to our grid again\n",
    "    copy!(V, Φ*vp.v_coeffs)\n",
    "    copy!(kpol, Φ*vp.k_coeffs)\n",
    "\n",
    "    return V\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Envelope Condition Value Function Iteration\n",
    "\n",
    "Very similar to the previous method. The insight of this algorithm is that since we are already approximating the value function and can evaluate its derivative, we can skip the numerical optimization piece of the update method and compute directly the policy using the envelope condition (hence the name).\n",
    "\n",
    "The envelope condition says:\n",
    "\n",
    "$$c^*(k_t, z_t) = u'^{-1} \\left( \\frac{\\partial V(k_t, z_t)}{\\partial k_t} (1 - \\delta + r)^{-1} \\right)$$\n",
    "\n",
    "so\n",
    "\n",
    "$$k^*(k_t, z_t) = z_t f(k_t) + (1-\\delta)k_t - c^*(k_t, z_t)$$\n",
    "\n",
    "The functions below compute the policy using the envelope condition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "env_condition_kp (generic function with 2 methods)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function env_condition_kp!(cp_out::Vector{Float64}, ncgm::NeoclassicalGrowth,\n",
    "                           vp::ValueCoeffs, k::Float64, z::Float64)\n",
    "    # Compute derivative of VF\n",
    "    dV = dot(vp.v_coeffs, complete_polynomial!(cp_out, [k, z], vp.d, Derivative{1}()))\n",
    "\n",
    "    # Consumption is then computed as\n",
    "    c = duinv(ncgm, dV / (1 - ncgm.δ + df(ncgm, k, z)))\n",
    "\n",
    "    return expendables_t(ncgm, k, z) - c\n",
    "end\n",
    "\n",
    "function env_condition_kp(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,\n",
    "                          k::Float64, z::Float64)\n",
    "    cp_out = Array{Float64}(n_complete(2, vp.d))\n",
    "    env_condition_kp!(cp_out, ncgm, vp, k, z)\n",
    "end\n",
    "\n",
    "function env_condition_kp(ncgm::NeoclassicalGrowth, vp::ValueCoeffs)\n",
    "    # Pull out k and z from grid\n",
    "    k = ncgm.grid[:, 1]\n",
    "    z = ncgm.grid[:, 2]\n",
    "\n",
    "    # Create basis matrix for entire grid\n",
    "    dPhi = complete_polynomial(ncgm.grid, vp.d, Derivative{1}())\n",
    "\n",
    "    # Compute consumption\n",
    "    c = duinv(ncgm, (dPhi*vp.v_coeffs) ./ (1-ncgm.δ+df(ncgm, k, z)))\n",
    "\n",
    "    return expendables_t(ncgm, k, z) .- c\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The update method is then very similar to other value iteration style methods, but avoids the numerical solver.\n",
    "\n",
    "* For each point, $(k_t, z_t)$ get $c^*(k_t, z_t)$ from the envelope condition\n",
    "* Define $k^*(k_t, z_t) = (1 - \\delta) k_t + z_t f(k_t) - c^*(k_t, z_t)$\n",
    "* Update value function according to $\\hat{V}(k_t, z_t) = u(c^*(k_t, z_t)) + \\beta E \\left[ V(k^*(k_t, z_t), z_{t+1}) \\right]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 4 methods)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{VFI_ECM},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid;\n",
    "    nk, nz = length(kgrid), length(zgrid)\n",
    "\n",
    "    # Iterate over all states\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for ik in 1:nk, iz in 1:nz\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        k = kgrid[ik]\n",
    "        z = zgrid[iz]\n",
    "\n",
    "        # Policy from envelope condition\n",
    "        kp = env_condition_kp!(temp, ncgm, vp, k, z)\n",
    "        c = expendables_t(ncgm, k, z) - kp\n",
    "        kpol[ikiz_index] = kp\n",
    "\n",
    "        # New value\n",
    "        EV = compute_EV!(temp, ncgm, vp, kp, iz)\n",
    "        V[ikiz_index] = u(ncgm, c) + ncgm.β*EV\n",
    "    end\n",
    "\n",
    "    # Update coefficients\n",
    "    update_v!(vp, Φ \\ V, 1.0)\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "\n",
    "    return V\n",
    "end\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Envelope Condition Derivative Value Function Iteration\n",
    "\n",
    "This method uses the same insight of the \"Envelope Condition Value Function Iteration,\" but, rather than iterate directly on the value function, it iterates on the derivative of the value function. The update steps are\n",
    "\n",
    "* For each point, $(k_t, z_t)$ get $c^*(k_t, z_t)$ from the envelope condition (which only depends on the derivative of the value function!)\n",
    "* Define $k^*(k_t, z_t) = (1 - \\delta) k_t + z_t f(k_t) - c^*(k_t, z_t)$\n",
    "* Update value function according to $\\hat{V}_1(k_t, z_t) = \\beta (1 - \\delta + z_t f'(k_t)) E \\left[ V_1(k^*(k_t, z_t), z_{t+1}) \\right]$\n",
    "\n",
    "Once it has converged, you use the implied policy rule and iterate to convergence using the \"iterate to convergence (given policy)\" method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 5 methods)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(dV::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{dVFI_ECM},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;\n",
    "    nk, nz, ns = length(kgrid), length(zgrid), size(grid, 1)\n",
    "\n",
    "    # Iterate over all states\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for iz=1:nz, ik=1:nk\n",
    "        k = kgrid[ik]; z = zgrid[iz];\n",
    "\n",
    "        # Envelope condition implies optimal kp\n",
    "        kp = env_condition_kp!(temp, ncgm, vp, k, z)\n",
    "        c = expendables_t(ncgm, k, z) - kp\n",
    "\n",
    "        # New value\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        dEV = compute_dEV!(temp, ncgm, vp, kp, iz)\n",
    "        dV[ikiz_index] = (1-ncgm.δ+df(ncgm, k, z))*ncgm.β*dEV\n",
    "        kpol[ikiz_index] = kp\n",
    "    end\n",
    "\n",
    "    # Get new coeffs\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "    update_v!(vp, dΦ \\ dV, 1.0)\n",
    "\n",
    "    return dV\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conventional Policy Function Iteration\n",
    "\n",
    "Policy function iteration is different than value function iteration in that it starts with a policy function, then updates the value function, and finally finds the new optimal policy function. Given a policy $c(k_t, z_t)$ and for each pair $(k_t, z_t)$\n",
    "\n",
    "* Define $k(k_t, z_t) = (1 - \\delta) k_t + z_t f(k_t) - c(k_t, z_t)$\n",
    "* Find fixed point of $V(k_t, z_t) = u(c(k_t, z_t)) + \\beta E \\left[ V(k(k_t, z_t), z_t) \\right]$ (Iterate to convergence given policy)\n",
    "* Given $V(k_t, z_t)$, numerically solve for new policy $c^*(k_t, z_t)$ -- Stop when $c(k_t, z_t) \\approx c^*(k_t, z_t)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 6 methods)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{PFI},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;\n",
    "    nk, nz = length(kgrid), length(zgrid)\n",
    "\n",
    "    # Copy valuecoeffs object and use to iterate to\n",
    "    # convergence given a policy\n",
    "    vp_igp = copy(vp, IterateOnPolicy())\n",
    "    solve(ncgm, vp_igp; nskipprint=1000, maxiter=5000, verbose=false)\n",
    "\n",
    "    # Update the policy and values\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for ik in 1:nk, iz in 1:nz\n",
    "        k = kgrid[ik]; z = zgrid[iz];\n",
    "\n",
    "        # Define an objective function (negative for minimization)\n",
    "        y = expendables_t(ncgm, k, z)\n",
    "        solme(kp) = du(ncgm, y - kp) - ncgm.β*compute_dEV!(temp, ncgm, vp, kp, iz)\n",
    "\n",
    "        # Find minimum of objective\n",
    "        kp = brent(solme, 1e-8, y-1e-8; rtol=1e-12)\n",
    "\n",
    "        # Update policy function\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        kpol[ikiz_index] = kp\n",
    "    end\n",
    "\n",
    "    # Get new coeffs\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "    update_v!(vp, vp_igp.v_coeffs, 1.0)\n",
    "\n",
    "    # Update all elements of value\n",
    "    copy!(V, Φ*vp.v_coeffs)\n",
    "\n",
    "    return V\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Envelope Condition Policy Function Iteration\n",
    "\n",
    "Similar to policy function iteration, but, rather than numerically solve for new policies, it uses the envelope condition to directly compute them. Given a starting policy $c(k_t, z_t)$ and for each pair $(k_t, z_t)$\n",
    "\n",
    "* Define $k(k_t, z_t) = (1 - \\delta) k_t + z_t f(k_t) - c(k_t, z_t)$\n",
    "* Find fixed point of $V(k_t, z_t) = u(c(k_t, z_t)) + \\beta E \\left[ V(k(k_t, z_t), z_t) \\right]$ (Iterate to convergence given policy)\n",
    "* Given $V(k_t, z_t)$ find $c^*(k_t, z_t)$ using envelope condition -- Stop when $c(k_t, z_t) \\approx c^*(k_t, z_t)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 7 methods)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{PFI_ECM},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Copy valuecoeffs object and use to iterate to\n",
    "    # convergence given a policy\n",
    "    vp_igp = copy(vp, IterateOnPolicy())\n",
    "    solve(ncgm, vp_igp; nskipprint=1000, maxiter=5000, verbose=false)\n",
    "\n",
    "    # Update the policy and values\n",
    "    kp = env_condition_kp(ncgm, vp)\n",
    "    update_k!(vp, Φ \\ kp, 1.0)\n",
    "    update_v!(vp, vp_igp.v_coeffs, 1.0)\n",
    "\n",
    "    # Update all elements of value\n",
    "    copy!(V, Φ*vp.v_coeffs)\n",
    "    copy!(kpol, kp)\n",
    "\n",
    "    return V\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Euler Equation Method\n",
    "\n",
    "Euler equation methods operate directly on the Euler equation: $u'(c_t) = \\beta E \\left[ u'(c_{t+1}) (1 - \\delta + z_t f'(k_t)) \\right]$.\n",
    "\n",
    "Given an initial policy $c(k_t, z_t)$ for each grid point $(k_t, z_t)$\n",
    "\n",
    "* Find $k(k_t, z_t) = (1-\\delta)k_t + z_t f(k_t) - c(k_t, z_t)$\n",
    "* Let $c_{t+1} = c(k(k_t, z_t), z_t)$\n",
    "* Numerically solve for a $c^*$ that satisfies the Euler equation i.e. $u'(c^*) = \\beta E \\left[ u'(c_{t+1}) (1 - \\delta + z_t f'(k_t)) \\right]$\n",
    "* Stop when $c^*(k_t, z_t) \\approx c(k_t, z_t)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "update! (generic function with 8 methods)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function update!(dV::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{dVFI_ECM},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    kgrid = ncgm.kgrid; zgrid = ncgm.zgrid; grid = ncgm.grid;\n",
    "    nk, nz, ns = length(kgrid), length(zgrid), size(grid, 1)\n",
    "\n",
    "    # Iterate over all states\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for iz=1:nz, ik=1:nk\n",
    "        k = kgrid[ik]; z = zgrid[iz];\n",
    "\n",
    "        # Envelope condition implies optimal kp\n",
    "        kp = env_condition_kp!(temp, ncgm, vp, k, z)\n",
    "        c = expendables_t(ncgm, k, z) - kp\n",
    "\n",
    "        # New value\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        dEV = compute_dEV!(temp, ncgm, vp, kp, iz)\n",
    "        dV[ikiz_index] = (1-ncgm.δ+df(ncgm, k, z))*ncgm.β*dEV\n",
    "        kpol[ikiz_index] = kp\n",
    "    end\n",
    "\n",
    "    # Get new coeffs\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "    update_v!(vp, dΦ \\ dV, 1.0)\n",
    "\n",
    "    return dV\n",
    "end\n",
    "\n",
    "# Conventional Euler equation method\n",
    "function update!(V::Vector{Float64}, kpol::Vector{Float64},\n",
    "                 ncgm::NeoclassicalGrowth, vp::ValueCoeffs{EulEq},\n",
    "                 Φ::Matrix{Float64}, dΦ::Matrix{Float64})\n",
    "    # Get sizes and allocate for complete_polynomial\n",
    "    @unpack kgrid, zgrid, weights, z1 = ncgm\n",
    "    nz1, nz = size(z1)\n",
    "    nk = length(kgrid)\n",
    "\n",
    "    # Iterate over all states\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for iz in 1:nz, ik in 1:nk\n",
    "        k = kgrid[ik]; z = zgrid[iz];\n",
    "\n",
    "        # Create current polynomial\n",
    "        complete_polynomial!(temp, [k, z], vp.d)\n",
    "\n",
    "        # Compute what capital will be tomorrow according to policy\n",
    "        kp = dot(temp, vp.k_coeffs)\n",
    "\n",
    "        # Compute RHS of EE\n",
    "        rhs_ee = 0.0\n",
    "        for iz1 in 1:nz1\n",
    "            # Possible z in t+1\n",
    "            zp = z1[iz1, iz]\n",
    "\n",
    "            # Policy for k_{t+2}\n",
    "            complete_polynomial!(temp, [kp, zp], vp.d)\n",
    "            kpp = dot(temp, vp.k_coeffs)\n",
    "\n",
    "            # Implied t+1 consumption\n",
    "            cp = expendables_t(ncgm, kp, zp) - kpp\n",
    "\n",
    "            # Add to running expectation\n",
    "            rhs_ee += ncgm.β*weights[iz1]*du(ncgm, cp)*(1-ncgm.δ+df(ncgm, kp, zp))\n",
    "        end\n",
    "\n",
    "        # The rhs of EE implies consumption and investment in t\n",
    "        c = duinv(ncgm, rhs_ee)\n",
    "        kp_star = expendables_t(ncgm, k, z) - c\n",
    "\n",
    "        # New value\n",
    "        ikiz_index = sub2ind((nk, nz), ik, iz)\n",
    "        EV = compute_EV!(temp, ncgm, vp, kp_star, iz)\n",
    "        V[ikiz_index] = u(ncgm, c) + ncgm.β*EV\n",
    "        kpol[ikiz_index] = kp_star\n",
    "    end\n",
    "\n",
    "    # Update coefficients\n",
    "    update_v!(vp, Φ \\ V, 1.0)\n",
    "    update_k!(vp, Φ \\ kpol, 1.0)\n",
    "\n",
    "    return V\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulation and Euler Error Methods\n",
    "\n",
    "The following functions to simulate and compute Euler errors are easily defined given our model type and the solution type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ee_residuals (generic function with 2 methods)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"\"\"\n",
    "Simulates the neoclassical growth model for a given set of solution\n",
    "coefficients. It simulates for `capT` periods and discards first\n",
    "`nburn` observations.\n",
    "\"\"\"\n",
    "function simulate(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,\n",
    "                  shocks::Vector{Float64}; capT::Int=10_000,\n",
    "                  nburn::Int=200)\n",
    "    # Unpack parameters\n",
    "    kp = 0.0  # Policy holder\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "\n",
    "    # Allocate space for k and z\n",
    "    ksim = Array{Float64}(capT+nburn)\n",
    "    zsim = Array{Float64}(capT+nburn)\n",
    "\n",
    "    # Initialize both k and z at 1\n",
    "    ksim[1] = 1.0\n",
    "    zsim[1] = 1.0\n",
    "\n",
    "    # Simulate\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "    for t in 2:capT+nburn\n",
    "        # Evaluate k_t given yesterday's (k_{t-1}, z_{t-1})\n",
    "        kp = env_condition_kp!(temp, ncgm, vp, ksim[t-1], zsim[t-1])\n",
    "\n",
    "        # Draw new z and update k using policy above\n",
    "        zsim[t] = zsim[t-1]^ncgm.ρ * exp(ncgm.σ*shocks[t])\n",
    "        ksim[t] = kp\n",
    "    end\n",
    "\n",
    "    return ksim[nburn+1:end], zsim[nburn+1:end]\n",
    "end\n",
    "\n",
    "function simulate(ncgm::NeoclassicalGrowth, vp::ValueCoeffs;\n",
    "                  capT::Int=10_000, nburn::Int=200, seed=42)\n",
    "    srand(seed)  # Set specific seed\n",
    "    shocks = randn(capT + nburn)\n",
    "\n",
    "    return simulate(ncgm, vp, shocks; capT=capT, nburn=nburn)\n",
    "end\n",
    "\n",
    "\"\"\"\n",
    "This function evaluates the Euler Equation residual for a single point (k, z)\n",
    "\"\"\"\n",
    "function EulerEquation!(out::Vector{Float64}, ncgm::NeoclassicalGrowth,\n",
    "                        vp::ValueCoeffs, k::Float64, z::Float64,\n",
    "                        nodes::Vector{Float64}, weights::Vector{Float64})\n",
    "    # Evaluate consumption today\n",
    "    k1 = env_condition_kp!(out, ncgm, vp, k, z)\n",
    "    c = expendables_t(ncgm, k, z) - k1\n",
    "    LHS = du(ncgm, c)\n",
    "\n",
    "    # For each of realizations tomorrow, evaluate expectation on RHS\n",
    "    RHS = 0.0\n",
    "    for (eps, w) in zip(nodes, weights)\n",
    "        # Compute ztp1\n",
    "        z1 = z^ncgm.ρ * exp(eps)\n",
    "\n",
    "        # Evaluate the ktp2\n",
    "        ktp2 = env_condition_kp!(out, ncgm, vp, k1, z1)\n",
    "\n",
    "        # Get c1\n",
    "        c1 = expendables_t(ncgm, k1, z1) - ktp2\n",
    "\n",
    "        # Update RHS of equation\n",
    "        RHS = RHS + w*du(ncgm, c1)*(1 - ncgm.δ + df(ncgm, k1, z1))\n",
    "    end\n",
    "\n",
    "    return abs(ncgm.β*RHS/LHS - 1.0)\n",
    "end\n",
    "\n",
    "\"\"\"\n",
    "Given simulations for k and z, it computes the euler equation residuals\n",
    "along the entire simulation. It reports the mean and max values in\n",
    "log10.\n",
    "\"\"\"\n",
    "function ee_residuals(ncgm::NeoclassicalGrowth, vp::ValueCoeffs,\n",
    "                      ksim::Vector{Float64}, zsim::Vector{Float64}; Qn::Int=10)\n",
    "    # Figure out how many periods we simulated for and make sure k and z\n",
    "    # are same length\n",
    "    capT = length(ksim)\n",
    "    @assert length(zsim) == capT\n",
    "\n",
    "    # Finer integration nodes\n",
    "    eps_nodes, weight_nodes = qnwnorm(Qn, 0.0, ncgm.σ^2)\n",
    "    temp = Array{Float64}(n_complete(2, vp.d))\n",
    "\n",
    "    # Compute EE for each period\n",
    "    EE_resid = Array{Float64}(capT)\n",
    "    for t=1:capT\n",
    "        # Pull out current state\n",
    "        k, z = ksim[t], zsim[t]\n",
    "\n",
    "        # Compute residual of Euler Equation\n",
    "        EE_resid[t] = EulerEquation!(temp, ncgm, vp, k, z, eps_nodes, weight_nodes)\n",
    "    end\n",
    "\n",
    "    return EE_resid\n",
    "end\n",
    "\n",
    "function ee_residuals(ncgm::NeoclassicalGrowth, vp::ValueCoeffs; Qn::Int=10)\n",
    "    # Simulate and then call other ee_residuals method\n",
    "    ksim, zsim = simulate(ncgm, vp)\n",
    "\n",
    "    return ee_residuals(ncgm, vp, ksim, zsim; Qn=Qn)\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A Horse Race\n",
    "\n",
    "We can now run a horse race to compare the methods in terms of both accuracy and speed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "main (generic function with 3 methods)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function main(sm::SolutionMethod, nd::Int=5, shocks=randn(capT+nburn);\n",
    "              capT=10_000, nburn=200, tol=1e-9, maxiter=2500,\n",
    "              nskipprint=25, verbose=true)\n",
    "    # Create model\n",
    "    ncgm = NeoclassicalGrowth()\n",
    "\n",
    "    # Create initial quadratic guess\n",
    "    vp = ValueCoeffs(ncgm, Val{2}, IterateOnPolicy())\n",
    "    solve(ncgm, vp; tol=1e-6, verbose=false)\n",
    "\n",
    "    # Allocate memory for timings\n",
    "    times = Array{Float64}(nd-1)\n",
    "    sols = Array{ValueCoeffs}(nd-1)\n",
    "    mean_ees = Array{Float64}(nd-1)\n",
    "    max_ees = Array{Float64}(nd-1)\n",
    "\n",
    "    # Solve using the solution method for degree 2 to 5\n",
    "    vp = copy(vp, sm)\n",
    "    for d in 2:nd\n",
    "        # Change degree of solution method\n",
    "        vp = copy(ncgm, vp, Val{d})\n",
    "\n",
    "        # Time the current method\n",
    "        start_time = time()\n",
    "        solve(ncgm, vp; tol=tol, maxiter=maxiter, nskipprint=nskipprint,\n",
    "              verbose=verbose)\n",
    "        end_time = time()\n",
    "\n",
    "        # Save the time and solution\n",
    "        times[d-1] = end_time - start_time\n",
    "        sols[d-1] = vp\n",
    "\n",
    "        # Simulate and compute EE\n",
    "        ks, zs = simulate(ncgm, vp, shocks; capT=capT, nburn=nburn)\n",
    "        resids = ee_residuals(ncgm, vp, ks, zs; Qn=10)\n",
    "        mean_ees[d-1] = log10.(mean(abs.(resids)))\n",
    "        max_ees[d-1] = log10.(maximum(abs, resids))\n",
    "    end\n",
    "\n",
    "    return sols, times, mean_ees, max_ees\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution Method: VFI()\n",
      "\tDegree 2 took time 1.5208711624145508\n",
      "\tMean & Max EE are-3.803 & -2.875\n",
      "\tDegree 3 took time 1.144477128982544\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 1.2640371322631836\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 1.1280219554901123\n",
      "\tMean & Max EE are-6.916 & -4.942\n",
      "Solution Method: VFI_EGM()\n",
      "\tDegree 2 took time 0.6876471042633057\n",
      "\tMean & Max EE are-3.803 & -2.876\n",
      "\tDegree 3 took time 0.4869670867919922\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 0.5630497932434082\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 0.3711881637573242\n",
      "\tMean & Max EE are-6.916 & -4.942\n",
      "Solution Method: VFI_ECM()\n",
      "\tDegree 2 took time 0.282289981842041\n",
      "\tMean & Max EE are-3.803 & -2.875\n",
      "\tDegree 3 took time 0.21788907051086426\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 0.2694680690765381\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 0.17006278038024902\n",
      "\tMean & Max EE are-6.916 & -4.942\n",
      "Solution Method: dVFI_ECM()\n",
      "\tDegree 2 took time 0.577113151550293\n",
      "\tMean & Max EE are-3.803 & -2.876\n",
      "\tDegree 3 took time 0.82600998878479\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 0.9907219409942627\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 1.2688839435577393\n",
      "\tMean & Max EE are-6.916 & -4.942\n",
      "Solution Method: PFI()\n",
      "\tDegree 2 took time 0.30478405952453613\n",
      "\tMean & Max EE are-3.803 & -2.876\n",
      "\tDegree 3 took time 1.19022798538208\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 1.6783649921417236\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 1.3011729717254639\n",
      "\tMean & Max EE are-6.916 & -4.942\n",
      "Solution Method: PFI_ECM()\n",
      "\tDegree 2 took time 0.29001903533935547\n",
      "\tMean & Max EE are-3.803 & -2.876\n",
      "\tDegree 3 took time 0.24599814414978027\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 0.29113221168518066\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 0.16288995742797852\n",
      "\tMean & Max EE are-6.916 & -4.942\n",
      "Solution Method: EulEq()\n",
      "\tDegree 2 took time 0.37535905838012695\n",
      "\tMean & Max EE are-3.803 & -2.876\n",
      "\tDegree 3 took time 0.2786898612976074\n",
      "\tMean & Max EE are-4.914 & -3.487\n",
      "\tDegree 4 took time 0.32964205741882324\n",
      "\tMean & Max EE are-5.978 & -4.226\n",
      "\tDegree 5 took time 0.19198179244995117\n",
      "\tMean & Max EE are-6.916 & -4.942\n"
     ]
    }
   ],
   "source": [
    "srand(52)\n",
    "shocks = randn(10200)\n",
    "\n",
    "for sol_method in [VFI(), VFI_EGM(), VFI_ECM(), dVFI_ECM(),\n",
    "                   PFI(), PFI_ECM(), EulEq()]\n",
    "    # Make sure everything is compiled\n",
    "    main(sol_method, 5, shocks; maxiter=2, verbose=false)\n",
    "\n",
    "    # Run for real\n",
    "    s_sm, t_sm, mean_eem, max_eem = main(sol_method, 5, shocks;\n",
    "                                         tol=1e-8, verbose=false)\n",
    "\n",
    "    println(\"Solution Method: $sol_method\")\n",
    "    for (d, t) in zip([2, 3, 4, 5], t_sm)\n",
    "        println(\"\\tDegree $d took time $t\")\n",
    "        println(\"\\tMean & Max EE are\" *\n",
    "                \"$(round(mean_eem[d-1], 3)) & $(round(max_eem[d-1], 3))\")\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.6.1-pre",
   "language": "julia",
   "name": "julia-0.6"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "0.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}