Form.h

// Copyright (C) 2019-2023 Garth N. Wells and Chris Richardson
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "FunctionSpace.h"
#include "traits.h"
#include <algorithm>
#include <array>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/Mesh.h>
#include <functional>
#include <memory>
#include <span>
#include <string>
#include <tuple>
#include <vector>
 
namespace dolfinx::fem
{
 
template <dolfinx::scalar T>
class Constant;
template <dolfinx::scalar T, std::floating_point U>
class Function;
 
enum class IntegralType : std::int8_t
{
  cell = 0,           
  exterior_facet = 1, 
  interior_facet = 2, 
  vertex = 3          
};
 
template <dolfinx::scalar T, std::floating_point U = scalar_value_type_t<T>>
struct integral_data
{
  template <typename K, typename V, typename W>
    requires std::is_convertible_v<
                 std::remove_cvref_t<K>,
                 std::function<void(T*, const T*, const T*, const U*,
                                    const int*, const uint8_t*)>>
                 and std::is_convertible_v<std::remove_cvref_t<V>,
                                           std::vector<std::int32_t>>
                 and std::is_convertible_v<std::remove_cvref_t<W>,
                                           std::vector<int>>
  integral_data(int id, K&& kernel, V&& entities, W&& coeffs)
      : id(id), kernel(std::forward<K>(kernel)),
        entities(std::forward<V>(entities)), coeffs(std::forward<W>(coeffs))
  {
  }
 
  template <typename K, typename W>
    requires std::is_convertible_v<
                 std::remove_cvref_t<K>,
                 std::function<void(T*, const T*, const T*, const U*,
                                    const int*, const uint8_t*)>>
                 and std::is_convertible_v<std::remove_cvref_t<W>,
                                           std::vector<int>>
  integral_data(int id, K&& kernel, std::span<const std::int32_t> entities,
                W&& coeffs)
      : id(id), kernel(std::forward<K>(kernel)),
        entities(entities.begin(), entities.end()),
        coeffs(std::forward<W>(coeffs))
  {
  }
 
  int id;
 
  std::function<void(T*, const T*, const T*, const U*, const int*,
                     const uint8_t*)>
      kernel;
 
  std::vector<std::int32_t> entities;
 
  std::vector<int> coeffs;
};
 
template <dolfinx::scalar T,
          std::floating_point U = dolfinx::scalar_value_type_t<T>>
class Form
{
public:
  using scalar_type = T;
 
  using geometry_type = U;
 
  template <typename X>
    requires std::is_convertible_v<
                 std::remove_cvref_t<X>,
                 std::map<IntegralType, std::vector<integral_data<
                                            scalar_type, geometry_type>>>>
  Form(
      const std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>& V,
      X&& integrals,
      const std::vector<
          std::shared_ptr<const Function<scalar_type, geometry_type>>>&
          coefficients,
      const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
          constants,
      bool needs_facet_permutations,
      const std::map<std::shared_ptr<const mesh::Mesh<geometry_type>>,
                     std::span<const std::int32_t>>& entity_maps,
      std::shared_ptr<const mesh::Mesh<geometry_type>> mesh = nullptr)
      : _function_spaces(V), _coefficients(coefficients), _constants(constants),
        _mesh(mesh), _needs_facet_permutations(needs_facet_permutations)
  {
    // Extract _mesh from FunctionSpace, and check they are the same
    if (!_mesh and !V.empty())
      _mesh = V[0]->mesh();
    for (auto& space : V)
    {
      if (_mesh != space->mesh()
          and entity_maps.find(space->mesh()) == entity_maps.end())
      {
        throw std::runtime_error(
            "Incompatible mesh. entity_maps must be provided.");
      }
    }
    if (!_mesh)
      throw std::runtime_error("No mesh could be associated with the Form.");
 
    // Store kernels, looping over integrals by domain type (dimension)
    for (auto&& [domain_type, data] : integrals)
    {
      if (!std::ranges::is_sorted(data,
                                  [](auto& a, auto& b) { return a.id < b.id; }))
      {
        throw std::runtime_error("Integral IDs not sorted");
      }
 
      std::vector<integral_data<scalar_type, geometry_type>>& itg
          = _integrals[static_cast<std::size_t>(domain_type)];
      for (auto&& [id, kern, e, c] : data)
        itg.emplace_back(id, kern, std::move(e), std::move(c));
    }
 
    // Store entity maps
    for (auto [msh, map] : entity_maps)
      _entity_maps.insert({msh, std::vector(map.begin(), map.end())});
  }
 
