Support evaluation of facets in dolfinx.fem.Expression

cpp/dolfinx/fem/Expression.h

@@ -141,19 +141,35 @@ class Expression
     return n;
   }
 
-  /// @brief Evaluate Expression on cells.
+  /// @brief Evaluate Expression on cells or facets.
   /// @param[in] mesh Cells on which to evaluate the Expression.
-  /// @param[in] cells Cells on which to evaluate the Expression.
+  /// @param[in] entities List of entities to evaluate the expression on. This
+  /// could be either a list of cells or a list of (cell, local facet index)
+  /// tuples. Array is flattened per entity.
   /// @param[out] values A 2D array to store the result. Caller
   /// is responsible for correct sizing which should be `(num_cells,
   /// num_points * value_size * num_all_argument_dofs columns)`.
   /// @param[in] vshape The shape of @p values (row-major storage).
   void eval(const mesh::Mesh<geometry_type>& mesh,
-            std::span<const std::int32_t> cells, std::span<scalar_type> values,
+            std::span<const std::int32_t> entities,
+            std::span<scalar_type> values,
             std::array<std::size_t, 2> vshape) const
   {
+    std::size_t estride;
+    if (mesh.topology()->dim() == _x_ref.second[1])
+    {
+      estride = 1;
+    }
+    else if (mesh.topology()->dim() == _x_ref.second[1] + 1)
+    {
+      estride = 2;
+    }
+    else
+    {
+      throw std::runtime_error("Invalid dimension of evaluation points.");
+    }
     // Prepare coefficients and constants
-    const auto [coeffs, cstride] = pack_coefficients(*this, cells);
+    const auto [coeffs, cstride] = pack_coefficients(*this, entities, estride);
     const std::vector<scalar_type> constant_data = pack_constants(*this);
     auto fn = this->get_tabulate_expression();
 
@@ -198,29 +214,44 @@ class Expression
       }
     }
 
+    // Create get entity index function
+    std::function<const std::int32_t*(std::span<const std::int32_t>,
+                                      std::size_t)>
+        get_entity_index
+        = []([[maybe_unused]] std::span<const std::int32_t> entities,
+             [[maybe_unused]] std::size_t idx) { return nullptr; };
+    if (estride == 2)
+    {
+      get_entity_index
+          = [](std::span<const std::int32_t> entities, std::size_t idx)
+      { return entities.data() + 2 * idx + 1; };
+    }
+
     // Iterate over cells and 'assemble' into values
     const int size0 = _x_ref.second[0] * value_size();
     std::vector<scalar_type> values_local(size0 * num_argument_dofs, 0);
-    for (std::size_t c = 0; c < cells.size(); ++c)
+    for (std::size_t e = 0; e < entities.size() / estride; ++e)
     {
-      const std::int32_t cell = cells[c];
+      const std::int32_t entity = entities[e * estride];
       auto x_dofs = MDSPAN_IMPL_STANDARD_NAMESPACE::
           MDSPAN_IMPL_PROPOSED_NAMESPACE::submdspan(
-              x_dofmap, cell, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent);
+              x_dofmap, entity, MDSPAN_IMPL_STANDARD_NAMESPACE::full_extent);
       for (std::size_t i = 0; i < x_dofs.size(); ++i)
       {
         std::copy_n(std::next(x_g.begin(), 3 * x_dofs[i]), 3,
                     std::next(coord_dofs.begin(), 3 * i));
       }
 
-      const scalar_type* coeff_cell = coeffs.data() + c * cstride;
+      const scalar_type* coeff_cell = coeffs.data() + e * cstride;
+      const int* entity_index = get_entity_index(entities, e);
+
       std::fill(values_local.begin(), values_local.end(), 0);
       _fn(values_local.data(), coeff_cell, constant_data.data(),
-          coord_dofs.data(), nullptr, nullptr);
+          coord_dofs.data(), entity_index, nullptr);
 
-      post_dof_transform(values_local, cell_info, c, size0);
+      post_dof_transform(values_local, cell_info, e, size0);
       for (std::size_t j = 0; j < values_local.size(); ++j)
