Merge pull request #233 from firedrakeproject/ReubenHill/tsfc-interp-tidy

Tidy up interpolation tsfc<->firedrake interface

Author: dham

tests/test_dual_evaluation.py

@@ -0,0 +1,61 @@
+import pytest
+import ufl
+from tsfc.finatinterface import create_element
+from tsfc import compile_expression_dual_evaluation
+
+
+def test_ufl_only_simple():
+    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    v = ufl.Coefficient(V)
+    expr = ufl.inner(v, v)
+    W = V
+    to_element = create_element(W.ufl_element())
+    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)
+    assert first_coeff_fake_coords is False
+
+
+def test_ufl_only_spatialcoordinate():
+    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    x, y = ufl.SpatialCoordinate(mesh)
+    expr = x*y - y**2 + x
+    W = V
+    to_element = create_element(W.ufl_element())
+    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)
+    assert first_coeff_fake_coords is True
+
+
+def test_ufl_only_from_contravariant_piola():
+    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
+    v = ufl.Coefficient(V)
+    expr = ufl.inner(v, v)
+    W = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    to_element = create_element(W.ufl_element())
+    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)
+    assert first_coeff_fake_coords is True
+
+
+def test_ufl_only_to_contravariant_piola():
+    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
+    v = ufl.Coefficient(V)
+    expr = ufl.as_vector([v, v])
+    W = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
+    to_element = create_element(W.ufl_element())
+    ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)
+    assert first_coeff_fake_coords is True
+
+
+def test_ufl_only_shape_mismatch():
+    mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
+    V = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
+    v = ufl.Coefficient(V)
+    expr = ufl.inner(v, v)
+    assert expr.ufl_shape == ()
+    W = V
+    to_element = create_element(W.ufl_element())
+    assert to_element.value_shape == (2,)
+    with pytest.raises(ValueError):
+        ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, _ = compile_expression_dual_evaluation(expr, to_element, coffee=False)

tsfc/driver.py

@@ -267,7 +267,7 @@ def name_multiindex(multiindex, name):
     return builder.construct_kernel(kernel_name, impero_c, index_names, quad_rule)
 
 
-def compile_expression_dual_evaluation(expression, to_element, coordinates, *,
+def compile_expression_dual_evaluation(expression, to_element, *,
                                        domain=None, interface=None,
                                        parameters=None, coffee=False):
     """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis.
@@ -276,8 +276,7 @@ def compile_expression_dual_evaluation(expression, to_element, coordinates, *,
 
     :arg expression: UFL expression
     :arg to_element: A FInAT element for the target space
-    :arg coordinates: the coordinate function
-    :arg domain: optional UFL domain the expression is defined on (useful when expression contains no domain).
+    :arg domain: optional UFL domain the expression is defined on (required when expression contains no domain).
     :arg interface: backend module for the kernel interface
     :arg parameters: parameters object
     :arg coffee: compile coffee kernel instead of loopy kernel
@@ -328,25 +327,29 @@ def compile_expression_dual_evaluation(expression, to_element, coordinates, *,
     argument_multiindices = tuple(builder.create_element(arg.ufl_element()).get_indices()
                                   for arg in arguments)
 
-    # Replace coordinates (if any)
-    domain = expression.ufl_domain()
-    if domain:
-        assert coordinates.ufl_domain() == domain
-        builder.domain_coordinate[domain] = coordinates
-        builder.set_cell_sizes(domain)
+    # Replace coordinates (if any) unless otherwise specified by kwarg
+    if domain is None:
+        domain = expression.ufl_domain()
+    assert domain is not None
 
     # Collect required coefficients
+    first_coefficient_fake_coords = False
     coefficients = extract_coefficients(expression)
     if has_type(expression, GeometricQuantity) or any(fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients):
-        coefficients = [coordinates] + coefficients
+        # Create a fake coordinate coefficient for a domain.
+        coords_coefficient = ufl.Coefficient(ufl.FunctionSpace(domain, domain.ufl_coordinate_element()))
+        builder.domain_coordinate[domain] = coords_coefficient
+        builder.set_cell_sizes(domain)
+        coefficients = [coords_coefficient] + coefficients
+        first_coefficient_fake_coords = True
     builder.set_coefficients(coefficients)
 
     # Split mixed coefficients
     expression = ufl_utils.split_coefficients(expression, builder.coefficient_split)
 
     # Translate to GEM
     kernel_cfg = dict(interface=builder,
-                      ufl_cell=coordinates.ufl_domain().ufl_cell(),
+                      ufl_cell=domain.ufl_cell(),
                       argument_multiindices=argument_multiindices,
                       index_cache={},
                       scalar_type=parameters["scalar_type"])
@@ -431,7 +434,7 @@ def compile_expression_dual_evaluation(expression, to_element, coordinates, *,
     # Handle kernel interface requirements
     builder.register_requirements([ir])
     # Build kernel tuple
-    return builder.construct_kernel(return_arg, impero_c, index_names)
+    return builder.construct_kernel(return_arg, impero_c, index_names, first_coefficient_fake_coords)
 
 
 def lower_integral_type(fiat_cell, integral_type):

tsfc/kernel_interface/firedrake_loopy.py

@@ -17,7 +17,7 @@
 
 
 # Expression kernel description type
-ExpressionKernel = namedtuple('ExpressionKernel', ['ast', 'oriented', 'needs_cell_sizes', 'coefficients', 'tabulations'])
+ExpressionKernel = namedtuple('ExpressionKernel', ['ast', 'oriented', 'needs_cell_sizes', 'coefficients', 'first_coefficient_fake_coords', 'tabulations'])
 
 
 def make_builder(*args, **kwargs):
@@ -153,12 +153,14 @@ def register_requirements(self, ir):
         provided by the kernel interface."""
         self.oriented, self.cell_sizes, self.tabulations = check_requirements(ir)
 
-    def construct_kernel(self, return_arg, impero_c, index_names):
+    def construct_kernel(self, return_arg, impero_c, index_names, first_coefficient_fake_coords):
         """Constructs an :class:`ExpressionKernel`.
 
         :arg return_arg: loopy.GlobalArg for the return value
         :arg impero_c: gem.ImperoC object that represents the kernel
         :arg index_names: pre-assigned index names
+        :arg first_coefficient_fake_coords: If true, the kernel's first
+            coefficient is a constructed UFL coordinate field
         :returns: :class:`ExpressionKernel` object
         """
         args = [return_arg]
@@ -173,7 +175,8 @@ def construct_kernel(self, return_arg, impero_c, index_names):
         loopy_kernel = generate_loopy(impero_c, args, self.scalar_type,
                                       "expression_kernel", index_names)
         return ExpressionKernel(loopy_kernel, self.oriented, self.cell_sizes,
-                                self.coefficients, self.tabulations)
+                                self.coefficients, first_coefficient_fake_coords,
+                                self.tabulations)
 
 
 class KernelBuilder(KernelBuilderBase):