finite-element.h

// Copyright (c) 2020-2024 Chris Richardson, Matthew Scroggs and Garth . Wells
// FEniCS Project
// SPDX-License-Identifier:    MIT
 
#pragma once
 
#include "cell.h"
#include "element-families.h"
#include "maps.h"
#include "mdspan.hpp"
#include "polyset.h"
#include "precompute.h"
#include "sobolev-spaces.h"
#include <array>
#include <concepts>
#include <cstdint>
#include <functional>
#include <map>
#include <numeric>
#include <span>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

namespace basix
{
namespace impl
{
template <typename T, std::size_t d>
using mdspan_t = md::mdspan<T, md::dextents<std::size_t, d>>;
template <typename T, std::size_t d>
using mdarray_t
    = md::MDSPAN_IMPL_PROPOSED_NAMESPACE::mdarray<T,
                                                  md::dextents<std::size_t, d>>;

template <typename T>
std::array<std::vector<mdspan_t<const T, 2>>, 4>
to_mdspan(std::array<std::vector<mdarray_t<T, 2>>, 4>& x)
{
  std::array<std::vector<mdspan_t<const T, 2>>, 4> x1;
  for (std::size_t i = 0; i < x.size(); ++i)
    for (std::size_t j = 0; j < x[i].size(); ++j)
      x1[i].emplace_back(x[i][j].data(), x[i][j].extents());
 
  return x1;
}

template <typename T>
std::array<std::vector<mdspan_t<const T, 4>>, 4>
to_mdspan(std::array<std::vector<mdarray_t<T, 4>>, 4>& M)
{
  std::array<std::vector<mdspan_t<const T, 4>>, 4> M1;
  for (std::size_t i = 0; i < M.size(); ++i)
    for (std::size_t j = 0; j < M[i].size(); ++j)
      M1[i].emplace_back(M[i][j].data(), M[i][j].extents());
 
  return M1;
}

template <typename T>
std::array<std::vector<mdspan_t<const T, 2>>, 4>
to_mdspan(const std::array<std::vector<std::vector<T>>, 4>& x,
          const std::array<std::vector<std::array<std::size_t, 2>>, 4>& shape)
{
  std::array<std::vector<mdspan_t<const T, 2>>, 4> x1;
  for (std::size_t i = 0; i < x.size(); ++i)
    for (std::size_t j = 0; j < x[i].size(); ++j)
      x1[i].push_back(mdspan_t<const T, 2>(x[i][j].data(), shape[i][j]));
 
  return x1;
}

template <typename T>
std::array<std::vector<mdspan_t<const T, 4>>, 4>
to_mdspan(const std::array<std::vector<std::vector<T>>, 4>& M,
          const std::array<std::vector<std::array<std::size_t, 4>>, 4>& shape)
{
  std::array<std::vector<mdspan_t<const T, 4>>, 4> M1;
  for (std::size_t i = 0; i < M.size(); ++i)
    for (std::size_t j = 0; j < M[i].size(); ++j)
      M1[i].push_back(mdspan_t<const T, 4>(M[i][j].data(), shape[i][j]));
 
  return M1;
}
 
} // namespace impl
 
namespace element
{
template <typename T, std::size_t d>
using mdspan_t = impl::mdspan_t<T, d>;

template <std::floating_point T>
std::tuple<std::array<std::vector<std::vector<T>>, 4>,
           std::array<std::vector<std::array<std::size_t, 2>>, 4>,
           std::array<std::vector<std::vector<T>>, 4>,
           std::array<std::vector<std::array<std::size_t, 4>>, 4>>
make_discontinuous(const std::array<std::vector<mdspan_t<const T, 2>>, 4>& x,
                   const std::array<std::vector<mdspan_t<const T, 4>>, 4>& M,
                   std::size_t tdim, std::size_t value_size);
 
} // namespace element

template <std::floating_point F>
class FiniteElement
{
  template <typename T, std::size_t d>
  using mdspan_t = md::mdspan<T, md::dextents<std::size_t, d>>;
 
public:
  using scalar_type = F;