-        values[c * vshape[1] + j] = values_local[j];
+        values[e * vshape[1] + j] = values_local[j];
     }
   }
 

cpp/dolfinx/fem/utils.h

@@ -1091,10 +1091,10 @@ Expression<T, U> create_expression(
                              "function space was provided.");
   }
 
-  std::vector<U> X(e.points, e.points + e.num_points * e.topological_dimension);
+  std::vector<U> X(e.points, e.points + e.num_points * e.entity_dimension);
   std::array<std::size_t, 2> Xshape
       = {static_cast<std::size_t>(e.num_points),
-         static_cast<std::size_t>(e.topological_dimension)};
+         static_cast<std::size_t>(e.entity_dimension)};
   std::vector<int> value_shape(e.value_shape, e.value_shape + e.num_components);
   std::function<void(T*, const T*, const T*,
                      const typename scalar_value_type<T>::value_type*,
@@ -1193,15 +1193,17 @@ void pack_coefficients(const Form<T, U>& form,
 }
 
 /// @brief Pack coefficients of a Expression u for a give list of active
-/// cells.
+/// entities.
 ///
 /// @param[in] e The Expression
-/// @param[in] cells A list of active cells
+/// @param[in] entities A list of active entities
+/// @param[in] estride Stride for each entity in active entities (1 for cells, 2
+/// for facets)
 /// @return A pair of the form (coeffs, cstride)
 template <dolfinx::scalar T, std::floating_point U>
 std::pair<std::vector<T>, int>
 pack_coefficients(const Expression<T, U>& e,
-                  std::span<const std::int32_t> cells)
+                  std::span<const std::int32_t> entities, std::size_t estride)
 {
   // Get form coefficient offsets and dofmaps
   const std::vector<std::shared_ptr<const Function<T, U>>>& coeffs
@@ -1210,7 +1212,7 @@ pack_coefficients(const Expression<T, U>& e,
 
   // Copy data into coefficient array
   const int cstride = offsets.back();
-  std::vector<T> c(cells.size() * offsets.back());
+  std::vector<T> c(entities.size() / estride * offsets.back());
   if (!coeffs.empty())
   {
     std::span<const std::uint32_t> cell_info
@@ -1220,7 +1222,7 @@ pack_coefficients(const Expression<T, U>& e,
     for (std::size_t coeff = 0; coeff < coeffs.size(); ++coeff)
     {
       impl::pack_coefficient_entity(
-          std::span(c), cstride, *coeffs[coeff], cell_info, cells, 1,
+          std::span(c), cstride, *coeffs[coeff], cell_info, entities, estride,
           [](auto entity) { return entity[0]; }, offsets[coeff]);
     }
   }

python/dolfinx/fem/function.py

@@ -50,7 +50,9 @@ class Constant(ufl.Constant):
     ]
 
     def __init__(
-        self, domain, c: typing.Union[np.ndarray, typing.Sequence, np.floating, np.complexfloating]
+        self,
+        domain,
+        c: typing.Union[np.ndarray, typing.Sequence, np.floating, np.complexfloating],
     ):
         """A constant with respect to a domain.
 
@@ -165,7 +167,10 @@ def __init__(
             form_compiler_options = dict()
         form_compiler_options["scalar_type"] = dtype
         self._ufcx_expression, module, self._code = jit.ffcx_jit(
-            comm, (e, _X), form_compiler_options=form_compiler_options, jit_options=jit_options
+            comm,
+            (e, _X),
+            form_compiler_options=form_compiler_options,
+            jit_options=jit_options,
         )
         self._ufl_expression = e
 
@@ -209,31 +214,43 @@ def _create_expression(dtype):
         )
 
     def eval(
-        self, mesh: Mesh, cells: np.ndarray, values: typing.Optional[np.ndarray] = None
+        self,
+        mesh: Mesh,
+        entities: np.ndarray,
+        values: typing.Optional[np.ndarray] = None,
     ) -> np.ndarray:
-        """Evaluate Expression in cells.
+        """Evaluate Expression on entities.
 
         Args:
             mesh: Mesh to evaluate Expression on.
-            cells: Cells of `mesh`` to evaluate Expression on.
+            entities: Either an array of cells (index local to process) or an array of
+                integral tuples (cell index, local facet index). The array is flattened.
             values: Array to fill with evaluated values. If ``None``,
                 storage will be allocated. Otherwise must have shape
-                ``(len(cells), num_points * value_size *
+                ``(num_entities, num_points * value_size *
                 num_all_argument_dofs)``
 
         Returns:
-            Expression evaluated at points for `cells``.
+            Expression evaluated at points for `entities`.
 
         """
-        _cells = np.asarray(cells, dtype=np.int32)
+        _entities = np.asarray(entities, dtype=np.int32)
         if self.argument_function_space is None:
             argument_space_dimension = 1
         else:
             argument_space_dimension = self.argument_function_space.element.space_dimension
-        values_shape = (
-            _cells.shape[0],
-            self.X().shape[0] * self.value_size * argument_space_dimension,
-        )
+        if (tdim := mesh.topology.dim) != (expr_dim := self._cpp_object.X().shape[1]):
+            assert expr_dim == tdim - 1
+            assert _entities.shape[0] % 2 == 0
+            values_shape = (
+                _entities.shape[0] // 2,
+                self.X().shape[0] * self.value_size * argument_space_dimension,
+            )
+        else:
+            values_shape = (
+                _entities.shape[0],
+                self.X().shape[0] * self.value_size * argument_space_dimension,
+            )
 
         # Allocate memory for result if u was not provided
         if values is None:
@@ -243,8 +260,7 @@ def eval(
                 raise TypeError("Passed array values does not have correct shape.")
             if values.dtype != self.dtype:
                 raise TypeError("Passed array values does not have correct dtype.")
-
-        self._cpp_object.eval(mesh._cpp_object, cells, values)
+        self._cpp_object.eval(mesh._cpp_object, _entities, values)
         return values
 
     def X(self) -> np.ndarray:
@@ -612,7 +628,10 @@ def functionspace(
         form_compiler_options = dict()
     form_compiler_options["scalar_type"] = dtype
     (ufcx_element, ufcx_dofmap), module, code = jit.ffcx_jit(
-        mesh.comm, ufl_e, form_compiler_options=form_compiler_options, jit_options=jit_options
+        mesh.comm,
+        ufl_e,
+        form_compiler_options=form_compiler_options,
+        jit_options=jit_options,
     )
 
     ffi = module.ffi
@@ -626,7 +645,10 @@ def functionspace(
         )
 
     cpp_dofmap = _cpp.fem.create_dofmap(
-        mesh.comm, ffi.cast("uintptr_t", ffi.addressof(ufcx_dofmap)), mesh.topology, cpp_element
+        mesh.comm,
+        ffi.cast("uintptr_t", ffi.addressof(ufcx_dofmap)),
+        mesh.topology,
+        cpp_element,
     )
 
     assert np.issubdtype(

python/test/unit/fem/test_expression.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2019 Michal Habera
+# Copyright (C) 2019-2024 Michal Habera and Jørgen S. Dokken
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -334,3 +334,92 @@ def test_expression_comm(dtype):
     u = Function(functionspace(mesh, ("Lagrange", 1)), dtype=dtype)
     Expression(v, u.function_space.element.interpolation_points(), comm=MPI.COMM_WORLD)
     Expression(v, u.function_space.element.interpolation_points(), comm=MPI.COMM_SELF)
+
+
+def compute_exterior_facet_entities(mesh, facets):
+    """Helper function to compute (cell, local_facet_index) pairs for exterior facets"""
+    tdim = mesh.topology.dim
+    mesh.topology.create_connectivity(tdim - 1, tdim)
+    mesh.topology.create_connectivity(tdim, tdim - 1)
+    c_to_f = mesh.topology.connectivity(tdim, tdim - 1)
+    f_to_c = mesh.topology.connectivity(tdim - 1, tdim)
+    integration_entities = np.empty(2 * len(facets), dtype=np.int32)
+    for i, facet in enumerate(facets):
+        cells = f_to_c.links(facet)
+        assert len(cells) == 1
+        cell = cells[0]
+        local_facets = c_to_f.links(cell)
+        local_pos = np.flatnonzero(local_facets == facet)
+        integration_entities[2 * i] = cell
+        integration_entities[2 * i + 1] = local_pos[0]
+    return integration_entities
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.complex64, np.complex128])
+def test_facet_expression(dtype):
+    xtype = dtype(0).real.dtype
+    mesh = create_unit_square(MPI.COMM_WORLD, 4, 3, dtype=xtype)
+    n = ufl.FacetNormal(mesh)
+
+    tdim = mesh.topology.dim
+    mesh.topology.create_connectivity(tdim - 1, tdim)
+    facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
+
+    boundary_entities = compute_exterior_facet_entities(mesh, facets)
+
+    # Compute facet normal at midpoint of facet
+    reference_midpoint, _ = basix.quadrature.make_quadrature(
+        basix.cell.CellType.interval,
+        1,
+        basix.quadrature.QuadratureType.Default,
+        basix.quadrature.PolysetType.standard,
+    )
+    normal_expr = Expression(n, reference_midpoint, dtype=dtype)
+    facet_normals = normal_expr.eval(mesh, boundary_entities)
+
+    # Check facet normal by using midpoint to determine what exterior cell we are at
+    facet_midpoints = dolfinx.mesh.compute_midpoints(mesh, tdim - 1, facets)
+    atol = 100 * np.finfo(dtype).resolution
+    for midpoint, normal in zip(facet_midpoints, facet_normals):
+        if np.isclose(midpoint[0], 0, atol=atol):
+            assert np.allclose(normal, [-1, 0])
+        elif np.isclose(midpoint[0], 1, atol=atol):
+            assert np.allclose(normal, [1, 0], atol=atol)
+        elif np.isclose(midpoint[1], 0):
+            assert np.allclose(normal, [0, -1], atol=atol)
+        elif np.isclose(midpoint[1], 1, atol=atol):
+            assert np.allclose(normal, [0, 1])
+        else:
+            raise ValueError("Invalid midpoint")
+
+    # Check expression with coefficients from mixed space
+    el_v = basix.ufl.element("Lagrange", "triangle", 2, shape=(2,), dtype=xtype)
+    el_p = basix.ufl.element("Lagrange", "triangle", 1, dtype=xtype)
+    mixed_el = basix.ufl.mixed_element([el_v, el_p])
+    W = dolfinx.fem.functionspace(mesh, mixed_el)
+    w = dolfinx.fem.Function(W, dtype=dtype)
+    w.sub(0).interpolate(lambda x: (x[1] ** 2 + 3 * x[0] ** 2, -5 * x[1] ** 2 - 7 * x[0] ** 2))
+    w.sub(1).interpolate(lambda x: 2 * (x[1] + x[0]))
+    u, p = ufl.split(w)
+    n = ufl.FacetNormal(mesh)
+    mixed_expr = p * ufl.dot(ufl.grad(u), n)
+    facet_expression = dolfinx.fem.Expression(
+        mixed_expr, np.array([[0.5]], dtype=dtype), dtype=dtype
+    )
+    subset_values = facet_expression.eval(mesh, boundary_entities)
+    for values, midpoint in zip(subset_values, facet_midpoints):
+        grad_u = np.array(
+            [[6 * midpoint[0], 2 * midpoint[1]], [-14 * midpoint[0], -10 * midpoint[1]]],
+            dtype=dtype,
+        )
+        if np.isclose(midpoint[0], 0, atol=atol):
+            exact_n = [-1, 0]
+        elif np.isclose(midpoint[0], 1, atol=atol):
+            exact_n = [1, 0]
+        elif np.isclose(midpoint[1], 0):
+            exact_n = [0, -1]
+        elif np.isclose(midpoint[1], 1, atol=atol):
+            exact_n = [0, 1]
+
+        exact_expr = 2 * (midpoint[1] + midpoint[0]) * np.dot(grad_u, exact_n)
+        assert np.allclose(values, exact_expr, atol=atol)