  Form(const Form& form) = delete;
 
  Form(Form&& form) = default;
 
  virtual ~Form() = default;
 
  int rank() const { return _function_spaces.size(); }
 
  std::shared_ptr<const mesh::Mesh<geometry_type>> mesh() const
  {
    return _mesh;
  }
 
  const std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>&
  function_spaces() const
  {
    return _function_spaces;
  }
 
  std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
                     const geometry_type*, const int*, const uint8_t*)>
  kernel(IntegralType type, int i) const
  {
    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
    auto it = std::ranges::lower_bound(integrals, i, std::less<>{},
                                       [](const auto& a) { return a.id; });
    if (it != integrals.end() and it->id == i)
      return it->kernel;
    else
      throw std::runtime_error("No kernel for requested domain index.");
  }
 
  std::set<IntegralType> integral_types() const
  {
    std::set<IntegralType> set;
    for (std::size_t i = 0; i < _integrals.size(); ++i)
    {
      if (!_integrals[i].empty())
        set.insert(static_cast<IntegralType>(i));
    }
 
    return set;
  }
 
  int num_integrals(IntegralType type) const
  {
    return _integrals[static_cast<std::size_t>(type)].size();
  }
 
  std::vector<int> active_coeffs(IntegralType type, std::size_t i) const
  {
    assert(i < _integrals[static_cast<std::size_t>(type)].size());
    return _integrals[static_cast<std::size_t>(type)][i].coeffs;
  }
 
  std::vector<int> integral_ids(IntegralType type) const
  {
    std::vector<int> ids;
    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
    std::ranges::transform(integrals, std::back_inserter(ids),
                           [](auto& integral) { return integral.id; });
    return ids;
  }
 
  std::span<const std::int32_t> domain(IntegralType type, int i) const
  {
    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
    auto it = std::ranges::lower_bound(integrals, i, std::less<>{},
                                       [](const auto& a) { return a.id; });
    if (it != integrals.end() and it->id == i)
      return it->entities;
    else
      throw std::runtime_error("No mesh entities for requested domain index.");
  }
 