  FiniteElement(element::family family, cell::type cell_type,
                polyset::type poly_type, int degree,
                const std::vector<std::size_t>& value_shape,
                mdspan_t<const F, 2> wcoeffs,
                const std::array<std::vector<mdspan_t<const F, 2>>, 4>& x,
                const std::array<std::vector<mdspan_t<const F, 4>>, 4>& M,
                int interpolation_nderivs, maps::type map_type,
                sobolev::space sobolev_space, bool discontinuous,
                int embedded_subdegree, int embedded_superdegree,
                element::lagrange_variant lvariant,
                element::dpc_variant dvariant,
                std::vector<int> dof_ordering = {});

  FiniteElement(const FiniteElement& element) = default;

  FiniteElement(FiniteElement&& element) = default;

  ~FiniteElement() = default;

  FiniteElement& operator=(const FiniteElement& element) = default;

  FiniteElement& operator=(FiniteElement&& element) = default;

  bool operator==(const FiniteElement& e) const;

  std::size_t hash() const;

  std::array<std::size_t, 4> tabulate_shape(std::size_t nd,
                                            std::size_t num_points) const
  {
    std::size_t ndsize = 1;
    for (std::size_t i = 1; i <= nd; ++i)
      ndsize *= (_cell_tdim + i);
    for (std::size_t i = 1; i <= nd; ++i)
      ndsize /= i;
    std::size_t vs = std::accumulate(_value_shape.begin(), _value_shape.end(),
                                     1, std::multiplies{});
    std::size_t ndofs = _coeffs.second[0];
    return {ndsize, num_points, ndofs, vs};
  }

  std::pair<std::vector<F>, std::array<std::size_t, 4>>
  tabulate(int nd, impl::mdspan_t<const F, 2> x) const;

  std::pair<std::vector<F>, std::array<std::size_t, 4>>
  tabulate(int nd, std::span<const F> x,
           std::array<std::size_t, 2> shape) const;

  void tabulate(int nd, impl::mdspan_t<const F, 2> x,
                mdspan_t<F, 4> basis) const;

  void tabulate(int nd, std::span<const F> x, std::array<std::size_t, 2> xshape,
                std::span<F> basis) const;

  cell::type cell_type() const { return _cell_type; }

  polyset::type polyset_type() const { return _poly_type; }

  int degree() const { return _degree; }

  int embedded_superdegree() const { return _embedded_superdegree; }

  int embedded_subdegree() const { return _embedded_subdegree; }

  const std::vector<std::size_t>& value_shape() const { return _value_shape; }

  int dim() const { return _coeffs.second[0]; }

  element::family family() const { return _family; }

  element::lagrange_variant lagrange_variant() const
  {
    return _lagrange_variant;
  }

  element::dpc_variant dpc_variant() const { return _dpc_variant; }

  maps::type map_type() const { return _map_type; }

  sobolev::space sobolev_space() const { return _sobolev_space; }

  bool discontinuous() const { return _discontinuous; }

  bool dof_transformations_are_permutations() const
  {
    return _dof_transformations_are_permutations;
  }

  bool dof_transformations_are_identity() const
  {
    return _dof_transformations_are_identity;
  }

  std::pair<std::vector<F>, std::array<std::size_t, 3>>
  push_forward(impl::mdspan_t<const F, 3> U, impl::mdspan_t<const F, 3> J,
               std::span<const F> detJ, impl::mdspan_t<const F, 3> K) const;

  std::pair<std::vector<F>, std::array<std::size_t, 3>>
  pull_back(impl::mdspan_t<const F, 3> u, impl::mdspan_t<const F, 3> J,
            std::span<const F> detJ, impl::mdspan_t<const F, 3> K) const;

  template <typename O, typename P, typename Q, typename R>
  std::function<void(O&, const P&, const Q&, F, const R&)> map_fn() const
  {
    switch (_map_type)
    {
    case maps::type::identity:
      return [](O& u, const P& U, const Q&, F, const R&)
      {
        assert(U.extent(0) == u.extent(0));
        assert(U.extent(1) == u.extent(1));
        for (std::size_t i = 0; i < U.extent(0); ++i)
          for (std::size_t j = 0; j < U.extent(1); ++j)
            u(i, j) = U(i, j);
      };
    case maps::type::covariantPiola:
      return [](O& u, const P& U, const Q& J, F detJ, const R& K)
      { maps::covariant_piola(u, U, J, detJ, K); };
    case maps::type::contravariantPiola:
      return [](O& u, const P& U, const Q& J, F detJ, const R& K)
      { maps::contravariant_piola(u, U, J, detJ, K); };
    case maps::type::doubleCovariantPiola:
      return [](O& u, const P& U, const Q& J, F detJ, const R& K)
      { maps::double_covariant_piola(u, U, J, detJ, K); };
    case maps::type::doubleContravariantPiola:
      return [](O& u, const P& U, const Q& J, F detJ, const R& K)
      { maps::double_contravariant_piola(u, U, J, detJ, K); };
    default:
      throw std::runtime_error("Map not implemented");
    }
  }

