Fixing errors related to default L2 grad in adjoints

firedrake/adjoint_utils/function.py

@@ -225,24 +225,22 @@ def _ad_convert_riesz(self, value, options=None):
         from firedrake import Function, Cofunction
 
         options = {} if options is None else options
-        riesz_representation = options.get("riesz_representation", "l2")
+        riesz_representation = options.get("riesz_representation", "L2")
         solver_options = options.get("solver_options", {})
         V = options.get("function_space", self.function_space())
 
-        if riesz_representation != "l2" and not isinstance(value, Cofunction):
-            raise TypeError("Expected a Cofunction")
-        elif not isinstance(value, (Number, Cofunction, Function)):
+        if not isinstance(value, (Number, Cofunction, Function)):
             raise TypeError("Expected a Cofunction, Function or a float")
 
-        if riesz_representation == "l2":
-            if isinstance(value, (Cofunction, Function)):
-                return Function(V, val=value.dat)
-            else:
-                f = Function(V)
-                with stop_annotating():
-                    f.assign(value)
-                return f
+        if isinstance(value, Number):
+            # Default of a function datatype is zero
+            # Only works for l2 riesz representation
+            f = Function(V)
+            f.assign(value)
+            return f
 
+        if riesz_representation == "l2":
+            return Function(V, val=value.dat)
         elif riesz_representation in ("L2", "H1"):
             ret = Function(V)
             a = self._define_riesz_map_form(riesz_representation, V)
@@ -277,6 +275,8 @@ def _define_riesz_map_form(self, riesz_representation, V):
     def _ad_convert_type(self, value, options=None):
         # `_ad_convert_type` is not annotated, unlike `_ad_convert_riesz`
         options = {} if options is None else options
+        if "riesz_representation" not in options:
+            options = {"riesz_representation": "L2"}
         riesz_representation = options.get("riesz_representation", "L2")
         if riesz_representation is None:
             return value
@@ -317,7 +317,7 @@ def _ad_dot(self, other, options=None):
         from firedrake import assemble
 
         options = {} if options is None else options
-        riesz_representation = options.get("riesz_representation", "l2")
+        riesz_representation = options.get("riesz_representation", "L2")
         if riesz_representation == "l2":
             return self.dat.inner(other.dat)
         elif riesz_representation == "L2":

firedrake/ml/pytorch/fem_operator.py

@@ -81,9 +81,8 @@ def backward(ctx, grad_output):
         if isinstance(adj_input, Constant) and adj_input.ufl_shape == ():
             # This will later on result in an `AdjFloat` adjoint input instead of a Constant
             adj_input = float(adj_input)
-
         # Compute adjoint model of `F`: delegated to pyadjoint.ReducedFunctional
-        adj_output = F.derivative(adj_input=adj_input)
+        adj_output = F.derivative(adj_input=adj_input, options={"riesz_representation": None})
 
         # Tuplify adjoint output
         adj_output = (adj_output,) if not isinstance(adj_output, collections.abc.Sequence) else adj_output

tests/regression/test_adjoint_operators.py

@@ -744,7 +744,8 @@ def test_copy_function():
 def test_consecutive_nonlinear_solves():
     mesh = UnitSquareMesh(1, 1)
     V = FunctionSpace(mesh, "CG", 1)
-    uic = Constant(2.0, domain=mesh)
+    uic = Function(V)
+    uic.assign(2.0)
     u1 = Function(V).assign(uic)
     u0 = Function(u1)
     v = TestFunction(V)
@@ -756,8 +757,7 @@ def test_consecutive_nonlinear_solves():
         solver.solve()
     J = assemble(u1**16*dx)
     rf = ReducedFunctional(J, Control(uic))
-    h = Constant(0.01, domain=mesh)
-    assert taylor_test(rf, uic, h) > 1.9
+    assert taylor_test(rf, uic, Function(V).assign(0.1)) > 1.9
 
 
 @pytest.mark.skipcomplex
@@ -892,10 +892,37 @@ def test_cofunction_subfunctions_with_adjoint():
 def test_none_riesz_representation_to_derivative():
     mesh = UnitIntervalMesh(1)
     space = FunctionSpace(mesh, "Lagrange", 1)
-    u = Function(space).interpolate(SpatialCoordinate(mesh)[0])
-    J = assemble((u ** 2) * dx)
-    rf = ReducedFunctional(J, Control(u))
-    assert isinstance(rf.derivative(), Function)
-    assert isinstance(rf.derivative(options={"riesz_representation": "H1"}), Function)
-    assert isinstance(rf.derivative(options={"riesz_representation": "L2"}), Function)
-    assert isinstance(rf.derivative(options={"riesz_representation": None}), Cofunction)
+    f = Function(space).interpolate(SpatialCoordinate(mesh)[0])
+    J = assemble((f ** 2) * dx)
+    rf = ReducedFunctional(J, Control(f))
+    with stop_annotating():
+        v = TestFunction(space)
+        u = TrialFunction(space)
+        dJdu_cofunction = assemble(2 * inner(f, v) * dx)
+
+        # Riesz representation with l2
+        dJdu_function_l2 = Function(space, val=dJdu_cofunction.dat)
+
+        # Riesz representation with H1
+        a = firedrake.inner(u, v)*firedrake.dx \
+            + firedrake.inner(firedrake.grad(u), firedrake.grad(v))*firedrake.dx
+        dJdu_function_H1 = Function(space)
+        solve(a == dJdu_cofunction, dJdu_function_H1)
+
+        # Riesz representation with L2
+        a = firedrake.inner(u, v)*firedrake.dx
+        dJdu_function_L2 = Function(space)
+        solve(a == dJdu_cofunction, dJdu_function_L2)
+
+    dJdu_none = rf.derivative(options={"riesz_representation": None})
+    dJdu_l2 = rf.derivative(options={"riesz_representation": "l2"})
+    dJdu_H1 = rf.derivative(options={"riesz_representation": "H1"})
+    dJdu_default_L2 = rf.derivative()
+    assert (
+        isinstance(dJdu_none, Cofunction) and isinstance(dJdu_function_l2, Function)
+        and isinstance(dJdu_H1, Function) and isinstance(dJdu_default_L2, Function)
+        and np.allclose(dJdu_none.dat.data, dJdu_cofunction.dat.data)
+        and np.allclose(dJdu_l2.dat.data, dJdu_function_l2.dat.data)
+        and np.allclose(dJdu_H1.dat.data, dJdu_function_H1.dat.data)
+        and np.allclose(dJdu_default_L2.dat.data, dJdu_function_L2.dat.data)
+    )