Iterating a custom function efficiently in Julia

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

1. V is a 500x7 matrix and P a 7x7 transition matrix (i.e. each row sums one)
2. kgrid is a grid of length 500 and zgrid is a grid of length 7
3. β 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!