[SciMLOperators] Idiomatic way to define a stateful FunctionOperator

Specific Domains Modelling & Simulations
,
1

Hi,
I’m trying to implement the 2D laplacian of a scalar field as a SciMLOperators.FunctionOperator, using centered finite differences. The op function needs to know the physical size of the grid (Lx, Ly) as well as the grid size (Nx, Ny).
Using LinearMaps, this is how I would proceed

struct LaplaceOperator <: LinearMaps.LinearMap{Float64}
    Lx::Float64
    Ly::Float64
    Nx::Int
    Ny::Int
    size::Dims{2}
    Δx²::Float64
    Δy²::Float64
    function LaplaceOperator(Lx::Float64, Ly::Float64, Nx::Int, Ny::Int)
        ncells = Nx * Ny
        return new{T}(Lx, Ly, Nx, Ny, Dims([ncells, ncells]), (Lx / Nx)^2, (Ly / Ny)^2)
    end
end

Base.size(Δ::LaplaceOperator) = Δ.size

function LinearMaps._unsafe_mul!(v, Δ::LaplaceOperator, u::AbstractVector)
    u_arr = reshape(u, Δ.Nx, Δ.Ny)
    v_arr = reshape(v, Δ.Nx, Δ.Ny)
    for j₀ ∈ 1:Δ.Ny
        j₋₁ = j₀ == 1 ? Δ.Ny : j₀ - 1
        j₊₁ = j₀ == Δ.Ny ? 1 : j₀ + 1
        for i₀ ∈ 1:Δ.Nx
            i₋₁ = i₀ == 1 ? Δ.Nx : i₀ - 1
            i₊₁ = i₀ == Δ.Nx ? 1 : i₀ + 1
            v_arr[i₀, j₀] = (
                (u_arr[i₊₁, j₀] - 2 * u_arr[i₀, j₀] + u_arr[i₋₁, j₀]) / Δ.Δx² +
                    (u_arr[i₀, j₊₁] - 2 * u_arr[i₀, j₀] + u_arr[i₀, j₋₁]) / Δ.Δy²
            )
        end
    end
    return v
end

The resulting operator would then be created as follows

Lx = 2.5
Ly = 5.0
Nx = 4
Ny = 8

Δ_ref = LaplaceOperator(Lx, Ly, Nx, Ny)

ux = [sin(π * (n + 0.5) / Nx) for n = 0:(Nx-1)]
uy = [sin(π * (n + 0.5) / Ny) for n = 0:(Ny-1)]
u = ux .* uy'
u_vec = reshape(u, :)
v_ref = reshape(Δ_ref * u_vec, Nx, Ny)

With SciMLOperators, this is what I ended up with

struct LaplaceOperatorSciML
    Lx::Float64
    Ly::Float64
    Nx::Int
    Ny::Int
    Δx²::Float64
    Δy²::Float64
    function LaplaceOperatorSciML(Lx::Float64, Ly::Float64, Nx::Int, Ny::Int)
        return new(Lx, Ly, Nx, Ny, (Lx / Nx)^2, (Ly / Ny)^2)
    end
end

function (Δ::LaplaceOperatorSciML)(v, u, p, t)
    u_arr = reshape(u, Δ.Nx, Δ.Ny)
    v_arr = reshape(v, Δ.Nx, Δ.Ny)
    for j₀ ∈ 1:Δ.Ny
        j₋₁ = j₀ == 1 ? Δ.Ny : j₀ - 1
        j₊₁ = j₀ == Δ.Ny ? 1 : j₀ + 1
        for i₀ ∈ 1:Δ.Nx
            i₋₁ = i₀ == 1 ? Δ.Nx : i₀ - 1
            i₊₁ = i₀ == Δ.Nx ? 1 : i₀ + 1
            v_arr[i₀, j₀] = (
                (u_arr[i₊₁, j₀] - 2 * u_arr[i₀, j₀] + u_arr[i₋₁, j₀]) / Δ.Δx² +
                    (u_arr[i₀, j₊₁] - 2 * u_arr[i₀, j₀] + u_arr[i₀, j₋₁]) / Δ.Δy²
            )
        end
    end
    return v
end

function laplace_operator_SciML(Lx::Float64, Ly::Float64, Nx::Int, Ny::Int)
    Δ = LaplaceOperatorSciML(Lx, Ly, Nx, Ny)
    u = zeros(Float64, Nx * Ny)
    v = similar(u)
    return FunctionOperator(Δ, u, v; isinplace=true)
end

Which is then used as follows

Δ_SciML = laplace_operator_SciML(Lx, Ly, Nx, Ny)
v_SciML = reshape(Δ_SciML * u_vec, Nx, Ny)

And it can be verified that

@assert all(isapprox.(v_SciML, v_ref, atol = 1e-15, rtol=1e-12))

I find the SciMLOperators implementation a bit awkward. Could you suggest a more idiomatic way? Thanks