Using DifferentialEquations.jl with ghost cells

I’m trying to solve a few differential equations using the method of lines. I’m using a five-point laplacian and ghost cells. I will add some code below, for the case of an imaginary-time Gross-Pitaevskii equation with chemical potential.

This solver works fine with an explicit method like Tsit5(), but doesn’t with implicit solvers. One issue is that I used OffsetArrays, which the standard automatic differentiation backend doesn’t like. But even using finite differences, it is not working; I suspect because the ghost cells make the Jacobian singular.

How can this problem be avoided? I know that in this case, it is easy to omit the ghost cells and replace them by a separate stencil at the edges, but I would like to implement more complicated boundary conditions in the future, in which the ghost cell approach is most practical.

MWE:


using DifferentialEquations
using OffsetArrays
using NonlinearSolve: AutoFiniteDiff          # finite-difference Jacobian for the stiff solver
using LinearAlgebra: norm
using Printf

# ---- boundary conditions: fill the ghost layer from the interior ------------
function apply_bcs!(u, N, bc)
    if bc === :dirichlet                      # ψ = 0 on the boundary
        u[0, :]   .= 0;
        u[N+1, :] .= 0
        u[:, 0]   .= 0;
        u[:, N+1] .= 0
    elseif bc === :neumann                    # zero normal gradient: ghost = nearest interior
        u[0, :]   .= u[1, :];  u[N+1, :] .= u[N, :]
        u[:, 0]   .= u[:, 1];  u[:, N+1] .= u[:, N]
    else
        error("unknown bc $bc")
    end
    return nothing
end

# ---- RHS: imaginary-time GPE gradient flow ----------------------------------
function gpe_imag!(du, u, p, t)
    (; N, dx, g, μ, V, bc) = p
    if !(u isa OffsetArray)
        u  = OffsetArray(u,  0:N+1, 0:N+1)
        du = OffsetArray(du, 0:N+1, 0:N+1)
    end
    apply_bcs!(u, N, bc)
    @inbounds for j in 1:N, i in 1:N
        Δψ = (u[i-1, j] + u[i+1, j] + u[i, j+1] + u[i, j-1] - 4u[i, j]) / dx^2
        du[i, j] = 0.5Δψ - V[i, j] * u[i, j] - g * u[i, j]^3 + μ * u[i, j]
    end
    # du at the ghost cells is deliberately left untouched (exactly as f_steady! does).
    return nothing
end

# ghost-padded state from an interior field: indices 0:N+1, interior 1:N
function padded(interior, N)
    u = OffsetArray(zeros(N + 2, N + 2), 0:N+1, 0:N+1)
    u[1:N, 1:N] .= interior
    return u
end

# ---- problem setup (no global constants — everything goes into p) -----------
function build_problem(bc; N = 40, L = 12.0, g = 1.0, μ = 1.0, τ_F = 100.0)
    dx = L / (N - 1)
    xs = range(-L/2, L/2, length = N)                  # interior coordinates
    V  = [0.5 * (x^2 + y^2) for x in xs, y in xs]      # harmonic trap on the interior
    u0 = padded([exp(-(x^2 + y^2) / 2) for x in xs, y in xs], N)   # Gaussian seed
    p  = (; N, L, dx, g, μ, V, bc)
    return ODEProblem(gpe_imag!, u0, (0.0, τ_F), p), p
end

# ---- run one solver, report what happens ------------------------------------
function try_solver(label, alg, prob, N)
    println(label); flush(stdout)
    try
        sol = solve(prob, alg; save_everystep = false, abstol = 1e-6, reltol = 1e-6)
        uf  = sol.u[end]
        @printf("ran    retcode=%-9s  ‖ψ_final‖=%.4f\n", sol.retcode, norm(uf[1:N, 1:N]))
    catch e
        # msg = replace(sprint(showerror, e), '\n' => " ")
        # @printf("THREW  %s\n%50s%s\n", nameof(typeof(e)), "", first(msg, 200))
        println(sprint(showerror, e))
    end
    flush(stdout)
end

function run_all(bc)
    prob, p = build_problem(bc)
    @printf("\n========== bc = :%s   (N=%d, state size %d = (N+2)² ghost-padded, μ=%.1f) ==========\n",
            bc, p.N, (p.N + 2)^2, p.μ)
    println("── implicit (stiff), DEFAULT autodiff = ForwardDiff ──")
    try_solver("Rosenbrock23()  [ForwardDiff]", Rosenbrock23(), prob, p.N)
    println("\n── implicit (stiff), autodiff = AutoFiniteDiff() ──")
    try_solver("Rosenbrock23(autodiff = AutoFiniteDiff())", Rosenbrock23(autodiff = AutoFiniteDiff()), prob, p.N)
    println("\n── explicit (no Jacobian) ──")
    try_solver("Tsit5()", Tsit5(), prob, p.N)
end

run_all(:dirichlet)
run_all(:neumann)

When using the implicit method with automatic differentiation, I get the error “First call to automatic differentiation for the Jacobian
failed. This means that the user f function is not compatible
with automatic differentiation …”. Using the finite differences backend gives a warning “dt was forced below floating point epsilon” in case of dirichlet boundary conditions but seems fine here for neumann boundary conditions. The Tsit5() solver always works.