In [1]:
using BenchmarkTools
using Plots
In [2]:
type Model
    # primitive parameter
    beta::Float64 #subjective discount factor
    sigma::Float64 # relative riskb aversion
    delta::Float64 #depriciation rate
    alpha::Float64 # capital share

    # discretize asset space
    agrid::Vector{Float64}
end


function argmax(mat)
    values, indices = findmax(mat,2)
    return ind2sub(size(mat),vec(indices))[2]
end


function VFI(m::Model; max_iter::Int=1000, tol::Float64=1e-5)
    const penalty = -999999999.9
    const na = size(m.agrid, 1)
   
    #initialize value function and so on
    v  = zeros(na, na) # temp value function
    c  = zeros(na, na) # consuption matrix
    util = zeros(na, na) # utility matrix
    
    v0 = zeros(na, 1) # initial guess of value function
    Tv = zeros(na, 1) # update value function
    pol_a = zeros(na, 1) # policy function
    
    #create consuption and utility matrix        
    for i in 1:na
        for j in 1:na
            @inbounds c[j,i] = m.agrid[j]^m.alpha + (1-m.delta) * m.agrid[j] - m.agrid[i]
            @inbounds util[j,i] =(c[j,i]^(1-m.sigma)) / (1-m.sigma)
            if c[j,i] <= 0
                @inbounds util[j,i] = penalty # penalty
            end
        end
    end
    
    # value function iteration
    err = 0.0
    for t in 1:max_iter
        # calculate temp value function
        for i in 1:na
            for j in 1:na
                @inbounds v[j, i] = util[j, i] + m.beta * v0[i]
            end
        end
        
        Tv = maximum(v, 2) # obtain new value funtion
        err = maximum(abs.(Tv - v0)) # update error
        v0 = Tv # update value function
        
        if err < tol
            # obtain policy function
            a_index = argmax(v)
            for i in 1:na
                @inbounds pol_a[i] = m.agrid[a_index[i]]
            end
            break
        end
    end
    
    if err >= tol
        println("VFI does not converge in $(max_iter) times")
    end
    
    return(m.agrid, v0, pol_a)
end
Out[2]:
VFI (generic function with 1 method)
In [3]:
beta = 0.95 #subjective discount factor
sigma = 2.0 # relative riskb aversion
delta = 0.1 #depriciation rate
alpha = 0.33 # capital share

# Steady state
aterm = 1.0/beta -(1.0 -delta)
kstar = alpha/aterm
kstar = kstar^(1.0/(1.0-alpha))

amin = 0.1 * kstar
amax = 2 * kstar
na   = 250

model = Model(beta, sigma, delta, alpha, linspace(amin, amax, na))
Out[3]:
Model(0.95, 2.0, 0.1, 0.33, [0.316086, 0.340205, 0.364324, 0.388443, 0.412562, 0.436681, 0.4608, 0.484919, 0.509038, 0.533157  …  6.10465, 6.12877, 6.15289, 6.17701, 6.20113, 6.22524, 6.24936, 6.27348, 6.2976, 6.32172])
In [4]:
agrid, v0, pol_a = VFI(model)
Out[4]:
([0.316086, 0.340205, 0.364324, 0.388443, 0.412562, 0.436681, 0.4608, 0.484919, 0.509038, 0.533157  …  6.10465, 6.12877, 6.15289, 6.17701, 6.20113, 6.22524, 6.24936, 6.27348, 6.2976, 6.32172], [-22.9349; -22.7696; … ; -15.6818; -15.6718], [0.484919; 0.509038; … ; 5.95994; 5.95994])
In [5]:
@benchmark VFI(model)
Out[5]:
BenchmarkTools.Trial: 
  memory estimate:  2.87 MiB
  allocs estimate:  1828
  --------------
  minimum time:     24.128 ms (0.00% GC)
  median time:      25.455 ms (0.00% GC)
  mean time:        26.098 ms (0.77% GC)
  maximum time:     33.859 ms (4.05% GC)
  --------------
  samples:          192
  evals/sample:     1
In [6]:
plot(agrid, pol_a, color="blue", linewidth=1.5, label="Policy Function")
plot!(agrid, agrid, color="red", linewidth=1.5, label="45 degree Line")