{
 "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 Python packages. We recommend downloading [Anaconda](https://www.continuum.io/downloads) and/or following the instructions on [quantecon.org](https://lectures.quantecon.org/py/getting_started.html). Once your Python installation is up and running, there are a few additional packages you will need in order to run the code here. To do this, you should run the following commands in your terminal\n",
    "\n",
    "```\n",
    "pip install interpolation\n",
    "pip install quantecon\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.linalg as la\n",
    "import scipy.optimize as opt\n",
    "import time\n",
    "import quantecon as qe\n",
    "\n",
    "from collections import namedtuple\n",
    "from interpolation.complete_poly import (CompletePolynomial,\n",
    "                                         n_complete, complete_polynomial,\n",
    "                                         complete_polynomial_der,\n",
    "                                         _complete_poly_impl,\n",
    "                                         _complete_poly_impl_vec,\n",
    "                                         _complete_poly_der_impl,\n",
    "                                         _complete_poly_der_impl_vec)\n",
    "from numba import jit, vectorize"
   ]
  },
  {
   "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": [
    "## Python Code\n",
    "\n",
    "We begin by defining a `namedtuple` that contains the parameters of our model. This is useful to pass the parameters around to functions that are just-in-time (JIT) compiled by `numba`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#\n",
    "# Create a named tuple type that we can pass into the jitted functions\n",
    "# so that we don't have to pass parameters one by one\n",
    "#\n",
    "Params = namedtuple(\"Params\", [\"A\", \"alpha\", \"beta\", \"delta\", \"gamma\",\n",
    "                               \"rho\", \"sigma\"])\n",
    "\n",
    "@jit(nopython=True)\n",
    "def param_unpack(params):\n",
    "    \"Unpack parameters from the Params type\"\n",
    "    out = (params.A, params.alpha, params.beta, params.delta,\n",
    "           params.gamma, params.rho, params.sigma)\n",
    "\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will then define various helper functions to ensure that we [don't repeat ourselves](https://lectures.quantecon.org/py/writing_good_code.html#don-t-repeat-yourself) and that the inner functions can be JIT compiled."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#\n",
    "# Helper functions to make sure things are jitted\n",
    "#\n",
    "@vectorize(nopython=True)\n",
    "def u(c, gamma):\n",
    "    \"CRRA utility function\"\n",
    "    return -1e10 if c < 1e-10 else (c**(1-gamma) - 1.0)/(1-gamma)\n",
    "\n",
    "@vectorize(nopython=True)\n",
    "def du(c, gamma):\n",
    "    \"Derivative of CRRA utility function\"\n",
    "    return 1e10 if c < 1e-10 else c**(-gamma)\n",
    "\n",
    "@vectorize(nopython=True)\n",
    "def duinv(u, gamma):\n",
    "    \"Inverse of the derivative of the CRRA utility function\"\n",
    "    return u**(-1.0 / gamma)\n",
    "\n",
    "@vectorize(nopython=True)\n",
    "def f(k, z, A, alpha):\n",
    "    \"C-D production function\"\n",
    "    return A * z * k**alpha\n",
    "\n",
    "@vectorize(nopython=True)\n",
    "def df(k, z, A, alpha):\n",
    "    \"Derivative of C-D production function\"\n",
    "    return alpha*A*z*k**(alpha - 1.0)\n",
    "\n",
    "@vectorize(nopython=True)\n",
    "def expendables_t(k, z, A, alpha, delta):\n",
    "    \"Budget constraint\"\n",
    "    return (1-delta)*k + f(k, z, A, alpha)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code block contains other functions that are useful for computing the policy using the envelope condition (to be discussed later), simulating, and computing Euler errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@jit(nopython=True)\n",
    "def env_cond_kp(temp, params, degree, v_coeffs, kt, zt):\n",
    "    # Unpack parameters\n",
    "    A, alpha, beta, delta, gamma, rho, sigma = param_unpack(params)\n",
    "\n",
    "    # Compute derivative of VF wrt k\n",
    "    _complete_poly_der_impl_vec(np.array([kt, zt]), degree, 0, temp)\n",
    "\n",
    "    c = duinv(np.dot(temp, v_coeffs) / (1.0-delta+df(kt, zt, A, alpha)), gamma)\n",
    "\n",
    "    return expendables_t(kt, zt, A, alpha, delta) - c\n",
    "\n",
    "\n",
    "@jit(nopython=True)\n",
    "def jit_simulate_ncgm(params, degree, v_coeffs, T, nburn, shocks):\n",
    "    \"Simulates economy using envelope condition as policy rule\"\n",
    "    # Unpack parameters\n",
    "    A, alpha, beta, delta, gamma, rho, sigma = param_unpack(params)\n",
    "\n",
    "    # Allocate space for output\n",
    "    ksim = np.empty(T+nburn)\n",
    "    zsim = np.empty(T+nburn)\n",
    "    ksim[0], zsim[0] = 1.0, 1.0\n",
    "\n",
    "    # Allocate space for temporary vector to fill with complete polynomials\n",
    "    temp = np.empty(n_complete(2, degree))\n",
    "\n",
    "    # Simulate\n",
    "    for t in range(1, T+nburn):\n",
    "        # Evaluate policy for today given yesterdays state\n",
    "        kp = env_cond_kp(temp, params, degree, v_coeffs, ksim[t-1], zsim[t-1])\n",
    "\n",
    "        # Draw new z and update k using policy from above\n",
    "        zsim[t] = zsim[t-1]**rho * np.exp(sigma*shocks[t])\n",
    "        ksim[t] = kp\n",
    "\n",
    "    return ksim[nburn:], zsim[nburn:]\n",
    "\n",
    "\n",
    "@jit(nopython=True)\n",
    "def jit_ee(params, degree, v_coeffs, nodes, weights, ks, zs):\n",
    "    # Unpack parameters\n",
    "    A, alpha, beta, delta, gamma, rho, sigma = param_unpack(params)\n",
    "\n",
    "    # Allocate space for temporary vector to fill with complete polynomials\n",
    "    temp = np.empty(n_complete(2, degree))\n",
    "    T = ks.size\n",
    "    Qn = weights.size\n",
    "\n",
    "    # Allocate space to store euler errors\n",
    "    ee = np.empty(T)\n",
    "\n",
    "    # Iterate over all ks and zs\n",
    "    for t in range(T):\n",
    "        # Current states\n",
    "        k, z = ks[t], zs[t]\n",
    "\n",
    "        # Compute decision for kp and implied c\n",
    "        k1 = env_cond_kp(temp, params, degree, v_coeffs, k, z)\n",
    "        c = expendables_t(k, z, A, alpha, delta) - k1\n",
    "\n",
    "        # Compute euler error for period t\n",
    "        lhs = du(c, gamma)\n",
    "        rhs = 0.0\n",
    "        for i in range(Qn):\n",
    "            # Get productivity tomorrow\n",
    "            z1 = z**rho * np.exp(nodes[i])\n",
    "\n",
    "            # Compute decision for kpp and implied c\n",
    "            k2 = env_cond_kp(temp, params, degree, v_coeffs, k1, z1)\n",
    "            c1 = expendables_t(k1, z1, A, alpha, delta) - k2\n",
    "            rhs = rhs + weights[i]*du(c1, gamma)*(1-delta+df(k1, z1, A, alpha))\n",
    "\n",
    "        ee[t] = np.abs(1.0 - beta*rhs/lhs)\n",
    "\n",
    "    return ee\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also now define a class that contains\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",
    "This again helps us maintain all of the relevant information in one place and simplifies passing it to other functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class NeoclassicalGrowth(object):\n",
    "    \"\"\"\n",
    "    The stochastic Neoclassical Growth model contains\n",
    "    parameters which include\n",
    "\n",
    "    * alpha: Capital share in output\n",
    "    * beta: discount factor\n",
    "    * delta: depreciation rate\n",
    "    * gamma: risk aversion\n",
    "    * rho: persistence of the log productivity level\n",
    "    * sigma: standard deviation of shocks to log productivity\n",
    "    \"\"\"\n",
    "    def __init__(self, alpha=0.36, beta=0.99, delta=0.02, gamma=2.0,\n",
    "                 rho=0.95, sigma=0.01, kmin=0.9, kmax=1.1, nk=10,\n",
    "                 zmin=0.9, zmax=1.1, nz=10, Qn=5):\n",
    "\n",
    "        # Household parameters\n",
    "        self.beta, self.gamma = beta, gamma\n",
    "\n",
    "        # Firm/technology parameters\n",
    "        self.alpha, self.delta, self.rho, self.sigma = alpha, delta, rho, sigma\n",
    "\n",
    "        # Make A such that CE steady state k is roughly 1\n",
    "        self.A = (1.0/beta - (1-delta)) / alpha\n",
    "\n",
    "        # Create t grids\n",
    "        self.kgrid = np.linspace(kmin, kmax, nk)\n",
    "        self.zgrid = np.linspace(zmin, zmax, nz)\n",
    "        self.grid = qe.gridtools.cartesian([self.kgrid, self.zgrid])\n",
    "        k, z = self.grid[:, 0], self.grid[:, 1]\n",
    "\n",
    "        # Create t+1 grids\n",
    "        self.ns, self.Qn = nz*nk, Qn\n",
    "        eps_nodes, weights = qe.quad.qnwnorm(Qn, 0.0, sigma**2)\n",
    "        self.weights = weights\n",
    "        self.z1 = z[:, None]**rho * np.exp(eps_nodes[None, :])\n",
    "\n",
    "    def _unpack_params(self):\n",
    "        out = (self.A, self.alpha, self.beta, self.delta,\n",
    "               self.gamma, self.rho, self.sigma)\n",
    "        return out\n",
    "\n",
    "    def _unpack_grids(self):\n",
    "        out = (self.kgrid, self.zgrid, self.grid, self.ztp1, self.weights)\n",
    "        return out\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 a class for each solution method that inherit various properties from a general solution parent class. The parent class will contain methods which apply to all of the other solution methods such as a general solve method, computing expectations, simulating, etc... This implementation may feel strange at first if one isn't familiar with object oriented programming, but these concepts can be powerful when properly used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class GeneralSolution:\n",
    "    \"\"\"\n",
    "    This is a general solution method. We define this, so that we can\n",
    "    sub-class it and define specific update methods for each particular\n",
    "    solution method\n",
    "    \"\"\"\n",
    "    def __init__(self, ncgm, degree, prev_sol=None):\n",
    "        # Save model and approximation degree\n",
    "        self.ncgm, self.degree = ncgm, degree\n",
    "\n",
    "        # Unpack some info from ncgm\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "        k = grid[:, 0]\n",
    "        z = grid[:, 1]\n",
    "\n",
    "        # Use parameter values from model to create a namedtuple with\n",
    "        # parameters saved inside\n",
    "        self.params = Params(A, alpha, beta, delta, gamma, rho, sigma)\n",
    "\n",
    "        self.Phi = complete_polynomial(grid.T, degree).T\n",
    "        self.dPhi = complete_polynomial_der(grid.T, degree, 0).T\n",
    "\n",
    "        # Update to fill initial value and policy matrices\n",
    "        # If we give it another solution type then use it to\n",
    "        # generate values and policies\n",
    "        if issubclass(type(prev_sol), GeneralSolution):\n",
    "            oldPhi = complete_polynomial(ncgm.grid.T, prev_sol.degree).T\n",
    "            self.VF = oldPhi @ prev_sol.v_coeffs\n",
    "            self.KP = oldPhi @ prev_sol.k_coeffs\n",
    "        # If we give it a tuple then assume it is (policy, value) pair\n",
    "        elif type(prev_sol) is tuple:\n",
    "            self.KP = prev_sol[0]\n",
    "            self.VF = prev_sol[1]\n",
    "        # Otherwise guess a constant value function and a policy\n",
    "        # of roughly steady state\n",
    "        else:\n",
    "            # VF is 0 everywhere\n",
    "            self.VF = np.zeros(ncgm.ns)\n",
    "\n",
    "            # Roughly ss policy\n",
    "            c_pol = f(k, z, A, alpha) * (A-delta)/A\n",
    "            self.KP = expendables_t(k, z, A, alpha, delta) - c_pol\n",
    "\n",
    "        # Coefficients based on guesses\n",
    "        self.v_coeffs = la.lstsq(self.Phi, self.VF)[0]\n",
    "        self.k_coeffs = la.lstsq(self.Phi, self.KP)[0]\n",
    "\n",
    "    def _unpack_params(self):\n",
    "        return self.ncgm._unpack_params()\n",
    "\n",
    "    def build_VF(self):\n",
    "        \"\"\"\n",
    "        Using the current coefficients, this builds the value function\n",
    "        for all states\n",
    "        \"\"\"\n",
    "        VF = self.Phi @ self.v_coeffs\n",
    "\n",
    "        return VF\n",
    "\n",
    "    def build_KP(self):\n",
    "        \"\"\"\n",
    "        Using the current coefficients, this builds the value function\n",
    "        for all states\n",
    "        \"\"\"\n",
    "        KP = self.Phi @ self.k_coeffs\n",
    "\n",
    "        return KP\n",
    "\n",
    "    def update_v(self, new_v_coeffs, dampen=1.0):\n",
    "        \"\"\"\n",
    "        Updates the coefficients for the value function\n",
    "        \"\"\"\n",
    "        self.v_coeffs = (1-dampen)*self.v_coeffs + dampen*new_v_coeffs\n",
    "\n",
    "        return None\n",
    "\n",
    "    def update_k(self, new_k_coeffs, dampen=1.0):\n",
    "        \"\"\"\n",
    "        Updates the coefficients for the policy function\n",
    "        \"\"\"\n",
    "        self.k_coeffs = (1-dampen)*self.k_coeffs + dampen*new_k_coeffs\n",
    "\n",
    "        return None\n",
    "\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Given the current state of everything in solution, update the\n",
    "        value and policy coefficients\n",
    "        \"\"\"\n",
    "        emsg = \"The update method is implemented in solution specific classes\"\n",
    "        emsg += \"\\nand cannot be called from `GeneralSolution`\"\n",
    "        raise ValueError(emsg)\n",
    "\n",
    "    def compute_coefficients(self, kp, VF):\n",
    "        \"\"\"\n",
    "        Given a policy and value return corresponding coefficients\n",
    "        \"\"\"\n",
    "        new_k_coeffs = la.lstsq(self.Phi, kp)[0]\n",
    "        new_v_coeffs = la.lstsq(self.Phi, VF)[0]\n",
    "\n",
    "        return new_k_coeffs, new_v_coeffs\n",
    "\n",
    "    def compute_EV_scalar(self, istate, kp):\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "\n",
    "        # All possible exogenous states tomorrow\n",
    "        z1 = self.ncgm.z1[istate, :]\n",
    "        phi = complete_polynomial(np.vstack([np.ones(self.ncgm.Qn)*kp, z1]),\n",
    "                                  self.degree).T\n",
    "        val = self.ncgm.weights@(phi@self.v_coeffs)\n",
    "\n",
    "        return val\n",
    "\n",
    "    def compute_dEV_scalar(self, istate, kp):\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "\n",
    "        # All possible exogenous states tomorrow\n",
    "        z1 = self.ncgm.z1[istate, :]\n",
    "        phi = complete_polynomial_der(np.vstack([np.ones(self.ncgm.Qn)*kp, z1]),\n",
    "                                      self.degree, 0).T\n",
    "        val = self.ncgm.weights@(phi@self.v_coeffs)\n",
    "\n",
    "        return val\n",
    "\n",
    "    def compute_EV(self, kp=None):\n",
    "        \"\"\"\n",
    "        Compute the expected value\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "        ns, Qn = self.ncgm.ns, self.ncgm.Qn\n",
    "\n",
    "        # Use policy to compute kp and c\n",
    "        if kp is None:\n",
    "            kp = self.Phi @ self.k_coeffs\n",
    "\n",
    "        # Evaluate E[V_{t+1}]\n",
    "        Vtp1 = np.empty((Qn, grid.shape[0]))\n",
    "        for iztp1 in range(Qn):\n",
    "            grid_tp1 = np.vstack([kp, self.ncgm.z1[:, iztp1]])\n",
    "            Phi_tp1 = complete_polynomial(grid_tp1, self.degree).T\n",
    "            Vtp1[iztp1, :] = Phi_tp1 @ self.v_coeffs\n",
    "\n",
    "        EV = self.ncgm.weights @ Vtp1\n",
    "\n",
    "        return EV\n",
    "\n",
    "    def compute_dEV(self, kp=None):\n",
    "        \"\"\"\n",
    "        Compute the derivative of the expected value\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "        ns, Qn = self.ncgm.ns, self.ncgm.Qn\n",
    "\n",
    "        # Use policy to compute kp and c\n",
    "        if kp is None:\n",
    "            kp = self.Phi @ self.k_coeffs\n",
    "\n",
    "        # Evaluate E[V_{t+1}]\n",
    "        dVtp1 = np.empty((Qn, grid.shape[0]))\n",
    "        for iztp1 in range(Qn):\n",
    "            grid_tp1 = np.vstack([kp, self.ncgm.z1[:, iztp1]])\n",
    "            dPhi_tp1 = complete_polynomial_der(grid_tp1, self.degree, 0).T\n",
    "            dVtp1[iztp1, :] = dPhi_tp1 @ self.v_coeffs\n",
    "\n",
    "        dEV = self.ncgm.weights @ dVtp1\n",
    "\n",
    "        return dEV\n",
    "\n",
    "    def envelope_policy(self):\n",
    "        \"\"\"\n",
    "        Applies the envelope condition to compute the policy for\n",
    "        k_{t+1} at every point on the grid\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "        k, z = grid[:, 0], grid[:, 1]\n",
    "\n",
    "        dV = self.dPhi@self.v_coeffs\n",
    "\n",
    "        # Compute the consumption\n",
    "        temp = dV / (1 - delta + df(k, z, A, alpha))\n",
    "        c = duinv(temp, gamma)\n",
    "\n",
    "        return expendables_t(k, z, A, alpha, delta) - c\n",
    "\n",
    "    def compute_distance(self, kp, VF):\n",
    "        \"\"\"\n",
    "        Computes distance between policy functions\n",
    "        \"\"\"\n",
    "        return np.max(np.abs(1.0 - kp/self.KP))\n",
    "\n",
    "    def solve(self, tol=1e-6, maxiter=2500, verbose=False, nskipprint=25):\n",
    "        # Iterate until convergence\n",
    "        dist, it = 10.0, 0\n",
    "        while (dist>tol) and (it<maxiter):\n",
    "            # Run solution specific update code\n",
    "            kp, VF = self.update()\n",
    "\n",
    "            # Compute new policy and value coeffs\n",
    "            new_k_coeffs, new_v_coeffs = self.compute_coefficients(kp, VF)\n",
    "\n",
    "            # Update distance and iterations\n",
    "            dist = self.compute_distance(kp, VF)\n",
    "            self.KP, self.VF = kp, VF\n",
    "            it += 1\n",
    "            if verbose and it%nskipprint == 0:\n",
    "                print(it, dist)\n",
    "\n",
    "            # Update all coefficients\n",
    "            self.update_k(new_k_coeffs)\n",
    "            self.update_v(new_v_coeffs)\n",
    "\n",
    "        # After finishing iteration, iterate to convergence using policy\n",
    "        if not isinstance(self, IterateOnPolicy):\n",
    "            sol_iop = IterateOnPolicy(self.ncgm, self.degree, self)\n",
    "            kp, VF = sol_iop.solve(tol=1e-10)\n",
    "\n",
    "            # Save final versions of everything\n",
    "            self.KP, self.VF = kp, VF\n",
    "            new_k_coeffs, new_v_coeffs = sol_iop.compute_coefficients(kp, VF)\n",
    "            self.update_k(new_k_coeffs)\n",
    "            self.update_v(new_v_coeffs)\n",
    "\n",
    "        return self.KP, self.VF\n",
    "\n",
    "    def simulate(self, T=10000, nburn=200, shocks=None, seed=42):\n",
    "        \"\"\"\n",
    "        Simulates the neoclassical growth model with policy function\n",
    "        given by self.KP. Simulates for `T` periods and discarsd first\n",
    "        nburn `observations`\n",
    "        \"\"\"\n",
    "        if shocks is None:\n",
    "            np.random.seed(seed)\n",
    "            shocks = np.random.randn(T+nburn)\n",
    "\n",
    "        return jit_simulate_ncgm(self.params, self.degree, self.v_coeffs,\n",
    "                                 T, nburn, shocks)\n",
    "\n",
    "    def ee_residuals(self, ksim=None, zsim=None, Qn=10, seed=42):\n",
    "        \"\"\"\n",
    "        Computes the Euler Equation residuals of a simulated capital and\n",
    "        productivity levels (ksim and zsim) and uses Qn nodes for computing\n",
    "        the expectation\n",
    "        \"\"\"\n",
    "        if (ksim is None) or (zsim is None):\n",
    "            ksim, zsim = self.simulate(T=10000, nburn=200, seed=seed)\n",
    "\n",
    "        nodes, weights = qe.quad.qnwnorm(Qn, 0.0, self.ncgm.sigma**2)\n",
    "        ee = jit_ee(self.params, self.degree, self.v_coeffs,\n",
    "                    nodes, weights, ksim, zsim)\n",
    "\n",
    "        return np.log10(np.mean(ee)), np.log10(np.max(ee))\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]$\n",
    "\n",
    "We override two of the methods from `GeneralSolution`\n",
    "\n",
    "* `compute_distance` because when we are iterating to convergence on the value function we want to check distnace using value function rather than policy function\n",
    "* `compute_coefficients` because we don't need to update the policy functions coefficients.\n",
    "\n",
    "The `update` method just repeatedly applies a particular policy function to update the value function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class IterateOnPolicy(GeneralSolution):\n",
    "    \"\"\"\n",
    "    Subclass of the general solution method. The update method for this\n",
    "    class simply computes the fixed point of the value function given\n",
    "    a specific policy\n",
    "    \"\"\"\n",
    "    def compute_distance(self, kp, VF):\n",
    "        \"\"\"\n",
    "        Computes distance between policy functions. When we are\n",
    "        iterating on a specific policy, we would like to compute\n",
    "        distances by the difference between VFs\n",
    "        \"\"\"\n",
    "        return np.max(np.abs(1.0 - VF/self.VF))\n",
    "\n",
    "    def compute_coefficients(self, kp, VF):\n",
    "        \"\"\"\n",
    "        Given a policy and value return corresponding coefficients.\n",
    "        When we are iterating on a specific policy, we don't want to\n",
    "        update the policy coefficients.\n",
    "        \"\"\"\n",
    "        new_v_coeffs = la.lstsq(self.Phi, VF)[0]\n",
    "\n",
    "        return self.k_coeffs, new_v_coeffs\n",
    "\n",
    "    def update(self):\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "\n",
    "        # Get the policy and update value function\n",
    "        c = expendables_t(grid[:, 0], grid[:, 1], A, alpha, delta) - self.KP\n",
    "\n",
    "        # Update the value function\n",
    "        VF = u(c, gamma) + beta*self.compute_EV(self.KP)\n",
    "\n",
    "        return self.KP, VF\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": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class VFI(GeneralSolution):\n",
    "    \"\"\"\n",
    "    Updates the coefficients and value functions using the VFI\n",
    "    method\n",
    "    \"\"\"\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Updates the coefficients and value functions using the VFI_ECM\n",
    "        method\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "        n_state = grid.shape[0]\n",
    "\n",
    "        # Get the policy and update it\n",
    "        kp = np.empty(n_state)\n",
    "        VF = np.empty(n_state)\n",
    "        for i_s in range(n_state):\n",
    "            # Pull out current vals\n",
    "            k, z = grid[i_s, :]\n",
    "\n",
    "            # Compute possible expendables\n",
    "            expend_t = expendables_t(k, z, A, alpha, delta)\n",
    "\n",
    "            # Negative of the objective function since we are minimizing\n",
    "            _f = lambda kp: (du(expend_t-kp, gamma) -\n",
    "                             beta*self.compute_dEV_scalar(i_s, kp))\n",
    "            _kp = opt.brentq(_f, 0.25, expend_t - 1e-2, rtol=1e-12)\n",
    "\n",
    "            kp[i_s] = _kp\n",
    "            VF[i_s] = u(expend_t-_kp, gamma)+beta*self.compute_EV_scalar(i_s, _kp)\n",
    "\n",
    "        return kp, VF\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": 9,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class VFI_EGM(GeneralSolution):\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Updates the coefficients and value functions using the VFI_ECM\n",
    "        method\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        ns = self.ncgm.ns\n",
    "        grid = self.ncgm.grid\n",
    "\n",
    "        # Use the grid as capital tomorrow instead of capital today\n",
    "        k1 = grid[:, 0]\n",
    "        z = grid[:, 1]\n",
    "        dEV = self.compute_dEV(kp=k1)\n",
    "        c = duinv(beta*dEV, gamma)\n",
    "\n",
    "        # Now find corresponding capital today such that kp is optimal\n",
    "        kt = np.empty(ns)\n",
    "        for i in range(ns):\n",
    "            ct = c[i]\n",
    "            ktp1 = k1[i]\n",
    "            zt = z[i]\n",
    "            obj = lambda k: expendables_t(k, zt, A, alpha, delta) - ct - ktp1\n",
    "            kt[i] = opt.bisect(obj, 0.25, 2.0)\n",
    "\n",
    "        # Compute value today\n",
    "        EV = self.compute_EV(kp=k1)\n",
    "        V = u(c, gamma) + beta*EV\n",
    "\n",
    "        # Get implied coefficients and evaluate kp and VF\n",
    "        phi = complete_polynomial(np.vstack([kt, z]), self.degree).T\n",
    "        temp_k_coeffs = la.lstsq(phi, k1)[0]\n",
    "        temp_v_coeffs = la.lstsq(phi, V)[0]\n",
    "        kp = self.Phi@temp_k_coeffs\n",
    "        VF = self.Phi@temp_v_coeffs\n",
    "\n",
    "        return kp, VF\n"
   ]
  },
  {
   "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",
    "We defined functions which get the policy using the envelope condition in the helper function section and inside the `GeneralSolution` class"
   ]
  },
  {
   "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": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class VFI_ECM(GeneralSolution):\n",
    "    \"\"\"\n",
    "    This class implements the Envelope Condition Method for Value\n",
    "    Function Iteration.\n",
    "    \"\"\"\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Updates the coefficients and value functions using the VFI_ECM\n",
    "        method\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "\n",
    "        # Get the policy and update it\n",
    "        kp = self.envelope_policy()\n",
    "        c = expendables_t(grid[:, 0], grid[:, 1], A, alpha, delta) - kp\n",
    "\n",
    "        # Update the value function\n",
    "        VF = u(c, gamma) + beta*self.compute_EV(kp)\n",
    "\n",
    "        return kp, VF\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.\n",
    "\n",
    "Here we override the `compute_coefficients` method to ensure that we are using the derivative of basis matrices to update the coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class dVFI_ECM(GeneralSolution):\n",
    "    def compute_coefficients(self, kp, dVF):\n",
    "        \"\"\"\n",
    "        Given a policy and derivative of the value function\n",
    "        return corresponding coefficients\n",
    "        \"\"\"\n",
    "        new_k_coeffs = la.lstsq(self.Phi, kp)[0]\n",
    "        new_v_coeffs = la.lstsq(self.dPhi, dVF)[0]\n",
    "\n",
    "        return new_k_coeffs, new_v_coeffs\n",
    "\n",
    "    def update(self):\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        grid = self.ncgm.grid\n",
    "\n",
    "        # Get the policy and update it\n",
    "        kp = self.envelope_policy()\n",
    "        c = expendables_t(grid[:, 0], grid[:, 1], A, alpha, delta) - kp\n",
    "\n",
    "        dEV = self.compute_dEV(kp=kp)\n",
    "        dV = (1-delta+df(grid[:, 0], grid[:, 1], A, alpha))*beta*dEV\n",
    "\n",
    "        return kp, dV\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": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class PFI(GeneralSolution):\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Updates the coefficients and value functions using the PFI\n",
    "        method\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        ns = self.ncgm.ns\n",
    "        grid = self.ncgm.grid\n",
    "\n",
    "        # Update policy and use new policy to iterate to convergence\n",
    "        kp = np.empty(ns)\n",
    "        for i_s in range(ns):\n",
    "            # Pull out current vals\n",
    "            k, z = grid[i_s, :]\n",
    "\n",
    "            # Compute possible expendables\n",
    "            expend_t = expendables_t(k, z, A, alpha, delta)\n",
    "\n",
    "            # Negative of the objective function since we are minimizing\n",
    "            _f = lambda kp: (du(expend_t-kp, gamma) -\n",
    "                             beta*self.compute_dEV_scalar(i_s, kp))\n",
    "            _kp = opt.brentq(_f, 0.25, expend_t - 1e-2, rtol=1e-12)\n",
    "            kp[i_s] = _kp\n",
    "\n",
    "        sol_itp = IterateOnPolicy(self.ncgm, self.degree, (kp, self.VF))\n",
    "        _, VF = sol_itp.solve(maxiter=5000)\n",
    "\n",
    "        return kp, VF\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": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class PFI_ECM(GeneralSolution):\n",
    "    \"\"\"\n",
    "    Updates the coefficients and value functions using the PFI_ECM\n",
    "    method\n",
    "    \"\"\"\n",
    "    def update(self):\n",
    "        \"\"\"\n",
    "        Updates the coefficients and value functions using the PFI_ECM\n",
    "        method\n",
    "        \"\"\"\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        ns = self.ncgm.ns\n",
    "        grid = self.ncgm.grid\n",
    "\n",
    "        # Update policy and use new policy to iterate to convergence\n",
    "        kp = self.envelope_policy()\n",
    "        sol_itp = IterateOnPolicy(self.ncgm, self.degree, (kp, self.VF))\n",
    "        _, VF = sol_itp.solve(maxiter=5000)\n",
    "\n",
    "        return kp, VF\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": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class EulEq(GeneralSolution):\n",
    "    def update(self):\n",
    "        # Unpack parameters\n",
    "        A, alpha, beta, delta, gamma, rho, sigma = self._unpack_params()\n",
    "        ns, Qn = self.ncgm.ns, self.ncgm.Qn\n",
    "        ncgm, grid = self.ncgm, self.ncgm.grid\n",
    "        k, z = grid[:, 0], grid[:, 1]\n",
    "\n",
    "        # Evaluate RHS of EE\n",
    "        temp = np.empty((n_complete(2, self.degree), ns))\n",
    "        rhs_ee = 0.0\n",
    "        for iz1 in range(Qn):\n",
    "            # Build t+1 decisions\n",
    "            zp = ncgm.z1[:, iz1]\n",
    "            _complete_poly_impl(np.vstack([self.KP, zp]), self.degree, temp)\n",
    "\n",
    "            # Compute k_{t+2} and the corresponding c_{t+1}\n",
    "            kpp = self.k_coeffs@temp\n",
    "            cp = expendables_t(self.KP, zp, A, alpha, delta) - kpp\n",
    "            foo = beta*ncgm.weights[iz1]*(1.0-delta+df(self.KP, zp, A, alpha))\n",
    "            foo *= du(cp, gamma)\n",
    "            rhs_ee += foo\n",
    "\n",
    "        # Determine c_t using euler equation\n",
    "        c = duinv(rhs_ee, gamma)\n",
    "        kp = expendables_t(k, z, A, alpha, delta) - c\n",
    "        VF = u(c, gamma) + beta*self.compute_EV(kp=kp)\n",
    "\n",
    "        return kp, VF\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulation and Euler Error Methods\n",
    "\n",
    "We have already defined functions which simulate and compute Euler errors for these solution methods."
   ]
  },
  {
   "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": 18,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/chase/Programming/anaconda3/lib/python3.6/site-packages/ipykernel/__main__.py:13: RuntimeWarning: divide by zero encountered in true_divide\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution Method: <class '__main__.VFI'>\n",
      "\n",
      "\tDegree 2 took 26.38953447341919\n",
      "\n",
      "\t\tMean and Max EE are -3.8282272985404875 & -2.762117203856586\n",
      "\n",
      "\tDegree 3 took 25.3867290019989\n",
      "\n",
      "\t\tMean and Max EE are -4.974626763230692 & -3.322176647973835\n",
      "\n",
      "\tDegree 4 took 23.616097450256348\n",
      "\n",
      "\t\tMean and Max EE are -6.060502091280513 & -4.026228854101265\n",
      "\n",
      "\tDegree 5 took 19.283337593078613\n",
      "\n",
      "\t\tMean and Max EE are -7.000247002015084 & -4.702989176950847\n",
      "\n",
      "Solution Method: <class '__main__.VFI_ECM'>\n",
      "\n",
      "\tDegree 2 took 1.1627769470214844\n",
      "\n",
      "\t\tMean and Max EE are -3.828224462040953 & -2.7620824119928944\n",
      "\n",
      "\tDegree 3 took 0.7431914806365967\n",
      "\n",
      "\t\tMean and Max EE are -4.974628189256603 & -3.3221833623376016\n",
      "\n",
      "\tDegree 4 took 0.8929553031921387\n",
      "\n",
      "\t\tMean and Max EE are -6.060501451314666 & -4.026228224602689\n",
      "\n",
      "\tDegree 5 took 0.4698905944824219\n",
      "\n",
      "\t\tMean and Max EE are -7.000246854330695 & -4.702989013989957\n",
      "\n",
      "Solution Method: <class '__main__.VFI_EGM'>\n",
      "\n",
      "\tDegree 2 took 21.106467247009277\n",
      "\n",
      "\t\tMean and Max EE are -3.8282371764027534 & -2.762131223939431\n",
      "\n",
      "\tDegree 3 took 15.00942611694336\n",
      "\n",
      "\t\tMean and Max EE are -4.974623148141401 & -3.3222165694632704\n",
      "\n",
      "\tDegree 4 took 13.397120237350464\n",
      "\n",
      "\t\tMean and Max EE are -6.060493787884856 & -4.026220770555142\n",
      "\n",
      "\tDegree 5 took 10.941481828689575\n",
      "\n",
      "\t\tMean and Max EE are -7.000269984161648 & -4.702994272397169\n",
      "\n",
      "Solution Method: <class '__main__.PFI'>\n",
      "\n",
      "\tDegree 2 took 27.039164066314697\n",
      "\n",
      "\t\tMean and Max EE are -3.828227298733333 & -2.762117204203986\n",
      "\n",
      "\tDegree 3 took 24.623878479003906\n",
      "\n",
      "\t\tMean and Max EE are -4.974626792857059 & -3.322176696671503\n",
      "\n",
      "\tDegree 4 took 31.325947999954224\n",
      "\n",
      "\t\tMean and Max EE are -6.060502089951161 & -4.026228852900712\n",
      "\n",
      "\tDegree 5 took 22.670735120773315\n",
      "\n",
      "\t\tMean and Max EE are -7.000246984163237 & -4.702989165344325\n",
      "\n",
      "Solution Method: <class '__main__.PFI_ECM'>\n",
      "\n",
      "\tDegree 2 took 1.586700201034546\n",
      "\n",
      "\t\tMean and Max EE are -3.828224462146159 & -2.7620824121954635\n",
      "\n",
      "\tDegree 3 took 1.2643656730651855\n",
      "\n",
      "\t\tMean and Max EE are -4.974628109204482 & -3.322183270022211\n",
      "\n",
      "\tDegree 4 took 1.5040409564971924\n",
      "\n",
      "\t\tMean and Max EE are -6.060501451585295 & -4.026228224939704\n",
      "\n",
      "\tDegree 5 took 0.9609415531158447\n",
      "\n",
      "\t\tMean and Max EE are -7.000246877345056 & -4.702989032703946\n",
      "\n",
      "Solution Method: <class '__main__.dVFI_ECM'>\n",
      "\n",
      "\tDegree 2 took 2.5871691703796387\n",
      "\n",
      "\t\tMean and Max EE are -3.8282295848809236 & -2.7627926640390177\n",
      "\n",
      "\tDegree 3 took 2.9016075134277344\n",
      "\n",
      "\t\tMean and Max EE are -4.974479798690568 & -3.322298622302441\n",
      "\n",
      "\tDegree 4 took 2.8114852905273438\n",
      "\n",
      "\t\tMean and Max EE are -6.060474386137703 & -4.026206968896713\n",
      "\n",
      "\tDegree 5 took 3.3714330196380615\n",
      "\n",
      "\t\tMean and Max EE are -7.000160583421296 & -4.702983153468082\n",
      "\n",
      "Solution Method: <class '__main__.EulEq'>\n",
      "\n",
      "\tDegree 2 took 1.5268235206604004\n",
      "\n",
      "\t\tMean and Max EE are -3.8233744972605015 & -2.751850845246666\n",
      "\n",
      "\tDegree 3 took 1.0108237266540527\n",
      "\n",
      "\t\tMean and Max EE are -4.973608649648261 & -3.3232579605621035\n",
      "\n",
      "\tDegree 4 took 0.8354387283325195\n",
      "\n",
      "\t\tMean and Max EE are -6.060064535170601 & -4.025869675459791\n",
      "\n",
      "\tDegree 5 took 0.44197678565979004\n",
      "\n",
      "\t\tMean and Max EE are -7.000370717683767 & -4.703035259874359\n",
      "\n"
     ]
    }
   ],
   "source": [
    "ncgm = NeoclassicalGrowth()\n",
    "\n",
    "# First guess\n",
    "vp = IterateOnPolicy(ncgm, 2)\n",
    "vp.solve(tol=1e-9)\n",
    "\n",
    "np.random.seed(61089)\n",
    "shocks = np.random.randn(10200)\n",
    "\n",
    "for sol_method in [VFI, VFI_ECM, VFI_EGM, PFI, PFI_ECM, dVFI_ECM, EulEq]:\n",
    "\n",
    "    # Set prev sol as iterate on policy\n",
    "    new_sol = vp\n",
    "    print(\"Solution Method: {}\\n\".format(sol_method))\n",
    "\n",
    "    for d in range(2, 6):\n",
    "        new_sol = sol_method(ncgm, d, new_sol)\n",
    "        ts = time.time()\n",
    "        new_sol.solve(tol=1e-9, verbose=False, nskipprint=25)\n",
    "        time_took = time.time() - ts\n",
    "        print(\"\\tDegree {} took {}\\n\".format(d, time_took))\n",
    "\n",
    "        # Compute Euler Errors\n",
    "        ks, zs = new_sol.simulate(10000, 200, shocks=shocks)\n",
    "        mean_ee, max_ee = new_sol.ee_residuals(ks, zs, Qn=10)\n",
    "        print(\"\\t\\tMean and Max EE are {} & {}\\n\".format(mean_ee, max_ee))\n",
    "        new_sol = new_sol"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}