TensorProduct(PhysicallyMapped)

finat/cube.py

@@ -75,7 +75,8 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
         if entity is None:
             entity = (self.cell.get_spatial_dimension(), 0)
 
-        return self.product.basis_evaluation(order, ps, self._unflatten[entity])
+        return self.product.basis_evaluation(order, ps, self._unflatten[entity],
+                                             coordinate_mapping)
 
     def point_evaluation(self, order, point, entity=None, coordinate_mapping=None):
         if entity is None:

finat/tensor_product.py

@@ -1,5 +1,7 @@
+from collections.abc import Sequence
 from itertools import chain, product
 from operator import methodcaller
+from typing import Any
 
 import numpy
 
@@ -12,11 +14,22 @@
 from gem.utils import cached_property
 
 from finat.finiteelementbase import FiniteElementBase
+from finat.physically_mapped import (NeedsCoordinateMappingElement,
+                                     PhysicalGeometry,
+                                     PhysicallyMappedElement,
+                                     determinant, identity, inverse)
 from finat.point_set import PointSingleton, PointSet, TensorPointSet
 
 
 class TensorProductElement(FiniteElementBase):
 
+    def __new__(cls, factors: Sequence[FiniteElementBase]) -> "TensorProductElement":
+        factors = tuple(factors)
+        if cls is TensorProductElement and any(
+                isinstance(fe, PhysicallyMappedElement) for fe in factors):
+            return super().__new__(TensorProductPhysicallyMappedElement)
+        return super().__new__(cls)
+
     def __init__(self, factors):
         super(TensorProductElement, self).__init__()
         self.factors = tuple(factors)
@@ -138,8 +151,10 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
 
         ps_factors = factor_point_set(self.cell, entity_dim, ps)
 
-        factor_results = [fe.basis_evaluation(order, ps_, e)
-                          for fe, ps_, e in zip(self.factors, ps_factors, entities)]
+        factor_mappings = self._factor_mappings(coordinate_mapping)
+        factor_results = [fe.basis_evaluation(order, ps_, e, coordinate_mapping=mapping)
+                          for fe, ps_, e, mapping
+                          in zip(self.factors, ps_factors, entities, factor_mappings)]
 
         return self._merge_evaluations(factor_results)
 
@@ -161,8 +176,10 @@ def point_evaluation(self, order, point, entity=None, coordinate_mapping=None):
             ))
 
         # Subelement results
-        factor_results = [fe.point_evaluation(order, p_, e)
-                          for fe, p_, e in zip(self.factors, point_factors, entities)]
+        factor_mappings = self._factor_mappings(coordinate_mapping)
+        factor_results = [fe.point_evaluation(order, p_, e, coordinate_mapping=mapping)
+                          for fe, p_, e, mapping
+                          in zip(self.factors, point_factors, entities, factor_mappings)]
 
         return self._merge_evaluations(factor_results)
 
@@ -192,6 +209,152 @@ def mapping(self):
         else:
             return None
 
