ipynb/john_stachurski-python_vs_julia.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# No, Python is Not Too Slow for Computational Economics\n",
                "\n",
                "### John Stachurski\n",
                "\n",
                "#### Australian National University\n",
                "john.stachurski@anu.edu.au"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Fri 28 Sep 2018 11:18:00 EDT\r\n"
                    ]
                }
            ],
            "source": [
                "!date"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Acknowledgements: Thanks to Quentin Batista and Natasha Watkins for helpful comments and suggestions."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Introduction"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "In [this paper](https://github.com/jstac/julia_python_comparison/blob/master/Update_March_23_2018.pdf), S. Boragan Aruoba and Jesus Fernandez-Villaverde (AFV) study the relative speed of a range of programming lanugages by testing them on value function iteration.  In the abstract they state:\n",
                "\n",
                "**\"The central conclusions of our original paper remain unaltered: C++ is the fastest alternative, Julia offers a great balance of speed and ease of use, and Python is too slow.\"**\n",
                "\n",
                "The last claim is strange, since the authors' own data do not support their conclusion.  In their table 1, Python combined with its scientific libraries is actually reported as slightly faster than Julia (2.31 seconds for Python vs 2.35 for Julia).\n",
                "\n",
                "We rerun the code below and find a similar outcome: The Python code runs between 10% and 30% faster.  (Measurements vary depending on whether or not we include compile time and so on.  Further discussion of timing is given below.) The execution environment is \n",
                "\n",
                "```\n",
                "Architecture:        x86_64\n",
                "CPU op-mode(s):      32-bit, 64-bit\n",
                "Byte Order:          Little Endian\n",
                "CPU(s):              4\n",
                "On-line CPU(s) list: 0-3\n",
                "Thread(s) per core:  2\n",
                "Core(s) per socket:  2\n",
                "Socket(s):           1\n",
                "NUMA node(s):        1\n",
                "Vendor ID:           GenuineIntel\n",
                "CPU family:          6\n",
                "Model:               142\n",
                "Model name:          Intel(R) Core(TM) i5-7300U CPU @ 2.60GHz\n",
                "```\n",
                "\n",
                "We are using the latest version of each language and scientific ecosystem at the time of writing: Julia 1.0 and Anaconda Python 5.2.\n",
                "\n",
                "Of course, we are comparing only Julia vs Python plus its scientific libraries (as opposed to Julia vs pure Python), since no scientist would consider foregoing the scientific libraries when implementing computationally intensive algorithms in Python.\n",
                "\n",
                "Here \"Python plus its scientific libraries\" means the computing environment contained in the popular Anaconda Python distribution.\n",
                "\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Additional Comments in Favor of Python"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "1. Python is massively popular for both general purpose and scientific computing (see, e.g., [this ranking](https://spectrum.ieee.org/static/interactive-the-top-programming-languages-2018)).  Its popularity is partly due to (and partly caused by) the huge range of third party libraries.  Within the scientific domain, these libraries include high quality array processing software and a powerful JIT compiler.  If Python is \"too slow\" for scientific computing in economics, why is that not the case in, say, astrophysics or machine learning?\n",
                "\n",
                "2. There has been some discussion in the past about making parts of Julia interpreted rather than compiled, since JIT compilation makes debugging harder and adds other complications.  Moreover, if one considers, say, a large plotting library, then very little of the code base will be critical to performance.  So why compile it?  In a sense, compiling the entire code base is a form of premature optimization --- decried as the \"root of all evil\" by [Donald Knuth](https://en.wikipedia.org/wiki/Donald_Knuth), one of the greats of scientific computing.  In contrast, Python is interpreted and JIT compilation is judiciously applied to small parts of one's code base.  Perhaps this is the sweet spot for scientific computing that Python has somehow stumbled into?\n",
                "\n",
                "3. AFV do not display any profiling of their codes in the different languages. No memory footprint, no flop rates, etc. are shown. This matters because memory access patterns can always be optimized, especially in small problems as treated here.  One would have to take into account all the individual features of the programming languages to make a fair comparison. (Most likely there would be no surprise, as their performance behavior is known.)  We follow AFV in simply reporting timings but a useful reference for the topics listed above is: S. Goedecker and A. Hoisie. Performance Optimization of Numerically Intensive Codes (SIAM), 2001. ISBN 978-0898714845. "
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Additional Comments in Favor of Julia"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "1. All of Julia and its scientific libraries can be JIT compiled.  The same is not true of Python and its scientific libraries (although progress towards this goal has been rapid).\n",
                "\n",
                "2. As an aside, the author thinks Julia is a great project and hopes that it succeeds.  Python wouldn't be where it is now without the flow of ideas (and competitive pressure) from Julia.\n",
                "\n",
                "3. Python's primary scientific JIT compiler (Numba) and Julia's JIT compiler are built on the same architecture (LLVM).  Similar execution times are to be expected.  From this perspective it is surprising that the Python code runs faster on this machine.  Other users might experience different outcomes."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Timings"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "The code below is exactly as in the [GitHub](https://github.com/jesusfv/Comparison-Programming-Languages-Economics) repo, except for the following changes, consistent with comments made in the 2018 paper. \n",
                "\n",
                "1. Doubled the length of the capital grid.\n",
                "2. Convergence tolerance is 1e-8.\n",
                "3. Print out status after every 20 iterations.\n",
                "4. Changed print statements to print functions (Python 3 compliance)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Output =  0.5627314338711378  Capital =  0.178198287392527  Consumption =  0.3845331464786108\n",
                        " Iteration =  1 , Sup Diff =  0.05274159340733661\n",
                        " Iteration =  20 , Sup Diff =  0.018703459886607154\n",
                        " Iteration =  40 , Sup Diff =  0.006668541708074516\n",
                        " Iteration =  60 , Sup Diff =  0.0023813118038720216\n",
                        " Iteration =  80 , Sup Diff =  0.0008513397746897633\n",
                        " Iteration =  100 , Sup Diff =  0.00030462324420166276\n",
                        " Iteration =  120 , Sup Diff =  0.00010906950872124899\n",
                        " Iteration =  140 , Sup Diff =  3.907108211809174e-05\n",
                        " Iteration =  160 , Sup Diff =  1.400864463663165e-05\n",
                        " Iteration =  180 , Sup Diff =  5.026474537817016e-06\n",
                        " Iteration =  200 , Sup Diff =  1.8035522479920019e-06\n",
                        " Iteration =  220 , Sup Diff =  6.471316942313621e-07\n",
                        " Iteration =  240 , Sup Diff =  2.3219657907525004e-07\n",
                        " Iteration =  260 , Sup Diff =  8.331409384609856e-08\n",
                        " Iteration =  280 , Sup Diff =  2.989378011797328e-08\n",
                        " Iteration =  300 , Sup Diff =  1.0726129207050406e-08\n",
                        " Iteration =  302 , Sup Duff =  9.681200952016411e-09\n",
                        " \n",
                        " My Check =  0.14404414369631846\n",
                        " \n",
                        "Elapse time = is  2.7670769691467285\n"
                    ]
                }
            ],
            "source": [
                "# Basic RBC model with full depreciation (Alternate 1)\n",
                "#\n",
                "# Jesus Fernandez-Villaverde\n",
                "# Haverford, July 3, 2013\n",
                "\n",
                "import numpy as np\n",
                "import math\n",
                "import time\n",
                "from numba import jit\n",
                "\n",
                "# - Start Inner Loop - #\n",
                "# - bbeta                   float\n",
                "# - nGridCapital:           int64\n",
                "# - gridCapitalNextPeriod:  int64\n",
                "# - mOutput:                float (17820 x 5)\n",
                "# - nProductivity:          int64\n",
                "# - vGridCapital:           float (17820, )\n",
                "# - mValueFunction:         float (17820 x 5)\n",
                "# - mPolicyFunction:        float (17820 x 5)\n",
                "\n",
                "@jit\n",
                "def innerloop(bbeta, nGridCapital, gridCapitalNextPeriod, mOutput, nProductivity, vGridCapital, expectedValueFunction, mValueFunction, mValueFunctionNew, mPolicyFunction):\n",
                "\n",
                "    for nCapital in range(nGridCapital):\n",
                "        valueHighSoFar = -100000.0\n",
                "        capitalChoice  = vGridCapital[0]\n",
                "        \n",
                "        for nCapitalNextPeriod in range(gridCapitalNextPeriod, nGridCapital):\n",
                "            consumption = mOutput[nCapital,nProductivity] - vGridCapital[nCapitalNextPeriod]\n",
                "            valueProvisional = (1-bbeta)*np.log(consumption)+bbeta*expectedValueFunction[nCapitalNextPeriod,nProductivity];\n",
                "\n",
                "            if  valueProvisional > valueHighSoFar:\n",
                "                valueHighSoFar = valueProvisional\n",
                "                capitalChoice = vGridCapital[nCapitalNextPeriod]\n",
                "                gridCapitalNextPeriod = nCapitalNextPeriod\n",
                "            else:\n",
                "                break \n",
                "\n",
                "        mValueFunctionNew[nCapital,nProductivity] = valueHighSoFar\n",
                "        mPolicyFunction[nCapital,nProductivity]   = capitalChoice\n",
                "\n",
                "    return mValueFunctionNew, mPolicyFunction\n",
                "\n",
                "def main_func():\n",
                "\n",
                "    #  1. Calibration\n",
                "\n",
                "    aalpha = 1.0/3.0     # Elasticity of output w.r.t. capital\n",
                "    bbeta  = 0.95        # Discount factor\n",
                "\n",
                "    # Productivity values\n",
                "    vProductivity = np.array([0.9792, 0.9896, 1.0000, 1.0106, 1.0212],float)\n",
                "\n",
                "    # Transition matrix\n",
                "    mTransition   = np.array([[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],\n",
                "                     [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],\n",
                "                     [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],\n",
                "                     [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],\n",
                "                     [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]],float)\n",
                "\n",
                "    ## 2. Steady State\n",
                "\n",
                "    capitalSteadyState     = (aalpha*bbeta)**(1/(1-aalpha))\n",
                "    outputSteadyState      = capitalSteadyState**aalpha\n",
                "    consumptionSteadyState = outputSteadyState-capitalSteadyState\n",
                "\n",
                "    print(\"Output = \", outputSteadyState, \" Capital = \", capitalSteadyState, \" Consumption = \", consumptionSteadyState) \n",
                "\n",
                "    # We generate the grid of capital\n",
                "    vGridCapital           = np.arange(0.5*capitalSteadyState,1.5*capitalSteadyState,0.000005)\n",
                "\n",
                "    nGridCapital           = len(vGridCapital)\n",
                "    nGridProductivity      = len(vProductivity)\n",
                "\n",
                "    ## 3. Required matrices and vectors\n",
                "\n",
                "    mOutput           = np.zeros((nGridCapital,nGridProductivity),dtype=float)\n",
                "    mValueFunction    = np.zeros((nGridCapital,nGridProductivity),dtype=float)\n",
                "    mValueFunctionNew = np.zeros((nGridCapital,nGridProductivity),dtype=float)\n",
                "    mPolicyFunction   = np.zeros((nGridCapital,nGridProductivity),dtype=float)\n",
                "    expectedValueFunction = np.zeros((nGridCapital,nGridProductivity),dtype=float)\n",
                "\n",
                "    # 4. We pre-build output for each point in the grid\n",
                "\n",
                "    for nProductivity in range(nGridProductivity):\n",
                "        mOutput[:,nProductivity] = vProductivity[nProductivity]*(vGridCapital**aalpha)\n",
                "\n",
                "    ## 5. Main iteration\n",
                "\n",
                "    maxDifference = 10.0\n",
                "    tolerance = 0.00000001\n",
                "    iteration = 0\n",
                "\n",
                "    log = math.log\n",
                "    zeros = np.zeros\n",
                "    dot = np.dot\n",
                "\n",
                "    while(maxDifference > tolerance):\n",
                "\n",
                "        expectedValueFunction = dot(mValueFunction,mTransition.T)\n",
                "\n",
                "        for nProductivity in range(nGridProductivity):\n",
                "\n",
                "            # We start from previous choice (monotonicity of policy function)\n",
                "            gridCapitalNextPeriod = 0\n",
                "\n",
                "            # - Start Inner Loop - #\n",
                "\n",
                "            mValueFunctionNew, mPolicyFunction = innerloop(bbeta, nGridCapital, gridCapitalNextPeriod, mOutput, nProductivity, vGridCapital, expectedValueFunction, mValueFunction, mValueFunctionNew, mPolicyFunction)\n",
                "\n",
                "            # - End Inner Loop - #\n",
                "\n",
                "        maxDifference = (abs(mValueFunctionNew-mValueFunction)).max()\n",
                "\n",
                "        mValueFunction    = mValueFunctionNew\n",
                "        mValueFunctionNew = zeros((nGridCapital,nGridProductivity),dtype=float)\n",
                "\n",
                "        iteration += 1\n",
                "        if(iteration%20 == 0 or iteration == 1):\n",
                "            print(\" Iteration = \", iteration, \", Sup Diff = \", maxDifference)\n",
                "\n",
                "    return (maxDifference, iteration, mValueFunction, mPolicyFunction)\n",
                "\n",
                "if __name__ == '__main__':\n",
                "    # - Start Timer - #\n",
                "    t1=time.time()\n",
                "    # - Call Main Function - #\n",
                "    maxDiff, iterate, mValueF, mPolicyFunction = main_func()\n",
                "    # - End Timer - #\n",
                "    t2 = time.time()\n",
                "    print(\" Iteration = \", iterate, \", Sup Duff = \", maxDiff)\n",
                "    print(\" \")\n",
                "    print(\" My Check = \", mPolicyFunction[1000-1,3-1])\n",
                "    print(\" \")\n",
                "    print(\"Elapse time = is \", t2-t1)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Let's run once more to eliminate compile time."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Output =  0.5627314338711378  Capital =  0.178198287392527  Consumption =  0.3845331464786108\n",
                        " Iteration =  1 , Sup Diff =  0.05274159340733661\n",
                        " Iteration =  20 , Sup Diff =  0.018703459886607154\n",
                        " Iteration =  40 , Sup Diff =  0.006668541708074516\n",
                        " Iteration =  60 , Sup Diff =  0.0023813118038720216\n",
                        " Iteration =  80 , Sup Diff =  0.0008513397746897633\n",
                        " Iteration =  100 , Sup Diff =  0.00030462324420166276\n",
                        " Iteration =  120 , Sup Diff =  0.00010906950872124899\n",
                        " Iteration =  140 , Sup Diff =  3.907108211809174e-05\n",
                        " Iteration =  160 , Sup Diff =  1.400864463663165e-05\n",
                        " Iteration =  180 , Sup Diff =  5.026474537817016e-06\n",
                        " Iteration =  200 , Sup Diff =  1.8035522479920019e-06\n",
                        " Iteration =  220 , Sup Diff =  6.471316942313621e-07\n",
                        " Iteration =  240 , Sup Diff =  2.3219657907525004e-07\n",
                        " Iteration =  260 , Sup Diff =  8.331409384609856e-08\n",
                        " Iteration =  280 , Sup Diff =  2.989378011797328e-08\n",
                        " Iteration =  300 , Sup Diff =  1.0726129207050406e-08\n",
                        " Iteration =  302 , Sup Duff =  9.681200952016411e-09\n",
                        " \n",
                        " My Check =  0.14404414369631846\n",
                        " \n",
                        "Elapse time = is  2.024168014526367\n"
                    ]
                }
            ],
            "source": [
                "\n",
                "# - Start Timer - #\n",
                "t1=time.time()\n",
                "# - Call Main Function - #\n",
                "maxDiff, iterate, mValueF, mPolicyFunction = main_func()\n",
                "# - End Timer - #\n",
                "t2 = time.time()\n",
                "print(\" Iteration = \", iterate, \", Sup Duff = \", maxDiff)\n",
                "print(\" \")\n",
                "print(\" My Check = \", mPolicyFunction[1000-1,3-1])\n",
                "print(\" \")\n",
                "print(\"Elapse time = is \", t2-t1)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Julia version\n",
                "\n",
                "To run the Julia code you will need to download it from [here](https://github.com/jstac/julia_python_comparison) and place it in the present working directory. We remove the print statements and add a warm-up iteration to be consistent with the Python code."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "First run (warm-up): \n",
                        "  3.519536 seconds (2.62 M allocations: 1.337 GiB, 4.94% gc time)\n",
                        "Second run: \n",
                        "  2.453314 seconds (1.84 k allocations: 1.211 GiB, 3.78% gc time)\n"
                    ]
                }
            ],
            "source": [
                "!julia RBC_Julia.jl"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## An Alternate Python Version"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Incidentally, we would write the code a little differently.  For those who are interested, here's a version that's conforms more closely to PEP8 and other Python conventions.  \n",
                "\n",
                "Parallelization of the Bellman maximization loop is included, which involves switching one `range` call to `prange`.  There is a small speed gain in the second run.  The gain is only small because the algorithm exploits several features of the model that save time in serial execution but make parallelization more challenging (due to different wait times across threads).\n",
                "\n",
                "Variables are global for simplicity.  A natural next step would be to put them in a small class.  This would help with comparative statics."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {},
            "outputs": [],
            "source": [
                "from numba import prange, njit"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {},
            "outputs": [],
            "source": [
                "β = 0.95\n",
                "α = 1/3\n",
                "A = np.array([0.9792, 0.9896, 1.0000, 1.0106, 1.0212])  # Productivity\n",
                "\n",
                "# Transition matrix\n",
                "π = np.array([[0.9727, 0.0273, 0.0000, 0.0000, 0.0000],\n",
                "              [0.0041, 0.9806, 0.0153, 0.0000, 0.0000],\n",
                "              [0.0000, 0.0082, 0.9837, 0.0082, 0.0000],\n",
                "              [0.0000, 0.0000, 0.0153, 0.9806, 0.0041],\n",
                "              [0.0000, 0.0000, 0.0000, 0.0273, 0.9727]])\n",
                "\n",
                "# Steady state values\n",
                "k_star = (α * β)**(1 / (1 - α))  # Steady state capital\n",
                "y_star = k_star**α               # Steady state output\n",
                "c_star = y_star - k_star         # Steady state consumption\n",
                "\n",
                "# Set up capital grid\n",
                "n = 35640                        # Number of grid points\n",
                "k_grid = np.linspace(0.5 * k_star, 1.5 * k_star, n)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [],
            "source": [
                "@njit(parallel=True)\n",
                "def inner_loop(k_grid, v, A, π, β):\n",
                "    \n",
                "    E_v = v @ π.T                    # Expected value function\n",
                "    y_grid = np.outer(k_grid**α, A)  # Output grid\n",
                "    \n",
                "    v_new = np.zeros_like(y_grid)\n",
                "    σ = np.zeros_like(y_grid)\n",
                "    \n",
                "    for a in prange(len(A)):\n",
                "        \n",
                "        k_next = 0\n",
                "        \n",
                "        for i in range(len(k_grid)):\n",
                "\n",
                "            v_max = -np.inf\n",
                "            k_prime = k_grid[0]\n",
                "\n",
                "            for j in range(k_next, len(k_grid)):  \n",
                "                y = y_grid[i, a]\n",
                "                k = k_grid[j]\n",
                "                c = y - k\n",
                "                v_temp = (1 - β) * np.log(c) + β * E_v[j, a];\n",
                "\n",
                "                if v_temp > v_max:\n",
                "                    v_max = v_temp\n",
                "                    k_prime = k_grid[j]\n",
                "                    k_next = j\n",
                "                else:\n",
                "                    break \n",
                "\n",
                "            v_new[i, a] = v_max\n",
                "            σ[i, a] = k_prime\n",
                "\n",
                "    return v_new, σ"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {},
            "outputs": [],
            "source": [
                "def solve_model(k_grid, \n",
                "                A, \n",
                "                π,\n",
                "                α=1/3, \n",
                "                β=0.95,\n",
                "                tol=1e-8,\n",
                "                maxiter=1000):\n",
                "\n",
                "    # Initialize values\n",
                "    diff = 1e3\n",
                "    i = 0\n",
                "    v = np.zeros((len(k_grid), len(A)))\n",
                "\n",
                "    while diff > tol and i < maxiter:       \n",
                "            \n",
                "        # Update value and policy functions for given productivity\n",
                "        v_new, σ = inner_loop(k_grid, v, A, π, β)\n",
                "\n",
                "        diff = np.abs(v_new - v).max()\n",
                "        v = v_new\n",
                "\n",
                "        i += 1\n",
                "        if(i % 20 == 0 or i == 1):\n",
                "            print(f\"Iteration {i}: diff = {diff}\")\n",
                "            \n",
                "    print(f'Coverged in {i} iterations')\n",
                "            \n",
                "    return v_new, σ"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Iteration 1: diff = 0.05274159340733661\n",
                        "Iteration 20: diff = 0.018703459885652807\n",
                        "Iteration 40: diff = 0.006668541708059972\n",
                        "Iteration 60: diff = 0.0023813118038562564\n",
                        "Iteration 80: diff = 0.0008513397746838791\n",
                        "Iteration 100: diff = 0.0003046232441998864\n",
                        "Iteration 120: diff = 0.00010906950872069388\n",
                        "Iteration 140: diff = 3.9071082117980716e-05\n",
                        "Iteration 160: diff = 1.4008644636520629e-05\n",
                        "Iteration 180: diff = 5.026474537705994e-06\n",
                        "Iteration 200: diff = 1.8035522478809796e-06\n",
                        "Iteration 220: diff = 6.471316943423844e-07\n",
                        "Iteration 240: diff = 2.3219657918627234e-07\n",
                        "Iteration 260: diff = 8.331409395712086e-08\n",
                        "Iteration 280: diff = 2.989378011797328e-08\n",
                        "Iteration 300: diff = 1.0726129207050406e-08\n",
                        "Coverged in 302 iterations\n",
                        "CPU times: user 5.29 s, sys: 209 ms, total: 5.5 s\n",
                        "Wall time: 3.27 s\n"
                    ]
                },
                {
                    "data": {
                        "text/plain": [
                            "(array([[-0.99728785, -0.98552127, -0.97408174, -0.96027354, -0.94819576],\n",
                            "        [-0.99728648, -0.9855199 , -0.97408037, -0.96027217, -0.9481944 ],\n",
                            "        [-0.99728512, -0.98551853, -0.974079  , -0.9602708 , -0.94819303],\n",
                            "        ...,\n",
                            "        [-0.97049334, -0.95872675, -0.947286  , -0.93347902, -0.92140126],\n",
                            "        [-0.97049289, -0.9587263 , -0.94728554, -0.93347856, -0.9214008 ],\n",
                            "        [-0.97049243, -0.95872584, -0.94728509, -0.93347811, -0.92140034]]),\n",
                            " array([[0.13849505, 0.13996508, 0.14145011, 0.14293514, 0.14443516],\n",
                            "        [0.13849505, 0.13997008, 0.14145011, 0.14294014, 0.14443516],\n",
                            "        [0.13850006, 0.13997008, 0.14145511, 0.14294014, 0.14444016],\n",
                            "        ...,\n",
                            "        [0.19974118, 0.20186122, 0.20400626, 0.2061463 , 0.20830634],\n",
                            "        [0.19974118, 0.20186622, 0.20400626, 0.2061463 , 0.20831134],\n",
                            "        [0.19974118, 0.20186622, 0.20400626, 0.2061513 , 0.20831134]]))"
                        ]
                    },
                    "execution_count": 9,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "%time solve_model(k_grid, A, π, α, β)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Iteration 1: diff = 0.05274159340733661\n",
                        "Iteration 20: diff = 0.018703459885652807\n",
                        "Iteration 40: diff = 0.006668541708059972\n",
                        "Iteration 60: diff = 0.0023813118038562564\n",
                        "Iteration 80: diff = 0.0008513397746838791\n",
                        "Iteration 100: diff = 0.0003046232441998864\n",
                        "Iteration 120: diff = 0.00010906950872069388\n",
                        "Iteration 140: diff = 3.9071082117980716e-05\n",
                        "Iteration 160: diff = 1.4008644636520629e-05\n",
                        "Iteration 180: diff = 5.026474537705994e-06\n",
                        "Iteration 200: diff = 1.8035522478809796e-06\n",
                        "Iteration 220: diff = 6.471316943423844e-07\n",
                        "Iteration 240: diff = 2.3219657918627234e-07\n",
                        "Iteration 260: diff = 8.331409395712086e-08\n",
                        "Iteration 280: diff = 2.989378011797328e-08\n",
                        "Iteration 300: diff = 1.0726129207050406e-08\n",
                        "Coverged in 302 iterations\n",
                        "CPU times: user 3.67 s, sys: 153 ms, total: 3.82 s\n",
                        "Wall time: 1.58 s\n"
                    ]
                }
            ],
            "source": [
                "%time v_new, σ = solve_model(k_grid, A, π)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "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.4"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}