Allow expressions with submeshes (#4140)

* Make submesh expressions work (codim 0 for now, missing test for codim 1).

* Minor cleanup

* Add codim 1 test and fix bug

* Change fffcx ref

* Ruff formatting

* Mypy

* Sort imports

* Remove file

* Fix C++ demo and test and minor logic change for when to check for coordinate element hash

* Try fixing create_expression

* Reorder attributes

* Add docs

* Fix import for docs

* Copy data as it is strided when extracted from python

* Fix for parallel test

* Further parallel fixes

* Revert to main branch of FFCx

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Update API

* Add some typing

* Improve docstrings

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Docstring improvement part 2

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Add constexpr to else if

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Clang format

* Use switch

* Refactor to use extraction of coefficient cells.

* Add docs

* Minor doc change to rerun CI

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Use reference

Co-authored-by: Paul T. Kühner <56360279+schnellerhase@users.noreply.github.com>

* Address comments by @schnellerhase

* Simpler to vector trafo

* Fix: left facet marking

* Fix: test use fdim and analytical expression

* Check me: use extent not size

* Expose facet permutations to expression on facets

* Ruff format

* More ruff

* Extend tests

* Test ffcx branch

* Reformat imports in test

* Apply suggestions from code review

Co-authored-by: Jørgen Schartum Dokken <dokken92@gmail.com>

* Merge with main

* remove new file

---------

Co-authored-by: Paul T. Kühner <56360279+schnellerhase@users.noreply.github.com>

Author: jorgensd

cpp/demo/hyperelasticity/main.cpp

@@ -326,7 +326,7 @@ int main(int argc, char* argv[])
     // Compute Cauchy stress. Construct appropriate Basix element for
     // stress.
     fem::Expression sigma_expression = fem::create_expression<T, U>(
-        *expression_hyperelasticity_sigma, {{"u", u}}, {});
+        *expression_hyperelasticity_sigma, {{"u", u}}, {}, {});
 
     constexpr auto family = basix::element::family::P;
     auto cell_type

cpp/dolfinx/fem/Expression.h

@@ -1,4 +1,5 @@
-// Copyright (C) 2020-2025 Jack S. Hale, Michal Habera and Garth N. Wells
+// Copyright (C) 2020-2026 Jack S. Hale, Michal Habera, Garth N. Wells and
+// Jørgen S. Dokken
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -12,6 +13,7 @@
 #include <array>
 #include <concepts>
 #include <dolfinx/common/types.h>
+#include <dolfinx/mesh/EntityMap.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <functional>
 #include <memory>
@@ -62,38 +64,35 @@ class Expression
   /// @param[in] fn Function for tabulating the Expression.
   /// @param[in] value_shape Shape of Expression evaluated at single
   /// point.
+  /// @param[in] entity_maps Maps between entities of different meshes.
+  /// @param[in] coordinate_element_hash Hash for coordinate element used
+  /// to create the expression.
   /// @param[in] argument_space Function space for Argument. Used to
   /// computed a 1-form expression, e.g. can be used to create a matrix
   /// that when applied to a degree-of-freedom vector gives the
   /// expression values at the evaluation points.
-  Expression(const std::vector<std::shared_ptr<
-                 const Function<scalar_type, geometry_type>>>& coefficients,
-             const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
-                 constants,
-             std::span<const geometry_type> X,
-             std::array<std::size_t, 2> Xshape,
-             std::function<void(scalar_type*, const scalar_type*,
-                                const scalar_type*, const geometry_type*,
-                                const int*, const uint8_t*, void*)>
-                 fn,
-             const std::vector<std::size_t>& value_shape,
-             std::shared_ptr<const FunctionSpace<geometry_type>> argument_space
-             = nullptr)
+  Expression(
+      const std::vector<std::shared_ptr<
+          const Function<scalar_type, geometry_type>>>& coefficients,
+      const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
+          constants,
+      std::span<const geometry_type> X, std::array<std::size_t, 2> Xshape,
+      std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
+                         const geometry_type*, const int*, const uint8_t*,
+                         void*)>
+          fn,
+      const std::vector<std::size_t>& value_shape,
+      const std::vector<std::reference_wrapper<const dolfinx::mesh::EntityMap>>&
+          entity_maps,
+
+      std::uint64_t coordinate_element_hash,
+      std::shared_ptr<const FunctionSpace<geometry_type>> argument_space
+      = nullptr)
       : _argument_space(argument_space), _coefficients(coefficients),
         _constants(constants), _fn(fn), _value_shape(value_shape),
