Add permute_subentity_closure function

cpp/basix/finite-element.cpp

@@ -1133,6 +1133,259 @@ FiniteElement<F>::FiniteElement(
     }
   }
 
+  // If DOF transformations are permutations, compute the subentity closure
+  // permutations
+  if (_dof_transformations_are_permutations)
+  {
+    if (_cell_tdim > 1)
+    {
+      // interval
+      int dof_n = 0;
+      const auto conn = cell::sub_entity_connectivity(_cell_type)[1][0];
+      std::vector<std::vector<std::vector<int>>> dofs;
+      dofs.resize(2);
+      for (int dim = 0; dim <= 1; ++dim)
+      {
+        dofs[dim].resize(conn[dim].size());
+        for (std::size_t i = 0; i < conn[dim].size(); ++i)
+        {
+          const int e = conn[dim][i];
+          for (int j : _edofs[dim][e])
+          {
+            std::ignore = j;
+            dofs[dim][i].push_back(dof_n++);
+          }
+        }
+      }
+
+      std::vector<std::size_t> ref;
+      for (int i : dofs[0][1])
+        ref.push_back(i);
+      for (int i : dofs[0][0])
+        ref.push_back(i);
+      // Edges
+      for (int i : dofs[1][0])
+        ref.push_back(i);
+
+      if (!_dof_transformations_are_identity)
+      {
+        auto& trans1 = _eperm.at(cell::type::interval)[0];
+        if (dofs[1][0].size() > 0)
+        {
+          precompute::apply_permutation(trans1, std::span(ref), dofs[1][0][0]);
+        }
+      }
+
+      precompute::prepare_permutation(ref);
+
+      auto& secp = _subentity_closure_perm.try_emplace(cell::type::interval)
+                       .first->second;
+      secp.push_back(ref);
+    }
+    if (_cell_type == cell::type::tetrahedron || cell_type == cell::type::prism
+        || cell_type == cell::type::pyramid)
+    {
+      // triangle
+      const int face_n = cell_type == cell::type::pyramid ? 1 : 0;
+
+      int dof_n = 0;
+      const auto conn = cell::sub_entity_connectivity(_cell_type)[2][face_n];
+      std::vector<std::vector<std::vector<int>>> dofs;
+      dofs.resize(3);
+      for (int dim = 0; dim <= 2; ++dim)
+      {
+        dofs[dim].resize(conn[dim].size());
+        for (std::size_t i = 0; i < conn[dim].size(); ++i)
+        {
+          const int e = conn[dim][i];
+          for (int j : _edofs[dim][e])
+          {
+            std::ignore = j;
+            dofs[dim][i].push_back(dof_n++);
+          }
+        }
+      }
+
+      std::vector<std::size_t> rot;
+      // Vertices
+      for (int i : dofs[0][1])
+        rot.push_back(i);
+      for (int i : dofs[0][2])
+        rot.push_back(i);
+      for (int i : dofs[0][0])
+        rot.push_back(i);
+      // Edges
+      for (int i : dofs[1][1])
+        rot.push_back(i);
+      for (int i : dofs[1][2])
+        rot.push_back(i);
+      for (int i : dofs[1][0])
+        rot.push_back(i);
+      // Face
+      for (int i : dofs[2][0])
+        rot.push_back(i);
+
+      std::vector<std::size_t> ref;
+      for (int i : dofs[0][0])
+        ref.push_back(i);
+      for (int i : dofs[0][2])
+        ref.push_back(i);
+      for (int i : dofs[0][1])
+        ref.push_back(i);
+      // Edges
+      for (int i : dofs[1][0])
+        ref.push_back(i);
+      for (int i : dofs[1][2])
+        ref.push_back(i);
+      for (int i : dofs[1][1])
+        ref.push_back(i);
+      // Face
+      for (int i : dofs[2][0])
+        ref.push_back(i);
+
+      if (!_dof_transformations_are_identity)
+      {
+        auto& trans1 = _eperm.at(cell::type::interval)[0];
+        auto& trans2 = _eperm.at(cell::type::triangle);
+
+        if (dofs[1][0].size() > 0)
+        {
+          precompute::apply_permutation(trans1, std::span(rot), dofs[1][0][0]);
+        }
+        if (dofs[1][1].size() > 0)
+        {
+          precompute::apply_permutation(trans1, std::span(rot), dofs[1][1][0]);
+        }
+        if (dofs[2][0].size() > 0)
+        {
+          precompute::apply_permutation(trans2[0], std::span(rot),
+                                        dofs[2][0][0]);
+        }
+        if (dofs[1][0].size() > 0)
+        {
+          precompute::apply_permutation(trans1, std::span(ref), dofs[1][0][0]);
+        }
+        if (dofs[2][0].size() > 0)
+        {
+          precompute::apply_permutation(trans2[1], std::span(ref),
+                                        dofs[2][0][0]);
+        }
+      }
+
+      precompute::prepare_permutation(rot);
+      precompute::prepare_permutation(ref);
+
+      auto& secp = _subentity_closure_perm.try_emplace(cell::type::triangle)
+                       .first->second;
+      secp.push_back(rot);
+      secp.push_back(ref);
+    }
+    if (_cell_type == cell::type::hexahedron || cell_type == cell::type::prism
+        || cell_type == cell::type::pyramid)
+    {
+      // quadrilateral
+      const int face_n = cell_type == cell::type::prism ? 1 : 0;
+
+      int dof_n = 0;
+      const auto conn = cell::sub_entity_connectivity(_cell_type)[2][face_n];
+      std::vector<std::vector<std::vector<int>>> dofs;
+      dofs.resize(3);
+      for (int dim = 0; dim <= 2; ++dim)
+      {
+        dofs[dim].resize(conn[dim].size());
+        for (std::size_t i = 0; i < conn[dim].size(); ++i)
+        {
+          const int e = conn[dim][i];
+          for (int j : _edofs[dim][e])
+          {
+            std::ignore = j;
+            dofs[dim][i].push_back(dof_n++);
+          }
+        }
+      }
+
+      std::vector<std::size_t> rot;
+      // Vertices
+      for (int i : dofs[0][1])
+        rot.push_back(i);
+      for (int i : dofs[0][3])
+        rot.push_back(i);
+      for (int i : dofs[0][0])
+        rot.push_back(i);
+      for (int i : dofs[0][2])
+        rot.push_back(i);
+      // Edges
+      for (int i : dofs[1][2])
+        rot.push_back(i);
+      for (int i : dofs[1][0])
+        rot.push_back(i);
+      for (int i : dofs[1][3])
+        rot.push_back(i);
+      for (int i : dofs[1][1])
+        rot.push_back(i);
+      // Face
+      for (int i : dofs[2][0])
+        rot.push_back(i);
+
+      std::vector<std::size_t> ref;
+      for (int i : dofs[0][0])
+        ref.push_back(i);
+      for (int i : dofs[0][2])
+        ref.push_back(i);
+      for (int i : dofs[0][1])
+        ref.push_back(i);
+      for (int i : dofs[0][3])
+        ref.push_back(i);
+      // Edges
+      for (int i : dofs[1][1])
+        ref.push_back(i);
+      for (int i : dofs[1][0])
+        ref.push_back(i);
+      for (int i : dofs[1][3])
+        ref.push_back(i);
+      for (int i : dofs[1][2])
+        ref.push_back(i);
+      // Face
+      for (int i : dofs[2][0])
+        ref.push_back(i);
+
+      if (!_dof_transformations_are_identity)
+      {
+        auto& trans1 = _eperm.at(cell::type::interval)[0];
+        auto& trans2 = _eperm.at(cell::type::quadrilateral);
+
+        if (dofs[1][1].size() > 0)
+        {
+          precompute::apply_permutation(trans1, std::span(rot), dofs[1][1][0]);
+        }
+        if (dofs[1][2].size() > 0)
+        {
+          precompute::apply_permutation(trans1, std::span(rot), dofs[1][2][0]);
+        }
+        if (dofs[2][0].size() > 0)
+        {
+          precompute::apply_permutation(trans2[0], std::span(rot),
+                                        dofs[2][0][0]);
+        }
+
+        if (dofs[2][0].size() > 0)
+        {
+          precompute::apply_permutation(trans2[1], std::span(ref),
+                                        dofs[2][0][0]);
+        }
+      }
+
+      precompute::prepare_permutation(rot);
+      precompute::prepare_permutation(ref);
+
+      auto& secp
+          = _subentity_closure_perm.try_emplace(cell::type::quadrilateral)
+                .first->second;
+      secp.push_back(rot);
+      secp.push_back(ref);
+    }
+  }
+
   // Check if interpolation matrix is the identity
   mdspan_t<const F, 2> matM(_matM.first.data(), _matM.second);
   _interpolation_is_identity = matM.extent(0) == matM.extent(1);

