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 <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>
    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>>
  integral_data(int id, K&& kernel, V&& entities)
      : id(id), kernel(std::forward<K>(kernel)),
        entities(std::forward<V>(entities))
  {
  }
 
  template <typename K>
    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*)>>
  integral_data(int id, K&& kernel, std::span<const std::int32_t> e)
      : id(id), kernel(std::forward<K>(kernel)), entities(e.begin(), e.end())
  {
  }
 
  int id;
 
  std::function<void(T*, const T*, const T*, const U*, const int*,
                     const uint8_t*)>
      kernel;
 
  std::vector<std::int32_t> entities;
};
 
template <dolfinx::scalar T,
          std::floating_point U = dolfinx::scalar_value_type_t<T>>
class Form
{
public:
  using scalar_type = T;
 
  template <typename X>
    requires std::is_convertible_v<
                 std::remove_cvref_t<X>,
                 std::map<IntegralType, std::vector<integral_data<T, U>>>>
  Form(const std::vector<std::shared_ptr<const FunctionSpace<U>>>& V,
       X&& integrals,
       const std::vector<std::shared_ptr<const Function<scalar_type, U>>>&
           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<U>>,
                      std::span<const std::int32_t>>& entity_maps,
       std::shared_ptr<const mesh::Mesh<U>> 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::is_sorted(data.begin(), data.end(),
                          [](auto& a, auto& b) { return a.id < b.id; }))
      {
        throw std::runtime_error("Integral IDs not sorted");
      }
 
      std::vector<integral_data<T, U>>& itg
          = _integrals[static_cast<std::size_t>(domain_type)];
      for (auto&& [id, kern, e] : data)
        itg.emplace_back(id, kern, std::move(e));
    }
 
    // 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<U>> mesh() const { return _mesh; }
 
  const std::vector<std::shared_ptr<const FunctionSpace<U>>>&
  function_spaces() const
  {
    return _function_spaces;
  }
 
  std::function<void(T*, const T*, const T*, const U*, const int*,
                     const uint8_t*)>
  kernel(IntegralType type, int i) const
  {
    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
    auto it = std::lower_bound(integrals.begin(), integrals.end(), i,
                               [](auto& itg_data, int i)
                               { return itg_data.id < i; });
    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> integral_ids(IntegralType type) const
  {
    std::vector<int> ids;
    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
    std::transform(integrals.begin(), integrals.end(), 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::lower_bound(integrals.begin(), integrals.end(), i,
                               [](auto& itg_data, int i)
                               { return itg_data.id < i; });
    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<U>& mesh) const
  {
    // Hack to avoid passing shared pointer to this function
    std::shared_ptr<const mesh::Mesh<U>> msh_ptr(&mesh,
                                                 [](const mesh::Mesh<U>*) {});
 
    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::transform(entities.begin(), entities.end(),
                       std::back_inserter(mapped_entities),
                       [&entity_map](auto e) { return entity_map[e]; });
        break;
      }
      case IntegralType::exterior_facet:
        // Exterior and interior facets are treated the same
        [[fallthrough]];
      case IntegralType::interior_facet:
      {
        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]});
        }
        break;
      }
      default:
        throw std::runtime_error("Integral type not supported.");
      }
 
      return mapped_entities;
    }
  }
 
  const std::vector<std::shared_ptr<const Function<T, U>>>& 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<T>>>& constants() const
  {
    return _constants;
  }
 
private:
  // Function spaces (one for each argument)
  std::vector<std::shared_ptr<const FunctionSpace<U>>> _function_spaces;
 
  // Integrals. Array index is
  // static_cast<std::size_t(IntegralType::foo)
  std::array<std::vector<integral_data<T, U>>, 4> _integrals;
 
  // Form coefficients
  std::vector<std::shared_ptr<const Function<T, U>>> _coefficients;
 
  // Constants associated with the Form
  std::vector<std::shared_ptr<const Constant<T>>> _constants;
 
  // The mesh
  std::shared_ptr<const mesh::Mesh<U>> _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<U>>, std::vector<std::int32_t>>
      _entity_maps;
}; // namespace dolfinx::fem
} // namespace dolfinx::fem