KSfunctions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Functions for replicating Maliar, Maliar, and Valli (2010, JEDC) to solve Krusell and Smith (1998, JPE) model using Julia\n",
    "\n",
    "\n",
    "By [Shunsuke Hori](https://github.com/Shunsuke-Hori)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview of the notebook\n",
    "This notebook collects functions to solve the model of [Krusell and Smith (1998, JPE)](https://www.journals.uchicago.edu/doi/10.1086/250034) and succesfully replicating the result of [Maliar, Maliar, and Valli (2010, JEDC)](https://www.sciencedirect.com/science/article/pii/S0165188909001328).\n",
    "\n",
    "The solution strategy is as follows\n",
    "\n",
    "1. Solve the individual problem by Euler equation method or value function iteration (VFI) with 2D interpolation\n",
    "    - Agents are boundedly rational. In the code, they take into account the information about the mean of capital\n",
    "    - Aggregate law of motion is approximated by log-linear relation, i.e. $\\log(K_{t+1})=B1+B2\\log(K_{t})$ for good aggregate state and $\\log(K_{t+1})=B3+B4\\log(K_{t})$ for bad aggregate state \n",
    "    - If specified, Howard's policy iteration is used\n",
    "1. Compute the path of aggregate capital using the policy function obtained by 1. There are two ways of simulation:\n",
    "    - Monte Carlo following Krusell and Smith (1998). That is, aggregate technology shocks and idiosyncratic employment shocks are drawn for many agents and many periods. Then, using the LLN, the aggregate capital is computed by aggregating all agents for all period.\n",
    "    - Non-stochastic method following [Young (2010, JEDC)](https://www.sciencedirect.com/science/article/pii/S0165188909001316).\n",
    "1. Update the coefficient of aggregate capital law of motion, $B1$, $B2$, $B3$ and $B4$, by regression\n",
    "1. Check convergence of $B1$, $B2$, $B3$ and $B4$\n",
    "\n",
    "NOTE: Regarding interpolation, Krusell and Smith uses various interpolation scheme depending on the purpose, including polynomial interpolation. Maliar, Maliar, and Valli uses spline interpolation in their paper. This notebook only uses linear interpolation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code to solve models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First thing to do is importing some packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Interpolations # to use interpolation \n",
    "using Random, LinearAlgebra\n",
    "using QuantEcon  # to use `gridmake`, `<:AbstractUtility`\n",
    "using Optim      # to use minimization routine to maximize RHS of bellman equation\n",
    "using GLM        # to regress\n",
    "using JLD2       # to save the result\n",
    "using ProgressMeter # to show progress of iterations\n",
    "using Parameters # to use type with keyword arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Setup\n",
    "\n",
    "Functions in this section are prepared for model parameters and initial guess of the solutions "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `TransitionMatrix`\n",
    "Collection of transition matrices. Each row corresponds to the current state and each column is a probability to jump to that state.\n",
    "\n",
    "##### Fields\n",
    "- `P::Matrix{Float64}`: transition matrix of shocks. States are ordered as (good, employed), (bad, employed), (good, unemployed), (bad, unemployed).\n",
    "- `Pz::Matrix{Float64}`: transition matrix of aggregate shock\n",
    "- `Peps_gg::Matrix{Float64}`: transition matrix of idiosyncratic shock conditional on good to good\n",
    "- `Peps_bb::Matrix{Float64}`: transition matrix of idiosyncratic shock conditional on bad to bad\n",
    "- `Peps_gb::Matrix{Float64}`: transiiton matrix of idiosyncratic shock conditional on good to bad\n",
    "- `Peps_bg::Matrix{Float64}`: transition matrix of idiosyncratic shock conditional on bad to good"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct TransitionMatrix\n",
    "    P::Matrix{Float64}       # 4x4 \n",
    "    Pz::Matrix{Float64}      # 2x2 aggregate shock\n",
    "    Peps_gg::Matrix{Float64} # 2x2 idiosyncratic shock conditional on good to good\n",
    "    Peps_bb::Matrix{Float64} # 2x2 idiosyncratic shock conditional on bad to bad\n",
    "    Peps_gb::Matrix{Float64} # 2x2 idiosyncratic shock conditional on good to bad\n",
    "    Peps_bg::Matrix{Float64} # 2x2 idiosyncratic shock conditional on bad to good\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `UMPSolutionMethod`\n",
    "Abstract type which has `EulerMethod` and `VFI` as subtypes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abstract type UMPSolutionMethod end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `EulerMethod <: UMPSolutionMethod`\n",
    "Specify the update parameter for Euler equation method. Thanks to the multiple dispatch, individual UMP is solved by Euler method by passing `EulerMethod` to `solve_ump!`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@with_kw struct EulerMethod <: UMPSolutionMethod\n",
    "    update_k::Float64 = 0.7\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `VFI<: UMPSolutionMethod`\n",
    "Specify whether Howard policy iteration is implemented or not and the number of the policy iterations. Thanks to the multiple dispatch, individual UMP is solved by value function iteration by passing `VFI` to `solve_ump!`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@with_kw struct VFI <: UMPSolutionMethod\n",
    "    Howard_on::Bool = false\n",
    "    Howard_n_iter::Int = 20\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `create_transition_matrix`\n",
    "Create transition matrices for aggregate shock, idiosyncratic shock, and shock state\n",
    "\n",
    "##### Arguments\n",
    "- `ug::Real`: unemployment rate in good state\n",
    "- `ub::Real`: unemployment rate in bad state\n",
    "- `zg_ave_dur::Real`: average duration of good state\n",
    "- `zb_ave_dur::Real`: average duration of bad state\n",
    "- `ug_ave_dur::Real`: average duration of unemployment in good state\n",
    "- `ub_ave_dur::Real`: average duration of unemployment in bad state\n",
    "- `puu_rel_gb2bb::Real`: prob. of u to u cond. on g to b relative to that of b to b \n",
    "- `puu_rel_bg2gg::Real`: prob. of u to u cond. on b to g relative to that of g to g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function create_transition_matrix(ug::Real, ub::Real,\n",
    "        zg_ave_dur::Real, zb_ave_dur::Real,\n",
    "        ug_ave_dur::Real, ub_ave_dur::Real,\n",
    "        puu_rel_gb2bb::Real, puu_rel_bg2gg::Real)\n",
    "    \n",
    "    # probability of remaining in good state\n",
    "    pgg = 1-1/zg_ave_dur\n",
    "    # probability of remaining in bad state\n",
    "    pbb = 1-1/zb_ave_dur\n",
    "    # probability of changing from g to b\n",
    "    pgb = 1-pgg\n",
    "    # probability of changing from b to g\n",
    "    pbg = 1-pbb  \n",
    "    \n",
    "    # prob. of 0 to 0 cond. on g to g\n",
    "    p00_gg = 1-1/ug_ave_dur\n",
    "    # prob. of 0 to 0 cond. on b to b\n",
    "    p00_bb = 1-1/ub_ave_dur\n",
    "    # prob. of 0 to 1 cond. on g to g\n",
    "    p01_gg = 1-p00_gg\n",
    "    # prob. of 0 to 1 cond. on b to b\n",
    "    p01_bb = 1-p00_bb\n",
    "    \n",
    "    # prob. of 0 to 0 cond. on g to b\n",
    "    p00_gb=puu_rel_gb2bb*p00_bb\n",
    "    # prob. of 0 to 0 cond. on b to g\n",
    "    p00_bg=puu_rel_bg2gg*p00_gg\n",
    "    # prob. of 0 to 1 cond. on g to b\n",
    "    p01_gb=1-p00_gb\n",
    "    # prob. of 0 to 1 cond. on b to g\n",
    "    p01_bg=1-p00_bg\n",
    "    \n",
    "    # prob. of 1 to 0 cond. on  g to g\n",
    "    p10_gg=(ug - ug*p00_gg)/(1-ug)\n",
    "    # prob. of 1 to 0 cond. on b to b\n",
    "    p10_bb=(ub - ub*p00_bb)/(1-ub)\n",
    "    # prob. of 1 to 0 cond. on g to b\n",
    "    p10_gb=(ub - ug*p00_gb)/(1-ug)\n",
    "    # prob. of 1 to 0 cond on b to g\n",
    "    p10_bg=(ug - ub*p00_bg)/(1-ub)\n",
    "    # prob. of 1 to 1 cond. on  g to g\n",
    "    p11_gg= 1-p10_gg\n",
    "    # prob. of 1 to 1 cond. on b to b\n",
    "    p11_bb= 1-p10_bb\n",
    "    # prob. of 1 to 1 cond. on g to b\n",
    "    p11_gb= 1-p10_gb\n",
    "    # prob. of 1 to 1 cond on b to g\n",
    "    p11_bg= 1-p10_bg\n",
    "    \n",
    "    #   (g1)         (b1)        (g0)       (b0)\n",
    "    P=[pgg*p11_gg pgb*p11_gb pgg*p10_gg pgb*p10_gb;\n",
    "       pbg*p11_bg pbb*p11_bb pbg*p10_bg pbb*p10_bb;\n",
    "       pgg*p01_gg pgb*p01_gb pgg*p00_gg pgb*p00_gb;\n",
    "       pbg*p01_bg pbb*p01_bb pbg*p00_bg pbb*p00_bb]\n",
    "    Pz=[pgg pgb;\n",
    "        pbg pbb]\n",
    "    Peps_gg=[p11_gg p10_gg\n",
    "             p01_gg p00_gg]\n",
    "    Peps_bb=[p11_bb p10_bb\n",
    "             p01_bb p00_bb]\n",
    "    Peps_gb=[p11_gb p10_gb\n",
    "             p01_gb p00_gb]\n",
    "    Peps_bg=[p11_bg p10_bg\n",
    "             p01_bg p00_bg]\n",
    "    transmat=TransitionMatrix(P, Pz, Peps_gg, Peps_bb, Peps_gb, Peps_bg)\n",
    "    return transmat\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `KSParameter`\n",
    "Create `NamedTuple` specifying the parameters of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function KSParameter(;\n",
    "            beta::AbstractFloat=0.99,\n",
    "            alpha::AbstractFloat=0.36,\n",
    "            delta::Real=0.025,\n",
    "            theta::Real=1,\n",
    "            k_min::Real=0,\n",
    "            k_max::Real=1000,\n",
    "            k_size::Integer=100,\n",
    "            K_min::Real=30,\n",
    "            K_max::Real=50,\n",
    "            K_size::Integer=4,\n",
    "            z_min::Real=0.99,\n",
    "            z_max::Real=1.01,\n",
    "            z_size::Integer=2,\n",
    "            eps_min::Real=0.0,\n",
    "            eps_max::Real=1.0,\n",
    "            eps_size::Integer=2,\n",
    "            ug::AbstractFloat=0.04,\n",
    "            ub::AbstractFloat=0.1,\n",
    "            zg_ave_dur::Real=8,\n",
    "            zb_ave_dur::Real=8,\n",
    "            ug_ave_dur::Real=1.5,\n",
    "            ub_ave_dur::Real=2.5,\n",
    "            puu_rel_gb2bb::Real=1.25,\n",
    "            puu_rel_bg2gg::Real=0.75,\n",
    "            mu::Real=0, \n",
    "            degree::Real=7)\n",
    "    if theta == 1\n",
    "        u = LogUtility()\n",
    "    else\n",
    "        u = CRRAUtility(theta)\n",
    "    end\n",
    "    l_bar=1/(1-ub)\n",
    "    # individual capital grid\n",
    "    k_grid=\n",
    "        (range(0, stop=k_size-1, length=k_size)/(k_size-1)).^degree*(k_max-k_min).+k_min   \n",
    "    k_grid[1] = k_min; k_grid[end] = k_max; # adjust numerical error\n",
    "    # aggregate capital grid\n",
    "    K_grid=range(K_min, stop=K_max, length=K_size)\n",
    "    # aggregate technology shock\n",
    "    z_grid=range(z_max, stop=z_min, length=z_size)\n",
    "    # idiosyncratic employment shock grid\n",
    "    eps_grid=range(eps_max, stop=eps_min, length=eps_size)\n",
    "    s_grid=gridmake(z_grid, eps_grid)               # shock grid\n",
    "    # collection of transition matrices\n",
    "    transmat=create_transition_matrix(ug,ub,\n",
    "        zg_ave_dur,zb_ave_dur,\n",
    "        ug_ave_dur,ub_ave_dur,\n",
    "        puu_rel_gb2bb,puu_rel_bg2gg)\n",
    "\n",
    "    ksp=(u=u, beta=beta, alpha=alpha, delta=delta, theta=theta,\n",
    "         l_bar=l_bar, k_min=k_min, k_max=k_max, k_grid=k_grid,\n",
    "         K_min=K_min, K_max=K_max, K_grid=K_grid, z_grid=z_grid,\n",
    "         eps_grid=eps_grid, s_grid=s_grid, k_size=k_size, K_size=K_size,\n",
    "         z_size=z_size, eps_size=eps_size, s_size=z_size*eps_size, \n",
    "         ug=ug, ub=ub, transmat=transmat, mu=mu)\n",
    "\n",
    "    return ksp\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `r`\n",
    "Compute interest rate (=marginal product of capital) given aggregate capital, labor, and productivity\n",
    "### `w`\n",
    "Compute wage given aggregate capital, labor, and productivity\n",
    "\n",
    "##### Arguments\n",
    "- `alpha::Real`: capital share\n",
    "- `z::Real`: aggregate shock\n",
    "- `K::Real`: aggregate capital\n",
    "- `L::Real`: aggregate labor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r(alpha::Real, z::Real, K::Real, L::Real)=alpha*z*K^(alpha-1)*L^(1-alpha)\n",
    "w(alpha::Real,z::Real,K::Real,L::Real)=(1-alpha)*z*K^(alpha)*L^(-alpha)    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `KSSolution` (mutable type)\n",
    "Collection of KS solutions. Specifically, policy function of individual agents and the coefficients of aggregate capital law of motion.\n",
    "\n",
    "##### Fields\n",
    "- `k_opt::Array{Float64,3}`: individual policy function. First dimension is individual capital, second one is aggregate capital, third one is shock state.\n",
    "- `value::Array{Float64,3}`: individual value function. First dimension is individual capital, second one is aggregate capital, third one is shock state.\n",
    "- `B::Vector{Float64}`: Coefficients on approximate aggregate capital law of motion.\n",
    "- `R2::Vector{Float64}`: R square of approximate aggregate capital law of motion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutable struct KSSolution\n",
    "    k_opt::Array{Float64,3}\n",
    "    value::Array{Float64,3}\n",
    "    B::Vector{Float64}\n",
    "    R2::Vector{Float64}\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `KSSolution` (function)\n",
    "Create `KSSolution` from `ksp::NamedTuple`. If other results are saved in `.jld2` format, you can load it as initial guess by passing `filename` and setting `load_value=true` or `load_B=true`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function KSSolution(ksp::NamedTuple;\n",
    "                    load_value::Bool=false,\n",
    "                    load_B::Bool=false,\n",
    "                    filename::String=\"result.jld2\")\n",
    "    if load_value || load_B\n",
    "        result=load(filename)\n",
    "        kss_temp=result[\"kss\"]\n",
    "    end\n",
    "    if load_value\n",
    "        k_opt=kss_temp.k_opt\n",
    "        value=kss_temp.value\n",
    "    else\n",
    "        k_opt=ksp.beta*repeat(ksp.k_grid,outer=[1,ksp.K_size,ksp.s_size])\n",
    "        k_opt=0.9*repeat(ksp.k_grid,outer=[1,ksp.K_size,ksp.s_size])\n",
    "        k_opt .= clamp.(k_opt, ksp.k_min, ksp.k_max)\n",
    "        value=ksp.u.(0.1/0.9*k_opt)/(1-ksp.beta)\n",
    "    end\n",
    "    if load_B\n",
    "        B = kss_temp.B\n",
    "    else\n",
    "        B = [0.0, 1.0, 0.0, 1.0]\n",
    "    end\n",
    "    kss = KSSolution(k_opt, value, B, [0.0, 0.0])\n",
    "    return kss\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shock Generation\n",
    "The function in this section is used to draw aggregate and idiosyncratic shocks\n",
    "### `generate_shocks`\n",
    "Generate aggregate and idiosyncratic shock\n",
    "\n",
    "##### Arguments\n",
    "- `ksp::NamedTuple` : `NamedTuple` generated by `KSParameter`\n",
    "- `z_shock_size::Integer` : size of aggregate shock\n",
    "- `population::Integer` : size idiosyncratic shock in one period"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function generate_shocks(ksp::NamedTuple;\n",
    "                         z_shock_size::Integer = 1100,\n",
    "                         population::Integer = 10000)\n",
    "\n",
    "    # unpack parameters\n",
    "    Peps_gg = ksp.transmat.Peps_gg\n",
    "    Peps_bg = ksp.transmat.Peps_bg\n",
    "    Peps_gb = ksp.transmat.Peps_gb\n",
    "    Peps_bb = ksp.transmat.Peps_bb\n",
    "    \n",
    "    # draw aggregate shock\n",
    "    zi_shock = simulate(MarkovChain(ksp.transmat.Pz), z_shock_size)\n",
    "    \n",
    "    ### Let's draw individual shock ###\n",
    "    epsi_shock = Array{Int}(undef, z_shock_size, population) # preallocation\n",
    "    \n",
    "    # first period\n",
    "    rand_draw=rand(population)\n",
    "    # recall: index 1 of eps is employed, index 2 of eps is unemployed\n",
    "    if zi_shock[1] == 1 # if good\n",
    "        epsi_shock[1, :] .= (rand_draw .< ksp.ug) .+ 1 # if draw is higher, become employed \n",
    "    elseif zi_shock[1] == 2 # if bad\n",
    "        epsi_shock[1, :] .= (rand_draw .< ksp.ub) .+ 1 # if draw is higher, become employed\n",
    "    else\n",
    "        error(\"the value of z_shocks[1] (=$(z_shocks[1])) is strange\")\n",
    "    end\n",
    "    \n",
    "    # from second period ...   \n",
    "    for t = 2:z_shock_size\n",
    "        draw_eps_shock!(Val(zi_shock[t]), Val(zi_shock[t-1]),\n",
    "                        view(epsi_shock, t, :), epsi_shock[t-1, :], ksp.transmat)\n",
    "    end\n",
    "    \n",
    "    # adjustment\n",
    "    for t=1:z_shock_size\n",
    "        n_e = count(epsi_shock[t,:].==1) # count number of employed\n",
    "        empl_rate_ideal = ifelse(zi_shock[t] == 1, 1.0-ksp.ug, 1.0-ksp.ub)\n",
    "        gap = round(Int, empl_rate_ideal*population) - n_e\n",
    "        if gap > 0\n",
    "            become_employed_i = rand(findall(2 .== epsi_shock[t,:]), gap)\n",
    "            epsi_shock[t, become_employed_i] .= 1\n",
    "        elseif gap < 0\n",
    "            become_unemployed_i = rand(findall(1 .== epsi_shock[t, :]), -gap)\n",
    "            epsi_shock[t,become_unemployed_i] .= 2\n",
    "        end \n",
    "    end\n",
    "            \n",
    "    return zi_shock, epsi_shock    \n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `draw_eps_shock!` \n",
    "Draw idiosyncratic shock given previous idiosyncratic shock and \n",
    "transition matrix.\n",
    "The transition matrix must be consistent with aggregate shock\n",
    "\n",
    "##### Arguments\n",
    "- `epsi_shocks` : preallocated vector. current period shock is stored in it\n",
    "- `epsi_shock_before` : previous period idiosyncratic shock\n",
    "- `Peps` : transition matrix of the current period\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "draw_eps_shock!(zi::Val{1}, zi_lag::Val{1}, epsi, \n",
    "                epsi_lag::AbstractVector, transmat::TransitionMatrix) = \n",
    "    draw_eps_shock!(epsi, epsi_lag, transmat.Peps_gg)\n",
    "draw_eps_shock!(zi::Val{1}, zi_lag::Val{2}, epsi, \n",
    "                epsi_lag::AbstractVector, transmat) = \n",
    "    draw_eps_shock!(epsi, epsi_lag, transmat.Peps_bg)\n",
    "draw_eps_shock!(zi::Val{2}, zi_lag::Val{1}, epsi, \n",
    "                epsi_lag::AbstractVector, transmat) = \n",
    "    draw_eps_shock!(epsi, epsi_lag, transmat.Peps_gb)\n",
    "draw_eps_shock!(zi::Val{2}, zi_lag::Val{2}, epsi, \n",
    "                epsi_lag::AbstractVector, transmat) = \n",
    "    draw_eps_shock!(epsi, epsi_lag, transmat.Peps_bb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function draw_eps_shock!(epsi_shocks,\n",
    "                         epsi_shock_before,\n",
    "                         Peps::AbstractMatrix)\n",
    "    # loop over entire population\n",
    "    for i=1:length(epsi_shocks)\n",
    "        rand_draw=rand()\n",
    "        epsi_shocks[i]=ifelse(epsi_shock_before[i] == 1,\n",
    "                              (Peps[1, 1] < rand_draw)+1,  # if employed before\n",
    "                              (Peps[2, 1] < rand_draw)+1)  # if unemployed before\n",
    "    end\n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Functions for Euler equation method\n",
    "\n",
    "The functions in this section are used to solve individual problem by Euler equation method\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `solve_ump!`\n",
    "Solve individual problem by Euler equation method, namely, iterations of the Euler equation:\n",
    "$$\\left(c\\right)^{-\\theta}=\\beta E\\left[\\left(c'\\right)^{-\\theta}(1-\\delta+r')\\right].$$\n",
    "\n",
    "##### Arguments\n",
    "- `umpsm::EulerMethod`: The field of `umpsm`, `update_k`, specifies the degree of update of policy function\n",
    "- `ksp::NamedTuple`: `NamedTuple` generated by `KSParameter`\n",
    "- `kss::KSSolution`: `KSSolution` containing the guess of ALM coefficients\n",
    "- `max_iter::Integer`: maximum number of iteration of Euler equation method\n",
    "- `tol::AbstractFloat`: tolerance of policy function convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function solve_ump!(umpsm::EulerMethod, \n",
    "                    ksp::NamedTuple, kss::KSSolution;\n",
    "                    max_iter::Integer=10000,\n",
    "                    tol::AbstractFloat=1e-8)\n",
    "    alpha, beta, delta, theta, l_bar, mu = \n",
    "        ksp.alpha, ksp.beta, ksp.delta, ksp.theta, ksp.l_bar, ksp.mu\n",
    "    k_grid, k_size = ksp.k_grid, ksp.k_size\n",
    "    K_grid, K_size = ksp.K_grid, ksp.K_size\n",
    "    s_grid, s_size = ksp.s_grid, ksp.s_size\n",
    "    k_min, k_max = ksp.k_min, ksp.k_max\n",
    "    global counter = 0\n",
    "    k_opt_n = similar(kss.k_opt)\n",
    "    prog = ProgressThresh(tol, \"Solving individual UMP by Euler method: \")\n",
    "    while true\n",
    "        global counter += 1\n",
    "        for s_i = 1:s_size\n",
    "            z, eps = s_grid[s_i, 1], s_grid[s_i, 2]\n",
    "            for (K_i, K) = enumerate(K_grid)\n",
    "                Kp, L = compute_Kp_L(K,s_i,kss.B,ksp)\n",
    "                for (k_i, k) = enumerate(k_grid)\n",
    "                    wealth = (r(alpha,z,K,L)+1-delta)*k+\n",
    "                                w(alpha,z,K,L)*(eps*l_bar + mu*(1-eps))\n",
    "                    expec=compute_expectation_FOC(kss.k_opt[k_i, K_i, s_i], Kp, s_i, ksp)\n",
    "                    cn = (beta*expec)^(-1.0/theta)\n",
    "                    k_opt_n[k_i, K_i, s_i] = wealth-cn\n",
    "                end\n",
    "            end\n",
    "        end\n",
    "        k_opt_n .= clamp.(k_opt_n, k_min, k_max)\n",
    "        dif_k = maximum(abs, k_opt_n - kss.k_opt)\n",
    "        ProgressMeter.update!(prog, dif_k)\n",
    "        if dif_k < tol\n",
    "            break\n",
    "        end\n",
    "        if counter >= max_iter\n",
    "            @warn \"Euler method failed to converge with $counter iterations (dif = $dif_k)\"\n",
    "            break\n",
    "        end\n",
    "        kss.k_opt .= umpsm.update_k*k_opt_n .+ (1-umpsm.update_k)*kss.k_opt\n",
    "    end\n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `compute_expectation_FOC`\n",
    "Compute expectation term in Euler equation: $E\\left[\\left(c'\\right)^{-\\theta}(1-\\delta+r')\\right].$\n",
    "\n",
    "##### Arguments\n",
    "- `kp::Real`: next period individual capital holding\n",
    "- `Kp::Real`: next period aggregate capital\n",
    "- `s_i::Integer`: current state of shock\n",
    "- `ksp::NamedTuple` : `NamedTuple` generated by `KSParameter`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function compute_expectation_FOC(kp::Real,\n",
    "                                 Kp::Real,\n",
    "                                 s_i::Integer,\n",
    "                                 ksp::NamedTuple)\n",
    "    alpha, theta, delta, l_bar, mu, P =\n",
    "        ksp.alpha, ksp.theta, ksp.delta, ksp.l_bar, ksp.mu, ksp.transmat.P\n",
    "    global expec = 0.0\n",
    "    for s_n_i = 1:ksp.s_size\n",
    "        zp, epsp = ksp.s_grid[s_n_i, 1], ksp.s_grid[s_n_i, 2]\n",
    "        Kpp, Lp = compute_Kp_L(Kp, s_n_i, kss.B, ksp)\n",
    "        rn = r(alpha, zp, Kp, Lp)\n",
    "        kpp = interpolate((ksp.k_grid, ksp.K_grid), kss.k_opt[:, :, s_n_i], Gridded(Linear()))\n",
    "        cp = (rn+1-delta)*kp + w(alpha, zp ,Kp, Lp)*(epsp*l_bar+mu*(1.0-epsp))-kpp(kp, Kp)\n",
    "        global expec = expec + P[s_i, s_n_i]*(cp)^(-theta)*(1-delta+rn)\n",
    "    end \n",
    "    return expec\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `compute_Kp_L`\n",
    "Compute next period aggregate capital and labor\n",
    "\n",
    "##### Arguments\n",
    "- `K::Real` : Current aggregate capital\n",
    "- `s_i::Integer` : current shock index\n",
    "- `B::AbstractVector` : coefficient of ALM for capital\n",
    "- `ksp::NamedTuple` : `NamedTuple` generated by `KSParameter`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function compute_Kp_L(K::Real, s_i::Integer,\n",
    "                      B::AbstractVector, ksp::NamedTuple)\n",
    "    Kp, L=ifelse(s_i%ksp.eps_size == 1,\n",
    "                (exp(B[1]+B[2]*log(K)), ksp.l_bar*(1-ksp.ug)), # if good\n",
    "                (exp(B[3]+B[4]*log(K)), ksp.l_bar*(1-ksp.ub))) # if bad\n",
    "    Kp = clamp(Kp, ksp.K_min, ksp.K_max)\n",
    "    return Kp, L\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bellman equation and VFI\n",
    "The functions in this cell are used to solve individual household problem by VFI\n",
    "- Methods\n",
    "    - rhs_bellman: evaluate RHS of bellman equation given the guess of value function\n",
    "    - compute_expectation_VFI: compute expectation term of Bellman equation given current shock state\n",
    "    - solve_bellman_once!: maximize RHS of Bellman equation for an agent of state (k_i,K_i,s_i)\n",
    "    - solve_bellman!: maximize RHS of Bellman equation for all state until convergence\n",
    "    - iterate_policy!: iterating policy and compute values under the policy. used for Howard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `solve_ump!`\n",
    "Solve individual problem by value function iteration, namely, iterations of the Bellman equation: \n",
    "$$V(k, K, \\epsilon, z)=u(c)+\\beta EV(k', K', \\epsilon', z').$$\n",
    "\n",
    "##### Arguments\n",
    "- `ksp` : KSParameter\n",
    "- `kss` : KSSolution\n",
    "- `tol` : tolerance of value function difference\n",
    "- `max_iter` : maximum number of iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function solve_ump!(umpsm::VFI,\n",
    "            ksp::NamedTuple,\n",
    "            kss::KSSolution;\n",
    "            max_iter::Integer=100,\n",
    "            tol::AbstractFloat=1e-8,\n",
    "            print_skip::Integer=10)\n",
    "    \n",
    "    Howard, Howard_n_iter = umpsm.Howard_on, umpsm.Howard_n_iter\n",
    "    global counter_VFI=0  # counter\n",
    "    prog = ProgressThresh(tol, \"Solving individual UMP by VFI: \")\n",
    "    while true\n",
    "        global counter_VFI += 1\n",
    "        value_old=copy(kss.value) # guessed value\n",
    "        # maximize value for all state\n",
    "        [maximize_rhs!(k_i, K_i, s_i, ksp,kss)\n",
    "          for k_i in 1:ksp.k_size, K_i in 1:ksp.K_size, s_i in 1:ksp.s_size]\n",
    "        # Howard's policy iteration\n",
    "        !Howard || iterate_policy!(ksp, kss, n_iter=Howard_n_iter)\n",
    "        # difference of guessed and new value\n",
    "        dif=maximum(abs, value_old-kss.value)\n",
    "        # progress meter of covergence process\n",
    "        ProgressMeter.update!(prog, dif)\n",
    "        # if difference is sufficiently small\n",
    "        if dif<tol\n",
    "            println(\" ** VFI converged successfully!! dif = $dif\")\n",
    "            break\n",
    "        elseif counter_VFI >= max_iter\n",
    "            println(\"VFI reached its maximum iteration : $max_iter\")\n",
    "            break\n",
    "        end\n",
    "    end\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `maximize_rhs!`\n",
    "Maximize the RHS of the Bellman equation: $u(c)+\\beta EV(k', K', \\epsilon', z').$\n",
    "\n",
    "##### Arguments\n",
    "- `k_i::Integer` : individual capital state index\n",
    "- `K_i::Integer` : aggregate capital state index\n",
    "- `s_i::Integer` : shock state index\n",
    "- `ksp` : KSParameter\n",
    "- `kss` : KSSolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function maximize_rhs!(\n",
    "            k_i::Integer,\n",
    "            K_i::Integer,\n",
    "            s_i::Integer,\n",
    "            ksp::NamedTuple,\n",
    "            kss::KSSolution,\n",
    "            )\n",
    "    # obtain minimum and maximum of grid\n",
    "    k_min, k_max = ksp.k_grid[1], ksp.k_grid[end]\n",
    "    \n",
    "    # unpack parameters\n",
    "    alpha,delta,l_bar, mu = \n",
    "        ksp.alpha, ksp.delta, ksp.l_bar, ksp.mu\n",
    "    \n",
    "    # obtain state value\n",
    "    k=ksp.k_grid[k_i]   # obtain individual capital value\n",
    "    K=ksp.K_grid[K_i]   # obtain aggregate capital value\n",
    "    z, eps = ksp.s_grid[s_i, 1], ksp.s_grid[s_i, 2]\n",
    "    Kp, L=compute_Kp_L(K,s_i,kss.B,ksp) # next aggregate capital and current aggregate labor\n",
    "    # if kp>k_c_pos, consumption is negative \n",
    "    k_c_pos=(r(alpha,z,K,L)+1-delta)*k+\n",
    "        w(alpha,z,K,L)*(eps*l_bar+(1-eps)*mu)\n",
    "    obj(kp)=-rhs_bellman(ksp,kp,kss.value,k,K,s_i) # objective function\n",
    "    res=optimize(obj, k_min, min(k_c_pos,k_max)) # maximize value\n",
    "    # obtain result\n",
    "    kss.k_opt[k_i,K_i,s_i]=Optim.minimizer(res) \n",
    "    kss.value[k_i,K_i,s_i]=-Optim.minimum(res)\n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `rhs_bellman`\n",
    "Compute the right hand side of bellman equation: $u(c)+\\beta EV(k', K', \\epsilon', z').$\n",
    "\n",
    "##### Arguments\n",
    "- `kp` : next period capital\n",
    "- `value::Array{Float64, 3}`: \n",
    "- `k::Real` : current individual capital\n",
    "- `K::Real` : current aggregate capital\n",
    "- `s_i::Integer` : "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function rhs_bellman(ksp::NamedTuple,\n",
    "            kp::Real, value::Array{Float64, 3},\n",
    "            k::Real, K::Real, s_i::Integer)\n",
    "    u,s_grid,beta,alpha,l_bar,delta, mu =\n",
    "        ksp.u, ksp.s_grid, ksp.beta, ksp.alpha, ksp.l_bar, ksp.delta, ksp.mu\n",
    "    z, eps = s_grid[s_i,1], s_grid[s_i,2]\n",
    "    Kp,L = compute_Kp_L(K,s_i,kss.B,ksp) # Next period aggregate capital and current aggregate labor\n",
    "    c = (r(alpha,z,K,L)+1-delta)*k+\n",
    "         w(alpha,z,K,L)*(eps*l_bar+(1.0-eps)*mu)-kp # current consumption \n",
    "    expec = compute_expectation(kp,Kp,value,s_i,ksp)\n",
    "    return u(c)+beta*expec\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `compute_expectation`\n",
    "Compute expectation term in the Bellman equation: $EV(k', K', \\epsilon', z').$\n",
    "\n",
    "##### Arguments\n",
    "- `kp::Real` : next period individual capital\n",
    "- `Kp::Real` : next period aggregate capital\n",
    "- `value` : given value \n",
    "- `s_i` : current shock state\n",
    "- `ksp` : KSParameter instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function compute_expectation(\n",
    "              kp::Real,  # next period indicidual capital\n",
    "              Kp::Real,  # next period aggragte capital\n",
    "              value::Array{Float64,3}, # next period value\n",
    "              s_i::Integer, # index of current state,\n",
    "              ksp::NamedTuple)\n",
    "    k_grid, K_grid = ksp.k_grid, ksp.K_grid # unpack grid\n",
    "    beta, P = ksp.beta, ksp.transmat.P      # unpack parameters\n",
    "    \n",
    "    # compute expectations by summing up\n",
    "    global expec=0.0\n",
    "    for s_n_i=1:ksp.s_size\n",
    "        value_itp=interpolate((k_grid, K_grid), value[:, :, s_n_i], Gridded(Linear()))\n",
    "        global expec += P[s_i,s_n_i]*value_itp(kp, Kp)\n",
    "    end\n",
    "    return expec\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `iterate_policy!`\n",
    "Iterate the value function with a fixed policy function and update value function.\n",
    "\n",
    "##### Arguments\n",
    "- `ksp` : KSParameter instance\n",
    "- `kss` : KSSolution instance\n",
    "- `n_iter::Integer` : number of iterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "function iterate_policy!(ksp::NamedTuple,\n",
    "            kss::KSSolution; n_iter::Integer=20)\n",
    "    value=similar(kss.value)\n",
    "    for i=1:n_iter\n",
    "        # update value using policy\n",
    "        value .= \n",
    "            [rhs_bellman(ksp,\n",
    "                kss.k_opt[k_i, K_i, s_i], kss.value,\n",
    "                ksp.k_grid[k_i], ksp.K_grid[K_i], s_i)\n",
    "                for k_i in 1:ksp.k_size,\n",
    "                    K_i in 1:ksp.K_size,\n",
    "                    s_i in 1:ksp.s_size]\n",
    "        kss.value.=copy(value)\n",
    "    end\n",
    "    return nothing\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate path of aggregate capital"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `SimulationMethod`\n",
    "Abstract type which has `Stochastic` and `NonStochastic` as subtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "abstract type SimulationMethod end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `Stochastic` (`struct`)\n",
    "- `epsi_shocks::Matrix{Int}`: individual shocks\n",
    "- `k_population::Vector{Float64}`: capital holding of population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct Stochastic <: SimulationMethod\n",
    "    epsi_shocks::Matrix{Int}\n",
    "    k_population::Vector{Float64}\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Stochastic (function)\n",
    "- `epsi_shocks::Matrix{Int}`: individual shocks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Stochastic(epsi_shocks::Matrix{Int}) = \n",
    "    Stochastic(epsi_shocks, fill(40, size(epsi_shocks, 2)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `NonStochastic` (struct)\n",
    "- `k_dens::Matrix{Float64}`: density function. first and second dimensions are respectively individual capital and employment status."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct NonStochastic <: SimulationMethod\n",
    "    k_dens::Matrix{Float64}\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `NonStochastic` (function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function NonStochastic(ksp, zi)\n",
    "    k_dens = zeros(ksp.k_size, ksp.eps_size)\n",
    "    id = findlast(ksp.k_grid .< mean(ksp.K_grid))\n",
    "    if zi == 1\n",
    "        k_dens[id, 1] = ksp.ug\n",
    "        k_dens[id, 2] = 1-ksp.ug\n",
    "    elseif zi == 2\n",
    "        k_dens[id, 1] = ksp.ub\n",
    "        k_dens[id, 2] = 1-ksp.ub\n",
    "    end\n",
    "    return NonStochastic(k_dens)\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `simulate_aggregate_path!`\n",
    "Simulate aggregate capital's path using the policy function obtained by UMP and generated aggregate and idiosyncratic shocks\n",
    "\n",
    "##### Arguments\n",
    "- `ksp::NamedTuple`: `NamedTuple` generated by `KSParameter`\n",
    "- `kss::KSSolution`: `KSSolution` containing the solution of individual UMP\n",
    "- `zi_shocks::AbstractVector`: aggregate shocks\n",
    "- `K_ts::Vector`: `Vector` to store the sequence of aggregate \n",
    "- `sm::Stochastic`: `Stochastic` which has individual shocks and initial capital holding as fields."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function simulate_aggregate_path!(ksp::NamedTuple, kss::KSSolution,\n",
    "        zi_shocks::AbstractVector, K_ts::Vector, sm::Stochastic)\n",
    "    epsi_shocks, k_population = sm.epsi_shocks, sm.k_population\n",
    "    \n",
    "    T = length(zi_shocks)   # simulated duration\n",
    "    N=size(epsi_shocks, 2) # number of agents\n",
    "    \n",
    "    # loop over T periods\n",
    "    @showprogress 0.5 \"simulating aggregate path ...\" for (t, z_i) = enumerate(zi_shocks)\n",
    "        K_ts[t] = mean(k_population) # current aggrgate capital\n",
    "        \n",
    "        # loop over individuals\n",
    "        for (i, k) in enumerate(k_population)\n",
    "            eps_i = epsi_shocks[t, i]   # idiosyncratic shock\n",
    "            s_i = epsi_zi_to_si(eps_i, z_i, ksp.z_size) # transform (z_i, eps_i) to s_i\n",
    "            # obtain next capital holding by interpolation\n",
    "            itp_pol = interpolate((ksp.k_grid, ksp.K_grid), kss.k_opt[:, :, s_i], Gridded(Linear()))\n",
    "            k_population[i] = itp_pol(k, K_ts[t])\n",
    "        end\n",
    "    end\n",
    "    return nothing\n",
    "end\n",
    "epsi_zi_to_si(eps_i::Integer, z_i::Integer, z_size::Integer) = z_i + ksp.z_size*(eps_i-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `simulate_aggregate_path!`\n",
    "simulate aggregate capital path following Young (2008).\n",
    "##### Arguments\n",
    "- `ksp::NamedTuple`: `NamedTuple` generated by `KSParameter`\n",
    "- `kss::KSSolution`: `KSSolution` containing the solution of individual UMP\n",
    "- `zi_shocks::AbstractVector`: aggregate shocks\n",
    "- `K_ts::Vector`: `Vector` to store time series of capital.\n",
    "- `sm::NonStochastic`: `NonStochastic` which has initial density as a field."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function simulate_aggregate_path!(ksp::NamedTuple, kss::KSSolution,\n",
    "        zi_shocks::AbstractVector, K_ts::Vector, sm::NonStochastic)\n",
    "    k_dens = sm.k_dens\n",
    "    @showprogress 0.5 \"simulating aggregate path ...\" for (t, z_i) = enumerate(zi_shocks[1:end-1])\n",
    "        k_dens_n = zeros(size(k_dens))\n",
    "        K_ts[t] = dot(ksp.k_grid, k_dens[:, 1]) + dot(ksp.k_grid, k_dens[:, 2])\n",
    "        Peps = get_Peps(Val(zi_shocks[t+1]), Val(z_i), ksp.transmat)\n",
    "        for eps_i in 1:ksp.eps_size\n",
    "            for (k_i, k) in enumerate(ksp.k_grid)\n",
    "                s_i = epsi_zi_to_si(eps_i, z_i, ksp.z_size)\n",
    "                itp = interpolate((ksp.k_grid, ksp.K_grid), kss.k_opt[:, :, s_i], Gridded(Linear()))\n",
    "                kp = itp(k, K_ts[t])\n",
    "                kpi_u = findfirst(kp .<= ksp.k_grid)\n",
    "                kpi_l = findlast(kp .>= ksp.k_grid)\n",
    "                weight_l = ifelse(kpi_u==kpi_l, 1,\n",
    "                            (ksp.k_grid[kpi_u]-kp)/(ksp.k_grid[kpi_u]-ksp.k_grid[kpi_l]))\n",
    "                for epsp_i in 1:ksp.eps_size\n",
    "                    k_dens_n[kpi_u, epsp_i] += (1-weight_l)*Peps[eps_i, epsp_i]*k_dens[k_i, eps_i]\n",
    "                    k_dens_n[kpi_l, epsp_i] += weight_l*Peps[eps_i, epsp_i]*k_dens[k_i, eps_i]\n",
    "                end\n",
    "            end\n",
    "        end \n",
    "        k_dens_n .= k_dens_n/sum(k_dens_n)\n",
    "        k_dens .= k_dens_n \n",
    "    end\n",
    "    K_ts[end] = dot(ksp.k_grid, k_dens[:, 1]) + dot(ksp.k_grid, k_dens[:, 2])\n",
    "    return nothing\n",
    "end\n",
    "get_Peps(zpi::Val{1}, zi::Val{1}, transmat::TransitionMatrix) = transmat.Peps_gg\n",
    "get_Peps(zpi::Val{1}, zi::Val{2}, transmat::TransitionMatrix) = transmat.Peps_bg\n",
    "get_Peps(zpi::Val{2}, zi::Val{1}, transmat::TransitionMatrix) = transmat.Peps_gb\n",
    "get_Peps(zpi::Val{2}, zi::Val{2}, transmat::TransitionMatrix) = transmat.Peps_bb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Obtaining aggregate law of motion coefficient\n",
    "\n",
    "The functions in this cell are used to obtain the coefficient of approximate aggregate capital law of motion (ALM)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `regress_ALM!`\n",
    "Obtain new aggregate law of motion coefficients using the sequence of aggregate capital \n",
    "\n",
    "##### Arguments\n",
    "- `ksp::NamedTuple`: `NamedTuple` generated by `KSParameter`\n",
    "- `kss::KSSolution`: `KSSolution` containing the solution of individual UMP\n",
    "- `zi_shocks`: aggregate shocks\n",
    "- `K_ts`: aggregate capital flaw\n",
    "- `n_discard`: number of discarded samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function regress_ALM!(ksp::NamedTuple, kss::KSSolution,\n",
    "                      zi_shocks::Vector, K_ts::Vector;\n",
    "                      T_discard::Integer=100)\n",
    "    n_g=count(zi_shocks[T_discard+1:end-1] .== 1)\n",
    "    n_b=count(zi_shocks[T_discard+1:end-1] .== 2)\n",
    "    B_n=Vector{Float64}(undef, 4)\n",
    "    x_g=Vector{Float64}(undef, n_g)\n",
    "    y_g=Vector{Float64}(undef, n_g)\n",
    "    x_b=Vector{Float64}(undef, n_b)\n",
    "    y_b=Vector{Float64}(undef, n_b)\n",
    "    global i_g = 0\n",
    "    global i_b = 0\n",
    "    for t = T_discard+1:length(zi_shocks)-1\n",
    "        if zi_shocks[t] == 1\n",
    "            global i_g = i_g+1\n",
    "            x_g[i_g] = log(K_ts[t])\n",
    "            y_g[i_g] = log(K_ts[t+1])\n",
    "        else\n",
    "            global i_b = i_b+1\n",
    "            x_b[i_b] = log(K_ts[t])\n",
    "            y_b[i_b] = log(K_ts[t+1])\n",
    "        end\n",
    "    end\n",
    "    resg = lm([ones(n_g) x_g], y_g)\n",
    "    resb = lm([ones(n_b) x_b], y_b)\n",
    "    kss.R2 = [r2(resg), r2(resb)]\n",
    "    B_n[1], B_n[2] = coef(resg)\n",
    "    B_n[3], B_n[4] = coef(resb)\n",
    "    dif_B = maximum(abs, B_n-kss.B)\n",
    "    println(\"difference of ALM coefficient is $dif_B and B = $B_n\")\n",
    "    return B_n, dif_B\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `find_ALM_coef!`\n",
    "This function finds coefficients of aggregate law of motion.\n",
    "\n",
    "The algorithm constitutes of\n",
    "1. solve individual utility maximization problem (using a guess of aggregate law of motion coefficients)\n",
    "    - `solve_ump!` takes the advantage of multiple dispatch by giving `umpsm`\n",
    "1. simulation of aggregate capital\n",
    "1. obtain coefficients by regression \n",
    "1. iterate until convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function find_ALM_coef!(umpsm::UMPSolutionMethod, sm::SimulationMethod,\n",
    "                       ksp::NamedTuple, kss::KSSolution,\n",
    "                       zi_shocks::Vector{Int};\n",
    "                       tol_ump::AbstractFloat=1e-8,\n",
    "                       max_iter_ump::Integer=100,\n",
    "                       tol_B::AbstractFloat=1e-8,\n",
    "                       max_iter_B::Integer=20,\n",
    "                       update_B::AbstractFloat=0.3,\n",
    "                       T_discard::Integer=100)\n",
    "\n",
    "    K_ts = Vector{Float64}(undef, length(zi_shocks))\n",
    "    global counter_B = 0\n",
    "    while true\n",
    "        global counter_B = counter_B+1\n",
    "        println(\" --- Iteration over ALM coefficient: $counter_B ---\")\n",
    "\n",
    "        # solve individual problem\n",
    "        solve_ump!(umpsm, ksp, kss, max_iter=max_iter_ump, tol=tol_ump)\n",
    "\n",
    "        # compute aggregate path of capital\n",
    "        simulate_aggregate_path!(ksp, kss, zi_shocks, K_ts, sm)\n",
    "        \n",
    "        # obtain new ALM coefficient by regression\n",
    "        B_n, dif_B = regress_ALM!(ksp, kss, zi_shocks, K_ts, T_discard=T_discard)\n",
    "        \n",
    "        # check convergence\n",
    "        if dif_B < tol_B\n",
    "            println(\"-----------------------------------------------------\")\n",
    "            println(\"ALM coefficient successfully converged : dif = $dif_B\")\n",
    "            println(\"-----------------------------------------------------\")\n",
    "            break\n",
    "        elseif counter_B == max_iter_B\n",
    "            println(\"----------------------------------------------------------------\")\n",
    "            println(\"Iteration over ALM coefficient reached its maximum ($max_iter_B)\")\n",
    "            println(\"----------------------------------------------------------------\")\n",
    "            break\n",
    "        end\n",
    "        \n",
    "        # Update B\n",
    "        kss.B .= update_B .* B_n .+ (1-update_B) .* kss.B\n",
    "    end\n",
    "    return K_ts\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting\n",
    "The function is used to plot the paths of true and approximate aggregate capital.\n",
    "- Methods\n",
    "    - `plot_ALM`: plot the path of true path of aggregate capital and approximated one"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `plot_ALM`\n",
    "Plot true and approximated ALM of capital \n",
    "\n",
    "##### Arguments\n",
    "- `z_grid::AbstractVector`: aggregate shock grid\n",
    "- `zi_shocks::Vector` : draw of aggregate shock\n",
    "- `B::Vector` : ALM coefficient\n",
    "- `K_ts::Vector` : actual path of capital\n",
    "- `T_discard::Integer`: number of discarded samples in the plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function plot_ALM(z_grid::AbstractVector, zi_shocks::Vector,\n",
    "                  B::Vector, K_ts::Vector;\n",
    "                  T_discard::Integer = 100)\n",
    "\n",
    "    compute_approxKprime(K, z::Val{1}, B) = exp(B[1]+B[2]*log(K))\n",
    "    compute_approxKprime(K, z::Val{2}, B) = exp(B[3]+B[4]*log(K))\n",
    "    K_ts_approx = similar(K_ts) # preallocation\n",
    "\n",
    "    # compute approximate ALM for capital\n",
    "    K_ts_approx[T_discard]=K_ts[T_discard]\n",
    "\n",
    "    for t=T_discard:length(zi_shocks)-1\n",
    "        K_ts_approx[t+1] = \n",
    "            compute_approxKprime(K_ts_approx[t], Val(zi_shocks[t]), B)\n",
    "    end\n",
    "\n",
    "    p = plot(T_discard+1:length(K_ts), K_ts[T_discard+1:end],lab=\"true\",color=:red,line=:solid)\n",
    "    plot!(p, T_discard+1:length(K_ts), K_ts_approx[T_discard+1:end],lab=\"approximation\",color=:blue,line=:dash)\n",
    "    title!(p, \"aggregate law of motion for capital\")\n",
    "    return p\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `plot_Fig1` and `plot_Fig2`\n",
    "Replicate figure 1 and 2 of the Krusell and Smith."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function plot_Fig1(ksp, kss, K_ts)\n",
    "    K_min, K_max = minimum(K_ts), maximum(K_ts)\n",
    "    K_lim = range(K_min, stop=K_max, length=100)\n",
    "    Kp_g = exp.(kss.B[1] .+ kss.B[2]*log.(K_lim))\n",
    "    Kp_b = exp.(kss.B[3] .+ kss.B[4]*log.(K_lim))\n",
    "    \n",
    "    p = plot(K_lim, Kp_g, linestyle=:solid, lab=\"Good\")\n",
    "    plot!(p, K_lim, Kp_b, linestyle=:solid, lab=\"Bad\")\n",
    "    plot!(p, K_lim, K_lim, color=:black, linestyle=:dash, lab=\"45 degree\", width=0.5)\n",
    "    title!(p, \"FIG1: Tomorrow's vs. today's aggregate capital\")\n",
    "    return p\n",
    "end\n",
    "\n",
    "function plot_Fig2(ksp, kss, K_eval_point)\n",
    "    k_lim = range(0, stop=80, length=1000)\n",
    "    itp_e = interpolate((ksp.k_grid, ksp.K_grid), kss.k_opt[:, :, 1], Gridded(Linear()))\n",
    "    itp_u = interpolate((ksp.k_grid, ksp.K_grid), kss.k_opt[:, :, 3], Gridded(Linear()))\n",
    "    \n",
    "    kp_e(k) = itp_e(k, K_eval_point)\n",
    "    kp_u(k) = itp_u(k, K_eval_point)\n",
    "    \n",
    "    p = plot(k_lim, kp_e.(k_lim), linestyle=:solid, lab=\"employed\")\n",
    "    plot!(p, k_lim, kp_u.(k_lim), linestyle=:solid, lab=\"unemployed\")\n",
    "    plot!(p, k_lim, k_lim, color=:black, linestyle=:dash, lab=\"45 degree\", width=0.5)\n",
    "    title!(p, \"FIG2: Individual policy function \\n at K=$K_eval_point when good state\")\n",
    "    return p\n",
    "end"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "kernelspec": {
   "display_name": "Julia 1.0.2",
   "language": "julia",
   "name": "julia-1.0"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.0.2"
  },
  "toc": {
   "colors": {
    "hover_highlight": "#DAA520",
    "running_highlight": "#FF0000",
    "selected_highlight": "#FFD700"
   },
   "moveMenuLeft": true,
   "nav_menu": {
    "height": "31px",
    "width": "252px"
   },
   "navigate_menu": true,
   "number_sections": true,
   "sideBar": true,
   "threshold": 4,
   "toc_cell": false,
   "toc_section_display": "block",
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}