cpp/basix/finite-element.h

@@ -837,6 +837,64 @@ class FiniteElement
       permute_data<std::int32_t, true>(d, 1, cell_info, _eperm_rev);
   }
 
+  /// @brief Permute indices associated with degree-of-freedoms on the
+  /// closure of a sub-entity of the reference element
+  ///
+  /// This function performs a similar permutation to permute() but
+  /// additionally permutes the positions of vertices and edges
+  ///
+  /// @note This function is designed to be called at runtime, so its
+  /// performance is critical.
+  ///
+  /// @param[in,out] d Indices associated with each reference element
+  /// degree-of-freedom (in). Indices associated with each physical
+  /// element degree-of-freedom (out).
+  /// @param cell_info Permutation info for the cell
+  /// @param entity_type The cell type of the sub-entity
+  void permute_subentity_closure(std::span<std::int32_t> d,
+                                 std::uint32_t cell_info,
+                                 cell::type entity_type) const
+  {
+    if (!_dof_transformations_are_permutations)
+    {
+      throw std::runtime_error(
+          "The DOF transformations for this element are not permutations");
+    }
+
+    const int entity_dim = cell::topological_dimension(entity_type);
+
+    if (entity_dim == 0)
+      return;
+
+    auto& perm = _subentity_closure_perm.at(entity_type);
+    if (entity_dim == 1)
+    {
+      if (cell_info & 1)
+      {
+        precompute::apply_permutation(perm[0], d);
+      }
+    }
+    else if (entity_dim == 2)
+    {
+      // Rotate a face
+      for (std::uint32_t r = 0; r < (cell_info >> 1 & 3); ++r)
+      {
+        precompute::apply_permutation(perm[0], d);
+      }
+
+      // Reflect a face (post rotate)
+      if (cell_info & 1)
+      {
+        precompute::apply_permutation(perm[1], d);
+      }
+    }
+    else
+    {
+      throw std::runtime_error(
+          "Invalid dimension for permute_subentity_closure");
+    }
+  }
+
   /// @brief Transform basis functions from the reference element
   /// ordering and orientation to the globally consistent physical
   /// element ordering and orientation.
