3.3 Vectorized version

function vectorized(tf, nt, mesh::Mesh)

    dt = tf/nt

    
    f = exact(0.0, mesh)

    exky = exp.( 1im*tan(dt/2) .* mesh.x  .* mesh.ky')
    ekxy = exp.(-1im*sin(dt)   .* mesh.y' .* mesh.kx )
    
    for n = 1:nt
        f .= real(ifft(exky .* fft(f, 2), 2)) # df / dt = -x * df / dy  in [t,t+dt/2]
        f .= real(ifft(ekxy .* fft(f, 1), 1)) # df / dt = y * df / dx in [t,t+dt]
        f .= real(ifft(exky .* fft(f, 2), 2)) # df / dt = -x * df / dy in [t+dt/2,t+dt]
    end

    f
end
nt, tf = 1000, 200π
mesh = Mesh(-π, π, 256, -π, π, 512)
f = vectorized(0.1, 1, mesh) # trigger build
@time norm(vectorized(tf, nt, mesh) .- exact(tf, mesh))
julia --check-bounds=no -O3 --track-allocation=user program.jl
        - function vectorized(tf, nt, mesh::Mesh)
        - 
234240560     dt = tf/nt
        - 
        0     f = exact(0.0, mesh)
        - 
  1048656     exky = exp.( 1im*tan(dt/2) .* mesh.x  .* mesh.ky')
  1048656     ekxy = exp.(-1im*sin(dt)   .* mesh.y' .* mesh.kx )
        -     
        0     for n = 1:nt
2621712000         f .= real(ifft(exky .* fft(f, 2), 2))
2621712000         f .= real(ifft(ekxy .* fft(f, 1), 1))
2621712000         f .= real(ifft(exky .* fft(f, 2), 2))
        -     end
        - 
        0     f
        - end


CC BY-SA 4.0 Pierre Navaro