KS_nolabour.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#= Krussel-Smith using sequence-space Jacobian.\n",
    "\n",
    "Code will have three parts:\n",
    "\n",
    "(a) Decompose model into blocks. Same as KS example in Auclert et al., but with endogenous labour supply.\n",
    "\n",
    "(b) Finding steady-state values of KS (no calibration). Need to translate ks_ss into julia and add labour supply choice.\n",
    "\n",
    "(c) Accumulating Jacobians. In principle straightforward once we have done everything else. Need to get 300x300 matrices\n",
    "Jacobians of each output with respect to each input. For example: in firm block, need matrix of how \"w\" and \"r\" change\n",
    "with respect to \"Z\", \"K\", and \"L\". Store these in dictionary.\n",
    "\n",
    "Extra: depending on how things go, can add calibration step, estimation step, or non-linear dynamics. For more complicated\n",
    "models, can supply jacobian of \"hetblock\" to SHADE code directly in dictionary form and let their code do the accumulation\n",
    "automatically.\n",
    "\n",
    "Can use parameters from Krussel-Smith parameters in code. Other parameter would be elasticity \"v\" and disutility \"ϕ\" of labour.\n",
    "Can use something reasonable.\n",
    "\n",
    "=#\n",
    "\n",
    "#TODOs\n",
    "\n",
    "## have interpolation routines in-house (interest rate is sensitive to type of interpolation method used)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Defining parameters\n",
    "\n",
    "Here we define the parameter struct that will be passed to all subsequent functions. We also use the rowenhorst method to discretize the income process and construct asset grids. All parameters here are identical to those chosen by Auclert et al., in their IPYNB file \"krussel-smith\" (available on their Github page).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Param\n",
       "  apoints: Int64 500\n",
       "  amax: Float64 200.0\n",
       "  beta: Float64 0.9819527880123726\n",
       "  alpha: Float64 0.11\n",
       "  deprec: Float64 0.025\n",
       "  gamma: Float64 1.0\n",
       "  rho: Float64 0.966\n",
       "  num_states: Int64 7\n",
       "  sd: Float64 0.5\n",
       "  mean: Float64 0.0\n",
       "  uncond_sd: Float64 0.5\n",
       "  y_grid: Array{Float64}(()) fill(0.0)\n",
       "  P_trans: Array{Float64}(()) fill(0.0)\n",
       "  bc: Int64 0\n",
       "  Amat: Array{Float64}(()) fill(0.0)\n",
       "  Ymat: Array{Float64}(()) fill(0.0)\n"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Parameter struct\n",
    "\n",
    "using Parameters\n",
    "using LinearAlgebra\n",
    "using Dierckx #interpolation routines written in Fortran\n",
    "using NLsolve\n",
    "using Roots\n",
    "using Plots\n",
    "\n",
    "\n",
    "@with_kw mutable struct Param\n",
    "    apoints::Int = 500 #asset grid\n",
    "    amax::Float64 = 200  #asset max\n",
    "    beta::Float64 = 0.9819527880123726 #discount factor\n",
    "    alpha::Float64 = 0.11 #capital share\n",
    "    deprec::Float64 = 0.025 #depreciation rate\n",
    "    gamma::Float64 = 1 #elasticity of substitution\n",
    "    rho::Float64 = 0.966  #autocorr of income process\n",
    "    num_states::Int = 7 #number of states for income process\n",
    "    sd::Float64 = 0.5 #stand. dev. of deviation process\n",
    "    mean::Float64 = 0 #mean of income process\n",
    "    uncond_sd::Float64 = 0.5 # unconditional sd\n",
    "    y_grid::Array{Float64} = zeros() # grid for income process\n",
    "    P_trans::Array{Float64} = zeros() # transition matrix\n",
    "    bc::Int = 0 #borrowing constraint\n",
    "    # frisch::Float64 = 2/3 #elasticity of labour wrt wages\n",
    "    Amat::Array{Float64} = zeros() #asset grid\n",
    "    Ymat::Array{Float64}= zeros() #income grid\n",
    "    # phi::Float64 = 1 #constant in front of disutility of labour\n",
    "end\n",
    "  \n",
    "param = Param()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Param\n",
       "  apoints: Int64 500\n",
       "  amax: Float64 200.0\n",
       "  beta: Float64 0.9819527880123726\n",
       "  alpha: Float64 0.11\n",
       "  deprec: Float64 0.025\n",
       "  gamma: Float64 1.0\n",
       "  rho: Float64 0.966\n",
       "  num_states: Int64 7\n",
       "  sd: Float64 0.5\n",
       "  mean: Float64 0.0\n",
       "  uncond_sd: Float64 0.5\n",
       "  y_grid: Array{Float64}((7,)) [0.2595291268380845, 0.390378674741505, 0.5872000247124883, 0.8832548787421952, 1.3285748433063689, 1.998416489677544, 3.005979291521312]\n",
       "  P_trans: Array{Float64}((7, 7)) [0.9022379843199956 0.09361981119088468 … 8.374316586000039e-9 2.4137569000000128e-11; 0.015603301865147443 0.9035872016735433 … 4.0355125541400147e-7 1.3957194310000061e-9; … ; 1.3957194310000061e-9 4.0355125541400147e-7 … 0.9035872016735433 0.015603301865147443; 2.4137569000000125e-11 8.374316586000037e-9 … 0.09361981119088467 0.9022379843199955]\n",
       "  bc: Int64 0\n",
       "  Amat: Array{Float64}((500, 7)) [0.0 0.0 … 0.0 0.0; 0.40080160320641284 0.40080160320641284 … 0.40080160320641284 0.40080160320641284; … ; 199.59919839679358 199.59919839679358 … 199.59919839679358 199.59919839679358; 200.0 200.0 … 200.0 200.0]\n",
       "  Ymat: Array{Float64}((500, 7)) [0.2595291268380845 0.390378674741505 … 1.998416489677544 3.005979291521312; 0.2595291268380845 0.390378674741505 … 1.998416489677544 3.005979291521312; … ; 0.2595291268380845 0.390378674741505 … 1.998416489677544 3.005979291521312; 0.2595291268380845 0.390378674741505 … 1.998416489677544 3.005979291521312]\n"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#rowenhort's method to approximate AR(1) w Markov chain\n",
    "#also constructs asset grids and updates param\n",
    "\n",
    "function rowenhorst!(param::Param)\n",
    "    rho, sd, num_states, mean = param.rho, param.uncond_sd, param.num_states, param.mean\n",
    "    bc, apoints, amax = param.bc, param.apoints, param.amax\n",
    "    \n",
    "    #construct grids\n",
    "    step_r = sd*sqrt(num_states-1)\n",
    "    y_grid = -1:2/(num_states-1):1\n",
    "    y_grid = mean .+ step_r*y_grid\n",
    "    \n",
    "    #transition matrix \n",
    "    p = (rho+1)/2\n",
    "    q = p\n",
    "    \n",
    "    P_trans = [p 1-p; 1-q q] \n",
    "    \n",
    "    for i = 2:num_states -1\n",
    "        a1 = [P_trans zeros(i, 1); zeros(1, i+1)] \n",
    "        a2 = [zeros(i,1) P_trans; zeros(1, i+1)]\n",
    "        a3 = [zeros(1, i+1); P_trans zeros(i,1)]\n",
    "        a4 = [zeros(1, i+1); zeros(i,1) P_trans]\n",
    "        \n",
    "        P_trans = p*a1 + (1-p)*a2 + (1-q)*a3 + q*a4\n",
    "        P_trans[2:i, :] = P_trans[2:i, :]/2\n",
    "    end\n",
    "    \n",
    "    for i=1:num_states\n",
    "       P_trans[i,:] = P_trans[i,:]/sum(P_trans[i,:])\n",
    "    end\n",
    "\n",
    "    #get stationary distribution to normalize effective labour to L=1\n",
    "    pi = eigvecs(P_trans')[:,num_states]\n",
    "    \n",
    "    #normalize pi\n",
    "    pi = pi./sum(pi)\n",
    "    \n",
    "    #exponentiate\n",
    "    y_grid = exp.(y_grid)\n",
    "    \n",
    "    #normalize\n",
    "    y_grid = y_grid/sum(pi.*y_grid)\n",
    "    \n",
    "    param.y_grid = y_grid\n",
    "    param.P_trans = P_trans\n",
    "    \n",
    "    # construct asset grids\n",
    "    Amat = [i for i in range(bc, length = apoints, stop= amax), j in 1:length(y_grid)]\n",
    "    Ymat = [j for i=1:apoints, j in y_grid]\n",
    "    \n",
    "    param.Amat = Amat\n",
    "    param.Ymat = Ymat\n",
    "    \n",
    "    return param\n",
    "end\n",
    "\n",
    "param = rowenhorst!(param)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2 Function definitions to iterate value function (using EGM), construct transition matrix (Young's method), and get stationary distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "get_C (generic function with 1 method)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#marginal utilities and their inverses\n",
    "up(c, gamma) = c.^(gamma)\n",
    "invup(x, gamma) = x.^(-1/gamma)\n",
    "\n",
    "#vp(h) = param.phi.*h.^(1/param.frisch)\n",
    "#invvp(x) = (x./param.phi).^(param.frisch)\n",
    "    \n",
    "#define optimal labour supply\n",
    "    \n",
    "# function h_opt(c_curr,Ymat,w)\n",
    "#    u_c = up(c_curr)\n",
    "#    return invvp(u_c.*Ymat.*w)\n",
    "# end\n",
    "\n",
    "#define euler equation iteration\n",
    "function c_curr(beta,P_trans,gamma;r_future,c_next)\n",
    "    return invup(beta.*(1+r_future).*up(c_next,gamma)*P_trans',gamma)\n",
    "end\n",
    "\n",
    "#obtain current assets given consumption today defined on asset grid tomorrow\n",
    "function a_curr(;r_curr,w,c_curr,Amat,Ymat)\n",
    "    return 1/(1+r_curr).*(c_curr.+Amat.-w.*Ymat)\n",
    "end\n",
    "\n",
    "#find c at borrowing constraint \n",
    "\n",
    "# function c_binding(r,w,Amat,Ymat,bc)\n",
    "    \n",
    "#     #initial guess\n",
    "    \n",
    "#     c_guess = Amat.+Ymat.*w\n",
    "    \n",
    "#     #variable to be returned\n",
    "    \n",
    "#     c_bind = similar(Amat)\n",
    "    \n",
    "#     for i=1: length(Amat[1,:])\n",
    "#         for j = 1:length(Amat[:,1])\n",
    "\n",
    "#          g(c) = (1+r)*Amat[j,i] .+ Ymat[j,i].*w .- c .+bc\n",
    "    \n",
    "#         c_bind[j,i] = find_zero(g, c_guess[j,i])\n",
    "    \n",
    "#         end\n",
    "#     end\n",
    "    \n",
    "#     return c_bind\n",
    "# end\n",
    "\n",
    "# function get_capitaldis(;apoints::Int,num_states::Int,dist::Array{Float64}) \n",
    "#     #input is stationary distribution of dimension apoints*num_states\n",
    "#     #first apoints elements correspond to mass of \n",
    "#     #returns capital distribution (vector of size apoints)\n",
    "#     capital_dist = zeros(num_states*apoints)\n",
    "#     for i= 1:param.apoints\n",
    "#     capital_dist[i]=sum(dist[i:apoints:end]);\n",
    "#     end\n",
    "#     return capital_dist\n",
    "# end\n",
    "\n",
    "#get total EFFECTIVE labor\n",
    "\n",
    "# function get_L(;n_policy, apoints,num_states,Ymat,dist)\n",
    "#     #inputs are apoints x num_states policy matrix (policy function for labour) and dist - a stationary distribution\n",
    "#     # of dimension apoints*num_states.\n",
    "#     #output is total effective labour.\n",
    "#     L = reshape(n_policy, (apoints*num_states,1))\n",
    "#     L = L.*dist.*reshape(Ymat,(apoints*num_states,1))      \n",
    "#     L = sum(L)\n",
    "#     return L      \n",
    "# end\n",
    "\n",
    "#get capital of future period i.e. get assets saved in CURRENT period given distribution TODAY\n",
    "function get_K(;a_policy, dist)\n",
    "    K= reshape(a_policy, (length(dist),1))\n",
    "    K= sum(K.*dist)\n",
    "    return K\n",
    "end\n",
    "\n",
    "function get_C(;c_policy, dist)\n",
    "    C = reshape(c_policy, (length(dist),1)) \n",
    "    C = sum(C.*dist)\n",
    "    return C\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "egm_iterate (generic function with 1 method)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Note: function takes both current and future interest rates as arguments.\n",
    "\n",
    "function egm_iterate(param::Param;w::Float64,c_next::Array{Float64},c_bind::Array{Float64},r_curr::Float64,r_future::Float64)\n",
    "# Inputs: param: param struct\n",
    "#         w: wage rate\n",
    "#         c_next: consumption next period\n",
    "#         c_bind: consumption when borrowing constraint binds (to reduce computations)\n",
    "#         r_curr: current real interest rate\n",
    "#         r_future: interest rate in future period\n",
    "# Returns: c_it: na*ny consumption array\n",
    "#         a_it: na*ny savings array\n",
    "        \n",
    "#defining some parameters\n",
    "bc = param.bc\n",
    "apoints = param.apoints\n",
    "amax = param.amax\n",
    "y_grid = param.y_grid\n",
    "gamma = param.gamma\n",
    "beta = param.beta\n",
    "P_trans = param.P_trans\n",
    "#frisch = param.frisch\n",
    "Amat = param.Amat\n",
    "Ymat = param.Ymat\n",
    "#phi = param.phi\n",
    "\n",
    "#current policy functions on current grid\n",
    "c_it = c_curr(beta,P_trans,gamma; r_future=r_future,c_next=c_next)\n",
    "a_it = a_curr(r_curr=r_curr,w=w,c_curr=c_it,Amat=Amat,Ymat=Ymat)\n",
    "\n",
    "c_nonbinding = similar(Amat)\n",
    "    \n",
    "#get consumption policy function for current grid\n",
    "for i=1:length(y_grid)\n",
    "    c_nonbinding[:,i] = Spline1D(a_it[:,i],c_it[:,i],bc=\"extrapolate\",k=1)(Amat[:,i])\n",
    "end\n",
    "\n",
    "#update elements for which borrowing constraint does bind\n",
    "for j = 1:length(y_grid)\n",
    "   c_it[:,j] = (Amat[:,j].>a_it[1,j]).*c_nonbinding[:,j] .+ (Amat[:,j].<=a_it[1,j]).*c_bind[:,j]\n",
    "end\n",
    "    \n",
    "# # return labour function\n",
    "# h_it = h_opt(c_it,Ymat,w)\n",
    "  \n",
    "#return asset function    \n",
    "a_it = @. (1+r_curr)*Amat + w*Ymat - c_it\n",
    "\n",
    "return c_it, a_it\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "get_trans (generic function with 1 method)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#define transition function using Young's method\n",
    "function get_trans(param::Param; policyfun::Array{Float64})\n",
    "# Inputs: param: param struct\n",
    "#         policyfun: savings function, array na*ny\n",
    "# Returns: transition matrix: na*ny, na*ny\n",
    "\n",
    "    #some parameters    \n",
    "Amat = param.Amat\n",
    "num_states = param.num_states\n",
    "apoints = param.apoints\n",
    "P_trans = param.P_trans\n",
    "    \n",
    "#construct transition matrix without labour stochasticity\n",
    "Q = zeros(apoints, apoints, num_states)\n",
    "    \n",
    "#finds which element to interpolate to\n",
    "findnearest(A::AbstractArray,t) = findmin(abs.(A.-t))[2]\n",
    "\n",
    "next = similar(Amat)\n",
    "previous = similar(Amat)\n",
    "\n",
    "for j=1:num_states\n",
    "for k = 1:apoints\n",
    "    \n",
    "ind = findnearest(Amat[:,j],policyfun[k,j])\n",
    "    \n",
    "if policyfun[k,j] - Amat[ind,j] >0\n",
    "    previous[k,j] = Amat[ind,j]\n",
    "    if ind==apoints\n",
    "        next[k,j] = Amat[ind,j]\n",
    "    else         \n",
    "     next[k,j] = Amat[ind+1,j]\n",
    "    end\n",
    "            \n",
    "elseif policyfun[k,j] == Amat[ind,j]\n",
    "    previous[k,j]=policyfun[k,j] \n",
    "    next[k,j] = policyfun[k,]\n",
    "            \n",
    "else\n",
    "    next[k,j] = Amat[ind,j]\n",
    "    if ind==1\n",
    "        previous[k,j] = Amat[ind,j]\n",
    "    else\n",
    "        previous[k,j] = Amat[ind-1,j]\n",
    "    end\n",
    "    \n",
    "end  \n",
    "end\n",
    "end\n",
    "\n",
    "for k = 1:apoints\n",
    "    for j = 1:num_states\n",
    "        if next[k,j] == previous[k,j]\n",
    "            Q[Amat[:,1].==previous[k,j],k,j] .= 1\n",
    "        else\n",
    "            Q[Amat[:,1].==previous[k,j],k,j] .= (next[k,j]-policyfun[k,j])./(next[k,j]-previous[k,j])\n",
    "            Q[Amat[:,1].==next[k,j],k,j] .= 1 - (next[k,j]-policyfun[k,j])./(next[k,j]-previous[k,j])\n",
    "    \n",
    "        end\n",
    "    end\n",
    "end\n",
    "\n",
    "#construct matrix with labour stochasticity\n",
    "young_trans = zeros(apoints*num_states,apoints*num_states)\n",
    "\n",
    "for j=1:num_states\n",
    "for i=1:num_states\n",
    "young_trans[(i-1)*apoints+1:i*apoints,(j-1)*apoints+1:j*apoints]=Q[:,:,j].*P_trans[j,i];\n",
    "end\n",
    "end\n",
    "\n",
    "young_trans = young_trans'\n",
    "\n",
    "return young_trans\n",
    "    \n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "get_stat (generic function with 1 method)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# get the same results as auclert et al. if we relax tolerance to 1e-6\n",
    "\n",
    "function get_stat(param::Param;r, w, tol=1e-8, max_iter=1000)\n",
    "# Inputs: param: param struct\n",
    "#         r: steady-state interest rate\n",
    "#         w: steady-state wage rate\n",
    "# Returns: K: aggregate capital\n",
    "#          c_it: consumption function, array\n",
    "#          init: steady-state distribution vector na*ny\n",
    "#          a_it: savings function, array\n",
    "    \n",
    "    \n",
    "    num_states = param.num_states\n",
    "    apoints = param.apoints\n",
    "    Amat = param.Amat\n",
    "    Ymat = param.Ymat\n",
    "    bc = param.bc\n",
    "    \n",
    "    #initial guess for policy function iteration\n",
    "    c_next = @. r*param.Amat+param.Ymat*w\n",
    "    \n",
    "    #get binding consumption\n",
    "    c_bind = @. (1+r)*Amat + w*Ymat + bc\n",
    "    \n",
    "    #initial counters\n",
    "    dist1=1\n",
    "    iter=1\n",
    "    \n",
    "    while dist1>tol && iter<max_iter\n",
    "        c_it,_ = egm_iterate(param, w=w, r_curr=r, r_future = r, c_next=c_next, c_bind=c_bind)\n",
    "        dist1 = norm(c_it - c_next, Inf)\n",
    "        c_next = c_it\n",
    "        iter = iter+1\n",
    "    end\n",
    "    \n",
    "    #get policy functions after convergence\n",
    "    c_it,a_it = egm_iterate(param, w=w, r_curr=r,r_future=r, c_next = c_next,c_bind=c_bind)\n",
    "    \n",
    "    #initial guess for transition iteration\n",
    "    init = ones(apoints*num_states)\n",
    "    init = init./sum(init)\n",
    "    \n",
    "    #iterate\n",
    "    trans = get_trans(param,policyfun=a_it)\n",
    "\n",
    "    #initial counters\n",
    "    dist2 = 1\n",
    "    iter=1\n",
    "    \n",
    "    while dist2>tol/100 && iter< max_iter\n",
    "        next_init = (init'*trans)'\n",
    "        dist2 = norm(next_init - init,Inf)\n",
    "        init = next_init\n",
    "        iter = iter+1   \n",
    "    end\n",
    "    \n",
    "    #construct total effective labour\n",
    "    # L = get_L(n_policy = h_it, apoints=apoints, num_states=num_states,Ymat=Ymat,dist=init)\n",
    "    \n",
    "    #construct total assets\n",
    "    K = get_K(a_policy = a_it, dist = init)\n",
    "    \n",
    "    return  K, c_it, init, a_it\n",
    "    \n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3 Getting the steady-state\n",
    "\n",
    "We use the previous functions to get the steady-state interest rate that clears the capital markets. Our results are very similar to those of Auclert et al., but the steady-state interest rate differs by about 25 basis points. This is because we have a stricter tolerance for convergence than Auclert et al. for the consumption policy function (1e-8). If we relax the tolerance to (1e-6), we get a steady-state interest rate very close to one, which is what $\\beta$ is calibrated to target (in their example)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "Error on entry, no approximation returned. The following conditions\nmust hold:\n1<=k<=5\nx[1] < x[2] < ... < x[end]\nw[i] > 0.0 for all i\n\nAdditionally, if spline knots are given:\nlength(xknots) <= length(x) + k + 1\nx[1] < xknots[1] < xknots[k+2] < ... < xknots[end] < x[end]\nThe schoenberg-whitney conditions: there must be a subset of data points\nxx[j] such that t[j] < xx[j] < t[j+k+1] for j=1,2,...,n-k-1",
     "output_type": "error",
     "traceback": [
      "Error on entry, no approximation returned. The following conditions\nmust hold:\n1<=k<=5\nx[1] < x[2] < ... < x[end]\nw[i] > 0.0 for all i\n\nAdditionally, if spline knots are given:\nlength(xknots) <= length(x) + k + 1\nx[1] < xknots[1] < xknots[k+2] < ... < xknots[end] < x[end]\nThe schoenberg-whitney conditions: there must be a subset of data points\nxx[j] such that t[j] < xx[j] < t[j+k+1] for j=1,2,...,n-k-1",
      "",
      "Stacktrace:",
      " [1] error(::String) at .\\error.jl:33",
      " [2] Spline1D(::Array{Float64,1}, ::Array{Float64,1}; w::Array{Float64,1}, k::Int64, s::Float64, bc::String, periodic::Bool) at C:\\Users\\gniko\\.julia\\packages\\Dierckx\\0w6mc\\src\\Dierckx.jl:175",
      " [3] egm_iterate(::Param; w::Float64, c_next::Array{Float64,2}, c_bind::Array{Float64,2}, r_curr::Float64, r_future::Float64) at .\\In[25]:34",
      " [4] get_stat(::Param; r::Float64, w::Float64, tol::Float64, max_iter::Int64) at .\\In[7]:0",
      " [5] hh_ss!(::Param; r_guess::Float64, max_iter::Int64, tol::Float64) at .\\In[26]:19",
      " [6] hh_ss! at .\\In[26]:4 [inlined]",
      " [7] top-level scope at .\\timing.jl:174 [inlined]",
      " [8] top-level scope at .\\In[26]:0",
      " [9] include_string(::Function, ::Module, ::String, ::String) at .\\loading.jl:1091",
      " [10] execute_code(::String, ::String) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:27",
      " [11] execute_request(::ZMQ.Socket, ::IJulia.Msg) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:86",
      " [12] #invokelatest#1 at .\\essentials.jl:710 [inlined]",
      " [13] invokelatest at .\\essentials.jl:709 [inlined]",
      " [14] eventloop(::ZMQ.Socket) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\eventloop.jl:8",
      " [15] (::IJulia.var\"#15#18\")() at .\\task.jl:356"
     ]
    }
   ],
   "source": [
    "#note: r=0.0466024 is a good initial guess for algorithm\n",
    "\n",
    "function hh_ss!(param::Param;r_guess = 0.01,max_iter=100,tol=1e-6)\n",
    "    alpha = param.alpha\n",
    "    deprec = param.deprec\n",
    "    dist = 1\n",
    "    iter=1\n",
    "    \n",
    "    r_guess = r_guess\n",
    "    \n",
    "    while dist>tol && iter<max_iter\n",
    "    if r_guess>=1/(param.beta)-1\n",
    "        error(\"r too large for convergence\")\n",
    "    end\n",
    "    \n",
    "    Z = ((r_guess + deprec)/alpha)^alpha #normalize so Y=1\n",
    "    w = (1-alpha)*Z*(alpha*Z/(r_guess+deprec))^(alpha/(1-alpha))\n",
    "    \n",
    "    K,_ = get_stat(param,r=r_guess,w=w)\n",
    "    r_supply = Z*alpha*(1/K)^(1-alpha) - deprec\n",
    "    \n",
    "    # w_supply = (1-alpha)*(alpha/(r_supply+deprec))^(alpha/(1-alpha))\n",
    "    dist = abs(r_guess-r_supply)\n",
    "    r_guess= 0.9*r_guess+0.1*r_supply\n",
    "\n",
    "    iter = iter +1\n",
    "    end\n",
    "    \n",
    "    #get final results\n",
    "    Z = ((r_guess + deprec)/alpha)^alpha\n",
    "    w = (1-alpha)*Z*(alpha*Z/(r_guess+deprec))^(alpha/(1-alpha))\n",
    "    \n",
    "    K,c_ss,init,policyfun = get_stat(param,r=r_guess,w=w)\n",
    "    \n",
    "    return r_guess,K,c_ss,init,policyfun\n",
    "    \n",
    "end\n",
    "\n",
    "#get steady-state for jacobian later\n",
    "\n",
    "@time rss, Kss, css, dis, policyfun = hh_ss!(param)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "UndefVarError: rss not defined",
     "output_type": "error",
     "traceback": [
      "UndefVarError: rss not defined",
      "",
      "Stacktrace:",
      " [1] top-level scope at In[9]:6",
      " [2] include_string(::Function, ::Module, ::String, ::String) at .\\loading.jl:1091",
      " [3] execute_code(::String, ::String) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:27",
      " [4] execute_request(::ZMQ.Socket, ::IJulia.Msg) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:86",
      " [5] #invokelatest#1 at .\\essentials.jl:710 [inlined]",
      " [6] invokelatest at .\\essentials.jl:709 [inlined]",
      " [7] eventloop(::ZMQ.Socket) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\eventloop.jl:8",
      " [8] (::IJulia.var\"#15#18\")() at .\\task.jl:356"
     ]
    }
   ],
   "source": [
    "#check walras law\n",
    "\n",
    "let\n",
    "deprec = param.deprec\n",
    "alpha = param.alpha\n",
    "Z = ((rss + deprec)/alpha)^alpha\n",
    "\n",
    "Z*Kss^alpha - get_C(c_policy = css, dist = dis) - deprec*Kss\n",
    "\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "get_jac (generic function with 1 method)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "function expect_vec(param, policyfun, T; output)\n",
    "    #inputs are policyfunction (to construct transition matrix) and\n",
    "    # output vector with which transition matrix is MULTIPLIED. This\n",
    "    # is the policyfun if one wants assets, css for consumption.\n",
    "    \n",
    "    expv = Array{Array{Float64,1},1}(undef,T-1)\n",
    "\n",
    "    trans= get_trans(param, policyfun = policyfun)\n",
    "    vec = reshape(output, length(trans[:,1]))\n",
    "    \n",
    "    expv[1] = vec\n",
    "    \n",
    "    for i=2:T-1\n",
    "       expv[i] = (trans)*expv[i-1]\n",
    "    end\n",
    "    \n",
    "    return expv\n",
    "\n",
    "end\n",
    "\n",
    "function fake_news(T;expv::Array{Array{Float64,1},1},dD::Array{Array{Float64,1},1}, dY::Array{Float64,1})\n",
    "    F = zeros(T,T)\n",
    "    \n",
    "    for j=1:T\n",
    "        for i=1:T\n",
    "            if i == 1\n",
    "                F[i,j] = dY[j]\n",
    "            else\n",
    "                F[i,j] = expv[i-1]'*dD[j]\n",
    "            end\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    return F\n",
    "end\n",
    "\n",
    "function get_jac(F::Array{Float64})\n",
    "   J = copy(F)\n",
    "    for j = 2:length(F[:,1])\n",
    "        for i = 2:length(F[1,:])\n",
    "            J[i,j] = J[i-1,j-1] + F[i,j]\n",
    "        end\n",
    "    end\n",
    "    \n",
    "    return J\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "get_dYr (generic function with 1 method)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#getting jacobian with respect to r\n",
    "\n",
    "function get_dYr(param; T, diff, rss::Float64, css::Array{Float64}, dis::Vector{Float64})\n",
    "#inputs: time frame for backwards iteration, finite difference step range, equilibrium interest rate\n",
    "# equilbrium consumption (this is what is iterated), and apoints x num_states distribution (vector)\n",
    "# returns dY_0,s of chosen outputs (here, consumption and capital). \n",
    "\n",
    "alpha = param.alpha\n",
    "deprec = param.deprec\n",
    "Amat = param.Amat\n",
    "Ymat = param.Ymat\n",
    "bc = param.bc\n",
    "apoints = param.apoints\n",
    "num_states = param.num_states\n",
    "    \n",
    "r_up = rss + diff\n",
    "r_down = rss - diff\n",
    "    \n",
    "Z = ((rss + deprec)/alpha)^alpha\n",
    "w = (1-alpha)*Z*(alpha*Z/(rss+deprec))^(alpha/(1-alpha)) #keep w constant (partial Jacobian)\n",
    "    \n",
    "#vector that holds aggregate differences\n",
    "\n",
    "dK = zeros(T)\n",
    "dC = zeros(T)\n",
    "    \n",
    "# vector of one-dimensional vectors (must be column vectors!) to hold change in distributions\n",
    "    \n",
    "dD = Array{Array{Float64,1},1}(undef,T)\n",
    "    \n",
    "#first upward iteration\n",
    "c_bind = @. (1+r_up)*Amat + w*Ymat + bc\n",
    "cup,aup = egm_iterate(param, w = w, c_next = css, c_bind = c_bind, r_curr = r_up, r_future = rss)\n",
    "\n",
    "Λ_up = get_trans(param, policyfun = aup) #get transition function (up)\n",
    "\n",
    "#first downward iteration\n",
    "c_bind = @. (1+r_down)*Amat + w*Ymat + bc    \n",
    "cdown,adown = egm_iterate(param, w = w, c_next = css, c_bind = c_bind, r_curr = r_down, r_future = rss)\n",
    "    \n",
    "Λ_down = get_trans(param, policyfun = adown) #get transition function (down)\n",
    "\n",
    "#get capital and consumption\n",
    "dK[1] = get_K(a_policy=(aup.-adown)./2, dist=dis)\n",
    "dC[1] = get_C(c_policy = (cup.-cdown)./2, dist = dis)\n",
    "\n",
    "#get change in distribution\n",
    "dD[1] = (dis'*((Λ_up .- Λ_down)./2))'\n",
    "    \n",
    "#second upward iteration \n",
    "c_bind = @. (1+rss)*Amat + w*Ymat + bc\n",
    "cup,aup = egm_iterate(param, w = w, c_next = cup, c_bind = c_bind, r_curr = rss, r_future = r_up)\n",
    "    \n",
    "Λ_up = get_trans(param, policyfun = aup) #get transition function (up)\n",
    "\n",
    "#second downward iteration\n",
    "c_bind = @. (1+rss)*Amat + w*Ymat + bc\n",
    "cdown,adown = egm_iterate(param, w = w, c_next = cdown, c_bind = c_bind, r_curr = rss, r_future = r_down)\n",
    "    \n",
    "Λ_down = get_trans(param, policyfun = adown) #get transition function (down)\n",
    "    \n",
    "#get capital and consumption\n",
    "dK[2] = get_K(a_policy=(aup.-adown)./2, dist=dis)\n",
    "dC[2] = get_C(c_policy = (cup.-cdown)./2, dist = dis)\n",
    "\n",
    "#change in distribution\n",
    "dD[2] = (dis'*((Λ_up .- Λ_down)./2))'\n",
    "\n",
    "for i=1:T-2\n",
    "    #upward iteration\n",
    "    cup,aup = egm_iterate(param, w = w, c_next = cup, c_bind = c_bind, r_curr = rss, r_future = rss)\n",
    "    Λ_up = get_trans(param, policyfun = aup)\n",
    "    \n",
    "    #downward iteration\n",
    "    cdown,adown = egm_iterate(param, w = w, c_next = cdown, c_bind = c_bind, r_curr = rss, r_future = rss)\n",
    "    Λ_down = get_trans(param, policyfun = adown) #get transition function (down)\n",
    "    \n",
    "    #get capital, consumption, and distribution\n",
    "    dK[i+2] = get_K(a_policy=(aup.-adown)./2, dist=dis)\n",
    "    dC[i+2] = get_C(c_policy = (cup.-cdown)./2, dist = dis)\n",
    "    dD[i+2] = (dis'*((Λ_up .- Λ_down)./2))'\n",
    "end     \n",
    "    \n",
    "    return dK./diff, dC./diff, dD./diff\n",
    "\n",
    "end\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "get_dYw (generic function with 1 method)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#getting jacobian with respect to w\n",
    "\n",
    "function get_dYw(param; T, diff, rss::Float64, css::Array{Float64}, dis::Vector{Float64})\n",
    "#inputs: time frame for backwards iteration, finite difference step range, equilibrium interest rate\n",
    "# equilbrium consumption (this is what is iterated), and apoints x num_states distribution (vector)\n",
    "# returns dY_0,s of chosen outputs (here, consumption and capital). \n",
    "\n",
    "alpha = param.alpha\n",
    "deprec = param.deprec\n",
    "Amat = param.Amat\n",
    "Ymat = param.Ymat\n",
    "bc = param.bc\n",
    "apoints = param.apoints\n",
    "num_states = param.num_states\n",
    "    \n",
    "r = rss\n",
    "    \n",
    "Z = ((rss + deprec)/alpha)^alpha\n",
    "w = (1-alpha)*Z*(alpha*Z/(rss+deprec))^(alpha/(1-alpha)) #keep w constant (partial Jacobian)\n",
    "    \n",
    "wup = w + diff\n",
    "wdown = w - diff\n",
    "    \n",
    "#vector that holds aggregate differences\n",
    "\n",
    "dK = zeros(T)\n",
    "dC = zeros(T)\n",
    "    \n",
    "# vector of one-dimensional vectors (must be column vectors!) to hold change in distributions\n",
    "    \n",
    "dD = Array{Array{Float64,1},1}(undef,T)\n",
    "    \n",
    "#first upward iteration\n",
    "c_bind = @. (1+r)*Amat + wup*Ymat + bc\n",
    "cup,aup = egm_iterate(param, w = wup, c_next = css, c_bind = c_bind, r_curr = r, r_future = r)\n",
    "\n",
    "Λ_up = get_trans(param, policyfun = aup) #get transition function (up)\n",
    "\n",
    "#first downward iteration\n",
    "c_bind = @. (1+r)*Amat + wdown*Ymat + bc    \n",
    "cdown,adown = egm_iterate(param, w = wdown, c_next = css, c_bind = c_bind, r_curr = r, r_future = r)\n",
    "    \n",
    "Λ_down = get_trans(param, policyfun = adown) #get transition function (down)\n",
    "\n",
    "#get capital and consumption\n",
    "dK[1] = get_K(a_policy=(aup.-adown)./2, dist=dis)\n",
    "dC[1] = get_C(c_policy = (cup.-cdown)./2, dist = dis)\n",
    "\n",
    "#get change in distribution\n",
    "dD[1] = (dis'*((Λ_up .- Λ_down)./2))'\n",
    "    \n",
    "#second upward iteration \n",
    "c_bind = @. (1+r)*Amat + w*Ymat + bc\n",
    "cup,aup = egm_iterate(param, w = w, c_next = cup, c_bind = c_bind, r_curr = r, r_future = r)\n",
    "    \n",
    "Λ_up = get_trans(param, policyfun = aup) #get transition function (up)\n",
    "\n",
    "#second downward iteration\n",
    "c_bind = @. (1+r)*Amat + w*Ymat + bc\n",
    "cdown,adown = egm_iterate(param, w = w, c_next = cdown, c_bind = c_bind, r_curr = r, r_future = r)\n",
    "    \n",
    "Λ_down = get_trans(param, policyfun = adown) #get transition function (down)\n",
    "    \n",
    "#get capital and consumption\n",
    "dK[2] = get_K(a_policy=(aup.-adown)./2, dist=dis)\n",
    "dC[2] = get_C(c_policy = (cup.-cdown)./2, dist = dis)\n",
    "\n",
    "#change in distribution\n",
    "dD[2] = (dis'*((Λ_up .- Λ_down)./2))'\n",
    "\n",
    "for i=1:T-2\n",
    "    #upward iteration\n",
    "    cup,aup = egm_iterate(param, w = w, c_next = cup, c_bind = c_bind, r_curr = r, r_future = r)\n",
    "    Λ_up = get_trans(param, policyfun = aup)\n",
    "    \n",
    "    #downward iteration\n",
    "    cdown,adown = egm_iterate(param, w = w, c_next = cdown, c_bind = c_bind, r_curr = rss, r_future = rss)\n",
    "    Λ_down = get_trans(param, policyfun = adown) #get transition function (down)\n",
    "    \n",
    "    #get capital, consumption, and distribution\n",
    "    dK[i+2] = get_K(a_policy=(aup.-adown)./2, dist=dis)\n",
    "    dC[i+2] = get_C(c_policy = (cup.-cdown)./2, dist = dis)\n",
    "    dD[i+2] = (dis'*((Λ_up .- Λ_down)./2))'\n",
    "end     \n",
    "    \n",
    "    return dK./diff, dC./diff, dD./diff\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "UndefVarError: rss not defined",
     "output_type": "error",
     "traceback": [
      "UndefVarError: rss not defined",
      "",
      "Stacktrace:",
      " [1] top-level scope at .\\timing.jl:174 [inlined]",
      " [2] top-level scope at .\\In[13]:0",
      " [3] include_string(::Function, ::Module, ::String, ::String) at .\\loading.jl:1091",
      " [4] execute_code(::String, ::String) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:27",
      " [5] execute_request(::ZMQ.Socket, ::IJulia.Msg) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:86",
      " [6] #invokelatest#1 at .\\essentials.jl:710 [inlined]",
      " [7] invokelatest at .\\essentials.jl:709 [inlined]",
      " [8] eventloop(::ZMQ.Socket) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\eventloop.jl:8",
      " [9] (::IJulia.var\"#15#18\")() at .\\task.jl:356"
     ]
    }
   ],
   "source": [
    "function het_jac(param; rss::Float64, css::Array{Float64}, policyfun::Array{Float64}, dis::Vector{Float64}, T=5, diff = 1e-5)\n",
    "    \n",
    "    #get dY_0^s and dD_1^s\n",
    "    dKr, dCr, dDr = get_dYr(param; rss=rss, css=css, dis=dis, T=T, diff=diff)\n",
    "    dKw, dCw, dDw = get_dYw(param; rss=rss, css=css, dis=dis, T=T, diff=diff)\n",
    "    \n",
    "    #get vector of expectation vectors (size T-1)\n",
    "    expv_a = expect_vec(param, policyfun, T; output = policyfun)\n",
    "    expv_c = expect_vec(param, policyfun, T; output = css)\n",
    "    \n",
    "    #construct fake news matrix\n",
    "    F_ar = fake_news(T; expv = expv_a, dY = dKr, dD = dDr)\n",
    "    F_cr = fake_news(T; expv = expv_c, dY = dCr, dD = dDr)\n",
    "    \n",
    "    F_aw = fake_news(T; expv = expv_a, dY = dKw, dD = dDw)\n",
    "    F_cw = fake_news(T; expv = expv_c, dY = dCw, dD = dDw)\n",
    "    #get jacobian\n",
    "    J_ar = get_jac(F_ar)\n",
    "    J_cr = get_jac(F_cr)\n",
    "    \n",
    "    J_aw = get_jac(F_aw)\n",
    "    J_cw = get_jac(F_cw)\n",
    "\n",
    "    #construct dictionaries\n",
    "    J = Dict()\n",
    "    \n",
    "    J_a = Dict(\"r\" => J_ar, \"w\" => J_aw)\n",
    "    J_c = Dict(\"r\" => J_cr, \"w\" => J_cw)\n",
    "    \n",
    "    J[\"A\"] = J_a\n",
    "    J[\"C\"] = J_c\n",
    "    \n",
    "    return J\n",
    "end\n",
    "\n",
    "@time J = het_jac(param; rss = rss, css = css, policyfun = policyfun, dis = dis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "UndefVarError: J not defined",
     "output_type": "error",
     "traceback": [
      "UndefVarError: J not defined",
      "",
      "Stacktrace:",
      " [1] top-level scope at In[14]:1",
      " [2] include_string(::Function, ::Module, ::String, ::String) at .\\loading.jl:1091",
      " [3] execute_code(::String, ::String) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:27",
      " [4] execute_request(::ZMQ.Socket, ::IJulia.Msg) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\execute_request.jl:86",
      " [5] #invokelatest#1 at .\\essentials.jl:710 [inlined]",
      " [6] invokelatest at .\\essentials.jl:709 [inlined]",
      " [7] eventloop(::ZMQ.Socket) at C:\\Users\\gniko\\.julia\\packages\\IJulia\\rWZ9e\\src\\eventloop.jl:8",
      " [8] (::IJulia.var\"#15#18\")() at .\\task.jl:356"
     ]
    }
   ],
   "source": [
    "J[\"A\"][\"w\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.5.2",
   "language": "julia",
   "name": "julia-1.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}