Firedrake/Dolfin-adjoint integration

[dham]

This issue sketches out the plan for the sprint to integrate Firedrake and Dolfin-adjoint more closely.

The issues

Dolfin-adjoint subclasses some Firedrake objects. This causes issues if users manage to use unsubclassed objects. They can do this either by importing firedrake without realising they also have to import firedrake_adjoint, or in more subtle ways by subclassing Firedrake objects.

The proposed solution

We intend to refactor Dolfin-adjoint and Firedrake as follows:

• Dolfin-adjoint becomes a hard Firedrake dependency.
• Dolfin-adjoint ceases to import Firedrake.
• Dolfin-adjoint subclasses become mixins, and Firedrake always inherits from them.
• Dolfin-adjoint wrapper functions which annotate Firedrake operations (e.g. solve) become decorators which are applied in the Firedrake code.
• Annotatiation is switched off by default. This prevents unsuspecting users who are not solving an inverse problem from leaking memory. Users wishing to solve inverse problems will need to explicitly press "record" to start taping.

Potential issues

• Are there still places where Dolfin-adjoint still needs to call Firedrake?
• Will going through Dolfin-adjoint wrappers cause a measurable performance hit (hopefully no, because it just amounts to checking the value of a global variable).
• How much does this break the commonality of code between firedrake_adjoint and fenics_adjoint. Given the FEniCS API changes in the pipeline, is this even an issue?

[dham]

Function.split() currently does not get annotated. This is a Bad Thing.

[stephankramer]

Another thing that really needs fixing is interpolate it is currently not annotated at all (either the function or the method). It "works" when it's used as initialisation from an analytical expression (e.g. SpatialCoordinate dependent) as pyadjoint automatically checkpoints the function when it's first used and it's not as a block_variable on tape yet - but fails if the interpolate changes some intermediate state of a function, or when the function depends on other functions.

A similar story is true for reading from a DumbCheckPoint. Usually this is done as initialisation, and it works (on replay the checkpoint just sits in memory rather then being read from the actual hdf5 file) but when it is read into a function that's already been used as dependency previously (e.g. reading in a time series), this will not work.

[stephankramer]

Another one that came up in this context is appctx which is provided to assemble and solve to provide additional information to the solvers/preconditioners. Storing the appctx dictionary on tape has been fixed but the wider issue is that the appctx may contain UFL expressions that contain coefficients that need to be rewired to point to the relevant checkpointed functions and constants. Since knowledge about what is stored in the appctx is known to the specific preconditioner that picks it up it makes most sense to have it be responsible for replacing the coefficients when used (with some generic helper function). This should be much easier after the pyadjoint integration. Note that we're assuming that this is only used to improve the solver convergence, i.e. the provided appctx info does not influence the outcome of the solve / we are not lying about the actual system we are solving for. Otherwise the extra information should be stored as dependencies of the solve block but then we would also need the relevant adjoint/hessian hooks to compute the relevant contributions in the backward model.

[dham]

A big chunk of the interpolate issue would be solved if we merged #1453. That would give us the tangent linear and adjoint operators for interpolation. I will try to expedite.

[dham]

We need an xfail test for Function.split

[mc4117]

I have been looking into the adjoint issue with split(). It is defined in https://github.com/firedrakeproject/firedrake/blob/master/firedrake/function.py. I think the issue comes from the fact that the split function requires the use of ".dat" (Lines 103-107) and it is this that doesn't get annotated to the tape properly.

[dham]

@mc4117: I don't think that's right. The Dats are an implementation issue below the level that dolfin-adjoint works at. The issue here is that split() creates new Functions but the relationship between the new, split, Functions and the original Function is not taped.

[mc4117]

I have done some work on this and have created a branch on thetis "https://github.com/thetisproject/thetis/tree/fixing_split". Unfortunately I can't create a branch on pyadjoint because I don't have write access but the main file I changed was the one attached which is this file "firedrake/src/pyadjoint/firedrake_adjoint/types/function.py". However in my test example ("simple_split.txt") even though split seems to now be taping the end result from the reduced functional is the same so I don't know if I have actually managed to fix the problem. simple_split.txt

function.txt

[dham]

@mc4117 I have just invited you to Firedrake, which should enable you to push a branch with your proposed changes.

[wence-]

This happened, right?

[mc4117]

It did happen but function.split() still isn't passing tests on firedrake apparently (@dham can probably provide more information) so it's not fixed yet

[Ig-dolci]

@dham the issue of Function.split() does not get annotate was fixed?