Decide which function to use inside a loop based on a variable

New to Julia
1

I currently have some code that looks like this:

    for i = 1:Nₓ-1
        if algorithm == "2S"
           ψ[i+1, :] = T2(ψ[i, :], ω, dx, F, F̃)
        elseif algorithm == "4S"
            ψ[i+1, :] = T4S(ψ[i, :], ω, dx, F, F̃)
        elseif algorithm == "6S"
            ψ[i+1, :] = T6S(ψ[i, :], ω, dx, F, F̃)
        elseif algorithm == "8S"
            ψ[i+1, :] = T8S(ψ[i, :], ω, dx, F, F̃)
        else
            throw(ArgumentError("Algorithm type unknown, please check the documentation"))
        end
   end

Essentially the code uses the variable algorithm to decide which function to use. The problem is that this is done repeatedly inside an extremely long loop, hence degrading performance according to the profiler. It is also somewhat ugly. Is there any way around this like there is in some other languages?

EDIT:
I am updating this to show the form of F, F̃ and the function T2:

F = plan_fft!(sim.ψ[1, :]) # Plan
F̃ = plan_ifft!(sim.ψ[1, :]) # Inverse Plan

function T2(ψ, ω, dx, F, F̃)
    # Nonlinear (loop over vectorization for speed)
    for i in 1:length(ψ)
        ψ[i] *= exp(-im * dx/2 * (-1*abs2(ψ[i]))) 
    end
    # Kinetic
    F*ψ
    ψ .*= ifftshift(exp.(-im * dx * ω .^ 2 / 2)) 
    F̃*ψ

    # Nonlinear (loop over vectorization for speed)
    for i in 1:length(ψ)
        ψ[i] *= exp(-im * dx/2 * (-1*abs2(ψ[i]))) 
    end

    return ψ
end #T2
2

It’s hard to say without more context but:

  • It is quite unlikely that the string equality check has any performance impact. The allocation of ψ[i, :] should likely be much more expensive. Perhaps you can pass @view ψ[i, :] instead.
  • You can use enums instead of strings, they are cheaper to compare and their type make more sense (it makes it impossible to pass an invalid input)
    @enum Algorithms begin
    ALGORITHM_2S,
    ...
end
run(ALGORITHM_2S)
  • You can just pass the function directly as an argument. run(T2) instead of run("2S2").
4 Likes
3

I have tried moving to views but surprisingly, @time says I only have 2 allocations less and the code is actually slower by about 10%. I have updated the code above with the definition of T2 and F. When I use views it’s simply:

F = plan_fft!(sim.ψ[1, :]) # Plan
F̃ = plan_ifft!(sim.ψ[1, :]) # Inverse Plan
@views ψ[i+1, :] = T2(ψ[i, :], ω, dx, F, F̃)

I have moved my code over from strings to enums, thank you.

Could you possibly give me an example of this with T2’s arguments or point me to somewhere in the documentation? I’m not sure how to get this to work and all I see from ?run is stuff about external programs.

4

You must have somewhere a function you call like fun("2S", ψ, ω, dx), which contains this if-else block, and Kristoffer is suggesting you change this to fun(T2, ψ, ω, dx). Instead of testing the string, it would then just apply whatever function it receives.

In T2, you may want to do this:

    @inbounds for i in 1:length(ψ)
        ψ[i] *= cis(dx/2 *abs2(ψ[i])) 
    end

and to notice that exp.(-im * dx * ω .^ 2 / 2) allocates a few arrays (assuming ω is an array), you may want to pass T2 a scratch space to re-use, @. cache = cis(-dx * ω^2 / 2). And perhaps write a function like ifftshift_times!(ψ, cache)

3 Likes
5

What version of Julia are much faster in 1.5. (and they usually don’t allocate anymore)

6

What Oscar tried to say is that @view is non-allocating from Julia 1.5 on, so if you are using a previous version maybe you will not see a positive difference.

1 Like
7

(Note that this allocates multiple temporary arrays, e.g. ω .^ 2 / 2 allocates a temporary array for the / 2, and fftshift allocates another temporary array.)

4 Likes
8

Make use of multiple dispatch by creating types and structs:

# define types
abstract type Algorithm end

struct Algo2S <: Algorithm end
struct Algo4S <: Algorithm end
struct Algo6S <: Algorithm end
struct Algo8S <: Algorithm end

# then you define functions and methods

# default
run_algorithm(::Algorithm, ψ, ω, dx, F, F̃) = throw(ArgumentError("Not defined"))

run_algorithm(::Algo2S, ψ, ω, dx, F, F̃) = T2(ψ[i, :], ω, dx, F, F̃)
run_algorithm(::Algo4S, ψ, ω, dx, F, F̃) = T4S(ψ[i, :], ω, dx, F, F̃)
run_algorithm(::Algo6S, ψ, ω, dx, F, F̃) = T6S(ψ[i, :], ω, dx, F, F̃)
run_algorithm(::Algo8S, ψ, ω, dx, F, F̃) = T8S(ψ[i, :], ω, dx, F, F̃)

#lastly, in your program you just do

algorithm = Algo2S() 
for i = 1:Nₓ-1
    ψ[i+1, :] = run_algorithm(algorithm, ψ, ω, dx, F, F̃)
end

```
9

I am really not sure if this is a better way than just passing a function. If the receiving function had multiple distinct calls that change with the passed algorithm (i.e., you also wanted to query the name of the algorithm for the output, so you defined an algorithm_name too, and so on) then this would be the best solution, but seems overkill if the receiving function only needs to do a single thing differently because of the chosen internal algorithm.

5 Likes
10

@Henrique_Becker you are right!

It would be much easier to do:

#set algorithm value to the function name, e.g.:
algorithm = T4S

#then, algorithm is the function you want to call 
for i = 1:Nₓ-1
    ψ[i+1, :] = algorithm(ψ, ω, dx, F, F̃)
end
11

Thanks to everyone’s suggestions, I ended up with this:

function T2(ψ, W, dx, F, F̃)
    @inbounds for i in 1:length(ψ)
        ψ[i] *= cis(dx/2 * (-1*abs2(ψ[i]))) 
    end

    F*ψ
    @inbounds for i in 1:length(W)
        ψ[i] *= W[i]
    end
    F̃*ψ

    @inbounds for i in 1:length(ψ)
        ψ[i] *= cis(dx/2 * (-1*abs2(ψ[i]))) 
    end

    return ψ
end

function run(algorithm, ...)
    F = plan_fft!(sim.ψ[1, :]) 
    F̃ = plan_ifft!(sim.ψ[1, :]) 

    W = ifftshift(cis.(sim.box.dx * sim.box.ω .^ 2 / 2))

   for i = 1:sim.box.Nₓ-1
        ψ[i+1, :] = algorithm(ψ[i,:], W, dx, F, F̃)
   end #for

end

run(T2)

This has improved my performance by a factor of 3.5-4, which is fantastic. I would still like to get views working (currently they slow down my code on Julia 1.5), as well as multithreading the loops and FFTW (they SIGNIFICANTLY) slow down my code. I’ll ask that as a separate question.

1 Like