TAO interface

[jrmaddison]

Optimisation using PETSc/TAO.

Uses the pyadjoint OverloadedType interface to provide a generic implementation. However these interfaces don't provide enough information about data decomposition with MPI parallelism. This is worked around in this PR using OverloadedType._ad_assign_numpy, but this requires a temporary, but global, NumPy array. Fixing this requires a change to OverloadedType and then other updates elsewhere (in particular in Firedrake).

Related, OverloadedType._ad_to_list returns a list rather than a NumPy ndarray.

Not added:

• MaximisationProblem support.
• Equality or inequality constraints.
• Any optimisations to prevent extra forward or first adjoint reruns (e.g. with second order methods which might call TAOObjective.objective_gradient followed by TAOObjective.hessian).

[dham]

I think that even if you are going to go via Numpy, this should be done by giving OverloadedType an ._ad_as_petsc_vec method so that we can go back and do it properly for Firedrake Functions afterwards. As it stands, I think that the TAO interface would have to be substantially rewritten in order to support OverloadedTypes that had a native mechanism for casting to PETSc Vec, and that seems unfortunate.

[jrmaddison]

I think that even if you are going to go via Numpy, this should be done by giving OverloadedType an ._ad_as_petsc_vec method so that we can go back and do it properly for Firedrake Functions afterwards. As it stands, I think that the TAO interface would have to be substantially rewritten in order to support OverloadedTypes that had a native mechanism for casting to PETSc Vec, and that seems unfortunate.

I think the 'one-or-more' pyadjoint interfaces mean that copying is going to be needed somewhere. e.g. there's nothing to stop a mixture of AdjFloats and firedrake.Function controls. Maybe those could be PETSc copy operations if the backend supplies PETSc Vecs -- I suggested NumPy as an interface layer as that's close to what OverloadedTypes already supply.

[dham]

OK, I can see what's going on here. I actually think it would be an easy fix to do this without the global gather, it just requires a few very simple methods to be added to Overloaded type, but I guess that can be delayed. Please add a health warning about the global gather in the class docstring, because the comment on the PR will not be visible to users.

[jrmaddison]

OK, I can see what's going on here. I actually think it would be an easy fix to do this without the global gather, it just requires a few very simple methods to be added to Overloaded type, but I guess that can be delayed. Please add a health warning about the global gather in the class docstring, because the comment on the PR will not be visible to users.

It's not even a global gather, we just need to known the number of process local (owned) degrees of freedom, and OverloadedType doesn't provide this except indirectly via _ad_convert_numpy.

[dham]

OK, I can see what's going on here. I actually think it would be an easy fix to do this without the global gather, it just requires a few very simple methods to be added to Overloaded type, but I guess that can be delayed. Please add a health warning about the global gather in the class docstring, because the comment on the PR will not be visible to users.

It's not even a global gather, we just need to known the number of process local (owned) degrees of freedom, and OverloadedType doesn't provide this except indirectly via _ad_convert_numpy.

OK, I'm confused again. I think the current implementation does a global gather onto a single numpy array. Is that not correct?

[jrmaddison]

OK, I'm confused again. I think the current implementation does a global gather onto a single numpy array. Is that not correct?

You're correct. I'd missed (or forgotten) how OverloadedType._ad_assign_numpy and OverloadedType._to_list behave -- e.g. using a gather in OverloadedType._to_list.

This obviously won't scale, but I also don't see an alternative via the current OverloadedType interface.

[jrmaddison]

Firedrake test requires https://github.com/firedrakeproject/firedrake/pull/3657