Advection-diffusion-reaction

We have already seen time-dependent advection equations and a steady-state advection-diffusion equation. In this demo, we increase the level of complexity again and solve the Gray-Scott advection-diffusion-reaction equation, as described in [].

The test case consists of two tracer species, which experience different diffusivities and react with one another nonlinearly.

from firedrake import *

from goalie_adjoint import *

The problem is defined on a doubly periodic mesh of squares.

field_names = ["ab"]
mesh = PeriodicSquareMesh(65, 65, 2.5, quadrilateral=True, direction="both")

We solve for the tracer species using a mixed formulation, with a \(\mathbb P1\) approximation for both components.

def get_function_spaces(mesh):
    V = FunctionSpace(mesh, "CG", 1)
    return {"ab": V * V}

The initial conditions are localised within the region \([1, 1.5]^2\).

def get_initial_condition(mesh_seq):
    x, y = SpatialCoordinate(mesh_seq[0])
    fs = mesh_seq.function_spaces["ab"][0]
    ab_init = Function(fs)
    a_init, b_init = ab_init.subfunctions
    b_init.interpolate(
        conditional(
            And(And(1 <= x, x <= 1.5), And(1 <= y, y <= 1.5)),
            0.25 * sin(4 * pi * x) ** 2 * sin(4 * pi * y) ** 2,
            0,
        )
    )
    a_init.interpolate(1 - 2 * b_init)
    return {"ab": ab_init}

Since we are using a mixed formulation, the forms for each component equation are summed together.

def get_form(mesh_seq):
    def form(index):
        ab, ab_ = mesh_seq.fields["ab"]
        a, b = split(ab)
        a_, b_ = split(ab_)
        psi_a, psi_b = TestFunctions(mesh_seq.function_spaces["ab"][index])

        # Define constants
        R = FunctionSpace(mesh_seq[index], "R", 0)
        dt = Function(R).assign(mesh_seq.time_partition.timesteps[index])
        D_a = Function(R).assign(8.0e-05)
        D_b = Function(R).assign(4.0e-05)
        gamma = Function(R).assign(0.024)
        kappa = Function(R).assign(0.06)

        # Write the two equations in variational form
        F = (
            psi_a * (a - a_) * dx
            + dt * D_a * inner(grad(psi_a), grad(a)) * dx
            - dt * psi_a * (-a * b**2 + gamma * (1 - a)) * dx
            + psi_b * (b - b_) * dx
            + dt * D_b * inner(grad(psi_b), grad(b)) * dx
            - dt * psi_b * (a * b**2 - (gamma + kappa) * b) * dx
        )
        return {"ab": F}

    return form

For the solver, we just use the default configuration of Firedrake’s NonlinearVariationalSolver.

def get_solver(mesh_seq):
    def solver(index):
        ab, ab_ = mesh_seq.fields["ab"]

        # Setup solver objects
        F = mesh_seq.form(index)["ab"]
        nlvp = NonlinearVariationalProblem(F, ab)
        nlvs = NonlinearVariationalSolver(nlvp, ad_block_tag="ab")

        # Time integrate from t_start to t_end
        tp = mesh_seq.time_partition
        t_start, t_end = tp.subintervals[index]
        dt = tp.timesteps[index]
        t = t_start
        while t < t_end - 0.5 * dt:
            nlvs.solve()
            yield

            ab_.assign(ab)
            t += dt

    return solver

The term \(a * b ^ 2\) appears in both equations. By solving the adjoint for the QoI \(\int a(x,T) * b(x,T) * dx\) we consider sensitivities to this term.

def get_qoi(mesh_seq, index):
    def qoi():
        ab = mesh_seq.fields["ab"][0]
        a, b = split(ab)
        return a * b**2 * dx

    return qoi

This problem is multi-scale in time and requires spinning up by gradually increasing the timestep. This is straightforwardly done in Goalie using TimePartition.

end_time = 2000.0
dt = [0.0001, 0.001, 0.01, 0.1, (end_time - 1) / end_time]
num_subintervals = 5
dt_per_export = [10, 9, 9, 9, 10]
time_partition = TimePartition(
    end_time,
    num_subintervals,
    dt,
    field_names,
    num_timesteps_per_export=dt_per_export,
    subintervals=[
        (0.0, 0.001),
        (0.001, 0.01),
        (0.01, 0.1),
        (0.1, 1.0),
        (1.0, end_time),
    ],
)

As usual, define an appropriate MeshSeq and choose the qoi_type.

mesh_seq = AdjointMeshSeq(
    time_partition,
    mesh,
    get_function_spaces=get_function_spaces,
    get_initial_condition=get_initial_condition,
    get_form=get_form,
    get_solver=get_solver,
    get_qoi=get_qoi,
    qoi_type="end_time",
)
solutions = mesh_seq.solve_adjoint()

Finally, plot the outputs to be viewed in Paraview.

ic = mesh_seq.get_initial_condition()
for field, sols in solutions.items():
    fwd_outfile = VTKFile(f"gray_scott/{field}_forward.pvd")
    adj_outfile = VTKFile(f"gray_scott/{field}_adjoint.pvd")
    fwd_outfile.write(*ic[field].subfunctions)
    for i in range(num_subintervals):
        for sol in sols["forward"][i]:
            fwd_outfile.write(*sol.subfunctions)
        for sol in sols["adjoint"][i]:
            adj_outfile.write(*sol.subfunctions)
    adj_end = Function(ic[field]).assign(0.0)
    adj_outfile.write(*adj_end.subfunctions)

In the next demo, we consider solving the same problem, but splitting the solution field into multiple components.

This tutorial can be dowloaded as a Python script.

References