FiniteElement.h

// Copyright (C) 2020-2026 Garth N. Wells, Matthew W. Scroggs and Jørgen S.
// Dokken
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "traits.h"
#include <algorithm>
#include <array>
#include <basix/finite-element.h>
#include <basix/maps.h>
#include <cassert>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <dolfinx/mesh/cell_types.h>
#include <dolfinx/mesh/permutationcomputation.h>
#include <functional>
#include <memory>
#include <optional>
#include <span>
#include <stdexcept>
#include <utility>
#include <vector>
 
namespace dolfinx::fem
{
enum class doftransform : std::uint8_t
{
  standard = 0,          
  transpose = 1,         
  inverse = 2,           
  inverse_transpose = 3, 
};
 
template <std::floating_point T>
class FiniteElement;

std::vector<std::size_t>
compute_value_shape(basix::maps::type map_type,
                    std::span<const std::size_t> reference_value_shape,
                    std::size_t gdim);

template <std::floating_point T>
struct BasixElementData
{
  std::reference_wrapper<const basix::FiniteElement<T>>
      element; 
  std::optional<std::vector<std::size_t>> value_shape = std::nullopt;
  bool symmetry = false; 
};

template <typename U, typename V, typename W>
BasixElementData(U element, V bs, W symmetry)
    -> BasixElementData<typename std::remove_cvref<U>::type::scalar_type>;

template <std::floating_point T>
class FiniteElement
{
public:
  using geometry_type = T;

  FiniteElement(const basix::FiniteElement<geometry_type>& element,
                std::size_t gdim,
                const std::optional<std::vector<std::size_t>>& value_shape
                = std::nullopt,
                bool symmetric = false);

  FiniteElement(std::vector<BasixElementData<geometry_type>> elements,
                std::size_t gdim);

  FiniteElement(
      const std::vector<std::shared_ptr<const FiniteElement<geometry_type>>>&
          elements);

  FiniteElement(mesh::CellType cell_type, std::span<const geometry_type> points,
                std::array<std::size_t, 2> pshape,
                std::vector<std::size_t> value_shape = {},
                bool symmetric = false);

  FiniteElement(const FiniteElement& element) = delete;

  FiniteElement(FiniteElement&& element) = default;

  ~FiniteElement() = default;

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

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

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

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

  mesh::CellType cell_type() const noexcept;

  const std::string& signature() const noexcept;

  int space_dimension() const noexcept;

  int block_size() const noexcept;

  int value_size() const;

  std::span<const std::size_t> value_shape() const;

  int physical_base_value_size() const;

  int reference_value_size() const;

  std::span<const std::size_t> reference_value_shape() const;

  const std::vector<std::vector<std::vector<int>>>&
  entity_dofs() const noexcept;

  const std::vector<std::vector<std::vector<int>>>&
  entity_closure_dofs() const noexcept;

  bool symmetric() const;

  void tabulate(std::span<geometry_type> values,
                std::span<const geometry_type> X,
                std::array<std::size_t, 2> shape, int order) const;

  std::pair<std::vector<geometry_type>, std::array<std::size_t, 4>>
  tabulate(std::span<const geometry_type> X, std::array<std::size_t, 2> shape,
           int order) const;

  int num_sub_elements() const noexcept;

  bool is_mixed() const noexcept;

  const std::vector<std::shared_ptr<const FiniteElement<geometry_type>>>&
  sub_elements() const noexcept;

  std::shared_ptr<const FiniteElement<geometry_type>>
  extract_sub_element(const std::vector<int>& component) const;

  const basix::FiniteElement<geometry_type>& basix_element() const;

  basix::maps::type map_type() const;

  bool interpolation_ident() const noexcept;

  bool map_ident() const noexcept;

  std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>>
  interpolation_points() const;

  std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>>
  interpolation_operator() const;

  std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>>
  create_interpolation_operator(const FiniteElement& from) const;

  bool needs_dof_transformations() const noexcept;

  bool depends_on_facet_orientations() const noexcept;

  bool needs_dof_permutations() const noexcept;

