Merge pull request #262 from firedrakeproject/wence/feature/interpola…

…tion-mappings

Update interpolation API

tests/test_dual_evaluation.py

@@ -11,8 +11,8 @@ def test_ufl_only_simple():
     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)
-    assert first_coeff_fake_coords is False
+    kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element())
+    assert kernel.first_coefficient_fake_coords is False
 
 
 def test_ufl_only_spatialcoordinate():
@@ -22,8 +22,8 @@ def test_ufl_only_spatialcoordinate():
     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)
-    assert first_coeff_fake_coords is True
+    kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element())
+    assert kernel.first_coefficient_fake_coords is True
 
 
 def test_ufl_only_from_contravariant_piola():
@@ -33,8 +33,8 @@ def test_ufl_only_from_contravariant_piola():
     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)
-    assert first_coeff_fake_coords is True
+    kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element())
+    assert kernel.first_coefficient_fake_coords is True
 
 
 def test_ufl_only_to_contravariant_piola():
@@ -44,8 +44,8 @@ def test_ufl_only_to_contravariant_piola():
     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)
-    assert first_coeff_fake_coords is True
+    kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element())
+    assert kernel.first_coefficient_fake_coords is True
 
 
 def test_ufl_only_shape_mismatch():
@@ -58,4 +58,4 @@ def test_ufl_only_shape_mismatch():
     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)
+        compile_expression_dual_evaluation(expr, to_element, W.ufl_element())

tests/test_interpolation_factorisation.py

@@ -6,7 +6,7 @@
                  TensorElement, Coefficient,
                  interval, quadrilateral, hexahedron)
 
-from tsfc.driver import compile_expression_dual_evaluation
+from tsfc import compile_expression_dual_evaluation
 from tsfc.finatinterface import create_element
 
 
@@ -31,7 +31,7 @@ def flop_count(mesh, source, target):
     Vsource = FunctionSpace(mesh, source)
     to_element = create_element(Vtarget.ufl_element())
     expr = Coefficient(Vsource)
-    kernel = compile_expression_dual_evaluation(expr, to_element)
+    kernel = compile_expression_dual_evaluation(expr, to_element, Vtarget.ufl_element())
     return kernel.flop_count
 
 

tsfc/driver.py

@@ -5,6 +5,7 @@
 import sys
 from functools import reduce
 from itertools import chain
+from finat.physically_mapped import DirectlyDefinedElement, PhysicallyMappedElement
 
 from numpy import asarray
 
