Introduce MixedMesh + handle restrictions on MixedMesh

test/test_equals.py

@@ -4,6 +4,7 @@
 from ufl.finiteelement import FiniteElement
 from ufl.pullback import identity_pullback
 from ufl.sobolevspace import H1
+from ufl.exprcontainers import ExprList
 
 
 def test_comparison_of_coefficients():
@@ -69,6 +70,10 @@ def test_comparison_of_cofunctions():
     assert not v1 == u1
     assert not v2 == u2
 
+    # Objects in ExprList as happens when taking derivatives.
+    assert ExprList(v1, v1) == ExprList(v1, v1b)
+    assert not ExprList(v1, v2) == ExprList(v1, v1)
+
 
 def test_comparison_of_products():
     V = FiniteElement("Lagrange", triangle, 1, (), identity_pullback, H1)

test/test_external_operator.py

@@ -181,7 +181,7 @@ def test_differentiation_procedure_action(V1, V2):
 
 
 def test_extractions(domain_2d, V1):
-    from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients,
+    from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients_and_geometric_quantities,
                                          extract_base_form_operators, extract_coefficients, extract_constants)
 
     u = Coefficient(V1)
@@ -192,15 +192,15 @@ def test_extractions(domain_2d, V1):
 
     assert extract_coefficients(e) == [u]
     assert extract_arguments(e) == [vstar_e]
-    assert extract_arguments_and_coefficients(e) == ([vstar_e], [u])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e], [u], [])
     assert extract_constants(e) == [c]
     assert extract_base_form_operators(e) == [e]
 
     F = e * dx
 
     assert extract_coefficients(F) == [u]
     assert extract_arguments(e) == [vstar_e]
-    assert extract_arguments_and_coefficients(e) == ([vstar_e], [u])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e], [u], [])
     assert extract_constants(F) == [c]
     assert F.base_form_operators() == (e,)
 
@@ -209,14 +209,14 @@ def test_extractions(domain_2d, V1):
 
     assert extract_coefficients(e) == [u]
     assert extract_arguments(e) == [vstar_e, u_hat]
-    assert extract_arguments_and_coefficients(e) == ([vstar_e, u_hat], [u])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e, u_hat], [u], [])
     assert extract_base_form_operators(e) == [e]
 
     F = e * dx
 
     assert extract_coefficients(F) == [u]
     assert extract_arguments(e) == [vstar_e, u_hat]
-    assert extract_arguments_and_coefficients(e) == ([vstar_e, u_hat], [u])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(e) == ([vstar_e, u_hat], [u], [])
     assert F.base_form_operators() == (e,)
 
     w = Coefficient(V1)
@@ -225,14 +225,14 @@ def test_extractions(domain_2d, V1):
 
     assert extract_coefficients(e2) == [u, w]
     assert extract_arguments(e2) == [vstar_e2, u_hat]
-    assert extract_arguments_and_coefficients(e2) == ([vstar_e2, u_hat], [u, w])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(e2) == ([vstar_e2, u_hat], [u, w], [])
     assert extract_base_form_operators(e2) == [e, e2]
 
     F = e2 * dx
 
     assert extract_coefficients(e2) == [u, w]
     assert extract_arguments(e2) == [vstar_e2, u_hat]
-    assert extract_arguments_and_coefficients(e2) == ([vstar_e2, u_hat], [u, w])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(e2) == ([vstar_e2, u_hat], [u, w], [])
     assert F.base_form_operators() == (e, e2)
 
 

test/test_interpolate.py

@@ -8,7 +8,7 @@
 from ufl import (Action, Adjoint, Argument, Coefficient, FunctionSpace, Mesh, TestFunction, TrialFunction, action,
                  adjoint, derivative, dx, grad, inner, replace, triangle)
 from ufl.algorithms.ad import expand_derivatives
-from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients, extract_base_form_operators,
+from ufl.algorithms.analysis import (extract_arguments, extract_arguments_and_coefficients_and_geometric_quantities, extract_base_form_operators,
                                      extract_coefficients)
 from ufl.algorithms.expand_indices import expand_indices
 from ufl.core.interpolate import Interpolate
@@ -141,12 +141,12 @@ def test_extract_base_form_operators(V1, V2):
     # -- Interpolate(u, V2) -- #
     Iu = Interpolate(u, V2)
     assert extract_arguments(Iu) == [vstar]