  template <typename U>
  std::function<void(std::span<U>, std::span<const std::uint32_t>, std::int32_t,
                     int)>
  dof_transformation_fn(doftransform ttype, bool scalar_element = false) const
  {
    if (!needs_dof_transformations())
      return nullptr;
 
    if (!_sub_elements.empty())
    {
      if (!_reference_value_shape) // Mixed element
      {
        std::vector<std::function<void(
            std::span<U>, std::span<const std::uint32_t>, std::int32_t, int)>>
            sub_element_fns;
        std::vector<int> dims;
        sub_element_fns.reserve(_sub_elements.size());
        dims.reserve(_sub_elements.size());
        for (std::size_t i = 0; i < _sub_elements.size(); ++i)
        {
          sub_element_fns.push_back(
              _sub_elements[i]->template dof_transformation_fn<U>(ttype));
          dims.push_back(_sub_elements[i]->space_dimension());
        }
 
        return [dims = std::move(dims),
                sub_element_fns = std::move(sub_element_fns)](
                   std::span<U> data, std::span<const std::uint32_t> cell_info,
                   std::int32_t cell, int block_size)
        {
          std::size_t offset = 0;
          for (std::size_t e = 0; e < sub_element_fns.size(); ++e)
          {
            const std::size_t width = dims[e] * block_size;
            if (sub_element_fns[e])
            {
              sub_element_fns[e](data.subspan(offset, width), cell_info, cell,
                                 block_size);
            }
            offset += width;
          }
        };
      }
      else if (!scalar_element)
      {
        // Blocked element
        std::function<void(std::span<U>, std::span<const std::uint32_t>,
                           std::int32_t, int)>
            sub_fn
            = _sub_elements.front()->template dof_transformation_fn<U>(ttype);
        assert(sub_fn); // Consistent with needs_dof_transformations() above
        const int ebs = _bs;
        return [ebs, sub_fn = std::move(sub_fn)](
                   std::span<U> data, std::span<const std::uint32_t> cell_info,
                   std::int32_t cell, int data_block_size)
        { sub_fn(data, cell_info, cell, ebs * data_block_size); };
      }
    }
 
    switch (ttype)
    {
    case doftransform::inverse_transpose:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int block_size)
      { Tt_inv_apply(data, cell_info[cell], block_size); };
    case doftransform::transpose:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int block_size)
      { Tt_apply(data, cell_info[cell], block_size); };
    case doftransform::inverse:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int block_size)
      { Tinv_apply(data, cell_info[cell], block_size); };
    case doftransform::standard:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int block_size)
      { T_apply(data, cell_info[cell], block_size); };
    default:
      throw std::invalid_argument("Unknown transformation type");
    }
  }

  template <typename U>
  std::function<void(std::span<U>, std::span<const std::uint32_t>, std::int32_t,
                     int)>
  dof_transformation_right_fn(doftransform ttype,
                              bool scalar_element = false) const
  {
    if (!needs_dof_transformations())
      return nullptr;
    else if (!_sub_elements.empty())
    {
      if (!_reference_value_shape) // Mixed element
      {
        std::vector<std::function<void(
            std::span<U>, std::span<const std::uint32_t>, std::int32_t, int)>>
            sub_element_fns;
        std::vector<int> dims;
        sub_element_fns.reserve(_sub_elements.size());
        dims.reserve(_sub_elements.size());
        for (std::size_t i = 0; i < _sub_elements.size(); ++i)
        {
          sub_element_fns.push_back(
              _sub_elements[i]->template dof_transformation_right_fn<U>(ttype));
          dims.push_back(_sub_elements[i]->space_dimension());
        }
 
        return [dims = std::move(dims),
                sub_element_fns = std::move(sub_element_fns)](
                   std::span<U> data, std::span<const std::uint32_t> cell_info,
                   std::int32_t cell, int block_size)
        {
          // `data` is (block_size, ndofs), row-major. Sub-element `e`
          // owns columns [offset, offset + dims[e]) of every row, so the
          // rows are transformed one at a time.
          const std::size_t ndofs = data.size() / block_size;
          std::size_t offset = 0;
          for (std::size_t e = 0; e < sub_element_fns.size(); ++e)
          {
            if (sub_element_fns[e])
            {
              for (int i = 0; i < block_size; ++i)
              {
                sub_element_fns[e](data.subspan(i * ndofs + offset, dims[e]),
                                   cell_info, cell, 1);
              }
            }
            offset += dims[e];
          }
        };
      }
      else if (!scalar_element)
      {
        // Blocked element
        // The transformation from the left can be used here as blocked
        // elements use xyzxyzxyz ordering, and so applying the DOF
        // transformation from the right is equivalent to applying the DOF
        // transformation from the left to data using xxxyyyzzz ordering
        std::function<void(std::span<U>, std::span<const std::uint32_t>,
                           std::int32_t, int)>
            sub_fn
            = _sub_elements.front()->template dof_transformation_fn<U>(ttype);
        assert(sub_fn); // Consistent with needs_dof_transformations() above
        return [this, sub_fn = std::move(sub_fn)](
                   std::span<U> data, std::span<const std::uint32_t> cell_info,
                   std::int32_t cell, int data_block_size)
        {
          const int ebs = block_size();
          const std::size_t dof_count = data.size() / data_block_size;
          for (int block = 0; block < data_block_size; ++block)
          {
            sub_fn(data.subspan(block * dof_count, dof_count), cell_info, cell,
                   ebs);
          }
        };
      }
    }
 
    switch (ttype)
    {
    case doftransform::inverse_transpose:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int n)
      { Tt_inv_apply_right(data, cell_info[cell], n); };
    case doftransform::transpose:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int n)
      { Tt_apply_right(data, cell_info[cell], n); };
    case doftransform::inverse:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int n)
      { Tinv_apply_right(data, cell_info[cell], n); };
    case doftransform::standard:
      return [this](std::span<U> data, std::span<const std::uint32_t> cell_info,
                    std::int32_t cell, int n)
      { T_apply_right(data, cell_info[cell], n); };
    default:
      throw std::invalid_argument("Unknown transformation type");
    }
  }

  template <typename U>
  void T_apply(std::span<U> data, std::uint32_t cell_permutation, int n) const
  {
    assert(_element);
    _element->T_apply(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<false>(data, cell_permutation, n);
  }

  template <typename U>
  void Tt_inv_apply(std::span<U> data, std::uint32_t cell_permutation,
                    int n) const
  {
    assert(_element);
    _element->Tt_inv_apply(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<false>(data, cell_permutation, n);
  }

  template <typename U>
  void Tt_apply(std::span<U> data, std::uint32_t cell_permutation, int n) const
  {
    assert(_element);
    _element->Tt_apply(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<false>(data, cell_permutation, n);
  }

  template <typename U>
  void Tinv_apply(std::span<U> data, std::uint32_t cell_permutation,
                  int n) const
  {
    assert(_element);
    _element->Tinv_apply(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<false>(data, cell_permutation, n);
  }

  template <typename U>
  void T_apply_right(std::span<U> data, std::uint32_t cell_permutation,
                     int n) const
  {
    assert(_element);
    _element->T_apply_right(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<true>(data, cell_permutation, n);
  }

  template <typename U>
  void Tinv_apply_right(std::span<U> data, std::uint32_t cell_permutation,
                        int n) const
  {
    assert(_element);
    _element->Tinv_apply_right(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<true>(data, cell_permutation, n);
  }

  template <typename U>
  void Tt_apply_right(std::span<U> data, std::uint32_t cell_permutation,
                      int n) const
  {
    assert(_element);
    _element->Tt_apply_right(data, n, cell_permutation & ~_orientation_mask);
    flip_reversed<true>(data, cell_permutation, n);
  }

  template <typename U>
  void Tt_inv_apply_right(std::span<U> data, std::uint32_t cell_permutation,
                          int n) const
  {
    assert(_element);
    _element->Tt_inv_apply_right(data, n,
                                 cell_permutation & ~_orientation_mask);
    flip_reversed<true>(data, cell_permutation, n);
  }

  void permute(std::span<std::int32_t> doflist,
               std::uint32_t cell_permutation) const;

  void permute_inv(std::span<std::int32_t> doflist,
                   std::uint32_t cell_permutation) const;

  std::function<void(std::span<std::int32_t>, std::uint32_t)>
  dof_permutation_fn(bool inverse = false, bool scalar_element = false) const;
 
private:
  // Negate data on each facet marked with mesh::reversed_facet_bit, if
  // the basis depends on the facet signs (see
  // mesh::Topology::create_facet_orientations). A sign is its own
  // inverse and transpose, so every DOF transformation applies it the
  // same way.
  //
  // `post` selects the layout: the right-applied transformations store
  // their data as `(n, num_dofs)`, the others as `(num_dofs, n)`.
  template <bool post, typename U>
  void flip_reversed(std::span<U> data, std::uint32_t cell_permutation,
                     int n) const
  {
    if (!_depends_on_facet_orientations
        or !(cell_permutation & _orientation_mask)) [[likely]]
    {
      return;
    }
 
    // The mask is non-zero only on a cell of dimension 2, whose facets
    // are its edges
    assert(mesh::cell_dim(_cell_type) == 2 and _entity_dofs.size() > 1);
    const std::vector<std::vector<int>>& facet_dofs = _entity_dofs[1];
 
    assert(n > 0 and data.size() % n == 0);
    const std::size_t num_dofs = data.size() / n;
    for (std::size_t f = 0; f < facet_dofs.size(); ++f)
    {
      if (!(cell_permutation & mesh::reversed_facet_bit(f)))
        continue;
      for (int dof : facet_dofs[f])
      {
        for (int j = 0; j < n; ++j)
        {
          U& value = post ? data[j * num_dofs + dof] : data[dof * n + j];
          value = -value;
        }
      }
    }
  }
 
  // Value shape in physical space. For blocked elements this is larger
  // than _reference_value_shape. For non-blocked elements it equals
  // _reference_value_shape, except for a Piola-mapped element on a
  // manifold, where the trailing axes have extent gdim rather than
  // tdim. For mixed elements, it is std::nullopt.
  std::optional<std::vector<std::size_t>> _value_shape;
 
  // Block size for BlockedElements. This gives the number of DOFs
  // co-located at each dof 'point'.
  int _bs;
 
  // Element cell shape
  mesh::CellType _cell_type;
 
  // Element signature
  std::string _signature;
 
  // Dimension of the finite element space (accounting for any blocking)
  int _space_dim;
 
  // List of sub-elements (if any)
  std::vector<std::shared_ptr<const FiniteElement<geometry_type>>>
      _sub_elements;
 
  // Value shape on the reference cell, i.e. the shape Basix tabulates
  // in, e.g. {} for a scalar, {3, 3} for a tensor in 3D. For a mixed
  // element it is std::nullopt, which is the sentinel used by
  // is_mixed() and the dof transformation functions.
  std::optional<std::vector<std::size_t>> _reference_value_shape;
 
  // Basix Element (nullptr for mixed elements)
  std::unique_ptr<basix::FiniteElement<geometry_type>> _element;
 
  // Indicate whether this element represents a symmetric 2-tensor
  bool _symmetric;
 
  // Indicate whether the element needs permutations or transformations
  bool _needs_dof_permutations;
  bool _needs_dof_transformations;
 
  // Indicate whether the basis depends on the facet signs
  bool _depends_on_facet_orientations;
 
  // Bits of the cell permutation info holding facet signs for this
  // element's cell type, cleared before the info is passed to Basix
  std::uint32_t _orientation_mask = 0;
 
  std::vector<std::vector<std::vector<int>>> _entity_dofs;
  std::vector<std::vector<std::vector<int>>> _entity_closure_dofs;
 
  // Quadrature points of a quadrature element (0 dimensional array for
  // all elements except quadrature elements)
  std::pair<std::vector<geometry_type>, std::array<std::size_t, 2>> _points;
};
 
} // namespace dolfinx::fem