Seemingly innoculus change causes code to run 4 times slower

New to Julia
1

I have a simple script that solves an economic model. It solves my problem quickly but isn’t very readable. On making a small change to the code (to improve readability), essentially defining the simple function
rhs(i,j,value_func) leads to an almost four times slower runtime when benchmarked. So I’m obviously doing something stupid. I really want to know how I’ve managed to affect the performance so drastically. Any other advice on how to improve the code, performance or readability is always welcome. I’m afraid economists don’t make good programmers!

using Parameters
using Plots
using BenchmarkTools



Model = @with_kw (α = 0.3,
                  β = 0.99, #(1.04)^(-1/4),
                  γ = 1.0,
                  δ = 0.02,
                  k_ss = ((1/β + δ-1)*(1/α))^(1/(α-1)),
                  grid_points = 1000,
                  k_grid = range(0.9*k_ss, length = grid_points, 1.1*k_ss),
                  )



function u(c; γ = m.γ)
    if c > 0 && γ != 1
        return c^(1-γ)/(1-γ)
    elseif c > 0 && γ == 1
        return log(c)
    else
        return -Inf
    end
end

function rhs(i,j,value_func)
    @unpack α, β, δ, k_grid = m 
    return u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[j]) + β * value_func[j];
end

function faster_bellman(m, value_func, policy_func)
    @unpack α, β, δ, k_grid = m

    
    value = zeros(length(k_grid))
    policy = zeros(Int64, length(k_grid))

    js = 1
    w = zeros(3)
    
    for i in 1:length(k_grid)
        jmin = js;
        jmax = length(k_grid);

        while (jmax-jmin)>2       
            jl = floor(Int64, (jmin+jmax)/2); ju = jl+1;
            w[1] = u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[jl]) + β * value_func[jl]; #rhs(i,jl,value_func)
            w[2] = u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[ju]) + β *  value_func[ju]; #rhs(i,ju,value_func)
            if w[2] > w[1]
                jmin = jl;
            else
                jmax = ju;
            end
        end
        w[1] = u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[jmin]) + β * value_func[jmin]; #rhs(i,jmin,value_func)
        if jmax > jmin;
            w[2] = u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[jmin+1]) + β * value_func[jmin+1]; #rhs(i,jmin+1,value_func)
        else
            w[2] = w[1];
        end
        w[3] = u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[jmax]) + β * value_func[jmax]; #rhs(i,jmax,value_func)
        loc = findmax(w)[2];
        js = loc
        value[i] = w[js];
        js = jmin + js- 1;
        policy[i] = js;
    end
    return value, policy
end


function solve(m; tol = 1e-6, max_iter = 1500)
    @unpack α, β, δ, k_grid = m

    value_func = zeros(length(k_grid))
    policy_func = zeros(length(k_grid))
    error = Inf;
    iter = 1

    while error > tol && iter <= max_iter
        value_func_update, policy_func_update = faster_bellman(m, value_func, policy_func);
        error = maximum(abs.(value_func_update - value_func));

        value_func = value_func_update;
        policy_func = policy_func_update;
        iter += 1
        println("Iteration $iter, error = $error")
    end

    return value_func, policy_func
end

m = Model();


@benchmark solve(m)
2

I think the problem is that you need to pass your model m to rhs. Otherwise it is treated as a global variable. The default keyword in u might be problematic too.

using Parameters
using Plots
using BenchmarkTools



Model = @with_kw (α = 0.3,
                  β = 0.99, #(1.04)^(-1/4),
                  γ = 1.0,
                  δ = 0.02,
                  k_ss = ((1/β + δ-1)*(1/α))^(1/(α-1)),
                  grid_points = 1000,
                  k_grid = range(0.9*k_ss, length = grid_points, 1.1*k_ss),
                  )



function u(c; γ = m.γ)
    if c > 0 && γ != 1
        return c^(1-γ)/(1-γ)
    elseif c > 0 && γ == 1
        return log(c)
    else
        return -Inf
    end
end

function rhs(m, i,j,value_func)
    @unpack α, β, δ, k_grid = m 
    return u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[j]) + β * value_func[j];
end