@@ -265,7 +266,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, *,
+def compile_expression_dual_evaluation(expression, to_element, ufl_element, *,
                                        domain=None, interface=None,
                                        parameters=None):
     """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis.
@@ -274,6 +275,7 @@ def compile_expression_dual_evaluation(expression, to_element, *,
 
     :arg expression: UFL expression
     :arg to_element: A FInAT element for the target space
+    :arg ufl_element: The UFL element of the target space.
     :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
@@ -289,12 +291,10 @@ def compile_expression_dual_evaluation(expression, to_element, *,
     # Determine whether in complex mode
     complex_mode = is_complex(parameters["scalar_type"])
 
-    # Find out which mapping to apply
-    try:
-        mapping, = set((to_element.mapping,))
-    except ValueError:
+    if isinstance(to_element, (PhysicallyMappedElement, DirectlyDefinedElement)):
         raise NotImplementedError("Don't know how to interpolate onto zany spaces, sorry")
-    expression = apply_mapping(expression, mapping, domain)
+    # Map into reference space
+    expression = apply_mapping(expression, ufl_element, domain)
 
     # Apply UFL preprocessing
     expression = ufl_utils.preprocess_expression(expression,

tsfc/finatinterface.py

@@ -209,20 +209,13 @@ def convert_mixedelement(element, **kwargs):
 
 
 @convert.register(ufl.VectorElement)
-def convert_vectorelement(element, **kwargs):
-    scalar_elem, deps = _create_element(element.sub_elements()[0], **kwargs)
-    shape = (element.num_sub_elements(),)
-    shape_innermost = kwargs["shape_innermost"]
-    return (finat.TensorFiniteElement(scalar_elem, shape, not shape_innermost),
-            deps | {"shape_innermost"})
-
-
 @convert.register(ufl.TensorElement)
 def convert_tensorelement(element, **kwargs):
-    scalar_elem, deps = _create_element(element.sub_elements()[0], **kwargs)
+    inner_elem, deps = _create_element(element.sub_elements()[0], **kwargs)
     shape = element.reference_value_shape()
+    shape = shape[:len(shape) - len(inner_elem.value_shape)]
     shape_innermost = kwargs["shape_innermost"]
-    return (finat.TensorFiniteElement(scalar_elem, shape, not shape_innermost),
+    return (finat.TensorFiniteElement(inner_elem, shape, not shape_innermost),
             deps | {"shape_innermost"})
 
 

tsfc/ufl_utils.py

@@ -22,7 +22,7 @@
 from ufl.geometry import Jacobian, JacobianDeterminant, JacobianInverse
 from ufl.classes import (Abs, Argument, CellOrientation, Coefficient,
                          ComponentTensor, Expr, FloatValue, Division,
-                         Indexed, MixedElement, MultiIndex, Product,
+                         MixedElement, MultiIndex, Product,
                          ScalarValue, Sqrt, Zero, CellVolume, FacetArea)
 
 from gem.node import MemoizerArg
@@ -345,12 +345,19 @@ def simplify_abs(expression, complex_mode):
     return mapper(expression, False)
 
 
-def apply_mapping(expression, mapping, domain):
-    """
-    This applies the appropriate transformation to the
-    given expression for interpolation to a specific
-    element, according to the manner in which it maps
-    from the reference cell.
+def apply_mapping(expression, element, domain):
+    """Apply the inverse of the pullback for element to an expression.
+
+    :arg expression: An expression in physical space
+    :arg element: The element we're going to interpolate into, whose
+         value_shape must match the shape of the expression, and will
+         advertise the pullback to apply.
+    :arg domain: Optional domain to provide in case expression does
+         not contain a domain (used for constructing geometric quantities).
+    :returns: A new UFL expression with shape element.reference_value_shape()
+    :raises NotImplementedError: If we don't know how to apply the
+        inverse of the pullback.
+    :raises ValueError: If we get shape mismatches.
 
     The following is borrowed from the UFC documentation:
 
@@ -364,7 +371,7 @@ def apply_mapping(expression, mapping, domain):
     inverse of the Jacobian K = J^{-1}. Then we (currently) have the
     following four types of mappings:
 
-    'affine' mapping for g:
+    'identity' mapping for g:
 
       G(X) = g(x)
 
@@ -386,38 +393,76 @@ def apply_mapping(expression, mapping, domain):
 
       G(X) = det(J)^2 K g(x) K^T  i.e. G_il(X)=(detJ)^2 K_ij g_jk K_lk
 
-    If 'contravariant piola' or 'covariant piola' are applied to a
-    matrix-valued function, the appropriate mappings are applied row-by-row.
-
-    :arg expression: UFL expression
-    :arg mapping: a string indicating the mapping to apply
+    If 'contravariant piola' or 'covariant piola' (or their double
+    variants) are applied to a matrix-valued function, the appropriate
+    mappings are applied row-by-row.
     """
-
     mesh = expression.ufl_domain()
     if mesh is None:
         mesh = domain
     if domain is not None and mesh != domain:
         raise NotImplementedError("Multiple domains not supported")
-    rank = len(expression.ufl_shape)
-    if mapping == "affine":
-        return expression
+    if expression.ufl_shape != element.value_shape():
+        raise ValueError(f"Mismatching shapes, got {expression.ufl_shape}, expected {element.value_shape()}")
+    mapping = element.mapping().lower()
+    if mapping == "identity":
+        rexpression = expression
     elif mapping == "covariant piola":
         J = Jacobian(mesh)
