Iterating a custom function efficiently in Julia - performance

Question

I have a operator T_ implemented quite efficiently in Julia and I want to iterate using the while loop. My operator is given by:
% parameters
β = 0.987
δ = 0.012;
% grids
Kss = 48.1905148382166
kgrid = range(0.75*Kss, stop=1.25*Kss, length=500);
zgrid = [-0.06725382459813659, -0.044835883065424395, -0.0224179415327122, 0 , 0.022417941532712187, 0.04483588306542438, 0.06725382459813657]
% auxiliary functions to build my operator
F_(z,k) = exp(z) * (k^(1/3));
u_(c) = (c^(1-2) - 1)/(1-2)
% T_operator
function T_(V, P, kgrid, zgrid, β, δ)
E = V * P'
T1 = similar(V)
for i in axes(T1, 2)
for j in axes(T1, 1)
temp = F_(zgrid[i], kgrid[j]) + (1-δ)*kgrid[j]
aux = -Inf
for l in eachindex(kgrid)
c = max(0.0, temp - kgrid[l])
aux = max(aux, u_(c) + β * E[l, i])
end
T1[j,i] = aux
end
end
return T1
end
Explaining briefly. This operator has as input
V is a 500x7 matrix and P a 7x7 transition matrix (i.e. each row sums one)
kgrid is a grid of length 500 and zgrid is a grid of length 7
β and δ particular parameters
T_ returns a T1 (500x7) matrix. More details about this operator and the correct way to run this operator can be found in this other question that I asked: Tricks to improve the performance of a cunstom function in Julia
Running this operator only once, it takes very little time, almost instantly. However, I need to iterate this operator until I get an acceptable tolerance error, but my implementation results in an inefficient process taking a long time:
max_it = 1000
it = 1
tol = 1e-3
dist = tol +1
V0 = repeat(sqrt.(a_grid), outer = [1,7]);
while it < max_it && dist > tol
TV= T_(V0,P,kgrid, zgrid, β, δ)
dist = maximum(abs.(TV - V0)) % Computing distance or error
V0 = TV % update
it = it + 1 % Updating iterations
% Some information about the state of the iteration
if rem(it, 100) == 0
println("Current iteration:")
println(it)
println("Current norm:")
println(dist)
end
end
I think a more efficient solution is to incorporate the while loop directly into the implementation of the T_ operator, but I spent the whole day trying this out and couldn't do it. Help.
UPDATE
This the MATLAB version. It is more efficient
V0 = repmat(sqrt(kgrid), 1, 7); % Concave and increasing guess
max_it = 1000;
tol = 1e-3;
%% Iteration
tic
norm = tol + 1;
it = 1;
tic;
[K, Z, new_K] = meshgrid(kgrid, zgrid, kgrid);
K = permute(K, [2, 1, 3]);
Z = permute(Z, [2, 1, 3]);
new_K = permute(new_K, [2, 1, 3]);
% Computing consumption on each possible state and choice
C = max(f(Z,K) + (1-delta)*K - new_K,0);
% All possible utilities
U = u(C);
disp('Starting value function iteration through the good and old brute force...')
while it < max_it & norm > tol
EV = V0 * P';
EV = permute(repmat(EV, 1, 1, nk), [3, 2, 1]);
H = U + beta*EV;
[TV, index] = max(H, [], 3);
it = it + 1; % Updating iterations
norm = max(max(abs(TV - V0))); % Computing error
V0 = TV;
if rem(it, 100) == 0
disp('Current iteration:')
disp(it)
disp('Current norm:')
disp(norm)
end
end
V = TV;
toc;

Answer 1

