NKModel_CLMM_python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving a New Keynesian model with Python\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 following the instructions on [quantecon.org](https://lectures.quantecon.org/jl/getting_started.html) for getting a base python installation set up. Then to acquire additional packages used in this notebook, uncomment the lines in the cell below (delete the `#` and space at the beginning of the line) and then run the cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# !pip install git+https://github.com/EconForge/interpolation.py.git\n",
    "# !pip install git+https://github.com/naught101/sobol_seq.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "outputExpanded": false
   },
   "source": [
    "## Python Code\n",
    "\n",
    "The Python version of our algorithm is implemented as a few methods defined on\n",
    "a core class named `Model`. This class is itself composed of instances of three\n",
    "different classes that hold the model parameters, steady state, and grids\n",
    "needed to describe the numerical model. Before we get to the classes, we need\n",
    "to bring in some dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true,
    "outputExpanded": false
   },
   "outputs": [],
   "source": [
    "import math\n",
    "from math import sqrt\n",
    "import time as time\n",
    "from collections import namedtuple\n",
    "\n",
    "import numpy as np\n",
    "from numpy import exp\n",
    "from interpolation.complete_poly import (_complete_poly_impl_vec,\n",
    "                                         _complete_poly_impl,\n",
    "                                         complete_polynomial)\n",
    "\n",
    "import sobol_seq\n",
    "\n",
    "# set seed on random number generator to make results reproducible\n",
    "np.random.seed(42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "outputExpanded": false
   },
   "source": [
    "We will also need the following two functions, which use monomial rules to\n",
    "compute quadrature nodes and weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true,
    "outputExpanded": false
   },
   "outputs": [],
   "source": [
    "def qnwmonomial1(vcv):\n",
    "    n = vcv.shape[0]\n",
    "    n_nodes = 2*n\n",
    "\n",
    "    z1 = np.zeros((n_nodes, n))\n",
    "\n",
    "    # In each node, random variable i takes value either 1 or -1, and\n",
    "    # all other variables take value 0. For example, for N = 2,\n",
    "    # z1 = [1 0; -1 0; 0 1; 0 -1]\n",
    "    for i in range(n):\n",
    "        z1[2*i:2*(i+1), i] = [1, -1]\n",
    "\n",
    "\n",
    "    sqrt_vcv = np.linalg.cholesky(vcv)\n",
    "    R = np.sqrt(n)*sqrt_vcv\n",
    "    ϵj = z1 @ R\n",
    "    ωj = np.ones(n_nodes) / n_nodes\n",
    "\n",
    "    return ϵj, ωj\n",
    "\n",
    "\n",
    "def qnwmonomial2(vcv):\n",
    "    n = vcv.shape[0]\n",
    "    assert n == vcv.shape[1], \"Variance covariance matrix must be square\"\n",
    "    z0 = np.zeros((1, n))\n",
    "\n",
    "    z1 = np.zeros((2*n, n))\n",
    "    # In each node, random variable i takes value either 1 or -1, and\n",
    "    # all other variables take value 0. For example, for N = 2,\n",
    "    # z1 = [1 0; -1 0; 0 1; 0 -1]\n",
    "    for i in range(n):\n",
    "        z1[2*i:2*(i+1), i] = [1, -1]\n",
    "\n",
    "    z2 = np.zeros((2*n*(n-1), n))\n",
    "    i = 0\n",
    "\n",
    "    # In each node, a pair of random variables (p,q) takes either values\n",
    "    # (1,1) or (1,-1) or (-1,1) or (-1,-1), and all other variables take\n",
    "    # value 0. For example, for N = 2, `z2 = [1 1; 1 -1; -1 1; -1 1]`\n",
    "    for p in range(n-1):\n",
    "        for q in range(p+1, n):\n",
    "            z2[4*i:4*(i+1), p] = [1, -1, 1, -1]\n",
    "            z2[4*i:4*(i+1), q] = [1, 1, -1, -1]\n",
    "            i += 1\n",
    "\n",
    "    sqrt_vcv = np.linalg.cholesky(vcv)\n",
    "    R = np.sqrt(n+2)*sqrt_vcv\n",
    "    S = np.sqrt((n+2)/2)*sqrt_vcv\n",
    "    ϵj = np.row_stack([z0, z1 @ R, z2 @ S])\n",
    "\n",
    "    ωj = np.concatenate([2/(n+2) * np.ones(z0.shape[0]),\n",
    "                         (4-n)/(2*(n+2)**2) * np.ones(z1.shape[0]),\n",
    "                         1/(n+2)**2 * np.ones(z2.shape[0])])\n",
    "    return ϵj, ωj"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "outputExpanded": false
   },
   "source": [
    "## Classes\n",
    "\n",
    "First we have the `Params` class, which holds all the model parameters as well\n",
    "as the paramters that drive the algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true,
    "outputExpanded": false
   },
   "outputs": [],
   "source": [
    "SteadyState = namedtuple(\"SteadyState\",\n",
    "                         [\"Yn\", \"Y\", \"π\", \"δ\", \"L\", \"C\", \"F\", \"S\", \"R\", \"w\"])\n",
    "\n",
    "class Params(object):\n",
    "    def __init__(self, zlb=True, γ=1, β=0.99, ϑ=2.09, ϵ=4.45, ϕ_y=0.07,\n",
    "                 ϕ_π=2.21, μ=0.82, Θ=0.83, πstar=1, gbar=0.23,\n",
    "                 ρηR=0.0, ρηa=0.95, ρηL=0.25, ρηu=0.92, ρηB=0.0, ρηG=0.95,\n",
    "                 σηR=0.0028, σηa=0.0045, σηL=0.0500, σηu=0.0054, σηB=0.0010,\n",
    "                 σηG=0.0038, degree=2):\n",
    "\n",
    "        self.zlb = zlb      # whether or not the zlb should be imposed\n",
    "        self.γ = γ          # Utility-function parameter\n",
    "        self.β = β          # Discount factor\n",
    "        self.ϑ = ϑ          # Utility-function parameter\n",
    "        self.ϵ = ϵ          # Parameter in the Dixit-Stiglitz aggregator\n",
    "        self.ϕ_y = ϕ_y      # Parameter of the Taylor rule\n",
    "        self.ϕ_π = ϕ_π      # Parameter of the Taylor rule\n",
    "        self.μ = μ          # Parameter of the Taylor rule\n",
    "        self.Θ = Θ          # Share of non-reoptimizing firms (Calvo's pricing)\n",
    "        self.πstar = πstar  # Target (gross) inflation rate\n",
    "        self.gbar = gbar    # Steady-state share of government spending in output\n",
    "\n",
    "        # autocorrelation coefficients\n",
    "        self.ρηR = ρηR  # See process (28) in MM (2015)\n",
    "        self.ρηa = ρηa  # See process (22) in MM (2015)\n",
    "        self.ρηL = ρηL  # See process (16) in MM (2015)\n",
    "        self.ρηu = ρηu  # See process (15) in MM (2015)\n",
    "        self.ρηB = ρηB  # See process (17) in MM (2015)\n",
    "        self.ρηG = ρηG  # See process (26) in MM (2015)\n",
    "\n",
    "        # standard deviations\n",
    "        self.σηR = σηR  # See process (28) in MM (2015)\n",
    "        self.σηa = σηa  # See process (22) in MM (2015)\n",
    "        self.σηL = σηL  # See process (16) in MM (2015)\n",
    "        self.σηu = σηu  # See process (15) in MM (2015)\n",
    "        self.σηB = σηB  # See process (17) in MM (2015)\n",
    "        self.σηG = σηG  # See process (26) in MM (2015)\n",
    "\n",
    "        self.degree = degree\n",
    "\n",
    "    @property\n",
    "    def vcov(self):\n",
    "        return np.diag([self.σηR**2, self.σηa**2, self.σηL**2,\n",
    "                        self.σηu**2, self.σηB**2, self.σηG**2])\n",
    "\n",
    "    @property\n",
    "    def steady_state(self):\n",
    "        Yn_ss = exp(self.gbar)**(self.γ/(self.ϑ+self.γ))\n",
    "        Y_ss = Yn_ss\n",
    "        π_ss = 1.0\n",
    "        δ_ss = 1.0\n",
    "        L_ss = Y_ss/δ_ss\n",
    "        C_ss = (1-self.gbar)*Y_ss\n",
    "        F_ss = C_ss**(-self.γ)*Y_ss/(1-self.β*self.Θ*π_ss**(self.ϵ-1))\n",
    "        S_ss = L_ss**self.ϑ*Y_ss/(1-self.β*self.Θ*π_ss**self.ϵ)\n",
    "        R_ss = π_ss/self.β\n",
    "        w_ss = (L_ss**self.ϑ)*(C_ss**self.γ)\n",
    "\n",
    "        return SteadyState(Yn_ss, Y_ss, π_ss, δ_ss, L_ss, C_ss, F_ss, S_ss, R_ss, w_ss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "outputExpanded": false
   },
   "source": [
    "Notice that we have a namedtuple to hold the steady state of the model. Using\n",
    "the namedtuple infrastructure allows us to have convenient \"dot-style\" access\n",
    "to the steady state, without defining a full class.\n",
    "\n",
    "Given an instance of `Params` class, we can construct the grid on which we will\n",
    "solve the model.\n",
    "\n",
    "The `Grids` class holds this grid as well as matrices used to compute\n",
    "expectations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true,
    "outputExpanded": false
   },
   "outputs": [],
   "source": [
    "class Grids(object):\n",
    "\n",
    "    def __init__(self, p, m=200, kind=\"rands\"):\n",
    "\n",
    "        if kind == \"sobol\":\n",
    "            ub = np.array([\n",
    "                2 * p.σηR / sqrt(1 - p.ρηR**2),\n",
    "                2 * p.σηa / sqrt(1 - p.ρηa**2),\n",
    "                2 * p.σηL / sqrt(1 - p.ρηL**2),\n",
    "                2 * p.σηu / sqrt(1 - p.ρηu**2),\n",
    "                2 * p.σηB / sqrt(1 - p.ρηB**2),\n",
    "                2 * p.σηG / sqrt(1 - p.ρηG**2),\n",
    "                1.05,  # R\n",
    "                1.0    # δ\n",
    "             ])\n",
    "            lb = -ub\n",
    "            lb[[6, 7]] = [1.0, 0.95]  # adjust lower bound for R and δ\n",
    "            s = sobol_seq.i4_sobol_generate(8, m)\n",
    "            s *= (ub - lb)\n",
    "            s += lb\n",
    "\n",
    "            ηR = s[:, 0]\n",
    "            ηa = s[:, 1]\n",
    "            ηL = s[:, 2]\n",
    "            ηu = s[:, 3]\n",
    "            ηB = s[:, 4]\n",
    "            ηG = s[:, 5]\n",
    "            R = s[:, 6]\n",
    "            δ = s[:, 7]\n",
    "        else:\n",
    "            # Values of exogenous state variables are distributed uniformly\n",
    "            # in the interval +/- std/sqrt(1-rho_nu**2)\n",
    "            ηR = (-2*p.σηR + 4*p.σηR*np.random.rand(m)) / sqrt(1-p.ρηR**2)\n",
    "            ηa = (-2*p.σηa + 4*p.σηa*np.random.rand(m)) / sqrt(1-p.ρηa**2)\n",
    "            ηL = (-2*p.σηL + 4*p.σηL*np.random.rand(m)) / sqrt(1-p.ρηL**2)\n",
    "            ηu = (-2*p.σηu + 4*p.σηu*np.random.rand(m)) / sqrt(1-p.ρηu**2)\n",
    "            ηB = (-2*p.σηB + 4*p.σηB*np.random.rand(m)) / sqrt(1-p.ρηB**2)\n",
    "            ηG = (-2*p.σηG + 4*p.σηG*np.random.rand(m)) / sqrt(1-p.ρηG**2)\n",
    "\n",
    "            # Values of endogenous state variables are distributed uniformly\n",
    "            # in the intervals [1 1.05] and [0.95 1], respectively\n",
    "            R = 1 + 0.05*np.random.rand(m)\n",
    "            δ = 0.95 + 0.05*np.random.rand(m)\n",
    "\n",
    "        self.ηR = ηR\n",
    "        self.ηa = ηa\n",
    "        self.ηL = ηL\n",
    "        self.ηu = ηu\n",
    "        self.ηB = ηB\n",
    "        self.ηG = ηG\n",
    "        self.R = R\n",
    "        self.δ = δ\n",
    "\n",
    "        # shape (8, m)\n",
    "        self.X = np.vstack([np.log(R), np.log(δ), ηR, ηa, ηL, ηu, ηB, ηG])\n",
    "\n",
    "        # shape (n_complete(8, 2), m)\n",
    "        self.X0_G = {\n",
    "            1: complete_polynomial(self.X, 1),\n",
    "            p.degree: complete_polynomial(self.X, p.degree)\n",
    "        }\n",
    "\n",
    "        # shape (2*n=12, n=6)\n",
    "        self.ϵ_nodes, self.ω_nodes = qnwmonomial1(p.vcov)\n",
    "\n",
    "        # all shape (len(ϵ_nodes), m)\n",
    "        self.ηR1 = p.ρηR * ηR[None, :] + self.ϵ_nodes[:, None, 0]\n",
    "        self.ηa1 = p.ρηa * ηa[None, :] + self.ϵ_nodes[:, None, 1]\n",
    "        self.ηL1 = p.ρηL * ηL[None, :] + self.ϵ_nodes[:, None, 2]\n",
    "        self.ηu1 = p.ρηu * ηu[None, :] + self.ϵ_nodes[:, None, 3]\n",
    "        self.ηB1 = p.ρηB * ηB[None, :] + self.ϵ_nodes[:, None, 4]\n",
    "        self.ηG1 = p.ρηG * ηG[None, :] + self.ϵ_nodes[:, None, 5]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "outputExpanded": false
   },
   "source": [
    "Finally, we construct the Model class, which has an instance of Params,\n",
    "SteadyState and Grids as its three attributes.\n",
    "\n",
    "This block of code will be longer than the others because we also include\n",
    "routines to solve and simulate the model as methods on the Model class. These\n",
    "methods will be clearly marked and commented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true,
    "outputExpanded": false
   },
   "outputs": [],
   "source": [
    "class Model(object):\n",
    "    def __init__(self, p=Params(), g=None):\n",
    "        if g is None:\n",
    "            g = Grids(p)\n",
    "\n",
    "        self.p = p\n",
    "        self.g = g\n",
    "        self.s = self.p.steady_state\n",
    "\n",
    "    def init_coefs(self):\n",
    "        \"Iniital guess for coefs. We evaluate interpoland as coefs @ basis_mat\"\n",
    "        npol = self.g.X0_G[1].shape[0]\n",
    "        coefs = np.full((3, npol), 1e-5)\n",
    "        coefs[:, 0] = [self.s.S, self.s.F, self.s.C**(-self.p.γ)]\n",
    "        return coefs\n",
    "\n",
    "    def step(self, S, F, C, δ0, R0, ηG, ηa, ηL, ηR):\n",
    "        # simplify notation\n",
    "        Θ, ϵ, gbar, ϑ, γ = self.p.Θ, self.p.ϵ, self.p.gbar, self.p.ϑ, self.p.γ\n",
    "        β, μ, ϕ_π, ϕ_y, πs = self.p.β, self.p.μ, self.p.ϕ_π, self.p.ϕ_y, self.s.π\n",
    "\n",
    "        # Compute pie(t) from condition (35) in MM (2015)\n",
    "        π0 = ((1-(1-Θ)*(S/F)**(1-ϵ))/Θ)**(1/(ϵ-1))\n",
    "\n",
    "        # Compute delta(t) from condition (36) in MM (2015)\n",
    "        δ1 = ((1-Θ)*((1-Θ*π0**(ϵ-1))/(1-Θ))**(ϵ/(ϵ-1))+Θ*π0**ϵ/δ0)**(-1)\n",
    "\n",
    "        # Compute Y(t) from condition (38) in MM (2015)\n",
    "        Y0 = C/(1-gbar/exp(ηG))\n",
    "\n",
    "        # Compute L(t) from condition (37) in MM (2015)\n",
    "        L0 = Y0/exp(ηa)/δ1\n",
    "\n",
    "        #  Compute Yn(t) from condition (31) in MM (2015)\n",
    "        Yn0 = (exp(ηa)**(1+ϑ)*(1-gbar/exp(ηG))**(-γ)/exp(ηL))**(1/(ϑ+γ))\n",
    "\n",
    "        # Compute R(t) from conditions (27), (39) in MM (2015) -- Taylor rule\n",
    "        R1 = πs/β*(R0*β/πs)**μ*((π0/πs)**ϕ_π * (Y0/Yn0)**ϕ_y)**(1-μ)*exp(ηR)\n",
    "\n",
    "        return π0, δ1, Y0, L0, Yn0, R1\n",
    "\n",
    "    def solve(self, damp=0.1, tol=1e-7):\n",
    "        # rename self to m to make code below readable\n",
    "        m = self\n",
    "\n",
    "        n = len(m.g.ηR)\n",
    "        n_nodes = len(m.g.ω_nodes)\n",
    "\n",
    "        ## allocate memory\n",
    "        # euler equations\n",
    "        e = np.zeros((3, n))\n",
    "\n",
    "        # previous iteration S, F, C\n",
    "        S0_old_G = np.ones(n)\n",
    "        F0_old_G = np.ones(n)\n",
    "        C0_old_G = np.ones(n)\n",
    "\n",
    "        # current iteration S, F, C\n",
    "        S0_new_G = np.ones(n)\n",
    "        F0_new_G = np.ones(n)\n",
    "        C0_new_G = np.ones(n)\n",
    "\n",
    "        # future S, F, C\n",
    "        S1 = np.zeros((n_nodes, n))\n",
    "        F1 = np.zeros((n_nodes, n))\n",
    "        C1 = np.zeros((n_nodes, n))\n",
    "\n",
    "        for deg in [1, self.p.degree]:\n",
    "            # housekeeping\n",
    "            err = 1.0\n",
    "            X0_G = m.g.X0_G[deg]\n",
    "\n",
    "            if deg > 1:\n",
    "                # compute degree coefs using degree 1 coefs as guess\n",
    "                coefs = np.linalg.lstsq(X0_G.T, e.T)[0].T\n",
    "            else:\n",
    "                coefs = self.init_coefs()\n",
    "\n",
    "            while err > tol:\n",
    "                # Current choices (at t)\n",
    "                # ------------------------------\n",
    "                SFC0 = coefs @ X0_G\n",
    "                S0 = SFC0[0, :]                # Compute S(t) using coefs\n",
    "                F0 = SFC0[1, :]                # Compute F(t) using coefs\n",
    "                C0 = (SFC0[2, :])**(-1/m.p.γ)  # Compute C(t) using coefs\n",
    "                π0, δ1, Y0, L0, Yn0, R1 = self.step(S0, F0, C0, m.g.δ, m.g.R,\n",
    "                                                    m.g.ηG, m.g.ηa, m.g.ηL, m.g.ηR)\n",
    "\n",
    "                if self.p.zlb:\n",
    "                    R1 = np.maximum(R1, 1.0)\n",
    "\n",
    "                for u in range(n_nodes):\n",
    "                    # Form complete polynomial of degree \"Degree\" (at t+1) on future state\n",
    "                    grid1 = [np.log(R1), np.log(δ1), m.g.ηR1[u, :], m.g.ηa1[u, :],\n",
    "                             m.g.ηL1[u, :], m.g.ηu1[u, :], m.g.ηB1[u, :], m.g.ηG1[u, :]]\n",
    "\n",
    "                    X1 = complete_polynomial(grid1, deg)\n",
    "\n",
    "                    S1[u, :] = coefs[0, :] @ X1                # Compute S(t+1)\n",
    "                    F1[u, :] = coefs[1, :] @ X1                # Compute F(t+1)\n",
    "                    C1[u, :] = (coefs[2, :] @ X1)**(-1/m.p.γ)  # Compute C(t+1)\n",
    "\n",
    "                # Compute next-period π using condition\n",
    "                # (35) in MM (2015)\n",
    "                π1 = ((1-(1-m.p.Θ)*(S1/F1)**(1-m.p.ϵ))/m.p.Θ)**(1/(m.p.ϵ-1))\n",
    "\n",
    "                # Evaluate conditional expectations in the Euler equations\n",
    "                #---------------------------------------------------------\n",
    "                e[0, :] = exp(m.g.ηu)*exp(m.g.ηL)*L0**m.p.ϑ*Y0/exp(m.g.ηa) + m.g.ω_nodes @ (m.p.β*m.p.Θ*π1**m.p.ϵ*S1)\n",
    "                e[1, :] = exp(m.g.ηu)*C0**(-m.p.γ)*Y0 + m.g.ω_nodes @ (m.p.β*m.p.Θ*π1**(m.p.ϵ-1)*F1)\n",
    "                e[2, :] = m.p.β*exp(m.g.ηB)/exp(m.g.ηu)*R1 * (m.g.ω_nodes @ ((exp(m.g.ηu1)*C1**(-m.p.γ)/π1)))\n",
    "\n",
    "                # Variables of the current iteration\n",
    "                #-----------------------------------\n",
    "                np.copyto(S0_new_G, S0)\n",
    "                np.copyto(F0_new_G, F0)\n",
    "                np.copyto(C0_new_G, C0)\n",
    "\n",
    "                # Compute and update the coefficients of the decision functions\n",
    "                # -------------------------------------------------------------\n",
    "                coefs_hat = np.linalg.lstsq(X0_G.T, e.T)[0].T\n",
    "\n",
    "                # Update the coefficients using damping\n",
    "                coefs = damp*coefs_hat + (1-damp)*coefs\n",
    "\n",
    "                # Evaluate the percentage (unit-free) difference between the values\n",
    "                # on the grid from the previous and current iterations\n",
    "                # -----------------------------------------------------------------\n",
    "                # The convergence criterion is adjusted to the damping parameters\n",
    "                err = (np.mean(abs(1-S0_new_G/S0_old_G)) +\n",
    "                       np.mean(abs(1-F0_new_G/F0_old_G)) +\n",
    "                       np.mean(abs(1-C0_new_G/C0_old_G)))\n",
    "\n",
    "                # Store the obtained values for S(t), F(t), C(t) on the grid to\n",
    "                # be used on the subsequent iteration in Section 10.2.6\n",
    "                #-----------------------------------------------------------------------\n",
    "                np.copyto(S0_old_G, S0_new_G)\n",
    "                np.copyto(F0_old_G, F0_new_G)\n",
    "                np.copyto(C0_old_G, C0_new_G)\n",
    "\n",
    "        return coefs\n",
    "\n",
    "    def simulate(self, coefs=None, capT=10201):\n",
    "        if coefs is None:\n",
    "            coefs = self.solve()\n",
    "\n",
    "        # rename self to m to make code below readable\n",
    "        m = self\n",
    "\n",
    "        # create namedtuple to hold simulation results in an organized container\n",
    "        Simulation = namedtuple(\"Simulation\",\n",
    "                                [\"nuR\", \"nua\", \"nuL\", \"nuu\", \"nuB\", \"nuG\",\n",
    "                                 \"δ\", \"R\", \"S\", \"F\", \"C\", \"π\", \"Y\", \"L\", \"Yn\", \"w\"])\n",
    "\n",
    "        # 11. Simualating a time-series solution\n",
    "        #---------------------------------------\n",
    "\n",
    "        # Initialize the values of 6 exogenous shocks and draw innovations\n",
    "        #-----------------------------------------------------------------\n",
    "        nuR = np.zeros(capT)\n",
    "        nua = np.zeros(capT)\n",
    "        nuL = np.zeros(capT)\n",
    "        nuu = np.zeros(capT)\n",
    "        nuB = np.zeros(capT)\n",
    "        nuG = np.zeros(capT)\n",
    "\n",
    "        # Generate the series for shocks\n",
    "        #-------------------------------\n",
    "        rands = np.random.randn(capT-1, 6)\n",
    "\n",
    "        for t in range(capT-1):\n",
    "            nuR[t+1] = self.p.ρηR*nuR[t] + self.p.σηR*rands[t, 0]\n",
    "            nua[t+1] = self.p.ρηa*nua[t] + self.p.σηa*rands[t, 1]\n",
    "            nuL[t+1] = self.p.ρηL*nuL[t] + self.p.σηL*rands[t, 2]\n",
    "            nuu[t+1] = self.p.ρηu*nuu[t] + self.p.σηu*rands[t, 3]\n",
    "            nuB[t+1] = self.p.ρηB*nuB[t] + self.p.σηB*rands[t, 4]\n",
    "            nuG[t+1] = self.p.ρηG*nuG[t] + self.p.σηG*rands[t, 5]\n",
    "\n",
    "        δ = np.ones(capT+1) # Allocate memory for the time series of delta(t)\n",
    "        R = np.ones(capT+1) # Allocate memory for the time series of R(t)\n",
    "        S = np.ones(capT)   # Allocate memory for the time series of S(t)\n",
    "        F = np.ones(capT)   # Allocate memory for the time series of F(t)\n",
    "        C = np.ones(capT)   # Allocate memory for the time series of C(t)\n",
    "        π = np.ones(capT)   # Allocate memory for the time series of π(t)\n",
    "        Y = np.ones(capT)   # Allocate memory for the time series of Y(t)\n",
    "        L = np.ones(capT)   # Allocate memory for the time series of L(t)\n",
    "        Yn = np.ones(capT)   # Allocate memory for the time series of Yn(t)\n",
    "        w = np.ones(capT)\n",
    "\n",
    "        pol_bases = np.empty(coefs.shape[1])\n",
    "        states = np.empty(8)\n",
    "\n",
    "        for t in range(capT):\n",
    "            states[0] = math.log(R[t]); states[1] = math.log(δ[t])\n",
    "            states[2] = nuR[t]; states[3] = nua[t]\n",
    "            states[4] = nuL[t]; states[5] = nuu[t]\n",
    "            states[6] = nuB[t]; states[7] = nuG[t]\n",
    "\n",
    "            _complete_poly_impl_vec(states, self.p.degree, pol_bases)\n",
    "\n",
    "            vals = coefs @ pol_bases\n",
    "            S[t] = vals[0]\n",
    "            F[t] = vals[1]\n",
    "            C[t] = (vals[2])**(-1/m.p.γ)\n",
    "\n",
    "            π[t], δ[t+1], Y[t], L[t], Yn[t], R[t+1] = self.step(S[t], F[t], C[t],\n",
    "                                                                δ[t], R[t], nuG[t],\n",
    "                                                                nua[t], nuL[t],\n",
    "                                                                nuR[t])\n",
    "\n",
    "            # Compute real wage\n",
    "            w[t] = exp(nuL[t])*(L[t]**m.p.ϑ)*(C[t]**m.p.γ)\n",
    "\n",
    "            # If ZLB is imposed, set R(t)=1 if ZLB binds\n",
    "            if self.p.zlb:\n",
    "                R[t+1] = max(R[t+1], 1.0)\n",
    "\n",
    "        return Simulation(nuR, nua, nuL, nuu, nuB, nuG, δ, R, S, F, C, π, Y, L, Yn, w)\n",
    "\n",
    "    def residuals(self, coefs, sim, burn=200):\n",
    "        m = self  # rename self to m so the rest of this code is more readable\n",
    "        capT = len(sim.w)\n",
    "        resids = np.zeros((capT, 9))\n",
    "\n",
    "        # Integration method for evaluating accuracy\n",
    "        # ------------------------------------------\n",
    "        # Monomial integration rule with 2N**2+1 nodes\n",
    "        ϵ_nodes, ω_nodes = qnwmonomial2(m.p.vcov)\n",
    "        n_nodes = len(ω_nodes)\n",
    "\n",
    "        # Allocate for arrays needed in the loop\n",
    "        basis_mat = np.empty((8, n_nodes))\n",
    "        X1 = np.empty((coefs.shape[1], n_nodes))\n",
    "\n",
    "        nuR1 = np.empty(n_nodes)\n",
    "        nua1 = np.empty(n_nodes)\n",
    "        nuL1 = np.empty(n_nodes)\n",
    "        nuu1 = np.empty(n_nodes)\n",
    "        nuB1 = np.empty(n_nodes)\n",
    "        nuG1 = np.empty(n_nodes)\n",
    "\n",
    "        for t in range(capT):                 # For each given point,\n",
    "            # Take the corresponding value for shocks at t\n",
    "            #---------------------------------------------\n",
    "            nuR0 = sim.nuR[t]  # nuR(t)\n",
    "            nua0 = sim.nua[t]  # nua(t)\n",
    "            nuL0 = sim.nuL[t]  # nuL(t)\n",
    "            nuu0 = sim.nuu[t]  # nuu(t)\n",
    "            nuB0 = sim.nuB[t]  # nuB(t)\n",
    "            nuG0 = sim.nuG[t]  # nuG(t)\n",
    "\n",
    "            # Exctract time t values for all other variables (and t+1 for R, δ)\n",
    "            #------------------------------------------------------------------\n",
    "            R0 = sim.R[t]    # R(t-1)\n",
    "            δ0 = sim.δ[t]    # δ(t-1)\n",
    "            R1 = sim.R[t+1]  # R(t)\n",
    "            δ1 = sim.δ[t+1]  # δ(t)\n",
    "\n",
    "            L0 = sim.L[t]    # L(t)\n",
    "            Y0 = sim.Y[t]    # Y(t)\n",
    "            Yn0 = sim.Yn[t]  # Yn(t)\n",
    "            π0 = sim.π[t]    # π(t)\n",
    "            S0 = sim.S[t]    # S(t)\n",
    "            F0 = sim.F[t]    # F(t)\n",
    "            C0 = sim.C[t]    # C(t)\n",
    "\n",
    "            # Fill basis matrix with R1, δ1 and shocks\n",
    "            #-----------------------------------------\n",
    "            # Note that we do not premultiply by standard deviations as ϵ_nodes\n",
    "            # already include them. All these variables are vectors of length n_nodes\n",
    "            nuR1[:] = nuR0*m.p.ρηR + ϵ_nodes[:, 0]\n",
    "            nua1[:] = nua0*m.p.ρηa + ϵ_nodes[:, 1]\n",
    "            nuL1[:] = nuL0*m.p.ρηL + ϵ_nodes[:, 2]\n",
    "            nuu1[:] = nuu0*m.p.ρηu + ϵ_nodes[:, 3]\n",
    "            nuB1[:] = nuB0*m.p.ρηB + ϵ_nodes[:, 4]\n",
    "            nuG1[:] = nuG0*m.p.ρηG + ϵ_nodes[:, 5]\n",
    "\n",
    "            basis_mat[0, :] = np.log(R1)\n",
    "            basis_mat[1, :] = np.log(δ1)\n",
    "            basis_mat[2, :] = nuR1\n",
    "            basis_mat[3, :] = nua1\n",
    "            basis_mat[4, :] = nuL1\n",
    "            basis_mat[5, :] = nuu1\n",
    "            basis_mat[6, :] = nuB1\n",
    "            basis_mat[7, :] = nuG1\n",
    "\n",
    "            # Future choices at t+1\n",
    "            #----------------------\n",
    "            # Form a complete polynomial of degree \"Degree\" (at t+1) on future state\n",
    "            # variables; n_nodes-by-npol\n",
    "            _complete_poly_impl(basis_mat, self.p.degree, X1)\n",
    "\n",
    "            # Compute S(t+1), F(t+1) and C(t+1) in all nodes using coefs\n",
    "            S1 = coefs[0, :] @ X1\n",
    "            F1 = coefs[1, :] @ X1\n",
    "            C1 = (coefs[2, :] @ X1)**(-1/m.p.γ)\n",
    "\n",
    "            # Compute π(t+1) using condition (35) in MM (2015)\n",
    "            π1 = ((1-(1-m.p.Θ)*(S1/F1)**(1-m.p.ϵ))/m.p.Θ)**(1/(m.p.ϵ-1))\n",
    "\n",
    "            # Compute residuals for each of the 9 equilibrium conditions\n",
    "            #-----------------------------------------------------------\n",
    "            resids[t, 0] = 1-(ω_nodes @\n",
    "                (exp(nuu0)*exp(nuL0)*L0**m.p.ϑ*Y0/exp(nua0) +\n",
    "                 m.p.β*m.p.Θ*π1**m.p.ϵ*S1)/S0\n",
    "            )\n",
    "            resids[t, 1] = 1 - (ω_nodes @\n",
    "                (exp(nuu0)*C0**(-m.p.γ)*Y0 + m.p.β*m.p.Θ*π1**(m.p.ϵ-1)*F1)/F0\n",
    "            )\n",
    "            resids[t, 2] = 1.0 -(ω_nodes @\n",
    "                (m.p.β*exp(nuB0)/exp(nuu0)*R1*exp(nuu1)*C1**(-m.p.γ)/π1)/C0**(-m.p.γ)\n",
    "            )\n",
    "\n",
    "            resids[t, 3] = 1-((1-m.p.Θ*π0**(m.p.ϵ-1))/(1-m.p.Θ))**(1/(1-m.p.ϵ))*F0/S0\n",
    "            resids[t, 4] = 1-((1-m.p.Θ)*((1-m.p.Θ*π0**(m.p.ϵ-1))/(1-m.p.Θ))**(m.p.ϵ/(m.p.ϵ-1)) + m.p.Θ*π0**m.p.ϵ/δ0)**(-1)/δ1\n",
    "            resids[t, 5] = 1-exp(nua0)*L0*δ1/Y0\n",
    "            resids[t, 6] = 1-(1-m.p.gbar/exp(nuG0))*Y0/C0\n",
    "            resids[t, 7] = 1-(exp(nua0)**(1+m.p.ϑ)*(1-m.p.gbar/exp(nuG0))**(-m.p.γ)/exp(nuL0))**(1/(m.p.ϑ+m.p.γ))/Yn0\n",
    "            resids[t, 8] = 1-m.s.π/m.p.β*(R0*m.p.β/m.s.π)**m.p.μ*((π0/m.s.π)**m.p.ϕ_π * (Y0/Yn0)**m.p.ϕ_y)**(1-m.p.μ)*exp(nuR0)/R1   # Taylor rule\n",
    "\n",
    "            # If the ZLB is imposed and R>1, the residuals in the Taylor rule (the\n",
    "            # 9th equation) are zero\n",
    "            if m.p.zlb and R1 <= 1:\n",
    "                resids[t, 8] = 0.0\n",
    "\n",
    "        return resids[burn:, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "outputExpanded": false
   },
   "source": [
    "## Running the code\n",
    "\n",
    "Now that we've done all the hard work to define the model, its solution and\n",
    "simulation, and accuracy checks, let's put things together and run the code!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true,
    "outputExpanded": false
   },
   "outputs": [],
   "source": [
    "def main(m=Model(), file=None):\n",
    "    if file is None:\n",
    "        mprint = print\n",
    "    else:\n",
    "        def mprint(*x):\n",
    "            print(*x, file=file)\n",
    "\n",
    "    # solve the model\n",
    "    t1 = time.time()\n",
    "    coefs = m.solve()\n",
    "    solve_time = time.time() - t1\n",
    "\n",
    "    # simulate the model\n",
    "    t1 = time.time()\n",
    "    sim = m.simulate(coefs)\n",
    "    sim_time = time.time() - t1\n",
    "\n",
    "    # check accuracy\n",
    "    t1 = time.time()\n",
    "    resids = m.residuals(coefs, sim)\n",
    "    resids_time = time.time() - t1\n",
    "\n",
    "    tot_time = solve_time + sim_time + resids_time\n",
    "    mean_err = np.log10(abs(resids).mean())\n",
    "    max_err = np.log10(abs(resids).max())\n",
    "    max_err_eqn = np.log10(abs(resids).max(1) + 1e-16)\n",
    "    l1 = np.log10(abs(resids).max(0).sum())\n",
    "\n",
    "    mprint(\"Solver time (in seconds): \", solve_time)\n",
    "    mprint(\"Simulation time (in seconds): \", sim_time)\n",
    "    mprint(\"Residuals time (in seconds): \", resids_time)\n",
    "    mprint(\"Total time (in seconds): \", tot_time)\n",
    "    mprint(\"\\nAPPROXIMATION ERRORS (log10):\")\n",
    "    mprint(\"\\ta) mean error in the model equations: {:0.3f}\".format(mean_err))\n",
    "    mprint(\"\\tb) sum of max error per equation: {:0.3f}\".format(l1));\n",
    "    mprint(\"\\tc) max error in the model equations: {:0.3f}\".format(max_err))\n",
    "    mprint(\"\\td) max error by equation: \", max_err_eqn)\n",
    "    mprint(\"tex row:\", \"{:.2f} & {:.2f} & {:.2f}\".format(l1, max_err, tot_time))\n",
    "\n",
    "    return solve_time, sim_time, resids_time, coefs, sim, resids\n",
    "\n",
    "\n",
    "def build_paper_table():\n",
    "    with open(\"output.log\", \"w\") as f:\n",
    "        for params in (dict(πstar=1.0, σηL=0.1821, zlb=False),\n",
    "                       dict(πstar=1.0, σηL=0.4054, zlb=False),\n",
    "                       dict(πstar=1.0, σηL=0.1821, zlb=True)):\n",
    "\n",
    "\n",
    "            for grid_kind in [\"sobol\", \"random\"]:\n",
    "                p = Params(**params)\n",
    "                g = Grids(p, kind=grid_kind)\n",
    "                m = Model(p, g)\n",
    "\n",
    "                print(\"working with params:\", params, file=f)\n",
    "                print(\"And grid type:\", grid_kind, file=f)\n",
    "                main(m, f)\n",
    "                print(\"\\n\"*5, file=f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "build_paper_table()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solver time (in seconds):  1.620750904083252\n",
      "Simulation time (in seconds):  0.671532154083252\n",
      "Residuals time (in seconds):  0.826570987701416\n",
      "Total time (in seconds):  3.11885404586792\n",
      "\n",
      "APPROXIMATION ERRORS (log10):\n",
      "\ta) mean error in the model equations: -4.342\n",
      "\tb) sum of max error per equation: -1.375\n",
      "\tc) max error in the model equations: -1.704\n",
      "\td) max error by equation:  [-3.73802289 -3.31801009 -3.40755216 ..., -4.56959248 -3.7830201  -3.6766427 ]\n",
      "tex row: -1.37 & -1.70 & 3.12\n"
     ]
    }
   ],
   "source": [
    "main();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "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": 1
}