-        *i, j, k = indices(len(expression.ufl_shape) + 1)
-        expression = Indexed(expression, MultiIndex((*i, k)))
-        return as_tensor(J.T[j, k] * expression, (*i, j))
+        *k, i, j = indices(len(expression.ufl_shape) + 1)
+        kj = (*k, j)
+        rexpression = as_tensor(J[j, i] * expression[kj], (*k, i))
+    elif mapping == "l2 piola":
+        detJ = JacobianDeterminant(mesh)
+        rexpression = expression * detJ
     elif mapping == "contravariant piola":
         K = JacobianInverse(mesh)
         detJ = JacobianDeterminant(mesh)
-        *i, j, k = indices(len(expression.ufl_shape) + 1)
-        expression = Indexed(expression, MultiIndex((*i, k)))
-        return as_tensor(detJ * K[j, k] * expression, (*i, j))
-    elif mapping == "double covariant piola" and rank == 2:
+        *k, i, j = indices(len(expression.ufl_shape) + 1)
+        kj = (*k, j)
+        rexpression = as_tensor(detJ * K[i, j] * expression[kj], (*k, i))
+    elif mapping == "double covariant piola":
         J = Jacobian(mesh)
-        return J.T * expression * J
-    elif mapping == "double contravariant piola" and rank == 2:
+        *k, i, j, m, n = indices(len(expression.ufl_shape) + 2)
+        kmn = (*k, m, n)
+        rexpression = as_tensor(J[m, i] * expression[kmn] * J[n, j], (*k, i, j))
+    elif mapping == "double contravariant piola":
         K = JacobianInverse(mesh)
         detJ = JacobianDeterminant(mesh)
-        return (detJ)**2 * K * expression * K.T
+        *k, i, j, m, n = indices(len(expression.ufl_shape) + 2)
+        kmn = (*k, m, n)
+        rexpression = as_tensor(detJ**2 * K[i, m] * expression[kmn] * K[j, n], (*k, i, j))
+    elif mapping == "symmetries":
+        # This tells us how to get from the pieces of the reference
+        # space expression to the physical space one.
+        # We're going to apply the inverse of the physical to
+        # reference space mapping.
+        fcm = element.flattened_sub_element_mapping()
+        sub_elem = element.sub_elements()[0]
+        shape = expression.ufl_shape
+        flat = ufl.as_vector([expression[i] for i in numpy.ndindex(shape)])
+        vs = sub_elem.value_shape()
+        rvs = sub_elem.reference_value_shape()
+        seen = set()
+        rpieces = []
+        gm = int(numpy.prod(vs, dtype=int))
+        for gi, ri in enumerate(fcm):
+            # For each unique piece in reference space
+            if ri in seen:
+                continue
+            seen.add(ri)
+            # Get the physical space piece
+            piece = [flat[gm*gi + j] for j in range(gm)]
+            piece = as_tensor(numpy.asarray(piece).reshape(vs))
+            # get into reference space
+            piece = apply_mapping(piece, sub_elem, mesh)
+            assert piece.ufl_shape == rvs
+            # Concatenate with the other pieces
+            rpieces.extend([piece[idx] for idx in numpy.ndindex(rvs)])
+        # And reshape
+        rexpression = as_tensor(numpy.asarray(rpieces).reshape(element.reference_value_shape()))
     else:
-        raise NotImplementedError("Don't know how to handle mapping type %s for expression of rank %d" % (mapping, rank))
+        raise NotImplementedError(f"Don't know how to handle mapping type {mapping} for expression of rank {element.value_shape()}")
+    if rexpression.ufl_shape != element.reference_value_shape():
+        raise ValueError(f"Mismatching reference shapes, got {rexpression.ufl_shape} expected {element.reference_value_shape()}")
+    return rexpression