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

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
``````

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

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.

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

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

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

(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

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

`````````

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

@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
``````

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