Give each pull back an inverse

test/test_pullback_inverse.py

@@ -0,0 +1,98 @@
+"""Tests of the inverse pull backs."""
+
+import numpy as np
+import pytest
+from utils import FiniteElement, LagrangeElement, MixedElement, SymmetricElement
+
+from ufl import Cell, Coefficient, FunctionSpace, Mesh
+from ufl.algorithms.apply_derivatives import apply_derivatives
+from ufl.algorithms.cancel_jacobian_products import cancel_jacobian_products
+from ufl.algorithms.remove_component_tensors import remove_component_tensors
+from ufl.algorithms.renumbering import renumber_indices
+from ufl.classes import JacobianDeterminant, ReferenceValue
+from ufl.pullback import (
+    contravariant_piola,
+    covariant_contravariant_piola,
+    covariant_piola,
+    double_contravariant_piola,
+    double_covariant_piola,
+    l2_piola,
+    physical_pullback,
+    undefined_pullback,
+)
+from ufl.sobolevspace import H1, L2, HCurl, HDiv, HDivDiv, HEin
+
+cell = Cell("triangle")
+domain = Mesh(LagrangeElement(cell, 1, (2,)))
+
+U = LagrangeElement(cell, 1)
+Vd = FiniteElement("Raviart-Thomas", cell, 1, (2,), contravariant_piola, HDiv)
+Vc = FiniteElement("N1curl", cell, 1, (2,), covariant_piola, HCurl)
+Td = FiniteElement("Regge", cell, 1, (2, 2), double_covariant_piola, HEin)
+Tc = FiniteElement("HHJ", cell, 1, (2, 2), double_contravariant_piola, HDivDiv)
+Tcc = FiniteElement("CC", cell, 1, (2, 2), covariant_contravariant_piola, HDivDiv)
+S = SymmetricElement({(0, 0): 0, (1, 0): 1, (0, 1): 1, (1, 1): 2}, [U, U, U])
+M = MixedElement([U, Vd, Vc])
+
+
+def simplify(expr):
+    """Cancel the Jacobian products a round trip leaves behind."""
+    return cancel_jacobian_products(remove_component_tensors(apply_derivatives(expr)))
+
+
+@pytest.mark.parametrize(
+    "element",
+    [U, Vd, Vc, Td, Tc, Tcc, S, M],
+    ids=[
+        "identity",
+        "contravariant",
+        "covariant",
+        "double covariant",
+        "double contravariant",
+        "covariant contravariant",
+        "symmetric",
+        "mixed",
+    ],
+)
+def test_apply_inverse_undoes_apply(element):
+    """The inverse pull back returns a pushed-forward function unchanged."""
+    pullback = element.pullback
+    reference = ReferenceValue(Coefficient(FunctionSpace(domain, element)))
+    actual = simplify(pullback.apply_inverse(pullback.apply(reference)))
+
+    assert actual.ufl_shape == reference.ufl_shape
+    for idx in np.ndindex(reference.ufl_shape):
+        assert renumber_indices(actual[idx]) == renumber_indices(reference[idx])
+
+
+def test_l2_piola_apply_inverse():
+    """The L2 Piola scales by the Jacobian determinant.
+
+    The round trip leaves ``detJ / detJ`` standing, because cancelling a scalar
+    factor is not something ``cancel_jacobian_products`` does.
+    """
+    element = FiniteElement("Discontinuous Lagrange", cell, 1, (), l2_piola, L2)
+    reference = ReferenceValue(Coefficient(FunctionSpace(domain, element)))
+
+    assert l2_piola.apply_inverse(reference) == reference * JacobianDeterminant(domain)
+
+
+def test_apply_inverse_of_physical_value_shape():
+    """The inverse pull back maps a physical shape to a reference shape."""
+    for element in [U, Vd, Vc, Td, Tc, Tcc, S, M]:
+        pullback = element.pullback
+        physical = pullback.apply(ReferenceValue(Coefficient(FunctionSpace(domain, element))))
+        assert physical.ufl_shape == pullback.physical_value_shape(element, domain)
+        assert pullback.apply_inverse(physical).ufl_shape == element.reference_value_shape
+
+
+@pytest.mark.parametrize("pullback", [physical_pullback, undefined_pullback])
+def test_apply_inverse_is_not_defined(pullback):
+    """A pull back with no standard inverse says so."""
+    element = FiniteElement("Custom", cell, 1, (), pullback, H1)
+    reference = ReferenceValue(Coefficient(FunctionSpace(domain, element)))
+    if pullback is undefined_pullback:
+        with pytest.raises(BaseException):
+            pullback.apply_inverse(reference)
+    else:
+        assert pullback.apply_inverse(reference) == reference