-    assert extract_arguments_and_coefficients(Iu) == ([vstar], [u])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(Iu) == ([vstar], [u], [])
 
     F = Iu * dx
     # Form composition: Iu * dx <=> Action(v * dx, Iu(u; v*))
     assert extract_arguments(F) == []
-    assert extract_arguments_and_coefficients(F) == ([], [u])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(F) == ([], [u], [])
 
     for e in [Iu, F]:
         assert extract_coefficients(e) == [u]
@@ -155,7 +155,7 @@ def test_extract_base_form_operators(V1, V2):
     # -- Interpolate(u, V2) -- #
     Iv = Interpolate(uhat, V2)
     assert extract_arguments(Iv) == [vstar, uhat]
-    assert extract_arguments_and_coefficients(Iv) == ([vstar, uhat], [])
+    assert extract_arguments_and_coefficients_and_geometric_quantities(Iv) == ([vstar, uhat], [], [])
     assert extract_coefficients(Iv) == []
     assert extract_base_form_operators(Iv) == [Iv]
 

test/test_mixed_function_space_with_mixed_mesh.py

@@ -0,0 +1,109 @@
+from ufl import (triangle, Mesh, MixedMesh, FunctionSpace, TestFunction, TrialFunction, Coefficient,
+                 Measure, SpatialCoordinate, FacetNormal, CellVolume, FacetArea, inner, grad, div, split, )
+from ufl.algorithms import compute_form_data
+from ufl.finiteelement import FiniteElement, MixedElement
+from ufl.pullback import identity_pullback, contravariant_piola
+from ufl.sobolevspace import H1, HDiv, L2
+from ufl.domain import extract_domains
+
+
+def test_mixed_function_space_with_mixed_mesh_basic():
+    cell = triangle
+    elem0 = FiniteElement("Lagrange", cell, 1, (), identity_pullback, H1)
+    elem1 = FiniteElement("Brezzi-Douglas-Marini", cell, 1, (2, ), contravariant_piola, HDiv)
+    elem2 = FiniteElement("Discontinuous Lagrange", cell, 0, (), identity_pullback, L2)
+    elem = MixedElement([elem0, elem1, elem2])
+    mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=100)
+    mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=101)
+    mesh2 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=102)
+    domain = MixedMesh(mesh0, mesh1, mesh2)
+    V = FunctionSpace(domain, elem)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    f = Coefficient(V, count=1000)
+    g = Coefficient(V, count=2000)
+    u0, u1, u2 = split(u)
+    v0, v1, v2 = split(v)
+    f0, f1, f2 = split(f)
+    g0, g1, g2 = split(g)
+    dx1 = Measure("dx", mesh1)
+    ds2 = Measure("ds", mesh2)
+    x = SpatialCoordinate(mesh1)
+    form = x[1] * f0 * inner(grad(u0), v1) * dx1(999) + div(f1) * g2 * inner(u1, grad(v2)) * ds2(888)
+    fd = compute_form_data(form,
+                           do_apply_function_pullbacks=True,
+                           do_apply_integral_scaling=True,
+                           do_apply_geometry_lowering=True,
+                           preserve_geometry_types=(CellVolume, FacetArea),
+                           do_apply_restrictions=True,
+                           do_estimate_degrees=True,
+                           complex_mode=False)
+    id0, id1 = fd.integral_data
+    assert fd.preprocessed_form.arguments() == (v, u)
+    assert fd.reduced_coefficients == [f, g]
+    assert form.coefficients()[fd.original_coefficient_positions[0]] is f
+    assert form.coefficients()[fd.original_coefficient_positions[1]] is g
+    assert id0.domain is mesh1
+    assert id0.integral_type == 'cell'
+    assert id0.subdomain_id == (999, )
+    assert fd.original_form.domain_numbering()[id0.domain] == 0
+    assert id0.integral_coefficients == set([f])
+    assert id0.enabled_coefficients == [True, False]
+    assert id1.domain is mesh2
+    assert id1.integral_type == 'exterior_facet'
+    assert id1.subdomain_id == (888, )
+    assert fd.original_form.domain_numbering()[id1.domain] == 1
+    assert id1.integral_coefficients == set([f, g])
+    assert id1.enabled_coefficients == [True, True]
+
+
+def test_mixed_function_space_with_mixed_mesh_restriction():
+    cell = triangle
+    elem0 = FiniteElement("Lagrange", cell, 1, (), identity_pullback, H1)
+    elem1 = FiniteElement("Brezzi-Douglas-Marini", cell, 1, (2, ), contravariant_piola, HDiv)
+    elem2 = FiniteElement("Discontinuous Lagrange", cell, 0, (), identity_pullback, L2)
+    elem = MixedElement([elem0, elem1, elem2])
+    mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=100)
+    mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=101)
+    mesh2 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=102)
+    domain = MixedMesh(mesh0, mesh1, mesh2)
+    V = FunctionSpace(domain, elem)
+    V1 = FunctionSpace(mesh1, elem1)
+    V2 = FunctionSpace(mesh2, elem2)
+    u1 = TrialFunction(V1)
+    v2 = TestFunction(V2)
+    f = Coefficient(V, count=1000)
+    g = Coefficient(V, count=2000)
+    f0, f1, f2 = split(f)
+    g0, g1, g2 = split(g)
+    dS1 = Measure("dS", mesh1)
+    x2 = SpatialCoordinate(mesh2)
+    form = inner(x2, g1) * g2 * inner(u1('-'), grad(v2('|'))) * dS1(999)
+    fd = compute_form_data(form,
+                           do_apply_function_pullbacks=True,
+                           do_apply_integral_scaling=True,
+                           do_apply_geometry_lowering=True,
+                           preserve_geometry_types=(CellVolume, FacetArea),
+                           do_apply_restrictions=True,
+                           do_estimate_degrees=True,
+                           do_split_coefficients=(f, g),
+                           do_assume_single_integral_type=False,
+                           complex_mode=False)
+    integral_data, = fd.integral_data
+    assert integral_data.domain_integral_type_map[mesh1] == "interior_facet"
+    assert integral_data.domain_integral_type_map[mesh2] == "exterior_facet"
+
+
+def test_mixed_function_space_with_mixed_mesh_signature():
+    cell = triangle
+    mesh0 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=100)
+    mesh1 = Mesh(FiniteElement("Lagrange", cell, 1, (2, ), identity_pullback, H1), ufl_id=101)
+    dx0 = Measure("dx", mesh0)
+    dx1 = Measure("dx", mesh1)
+    n0 = FacetNormal(mesh0)
+    n1 = FacetNormal(mesh1)
+    form_a = inner(n1, n1) * dx0(999)
+    form_b = inner(n0, n0) * dx1(999)
+    assert form_a.signature() == form_b.signature()
+    assert extract_domains(form_a) == (mesh0, mesh1)
+    assert extract_domains(form_b) == (mesh1, mesh0)