+    def _factor_mappings(
+            self, coordinate_mapping: PhysicalGeometry | None
+    ) -> tuple[PhysicalGeometry | None, ...]:
+        if coordinate_mapping is None:
+            return (None,) * len(self.factors)
+        return tuple(TensorProductPhysicalGeometry(self, coordinate_mapping, i)
+                     if isinstance(fe, NeedsCoordinateMappingElement) else None
+                     for i, fe in enumerate(self.factors))
+
+
+class TensorProductPhysicallyMappedElement(TensorProductElement,
+                                           PhysicallyMappedElement):
+    """A tensor product with one or more physically mapped factors."""
+
+    def basis_transformation(self, coordinate_mapping: PhysicalGeometry) -> gem.Node:
+        matrices = []
+        for factor, mapping in zip(self.factors,
+                                   self._factor_mappings(coordinate_mapping)):
+            if isinstance(factor, PhysicallyMappedElement):
+                matrix = factor.basis_transformation(mapping).array
+            else:
+                matrix = identity(factor.space_dimension())
+            matrices.append(matrix)
+
+        result = matrices[0]
+        for matrix in matrices[1:]:
+            result = numpy.kron(result, matrix)
+        result = numpy.vectorize(gem.as_gem)(result)
+        return gem.ListTensor(result)
+
+
+class TensorProductPhysicalGeometry(PhysicalGeometry):
+    """Physical geometry restricted to one factor of a tensor product."""
+
+    def __init__(self, element: TensorProductElement,
+                 coordinate_mapping: PhysicalGeometry, factor: int) -> None:
+        self.element = element
+        self.coordinate_mapping = coordinate_mapping
+        self.factor = factor
+
+        cells = self.element.cell.cells
+        dimensions = [cell.get_spatial_dimension() for cell in cells]
+        self._column_slices = self.element.cell._split_slices(dimensions)
+
+        points = []
+        for cell in cells:
+            sd = cell.get_spatial_dimension()
+            point, = cell.make_points(sd, 0, sd + 1)
+            points.append(point)
+        self._points = points
+
+        J = self.coordinate_mapping.jacobian_at(self._embed_point(points[factor]))
+        physical_dimensions = list(dimensions)
+        # Extra ambient dimensions belong to the base manifold factor.
+        physical_dimensions[0] += J.shape[0] - sum(dimensions)
+        self._row_slices = self.element.cell._split_slices(physical_dimensions)
+
+    @property
+    def cell(self) -> Any:
+        return self.element.cell.cells[self.factor]
+
+    def _embed_point(self, point: Sequence[float]) -> tuple[float, ...]:
+        points = list(self._points)
+        points[self.factor] = point
+        return tuple(chain.from_iterable(points))
+
+    def cell_size(self) -> gem.Node:
+        sizes = gem.as_gem(self.coordinate_mapping.cell_size())
+        vertex_shape = tuple(len(cell.get_vertices())
+                             for cell in self.element.cell.cells)
+        indices = []
+        for vertex in range(vertex_shape[self.factor]):
+            multiindex = [0] * len(vertex_shape)
+            multiindex[self.factor] = vertex
+            if len(sizes.shape) == len(vertex_shape):
+                indices.append(tuple(multiindex))
+            else:
+                indices.append((numpy.ravel_multi_index(multiindex, vertex_shape),))
+        return gem.ListTensor([sizes[i] for i in indices])
+
+    def jacobian_at(self, point: Sequence[float]) -> gem.Node:
+        J = self.coordinate_mapping.jacobian_at(self._embed_point(point))
+        rows = range(*self._row_slices[self.factor].indices(J.shape[0]))
+        columns = range(*self._column_slices[self.factor].indices(J.shape[1]))
+        return gem.ListTensor([[J[i, j] for j in columns] for i in rows])
+
+    def detJ_at(self, point: Sequence[float]) -> gem.Node:
+        J = self.jacobian_at(point)
+        if J.shape[0] == J.shape[1]:
+            return determinant(J)
+        return gem.Power(determinant(J.T @ J), gem.Literal(0.5))
+
+    def reference_normals(self) -> gem.Node:
+        sd = self.cell.get_spatial_dimension()
+        return gem.Literal(numpy.asarray([
+            self.cell.compute_normal(i) for i in sorted(self.cell.get_topology()[sd-1])
+        ]))
+
+    def normalized_reference_edge_tangents(self) -> gem.Node:
+        return gem.Literal(numpy.asarray([
+            self.cell.compute_normalized_edge_tangent(i)
+            for i in sorted(self.cell.get_topology()[1])
+        ]))
+
+    def physical_tangents(self) -> gem.Node:
+        sd = self.cell.get_spatial_dimension()
+        point, = self.cell.make_points(sd, 0, sd + 1)
+        J = self.jacobian_at(point)
+        tangents = []
+        for edge in sorted(self.cell.get_topology()[1]):
+            tangent = J @ gem.Literal(self.cell.compute_edge_tangent(edge))
+            length = gem.Power(tangent @ tangent, gem.Literal(0.5))
+            tangents.append(tangent / length)
+        return gem.ListTensor(tangents)
+
+    def physical_normals(self) -> gem.Node:
+        sd = self.cell.get_spatial_dimension()
+        point, = self.cell.make_points(sd, 0, sd + 1)
+        J = self.jacobian_at(point)
+        gram_inv = inverse(J.T @ J)
+        normals = []
+        for face in sorted(self.cell.get_topology()[sd-1]):
+            normal = J @ (gram_inv @ gem.Literal(self.cell.compute_normal(face)))
+            length = gem.Power(normal @ normal, gem.Literal(0.5))
+            normals.append(normal / length)
+        return gem.ListTensor(normals)
+
+    def physical_edge_lengths(self) -> gem.Node:
+        sd = self.cell.get_spatial_dimension()
+        point, = self.cell.make_points(sd, 0, sd + 1)
+        J = self.jacobian_at(point)
+        lengths = []
+        for edge in sorted(self.cell.get_topology()[1]):
+            tangent = J @ gem.Literal(self.cell.compute_edge_tangent(edge))
+            lengths.append(gem.Power(tangent @ tangent, gem.Literal(0.5)))
+        return gem.ListTensor(lengths)
+
+    def physical_points(self, point_set: PointSet,
+                        entity: tuple[int, int] | None = None) -> gem.Node:
+        points = PointSet(numpy.asarray([self._embed_point(point)
+                                        for point in point_set.points]))
+        return self.coordinate_mapping.physical_points(points, entity=None)
+
+    def physical_vertices(self) -> gem.Node:
+        return self.physical_points(PointSet(numpy.asarray(self.cell.get_vertices())))
+
 
 def productise(factors, method):
     '''Tensor product the dict mapping topological entities to dofs across factors.

finat/ufl/tensorproductelement.py

@@ -72,10 +72,13 @@ def __repr__(self):
 
     def mapping(self):
         """Doc."""
-        if all(e.mapping() == "identity" for e in self._factor_elements):
+        mappings = {e.mapping() for e in self._factor_elements}
+        if mappings == {"identity"}:
             return "identity"
-        elif all(e.mapping() == "L2 Piola" for e in self._factor_elements):
+        elif mappings == {"L2 Piola"}:
             return "L2 Piola"
+        elif mappings <= {"custom", "identity"}:
+            return "custom"
         else:
             return "undefined"
 

test/finat/conftest.py

@@ -1,7 +1,10 @@
+from copy import deepcopy
+
 import pytest
 import FIAT
 import gem
 import numpy as np
+from FIAT.reference_element import TensorProductCell
 from finat.physically_mapped import PhysicalGeometry
 
 
@@ -128,6 +131,33 @@ def ref_to_phys(ref_el, phys_el):
     return {dim: MyMapping(ref_el[int(dim)], phys_el[dim]) for dim in phys_el}
 
 
+@pytest.fixture
+def extruded_ref_to_phys() -> dict[str, MyMapping]:
+    interval = FIAT.ufc_simplex(1)
+    triangle = FIAT.ufc_simplex(2)
+
+    physical_interval = deepcopy(interval)
+    physical_interval.vertices = ((0.2,), (1.7,))
+    physical_vertical = deepcopy(interval)
+    physical_vertical.vertices = ((-0.4,), (0.9,))
+    physical_triangle = deepcopy(triangle)
+    physical_triangle.vertices = ((0.1, -0.2),
+                                  (1.4, 0.1),
+                                  (-0.3, 1.6))
+
+    cells = {
+        "quadrilateral": (
+            TensorProductCell(interval, interval),
+            TensorProductCell(physical_interval, physical_vertical),
+        ),
+        "wedge": (
+            TensorProductCell(triangle, interval),
+            TensorProductCell(physical_triangle, physical_vertical),
+        ),
+    }
+    return {name: MyMapping(*pair) for name, pair in cells.items()}
+
+
 @pytest.fixture
 def scaled_ref_to_phys(ref_el):
     return {dim: [ScaledMapping(ref_el[dim], scaled_simplex(dim, 0.5**k)) for k in range(3)]

test/finat/test_ufl_elements.py

@@ -62,3 +62,16 @@ def test_symmetry(domain, family):
         symmetry = S.symmetry()
         assert isinstance(symmetry, dict)
         assert len(symmetry) == symmetry_size
+
+
+@pytest.mark.parametrize("families, cells, degrees", [
+    (("Hermite", "Real"), (ufl.interval, ufl.interval), (3, 0)),
+    (("Bell", "Hermite"), (ufl.triangle, ufl.interval), (5, 3)),
+])
+def test_tensor_product_custom_mapping(families, cells, degrees) -> None:
+    factors = [finat.ufl.FiniteElement(family, cell, degree)
+               for family, cell, degree in zip(families, cells, degrees)]
+    element = finat.ufl.TensorProductElement(*factors)
+
+    assert element.mapping() == "custom"
+    assert element.pullback is ufl.pullback.custom_pullback

test/finat/test_zany_mapping.py

@@ -1,3 +1,5 @@
+from itertools import product
+
 import FIAT
 import finat
 import numpy as np
@@ -87,13 +89,20 @@ def check_zany_mapping(element, ref_to_phys, *args, **kwargs):
     if isinstance(finat_element, PhysicallyMappedElement):
         Mgem = finat_element.basis_transformation(ref_to_phys)
         M = evaluate([Mgem])[0].arr
-        ref_vals_zany = np.tensordot(M, ref_vals_piola, (-1, 0))
     else:
         M = np.eye(num_dofs, num_bfs)
-        ref_vals_zany = ref_vals_piola
+
+    check_transformed_values(M, ref_vals_piola, phys_vals, num_dofs)
+
+
+def check_transformed_values(M: np.ndarray, ref_vals: np.ndarray,
+                             phys_vals: np.ndarray, num_dofs: int) -> None:
+    """Compare a proposed basis transformation with the numerical one."""
+    ref_vals_zany = np.tensordot(M, ref_vals, (-1, 0))
 
     # Solve for the basis transformation and compare results
-    Phi = ref_vals_piola.reshape(num_bfs, -1)
+    num_bfs = ref_vals.shape[0]
+    Phi = ref_vals.reshape(num_bfs, -1)
     phi = phys_vals.reshape(num_bfs, -1)
     Vh, residual, *_ = np.linalg.lstsq(Phi.T, phi.T)
     Mh = Vh.T
@@ -119,6 +128,60 @@ def check_zany_mapping(element, ref_to_phys, *args, **kwargs):
     assert np.allclose(ref_vals_zany, phys_vals[:num_dofs]), pp.pformat((np.round(error, 8).tolist(), *inds))
 
 
+def make_tensor_product_points(element) -> list[tuple[float, ...]]:
+    factors = (element.A, element.B)
+    point_factors = [make_unisolvent_points(factor, interior=True)
+                     for factor in factors]
+    return [tuple(x for point in points for x in point)
+            for points in product(*point_factors)]
+
+
+def check_tensor_product_zany_mapping(factors, ref_to_phys) -> None:
+    ref_cells = ref_to_phys.ref_cell.cells
+    phys_cells = ref_to_phys.phys_cell.cells
+    ref_factors = [constructor(cell) for constructor, cell in zip(factors, ref_cells)]
+    phys_factors = [constructor(cell).fiat_equivalent
+                    for constructor, cell in zip(factors, phys_cells)]
+
+    finat_element = finat.TensorProductElement(ref_factors)
+    ref_element = finat_element.fiat_equivalent
+    phys_element = FIAT.TensorProductElement(*phys_factors)
+
+    ref_points = make_tensor_product_points(ref_element)
+    phys_points = make_tensor_product_points(phys_element)
+    ref_key = (0,) * ref_element.get_reference_element().get_spatial_dimension()
+    phys_key = (0,) * phys_element.get_reference_element().get_spatial_dimension()
+    ref_vals = ref_element.tabulate(0, ref_points)[ref_key]
+    phys_vals = phys_element.tabulate(0, phys_points)[phys_key]
+
+    Mgem = finat_element.basis_transformation(ref_to_phys)
+    M = evaluate([Mgem])[0].arr
+    check_transformed_values(M, ref_vals, phys_vals,
+                             finat_element.space_dimension())
+
+
+def hermite(cell) -> finat.Hermite:
+    return finat.Hermite(cell)
+
+
+def real(cell) -> finat.Real:
+    return finat.Real(cell, 0)
+
+
+def bell(cell) -> finat.Bell:
+    return finat.Bell(cell)
+
+
+@pytest.mark.parametrize("cell, factors", [
+    ("quadrilateral", (hermite, real)),
+    ("wedge", (hermite, real)),
+    ("quadrilateral", (hermite, hermite)),
+    ("wedge", (bell, hermite)),
+])
+def test_tensor_product(extruded_ref_to_phys, cell, factors) -> None:
+    check_tensor_product_zany_mapping(factors, extruded_ref_to_phys[cell])
+
+
 @pytest.mark.parametrize("element", [
                          finat.Morley,
                          finat.Hermite,