  const std::vector<std::vector<std::vector<int>>>& entity_dofs() const
  {
    return _edofs;
  }

  const std::vector<std::vector<std::vector<int>>>& entity_closure_dofs() const
  {
    return _e_closure_dofs;
  }

  std::pair<std::vector<F>, std::array<std::size_t, 3>>
  base_transformations() const;

  std::map<cell::type, std::pair<std::vector<F>, std::array<std::size_t, 3>>>
  entity_transformations() const
  {
    return _entity_transformations;
  }

  void permute(std::span<std::int32_t> d, std::uint32_t cell_info) const
  {
    if (!_dof_transformations_are_permutations)
    {
      throw std::runtime_error(
          "The DOF transformations for this element are not permutations");
    }
 
    if (_dof_transformations_are_identity)
      return;
    else
      permute_data<std::int32_t, false>(d, 1, cell_info, _eperm);
  }

  void permute_inv(std::span<std::int32_t> d, std::uint32_t cell_info) const
  {
    if (!_dof_transformations_are_permutations)
    {
      throw std::runtime_error(
          "The DOF transformations for this element are not permutations");
    }
 
    if (_dof_transformations_are_identity)
      return;
    else
      permute_data<std::int32_t, true>(d, 1, cell_info, _eperm_inv);
  }

  void permute_subentity_closure(std::span<std::int32_t> d,
                                 std::uint32_t cell_info,
                                 cell::type entity_type, int entity_index) const
  {
    const int entity_dim = cell::topological_dimension(entity_type);
 
    int face_start = _cell_tdim == 3 ? 3 * _edofs[2].size() : 0;
 
    std::uint32_t entity_info = 0;
    switch (entity_dim)
    {
    case 0:
      return;
    case 1:
      entity_info = cell_info >> (face_start + entity_index) & 1;
      break;
    case 2:
      entity_info = cell_info >> (3 * entity_index) & 7;
      break;
    default:
      throw std::runtime_error("Unsupported cell dimension");
    }
    permute_subentity_closure(d, entity_info, entity_type);
  }

  void permute_subentity_closure_inv(std::span<std::int32_t> d,
                                     std::uint32_t cell_info,
                                     cell::type entity_type,
                                     int entity_index) const
  {
    const int entity_dim = cell::topological_dimension(entity_type);
 
    int face_start = _cell_tdim == 3 ? 3 * _edofs[2].size() : 0;
 
    std::uint32_t entity_info;
    switch (entity_dim)
    {
    case 0:
      return;
    case 1:
      entity_info = cell_info >> (face_start + entity_index) & 1;
      break;
    case 2:
      entity_info = cell_info >> (3 * entity_index) & 7;
      break;
    default:
      throw std::runtime_error("Unsupported cell dimension");
    }
    permute_subentity_closure_inv(d, entity_info, entity_type);
  }

 
  void permute_subentity_closure(std::span<std::int32_t> d,
                                 std::uint32_t entity_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 (entity_info & 1)
      {
        precompute::apply_permutation(perm[0], d);
      }
    }
    else if (entity_dim == 2)
    {
      // Rotate a face
      for (std::uint32_t r = 0; r < (entity_info >> 1 & 3); ++r)
      {
        precompute::apply_permutation(perm[0], d);
      }
 
      // Reflect a face (post rotate)
      if (entity_info & 1)
      {
        precompute::apply_permutation(perm[1], d);
      }
    }
    else
    {
      throw std::runtime_error(
          "Invalid dimension for permute_subentity_closure");
    }
  }

  void permute_subentity_closure_inv(std::span<std::int32_t> d,
                                     std::uint32_t entity_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_inv.at(entity_type);
    if (entity_dim == 1)
    {
      if (entity_info & 1)
      {
        precompute::apply_permutation(perm[0], d);
      }
    }
    else if (entity_dim == 2)
    {
      // Reflect a face (pre rotate)
      if (entity_info & 1)
      {
        precompute::apply_permutation(perm[1], d);
      }
 
      // Rotate a face
      for (std::uint32_t r = 0; r < (entity_info >> 1 & 3); ++r)
      {
        precompute::apply_permutation(perm[0], d);
      }
    }
    else
    {
      throw std::runtime_error(
          "Invalid dimension for permute_subentity_closure");
    }
  }

