Adjoint + submesh

firedrake/adjoint_utils/blocks/assembly.py

@@ -1,6 +1,6 @@
 import ufl
 import firedrake
-from ufl.domain import as_domain
+from ufl.domain import extract_domains
 from ufl.formatting.ufl2unicode import ufl2unicode
 from pyadjoint import Block, AdjFloat, create_overloaded_object
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
@@ -11,14 +11,15 @@ class AssembleBlock(Block):
     def __init__(self, form, ad_block_tag=None):
         super(AssembleBlock, self).__init__(ad_block_tag=ad_block_tag)
         self.form = form
-        try:
-            mesh = as_domain(form)
+        try:  # form can have multiple meshes
+            meshes = tuple(extract_domains(form))
         except AttributeError:
-            mesh = None
+            meshes = None
 
-        if mesh and not isinstance(self.form, ufl.Interpolate):
+        if meshes and not isinstance(self.form, ufl.Interpolate):
             # Interpolation differentiation wrt spatial coordinates is currently not supported.
-            self.add_dependency(mesh)
+            for mesh in meshes:  # add all meshes as dependency
+                self.add_dependency(mesh, no_duplicates=True)
 
         for c in self.form.coefficients():
             self.add_dependency(c, no_duplicates=True)
@@ -104,8 +105,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         arity_form = len(form.arguments())
 
         if isconstant(c):
-            mesh = as_domain(self.form)
-            space = c._ad_function_space(mesh)
+            space = c.function_space()
         elif isinstance(c, (firedrake.Function, firedrake.Cofunction)):
             space = c.function_space()
         elif isinstance(c, firedrake.MeshGeometry):
@@ -162,8 +162,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
         c1_rep = block_variable.saved_output
 
         if isconstant(c1):
-            mesh = as_domain(form)
-            space = c1._ad_function_space(mesh)
+            space = c1.function_space()
         elif isinstance(c1, (firedrake.Function, firedrake.Cofunction)):
             space = c1.function_space()
         elif isinstance(c1, firedrake.MeshGeometry):

firedrake/adjoint_utils/blocks/solving.py

@@ -1,5 +1,6 @@
 import numpy
 import ufl
+from ufl.domain import extract_domains
 from ufl import replace
 from ufl.formatting.ufl2unicode import ufl2unicode
 from enum import Enum
@@ -74,8 +75,17 @@ def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
         for bc in self.bcs:
             self.add_dependency(bc, no_duplicates=True)
 
-        mesh = self.lhs.ufl_domain()
-        self.add_dependency(mesh)
+        try:  # add all meshes as dependency
+            for mesh in tuple(extract_domains(self.lhs)):
+                self.add_dependency(mesh, no_duplicates=True)
+        except AttributeError:
+            pass
+
+        if isinstance(self.rhs, (ufl.Form, ufl.Cofunction)):
+            # add all meshes as dependency
+            for mesh in tuple(extract_domains(self.rhs)):
+                self.add_dependency(mesh, no_duplicates=True)
+
         self._init_solver_parameters(args, kwargs)
 
     def _init_solver_parameters(self, args, kwargs):
@@ -243,9 +253,8 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         c_rep = block_variable.saved_output
 
         if isconstant(c):
-            mesh = F_form.ufl_domain()
             trial_function = firedrake.TrialFunction(
-                c._ad_function_space(mesh)
+                c.function_space()
             )
         elif isinstance(c, (firedrake.Function, firedrake.Cofunction)):
             trial_function = firedrake.TrialFunction(c.function_space())
@@ -455,8 +464,7 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs,
             return [tmp_bc]
 
         if isconstant(c_rep):
-            mesh = F_form.ufl_domain()
-            W = c._ad_function_space(mesh)
+            W = c_rep.function_space()
         elif isinstance(c, firedrake.MeshGeometry):
             X = firedrake.SpatialCoordinate(c)
             W = c._ad_function_space()
@@ -759,7 +767,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
         if isinstance(c, (firedrake.Function, firedrake.Cofunction)):
             trial_function = firedrake.TrialFunction(c.function_space())
         elif isinstance(c, firedrake.Constant):
-            mesh = F_form.ufl_domain()
+            mesh = extract_domains(F_form)[0]  # I don't like this
             trial_function = firedrake.TrialFunction(
                 c._ad_function_space(mesh)
             )