  std::vector<std::int32_t> domain(IntegralType type, int i,
                                   const mesh::Mesh<geometry_type>& mesh) const
  {
    // Hack to avoid passing shared pointer to this function
    std::shared_ptr<const mesh::Mesh<geometry_type>> msh_ptr(
        &mesh, [](const mesh::Mesh<geometry_type>*) {});
 
    std::span<const std::int32_t> entities = domain(type, i);
    if (msh_ptr == _mesh)
      return std::vector(entities.begin(), entities.end());
    else
    {
      std::span<const std::int32_t> entity_map = _entity_maps.at(msh_ptr);
      std::vector<std::int32_t> mapped_entities;
      mapped_entities.reserve(entities.size());
      switch (type)
      {
      case IntegralType::cell:
      {
        std::ranges::transform(entities, std::back_inserter(mapped_entities),
                               [&entity_map](auto e) { return entity_map[e]; });
        break;
      }
      case IntegralType::exterior_facet:
      {
        // Get the codimension of the mesh
        const int tdim = _mesh->topology()->dim();
        const int codim = tdim - mesh.topology()->dim();
        assert(codim >= 0);
        if (codim == 0)
        {
          for (std::size_t i = 0; i < entities.size(); i += 2)
          {
            // Add cell and the local facet index
            mapped_entities.insert(mapped_entities.end(),
                                   {entity_map[entities[i]], entities[i + 1]});
          }
        }
        else if (codim == 1)
        {
          // In this case, the entity maps take facets in (`_mesh`) to cells in
          // `mesh`, so we need to get the facet number from the (cell,
          // local_facet pair) first.
          auto c_to_f = _mesh->topology()->connectivity(tdim, tdim - 1);
          assert(c_to_f);
          for (std::size_t i = 0; i < entities.size(); i += 2)
          {
            // Get the facet index
            const std::int32_t facet
                = c_to_f->links(entities[i])[entities[i + 1]];
            // Add cell and the local facet index
            mapped_entities.insert(mapped_entities.end(),
                                   {entity_map[facet], entities[i + 1]});
          }
        }
        else
          throw std::runtime_error("Codimension > 1 not supported.");
 
        break;
      }
      case IntegralType::interior_facet:
      {
        // Get the codimension of the mesh
        const int tdim = _mesh->topology()->dim();
        const int codim = tdim - mesh.topology()->dim();
        assert(codim >= 0);
        if (codim == 0)
        {
          for (std::size_t i = 0; i < entities.size(); i += 2)
          {
            // Add cell and the local facet index
            mapped_entities.insert(mapped_entities.end(),
                                   {entity_map[entities[i]], entities[i + 1]});
          }
        }
        else if (codim == 1)
        {
          // In this case, the entity maps take facets in (`_mesh`) to cells in
          // `mesh`, so we need to get the facet number from the (cell,
          // local_facet pair) first.
          auto c_to_f = _mesh->topology()->connectivity(tdim, tdim - 1);
          assert(c_to_f);
          for (std::size_t i = 0; i < entities.size(); i += 2)
          {
            // Get the facet index
            const std::int32_t facet
                = c_to_f->links(entities[i])[entities[i + 1]];
            // Add cell and the local facet index
            mapped_entities.insert(mapped_entities.end(),
                                   {entity_map[facet], entities[i + 1]});
          }
        }
        break;
      }
      default:
        throw std::runtime_error("Integral type not supported.");
      }
 
      return mapped_entities;
    }
  }
 
  const std::vector<
      std::shared_ptr<const Function<scalar_type, geometry_type>>>&
  coefficients() const
  {
    return _coefficients;
  }
 
  bool needs_facet_permutations() const { return _needs_facet_permutations; }
 
  std::vector<int> coefficient_offsets() const
  {
    std::vector<int> n = {0};
    for (auto& c : _coefficients)
    {
      if (!c)
        throw std::runtime_error("Not all form coefficients have been set.");
      n.push_back(n.back() + c->function_space()->element()->space_dimension());
    }
    return n;
  }
 
  const std::vector<std::shared_ptr<const Constant<scalar_type>>>&
  constants() const
  {
    return _constants;
  }
 
private:
  // Function spaces (one for each argument)
  std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>
      _function_spaces;
 
  // Integrals. Array index is
  // static_cast<std::size_t(IntegralType::foo)
  std::array<std::vector<integral_data<scalar_type, geometry_type>>, 4>
      _integrals;
 
  // Form coefficients
  std::vector<std::shared_ptr<const Function<scalar_type, geometry_type>>>
      _coefficients;
 
  // Constants associated with the Form
  std::vector<std::shared_ptr<const Constant<scalar_type>>> _constants;
 
  // The mesh
  std::shared_ptr<const mesh::Mesh<geometry_type>> _mesh;
 
  // True if permutation data needs to be passed into these integrals
  bool _needs_facet_permutations;
 
  // Entity maps (see Form documentation)
  std::map<std::shared_ptr<const mesh::Mesh<geometry_type>>,
           std::vector<std::int32_t>>
      _entity_maps;
}; // namespace dolfinx::fem
} // namespace dolfinx::fem