@@ -1323,6 +1381,11 @@ class FiniteElement
   // The inverse transpose entity transformations in precomputed form
   std::map<cell::type, trans_data_t> _etrans_invT;
 
+  // The subentity closure permutations (factorised). This will only be set if
+  // _dof_transformations_are_permutations is True
+  std::map<cell::type, std::vector<std::vector<std::size_t>>>
+      _subentity_closure_perm;
+
   // Indicates whether or not this is the discontinuous version of the
   // element
   bool _discontinuous;

python/wrapper.cpp

@@ -98,6 +98,15 @@ void declare_float(nb::module_& m, std::string type)
            })
       .def("__eq__", &FiniteElement<T>::operator==)
       .def("__hash__", &FiniteElement<T>::hash)
+      .def("permute_subentity_closure",
+           [](const FiniteElement<T>& self,
+              nb::ndarray<std::int32_t, nb::ndim<1>, nb::c_contig> d,
+              std::uint32_t cell_info, cell::type entity_type)
+           {
+             std::span<std::int32_t> _d(d.data(), d.shape(0));
+             self.permute_subentity_closure(_d, cell_info, entity_type);
+             return as_nbarray(std::move(_d));
+           })
       .def("push_forward",
            [](const FiniteElement<T>& self,
               nb::ndarray<const T, nb::ndim<3>, nb::c_contig> U,