Arellano_MK.jl

## Arellano (2008)
# refinement of QuantEcon code
# Min Kim (4/24/2021)

## Package
using QuantEcon
using Plots

##
struct Model
    β::Float64 # discount factor
    γ::Float64 # RRA
    r::Float64 # world interest rate
    ρ::Float64 # persistence of y in AR(1) 
    η::Float64 # std for the shock of y in AR(1)
    θ::Float64 # prob to regain access to financ. mkt
    ny::Int64 # # of grids for y
    nB::Int64 # # of grids for B
    Bgrid::Array{Float64}
    iBzero::Int64  # index for B = 0 or close to 0
    Π::Matrix{Float64} # ny x ny matrix. transpose is because current y is in in column 
    ygrid::Array{Float64} # in level
    ydefgrid::Array{Float64} # h(y): imposing upper bound on y as a punishment for default
    u::Function 
end

function construct_model()
    β = .953 # discount factor
    γ = 2. # RRA
    r = 0.017 # world interest rate
    ρ = 0.945 # persistence of y in AR(1) 
    η = 0.025 # std for the shock of y in AR(1)
    θ = 0.282 # prob to regain access to financ. mkt
    ny = 12 # # of grids for y
    nB = 50 # # of grids for B
    Bgrid = collect(range(-.6, .4, length = nB))
    iBzero = searchsortedfirst(Bgrid, 0.) # find index for B = 0 or close to 0
    mc = tauchen(ny, ρ, η)     # Discretize AR(1) of log(y)
    Π = mc.p # ny x ny matrix
    ygrid = exp.(mc.state_values) # in level
    ydefgrid = min.(.969 * mean(ygrid), ygrid) # h(y): imposing upper bound on y as a punishment for default

    u(c) = c.^(1 - γ) ./ (1 - γ)

    return Model(β, γ, r, ρ, η, θ, ny, nB, Bgrid, iBzero, Π, ygrid, ydefgrid, u)
end

function VFI(m::Model)
    # initialize
    vf = zeros(m.nB, m.ny) # value function 
    vd = zeros(1, m.ny) # vf when default (same for all B)
    vr = zeros(m.nB, m.ny) # vf when repay
    vfnew = copy(vf)
    vrtemp = zeros(m.nB)
    Bpolicy = zeros(m.nB, m.ny) # bond policy function 
    default_states = zeros(m.nB, m.ny) # initial default states (no default)
    Π = transpose(m.Π) # transpose is because current y is in in column 
    δ = default_states * Π
    q = (1 .- δ) ./ (1 + m.r) # initial bond price 1/(1+r) at all (B',y)

    diff = 1.0
    maxiter = 1000
    tol = 1e-8 # tolerance for convergence

    ## loop
    for iter = 1:maxiter

        # expected values
        Evf = vf * Π
        Evd = vd * Π

        for (iy, y) in enumerate(m.ygrid) # current state y

            ydef = m.ydefgrid[iy]
            vd[iy] = m.u(ydef) + m.β * (m.θ * Evf[m.iBzero, iy] + (1-m.θ) * Evd[iy]) # expected value when default

            for (iB, B) in enumerate(m.Bgrid) # current state B
                
                for (iBp, Bp) in enumerate(m.Bgrid) # next state B'
                    ctemp = y - q[iBp,iy] * Bp + B # current consumption
                    ctemp ≥ 0.0 ? # nonnegativity
                        c = ctemp : c = 0.0
                    vrtemp[iBp] = m.u(c) + m.β * Evf[iBp, iy]
                end 
            
                vr[iB, iy], idx = findmax(vrtemp) # expected value when repay
                Bpolicy[iB, iy] = m.Bgrid[idx]

                # overall value as a max
                
                if vr[iB, iy] ≥ vd[iy] 
                    vfnew[iB, iy] = vr[iB, iy] # repay
                else 
                    vfnew[iB, iy] = vd[iy] # default
                    default_states[iB, iy] = 1.0
                end
            end       
        end

        # update q
        δ =  default_states * Π 
        q = (1 .- δ) ./ (1 + m.r)
        
        diff = maximum(abs.(vfnew-vf))
        if iter % 20 == 0
            println("Iteration $iter: difference $diff")
        end

        # update value
        vf = copy(vfnew)
        
        if diff < tol
            println("Converged in iteration $iter with difference $diff")
            break
        end
    end

    return vf, q
end

## Construct and Solve
m = construct_model()
@time vf, q = VFI(m)

## Figures
# create "Y High" and "Y Low" values as 5% devs from mean
high, low = 1.05 * mean(m.ygrid), 0.95 * mean(m.ygrid)
iyhigh, iylow = searchsortedfirst(m.ygrid, high), searchsortedfirst(m.ygrid, low)

p = plot(m.Bgrid, vf[:,iylow], label = "y low")
plot!(p, m.Bgrid, vf[:,iyhigh], label = "y high")
plot!(p, title = "Value function")

idx = findall((-0.35 .≤ m.Bgrid) .& (m.Bgrid .≤ 0.0))
qlow = q[idx, iylow]
qhigh = q[idx, iyhigh]

# generate plot
p2 = plot(qlow, label = "y low")
plot!(p2, qhigh, label = "y high")
plot!(p2, title = "Bond price")

display(p)
display(p2)
savefig(p,"vf.svg")
savefig(p2,"q.svg")

a typo corrected and little improved version of QuantEcon code