diff --git a/test/test_apply_function_pullbacks.py b/test/test_apply_function_pullbacks.py index 58a6e9762..dbebddfc6 100755 --- a/test/test_apply_function_pullbacks.py +++ b/test/test_apply_function_pullbacks.py @@ -1,9 +1,35 @@ import numpy as np +import pytest from utils import FiniteElement, LagrangeElement, MixedElement, SymmetricElement -from ufl import Cell, Coefficient, FunctionSpace, Mesh, as_tensor, as_vector, dx, indices, triangle +from ufl import ( + Cell, + Coefficient, + FunctionSpace, + Mesh, + SpatialCoordinate, + TestFunction, + TrialFunction, + as_tensor, + as_vector, + dx, + grad, + indices, + inner, + triangle, +) +from ufl.algorithms.analysis import extract_type +from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering +from ufl.algorithms.apply_function_pullbacks import ( + apply_function_pullbacks, + apply_interpolate_pullbacks, + apply_inverse_pullback, +) +from ufl.algorithms.cancel_jacobian_products import cancel_jacobian_products +from ufl.algorithms.compute_form_data import compute_form_data +from ufl.algorithms.remove_component_tensors import remove_component_tensors from ufl.algorithms.renumbering import renumber_indices -from ufl.classes import Jacobian, JacobianDeterminant, JacobianInverse, ReferenceValue +from ufl.classes import Interpolate, Jacobian, JacobianDeterminant, JacobianInverse, ReferenceValue from ufl.pullback import ( contravariant_piola, covariant_piola, @@ -12,7 +38,7 @@ identity_pullback, l2_piola, ) -from ufl.sobolevspace import L2, HCurl, HDiv, HDivDiv, HEin +from ufl.sobolevspace import H1, L2, HCurl, HDiv, HDivDiv, HEin def check_single_function_pullback(g, mappings): @@ -495,3 +521,171 @@ def test_apply_single_function_pullbacks_triangle(): # Check the ridiculous mixed element W combining it all check_single_function_pullback(w, mappings) + + +def cancel_jacobians(expr): + """Cancel the Jacobian factors that a pull back and its inverse insert.""" + expr = apply_algebra_lowering(expr) + expr = remove_component_tensors(expr) + expr = cancel_jacobian_products(expr) + return renumber_indices(remove_component_tensors(expr)) + + +def check_inverse_pullback(element, domain): + """Map a physical function onto the reference cell and pull it back again.""" + physical_value_shape = element.pullback.physical_value_shape(element, domain) + P = FiniteElement("Lagrange", domain.ufl_cell(), 2, physical_value_shape, identity_pullback, H1) + g = Coefficient(FunctionSpace(domain, P)) + r = apply_inverse_pullback(g, element, domain) + assert r.ufl_shape == element.reference_value_shape + back = element.pullback.apply(r, domain) + assert back.ufl_shape == physical_value_shape + return g, back + + +def test_apply_inverse_pullback_inverts_the_pullback(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + vd = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv) + u0 = FiniteElement("Discontinuous Lagrange", cell, 0, (), identity_pullback, L2) + elements = [ + LagrangeElement(cell, 1, (2,)), + vd, + FiniteElement("N1curl", cell, 1, (2,), covariant_piola, HCurl), + FiniteElement("Regge", cell, 1, (2, 2), double_contravariant_piola, HDivDiv), + MixedElement([vd, u0]), + ] + for element in elements: + g, back = check_inverse_pullback(element, domain) + assert cancel_jacobians(back) == g + + +def test_apply_inverse_pullback_symmetric_element(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + u1 = FiniteElement("Discontinuous Lagrange", cell, 1, (), identity_pullback, L2) + element = SymmetricElement({(0, 0): 0, (0, 1): 1, (1, 0): 1, (1, 1): 2}, [u1, u1, u1]) + # The element carries the independent components only, so the map onto the + # reference cell is invertible on a symmetric function alone. + U = FunctionSpace(domain, u1) + a, b, c = Coefficient(U), Coefficient(U), Coefficient(U) + g = as_tensor([[a, b], [b, c]]) + r = apply_inverse_pullback(g, element, domain) + assert r.ufl_shape == element.reference_value_shape + back = element.pullback.apply(r, domain) + assert cancel_jacobians(back) == cancel_jacobians(g) + + +def test_apply_inverse_pullback_checks_shapes(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + element = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv) + scalar = Coefficient(FunctionSpace(domain, LagrangeElement(cell, 1))) + with pytest.raises(ValueError): + apply_inverse_pullback(scalar, element, domain) + + other_domain = Mesh(LagrangeElement(cell, 1, (2,))) + g = Coefficient(FunctionSpace(other_domain, LagrangeElement(cell, 1, (2,)))) + with pytest.raises(NotImplementedError): + apply_inverse_pullback(g, element, domain) + + +def test_apply_interpolate_pullbacks(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + element = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv) + V = FunctionSpace(domain, element) + W = FunctionSpace(domain, LagrangeElement(cell, 2, (2,))) + f = Coefficient(W) + + expr = Interpolate(f, V) + r = apply_interpolate_pullbacks(expr) + (interpolation,) = extract_type(r, Interpolate) + # The interpolation is evaluated on the reference cell of its target + # element, and its result pulled back for the expression that holds it. + assert renumber_indices(interpolation.ufl_operands[0]) == renumber_indices( + apply_inverse_pullback(f, element, domain) + ) + assert renumber_indices(r) == renumber_indices( + element.pullback.apply(ReferenceValue(interpolation), domain) + ) + # The dual argument and the physical value shape are those of the original. + assert interpolation.argument_slots()[0] is expr.argument_slots()[0] + assert r.ufl_shape == expr.ufl_shape + + +def test_apply_interpolate_pullbacks_identity_pullback(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + W = FunctionSpace(domain, LagrangeElement(cell, 2, (2,))) + f = Coefficient(W) + # Nothing to map: the interpolation is only taken to its reference value. + assert apply_interpolate_pullbacks(Interpolate(f, W)) == ReferenceValue(Interpolate(f, W)) + + +def test_apply_interpolate_pullbacks_nested(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + element = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv) + V = FunctionSpace(domain, element) + W = FunctionSpace(domain, LagrangeElement(cell, 2, (2,))) + f = Coefficient(W) + + r = apply_interpolate_pullbacks(Interpolate(Interpolate(f, W), V)) + interpolations = extract_type(r, Interpolate) + assert len(interpolations) == 2 + for interpolation in interpolations: + (operand,) = interpolation.ufl_operands + assert operand.ufl_shape == interpolation.ufl_element().reference_value_shape + + +def test_apply_interpolate_pullbacks_in_a_form(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + element = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv) + V = FunctionSpace(domain, element) + W = FunctionSpace(domain, LagrangeElement(cell, 2, (2,))) + f = Coefficient(W) + v = TestFunction(V) + + form = inner(Interpolate(f, V), v) * dx + r = apply_interpolate_pullbacks(form) + assert r.arguments() == form.arguments() + (integrand,) = (itg.integrand() for itg in r.integrals()) + (interpolation,) = extract_type(integrand, Interpolate) + (operand,) = interpolation.ufl_operands + assert operand.ufl_shape == element.reference_value_shape + # The lowered form still goes through the rest of the pipeline. + compute_form_data(form, do_apply_function_pullbacks=True, do_apply_geometry_lowering=True) + + +def test_apply_interpolate_pullbacks_reuses_untouched_expressions(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + f = Coefficient(FunctionSpace(domain, LagrangeElement(cell, 2))) + expr = inner(grad(f), grad(f)) + assert apply_interpolate_pullbacks(expr) is expr + + +def test_apply_function_pullbacks_in_a_form(): + cell = triangle + domain = Mesh(LagrangeElement(cell, 1, (2,))) + element = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv) + V = FunctionSpace(domain, element) + form = inner(TrialFunction(V), TestFunction(V)) * dx + + r = apply_function_pullbacks(form) + assert r.arguments() == form.arguments() + (integrand,) = (itg.integrand() for itg in r.integrals()) + # Every form argument is now reached through its reference value. + reference_values = extract_type(integrand, ReferenceValue) + assert {rv.ufl_operands[0] for rv in reference_values} == set(form.arguments()) + # Each of them is replaced by the pull back of that reference value. + for w in form.arguments(): + assert renumber_indices(apply_function_pullbacks(w)) == renumber_indices( + element.pullback.apply(ReferenceValue(w)) + ) + + # Terminals that are not form arguments are left alone. + x = SpatialCoordinate(domain) + assert apply_function_pullbacks(x) is x diff --git a/ufl/algorithms/apply_function_pullbacks.py b/ufl/algorithms/apply_function_pullbacks.py index 124c2a2dd..5736d9635 100644 --- a/ufl/algorithms/apply_function_pullbacks.py +++ b/ufl/algorithms/apply_function_pullbacks.py @@ -6,29 +6,41 @@ # # SPDX-License-Identifier: LGPL-3.0-or-later -from ufl.algorithms.map_integrands import map_integrand_dags -from ufl.classes import ReferenceValue -from ufl.corealg.multifunction import MultiFunction, memoized_handler +from __future__ import annotations +from functools import singledispatchmethod -class FunctionPullbackApplier(MultiFunction): +from ufl.algorithms.map_integrands import map_integrands +from ufl.classes import Argument, Coefficient, Expr, FormArgument, Interpolate, ReferenceValue +from ufl.corealg.dag_traverser import DAGTraverser +from ufl.domain import extract_unique_domain +from ufl.form import BaseForm + + +class FunctionPullbackApplier(DAGTraverser): """A pull back applier.""" - def __init__(self): - """Initalise.""" - MultiFunction.__init__(self) + @singledispatchmethod + def process(self, o: Expr) -> Expr: + """Process ``o``. + + Args: + o: `Expr` to be processed. + + Returns: + Processed `Expr`. - expr = MultiFunction.reuse_if_untouched + """ + return super().process(o) - def terminal(self, t): - """Apply to a terminal.""" - return t + @process.register(Expr) + def _(self, o: Expr) -> Expr: + """Handle Expr.""" + return self.reuse_if_untouched(o) - @memoized_handler - def form_argument(self, o): - """Apply to a form_argument.""" - # Represent 0-derivatives of form arguments on reference - # element + @process.register(FormArgument) + def _(self, o: Argument | Coefficient) -> Expr: + """Represent 0-derivatives of a form argument on the reference element.""" r = ReferenceValue(o) space = o.ufl_function_space() element = o.ufl_element() @@ -49,6 +61,92 @@ def form_argument(self, o): return f +class InterpolatePullbackApplier(DAGTraverser): + """A pull back applier for interpolation.""" + + @singledispatchmethod + def process(self, o: Expr | BaseForm) -> Expr | BaseForm: + """Process ``o``. + + Args: + o: `Expr` or `BaseForm` to be processed. + + Returns: + Processed `Expr` or `BaseForm`. + + """ + return super().process(o) + + @process.register(Expr) + @process.register(BaseForm) + def _(self, o: Expr | BaseForm) -> Expr | BaseForm: + """Handle Expr and BaseForm.""" + return self.reuse_if_untouched(o) + + @process.register(Interpolate) + @DAGTraverser.postorder + def _(self, o: Interpolate, operand: Expr) -> Expr: + """Evaluate an interpolation on the reference cell of its target element.""" + dual_arg, _ = o.argument_slots() + element = o.ufl_element() + domain = extract_unique_domain(operand) or dual_arg.ufl_function_space().ufl_domain() + # Build the node here rather than reconstructing o: the mapped operand + # no longer has the physical value shape that a subclass may check. + r = Interpolate(apply_inverse_pullback(operand, element, domain), dual_arg) + return element.pullback.apply(ReferenceValue(r), domain) + + +def apply_inverse_pullback(expr, element, domain=None): + """Map a physical expression onto the reference cell of an element. + + This is a rule on a single node, not a DAG traversal: the expression is + mapped as a whole, and an interpolation inside it is left alone for + `apply_interpolate_pullbacks` to lower. + + Args: + expr: An expression on a physical cell, whose shape must be the + physical value shape of the element + element: The element whose pull back is inverted + domain: The domain to use if the expression carries none + + Returns: + The expression on the reference cell, with shape + ``element.reference_value_shape`` + """ + mesh = extract_unique_domain(expr) or domain + if domain is not None and mesh != domain: + raise NotImplementedError("Multiple domains not supported") + physical_value_shape = element.pullback.physical_value_shape(element, mesh) + if expr.ufl_shape != physical_value_shape: + raise ValueError( + f"Expecting physical expression with shape '{physical_value_shape}', " + f"got '{expr.ufl_shape}'" + ) + r = element.pullback.apply_inverse(expr, mesh) + if r.ufl_shape != element.reference_value_shape: + raise ValueError( + f"Expecting reference expression with shape " + f"'{element.reference_value_shape}', got '{r.ufl_shape}'" + ) + return r + + +def apply_interpolate_pullbacks(expr): + """Change the representation of the interpolations in an expression. + + An interpolation is evaluated on the reference cell of its target element, + so its operand is mapped there and the result is pulled back to the + physical cell for the expression that holds it. + + Args: + expr: An Expr or Form + + Returns: + The expression with its interpolations on their reference cells + """ + return map_integrands(InterpolatePullbackApplier(), expr) + + def apply_function_pullbacks(expr): """Change representation of coefficients and arguments in an expression. @@ -58,4 +156,4 @@ def apply_function_pullbacks(expr): Args: expr: An Expression """ - return map_integrand_dags(FunctionPullbackApplier(), expr) + return map_integrands(FunctionPullbackApplier(), expr) diff --git a/ufl/algorithms/compute_form_data.py b/ufl/algorithms/compute_form_data.py index 85c8a6014..329b35c9e 100644 --- a/ufl/algorithms/compute_form_data.py +++ b/ufl/algorithms/compute_form_data.py @@ -13,7 +13,10 @@ # These are the main symbolic processing steps: from ufl.algorithms.apply_derivatives import apply_coordinate_derivatives, apply_derivatives -from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks +from ufl.algorithms.apply_function_pullbacks import ( + apply_function_pullbacks, + apply_interpolate_pullbacks, +) from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering from ufl.algorithms.apply_integral_scaling import apply_integral_scaling from ufl.algorithms.cancel_jacobian_products import cancel_jacobian_products @@ -139,6 +142,12 @@ def compute_form_data( # See 'reduced_coefficients' below. original_form = form + # Evaluate interpolations on the reference cell of their target element. + # This happens before any other lowering, so that the operand a form + # compiler sees is the one the target element dual-evaluates. + if do_apply_function_pullbacks: + form = apply_interpolate_pullbacks(form) + # --- Pass form integrands through some symbolic manipulation form = preprocess_form(form, complex_mode) diff --git a/ufl/core/interpolate.py b/ufl/core/interpolate.py index 4c56901e7..1ff9a51c2 100644 --- a/ufl/core/interpolate.py +++ b/ufl/core/interpolate.py @@ -181,9 +181,7 @@ def _value_parent_type(self): def __neg__(self): """Negate the interpolation result.""" - if self._value_parent_type() is Operator: - return Operator.__rmul__(self, -1) - return BaseForm.__neg__(self) + return self._value_parent_type().__rmul__(self, -1) def __mul__(self, other): """Multiply, agreeing with negation on which space the value is in."""