Just to get an idea of where just we're starting from, let's wrap your inital implementation in a function
function iterate_T_firstattempt(; max_it=1000, it=1, tol=1e-3, dist=tol+1)
V0 = repeat(sqrt.(kgrid), outer = [1,7]) # Assuming the `a_grid` was a typo from your comments
while it < max_it && dist > tol
TV = T_(V0, P, kgrid, zgrid, β, δ)
dist = maximum(abs.(TV - V0)) # Computing distance or error
V0 = TV # update
it += 1 # Updating iterations
# Some information about the state of the iteration
if rem(it, 100) == 0
println("Current iteration:")
println(it)
println("Current norm:")
println(dist)
end
end
end
and benchmark it with BenchmarkTools.jl
julia> #benchmark iterate_T_firstattempt()
sample with 1 evaluation.
Single result which took 7.056 s (0.00% GC) to evaluate,
with a memory estimate of 52.33 MiB, over 5875 allocations.
Oof, that's a lot of allocations. Some of these are coming from the use of global variables, others from type instability, yet others from the design of your functions. A few specific points:
The compiler's probably already making the right call, but we might as well add an #inline to your definition of u_(c) and F_(z,k) to make sure they get inlined. And why not on T_ itself too while we're at it.
You're doing a lot of indexing in the nested for loops, might as well throw an #inbounds on there given that there should be no way of getting out-of-bounds indexing.
One better: the loops in T_ look to be safely reorder-able, so we can go ahead and upgrade that #inbounds to a #turbo or #tturbo from LoopVectorization.jl for an even bigger speedup by using your CPU's SIMD instructions / Advanced Vector Extensions.
The calculation of dist = maximum(abs.(TV - V0)) involves at least two large allocations, we can avoid those with a simple mapreduce. Or to use those SIMD instructions again, vmapreduce, from LoopVectorization.jl
The line TV = T_(V0, P, kgrid, zgrid, β, δ) is also allocating, let's switch that out for an in-place version T_!.
As mentioned above, global variables are bad news. We can just move them into the function signature of iterate_T easily enough though, which should fix that problem.
While we're at it, let's also break out three-arg mul! from the LinearAlgebra stdlib for a non-allocating calculation of E = V * P'. And to get rid of one last sneaky source of type-instability (which was causing a final ~2k allocations), we should change that outer=[1,7] to outer=(1,7) -- a nice stable tuple instead of an array.
Putting it all together:
using LinearAlgebra, LoopVectorization
# parameters
β = 0.987
δ = 0.012
# grids
Kss = 48.1905148382166
kgrid = range(0.75*Kss, stop=1.25*Kss, length=500)
zgrid = [-0.06725382459813659, -0.044835883065424395, -0.0224179415327122, 0 , 0.022417941532712187, 0.04483588306542438, 0.06725382459813657]
P = rand(7,7)
P ./= sum(P,dims=2) # Rows sum to one
# auxiliary functions to build operator
#inline F_(z,k) = exp(z) * (k^(1/3))
#inline u_(c) = (c^(1-2) - 1)/(1-2)
# T_operator, in-place version
#inline function T_!(TV, E, V, P, kgrid, zgrid, β, δ)
mul!(E, V, P')
#tturbo for i in axes(TV, 2)
for j in axes(TV, 1)
temp = F_(zgrid[i], kgrid[j]) + (1-δ)*kgrid[j]
aux = -Inf
for l in eachindex(kgrid)
c = max(0.0, temp - kgrid[l])
aux = max(aux, u_(c) + β * E[l, i])
end
TV[j,i] = aux
end
end
return TV
end
function iterate_T(P, kgrid, zgrid, β, δ; max_it=1000, it=1, tol=1e-3, dist=tol+1)
V0 = repeat(sqrt.(kgrid), outer=(1,7))
# Preallocate temporary arrays
TV = similar(V0)
E = similar(V0)
# Iterate
for it = 1:max_it
# Non-allocating in-place T_!
TV = T_!(TV, E, V0, P, kgrid, zgrid, β, δ)
# Compute distance or error
dist = vmapreduce((a,b)->abs(a-b), max, TV, V0)
copyto!(V0, TV) # update
# # Some information about the state of the iteration
# if rem(it, 100) == 0
# println("Current iteration:")
# println(it)
# println("Current norm:")
# println(dist)
# end
(dist < tol) && break
end
return V0
end
we get
julia> #benchmark iterate_T($P, $kgrid, $zgrid, $β, $δ)
BenchmarkTools.Trial: 11 samples with 1 evaluation.
Range (min … max): 460.246 ms … 599.820 ms ┊ GC (min … max): 0.00% … 0.00%
Time (median): 474.826 ms ┊ GC (median): 0.00%
Time (mean ± σ): 486.661 ms ± 40.359 ms ┊ GC (mean ± σ): 0.00% ± 0.00%
█ █
█▁▁▇▁▇█▁▁▁▁▁▁▇▁▁▁▁▁▁▁▁▇▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▇ ▁
460 ms Histogram: frequency by time 600 ms <
Memory estimate: 86.42 KiB, allocs estimate: 9.
That's a bit more like it!