Should the subfunction merge block apply a mask on the adjoint evaluation?

[JHopeCollins]

Background

Adjoining subfunctions is hard because it has to know that it doesn't hold its own data, it just views the full function. This means that when a subfunction is used, it also needs to add the full function onto the tape as a dependency so that the correct data is used.
The pyadjoint.FloatingType base class does with two custom blocks:

• When a subfunction usub = u.subfunctions[i] for component i of a full function u is added as a dependency to block block0 on the tape, what actually happens is:
  1. a SubfunctionBlock is added to the tape with u as a dependency and usub as an output. This block just filters the usub component.
  2. usub is then added as a dependency of block0. Now block0 implicitly depends on the value of u after filtering out everything but usub.
• When usub is added as an output to block block1 on the tape, what actually happens is:
  1. usub is added as an output to block block1.
  2. a FunctionMergeBlock is added to the tape with usub as a dependency and u as an output. This block combines the data in usub with all data in u except that of component i. block1 now implicitly has u as an output, but only for the data in component i.

This test just creates a function in the R space, and tapes 1) assigning to it from a control 2) multiplying by 2 and 3) calculating the square. The tapes with/without using subfunctions are shown below, where you can see the extra blocks in the subfunction case. (The odd numbering is just so the code matches the tape).

with stop_annotating():
    mesh = UnitIntervalMesh(1)
    R = FunctionSpace(mesh, "R", 0)

    w2 = Function(R).assign(2.0)
    w4 = Function(R)
    w7 = u.subfunctions[0]

use_subfunctions = True

continue_annotation()
with set_working_tape() as tape:
    w4.assign(w2)
    if use_subfunctions:
        w7 *= 2
    else:
        w4 *= 2
    J = assemble(inner(w4, w4) * dx)
    rf = ReducedFunctional(J, Control(w2), tape=tape)
pause_annotation()

assert taylor_test(rf, w2, Constant(0.1)) > 1.9

Without using subfunctions:

Using subfunctions:

Problem

This test fails if using subfunctions!
We want the adj_value of the i-th component to depend only on what happens along the subfunction branch, and the adj_value of all other components to be whatever they were before the subfunctions split/merge.

Instead, the adjoint of FunctionMergeBlock was (correctly) outputting the i-th component of the adj_value along the subfunction branch of the tape, but (incorrectly) outputting all components of the adj_value along the full-function branch of the tape.
When these two branches recombined the adj_value for the i-th component then had contributions from both branches.

Solution

The merge out[i] = usub; out[not i] = u is a mask. The adjoint of a mask is also a mask, so FunctionMergeBlock now applies a mask before returning the adjoint outputs: output0 = adj_input[i]; output1 = adj_input[not i]

Changes added