# [SciMLOperators] Idiomatic way to define a stateful FunctionOperator

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