Using a mixed `Function` as a control

[jwallwork23]

I'm trying to invert for an initial condition using firedrake-adjoint and have noticed that an error arises if the control is a from mixed function space.

Consider the Thetis 2d shallow water demo. If the control is the initial tuple then the ReducedFunctional built on the square L2 norm of the final solution tuple gives inconsistent output. If the control is the initial elevation component then everything is okay.

control = Control(q_init)

Unrolling the tape,   J = 2.94130152e+07
Rerunning the solver, J = 3.04229596e+07

vs
uv_init, elev_init = q_init.split(); control = Control(elev_init)

Unrolling the tape,   J = 3.04229596e+07
Rerunning the solver, J = 3.04229596e+07

See below for code. This issue may be related to #1843.

from thetis import *
from firedrake_adjoint import *
import numpy as np

# Mesh
lx = 40e3
ly = 2e3
nx = 25
ny = 2
mesh2d = RectangleMesh(nx, ny, lx, ly)

# Bathymetry
P1_2d = FunctionSpace(mesh2d, 'CG', 1)
bathymetry_2d = Function(P1_2d, name='Bathymetry')
depth = 20.0
bathymetry_2d.assign(depth)

# Timestepping
t_end = 0.25 * 3600
t_export = 100.0

# Initial condition
V = VectorFunctionSpace(mesh2d, 'DG', 1)*FunctionSpace(mesh2d, 'CG', 2)
q_init = Function(V)
uv_init, elev_init = q_init.split()
xy = SpatialCoordinate(mesh2d)
gauss_width = 4000.
gauss_ampl = 2.0
gauss_expr = gauss_ampl * exp(-((xy[0]-lx/2)/gauss_width)**2)
elev_init.interpolate(gauss_expr)

def my_reduced_functional(init):
    """
    Given an initial surface, run the forward model and
    evaluate some functional of the final solution tuple.
    """
    solver_obj = solver2d.FlowSolver2d(mesh2d, bathymetry_2d)
    options = solver_obj.options
    options.element_family = 'dg-cg'
    options.simulation_export_time = t_export
    options.simulation_end_time = t_end
    options.timestepper_type = 'CrankNicolson'
    options.timestep = 50.0
    u_init, eta_init = init.split()
    solver_obj.assign_initial_conditions(uv=u_init, elev=eta_init)
    solver_obj.iterate()
    q = solver_obj.fields.solution_2d
    return assemble(inner(q, q)*dx)

# Choose a control, write to tape
control = Control(elev_init)  # NOTE: does not work if elev_init replaced with q_init
J = my_reduced_functional(q_init)
print("Initially, J = {:.8e}".format(J))
stop_annotating()
Jhat = ReducedFunctional(J, control)  # Create a ReducedFunctional object

# Change the initial condition and check for consistency
elev_init.assign(-elev_init)
J_unrolled = Jhat(elev_init)  # NOTE: does not work if elev_init replaced with q_init
J_rerun = my_reduced_functional(q_init)
print("Unrolling the tape,   J = {:.8e}".format(J_unrolled))
print("Rerunning the solver, J = {:.8e}".format(J_rerun))
assert np.isclose(J_unrolled, J_rerun)