function faster_bellman(m, value_func, policy_func)
    @unpack α, β, δ, k_grid = m

    
    value = zeros(length(k_grid))
    policy = zeros(Int64, length(k_grid))

    js = 1
    w = zeros(3)
    
    for i in 1:length(k_grid)
        jmin = js;
        jmax = length(k_grid);

        while (jmax-jmin)>2       
            jl = floor(Int64, (jmin+jmax)/2); ju = jl+1;
            w[1] = rhs(m, i,jl,value_func)
            w[2] = rhs(m, i,ju,value_func)
            if w[2] > w[1]
                jmin = jl;
            else
                jmax = ju;
            end
        end
        w[1] = rhs(m, i,jmin,value_func)
        if jmax > jmin;
            w[2] = rhs(m, i,jmin+1,value_func)
        else
            w[2] = w[1];
        end
        w[3] = rhs(m, i,jmax, value_func)
        loc = findmax(w)[2];
        js = loc
        value[i] = w[js];
        js = jmin + js- 1;
        policy[i] = js;
    end
    return value, policy
end


function solve(m; tol = 1e-6, max_iter = 1500)
    @unpack α, β, δ, k_grid = m

    value_func = zeros(length(k_grid))
    policy_func = zeros(length(k_grid))
    error = Inf;
    iter = 1

    while error > tol && iter <= max_iter
        value_func_update, policy_func_update = faster_bellman(m, value_func, policy_func);
        error = maximum(abs.(value_func_update - value_func));

        value_func = value_func_update;
        policy_func = policy_func_update;
        iter += 1
        #println("Iteration $iter, error = $error")
    end

    return value_func, policy_func
end

m = Model();


@benchmark solve(m)
2 Likes
3

I think thats exactly my problem. I meant to pass m to rhs but somehow I forgot. Sorry! Out of interest though, why would the keyword in u be problematic?

4

Keyword parameters add an extra layer of indirection and overhead. It is not much, but it is often avoided for functions that are both very fast (in which the overhead is high relative to the function time) and are called a lot (so they are responsible for a good part of the run time).

However, it seems like your code should be efficient if the call does not actually pass any keyword arguments. In other words, there should be overhead only if you actually use the keyword arguments, my base for this statements is the Julia documentation for compiler developers.

5

It is treated as a global variable because you are not directly passing it to the function. Try this:

function rhs(m, i,j,value_func)
    @unpack α, β, δ, k_grid = m 
    return u(k_grid[i]^α + (1-δ)*k_grid[i] - k_grid[j]; γ = m.γ) + β * value_func[j];
end

It runs about 3 times faster with that change.

1 Like
6

Thanks for your help, that’s a great tip. It’s a style that I use quite often so that one small change will make a big difference to my work. The code now runs sub 1 second.

1 Like
7

No problem. As the the other commenter noted, keyword arguments have small over head, but the problem in your case was that the default value was a global variable rather than a hard coded value, such as gamma = 1.0. The compiler has trouble optimizing code when a variable is in global scope. As you have seen, global variables can have huge performance hits. Not only is passing variables to functions is more efficient, but it is better practice from a coding perspective because it can avoid side effects and make debugging easier. Grouping related variables in a NamedTuple (as you did) or Struct is a great way to organize code and limit the number of arguments in a function. Another thing to avoid if possible is changing the variable type. Basically, if you follow a few rules, your code will be fast in Julia. See the performance tips for more details. Happy coding!

8

Ah, I completely missed that the keyword argument was referring to a global binding. I am so accustomed to my own code that if I look at a keyword argument accessing a field of an object I always assume the object was passed as a positional argument of the same function, in such case the code only has that small overhead when a different keyword argument is passed and no overhead when all keyword defaults are used.

I reccommed avoiding global bindings in general. But if you need to have a global, check if it can be a constant (i.e., never change) and put const in front of it, this will make a good difference to the performance. If it cannot be a constant, and you want it to be a global binding, you can yet use the following trick to improve performance:

m = const Ref(<your object construction or literal here>)

function f(...)
    ...
    # when you refer to `m` (for access or change) you need to use `[]`
    if m[].count < 0
        m[].count = 0
    end
    # you can even change the whole object, but only to
    # another object of the same type
    m[] = <your new object of the same type>
    ...
end
2 Likes
9

Thanks for the info, all very helpful. I do make a point to try to avoid global variables but I didn’t realise this was one. I thought defining gamma = 1 within Model would take care of that. The reason that I didn’t pass m to the function c was down to readability. Obviously if I knew the performance drop beforehand I wouldn’t have done it. The doctrine in economics is to have code that reads like an economics paper and maybe that’s poor advice for Julia. Basically, I’m trying to avoid functions with a list of many arguments. In the future I’m looking to use a named tuple to handle that but haven’t got there yet. Learning to walk before I run at the moment. Thanks

2 Likes
10

BTW, I think you forgot the const before m here

1 Like
11

Thanks, fixed.