Do not set domains of all active coefficients in domain_integral_type (#4803)

* Raise useful exception
* Revert #4775

Author: pbrubeck

firedrake/__init__.py

@@ -67,7 +67,7 @@ def init_petsc():
 from firedrake.cofunction import Cofunction, RieszMap  # noqa: F401
 from firedrake.constant import Constant  # noqa: F401
 from firedrake.deflation import DeflatedSNES, Deflation  # noqa: F401
-from firedrake.exceptions import ConvergenceError  # noqa: F401
+from firedrake.exceptions import ConvergenceError, MismatchingDomainError  # noqa: F401
 from firedrake.function import (  # noqa: F401
     Function, PointNotInDomainError,
     CoordinatelessFunction, PointEvaluator

firedrake/exceptions.py

@@ -1,3 +1,4 @@
+from tsfc.exceptions import MismatchingDomainError  # noqa: F401
 
 
 class ConvergenceError(Exception):

tests/firedrake/regression/test_multiple_domains.py

@@ -41,10 +41,10 @@ def test_mismatching_meshes_indexed_function(mesh1, mesh3):
     with pytest.raises(NotImplementedError):
         project(d1, target)
 
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(MismatchingDomainError):
         assemble(inner(d1, TestFunction(V2))*dx(domain=mesh3))
 
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(MismatchingDomainError):
         assemble(inner(d1, TestFunction(V2))*dx(domain=mesh1))
 
 
@@ -177,29 +177,29 @@ def test_multi_domain_assemble():
 
     for i, j in [(0, 1), (1, 0)]:
         a1 = inner(u[i], v[j])*dx(domain=mesh1)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(MismatchingDomainError):
             assemble(a1)
         a2 = inner(u[i], v[j])*dx(domain=mesh2)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(MismatchingDomainError):
             assemble(a2)
         l1 = inner(f[i], v[j])*dx(domain=mesh1)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(MismatchingDomainError):
             assemble(l1)
         l2 = inner(f[i], v[j])*dx(domain=mesh2)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(MismatchingDomainError):
             assemble(l2)
 
     for i, j in [(0, 0), (1, 1)]:
         a = inner(u[i], v[j])*dx(domain=mesh1)
         if i == 1:
-            with pytest.raises(NotImplementedError):
+            with pytest.raises(MismatchingDomainError):
                 assemble(a)
             continue
         A = assemble(a)
         assert A.M.values.shape == (V.dim(), V.dim())
 
     a = inner(u[0], v[0])*dx(domain=mesh1) + inner(u[0], v[1])*dx(domain=mesh2)
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(MismatchingDomainError):
         assemble(a)
 
     a = inner(u[0], v[0])*dx(domain=mesh1) + inner(u[1], v[1])*dx(domain=mesh2)

tests/firedrake/submesh/test_submesh_assemble.py

@@ -555,7 +555,9 @@ def expr(m):
 
     Q = FunctionSpace(mesh, "DG", 0)
     q = Function(Q).interpolate(expr(mesh))
-    A = assemble(inner(grad(usub) * q, grad(vsub))*dx(domain=submesh))
+
+    subdx = Measure("dx", submesh, intersect_measures=(Measure("dx", mesh),))
+    A = assemble(inner(grad(usub) * q, grad(vsub))*subdx)
 
     Qsub = FunctionSpace(submesh, "DG", 0)
     qsub = Function(Qsub).interpolate(expr(submesh))

tsfc/__init__.py

@@ -1,5 +1,6 @@
 from tsfc.driver import compile_form, compile_expression_dual_evaluation  # noqa: F401
 from tsfc.parameters import default_parameters  # noqa: F401
+from tsfc.exceptions import MismatchingDomainError  # noqa: F401
 
 
 def register_citations():

tsfc/driver.py

@@ -23,6 +23,8 @@
 from tsfc.parameters import default_parameters, is_complex
 from tsfc.ufl_utils import apply_mapping, extract_firedrake_constants
 import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy
+from tsfc.exceptions import MismatchingDomainError
+
 
 # To handle big forms. The various transformations might need a deeper stack
 sys.setrecursionlimit(3000)
@@ -90,6 +92,9 @@ def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), di
         complex_mode=complex_mode,
     )
     logger.info(GREEN % "compute_form_data finished in %g seconds.", time.time() - cpu_time)
