3.5 Inplace computation and fft plans

To apply fft plan to an array A, we use a preallocated output array  by calling mul!(Â, plan, A). The input array A must be a complex floating-point array like the output Â. The inverse-transform is computed inplace by applying inv(P) with ldiv!(A, P, Â).

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

    dt = tf/nt

    f  = zeros(Complex{Float64},(mesh.nx,mesh.ny))
    f .= exact(0.0, mesh)
    f̂  = similar(f)

    exky = exp.( 1im*tan(dt/2) .* mesh.x  .* mesh.ky')
    ekxy = exp.(-1im*sin(dt)   .* mesh.y' .* mesh.kx )

    Px = plan_fft(f, 1)    
    Py = plan_fft(f, 2)
        
    for n = 1:nt
        
        mul!(f̂, Py, f)
        f̂ .= f̂ .* exky
        ldiv!(f, Py, f̂)
        
        mul!(f̂, Px, f)
        f̂ .= f̂ .* ekxy 
        ldiv!(f, Px, f̂)
        
        mul!(f̂, Py, f)
        f̂ .= f̂ .* exky
        ldiv!(f, Py, f̂)
        
    end

    real(f)
end
f = with_fft_plans_inplace(0.1, 1, mesh) # trigger build
@time norm(with_fft_plans_inplace(tf, nt, mesh) .- exact(tf, mesh))


CC BY-SA 4.0 Pierre Navaro