FiniteElement.h

// Copyright (C) 2020-2024 Garth N. Wells and Matthew W. Scroggs
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "traits.h"
#include <array>
#include <basix/finite-element.h>
#include <concepts>
#include <cstdint>
#include <dolfinx/mesh/cell_types.h>
#include <functional>
#include <memory>
#include <optional>
#include <span>
#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>
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,
                const std::optional<std::vector<std::size_t>>& value_shape
                = std::nullopt,
                bool symmetric = false);

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

  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;

  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 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 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())
    {
      // If no permutation needed, return function that does nothing
      return [](std::span<U>, std::span<const std::uint32_t>, std::int32_t, int)
      {
        // Do nothing
      };
    }
 
    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;
        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, 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;
            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);
        const int ebs = _bs;
        return [ebs, 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::runtime_error("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())
    {
      // If no permutation needed, return function that does nothing
      return [](std::span<U>, std::span<const std::uint32_t>, std::int32_t, int)
      {
        // Do nothing
      };
    }
    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;
        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));
        }
 
        return [this, 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)
          {
            sub_element_fns[e](data.subspan(offset, data.size() - offset),
                               cell_info, cell, block_size);
            offset += _sub_elements[e]->space_dimension();
          }
        };
      }
      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);
        return [this, 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::runtime_error("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);
  }

  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);
  }

  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);
  }

  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);
  }

  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);
  }

  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);
  }

  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);
  }

  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);
  }

  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:
  // Value shape. For blocked elements this is larger than
  // _reference_value_shape. For non-blocked 'primal' elements it is
  // equal to _reference_value_shape. 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 space shape, e.g. {} for a scalar, {3, 3} for a tensor in 3D.
  // For a mixed element it is std::nullopt.
  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;
 
  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