Fixing errors related to default L2 grad in adjoints

firedrake/adjoint_utils/blocks/function.py

@@ -60,7 +60,8 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                 R = block_variable.output._ad_function_space(
                     prepared.function_space().mesh()
                 )
-                return self._adj_assign_constant(prepared, R)
+                adj_output = self._adj_assign_constant(prepared, R)
+                return adj_output.riesz_representation(riesz_map="l2")
             else:
                 adj_output = firedrake.Function(
                     block_variable.output.function_space())
@@ -86,7 +87,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                     ufl.derivative(
                         expr,
                         block_variable.saved_output,
-                        firedrake.Constant(1., domain=mesh)
+                        firedrake.Function(firedrake.FunctionSpace(mesh, "R", 0), val=1.0)
                     )
                 )
                 adj_output.assign(diff_expr)

firedrake/adjoint_utils/function.py

@@ -225,23 +225,24 @@ 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 isinstance(value, Number):
+            # It is applied l2 representation for this case.
+            f = Function(V)
+            f.assign(value)
+            return f
+
         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
+            return Function(V, val=value.dat)
+
+        if riesz_representation != "l2" and not isinstance(value, Cofunction):
+            raise TypeError("Expected a Cofunction")
 
         elif riesz_representation in ("L2", "H1"):
             ret = Function(V)
@@ -317,7 +318,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

@@ -83,7 +83,7 @@ def backward(ctx, grad_output):
             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": 'l2'})
 
         # Tuplify adjoint output
         adj_output = (adj_output,) if not isinstance(adj_output, collections.abc.Sequence) else adj_output

tests/regression/test_adjoint_operators.py

@@ -889,13 +889,41 @@ def test_cofunction_subfunctions_with_adjoint():
 
 
 @pytest.mark.skipcomplex  # Taping for complex-valued 0-forms not yet done
-def test_none_riesz_representation_to_derivative():
+def test_riesz_representation_for_adjoints():
+    # Test is the riesz representation norms for adjoints are working as expected.
     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)
+    )