Stencil operations and Symbolics.jl

Hi there,

I’m writing a little package that solves classical PDEs on Cartesian grids, and am seeking advice on how to leverage Symbolics.jl to speed up operations such as matrix-vector products.

As a MWE, I’ll focus on a 2d scalar Laplacian with Dirichlet boundary conditions

using LinearAlgebra, SparseArrays
n = (3, 3)
fdm = @. spdiagm(-1 => -ones(n-1), 0 => 2ones(n), 1 => -ones(n-1))
eye = I.(n)
A = kron(eye[2], fdm[1]) + kron(fdm[2], eye[1])

What I’m trying to figure out how to get this to generate an efficient routine for matrix-vector multiplication. For now, I went along the lines of building some expressions as follows,

using Symbolics
@variables u
@syms i::Int j::Int
U = [first(@variables getindex(u, i+ii, j+jj)) for ii in -1:1 for jj in -1:1]

which for example give this away from the array edges

julia> expr = getindex(Δ * U, 5)
4.0getindex(u, i, j) - getindex(u, i, j - 1) - getindex(u, i, 1 + j) - getindex(u, 1 + i, j) - getindex(u, i - 1, j)

and then build a function and apply it,

middle = eval(build_function(expr, u, i, j))
middle(rand(n...), 2, 2)

By doing some bookkeeping for corner cases and then some metaprogramming (@generated?) I would be able to get this to work and extend to more complex non-linear operators, but I wonder if this is the right way to go…

Any pointers/help would be much appreciated!

This is actually similar to what we are doing over in DiffEqOperators, there we use lazy representations of operators and GPU convolutions from NNlib.

There is currently also an effort in Symbolics.jl to allow a very similar kind of thing, to improve performance in MethodOfLines among other things.

If you are interested in writing PDE solving code, we’d really appreciate your contributions in any of these packages.