+
+    validate_domains(form_data.preprocessed_form)
+
     # Create local kernels.
     kernels = []
     for integral_data in form_data.integral_data:
@@ -135,27 +140,16 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F
             if coeff in form_data.coefficient_split:
                 coefficient_split[coeff] = form_data.coefficient_split[coeff]
             coefficient_numbers.append(form_data.original_coefficient_positions[i])
-
     mesh = integral_data.domain
     all_meshes = extract_domains(form_data.original_form)
     domain_number = all_meshes.index(mesh)
-    coefficient_meshes = chain.from_iterable(map(extract_domains, coefficients))
-
-    domain_integral_type_map = dict.fromkeys(all_meshes, None)
-    domain_integral_type_map.update(dict.fromkeys(coefficient_meshes, "cell"))
-    domain_integral_type_map.update(integral_data.domain_integral_type_map)
-
-    for arg in arguments:
-        if domain_integral_type_map[extract_unique_domain(arg)] is None:
-            raise NotImplementedError("Assembly of forms over unrelated meshes is not supported. "
-                                      "Try using Submeshes or cross-mesh interpolation.")
 
     integral_data_info = TSFCIntegralDataInfo(
         domain=integral_data.domain,
         integral_type=integral_data.integral_type,
         subdomain_id=integral_data.subdomain_id,
         domain_number=domain_number,
-        domain_integral_type_map=domain_integral_type_map,
+        domain_integral_type_map={mesh: integral_data.domain_integral_type_map.get(mesh, None) for mesh in all_meshes},
         arguments=arguments,
         coefficients=coefficients,
         coefficient_split=coefficient_split,
@@ -186,6 +180,31 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F
     return builder.construct_kernel(kernel_name, ctx, parameters["add_petsc_events"])
 
 
+def validate_domains(form):
+    if len(extract_domains(form)) == 1:
+        # Not a multi-domain form, we do not need to keep checking
+        return
+
+    for itg in form.integrals():
+        # Check that all domains are related to each other
+        domain = itg.ufl_domain()
+        for other_domain in itg.extra_domain_integral_type_map():
+            if domain.submesh_youngest_common_ancester(other_domain) is None:
+                raise MismatchingDomainError("Assembly of forms over unrelated meshes is not supported. "
+                                             "Try using Submeshes or cross-mesh interpolation.")
+
+        # Check that all Arguments and Coefficients are defined on the valid domains
+        valid_domains = set(itg.extra_domain_integral_type_map())
+        valid_domains.add(domain)
+
+        itg_domains = set(extract_domains(itg))
+        if len(itg_domains - valid_domains) > 0:
+            raise MismatchingDomainError("Argument or Coefficient domain not found in integral. "
+                                         "Possibly, the form contains coefficients on different meshes "
+                                         "and requires measure intersection, for example: "
+                                         'Measure("dx", argument_mesh, intersect_measures=[Measure("dx", coefficient_mesh)]).')
+
+
 def preprocess_parameters(parameters):
     if parameters is None:
         parameters = default_parameters()

tsfc/exceptions.py

@@ -0,0 +1,4 @@
+
+
+class MismatchingDomainError(Exception):
+    """Error raised for unsupported multidomain problems"""

tsfc/kernel_interface/common.py

@@ -57,6 +57,7 @@ def coefficient(self, ufl_coefficient, restriction):
         expressions."""
         kernel_arg = self.coefficient_map[ufl_coefficient]
         domain = extract_unique_domain(ufl_coefficient)
+        assert self._domain_integral_type_map[domain] is not None
         if ufl_coefficient.ufl_element().family() == 'Real':
             return kernel_arg
         elif not self._domain_integral_type_map[domain].startswith("interior_facet"):