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 <array>
#include <basix/finite-element.h>
#include <concepts>
#include <dolfinx/mesh/cell_types.h>
#include <functional>
#include <memory>
#include <span>
#include <utility>
#include <vector>
 
struct ufcx_finite_element;
 
namespace dolfinx::fem
{
template <std::floating_point T>
class FiniteElement
{
public:
  explicit FiniteElement(const ufcx_finite_element& e);
 
  FiniteElement(const basix::FiniteElement<T>& element,
                const std::vector<std::size_t>& value_shape);
 
  FiniteElement(const FiniteElement& element) = delete;
 
  FiniteElement(FiniteElement&& element) = default;
 
  virtual ~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 value_size() const;
 
  int reference_value_size() const;
 
  std::span<const std::size_t> value_shape() const noexcept;
 
  std::string family() const noexcept;
 
  void tabulate(std::span<T> values, std::span<const T> X,
                std::array<std::size_t, 2> shape, int order) const;
 
  std::pair<std::vector<T>, std::array<std::size_t, 4>>
  tabulate(std::span<const T> 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<T>>>&
  sub_elements() const noexcept;
 
  std::shared_ptr<const FiniteElement<T>>
  extract_sub_element(const std::vector<int>& component) const;
 
  const basix::FiniteElement<T>& basix_element() const;
 
  basix::maps::type map_type() const;
 
  bool interpolation_ident() const noexcept;
 
  bool map_ident() const noexcept;
 
  std::pair<std::vector<T>, std::array<std::size_t, 2>>
  interpolation_points() const;
 
  std::pair<std::vector<T>, std::array<std::size_t, 2>>
  interpolation_operator() const;
 
  std::pair<std::vector<T>, 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(const std::span<U>&, const std::span<const std::uint32_t>&,
                     std::int32_t, int)>
  get_dof_transformation_function(bool inverse = false, bool transpose = false,
                                  bool scalar_element = false) const
  {
    if (!needs_dof_transformations())
    {
      // If no permutation needed, return function that does nothing
      return [](const std::span<U>&, const std::span<const std::uint32_t>&,
                std::int32_t, int)
      {
        // Do nothing
      };
    }
 
    if (_sub_elements.size() != 0)
    {
      if (_bs == 1)
      {
        // Mixed element
        std::vector<std::function<void(const std::span<U>&,
                                       const 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 get_dof_transformation_function<U>(
                  inverse, transpose));
          dims.push_back(_sub_elements[i]->space_dimension());
        }
 
        return [dims, sub_element_functions](
                   const std::span<U>& data,
                   const 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)
      {
        // Vector element
        const std::function<void(const std::span<U>&,
                                 const std::span<const std::uint32_t>&,
                                 std::int32_t, int)>
            sub_function
            = _sub_elements[0]->template get_dof_transformation_function<U>(
                inverse, transpose);
        const int ebs = _bs;
        return
            [ebs, sub_function](const std::span<U>& data,
                                const 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); };
      }
    }
    if (transpose)
    {
      if (inverse)
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size)
        {
          apply_inverse_transpose_dof_transformation(data, cell_info[cell],
                                                     block_size);
        };
      }
      else
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size) {
          apply_transpose_dof_transformation(data, cell_info[cell], block_size);
        };
      }
    }
    else
    {
      if (inverse)
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size) {
          apply_inverse_dof_transformation(data, cell_info[cell], block_size);
        };
      }
      else
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size)
        { apply_dof_transformation(data, cell_info[cell], block_size); };
      }
    }
  }
 
  template <typename U>
  std::function<void(const std::span<U>&, const std::span<const std::uint32_t>&,
                     std::int32_t, int)>
  get_dof_transformation_to_transpose_function(bool inverse = false,
                                               bool transpose = false,
                                               bool scalar_element
                                               = false) const
  {
    if (!needs_dof_transformations())
    {
      // If no permutation needed, return function that does nothing
      return [](const std::span<U>&, const std::span<const std::uint32_t>&,
                std::int32_t, int)
      {
        // Do nothing
      };
    }
    else if (_sub_elements.size() != 0)
    {
      if (_bs == 1)
      {
        // Mixed element
        std::vector<std::function<void(const std::span<U>&,
                                       const 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 get_dof_transformation_to_transpose_function<U>(
                      inverse, transpose));
        }
 
        return [this, sub_element_functions](
                   const std::span<U>& data,
                   const 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)
      {
        // Vector element
        const std::function<void(const std::span<U>&,
                                 const std::span<const std::uint32_t>&,
                                 std::int32_t, int)>
            sub_function
            = _sub_elements[0]->template get_dof_transformation_function<U>(
                inverse, transpose);
        return [this,
                sub_function](const std::span<U>& data,
                              const 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);
          }
        };
      }
    }
 
    if (transpose)
    {
      if (inverse)
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size)
        {
          apply_inverse_transpose_dof_transformation_to_transpose(
              data, cell_info[cell], block_size);
        };
      }
      else
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size)
        {
          apply_transpose_dof_transformation_to_transpose(data, cell_info[cell],
                                                          block_size);
        };
      }
    }
    else
    {
      if (inverse)
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size)
        {
          apply_inverse_dof_transformation_to_transpose(data, cell_info[cell],
                                                        block_size);
        };
      }
      else
      {
        return [this](const std::span<U>& data,
                      const std::span<const std::uint32_t>& cell_info,
                      std::int32_t cell, int block_size) {
          apply_dof_transformation_to_transpose(data, cell_info[cell],
                                                block_size);
        };
      }
    }
  }
 
  template <typename U>
  void apply_dof_transformation(const std::span<U>& data,
                                std::uint32_t cell_permutation,
                                int block_size) const
  {
    assert(_element);
    _element->apply_dof_transformation(data, block_size, cell_permutation);
  }
 
  template <typename U>
  void
  apply_inverse_transpose_dof_transformation(const std::span<U>& data,
                                             std::uint32_t cell_permutation,
                                             int block_size) const
  {
    assert(_element);
    _element->apply_inverse_transpose_dof_transformation(data, block_size,
                                                         cell_permutation);
  }
 
  template <typename U>
  void apply_transpose_dof_transformation(const std::span<U>& data,
                                          std::uint32_t cell_permutation,
                                          int block_size) const
  {
    assert(_element);
    _element->apply_transpose_dof_transformation(data, block_size,
                                                 cell_permutation);
  }
 
  template <typename U>
  void apply_inverse_dof_transformation(const std::span<U>& data,
                                        std::uint32_t cell_permutation,
                                        int block_size) const
  {
    assert(_element);
    _element->apply_inverse_dof_transformation(data, block_size,
                                               cell_permutation);
  }
 
  template <typename U>
  void apply_dof_transformation_to_transpose(const std::span<U>& data,
                                             std::uint32_t cell_permutation,
                                             int block_size) const
  {
    assert(_element);
    _element->apply_dof_transformation_to_transpose(data, block_size,
                                                    cell_permutation);
  }
 
  template <typename U>
  void
  apply_inverse_dof_transformation_to_transpose(const std::span<U>& data,
                                                std::uint32_t cell_permutation,
                                                int block_size) const
  {
    assert(_element);
    _element->apply_inverse_dof_transformation_to_transpose(data, block_size,
                                                            cell_permutation);
  }
 
  template <typename U>
  void apply_transpose_dof_transformation_to_transpose(
      const std::span<U>& data, std::uint32_t cell_permutation,
      int block_size) const
  {
    assert(_element);
    _element->apply_transpose_dof_transformation_to_transpose(data, block_size,
                                                              cell_permutation);
  }
 
  template <typename U>
  void apply_inverse_transpose_dof_transformation_to_transpose(
      const std::span<U>& data, std::uint32_t cell_permutation,
      int block_size) const
  {
    assert(_element);
    _element->apply_inverse_transpose_dof_transformation_to_transpose(
        data, block_size, cell_permutation);
  }
 
  void permute_dofs(const std::span<std::int32_t>& doflist,
                    std::uint32_t cell_permutation) const;
 
  void unpermute_dofs(const std::span<std::int32_t>& doflist,
                      std::uint32_t cell_permutation) const;
 
  std::function<void(const std::span<std::int32_t>&, std::uint32_t)>
  get_dof_permutation_function(bool inverse = false,
                               bool scalar_element = false) const;
 
private:
  std::string _signature, _family;
 
  mesh::CellType _cell_shape;
 
  int _space_dim;
 
  // List of sub-elements (if any)
  std::vector<std::shared_ptr<const FiniteElement<T>>> _sub_elements;
 
  // Dimension of each value space
  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;
 
  // 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<T>> _element;
};
} // namespace dolfinx::fem