warp-lang 1.7.0__py3-none-manylinux_2_28_x86_64.whl

warp/fem/space/__init__.py

@@ -0,0 +1,228 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# isort: skip_file
+
+from enum import Enum
+from typing import Optional
+
+import warp.fem.domain as _domain
+import warp.fem.geometry as _geometry
+import warp.fem.polynomial as _polynomial
+
+from .function_space import FunctionSpace
+from .basis_function_space import CollocatedFunctionSpace, ContravariantFunctionSpace, CovariantFunctionSpace
+from .topology import SpaceTopology
+from .basis_space import BasisSpace, PointBasisSpace, ShapeBasisSpace, make_discontinuous_basis_space
+from .shape import ElementBasis, get_shape_function, ShapeFunction
+
+from .grid_2d_function_space import make_grid_2d_space_topology
+
+from .grid_3d_function_space import make_grid_3d_space_topology
+
+from .trimesh_function_space import make_trimesh_space_topology
+
+from .tetmesh_function_space import make_tetmesh_space_topology
+
+from .quadmesh_function_space import make_quadmesh_space_topology
+
+from .hexmesh_function_space import make_hexmesh_space_topology
+
+from .nanogrid_function_space import make_nanogrid_space_topology
+
+
+from .partition import SpacePartition, make_space_partition
+from .restriction import SpaceRestriction
+
+
+from .dof_mapper import DofMapper, IdentityMapper, SymmetricTensorMapper, SkewSymmetricTensorMapper
+
+
+def make_space_restriction(
+    space: Optional[FunctionSpace] = None,
+    space_partition: Optional[SpacePartition] = None,
+    domain: Optional[_domain.GeometryDomain] = None,
+    space_topology: Optional[SpaceTopology] = None,
+    device=None,
+    temporary_store: "Optional[warp.fem.cache.TemporaryStore]" = None,  # noqa: F821
+) -> SpaceRestriction:
+    """
+    Restricts a function space partition to a Domain, i.e. a subset of its elements.
+
+    One of `space_partition`, `space_topology`, or `space` must be provided (and will be considered in that order).
+
+    Args:
+        space: (deprecated) if neither `space_partition` nor `space_topology` are provided, the space defining the topology to restrict
+        space_partition: the subset of nodes from the space topology to consider
+        domain: the domain to restrict the space to, defaults to all cells of the space geometry or partition.
+        space_topology: the space topology to be restricted, if `space_partition` is ``None``.
+        device: device on which to perform and store computations
+        temporary_store: shared pool from which to allocate temporary arrays
+    """
+
+    if space_partition is None:
+        if space_topology is None:
+            assert space is not None
+            space_topology = space.topology
+
+        if domain is None:
+            domain = _domain.Cells(geometry=space_topology.geometry)
+
+        space_partition = make_space_partition(
+            space_topology=space_topology, geometry_partition=domain.geometry_partition
+        )
+    elif domain is None:
+        domain = _domain.Cells(geometry=space_partition.geo_partition)
+
+    return SpaceRestriction(
+        space_partition=space_partition, domain=domain, device=device, temporary_store=temporary_store
+    )
+
+
+def make_polynomial_basis_space(
+    geo: _geometry.Geometry,
+    degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
+    discontinuous: bool = False,
+    family: Optional[_polynomial.Polynomial] = None,
+) -> BasisSpace:
+    """
+    Equips a geometry with a polynomial basis.
+
+    Args:
+        geo: the Geometry on which to build the space
+        degree: polynomial degree of the per-element shape functions
+        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the constructed basis space
+    """
+
+    base_geo = geo.base
+
+    if element_basis is None:
+        element_basis = ElementBasis.LAGRANGE
+    elif element_basis == ElementBasis.SERENDIPITY and degree == 1:
+        # Degree-1 serendipity is always equivalent to Lagrange
+        element_basis = ElementBasis.LAGRANGE
+
+    shape = get_shape_function(geo.reference_cell(), geo.dimension, degree, element_basis, family)
+
+    if discontinuous or degree == 0 or element_basis == ElementBasis.NONCONFORMING_POLYNOMIAL:
+        return make_discontinuous_basis_space(geo, shape)
+
+    topology = None
+    if isinstance(base_geo, _geometry.Grid2D):
+        topology = make_grid_2d_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Grid3D):
+        topology = make_grid_3d_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Trimesh):
+        topology = make_trimesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Tetmesh):
+        topology = make_tetmesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Quadmesh):
+        topology = make_quadmesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Hexmesh):
+        topology = make_hexmesh_space_topology(geo, shape)
+    elif isinstance(base_geo, _geometry.Nanogrid) or isinstance(base_geo, _geometry.AdaptiveNanogrid):
+        topology = make_nanogrid_space_topology(geo, shape)
+
+    if topology is None:
+        raise NotImplementedError(f"Unsupported geometry type {geo.name}")
+
+    return ShapeBasisSpace(topology, shape)
+
+
+def make_collocated_function_space(
+    basis_space: BasisSpace, dtype: type = float, dof_mapper: Optional[DofMapper] = None
+) -> CollocatedFunctionSpace:
+    """
+    Constructs a function space from a scalar-valued basis space and a value type, such that all degrees of freedom of the value type are stored at each of the basis nodes.
+
+    Args:
+        geo: the Geometry on which to build the space
+        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
+        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
+
+    Returns:
+        the constructed function space
+    """
+
+    if basis_space.value != ShapeFunction.Value.Scalar:
+        raise ValueError("Collocated function spaces may only be constructed from scalar-valued basis")
+
+    return CollocatedFunctionSpace(basis_space, dtype=dtype, dof_mapper=dof_mapper)
+
+
+def make_covariant_function_space(
+    basis_space: BasisSpace,
+) -> CovariantFunctionSpace:
+    """
+    Constructs a covariant function space from a vector-valued basis space
+    """
+
+    if basis_space.value != ShapeFunction.Value.CovariantVector:
+        raise ValueError("Covariant function spaces may only be constructed from covariant vector-valued basis")
+    return CovariantFunctionSpace(basis_space)
+
+
+def make_contravariant_function_space(
+    basis_space: BasisSpace,
+) -> ContravariantFunctionSpace:
+    """
+    Constructs a contravariant function space from a vector-valued basis space
+    """
+
+    if basis_space.value != ShapeFunction.Value.ContravariantVector:
+        raise ValueError("Contravariant function spaces may only be constructed from contravariant vector-valued basis")
+    return ContravariantFunctionSpace(basis_space)
+
+
+def make_polynomial_space(
+    geo: _geometry.Geometry,
+    dtype: type = float,
+    dof_mapper: Optional[DofMapper] = None,
+    degree: int = 1,
+    element_basis: Optional[ElementBasis] = None,
+    discontinuous: bool = False,
+    family: Optional[_polynomial.Polynomial] = None,
+) -> CollocatedFunctionSpace:
+    """
+    Equips a geometry with a collocated, polynomial function space.
+    Equivalent to successive calls to :func:`make_polynomial_basis_space` then `make_collocated_function_space`, `make_covariant_function_space` or `make_contravariant_function_space`.
+
+    Args:
+        geo: the Geometry on which to build the space
+        dtype: value type the function space. If ``dof_mapper`` is provided, the value type from the DofMapper will be used instead.
+        dof_mapper: mapping from node degrees of freedom to function values, defaults to Identity. Useful for reduced coordinates, e.g. :py:class:`SymmetricTensorMapper` maps 2x2 (resp 3x3) symmetric tensors to 3 (resp 6) degrees of freedom.
+        degree: polynomial degree of the per-element shape functions
+        discontinuous: if True, use Discontinuous Galerkin shape functions. Discontinuous is implied if degree is 0, i.e, piecewise-constant shape functions.
+        element_basis: type of basis function for the individual elements
+        family: Polynomial family used to generate the shape function basis. If not provided, a reasonable basis is chosen.
+
+    Returns:
+        the constructed function space
+    """
+
+    basis_space = make_polynomial_basis_space(geo, degree, element_basis, discontinuous, family)
+
+    if basis_space.value == ShapeFunction.Value.CovariantVector:
+        return make_covariant_function_space(basis_space)
+    if basis_space.value == ShapeFunction.Value.ContravariantVector:
+        return make_contravariant_function_space(basis_space)
+
+    return make_collocated_function_space(basis_space, dtype=dtype, dof_mapper=dof_mapper)