-        _x_ref(std::vector<geometry_type>(X.begin(), X.end()), Xshape)
-
-  {
-    for (auto& c : _coefficients)
-    {
-      assert(c);
-      if (c->function_space()->mesh()
-          != _coefficients.front()->function_space()->mesh())
-      {
-        throw std::runtime_error("Coefficients not all defined on same mesh.");
-      }
-    }
-  }
+        _x_ref(std::vector<geometry_type>(X.begin(), X.end()), Xshape),
+        _entity_maps(entity_maps),
+        _coordinate_element_hash(coordinate_element_hash) {};
 
   /// Move constructor
   Expression(Expression&& e) = default;
@@ -168,6 +167,19 @@ class Expression
     return _x_ref;
   }
 
+  /// @brief Maps between entities of different meshes.
+  const std::vector<std::reference_wrapper<const dolfinx::mesh::EntityMap>>&
+  entity_maps() const
+  {
+    return _entity_maps;
+  }
+
+  /// @brief Hash for coordinate element used to create the expression.
+  std::uint64_t coordinate_element_hash() const
+  {
+    return _coordinate_element_hash;
+  }
+
 private:
   // Function space for Argument
   std::shared_ptr<const FunctionSpace<geometry_type>> _argument_space;
@@ -189,5 +201,13 @@ class Expression
 
   // Evaluation points on reference cell
   std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>> _x_ref;
+
+  // Set of bidirectional maps of mesh entities (cells, facets, ridges or peaks)
+  // between meshes and submeshes of any codimension
+  std::vector<std::reference_wrapper<const dolfinx::mesh::EntityMap>>
+      _entity_maps;
+
+  // Hash for coordinate element used to create the expression
+  std::uint64_t _coordinate_element_hash;
 };
 } // namespace dolfinx::fem

cpp/dolfinx/fem/assemble_expression_impl.h

@@ -58,6 +58,7 @@ namespace dolfinx::fem::impl
 /// @param[in] P0 Degree-of-freedom transformation function. Applied when
 /// expressions includes an argument function that requires a
 /// transformation.
+/// @param[in] perms Entity permutation information for use in `fn`.
 template <dolfinx::scalar T, std::floating_point U>
 void tabulate_expression(
     std::span<T> values, fem::FEkernel<T> auto fn,
@@ -68,7 +69,8 @@ void tabulate_expression(
     md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
     std::span<const T> constants, fem::MDSpan2 auto entities,
     std::span<const std::uint32_t> cell_info,
-    fem::DofTransformKernel<T> auto P0)
+    fem::DofTransformKernel<T> auto P0,
+    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms)
 {
   static_assert(entities.rank() == 1 or entities.rank() == 2);
 
@@ -99,15 +101,16 @@ void tabulate_expression(
     else
     {
       std::int32_t entity = entities(e, 0);
+      std::int32_t local_entity = entities(e, 1);
+      std::uint8_t perm = perms.empty() ? 0 : perms(entity, local_entity);
       auto x_dofs = md::submdspan(x_dofmap, entity, md::full_extent);
       for (std::size_t i = 0; i < x_dofs.size(); ++i)
       {
         std::copy_n(std::next(x.begin(), 3 * x_dofs[i]), 3,
                     std::next(coord_dofs.begin(), 3 * i));
       }
       fn(values_local.data(), &coeffs(e, 0), constants.data(),
-         coord_dofs.data(), &entities(e, 1), nullptr, nullptr);
-
+         coord_dofs.data(), &local_entity, &perm, nullptr);
       P0(values_local, cell_info, entity, size0);
     }
 
@@ -183,10 +186,22 @@ void tabulate_expression(
               doftransform::transpose);
     }
   }