ufl/pullback.py

@@ -78,6 +78,17 @@ def apply(self, expr: Expr, domain: AbstractDomain | None = None) -> Expr:
         """
         raise NonStandardPullbackException()
 
+    def apply_inverse(self, expr: Expr, domain: AbstractDomain | None = None) -> Expr:
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        raise NonStandardPullbackException()
+
 
 class IdentityPullback(AbstractPullback):
     """The identity pull back."""
@@ -102,6 +113,17 @@ def apply(self, expr, domain=None):
         """
         return expr
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        return expr
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -147,6 +169,25 @@ def apply(self, expr, domain=None):
         kj = (*k, j)
         return as_tensor(transform[i, j] * expr[kj], (*k, i))
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        from ufl.classes import Jacobian, JacobianDeterminant, JacobianInverse
+
+        domain = domain or extract_unique_domain(expr)
+        J = Jacobian(domain)
+        detJ = JacobianDeterminant(J)
+        K = JacobianInverse(domain)
+        *k, i, j = indices(len(expr.ufl_shape) + 1)
+        kj = (*k, j)
+        return as_tensor(detJ * K[i, j] * expr[kj], (*k, i))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -191,6 +232,23 @@ def apply(self, expr, domain=None):
         kj = (*k, j)
         return as_tensor(K[j, i] * expr[kj], (*k, i))
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        from ufl.classes import Jacobian
+
+        domain = domain or extract_unique_domain(expr)
+        J = Jacobian(domain)
+        *k, i, j = indices(len(expr.ufl_shape) + 1)
+        kj = (*k, j)
+        return as_tensor(J[j, i] * expr[kj], (*k, i))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -232,6 +290,21 @@ def apply(self, expr, domain=None):
         detJ = JacobianDeterminant(domain)
         return expr / detJ
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        from ufl.classes import JacobianDeterminant
+
+        domain = domain or extract_unique_domain(expr)
+        detJ = JacobianDeterminant(domain)
+        return expr * detJ
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -276,6 +349,24 @@ def apply(self, expr, domain=None):
         kmn = (*k, m, n)
         return as_tensor((1.0 / detJ) ** 2 * J[i, m] * expr[kmn] * J[j, n], (*k, i, j))
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        from ufl.classes import JacobianDeterminant, JacobianInverse
+
+        domain = domain or extract_unique_domain(expr)
+        detJ = JacobianDeterminant(domain)
+        K = JacobianInverse(domain)
+        *k, i, j, m, n = indices(len(expr.ufl_shape) + 2)
+        kmn = (*k, m, n)
+        return as_tensor(detJ**2 * K[i, m] * expr[kmn] * K[j, n], (*k, i, j))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -320,6 +411,23 @@ def apply(self, expr, domain=None):
         kmn = (*k, m, n)
         return as_tensor(K[m, i] * expr[kmn] * K[n, j], (*k, i, j))
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        from ufl.classes import Jacobian
+
+        domain = domain or extract_unique_domain(expr)
+        J = Jacobian(domain)
+        *k, i, j, m, n = indices(len(expr.ufl_shape) + 2)
+        kmn = (*k, m, n)
+        return as_tensor(J[m, i] * expr[kmn] * J[n, j], (*k, i, j))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -366,6 +474,25 @@ def apply(self, expr, domain=None):
         kmn = (*k, m, n)
         return as_tensor((1.0 / detJ) * K[m, i] * expr[kmn] * J[j, n], (*k, i, j))
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        from ufl.classes import Jacobian, JacobianDeterminant, JacobianInverse
+
+        domain = domain or extract_unique_domain(expr)
+        J = Jacobian(domain)
+        detJ = JacobianDeterminant(J)
+        K = JacobianInverse(domain)
+        *k, i, j, m, n = indices(len(expr.ufl_shape) + 2)
+        kmn = (*k, m, n)
+        return as_tensor(detJ * J[m, i] * expr[kmn] * K[j, n], (*k, i, j))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element.
 
