Differentiating through the solver

IncompressibleNavierStokes is reverse-mode differentiable, which means that you can back-propagate gradients through the code. Two AD libraries are currently supported:

• Zygote.jl: it is the default AD library in the Julia ecosystem and is the most widely used.

• Enzyme.jl: currently has low coverage over the Julia programming language, however it is usually the most efficient if applicable.

See also the manual page on differentiability.

Automatic differentiation with Zygote

Zygote does not support array mutation. To make code differentiable with Zygote, we must therefore use the non-mutating versions of the operators (without exclamation marks), such as divergence and timestep. These allocate intermediate velocity fields that are kept in memory for use in the backward pass.

Example: Gradient of kinetic energy

To differentiate outputs of a simulation with respect to the initial conditions, make a time stepping loop composed of differentiable operations:

using IncompressibleNavierStokes

n = 100
ax = range(0, 1, n + 1)
setup = Setup(;
    x = (ax, ax),
    boundary_conditions = (;
        u = ((PeriodicBC(), PeriodicBC()), (PeriodicBC(), PeriodicBC())),
    ),
)
psolver = default_psolver(setup)
method = LMWray3()
Δt = 0.001
nstep = 100
(; Iu) = setup
function final_energy(u)
    state = (; u)
    stepper = create_stepper(method; setup, psolver, state, t = 0.0)
    for it = 1:nstep
        stepper = timestep(method, navierstokes, stepper, Δt; params = (; viscosity = 2e-3))
    end
    (; u) = stepper.state
    E = sum(abs2, u[Iu[1], 1]) / 2n^2 + sum(abs2, u[Iu[2], 2]) / 2n^2
end

u = random_field(setup)
final_energy(u)

0.24874366405018974

Now compute the gradient with Zygote:

using Zygote
g, = Zygote.gradient(final_energy, u)

@show size(u) size(g)

(102, 102, 2)

Now g is the gradient of final_energy with respect to the initial conditions u, and consequently has the same size.

Note that every operation in the final_energy function is non-mutating and thus differentiable.

Automatic differentiation with Enzyme

Enzyme performs AD on optimized code, which allows it to meet or exceed the performance of state-of-the-art AD tools. The downside is that the function to differentiate must not require garbage collection or calls to BLAS/LAPACK. Mutation is supported however, meaning that the fast in-place operators of IncompressibleNavierStokes can be used.

Enzyme limitation: vector returns

Enzyme's autodiff function can only handle functions with scalar output. To implement pullbacks for array-valued functions, use a mutating function that returns nothing and stores its result in one of the arguments, which must be passed wrapped in Duplicated.

Example: Gradient of the right-hand side

In this example we differentiate the right-hand side of the Navier-Stokes equations with respect to the velocity field u:

using Enzyme
ax = range(0, 1, 101)
setup = Setup(;
    x = (ax, ax),
    boundary_conditions = (;
        u = ((PeriodicBC(), PeriodicBC()), (PeriodicBC(), PeriodicBC())),
    ),
)
psolver = default_psolver(setup)
u = random_field(setup)
dudt = similar(u)
t = 0.0
f! = right_hand_side!

right_hand_side! (generic function with 1 method)

Notice that we are using the mutating (in-place) version of the right-hand side function. This function can not be differentiated by Zygote, which requires the slower non-mutating version.

We then define the derivative parts of the input and output, required to store the adjoint values:

ddudt = Enzyme.make_zero(dudt) .+ 1;
du = Enzyme.make_zero(u);

Remember that the derivative of the output (also called the seed) has to be set to $1$ in order to compute the gradient. In this case the output is the force, that we store by mutating the value of dudt inside right_hand_side!.

Then we pack the parameters to be passed to right_hand_side!:

viscosity = 2e-3
params = setup, psolver, viscosity;
params_ref = Ref(params);

Now, we call the autodiff function from Enzyme:

Enzyme.autodiff(
    Enzyme.Reverse,
    f!,
    Duplicated(dudt, ddudt),
    Duplicated(u, du),
    Const(params_ref),
    Const(t),
)

((nothing, nothing, nothing, nothing),)

Since we have passed a Duplicated object, the gradient of u is stored in du.

Note that the output seed ddudt is consumed in the process: Enzyme zeroes it during the reverse pass.

Finally, we can also compare its value with the one obtained by Zygote differentiating the out-of-place (non-mutating) version of the right-hand side. We seed the Zygote pullback with a fresh copy of the original seed (a vector of ones):

f = create_right_hand_side(setup, psolver)
_, zpull = Zygote.pullback(f, u, (; viscosity), 0.0);
w = one.(dudt); # The same seed values that ddudt started with
@assert zpull(w)[1] ≈ du

Copy-pasteable code

Below is the full code for this example stripped of comments and output.

using IncompressibleNavierStokes

n = 100
ax = range(0, 1, n + 1)
setup = Setup(;
    x = (ax, ax),
    boundary_conditions = (;
        u = ((PeriodicBC(), PeriodicBC()), (PeriodicBC(), PeriodicBC())),
    ),
)
psolver = default_psolver(setup)
method = LMWray3()
Δt = 0.001
nstep = 100
(; Iu) = setup
function final_energy(u)
    state = (; u)
    stepper = create_stepper(method; setup, psolver, state, t = 0.0)
    for it = 1:nstep
        stepper = timestep(method, navierstokes, stepper, Δt; params = (; viscosity = 2e-3))
    end
    (; u) = stepper.state
    E = sum(abs2, u[Iu[1], 1]) / 2n^2 + sum(abs2, u[Iu[2], 2]) / 2n^2
end

u = random_field(setup)
final_energy(u)

using Zygote
g, = Zygote.gradient(final_energy, u)

@show size(u) size(g)

using Enzyme
ax = range(0, 1, 101)
setup = Setup(;
    x = (ax, ax),
    boundary_conditions = (;
        u = ((PeriodicBC(), PeriodicBC()), (PeriodicBC(), PeriodicBC())),
    ),
)
psolver = default_psolver(setup)
u = random_field(setup)
dudt = similar(u)
t = 0.0
f! = right_hand_side!

ddudt = Enzyme.make_zero(dudt) .+ 1;
du = Enzyme.make_zero(u);

viscosity = 2e-3
params = setup, psolver, viscosity;
params_ref = Ref(params);

Enzyme.autodiff(
    Enzyme.Reverse,
    f!,
    Duplicated(dudt, ddudt),
    Duplicated(u, du),
    Const(params_ref),
    Const(t),
)

f = create_right_hand_side(setup, psolver)
_, zpull = Zygote.pullback(f, u, (; viscosity), 0.0);
w = one.(dudt); # The same seed values that ddudt started with
@assert zpull(w)[1] ≈ du

---

This page was generated using Literate.jl.