FiniteElement.h

// Copyright (C) 2020-2021 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 <span>
#include <utility>
#include <vector>
 
struct ufcx_finite_element;
 
namespace dolfinx::fem
{
enum class doftransform
{
  standard = 0,          
  transpose = 1,         
  inverse = 2,           
  inverse_transpose = 3, 
};
 
template <std::floating_point T>
class FiniteElement
{
public:
  using geometry_type = T;
 
  explicit FiniteElement(const ufcx_finite_element& e);
 
  FiniteElement(const basix::FiniteElement<geometry_type>& element,
                const std::size_t block_size, const 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;
 
  std::string signature() const noexcept;
 
  mesh::CellType cell_shape() const noexcept;
 
  int space_dimension() const noexcept;
 
  int block_size() const noexcept;
 
  int reference_value_size() const;
 
  std::span<const std::size_t> reference_value_shape() const;
 
  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 (_is_mixed)
      {
        // Mixed element
        std::vector<std::function<void(
            std::span<U>, std::span<const std::uint32_t>, std::int32_t, int)>>
            sub_element_functions;
        std::vector<int> dims;
        for (std::size_t i = 0; i < _sub_elements.size(); ++i)
        {
          sub_element_functions.push_back(
              _sub_elements[i]->template dof_transformation_fn<U>(ttype));
          dims.push_back(_sub_elements[i]->space_dimension());
        }
 
        return [dims, sub_element_functions](
                   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_functions.size(); ++e)
          {
            const std::size_t width = dims[e] * block_size;
            sub_element_functions[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_function
            = _sub_elements[0]->template dof_transformation_fn<U>(ttype);
        const int ebs = _bs;
        return [ebs, sub_function](std::span<U> data,
                                   std::span<const std::uint32_t> cell_info,
                                   std::int32_t cell, int data_block_size)
        { sub_function(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.size() != 0)
    {
      if (_is_mixed)
      {
        // Mixed element
        std::vector<std::function<void(
            std::span<U>, std::span<const std::uint32_t>, std::int32_t, int)>>
            sub_element_functions;
        for (std::size_t i = 0; i < _sub_elements.size(); ++i)
        {
          sub_element_functions.push_back(
              _sub_elements[i]->template dof_transformation_right_fn<U>(ttype));
        }
 
        return [this, sub_element_functions](
                   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_functions.size(); ++e)
          {
            sub_element_functions[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_function
            = _sub_elements[0]->template dof_transformation_fn<U>(ttype);
        return [this, sub_function](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_function(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:
  std::string _signature;
 
  mesh::CellType _cell_shape;
 
  int _space_dim;
 
  // List of sub-elements (if any)
  std::vector<std::shared_ptr<const FiniteElement<geometry_type>>>
      _sub_elements;
 
  // Dimension of each value space
  std::vector<std::size_t> _reference_value_shape;
 
  // Block size for BlockedElements. This gives the number of DOFs
  // co-located at each dof 'point'.
  int _bs;
 
  // Indicate whether this is a mixed element
  bool _is_mixed;
 
  // 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;
 
  // Basix Element (nullptr for mixed elements)
  std::unique_ptr<basix::FiniteElement<geometry_type>> _element;
 
  // 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