Advection-diffusion-reaction with multiple prognostic variables

In the the previous demo, we solved the Gray-Scott equation using a mixed formulation for the two tracer species. Here, we instead use different fields for each of them, treating the corresponding equations separately. This considers an additional level of complexity compared with the split solid body rotation demo because the equations differ in both the diffusion and reaction terms.

from firedrake import *

from goalie_adjoint import *

This time, we have two fields instead of one, as well as two function spaces.

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


def get_function_spaces(mesh):
    return {
        "a": FunctionSpace(mesh, "CG", 1),
        "b": FunctionSpace(mesh, "CG", 1),
    }

Therefore, the initial condition must be constructed using separate Functions.

def get_initial_condition(mesh_seq):
    x, y = SpatialCoordinate(mesh_seq[0])
    fs_a = mesh_seq.function_spaces["a"][0]
    fs_b = mesh_seq.function_spaces["b"][0]
    b_init = assemble(
        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,
            ),
            fs_b,
        )
    )
    a_init = assemble(interpolate(1 - 2 * b_init, fs_a))
    return {"a": a_init, "b": b_init}

We now see why get_form() needs to provide a function whose return value is a dictionary: its keys correspond to the different equations being solved.

def get_form(mesh_seq):
    def form(index):
        a, a_ = mesh_seq.fields["a"]
        b, b_ = mesh_seq.fields["b"]
        psi_a = TestFunction(mesh_seq.function_spaces["a"][index])
        psi_b = TestFunction(mesh_seq.function_spaces["b"][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_a = (
            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
        )
        F_b = (
            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 {"a": F_a, "b": F_b}

    return form

Correspondingly the solver needs to be constructed from the two parts and must include two nonlinear solves at each timestep.

def get_solver(mesh_seq):
    def solver(index):
        a, a_ = mesh_seq.fields["a"]
        b, b_ = mesh_seq.fields["b"]

        # Setup solver objects
        forms = mesh_seq.form(index)
        F_a = forms["a"]
        F_b = forms["b"]
        nlvp_a = NonlinearVariationalProblem(F_a, a)
        nlvs_a = NonlinearVariationalSolver(nlvp_a, ad_block_tag="a")
        nlvp_b = NonlinearVariationalProblem(F_b, b)
        nlvs_b = NonlinearVariationalSolver(nlvp_b, ad_block_tag="b")

        # 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_a.solve()
            nlvs_b.solve()
            yield

            a_.assign(a)
            b_.assign(b)
            t += dt

    return solver

Let’s consider the same QoI, time partition, and mesh sequence as in the previous demo, so that the outputs can be straightforwardly compared.

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

    return qoi


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),
    ],
)

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()

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

This tutorial can be dowloaded as a Python script.