BUG: adj_value for Constants is not computed, but it used to work and works with FEniCS adjoint

[IvanYashchuk]

Something changed with adjoints for Constants and values are not computed anymore. It used to work some time ago and works as expected with FEniCS adjoint. Here's an example I'm testing with:

from firedrake import *
from firedrake.adjoint import *
import pyadjoint
continue_annotation()
# replace above with "from fenics import *; from fenics_adjoint import *" for FEniCS

mesh = UnitSquareMesh(3, 2)
V = FunctionSpace(mesh, "CG", 1)

def assemble_form(u, k0, k1):
    x = SpatialCoordinate(mesh)
    f = x[0]
    J_form = 0.5 * inner(k0 * grad(u), grad(u)) * dx - k1 * f * u * dx
    J = assemble(J_form)
    return J

u = interpolate(Constant(1.0), V)
k0 = Constant(0.5)
k1 = Constant(0.6)
J = assemble_form(u, k0, k1)
assert type(J) == pyadjoint.adjfloat.AdjFloat

tape = get_working_tape()
J.block_variable.adj_value = 1.0
# After tape.evaluate_adj the previous behavior was computing
# and assigning corresponding adjoints for k0 and k1, now it's just None
with tape.marked_nodes([u, k0, k1]):
    tape.evaluate_adj(markings=True)

print([x.block_variable.adj_value for x in (u, k0, k1)])

prints

[Cofunction(FiredrakeDualSpace(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f647c494390>, FiniteElement('Lagrange', triangle, 1), name=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 1)), 13),
 None,
 None]

Expected behavior

I expect adjoints for Constants to be computed. The last time I checked that it was working was Oct 26, 2022.

When I inspect the assemble block I see only 2 dependencies in Firedrake (mesh and P1 coefficient) and 4 dependencies in FEniCS (mesh, P1 coefficient, and two R coefficients):

from pprint import pprint
assemble_block = tape.get_blocks()[-1]
pprint([x.saved_output for x in assemble_block.get_dependencies()])

Firedrake:

[Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 1),
 Coefficient(WithGeometry(FunctionSpace(<firedrake.mesh.MeshTopology object at 0x7f647c494390>, FiniteElement('Lagrange', triangle, 1), name=None), Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 1)), 4)]

FEniCS:

[<fenics_adjoint.types.mesh.overloaded_mesh.<locals>.OverloadedMesh object at 0x7fe9bfad5430>,
 Coefficient(FunctionSpace(Mesh(VectorElement(FiniteElement('Lagrange', triangle, 1), dim=2), 0), FiniteElement('Lagrange', triangle, 1)), 13),
 Coefficient(FunctionSpace(None, FiniteElement('Real', None, 0)), 14),
 Coefficient(FunctionSpace(None, FiniteElement('Real', None, 0)), 15)]

Am I missing something in my script with Firedrake?

[connorjward]

We know of an issue with Constant derivative calculations. This may be related.

[JDBetteridge]

#3261 should fix the issue Connor mentioned, do you still see the same issue on the fix branch?

[IvanYashchuk]

Hi Jack, I tried your branch, and unfortunately, I still see the same issue. adj_value is not computed for Constants.

[connorjward]

This is clearly an issue that needs further investigation, but do things work if you use a real space function instead of a Constant?

[connorjward]

I think I know what's going on. When we add dependencies for the assemble block we loop over form.coefficients() (source). We changed Constant to no longer inherit from a UFL Coefficient so it is getting missed as a dependency.

[connorjward]

Actually, I've just attempted a fix and discovered that this is probably a feature and not a bug. In particular if I add the Constants as dependencies of the assemble block I get the following error (raised here):

Traceback (most recent call last):
  File "/home/connor/Code/firedrake-dev1/src/firedrake/mfe.py", line 28, in <module>
    tape.evaluate_adj(markings=True)
  File "/home/connor/Code/firedrake-dev1/src/pyadjoint/pyadjoint/tape.py", line 232, in evaluate_adj
    self._blocks[i].evaluate_adj(markings=markings)
  File "/home/connor/Code/firedrake-dev1/src/pyadjoint/pyadjoint/tape.py", line 109, in wrapper
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/connor/Code/firedrake-dev1/src/pyadjoint/pyadjoint/block.py", line 131, in evaluate_ad
j
    adj_output = self.evaluate_adj_component(inputs,
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/connor/Code/firedrake-dev1/src/firedrake/firedrake/adjoint_utils/blocks/assembly.py", 
line 121, in evaluate_adj_component
    if self.compat.isconstant(c):
       ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/connor/Code/firedrake-dev1/src/firedrake/firedrake/adjoint_utils/blocks/compat.py", li
ne 159, in isconstant
    raise ValueError("Firedrake Constant requires a domain to work with pyadjoint")
ValueError: Firedrake Constant requires a domain to work with pyadjoint

The solution is therefore probably to use a Function in the Real space in place of Constant.

@dham is certainly the person most knowledgeable about this.

[dham]

This reflects a change in Constant, which is a feature rather than a bug. As you can see from your output, FEniCS are actually creating variables in the R space when you instantiate a Constant. Firedrake has also moved to this position but for reasons of parallel safety, we require functions over R to be defined on a mesh. This means we now have two kinds of constant. Constants without domains are really just a performance feature to enable parameters to be passed into the assembly kernel without recompilation. You can't differentiate with respect to them and hence they don't work as adjoint variables. Constants with domains get mapped to functions over R. That's what you want in this case, so you should pass the mesh in after the value in the definitions of k0 and k1

[IvanYashchuk]

Thanks a lot for the investigation and explaining everything! It all makes sense. I've tried using R FunctionSpace instead and it works as expected.

[IvanYashchuk]

Maybe there should be a warning or an error about this when someone tries to access Constant.block_variable.adj_value?