Generalize sub-entity-jacobians (#890)

* Add edge jacobians for 3D-1D code compilation

* Update stubs

* Fix stubs

* Add tests

* Update copyright notice

* Remove cell dimension from jacobian

* Add prism and pyramid

* Ruff formatting

basix/_basixcpp.pyi

@@ -414,6 +414,8 @@ class SobolevSpace(enum.IntEnum):
 
 def cell_facet_jacobians(arg: CellType, /) -> Annotated[ArrayLike, dict(dtype='float64', )]: ...
 
+def cell_edge_jacobians(arg: CellType, /) -> Annotated[ArrayLike, dict(dtype='float64', )]: ...
+
 def cell_facet_normals(arg: CellType, /) -> Annotated[ArrayLike, dict(dtype='float64', )]: ...
 
 def cell_facet_orientations(arg: CellType, /) -> list[int]: ...

cpp/basix/cell.cpp

@@ -608,33 +608,62 @@ std::vector<std::vector<cell::type>> cell::subentity_types(cell::type cell_type)
 //-----------------------------------------------------------------------------
 template <std::floating_point T>
 std::pair<std::vector<T>, std::array<std::size_t, 3>>
-cell::facet_jacobians(cell::type cell_type)
+cell::entity_jacobians(cell::type cell_type, std::size_t e_dim)
 {
   std::size_t tdim = cell::topological_dimension(cell_type);
-  if (tdim != 2 and tdim != 3)
+  if ((e_dim >= tdim))
   {
     throw std::runtime_error(
-        "Facet jacobians not supported for this cell type.");
+        "Entity jacobians supported for entity of dimension "
+        + std::to_string(e_dim) + "for cell of dimension "
+        + std::to_string(tdim));
   }
 
   const auto [_x, xshape] = cell::geometry<T>(cell_type);
   mdspan_t<const T, 2> x(_x.data(), xshape);
-  std::vector<std::vector<int>> facets = topology(cell_type)[tdim - 1];
+  std::vector<std::vector<int>> entities = topology(cell_type)[e_dim];
 
-  std::array<std::size_t, 3> shape = {facets.size(), tdim, tdim - 1};
+  std::array<std::size_t, 3> shape = {entities.size(), tdim, e_dim};
   std::vector<T> jacobians(shape[0] * shape[1] * shape[2]);
   mdspan_t<T, 3> J(jacobians.data(), shape);
-  for (std::size_t f = 0; f < facets.size(); ++f)
+  for (std::size_t e = 0; e < entities.size(); ++e)
   {
-    const std::vector<int>& facet = facets[f];
-    for (std::size_t j = 0; j < tdim - 1; ++j)
+    const std::vector<int>& entity = entities[e];
+    for (std::size_t j = 0; j < e_dim; ++j)
       for (std::size_t k = 0; k < J.extent(1); ++k)
-        J(f, k, j) = x(facet[1 + j], k) - x(facet[0], k);
+        J(e, k, j) = x(entity[1 + j], k) - x(entity[0], k);
   }
 
   return {std::move(jacobians), std::move(shape)};
 }
 //-----------------------------------------------------------------------------
+template <std::floating_point T>
+std::pair<std::vector<T>, std::array<std::size_t, 3>>
+cell::facet_jacobians(cell::type cell_type)
+{
+  std::size_t tdim = cell::topological_dimension(cell_type);
+  if (tdim != 2 and tdim != 3)
+  {
+    throw std::runtime_error(
+        "Facet jacobians not supported for this cell type.");
+  }
+
+  return entity_jacobians<T>(cell_type, tdim - 1);
+}
+//-----------------------------------------------------------------------------
+template <std::floating_point T>
+std::pair<std::vector<T>, std::array<std::size_t, 3>>
+cell::edge_jacobians(cell::type cell_type)
+{
+  std::size_t tdim = cell::topological_dimension(cell_type);
+  if (tdim != 3)
+  {
+    throw std::runtime_error(
+        "Edge jacobians not supported for this cell type.");
+  }
+  return entity_jacobians<T>(cell_type, 1);
+}
+//-----------------------------------------------------------------------------
 
 /// @cond
 // Explicit instantiation for double and float
@@ -668,6 +697,10 @@ template std::pair<std::vector<float>, std::array<std::size_t, 3>>
     cell::facet_jacobians(cell::type);
 template std::pair<std::vector<double>, std::array<std::size_t, 3>>
     cell::facet_jacobians(cell::type);
+template std::pair<std::vector<float>, std::array<std::size_t, 3>>
+    cell::edge_jacobians(cell::type);
+template std::pair<std::vector<double>, std::array<std::size_t, 3>>
+    cell::edge_jacobians(cell::type);
 /// @endcond
 
 //-----------------------------------------------------------------------------

cpp/basix/cell.h

@@ -121,11 +121,29 @@ std::vector<T> facet_reference_volumes(cell::type cell_type);
 /// @return The subentity types. Indices are (tdim, entity)
 std::vector<std::vector<cell::type>> subentity_types(cell::type cell_type);
 
+/// Get the jacobians of a set of entities of a reference cell
+/// @param cell_type Type of cell
+/// @param e_dim Dimension of entities
+/// @return The jacobians of the facets (flattened) and the shape (nfacets,
+/// tdim, tdim - 1).
+template <std::floating_point T>
+std::pair<std::vector<T>, std::array<std::size_t, 3>>
+entity_jacobians(cell::type cell_type, std::size_t e_dim);
+
 /// Get the jacobians of the facets of a reference cell
 /// @param cell_type Type of cell
-/// @return The jacobians of the facets. Shape is (nfacets, gdim, gdim - 1)
+/// @return The jacobians of the facets (flattened) and the shape (nfacets,
+/// tdim, tdim - 1).
 template <std::floating_point T>
 std::pair<std::vector<T>, std::array<std::size_t, 3>>
 facet_jacobians(cell::type cell_type);
 
+/// Get the jacobians of the edegs of a reference cell
+/// @param cell_type Type of cell
+/// @return The jacobians of the edges (flattened) and the shape (nedges, tdim,
+/// tdim-2).
+template <std::floating_point T>
+std::pair<std::vector<T>, std::array<std::size_t, 3>>
+edge_jacobians(cell::type cell_type);
+
 } // namespace basix::cell