FiniteElement.h

// Copyright (C) 2008-2020 Anders Logg and Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <basix/finite-element.h>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/cell_types.h>
#include <functional>
#include <memory>
#include <numeric>
#include <vector>
#include <xtl/xspan.hpp>
 
struct ufc_finite_element;
 
namespace dolfinx::fem
{
class FiniteElement
{
public:
  explicit FiniteElement(const ufc_finite_element& ufc_element);
 
  FiniteElement(const FiniteElement& element) = delete;
 
  FiniteElement(FiniteElement&& element) = default;
 
  virtual ~FiniteElement() = default;
 
  FiniteElement& operator=(const FiniteElement& element) = delete;
 
  FiniteElement& operator=(FiniteElement&& element) = default;
 
  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 noexcept;
 
  int reference_value_size() const noexcept;
 
  int value_rank() const noexcept;
 
  int value_dimension(int i) const;
 
  std::string family() const noexcept;
 
  void tabulate(xt::xtensor<double, 4>& values, const xt::xtensor<double, 2>& X,
                int order) const;
 
  void transform_reference_basis(xt::xtensor<double, 3>& values,
                                 const xt::xtensor<double, 3>& reference_values,
                                 const xt::xtensor<double, 3>& J,
                                 const xtl::span<const double>& detJ,
                                 const xt::xtensor<double, 3>& K) const;
 
  int num_sub_elements() const noexcept;
 
  const std::vector<std::shared_ptr<const FiniteElement>>&
  sub_elements() const noexcept;
 
  std::size_t hash() const noexcept;
 
  std::shared_ptr<const FiniteElement>
  extract_sub_element(const std::vector<int>& component) const;
 
  bool interpolation_ident() const noexcept;
 
  const xt::xtensor<double, 2>& interpolation_points() const;
 
  template <typename T>
  constexpr void interpolate(const xt::xtensor<T, 2>& values,
                             xtl::span<T> dofs) const
  {
    if (!_element)
    {
      throw std::runtime_error("No underlying element for interpolation. "
                               "Cannot interpolate mixed elements directly.");
    }
 
    _element->interpolate(tcb::make_span(dofs), tcb::make_span(values), _bs);
  }
 
  bool needs_dof_transformations() const noexcept;
 
  bool needs_dof_permutations() const noexcept;
 
  template <typename T>
  std::function<void(const xtl::span<T>&, const xtl::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 xtl::span<T>&, const xtl::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 xtl::span<T>&,
                                       const xtl::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]->get_dof_transformation_function<T>(inverse,
                                                                   transpose));
          dims.push_back(_sub_elements[i]->space_dimension());
        }
 
        return [dims, sub_element_functions](
                   const xtl::span<T>& data,
                   const xtl::span<const std::uint32_t>& cell_info,
                   std::int32_t cell, int block_size)
        {
          std::size_t start = 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(start, width), cell_info,
                                     cell, block_size);
            start += width;
          }
        };
      }
      else if (!scalar_element)
      {
        // Vector element
        const std::function<void(const xtl::span<T>&,
                                 const xtl::span<const std::uint32_t>&,
                                 std::int32_t, int)>
            sub_function = _sub_elements[0]->get_dof_transformation_function<T>(
                inverse, transpose);
        const int ebs = _bs;
        return
            [ebs, sub_function](const xtl::span<T>& data,
                                const xtl::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 xtl::span<T>& data,
                      const xtl::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 xtl::span<T>& data,
                      const xtl::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 xtl::span<T>& data,
                      const xtl::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 xtl::span<T>& data,
                      const xtl::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 T>
  std::function<void(const xtl::span<T>&, const xtl::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 xtl::span<T>&, const xtl::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 xtl::span<T>&,
                                       const xtl::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]->get_dof_transformation_to_transpose_function<T>(
                  inverse, transpose));
        }
 
        return [this, sub_element_functions](
                   const xtl::span<T>& data,
                   const xtl::span<const std::uint32_t>& cell_info,
                   std::int32_t cell, int block_size)
        {
          std::size_t start = 0;
          for (std::size_t e = 0; e < sub_element_functions.size(); ++e)
          {
            const std::size_t width
                = _sub_elements[e]->space_dimension() * block_size;
            sub_element_functions[e](data.subspan(start, width), cell_info,
                                     cell, block_size);
            start += width;
          }
        };
      }
      else if (!scalar_element)
      {
        // Vector element
        const std::function<void(const xtl::span<T>&,
                                 const xtl::span<const std::uint32_t>&,
                                 std::int32_t, int)>
            sub_function = _sub_elements[0]->get_dof_transformation_function<T>(
                inverse, transpose);
        return [this,
                sub_function](const xtl::span<T>& data,
                              const xtl::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 xtl::span<T>& data,
                      const xtl::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 xtl::span<T>& data,
                      const xtl::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 xtl::span<T>& data,
                      const xtl::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 xtl::span<T>& data,
                      const xtl::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 T>
  void apply_dof_transformation(const xtl::span<T>& data,
                                std::uint32_t cell_permutation,
                                int block_size) const
  {
    assert(_element);
    _element->apply_dof_transformation(data, block_size, cell_permutation);
  }
 
  template <typename T>
  void
  apply_inverse_transpose_dof_transformation(const xtl::span<T>& 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 T>
  void apply_transpose_dof_transformation(const xtl::span<T>& data,
                                          std::uint32_t cell_permutation,
                                          int block_size) const
  {
    assert(_element);
    _element->apply_transpose_dof_transformation(data, block_size,
                                                 cell_permutation);
  }
 
  template <typename T>
  void apply_inverse_dof_transformation(const xtl::span<T>& data,
                                        std::uint32_t cell_permutation,
                                        int block_size) const
  {
    assert(_element);
    _element->apply_inverse_dof_transformation(data, block_size,
                                               cell_permutation);
  }
 
  template <typename T>
  void apply_dof_transformation_to_transpose(const xtl::span<T>& 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 T>
  void
  apply_inverse_dof_transformation_to_transpose(const xtl::span<T>& 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 T>
  void apply_transpose_dof_transformation_to_transpose(
      const xtl::span<T>& 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 T>
  void apply_inverse_transpose_dof_transformation_to_transpose(
      const xtl::span<T>& 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);
  }
 
  template <typename O, typename P, typename Q, typename T, typename S>
  void map_pull_back(const O& u, const P& J, const Q& detJ, const T& K,
                     S&& U) const
  {
    assert(_element);
    _element->map_pull_back_m(u, J, detJ, K, U);
  }
 
  void permute_dofs(const xtl::span<std::int32_t>& doflist,
                    std::uint32_t cell_permutation) const;
 
  void unpermute_dofs(const xtl::span<std::int32_t>& doflist,
                      std::uint32_t cell_permutation) const;
 
  std::function<void(const xtl::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 _tdim, _space_dim, _value_size, _reference_value_size;
 
  // List of sub-elements (if any)
  std::vector<std::shared_ptr<const FiniteElement>> _sub_elements;
 
  // Simple hash of the signature string
  std::size_t _hash;
 
  // Dimension of each value space
  std::vector<int> _value_dimension;
 
  // Block size for VectorElements and TensorElements. This gives the
  // number of DOFs colocated at each 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> _element;
};
} // namespace dolfinx::fem