warp/fem/space/basis_function_space.py

@@ -0,0 +1,468 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Optional
+
+import warp as wp
+from warp.fem import cache
+from warp.fem.geometry import Geometry
+from warp.fem.linalg import basis_element, generalized_inner, generalized_outer
+from warp.fem.types import NULL_QP_INDEX, Coords, ElementIndex, make_free_sample
+
+from .basis_space import BasisSpace
+from .dof_mapper import DofMapper, IdentityMapper
+from .function_space import FunctionSpace
+from .partition import SpacePartition, make_space_partition
+
+
+class CollocatedFunctionSpace(FunctionSpace):
+    """Function space where values are collocated at nodes"""
+
+    @wp.struct
+    class LocalValueMap:
+        pass
+
+    def __init__(self, basis: BasisSpace, dtype: type = float, dof_mapper: DofMapper = None):
+        self.dof_mapper = IdentityMapper(dtype) if dof_mapper is None else dof_mapper
+        self._basis = basis
+
+        super().__init__(topology=basis.topology)
+
+        self.dtype = self.dof_mapper.value_dtype
+        self.dof_dtype = self.dof_mapper.dof_dtype
+        self.VALUE_DOF_COUNT = self.dof_mapper.DOF_SIZE
+        self.NODE_DOF_COUNT = self.dof_mapper.DOF_SIZE
+
+        self.SpaceArg = self._basis.BasisArg
+        self.space_arg_value = self._basis.basis_arg_value
+
+        self.ORDER = self._basis.ORDER
+
+        self.node_basis_element = self._make_node_basis_element()
+        self.value_basis_element = self._make_value_basis_element()
+
+        self.node_coords_in_element = self._basis.make_node_coords_in_element()
+        self.node_quadrature_weight = self._basis.make_node_quadrature_weight()
+        self.element_inner_weight = self._basis.make_element_inner_weight()
+        self.element_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
+        self.element_outer_weight = self._basis.make_element_outer_weight()
+        self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
+
+        self.space_value = self._make_space_value()
+        self.space_gradient = self._make_space_gradient()
+        self.space_divergence = self._make_space_divergence()
+
+        self.node_dof_value = self._make_node_dof_value()
+
+        # For backward compatibility
+        if hasattr(basis, "node_grid"):
+            self.node_grid = basis.node_grid
+        if hasattr(basis, "node_triangulation"):
+            self.node_triangulation = basis.node_triangulation
+        if hasattr(basis, "node_tets"):
+            self.node_tets = basis.node_tets
+        if hasattr(basis, "node_hexes"):
+            self.node_hexes = basis.node_hexes
+        if hasattr(basis, "vtk_cells"):
+            self.vtk_cells = basis.vtk_cells
+
+    @property
+    def name(self):
+        return f"{self._basis.name}_{self.dof_mapper}".replace(".", "_")
+
+    def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
+        return self._basis.node_positions(out=out)
+
+    def make_field(
+        self,
+        space_partition: Optional[SpacePartition] = None,
+    ) -> "wp.fem.field.NodalField":
+        from warp.fem.field import NodalField
+
+        if space_partition is None:
+            space_partition = make_space_partition(space_topology=self.topology)
+
+        return NodalField(space=self, space_partition=space_partition)
+
+    def trace(self) -> "CollocatedFunctionSpace":
+        return CollocatedFunctionSpaceTrace(self)
+
+    def _make_node_basis_element(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_basis_element(dof_coord: int):
+            return basis_element(self.dof_dtype(0.0), dof_coord)
+
+        return node_basis_element
+
+    def _make_value_basis_element(self):
+        @cache.dynamic_func(suffix=self.name)
+        def value_basis_element(dof_coord: int, value_map: CollocatedFunctionSpace.LocalValueMap):
+            return self.dof_mapper.dof_to_value(self.node_basis_element(dof_coord))
+
+        return value_basis_element
+
+    @wp.func
+    def local_value_map_inner(
+        elt_arg: Any,
+        element_index: ElementIndex,
+        element_coords: Coords,
+    ):
+        return CollocatedFunctionSpace.LocalValueMap()
+
+    @wp.func
+    def local_value_map_outer(
+        elt_arg: Any,
+        element_index: ElementIndex,
+        element_coords: Coords,
+    ):
+        return CollocatedFunctionSpace.LocalValueMap()
+
+    def _make_space_value(self):
+        @cache.dynamic_func(suffix=self.name)
+        def value_func(
+            dof_value: self.dof_dtype,
+            node_weight: self._basis.weight_type,
+            local_value_map: self.LocalValueMap,
+        ):
+            return node_weight * self.dof_mapper.dof_to_value(dof_value)
+
+        return value_func
+
+    def _make_space_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def gradient_func(
+            dof_value: self.dof_dtype,
+            node_weight_gradient: self._basis.weight_gradient_type,
+            local_value_map: self.LocalValueMap,
+            grad_transform: Any,
+        ):
+            return generalized_outer(self.dof_mapper.dof_to_value(dof_value), node_weight_gradient * grad_transform)
+
+        return gradient_func
+
+    def _make_space_divergence(self):
+        @cache.dynamic_func(suffix=self.name)
+        def divergence_func(
+            dof_value: self.dof_dtype,
+            node_weight_gradient: self._basis.weight_gradient_type,
+            local_value_map: self.LocalValueMap,
+            grad_transform: Any,
+        ):
+            return generalized_inner(self.dof_mapper.dof_to_value(dof_value), node_weight_gradient * grad_transform)
+
+        return divergence_func
+
+    def _make_node_dof_value(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_dof_value(
+            elt_arg: self.ElementArg,
+            space_arg: self.SpaceArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+            space_value: self.dtype,
+        ):
+            return self.dof_mapper.value_to_dof(space_value)
+
+        return node_dof_value
+
+
+class CollocatedFunctionSpaceTrace(CollocatedFunctionSpace):
+    """Trace of a :class:`CollocatedFunctionSpace`"""
+
+    def __init__(self, space: CollocatedFunctionSpace):
+        self._space = space
+        super().__init__(space._basis.trace(), space.dtype, space.dof_mapper)
+
+    @property
+    def name(self):
+        return f"{self._space.name}_Trace"
+
+    def __eq__(self, other: "CollocatedFunctionSpaceTrace") -> bool:
+        return self._space == other._space
+
+
+class VectorValuedFunctionSpace(FunctionSpace):
+    """Function space whose values are vectors"""
+
+    def __init__(self, basis: BasisSpace):
+        self._basis = basis
+
+        super().__init__(topology=basis.topology)
+
+        self.dtype = cache.cached_vec_type(self.geometry.dimension, dtype=float)
+        self.dof_dtype = float
+
+        self.VALUE_DOF_COUNT = self.geometry.dimension
+        self.NODE_DOF_COUNT = 1
+
+        self.SpaceArg = self._basis.BasisArg
+        self.space_arg_value = self._basis.basis_arg_value
+
+        self.ORDER = self._basis.ORDER
+
+        self.LocalValueMap = cache.cached_mat_type(
+            shape=(self.geometry.dimension, self.geometry.cell_dimension), dtype=float
+        )
+
+        self.value_basis_element = self._make_value_basis_element()
+
+        self.node_coords_in_element = self._basis.make_node_coords_in_element()
+        self.node_quadrature_weight = self._basis.make_node_quadrature_weight()
+        self.element_inner_weight = self._basis.make_element_inner_weight()
+        self.element_inner_weight_gradient = self._basis.make_element_inner_weight_gradient()
+        self.element_outer_weight = self._basis.make_element_outer_weight()
+        self.element_outer_weight_gradient = self._basis.make_element_outer_weight_gradient()
+
+        self.space_value = self._make_space_value()
+        self.space_gradient = self._make_space_gradient()
+        self.space_divergence = self._make_space_divergence()
+
+        self.node_dof_value = self._make_node_dof_value()
+
+    @property
+    def name(self):
+        return self._basis.name
+
+    def node_positions(self, out: Optional[wp.array] = None) -> wp.array:
+        return self._basis.node_positions(out=out)
+
+    def make_field(
+        self,
+        space_partition: Optional[SpacePartition] = None,
+    ) -> "wp.fem.field.NodalField":
+        from warp.fem.field import NodalField
+
+        if space_partition is None:
+            space_partition = make_space_partition(space_topology=self.topology)
+
+        return NodalField(space=self, space_partition=space_partition)
+
+    @wp.func
+    def node_basis_element(dof_coord: int):
+        return 1.0
+
+    def _make_value_basis_element(self):
+        @cache.dynamic_func(suffix=self.name)
+        def value_basis_element(dof_coord: int, value_map: Any):
+            return value_map * basis_element(self.dtype(0.0), dof_coord)
+
+        return value_basis_element
+
+    def _make_space_value(self):
+        @cache.dynamic_func(suffix=self.name)
+        def value_func(
+            dof_value: self.dof_dtype,
+            node_weight: self._basis.weight_type,
+            local_value_map: self.LocalValueMap,
+        ):
+            return local_value_map * (node_weight * dof_value)
+
+        return value_func
+
+    def _make_space_gradient(self):
+        @cache.dynamic_func(suffix=self.name)
+        def gradient_func(
+            dof_value: self.dof_dtype,
+            node_weight_gradient: self._basis.weight_gradient_type,
+            local_value_map: self.LocalValueMap,
+            grad_transform: Any,
+        ):
+            return dof_value * local_value_map * node_weight_gradient * grad_transform
+
+        return gradient_func
+
+    def _make_space_divergence(self):
+        @cache.dynamic_func(suffix=self.name)
+        def divergence_func(
+            dof_value: self.dof_dtype,
+            node_weight_gradient: self._basis.weight_gradient_type,
+            local_value_map: self.LocalValueMap,
+            grad_transform: Any,
+        ):
+            return dof_value * wp.trace(local_value_map * node_weight_gradient * grad_transform)
+
+        return divergence_func
+
+    def _make_node_dof_value(self):
+        @cache.dynamic_func(suffix=self.name)
+        def node_dof_value(
+            elt_arg: self.ElementArg,
+            space_arg: self.SpaceArg,
+            element_index: ElementIndex,
+            node_index_in_elt: int,
+            space_value: self.dtype,
+        ):
+            coords = self.node_coords_in_element(elt_arg, space_arg, element_index, node_index_in_elt)
+            weight = self.element_inner_weight(
+                elt_arg, space_arg, element_index, coords, node_index_in_elt, NULL_QP_INDEX
+            )
+            local_value_map = self.local_value_map_inner(elt_arg, element_index, coords)
+
+            unit_value = local_value_map * weight
+            return wp.dot(space_value, unit_value) / wp.length_sq(unit_value)
+
+        return node_dof_value
+
+
+class CovariantFunctionSpace(VectorValuedFunctionSpace):
+    """Function space whose values are covariant vectors"""
+
+    def __init__(self, basis: BasisSpace):
+        super().__init__(basis)
+
+        self.local_value_map_inner = self._make_local_value_map()
+        self.local_value_map_outer = self.local_value_map_inner
+
+    def trace(self) -> "CovariantFunctionSpaceTrace":
+        return CovariantFunctionSpaceTrace(self)
+
+    def _make_local_value_map(self):
+        @cache.dynamic_func(suffix=self.name)
+        def local_value_map(
+            elt_arg: self.ElementArg,
+            element_index: ElementIndex,
+            element_coords: Coords,
+        ):
+            J = wp.transpose(
+                self.geometry.cell_inverse_deformation_gradient(
+                    elt_arg, make_free_sample(element_index, element_coords)
+                )
+            )
+            return J
+
+        return local_value_map
+
+
+class CovariantFunctionSpaceTrace(VectorValuedFunctionSpace):
+    """Trace of a :class:`CovariantFunctionSpace`"""
+
+    def __init__(self, space: VectorValuedFunctionSpace):
+        self._space = space
+        super().__init__(space._basis.trace())
+
+        self.local_value_map_inner = self._make_local_value_map_inner()
+        self.local_value_map_outer = self._make_local_value_map_outer()
+
+    @property
+    def name(self):
+        return f"{self._space.name}_Trace"
+
+    def __eq__(self, other: "CovariantFunctionSpaceTrace") -> bool:
+        return self._space == other._space
+
+    def _make_local_value_map_inner(self):
+        @cache.dynamic_func(suffix=self.name)
+        def local_value_map_inner(
+            elt_arg: self.ElementArg,
+            element_index: ElementIndex,
+            element_coords: Coords,
+        ):
+            return wp.transpose(
+                self.geometry.side_inner_inverse_deformation_gradient(
+                    elt_arg, make_free_sample(element_index, element_coords)
+                )
+            )
+
+        return local_value_map_inner
+
+    def _make_local_value_map_outer(self):
+        @cache.dynamic_func(suffix=self.name)
+        def local_value_map_outer(
+            elt_arg: self.ElementArg,
+            element_index: ElementIndex,
+            element_coords: Coords,
+        ):
+            return wp.transpose(
+                self.geometry.side_outer_inverse_deformation_gradient(
+                    elt_arg, make_free_sample(element_index, element_coords)
+                )
+            )
+
+        return local_value_map_outer
+
+
+class ContravariantFunctionSpace(VectorValuedFunctionSpace):
+    """Function space whose values are contravariant vectors"""
+
+    def __init__(self, basis: BasisSpace):
+        super().__init__(basis)
+
+        self.local_value_map_inner = self._make_local_value_map()
+        self.local_value_map_outer = self.local_value_map_inner
+
+    def trace(self) -> "ContravariantFunctionSpaceTrace":
+        return ContravariantFunctionSpaceTrace(self)
+
+    def _make_local_value_map(self):
+        @cache.dynamic_func(suffix=self.name)
+        def local_value_map(
+            elt_arg: self.ElementArg,
+            element_index: ElementIndex,
+            element_coords: Coords,
+        ):
+            F = self.geometry.cell_deformation_gradient(elt_arg, make_free_sample(element_index, element_coords))
+            return F / Geometry._element_measure(F)
+
+        return local_value_map
+
+
+class ContravariantFunctionSpaceTrace(VectorValuedFunctionSpace):
+    """Trace of a :class:`ContravariantFunctionSpace`"""
+
+    def __init__(self, space: ContravariantFunctionSpace):
+        self._space = space
+        super().__init__(space._basis.trace())
+
+        self.local_value_map_inner = self._make_local_value_map_inner()
+        self.local_value_map_outer = self._make_local_value_map_outer()
+
+    @property
+    def name(self):
+        return f"{self._space.name}_Trace"
+
+    def __eq__(self, other: "ContravariantFunctionSpaceTrace") -> bool:
+        return self._space == other._space
+
+    def _make_local_value_map_inner(self):
+        @cache.dynamic_func(suffix=self.name)
+        def local_value_map_inner(
+            elt_arg: self.ElementArg,
+            element_index: ElementIndex,
+            element_coords: Coords,
+        ):
+            cell_index = self.geometry.side_inner_cell_index(elt_arg, element_index)
+            cell_coords = self.geometry.side_inner_cell_coords(elt_arg, element_index, element_coords)
+            cell_arg = self.geometry.side_to_cell_arg(elt_arg)
+
+            F = self.geometry.cell_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+            return F / Geometry._element_measure(F)
+
+        return local_value_map_inner
+
+    def _make_local_value_map_outer(self):
+        @cache.dynamic_func(suffix=self.name)
+        def local_value_map_outer(
+            elt_arg: self.ElementArg,
+            element_index: ElementIndex,
+            element_coords: Coords,
+        ):
+            cell_index = self.geometry.side_outer_cell_index(elt_arg, element_index)
+            cell_coords = self.geometry.side_outer_cell_coords(elt_arg, element_index, element_coords)
+            cell_arg = self.geometry.side_to_cell_arg(elt_arg)
+
+            F = self.geometry.cell_deformation_gradient(cell_arg, make_free_sample(cell_index, cell_coords))
+            return F / Geometry._element_measure(F)
+
+        return local_value_map_outer