-
+  // An expression has no notion of requiring a facet permutation.
+  md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
+  if constexpr (entities.rank() == 2)
+  {
+    mesh::CellType cell_type = mesh.topology()->cell_types()[0];
+    int num_facets_per_cell
+        = mesh::cell_num_entities(cell_type, mesh.topology()->dim() - 1);
+    mesh.topology_mutable()->create_entity_permutations();
+    const std::vector<std::uint8_t>& p
+        = mesh.topology()->get_facet_permutations();
+    facet_perms = md::mdspan(p.data(), p.size() / num_facets_per_cell,
+                             num_facets_per_cell);
+  }
   tabulate_expression<T, U>(values, fn, Xshape, value_size, num_argument_dofs,
                             mesh.geometry().dofmaps().front(),
                             mesh.geometry().x(), coeffs, constants, entities,
-                            cell_info, post_dof_transform);
+                            cell_info, post_dof_transform, facet_perms);
 }
 } // namespace dolfinx::fem::impl

cpp/dolfinx/fem/assembler.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2018-2025 Garth N. Wells
+// Copyright (C) 2018-2026 Garth N. Wells and Jørgen S. Dokken
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -19,6 +19,7 @@
 #include <basix/mdspan.hpp>
 #include <cstdint>
 #include <dolfinx/common/types.h>
+#include <dolfinx/mesh/EntityMap.h>
 #include <memory>
 #include <optional>
 #include <span>
@@ -70,6 +71,12 @@ void tabulate_expression(
         std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
         element)
 {
+  // Check that domain is the same as mesh of the expression
+  if (e.coordinate_element_hash() != mesh.geometry().cmaps().front().hash())
+  {
+    throw std::runtime_error(
+        "Expression was created on a different mesh. Cannot tabulate.");
+  }
   auto [X, Xshape] = e.X();
   impl::tabulate_expression(values, e.kernel(), Xshape, e.value_size(), coeffs,
                             constants, mesh, entities, element);
@@ -95,6 +102,13 @@ template <dolfinx::scalar T, std::floating_point U>
 void tabulate_expression(std::span<T> values, const fem::Expression<T, U>& e,
                          const mesh::Mesh<U>& mesh, fem::MDSpan2 auto entities)
 {
+  // Check that domain is the same as mesh of the expression
+  if (e.coordinate_element_hash() != mesh.geometry().cmaps().front().hash())
+  {
+    throw std::runtime_error(
+        "Expression was created on a different mesh. Cannot tabulate.");
+  }
+
   std::optional<
       std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
       element = std::nullopt;
@@ -115,12 +129,13 @@ void tabulate_expression(std::span<T> values, const fem::Expression<T, U>& e,
     std::vector<std::reference_wrapper<const Function<T, U>>> c;
     std::ranges::transform(coefficients, std::back_inserter(c),
                            [](auto c) -> const Function<T, U>& { return *c; });
-    fem::pack_coefficients(c, coffsets, entities, std::span(coeffs));
+    fem::pack_coefficients(c, mesh, entities, e.entity_maps(), coffsets,
+                           std::span(coeffs));
   }
   std::vector<T> constants = fem::pack_constants(e);
 
   tabulate_expression<T, U>(
-      values, e, md::mdspan(coeffs.data(), entities.size(), cstride),
+      values, e, md::mdspan(coeffs.data(), entities.extent(0), cstride),
       std::span<const T>(constants), mesh, entities, element);
 }