Notebooks for examples and new numerical methods on HJBEs

continuous_time_methods/notes/discretized-differential-operator-derivation.aux

@@ -0,0 +1 @@
+\bibdata{hact}

continuous_time_methods/tutorials/growth-hjb-implicit.ipynb

@@ -0,0 +1,236 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "---\ntitle       : \"Solving HJB equation for neoclassical growth models\"\nauthor      : Chiyoung Ahn (@chiyahn)\ndate        : \n\n\n`j using Dates; print(Dates.today())`\n---\n\n### About this document\nPresented by Chiyoung Ahn (@chiyahn), written for `Weave.jl`, based on Ben Moll's [note](http://www.princeton.edu/~moll/HACTproject/HACT_Additional_Codes.pdf) and [code](http://www.princeton.edu/~moll/HACTproject/HJB_NGM_implicit.m)."
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "using LinearAlgebra, Parameters, Plots, BenchmarkTools\ngr(fmt = :png); # save plots in .png"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Model\nConsider finding the optimal consumption plan $c(t)$ for\n\n$$\n\\max_{\\{c(t)\\}_{t \\geq 0} } \\int e^{-\\rho t } u(c(t)) dt\n$$\n\nwith $u(c(t)) = c(t)^{1-\\gamma} / (1-\\gamma)$ for some $\\gamma \\neq 1$ and the following law of motion for $k$:\n\n$$\n\\dot k(t) = F(k(t)) - \\delta k(t) - c(t)\n$$\n\n## Numerical solution of HJBE\nThe corresponding HJB equation is\n\n$$\n\\rho v (k) = \\max_c u(c) + v'(k ) (F(k) - \\delta k - c)\n$$\nThe first order condition with respect to $c$ on the maximand function yields \n\n$$\nu'(c) = v'(k)\n$$\n\ni.e., $c = (u')^{-1} (v' (k))$ for each $k$ at optimal.\n\n\n### Upwind scheme and discretization of $\\nabla v$\nNote that $v$ is concave in $k$ and the drift of state variables can be either positive or negative. Basic idea of upwind scheme: use forward difference when the drift of the state variable $f(x, \\alpha)$ is positive, backward difference when it is negative. Let $v'_{i, B}$ and $v'_{i, F}$ be $v'$ computed from backward difference and forward difference respectively.\n\nDefine the followings:\n\n$$\ns_{i,F} = F(k_i) - \\delta k_i - (u')^{-1} (v'_{i,F}) \\\\\ns_{i,B} = F(k_i) - \\delta k_i - (u')^{-1} (v'_{i,B})\n$$\n\nUsing the upwind scheme, the HJB eqaution can be rewritten as\n\n$$\n\\rho v_i = u(c_i) + \\dfrac{v_{i+1} - v_i}{\\Delta k} s^+_{i,F} + \\dfrac{v_i - v_{i-1}}{\\Delta k} s^-_{i, B}\n$$\n\nfor each $i$. This can be written in a compact matrix form\n\n$$\n\\rho {\\mathbf{v}} = {\\mathbf{u}} + {\\mathbf{A}}({\\mathbf{v}}) {\\mathbf{v}}\n$$\n\nsuch that the $i$th row of ${\\mathbf{A}}({\\mathbf{v}})$ is defined as \n\n$$\n\\begin{bmatrix}\n0 & \\cdots & 0 &  - \\dfrac{s^-_{i,B}}{\\Delta k} & \\dfrac{s^-_{i,B}}{\\Delta k} - \\dfrac{s^-_{i,F}}{\\Delta k} & \\dfrac{s^-_{i,F}}{\\Delta k} & 0 & \\cdots & 0\n\\end{bmatrix}\n$$\n\nwhere the non-zero elements above are located in $i-1, i, i+1$th columns respectively. \n\n### Discretization of $u(c)$\nOne thing left to find is ${\\mathbf{u}}$ at optimal consumption plan $c$. \n\n[Achdou et al. (2017)](http://www.princeton.edu/~moll/HACT.pdf) and [Fernandez-Villaverde et al. (2018)](https://www.sas.upenn.edu/~jesusfv/Financial_Frictions_Wealth_Distribution.pdf) suggest using the derivative of $v$ by the FOC above, which yields \n\n$$u (c) = u ( (u')^{-1} (v'(k))  ) $$\n\nTo compute the derivative, Achdou et al. (2017) and Fernandez-Villaverde et al. (2018) use the following discretization scheme:\n\n$$\nv'_i = v'_{i,F} {\\mathbf{1}} ( s_{i, F} > 0) + v'_{i,B} {\\mathbf{1}} ( s_{i, B} < 0) + \\overline{v}'_i {\\mathbf{1}} ( s_{i, F} < 0 < s_{i, B})\n$$\nwhere $\\overline{v}'_i := u'(F(k_i) - \\delta k_i)$, i.e., the steady state such that $s_i = 0$. This gives ${\\mathbf{u}}$ by computing $u(c(v'_i))$ at each grid $i$.\n\nHere's my take on this. I argue that there is no need to arbitrarily define $\\overline{v}'_i$ to deal with the case when $s_{i}$ is taken to be zero.  Note that we need $\\overline {v}_i'$ only for $u(c)$ when the current state is close enough to the steady state, i.e., $s_i \\approx 0$ by taking $s_i = 0$. On the other hand, note that when $s_i = 0$, we can directly compute the consumption value by the law of motion $\\dot k (t) = 0$ without relying on $v'$:\n\n$$\nc_i = F(k_i) - \\delta k_i\n$$\nHence, instead of defining $v'_i$ and $u(c(v'_i))$ accordingly, we can use the following consumption approximation scheme:\n\n$$\nc_i = c(v'_{i,F}) {\\mathbf{1}} ( s_{i, F} > 0) + c(v'_{i,B}) {\\mathbf{1}} ( s_{i, B} < 0) + [F(k_i) - \\delta(k_i)] {\\mathbf{1}} ( s_{i, F} < 0 < s_{i, B})\n$$\nso that $u(c(v'_i))$ can be now replaced with $u(c_i)$, which ameliorates the need of defining an arbitary discretization scheme for $v'_i$ when $s_i$ is close to zero.\n\n\n## Setup\n### Utility function"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "γ = 2.0\nu(c) = c^(1-γ)/(1-γ) # payoff function by control variable (consumption)\nu_prime(c) = c^(-γ) # derivative of payoff function by control variable (consumption)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Law of motion\nDefine a production as follows first:"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "A_productivity = 1.0\nα = 0.3 \nF(k) = A_productivity*k^α;"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "The corresponding law of motion for `k` given current `k` (state) and `c` (control)"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "δ = 0.05\nf(k, c) = F(k) - δ*k - c; # law of motion for saving (state)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Consumption function by inverse (`v_prime`)"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "c(v_prime) = v_prime^(-1/γ); # consumption by derivative of value function at certain k"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Parameters and grids"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "# ρ: utility discount rate\n# δ: capital discount rate\n# γ: CRRA parameter\n# F: production function that maps k to a real number\n# u: utility function that maps c to a real number\n# u_prime: derivative of utility function that maps c to a real number\n# f: law of motion function that maps k (state), c (control) to the derivative of k\n# c: control by v_prime; maps v_prime (derivative of v at certain k) to consumption (control)\n# c_ss: consumption (control) when v' = 0 (i.e. steady state)\nparams = (ρ = 0.05, δ = δ, γ = γ, F = F, u = u, u_prime = u_prime, f = f, c = c, c_ss = 0.0)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "# ks: grids for states (k) -- assume uniform grids\n# Δv: step size for iteration on v\n# vs0: initial guess for vs\n# maxit: maximum number of iterations\n# threshold: threshold to be used for termination condition (maximum(abs.(vs-vs_new)) < threshold)\n# verbose: boolean that ables/disables a verbose option\nk_ss = (α*A_productivity/(params.ρ+params.δ))^(1/(1-α))\nks = range(0.001*k_ss, stop = 2*k_ss, length = 10000)\nsettings = (ks = ks,\n            Δv = 1000, \n            vs0 = (A_productivity .* ks .^ α) .^ (1-params.γ) / (1-params.γ) / params.ρ,\n            maxit = 100, threshold = 1e-8, verbose = false)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Optimal plan solver, using the methods from [Achdou et al. (2017)](http://www.princeton.edu/~moll/HACT.pdf) and [Fernandez-Villaverde et al. (2018)](https://www.sas.upenn.edu/~jesusfv/Financial_Frictions_Wealth_Distribution.pdf)"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "function compute_optimal_plans(params, settings)\n    @unpack ρ, δ, γ, F, u, u_prime, f, c, c_ss = params\n    @unpack ks, Δv, vs0, maxit, threshold, verbose = settings\n\n    P = length(ks) # size of grids\n    Δk = ks[2] - ks[1] # assume uniform grids\n\n    # initial guess\n    vs = vs0; \n    vs_history = zeros(P, maxit)\n    # save control (consumption) plan as well\n    cs = zeros(P) \n\n    # begin iterations\n    for n in 1:maxit\n        # compute derivatives by FD and BD\n        dv = diff(vs) ./ Δk\n        dv_f = [dv; dv[end]] # forward difference\n        dv_b = [dv[1]; dv] # backward difference\n        dv_0 = u_prime.(f.(ks, fill(c_ss, P)))\n\n        # define the corresponding drifts\n        drift_f = f.(ks, c.(dv_f)) \n        drift_b = f.(ks, c.(dv_b))\n\n        # steady states at boundary\n        drift_f[end] = 0.0\n        drift_b[1] = 0.0\n\n        # compute consumptions and corresponding u(v)\n        I_f = drift_f .> 0.0\n        I_b = drift_b .< 0.0\n        I_0 = 1 .- I_f-I_b\n\n        dv_upwind = dv_f.*I_f + dv_b.*I_b + dv_0.*I_0;\n        cs = c.(dv_upwind)\n        us = u.(cs)\n\n        # define the matrix A\n        drift_f_upwind = max.(drift_f, 0.0) ./ Δk\n        drift_b_upwind = min.(drift_b, 0.0) ./ Δk\n        A = LinearAlgebra.Tridiagonal(-drift_b_upwind[2:P], \n                (-drift_f_upwind + drift_b_upwind), \n                drift_f_upwind[1:(P-1)]) \n\n        # solve the corresponding system to get vs_{n+1}\n        vs_new = (Diagonal(fill((ρ + 1/Δv), P)) - A) \\ (us + vs / Δv)\n\n        # show distance between vs_{n+1} and vs_n if verbose option is one\n        if (verbose) @show maximum(abs.(vs - vs_new)) end\n        \n        # check termination condition \n        if (maximum(abs.(vs-vs_new)) < threshold)\n            if (verbose) println(\"Value function converged -- total number of iterations: $n\") end\n            return (vs = vs, cs = cs, vs_history = vs_history) \n        end\n        \n        # update vs_{n+1}\n        vs = vs_new\n        vs_history[:,n] = vs\n    end\n    return (vs = vs, cs = cs, vs_history = vs_history) \nend"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### Optimal plan solver, using my new method"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "function compute_optimal_plans(params, settings)\n    @unpack ρ, δ, γ, F, u, u_prime, f, c, c_ss = params\n    @unpack ks, Δv, vs0, maxit, threshold, verbose = settings\n\n    P = length(ks) # size of grids\n    Δk = ks[2] - ks[1] # assume uniform grids\n\n    # initial guess\n    vs = vs0; \n    vs_history = zeros(P, maxit)\n    # save control (consumption) plan as well\n    cs = zeros(P) \n\n    # begin iterations\n    for n in 1:maxit\n        # compute derivatives by FD and BD\n        dv = diff(vs) ./ Δk\n        dv_f = [dv; dv[end]] # forward difference\n        dv_b = [dv[1]; dv] # backward difference\n\n        # define the corresponding drifts\n        drift_f = f.(ks, c.(dv_f)) \n        drift_b = f.(ks, c.(dv_b))\n\n        # steady states at boundary\n        drift_f[end] = 0.0\n        drift_b[1] = 0.0\n\n        # compute consumptions and corresponding u(v)\n        I_f = drift_f .> 0.0\n        I_b = drift_b .< 0.0\n        I_0 = 1 .- I_f-I_b\n\n        dv_upwind = dv_f.*I_f + dv_b.*I_b\n        cs_upwind = c.(dv_upwind)\n        cs_0 = f.(ks, 0.0) # this gives consumption when the state is zero\n        cs = cs_upwind.*I_f + cs_upwind.*I_b + cs_0.*I_0;\n        \n        us = u.(cs)\n\n        # define the matrix A\n        drift_f_upwind = max.(drift_f, 0.0) ./ Δk\n        drift_b_upwind = min.(drift_b, 0.0) ./ Δk\n        A = LinearAlgebra.Tridiagonal(-drift_b_upwind[2:P], \n                (-drift_f_upwind + drift_b_upwind), \n                drift_f_upwind[1:(P-1)]) \n\n        # solve the corresponding system to get vs_{n+1}\n        vs_new = (Diagonal(fill((ρ + 1/Δv), P)) - A) \\ (us + vs / Δv)\n\n        # show distance between vs_{n+1} and vs_n if verbose option is one\n        if (verbose) @show maximum(abs.(vs - vs_new)) end\n        \n        # check termination condition \n        if (maximum(abs.(vs-vs_new)) < threshold)\n            if (verbose) println(\"Value function converged -- total number of iterations: $n\") end\n            return (vs = vs, cs = cs, vs_history = vs_history) \n        end\n        \n        # update vs_{n+1}\n        vs = vs_new\n        vs_history[:,n] = vs\n    end\n    return (vs = vs, cs = cs, vs_history = vs_history) \nend"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Solve"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "vs, cs, vs_history = @btime compute_optimal_plans(params, settings)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Plot for `v_n(k)` path by `n` (iteration step):"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "plot(ks, vs_history[:,1:3],\n    linewidth = 3,\n    title=\"Value function per iteration step v_n(k)\",xaxis=\"k\",yaxis=\"v(k)\",\n    label = string.(\"v_\",1:3))"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## Plots\n### `v(k)`"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "plot(ks, vs,\n    linewidth = 3,\n    title=\"Value function v(k)\",xaxis=\"k\",yaxis=\"v(k)\",legend=false)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### `c(k)`"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "plot(ks, cs,\n    lw = 3,\n    title=\"Optimal consumption plan c(k)\",xaxis=\"k\",yaxis=\"c(k)\",legend=false)"
+      ],
+      "metadata": {},
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "### `s(k)`"
+      ],
+      "metadata": {}
+    },
+    {
+      "outputs": [],
+      "cell_type": "code",
+      "source": [
+        "savings = f.(ks, cs) # Savings (states) according to optimal consumption plan\nplot(ks, savings,\n    linewidth = 3,\n    title=\"Optimal saving plan s(k)\",xaxis=\"k\",yaxis=\"s(k)\",legend=false)\nplot!([.0], st = :hline, linestyle = :dot, lw = 3) # zero saving line"
+      ],
+      "metadata": {},
+      "execution_count": null
+    }
+  ],
+  "nbformat_minor": 2,
+  "metadata": {
+    "language_info": {
+      "file_extension": ".jl",
+      "mimetype": "application/julia",
+      "name": "julia",
+      "version": "1.1.0"
+    },
+    "kernelspec": {
+      "name": "julia-1.1",
+      "display_name": "Julia 1.1.0",
+      "language": "julia"
+    }
+  },
+  "nbformat": 4
+}