Reverse-mode AD for PDEs discretized with VoronoiFVM.jl

General Usage
1

Hi everyone,

I’m working on simulating cyclic adsorption processes (TVSA / PSA-type) using VoronoiFVM.jl, which I use to semi-discretize a set of PDEs in space. The resulting system of ODEs (or DAEs) is integrated over time to simulate column dynamics such as direct breakthrough or full cyclic steady-state operation.

I’m now interested in differentiating through this time-dependent PDE system to compute gradients of scalar performance metrics (e.g. total CO₂ outflow over a cycle) with respect to model hyperparameters such as inlet velocity (u_feed), feed temperature, outlet pressure, etc.

I understand that SciMLSensitivity.jl provides efficient adjoint and reverse-mode sensitivity algorithms for computing gradients of scalar outputs with respect to parameters in ODE/DAE systems.

My questions are:

  1. Is it possible to use reverse-mode automatic differentiation through a VoronoiFVM.jl simulation to compute gradients of integrals of the PDE solution with respect to such parameters?
  2. What is the recommended way to structure this?

Here is a minimal sketch of what I’m doing

using VoronoiFVM, DifferentialEquations

function column_model(; u_feed, T_feed, P_out)
    # PDE setup ...
    data = (params = (; u_feed=u_feed, T_feed=T_feed, P_out=P_out), )
    sys = VoronoiFVM.System(grid;
        storage, flux, reaction, source, bcondition, boutflow,
        data, outflowboundaries=[data.Γ_in, data.Γ_out])
    problem = ODEProblem(sys, inival, (0, duration))
    sol = solve(problem, FBDF(); kwargs...)
    return sol, sys
end