ufl/__init__.py

@@ -62,6 +62,7 @@
 
     - AbstractDomain
     - Mesh
+    - MixedMesh
     - MeshView
 
 * Sobolev spaces::
@@ -256,7 +257,7 @@
 from ufl.core.external_operator import ExternalOperator
 from ufl.core.interpolate import Interpolate, interpolate
 from ufl.core.multiindex import Index, indices
-from ufl.domain import AbstractDomain, Mesh, MeshView
+from ufl.domain import AbstractDomain, Mesh, MixedMesh, MeshView
 from ufl.finiteelement import AbstractFiniteElement
 from ufl.form import BaseForm, Form, FormSum, ZeroBaseForm
 from ufl.formoperators import (action, adjoint, derivative, energy_norm, extract_blocks, functional, lhs, replace, rhs,
@@ -287,7 +288,7 @@
 __all__ = [
     'product',
     'as_cell', 'AbstractCell', 'Cell', 'TensorProductCell',
-    'AbstractDomain', 'Mesh', 'MeshView',
+    'AbstractDomain', 'Mesh', 'MixedMesh', 'MeshView',
     'L2', 'H1', 'H2', 'HCurl', 'HDiv', 'HInf', 'HEin', 'HDivDiv',
     'identity_pullback', 'l2_piola', 'contravariant_piola', 'covariant_piola',
     'double_contravariant_piola', 'double_covariant_piola',

ufl/action.py

@@ -188,8 +188,8 @@ def _get_action_form_arguments(left, right):
     elif isinstance(right, CoefficientDerivative):
         # Action differentiation pushes differentiation through
         # right as a consequence of Leibniz formula.
-        from ufl.algorithms.analysis import extract_arguments_and_coefficients
-        right_args, right_coeffs = extract_arguments_and_coefficients(right)
+        from ufl.algorithms.analysis import extract_arguments_and_coefficients_and_geometric_quantities
+        right_args, right_coeffs, _ = extract_arguments_and_coefficients_and_geometric_quantities(right)
         arguments = left_args + tuple(right_args)
         coefficients += tuple(right_coeffs)
     elif isinstance(right, (BaseCoefficient, Zero)):

ufl/algorithms/analysis.py

@@ -12,9 +12,11 @@
 from itertools import chain
 
 from ufl.algorithms.traversal import iter_expressions
+from ufl.domain import Mesh
 from ufl.argument import BaseArgument, Coargument
 from ufl.coefficient import BaseCoefficient
 from ufl.constant import Constant
+from ufl.geometry import GeometricQuantity
 from ufl.core.base_form_operator import BaseFormOperator
 from ufl.core.terminal import Terminal
 from ufl.corealg.traversal import traverse_unique_terminals, unique_pre_traversal
@@ -187,19 +189,20 @@ def extract_base_form_operators(a):
     return sorted_by_count(extract_type(a, BaseFormOperator))
 
 
-def extract_arguments_and_coefficients(a):
-    """Build two sorted lists of all arguments and coefficients in a.
+def extract_arguments_and_coefficients_and_geometric_quantities(a):
+    """Build three sorted lists of all arguments, coefficients, and geometric quantities in a.
 
-    This function is faster than extract_arguments + extract_coefficients
+    This function is faster than extract_arguments + extract_coefficients + extract_geometric_quantities
     for large forms, and has more validation built in.
 
     Args:
         a: A BaseForm, Integral or Expr
     """
     # Extract lists of all BaseArgument and BaseCoefficient instances
-    base_coeff_and_args = extract_type(a, (BaseArgument, BaseCoefficient))
-    arguments = [f for f in base_coeff_and_args if isinstance(f, BaseArgument)]
-    coefficients = [f for f in base_coeff_and_args if isinstance(f, BaseCoefficient)]
+    base_coeff_and_args_and_gq = extract_type(a, (BaseArgument, BaseCoefficient, GeometricQuantity))
+    arguments = [f for f in base_coeff_and_args_and_gq if isinstance(f, BaseArgument)]
+    coefficients = [f for f in base_coeff_and_args_and_gq if isinstance(f, BaseCoefficient)]
+    geometric_quantities = [f for f in base_coeff_and_args_and_gq if isinstance(f, GeometricQuantity)]
 
     # Build number,part: instance mappings, should be one to one
     bfnp = dict((f, (f.number(), f.part())) for f in arguments)
@@ -214,19 +217,31 @@ def extract_arguments_and_coefficients(a):
     if len(fcounts) != len(set(fcounts.values())):
         raise ValueError(
             "Found different coefficients with same counts.\n"
-            "The arguments found are:\n" + "\n".join(f"  {c}" for c in coefficients))
+            "The Coefficients found are:\n" + "\n".join(f"  {c}" for c in coefficients))
+
+    # Build count: instance mappings, should be one to one
+    gqcounts = {}
+    for gq in geometric_quantities:
+        if not isinstance(gq._domain, Mesh):
+            raise TypeError(f"{gq}._domain must be a Mesh: got {gq._domain}")
+        gqcounts[gq] = (type(gq).name, gq._domain._ufl_id)
+    if len(gqcounts) != len(set(gqcounts.values())):
+        raise ValueError(
+            "Found different geometric quantities with same (geometric_quantity_type, domain).\n"
+            "The GeometricQuantities found are:\n" + "\n".join(f"  {gq}" for gq in geometric_quantities))
 
     # Passed checks, so we can safely sort the instances by count
     arguments = _sorted_by_number_and_part(arguments)
     coefficients = sorted_by_count(coefficients)
+    geometric_quantities = list(sorted(geometric_quantities, key=lambda gq: (type(gq).name, gq._domain._ufl_id)))
 
-    return arguments, coefficients
+    return arguments, coefficients, geometric_quantities
 
 
 def extract_elements(form):
     """Build sorted tuple of all elements used in form."""
-    args = chain(*extract_arguments_and_coefficients(form))
-    return tuple(f.ufl_element() for f in args)
+    arguments, coefficients, _ = extract_arguments_and_coefficients_and_geometric_quantities(form)
+    return tuple(f.ufl_element() for f in arguments + coefficients)
 
 
 def extract_unique_elements(form):