How to implement a FFT based 2D laplacian in SciMLOperators.jl?

Specific Domains Modelling & Simulations
,
1

Hello!

I was trying to implement a FFT based 2D laplacian in SciMLOperators.jl, but I ran into some trouble. Here is my attempt, which was based in this example in the docs:

using SciMLOperators
using LinearAlgebra, FFTW

function fft_laplacian(u, Lx, Ly)
    P = plan_fft(u, (1, 2))
    fwd(v, u, p, t) = P * v
    bwd(v, u, p, t) = P \ v

    fwd(w, v, u, p, t) = mul!(w, P, v)
    bwd(w, v, u, p, t) = ldiv!(w, P, v)

    Nx, Ny = size(u)
    kx = fftfreq(Nx, 2π / Lx)
    ky = fftfreq(Ny, 2π / Ly)
    minus_k2 = map(x -> -sum(abs2, x), Iterators.product(kx, ky))
    op_minus_k2 = DiagonalOperator(minus_k2)

    F = FunctionOperator(fwd, u, minus_k2;
        T=eltype(u), op_adjoint=bwd,
        op_inverse=bwd,
        op_adjoint_inverse=fwd, islinear=true
    )

    ∇² = F \ op_minus_k2 * F # This errors due to dimension mismatch

    cache_operator(∇², v)
    op_minus_k2 * F
end

u = ones(ComplexF64, 5, 5)
∇² = fft_laplacian(u, 2π / 5, 2π / 5)

The problem happens when I try to compose F (the FFT operator) with op_minus_k2, which multiplies each point in the grid with the corresponding squared frequency. I get a dimensional mismatch because F is 25 x 25 while op_minus_k2 is 5 x 5. Nonetheless, it appears to me that both operators, individually, are performing exactly the operation that I want.

What would be the proper way to implement this? Thanks in advance!