@@ -438,6 +565,33 @@ def apply(self, expr, domain=None):
             )
         return f
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        domain = domain or extract_unique_domain(expr, expand_mesh_sequence=False)
+        gflat = [expr[idx] for idx in np.ndindex(expr.ufl_shape)]
+        r_components = []
+        offset = 0
+        # For each piece in physical space, apply the appropriate inverse pullback
+        for subelem, subdomain in zip(
+            self._element.sub_elements, domain.iterable_like(self._element)
+        ):
+            physical_shape = subelem.pullback.physical_value_shape(subelem, subdomain)
+            size = int(np.prod(physical_shape, dtype=int))
+            gsub = as_tensor(np.asarray(gflat[offset : offset + size]).reshape(physical_shape))
+            gmapped = subelem.pullback.apply_inverse(gsub, domain=subdomain)
+            # Flatten into the mapped expression for the whole thing
+            r_components.extend(gmapped[idx] for idx in np.ndindex(gmapped.ufl_shape))
+            offset += size
+        # And reshape appropriately
+        return as_tensor(np.asarray(r_components).reshape(self._element.reference_value_shape))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -527,6 +681,37 @@ def apply(self, expr, domain=None):
             )
         return f
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        domain = domain or extract_unique_domain(expr, expand_mesh_sequence=False)
+        subelem = self._element.sub_elements[0]
+        physical_shape = subelem.pullback.physical_value_shape(subelem, domain)
+        size = int(np.prod(physical_shape, dtype=int))
+        gflat = [expr[idx] for idx in np.ndindex(expr.ufl_shape)]
+        # Symmetry repeats a reference piece across several physical blocks, so
+        # map each piece once, from the first block that carries it.
+        r_pieces = {}
+        for block, component in enumerate(np.ndindex(self._block_shape)):
+            i = self._symmetry[component]
+            if i in r_pieces:
+                continue
+            gsub = as_tensor(
+                np.asarray(gflat[size * block : size * (block + 1)]).reshape(physical_shape)
+            )
+            r_pieces[i] = subelem.pullback.apply_inverse(gsub, domain=domain)
+        r_components = []
+        for i in sorted(r_pieces):
+            gmapped = r_pieces[i]
+            r_components.extend(gmapped[idx] for idx in np.ndindex(gmapped.ufl_shape))
+        return as_tensor(np.asarray(r_components).reshape(self._element.reference_value_shape))
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -569,6 +754,17 @@ def apply(self, expr, domain=None):
         """
         return expr
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        return expr
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.
 
@@ -608,6 +804,17 @@ def apply(self, expr, domain=None):
         """
         return expr
 
+    def apply_inverse(self, expr, domain=None):
+        """Apply the inverse of the pull back.
+
+        Args:
+            expr: A function on a physical cell
+            domain: The domain on which the function is defined
+
+        Returns: The function mapped to the reference cell
+        """
+        return expr
+
     def physical_value_shape(self, element, domain) -> tuple[int, ...]:
         """Get the physical value shape when this pull back is applied to an element on a domain.