TOI1_TA_221107_VFI_Coding_Julia.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Value Function Iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg\n",
    "# Pkg.add([\"LinearAlgebra\",\"Gadfly\",\"DataFrames\",\"DataFramesMeta\",\"Chain\"]) # Run this line when running for the first time\n",
    "using LinearAlgebra,Gadfly,DataFrames,DataFramesMeta,Chain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Bruteforce Algorithm\n",
    "For this demonstration I am using $\\beta = 0.95, z = 1, \\alpha = 0.3,$ and $\\delta = 1$. This full depreciation case is particularly useful for expositional purposes becuase it admits an analytical solution for the value function. Hence, in this case, we will be able to check our iterated numerical solution against the true solution.\n",
    "\n",
    "Given these parameter values the Bellman for our problem is,\n",
    "$$V(k) = \\max_{k'\\in[0,k^{0.3}]} \\left\\{ \\log\\left(k^{0.3} - k'\\right) + 0.95 V(k')\\right\\}$$\n",
    "\n",
    "Hence, the associated Bellman operator is,\n",
    "$$Tf(k) = \\max_{k'\\in[0,k^{0.3}]} \\left\\{ \\log\\left(k^{0.3} - k'\\right) + 0.95 f(k')\\right\\}$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#=\n",
    "Define the parameters for convenience. \n",
    "    In Julia you can type Greek letters by typing \\beta + tab.\n",
    "    The semicolons at the end suppress output to REPL.\n",
    "=#\n",
    "β = 0.95;\n",
    "z = 1;\n",
    "α = 0.3;\n",
    "δ = 1;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1: Create grid for the state space\n",
    "For our parameter values the bounded state space is given by,\n",
    "$$k\\in \\left[0,\\left(\\frac{\\delta}{z}\\right)^{\\frac{1}{\\alpha - 1}}\\right] = [0,1]$$\n",
    "\n",
    "Let's create a grid with 100 points between [0,1]."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = collect(0.01:0.01:1); #state space\n",
    "println(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2: Guess Initial Function $V_0(k)$\n",
    "\n",
    "Let's be really naive and set $V_0(k) = 0, \\forall k$. We need to evaluate this function for all the $k\\in X$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V0 = [0 for k in X]; #initial guess\n",
    "println(V0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3: Iterate Using the Bellman Operator\n",
    "\n",
    "As mentioned above the Bellman operator is,\n",
    "$$V_{n+1}(k) = \\max_{k'\\in[0,k^{0.3}]} \\left\\{ \\log\\left(k^{0.3} - k'\\right) + \\beta V_{n}(k')\\right\\}$$\n",
    "\n",
    "We want to have an expression for the new function $V_{n+1}(k)$ for each $k\\in X$. Therefore we will need to loop through each $k\\in X$, and then for each $k$ we need to,\n",
    "1. Evaluate $\\log\\left(k^{0.3} - k'\\right) + \\beta V_{n}(k')$ for each $k' \\in \\Gamma(k)$. *(In the algorithm below I evaluate the funciton for all $k'\\in X$, but return a very negative value for infeasible k', this will ensure they are not picked as a max.)*\n",
    "1. Take the max value from the previous step and assign it to $V_{n+1}(k)$. If you like, at this stage you can store the maximising k' (the $\\arg \\max$) as the value for the policy rule i.e. $g_{n+1}(k)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iter = 0\n",
    "iter_max = 1000\n",
    "error = 9\n",
    "err_tol = 1e-6\n",
    "\n",
    "X = collect(0.01:0.01:1); # State space\n",
    "V0 = fill(0.0,100); # Initial guess\n",
    "V1 = [0.0 for k in X]; # updated guess (this way of generating a list is called list comprehension in Julia and Python)\n",
    "#=\n",
    "Coding comments:\n",
    "    In Julia, sometimes it is important to use the exact type (e.g. a float number) to initiate a list, or there could be errors.\n",
    "=#\n",
    "\n",
    "mV = zeros(length(X),iter_max) # Create a matrix to save our guessed V in each iteration (for plotting later)\n",
    "\n",
    "while error > err_tol && iter_max <= iter_max # Iterate until error is small enough or the maximal number of iterations is reached.\n",
    "    \n",
    "    for i = 1:100 # Loop over grid points\n",
    "    \n",
    "        k = X[i] # k\n",
    "        objective = [X[j] <= k^α ? log(k^α - X[j]) + β*V0[j] : -Inf for j = 1:100] # Evaulate the obejective function at each k'\n",
    "        V1[i] = maximum(objective) # Find the maximum of the objective function\n",
    "    \n",
    "    end\n",
    "\n",
    "    error = sum((V1-V0).^2) # You could also use sup norm maximum(abs.(V1-V0))\n",
    "    #=\n",
    "    Coding comments:\n",
    "        The dot denotes \"broadcast\" in Julia. It applies a function that takes one element as an argument to each element of a vector/matrix.\n",
    "    =#\n",
    "    \n",
    "    V0 = copy(V1) # Update guess\n",
    "    #=\n",
    "    Coding comments:\n",
    "        We need to be careful here. Julia arrays are not copied when assigned to another variable. After A = B, changing elements of B will modify A as well.\n",
    "        https://docs.julialang.org/en/v1/manual/noteworthy-differences/\n",
    "    =#\n",
    "\n",
    "    iter = iter+1\n",
    "\n",
    "    # Below are additional codes to save data points for plotting\n",
    "    mV[:,iter] = V1\n",
    "end\n",
    "println(iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Some take-aways\n",
    "\n",
    "The code below plots each iteration. As you can see, the initial guess was pretty bad, but we still seem to have converged. If you like, play-around with the initial guess to see how the number of iterations changes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_out = DataFrame(mV[:,1:iter],:auto)\n",
    "@transform!(df_out, :k=X) # Add k values\n",
    "df_out = stack(df_out,Not(:k),variable_name=:iter,value_name=:V) # Reshape to long\n",
    "@transform!(df_out, :iter=parse.(Int64, SubString.(:iter,2))); # Convert iteration count to int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "println(\"Number of iterations: \", iter)\n",
    "plot(df_out,x=:k,y=:V, Geom.line, color=:iter, Theme(key_position = :none), Guide.title(\"Value Function Iterations\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numerical estimate vs. truth\n",
    "\n",
    "As mentioned at the beginning, in the full depreciation case we in fact have an analytical expression for the value function $V(k)$. Let's see how our numerical method compares to the true function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function v_hat(k)\n",
    "    a = 0.3\n",
    "    b = 0.95\n",
    "\n",
    "    return a/(1-a*b)*log(k) + 1/(1-b)*(a*b/(1-a*b)*log(a*b)+log(1-a*b))\n",
    "end\n",
    "\n",
    "V_true = [v_hat(k) for k in X]\n",
    "\n",
    "plot(layer(@subset(df_out,:iter.==iter),x=:k,y=:V, Geom.line, color=[\"Numerical Iteration\"],Theme(line_style=[:solid]), order=1),\n",
    "    layer(x=X,y=V_true, Geom.line, color=[\"True Function\"],Theme(line_style=[:dot]), order=2),\n",
    "    Guide.title(\"Value Functions\"))\n",
    "\n",
    "# It might be easier to tweak figures if you RCall ggplot2 (at least I feel this way at the moment)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.8.2",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}