Function.assign adjoint bugfix

[jrmaddison]

Description

The pattern

u0.assign(uic)
u1.assign(2 * u0 + uic)

with all variables Functions leads to an inconsistent adjoint calculation. The first assignment leads to the checkpoint property for u0 being set to uic, which means that the adjoint sees the expression 2 * uic + uic in the second assignment.

The PR fixes this by copying uic when setting u0.block_variable._checkpoint.

Checklist for author:

• I have linted the codebase (make lint in the firedrake source directory).
• My changes generate no new warnings.
• All of my functions and classes have appropriate docstrings.
• I have commented my code where its purpose may be unclear.
• I have included and updated any relevant documentation.
• Documentation builds locally (make linkcheck; make html; make latexpdf in firedrake/docs directory)
• I have added tests specific to the issues fixed in this PR.
• I have added tests that exercise the new functionality I have introduced
• Tests pass locally (pytest tests in the firedrake source directory) (useful, but not essential if you don't have suitable hardware).
• I have performed a self-review of my own code using the below guidelines.

Checklist for reviewer:

• Docstrings present
• New tests present
• Code correctly commented
• No bad "code smells"
• No issues in parallel
• No CI issues (excessive parallelism/memory usage/time/warnings generated)
• Upstream/dependent branches and PRs are ready

Feel free to add reviewers if you know there is someone who is already aware of this work.

Please open this PR initially as a draft and mark as ready for review once CI tests are passing.

[jrmaddison]

I think my fix works, but perhaps undermines the intent of DelegatedFunctionCheckpoint?

[Ig-dolci]

I think my fix works, but perhaps undermines the intent of DelegatedFunctionCheckpoint?

I am going to verify that.

[Ig-dolci]

DelegateCheckpoint avoids storing the checkpoint data twice in the cases that we have:

    uic = Function(V, name="uic").assign(1.0)
    u0 = Function(V, name="u0")
    u0.assign(uic)

I verified with your test script that the DelegateCheckpoint was sttil playing its role.

[Ig-dolci]

DelegateCheckpoint avoids storing the checkpoint data twice in the cases that we have:

    uic = Function(V, name="uic").assign(1.0)
    u0 = Function(V, name="u0")
    u0.assign(uic)

I verified with your test script that the DelegateCheckpoint was sttil playing its role.

Wrong affirmation. I carried out more tests and verified that DelegateCheckpoint loses its function.

[dham]

I think this isn't quite the right fix. I think that what is happening is that the target of the delegated checkpoint is a control. When the functional is re-evaluated, the saved value of the control is overwritten. Somehow I think this doesn't propagate correctly. The answer must be in there somehow. This change is effectively just removing delegated checkpoint.

[dham]

OK, I think I might have fixed this. When the adjoint is calculated, the Functions in the RHS of the assignment are replaced by the saved outputs of the adjoint variables. In order for this to be safe, it's necessary for the saved output returned by the delegated checkpoint to have its own unique symbolic identity.