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

10 posts - 7 participants

Read full topic