# Example target quantity: total outlet concentration
function performance_metric(u_feed, T_feed, P_out)
    sol, sys = column_model(; u_feed, T_feed, P_out)

    tf = TestFunctionFactory(sys)
    T = testfunction(tf, [1], [2])
    outflow = integrate(sys, T, sol; rate=false)
    total_outflow = - trapz(outflow.t, reduce(hcat, outflow.u)')

    return total_outflow[data.ic_CO2]
end

When I attempt to use reverse-mode AD via SciMLSensitivity.jl, I get the following error:

ERROR: `p` is not a SciMLStructure. This is required for adjoint sensitivity analysis.
For more information, see the documentation on SciMLStructures.jl for the definition
of the SciMLStructures interface. In particular, adjoint sensitivities only apply to `Tunable`.

As I understand it, the current SciMLSensitivity interface expects parameters to be passed as a parameter vector p threaded into the ODEProblem.
However, in my case, the parameters enter through data.params in the VoronoiFVM setup, not as a flat vector p.

For context, the ultimate goal is to enable gradient-based optimization of cyclic adsorption processes, where each cycle involves solving the PDE system multiple times under varying boundary conditions.

Any guidance on how to make this setup AD-compatible would be greatly appreciated!

2

In the moment I don’t think it works. However, for a start, I tried out some approaches with ForwardDiff, see VoronoiFVM.jl/examples/Example430_ParameterDerivativesStationary.jl at master · WIAS-PDELib/VoronoiFVM.jl · GitHub . For this case, the code is fully differentiable and switches to Sparspak as differentiable sparse matrix solver.

As it is, VoronoiFVM heavily relies on mutation, so it is probably Mooncake with DifferentiationInterface which could work. I also want to find time to to leverage the new SparseConnectivityTracer to have an alternative to the current sparse matrix build-up via local Jacobians.

Getting AD-based gradient based optimization running is also very interesting for us. I think it will take a couple of months to get to an implementation and interest from users increases the motivation to go for it.

So if you have a reasonable test example, please feel free to open a PR so we can see from there what the real issues are, and if and how they can be fixed.

2 Likes
3

Thanks a lot for your reply and reference to the example.

I’ve tried adapting that example to my transient simulation setup to run forward mode AD through the solver.

I set up forward differentiation through my simulation (8 variables, 1D column) using AutoForwardDiff(). However, I consistently get a StackOverflowError when the call reaches sol = solve(problem, solver; kwargs...). I tried reducing the problem to N = 4 spatial cells and t_end = 1s, but the error persists. Would you expect ForwardDiff to work for transient PDE problems like this, or is it inherently unsuitable because of the internal mutation and system size?

I also tried to experiment with Mooncake.jl to see if reverse-mode differentiation might work.
Even for a very simple transient advection–diffusion PDE, I encounter the following error:

ERROR: MooncakeRuleCompilationError: an error occured while Mooncake was compiling a rule to differentiate something. If the `caused by` error message below does not make it clear to you how the problem can be fixed, please open an issue at github.com/chalk-lab/Mooncake.jl describing your problem.
To replicate this error run the following:

Mooncake.build_rrule(Mooncake.MooncakeInterpreter(), Tuple{VoronoiFVM.var"##SystemState#96", ProblemData{Float64}, Vector{Float64}, Symbol, Type{VoronoiFVM.SystemState}, Type{Float64}, VoronoiFVM.System{Float64, Float64, Int32, Int64, Matrix{Int32}}}; debug_mode=false)

Here is the code to reproduce the above error:

using VoronoiFVM
using ExtendableGrids
using GridVisualize
using DifferentialEquations
using DifferentiationInterface
import Mooncake

import SciMLSensitivity as SMS

function trapz(ts, U)
    Δt = diff(ts)
    vals = (U[1:end-1, :] .+ U[2:end, :]) ./ 2
    if size(vals)[2] == 1
        return (Δt' * vals)[1, 1]
    else
        return vec(Δt' * vals)
    end
end

mutable struct ProblemData{Tv}
    D::Float64      ## Diffusion coefficient
    v::Tv           ## Velocity vector
end

# Bernoulli function used in the exponential fitting discretization
function bernoulli(x)
    if abs(x) < nextfloat(eps(typeof(x)))
        return 1
    end
    return x / (exp(x) - 1)
end

function exponential_flux!(f, u, edge, data)
    vh = project(edge, data.v)
    Bplus = data.D * bernoulli(vh / data.D)
    Bminus = data.D * bernoulli(-vh / data.D)
    f[1] = Bminus * u[1, 1] - Bplus * u[1, 2]
    return nothing
end

function outflow!(f, u, node, data)
    if node.region == 2
        f[1] = data.v * u[1]
    end
    return nothing
end

function main(; n = 10, D = 0.01, v = [1.0], tend = 100)
    # Tv = eltype(v)

    # Create a one-dimensional discretization
    h = 1.0 / n
    grid = simplexgrid(0:h:1)

    # Initialize problem parameters in data structure
    problem_data = ProblemData(D, v[1])

    sys = VoronoiFVM.System(
        grid,
        VoronoiFVM.Physics(;
            flux = exponential_flux!,
            breaction = outflow!,
            data = problem_data
        );
        # valuetype=Tv
    )

    # Add species 1 to region 1
    enable_species!(sys, 1, [1])

    # Set boundary conditions
    boundary_neumann!(sys, 1, 1, 0.0)

    # Create a solution array
    inival = unknowns(sys)
    # inival[1, :] .= map(x -> 1 - 2x, grid)
    xs = range(0.0, 1.0, length = size(inival, 2))
    inival[1, :] .= @. 1 - 2 * xs

    # Transient solution of the problem
    control = VoronoiFVM.SolverControl()
    control.Δt = 0.01 * h
    control.Δt_min = 0.01 * h
    control.Δt_max = 0.1 * tend
    tsol = solve(sys; inival, times = [0, tend], control)

    # problem = ODEProblem(sys, inival, (0, tend))
    # sol = solve(problem, Rodas5P())
    # tsol = reshape(sol, sys)

    tff = VoronoiFVM.TestFunctionFactory(sys)
    tfc = testfunction(tff, [1], [2])
    outflow = integrate(sys, tfc, tsol; rate=false)
    total_outflow = trapz(outflow.t, reduce(hcat, outflow.u)')
    return [total_outflow]
end


f(v) = main(;v)

backend = AutoMooncake()
x = ones(1)
prep = prepare_jacobian(f, backend, x)
jacobian(f, prep, backend, x)

If you have any ideas for how I should proceed, that would be very much appreciated!
Thanks again for your time and for maintaining this great package!