  template <typename T>
  void T_apply(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void Tt_apply(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void Tt_inv_apply(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void Tinv_apply(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void Tt_apply_right(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void T_apply_right(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void Tinv_apply_right(std::span<T> u, int n, std::uint32_t cell_info) const;

  template <typename T>
  void Tt_inv_apply_right(std::span<T> u, int n, std::uint32_t cell_info) const;

  const std::pair<std::vector<F>, std::array<std::size_t, 2>>& points() const
  {
    return _points;
  }

  const std::pair<std::vector<F>, std::array<std::size_t, 2>>&
  interpolation_matrix() const
  {
    return _matM;
  }

  const std::pair<std::vector<F>, std::array<std::size_t, 2>>&
  dual_matrix() const
  {
    return _dual_matrix;
  }

  const std::pair<std::vector<F>, std::array<std::size_t, 2>>& wcoeffs() const
  {
    return _wcoeffs;
  }

  const std::array<
      std::vector<std::pair<std::vector<F>, std::array<std::size_t, 2>>>, 4>&
  x() const
  {
    return _x;
  }

  const std::array<
      std::vector<std::pair<std::vector<F>, std::array<std::size_t, 4>>>, 4>&
  M() const
  {
    return _M;
  }

  const std::pair<std::vector<F>, std::array<std::size_t, 2>>&
  coefficient_matrix() const
  {
    return _coeffs;
  }

  bool has_tensor_product_factorisation() const
  {
    return !_tensor_factors.empty();
  }

  std::vector<std::vector<FiniteElement<F>>>
  get_tensor_product_representation() const
  {
    if (!has_tensor_product_factorisation())
      throw std::runtime_error("Element has no tensor product representation.");
    return _tensor_factors;
  }

  bool interpolation_is_identity() const { return _interpolation_is_identity; }

  int interpolation_nderivs() const { return _interpolation_nderivs; }

  const std::vector<int>& dof_ordering() const { return _dof_ordering; }
 
private:
  template <typename T, bool post>
  void permute_data(
      std::span<T> data, int block_size, std::uint32_t cell_info,
      const std::map<cell::type, std::vector<std::vector<std::size_t>>>& eperm)
      const;
 
  using array2_t = std::pair<std::vector<F>, std::array<std::size_t, 2>>;
  using array3_t = std::pair<std::vector<F>, std::array<std::size_t, 3>>;
  using trans_data_t
      = std::vector<std::pair<std::vector<std::size_t>, array2_t>>;

  template <typename T, bool post, typename OP>
  void
  transform_data(std::span<T> data, int block_size, std::uint32_t cell_info,
                 const std::map<cell::type, trans_data_t>& etrans, OP op) const;
 
  // Cell type
  cell::type _cell_type;
 
  // Polyset type
  polyset::type _poly_type;
 
  // Topological dimension of the cell
  std::size_t _cell_tdim;
 
  // Topological dimension of the cell
  std::vector<std::vector<cell::type>> _cell_subentity_types;
 
  // Finite element family
  element::family _family;
 
  // Lagrange variant
  element::lagrange_variant _lagrange_variant;
 
  // DPC variant
  element::dpc_variant _dpc_variant;
 
  // Degree that was input when creating the element
  int _degree;
 
  // Degree
  int _interpolation_nderivs;
 
  // Highest degree polynomial in element's polyset
  int _embedded_superdegree;
 
  // Highest degree space that is a subspace of element's polyset
  int _embedded_subdegree;
 
  // Value shape
  std::vector<std::size_t> _value_shape;

  maps::type _map_type;

  sobolev::space _sobolev_space;
 
  // Shape function coefficient of expansion sets on cell. If shape
  // function is given by @f$\psi_i = \sum_{k} \phi_{k}
  // \alpha^{i}_{k}@f$, then _coeffs(i, j) = @f$\alpha^i_k@f$. ie
  // _coeffs.row(i) are the expansion coefficients for shape function i
  // (@f$\psi_{i}@f$).
  std::pair<std::vector<F>, std::array<std::size_t, 2>> _coeffs;
 
  // Dofs associated with each cell (sub-)entity
  std::vector<std::vector<std::vector<int>>> _edofs;
 
  // Dofs associated with the closdure of each cell (sub-)entity
  std::vector<std::vector<std::vector<int>>> _e_closure_dofs;
 
  // Entity transformations
  std::map<cell::type, array3_t> _entity_transformations;
 
  // Set of points used for point evaluation
  // Experimental - currently used for an implementation of
  // "tabulate_dof_coordinates" Most useful for Lagrange. This may change or go
  // away. For non-Lagrange elements, these points will be used in combination
  // with _interpolation_matrix to perform interpolation
  std::pair<std::vector<F>, std::array<std::size_t, 2>> _points;
 
  // Interpolation points on the cell. The shape is (entity_dim, num
  // entities of given dimension, num_points, tdim)
  std::array<std::vector<std::pair<std::vector<F>, std::array<std::size_t, 2>>>,
             4>
      _x;

  std::pair<std::vector<F>, std::array<std::size_t, 2>> _matM;
 
  // Indicates whether or not the DOF transformations are all
  // permutations
  bool _dof_transformations_are_permutations;
 
  // Indicates whether or not the DOF transformations are all identity
  bool _dof_transformations_are_identity;
 
  // The entity permutations (factorised). This will only be set if
  // _dof_transformations_are_permutations is True and
  // _dof_transformations_are_identity is False
  std::map<cell::type, std::vector<std::vector<std::size_t>>> _eperm;
 
  // The reverse entity permutations (factorised). This will only be set
  // if _dof_transformations_are_permutations is True and
  // _dof_transformations_are_identity is False
  std::map<cell::type, std::vector<std::vector<std::size_t>>> _eperm_inv;
 
  // The entity transformations in precomputed form
  std::map<cell::type, trans_data_t> _etrans;
 
  // The transposed entity transformations in precomputed form
  std::map<cell::type, trans_data_t> _etransT;
 
  // The inverse entity transformations in precomputed form
  std::map<cell::type, trans_data_t> _etrans_inv;
 
  // 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;
 
  // The inverse 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_inv;
 
  // Indicates whether or not this is the discontinuous version of the
  // element
  bool _discontinuous;
 
  // The dual matrix
  std::pair<std::vector<F>, std::array<std::size_t, 2>> _dual_matrix;
 
  // Dof reordering for different element dof layout compatibility.
  // The reference basix layout is ordered by entity, i.e. dofs on
  // vertices, followed by edges, faces, then internal dofs.
  // _dof_ordering stores the map to the new order required, e.g.
  // for a P2 triangle, _dof_ordering=[0 3 5 1 2 4] will place
  // dofs 0, 3, 5 on the vertices and 1, 2, 4, on the edges.
  std::vector<int> _dof_ordering;
 
  // Tensor product representation
  // Entries of tuple are (list of elements on an interval, permutation
  // of DOF numbers)
  // @todo: For vector-valued elements, a tensor product type and a
  // scaling factor may additionally be needed.
  std::vector<std::vector<FiniteElement>> _tensor_factors;
 
  // Is the interpolation matrix an identity?
  bool _interpolation_is_identity;
 
  // The coefficients that define the polynomial set in terms of the
  // orthonormal polynomials
  std::pair<std::vector<F>, std::array<std::size_t, 2>> _wcoeffs;
 
  // Interpolation matrices for each entity
  using array4_t
      = std::vector<std::pair<std::vector<F>, std::array<std::size_t, 4>>>;
  std::array<array4_t, 4> _M;
};

template <std::floating_point T>
FiniteElement<T> create_custom_element(
    cell::type cell_type, const std::vector<std::size_t>& value_shape,
    impl::mdspan_t<const T, 2> wcoeffs,
    const std::array<std::vector<impl::mdspan_t<const T, 2>>, 4>& x,
    const std::array<std::vector<impl::mdspan_t<const T, 4>>, 4>& M,
    int interpolation_nderivs, maps::type map_type,
    sobolev::space sobolev_space, bool discontinuous, int embedded_subdegree,
    int embedded_superdegree, polyset::type poly_type);

template <std::floating_point T>
FiniteElement<T> create_element(element::family family, cell::type cell,
                                int degree, element::lagrange_variant lvariant,
                                element::dpc_variant dvariant,
                                bool discontinuous,
                                std::vector<int> dof_ordering = {});

std::vector<int> tp_dof_ordering(element::family family, cell::type cell,
                                 int degree, element::lagrange_variant lvariant,
                                 element::dpc_variant dvariant,
                                 bool discontinuous);

std::vector<int> lex_dof_ordering(element::family family, cell::type cell,
                                  int degree, element::lagrange_variant lvariant,
                                  element::dpc_variant dvariant,
                                  bool discontinuous);

template <std::floating_point T>
std::vector<std::vector<FiniteElement<T>>>
tp_factors(element::family family, cell::type cell, int degree,
           element::lagrange_variant lvariant, element::dpc_variant dvariant,
           bool discontinuous, const std::vector<int>& dof_ordering);

template <std::floating_point T>
FiniteElement<T>
create_tp_element(element::family family, cell::type cell, int degree,
                  element::lagrange_variant lvariant,
                  element::dpc_variant dvariant, bool discontinuous);

std::string version();
 
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T, bool post>
void FiniteElement<F>::permute_data(
    std::span<T> data, int block_size, std::uint32_t cell_info,
    const std::map<cell::type, std::vector<std::vector<std::size_t>>>& eperm)
    const
{
  if (_cell_tdim >= 2)
  {
    // This assumes 3 bits are used per face. This will need updating if 3D
    // cells with faces with more than 4 sides are implemented
    int face_start = _cell_tdim == 3 ? 3 * _edofs[2].size() : 0;
 
    // Permute DOFs on edges
    {
      auto& trans = eperm.at(cell::type::interval)[0];
      for (std::size_t e = 0; e < _edofs[1].size(); ++e)
      {
        // Reverse an edge
        if (cell_info >> (face_start + e) & 1)
        {
          precompute::apply_permutation_mapped(trans, data, _edofs[1][e],
                                               block_size);
        }
      }
    }
 
    if (_cell_tdim == 3)
    {
      // Permute DOFs on faces
      for (std::size_t f = 0; f < _edofs[2].size(); ++f)
      {
        auto& trans = eperm.at(_cell_subentity_types[2][f]);
 
        // Reflect a face (pre rotate)
        if (!post and cell_info >> (3 * f) & 1)
        {
          precompute::apply_permutation_mapped(trans[1], data, _edofs[2][f],
                                               block_size);
        }
 
        // Rotate a face
        for (std::uint32_t r = 0; r < (cell_info >> (3 * f + 1) & 3); ++r)
        {
          precompute::apply_permutation_mapped(trans[0], data, _edofs[2][f],
                                               block_size);
        }
 
        // Reflect a face (post rotate)
        if (post and cell_info >> (3 * f) & 1)
        {
          precompute::apply_permutation_mapped(trans[1], data, _edofs[2][f],
                                               block_size);
        }
      }
    }
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T, bool post, typename OP>
void FiniteElement<F>::transform_data(
    std::span<T> data, int block_size, std::uint32_t cell_info,
    const std::map<cell::type, trans_data_t>& etrans, OP op) const
{
  if (_cell_tdim >= 2)
  {
    // This assumes 3 bits are used per face. This will need updating if
    // 3D cells with faces with more than 4 sides are implemented
    int face_start = _cell_tdim == 3 ? 3 * _edofs[2].size() : 0;
    int dofstart = 0;
    for (auto& edofs0 : _edofs[0])
      dofstart += edofs0.size();
 
    // Transform DOFs on edges
    {
      auto& [v_size_t, matrix] = etrans.at(cell::type::interval)[0];
      for (std::size_t e = 0; e < _edofs[1].size(); ++e)
      {
        // Reverse an edge
        if (cell_info >> (face_start + e) & 1)
        {
          op(std::span(v_size_t),
             mdspan_t<const F, 2>(matrix.first.data(), matrix.second), data,
             dofstart, block_size);
        }
        dofstart += _edofs[1][e].size();
      }
    }
 
    if (_cell_tdim == 3)
    {
      // Permute DOFs on faces
      for (std::size_t f = 0; f < _edofs[2].size(); ++f)
      {
        auto& trans = etrans.at(_cell_subentity_types[2][f]);
 
        // Reflect a face (pre rotation)
        if (!post and cell_info >> (3 * f) & 1)
        {
          const auto& m = trans[1];
          const auto& v_size_t = std::get<0>(m);
          const auto& matrix = std::get<1>(m);
          op(std::span(v_size_t),
             mdspan_t<const F, 2>(matrix.first.data(), matrix.second), data,
             dofstart, block_size);
        }
 
        // Rotate a face
        for (std::uint32_t r = 0; r < (cell_info >> (3 * f + 1) & 3); ++r)
        {
          const auto& m = trans[0];
          const auto& v_size_t = std::get<0>(m);
          const auto& matrix = std::get<1>(m);
          op(std::span(v_size_t),
             mdspan_t<const F, 2>(matrix.first.data(), matrix.second), data,
             dofstart, block_size);
        }
 
        // Reflect a face (post rotation)
        if (post and cell_info >> (3 * f) & 1)
        {
          const auto& m = trans[1];
          const auto& v_size_t = std::get<0>(m);
          const auto& matrix = std::get<1>(m);
          op(std::span(v_size_t),
             mdspan_t<const F, 2>(matrix.first.data(), matrix.second), data,
             dofstart, block_size);
        }
 
        dofstart += _edofs[2][f].size();
      }
    }
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::T_apply(std::span<T> u, int n,
                               std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
 
  if (_dof_transformations_are_permutations)
    permute_data<T, false>(u, n, cell_info, _eperm);
  else
  {
    transform_data<T, false>(u, n, cell_info, _etrans,
                             precompute::apply_matrix<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::Tt_apply(std::span<T> u, int n,
                                std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
    permute_data<T, true>(u, n, cell_info, _eperm_inv);
  else
  {
    transform_data<T, true>(u, n, cell_info, _etransT,
                            precompute::apply_matrix<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::Tt_inv_apply(std::span<T> u, int n,
                                    std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
    permute_data<T, false>(u, n, cell_info, _eperm);
  else
  {
    transform_data<T, false>(u, n, cell_info, _etrans_invT,
                             precompute::apply_matrix<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::Tinv_apply(std::span<T> u, int n,
                                  std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
    permute_data<T, true>(u, n, cell_info, _eperm_inv);
  else
  {
    transform_data<T, true>(u, n, cell_info, _etrans_inv,
                            precompute::apply_matrix<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::Tt_apply_right(std::span<T> u, int n,
                                      std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
  {
    assert(u.size() % n == 0);
    const int step = u.size() / n;
    for (int i = 0; i < n; ++i)
    {
      std::span<T> dblock(u.data() + i * step, step);
      permute_data<T, false>(dblock, 1, cell_info, _eperm);
    }
  }
  else
  {
    transform_data<T, false>(u, n, cell_info, _etrans,
                             precompute::apply_tranpose_matrix_right<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::Tinv_apply_right(std::span<T> u, int n,
                                        std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
  {
    assert(u.size() % n == 0);
    const int step = u.size() / n;
    for (int i = 0; i < n; ++i)
    {
      std::span<T> dblock(u.data() + i * step, step);
      permute_data<T, false>(dblock, 1, cell_info, _eperm);
    }
  }
  else
  {
    transform_data<T, false>(u, n, cell_info, _etrans_invT,
                             precompute::apply_tranpose_matrix_right<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::T_apply_right(std::span<T> u, int n,
                                     std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
  {
    assert(u.size() % n == 0);
    const int step = u.size() / n;
    for (int i = 0; i < n; ++i)
    {
      std::span<T> dblock(u.data() + i * step, step);
      permute_data<T, true>(dblock, 1, cell_info, _eperm_inv);
    }
  }
  else
  {
    transform_data<T, true>(u, n, cell_info, _etransT,
                            precompute::apply_tranpose_matrix_right<F, T>);
  }
}
//-----------------------------------------------------------------------------
template <std::floating_point F>
template <typename T>
void FiniteElement<F>::Tt_inv_apply_right(std::span<T> u, int n,
                                          std::uint32_t cell_info) const
{
  if (_dof_transformations_are_identity)
    return;
  else if (_dof_transformations_are_permutations)
  {
    assert(u.size() % n == 0);
    const int step = u.size() / n;
    for (int i = 0; i < n; ++i)
    {
      std::span<T> dblock(u.data() + i * step, step);
      permute_data<T, true>(dblock, 1, cell_info, _eperm_inv);
    }
  }
  else
  {
    transform_data<T, true>(u, n, cell_info, _etrans_inv,
                            precompute::apply_tranpose_matrix_right<F, T>);
  }
}
//-----------------------------------------------------------------------------
 
} // namespace basix