Form.h

// Copyright (C) 2019-2025 Garth N. Wells, Chris Richardson, Joseph P. Dean and
// Jørgen S. Dokken
//
// 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 <basix/mdspan.hpp>
#include <cassert>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/EntityMap.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/cell_types.h>
#include <format>
#include <functional>
#include <map>
#include <memory>
#include <ranges>
#include <set>
#include <span>
#include <stdexcept>
#include <tuple>
#include <utility>
#include <variant>
#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,         
  ridge = 4           
};

constexpr int integral_entity_dim(IntegralType type, int tdim)
{
  switch (type)
  {
  case IntegralType::exterior_facet:
  case IntegralType::interior_facet:
    return tdim - 1;
  case IntegralType::ridge:
    return tdim - 2;
  case IntegralType::vertex:
    return 0;
  case IntegralType::cell:
    return tdim;
  }
 
  throw std::invalid_argument("Unknown integral type.");
}
 
namespace impl
{
inline md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>>
entity_permutations(mesh::Topology& topology, IntegralType type,
                    mesh::CellType cell_type)
{
  if (type == IntegralType::cell)
    return {};
 
  const int tdim = topology.dim();
  const int dim = integral_entity_dim(type, tdim);
 
  topology.create_entity_permutations(dim);
  const std::vector<std::uint8_t>& p = topology.get_entity_permutations(dim);
  const int num_entities_per_cell = mesh::cell_num_entities(cell_type, dim);
  return md::mdspan(p.data(), p.size() / num_entities_per_cell,
                    num_entities_per_cell);
}
} // namespace impl

template <dolfinx::scalar T, std::floating_point U = scalar_value_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*, void*)>>
                 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(K&& kernel, V&& entities, W&& coeffs)
      : kernel(std::forward<K>(kernel)), entities(std::forward<V>(entities)),
        coeffs(std::forward<W>(coeffs))
  {
  }

  std::function<void(T*, const T*, const T*, const U*, const int*,
                     const uint8_t*, void*)>
      kernel;

  std::vector<std::int32_t> entities;

  std::vector<int> coeffs;
};

template <dolfinx::scalar T, std::floating_point U = dolfinx::scalar_value_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<std::tuple<IntegralType, int, int>,
                          integral_data<scalar_type, geometry_type>>>
  Form(
      const std::vector<std::shared_ptr<const FunctionSpace<geometry_type>>>& V,
      X&& integrals, std::shared_ptr<const mesh::Mesh<geometry_type>> mesh,
      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::vector<std::reference_wrapper<const mesh::EntityMap>>&
          entity_maps)
      : _function_spaces(V), _integrals(std::forward<X>(integrals)),
        _mesh(mesh), _coefficients(coefficients), _constants(constants),
        _needs_facet_permutations(needs_facet_permutations)
  {
    if (!_mesh)
      throw std::invalid_argument("Form Mesh is null.");
 
    // `_mesh` is fixed for the remainder of construction, so its
    // topology and dimension are fetched once and reused below rather
    // than being re-fetched for every integral/coefficient.
    const mesh::Topology& topology = *_mesh->topology();
    const int tdim = topology.dim();
 
    // A helper function to find the correct entity map for a given mesh
    auto get_entity_map
        = [mesh, &entity_maps](auto& mesh0) -> const mesh::EntityMap&
    {
      auto it = std::ranges::find_if(
          entity_maps,
          [mesh, mesh0](const mesh::EntityMap& em)
          {
            return ((em.topology() == mesh0->topology()
                     and em.sub_topology() == mesh->topology()))
                   or ((em.sub_topology() == mesh0->topology()
                        and em.topology() == mesh->topology()));
          });
 
      if (it == entity_maps.end())
      {
        throw std::invalid_argument(
            "Incompatible mesh. argument entity_maps must be provided.");
      }
      return *it;
    };
 
    // A helper function to compute the (cell, local_entity) pairs in the
    // argument/coefficient domain from the (cell, local_entity) pairs in
    // `this->mesh()`.
    auto compute_entity_domains
        = [](const auto& int_ents_mesh, int codim, const auto& c_to_e,
             const auto& emap, bool inverse)
    {
      // TODO: This function would be much neater using
      // `std::views::stride(2)` from C++ 23
 
      // Get a list of entities to map to the argument/coefficient
      // domain
      std::vector<std::int32_t> entities;
      entities.reserve(int_ents_mesh.size() / 2);
      if (codim == 0)
      {
        // In the codim 0 case, we need to map from cells in
        // `this->mesh()` to cells in the argument/coefficient mesh, so
        // here we extract the cells.
        for (std::size_t i = 0; i < int_ents_mesh.size(); i += 2)
          entities.push_back(int_ents_mesh[i]);
      }
      else
      {
        // Otherwise the integration entities are sub-entities of the
        // cells of `this->mesh()` and are themselves cells of the
        // argument/coefficient mesh, so here we extract the entity
        // index using the cell-to-entity connectivity.
        for (std::size_t i = 0; i < int_ents_mesh.size(); i += 2)
        {
          entities.push_back(
              c_to_e->links(int_ents_mesh[i])[int_ents_mesh[i + 1]]);
        }
      }
 
      // Map from entity indices in `this->mesh()` to the corresponding
      // cell indices in the argument/coefficient mesh
      std::vector<std::int32_t> cells_mesh0
          = emap.sub_topology_to_topology(entities, inverse);
 
      // Create a list of (cell, local_entity_index) pairs in the
      // argument/coefficient domain. Only the cell column is meaningful.
      // For codim > 0 the entity is itself the cell, so it has
      // no local index. The colummn is never used later on, but written
      // like this for consistency for packing/assembly.
      std::vector<std::int32_t> e = int_ents_mesh;
      for (std::size_t i = 0; i < cells_mesh0.size(); ++i)
        e[2 * i] = cells_mesh0[i];
 
      return e;
    };
 
    // Map the integration entities of one integral to the
    // argument/coefficient domain, checking first that the mapping is
    // expressible: `compute_entity_domains` maps an integration entity
    // to a *cell* of that mesh, so unless the meshes have equal
    // dimension the integral's entity dimension must be that mesh's.
    auto map_entities = [tdim, &topology, &compute_entity_domains](
                            IntegralType type, const auto& entities,
                            const mesh::Mesh<geometry_type>& mesh0,
                            const mesh::EntityMap& emap,
                            bool inverse) -> std::vector<std::int32_t>
    {
      if (type == IntegralType::cell)
        return emap.sub_topology_to_topology(entities, inverse);
 
      if (type == IntegralType::vertex)
      {
        throw std::invalid_argument(
            "Vertex integrals are not supported for a form with an argument "
            "or coefficient on another mesh. Supported types are cell, "
            "exterior facet, interior facet and ridge.");
      }
 
      const int dim0 = mesh0.topology()->dim();
      const int codim = tdim - dim0;
      const int edim = integral_entity_dim(type, tdim);
      assert(codim >= 0);
      if (codim > 0 and edim != dim0)
      {
        throw std::invalid_argument(std::format(
            "Cannot map integration entities of dimension {} to cells of a "
            "mesh of dimension {}. An argument or coefficient on another mesh "
            "must live on the entities being integrated over.",
            edim, dim0));
      }
 
      std::shared_ptr<const graph::AdjacencyList<std::int32_t>> c_to_e
          = topology.connectivity(tdim, edim);
      assert(codim == 0 or c_to_e);
      return compute_entity_domains(entities, codim, c_to_e, emap, inverse);
    };
 
    _edata.reserve(_function_spaces.size());
    for (auto& space : _function_spaces)
    {
      // Working map: [integral type, integral_idx, kernel_idx]->entities
      std::map<std::tuple<IntegralType, int, int>,
               std::variant<std::vector<std::int32_t>,
                            std::span<const std::int32_t>>>
          vdata;
 
      if (auto mesh0 = space->mesh(); mesh0 == _mesh)
      {
        for (auto& [key, integral] : _integrals)
          vdata.insert({key, std::span(integral.entities)});
      }
      else
      {
        // Find correct entity map
        const mesh::EntityMap& emap = get_entity_map(mesh0);
 
        // Determine direction of the map. We need to map from
        // `this->mesh()` to `mesh0`, so if `emap->sub_topology()` isn't
        // the source topology, we need the inverse map
        bool inverse = emap.sub_topology() == mesh0->topology();
        for (auto& [key, itg] : _integrals)
        {
          auto [type, idx, kernel_idx] = key;
          std::vector<std::int32_t> e;
          assert(mesh0);
          e = map_entities(type, itg.entities, *mesh0, emap, inverse);
 
          vdata.insert({key, std::move(e)});
        }
      }
 
      _edata.push_back(std::move(vdata));
    }
 
    for (auto& [key, integral] : _integrals)
    {
      auto [type, idx, kernel_idx] = key;
      for (int c : integral.coeffs)
      {
        if (auto mesh0 = coefficients.at(c)->function_space()->mesh();
            mesh0 == _mesh)
        {
          _cdata.insert({{type, idx, c}, std::span(integral.entities)});
        }
        else
        {
          // Find correct entity map and determine direction of the map
          const mesh::EntityMap& emap = get_entity_map(mesh0);
          bool inverse = emap.sub_topology() == mesh0->topology();
 
          std::vector<std::int32_t> e;
          assert(mesh0);
          e = map_entities(type, integral.entities, *mesh0, emap, inverse);
          _cdata.insert({{type, idx, c}, std::move(e)});
        }
      }
    }
  }
 
  // Copy constructor (deleted). _edata and _cdata cache std::spans
  // aliasing the entity vectors owned by _integrals; a shallow copy
  // would leave the copy's spans pointing into the original's data.
  Form(const Form& form) = delete;

#ifdef _MSC_VER
  Form(Form&& form) noexcept = default;
#else
  Form(Form&& form) = default;
#endif

  ~Form() = default;
 
  // Copy assignment (deleted). Same aliasing reason as the copy
  // constructor.
  Form& operator=(const Form& form) = delete;

  Form& operator=(Form&& 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*, void*)>
  kernel(IntegralType type, int idx, int kernel_idx) const
  {
    auto it = _integrals.find({type, idx, kernel_idx});
    if (it == _integrals.end())
      throw std::out_of_range("Requested integral kernel not found.");
    return it->second.kernel;
  }

  std::set<IntegralType> integral_types() const
  {
    std::set<IntegralType> types;
    for (auto& [key, integral] : _integrals)
      types.insert(std::get<0>(key));
    return types;
  }

  std::vector<int> active_coeffs(IntegralType type, int idx) const
  {
    auto it = std::ranges::find_if(_integrals,
                                   [type, idx](auto& x)
                                   {
                                     auto [t, idx_, kernel_idx] = x.first;
                                     return t == type and idx_ == idx;
                                   });
    if (it == _integrals.end())
      throw std::out_of_range("Could not find active coefficient list.");
    return it->second.coeffs;
  }

  int num_integrals(IntegralType type, int kernel_idx) const
  {
    return std::ranges::count_if(_integrals,
                                 [type, kernel_idx](auto& x)
                                 {
                                   auto [t, id, k_idx] = x.first;
                                   return t == type and k_idx == kernel_idx;
                                 });
  }

  std::span<const std::int32_t> domain(IntegralType type, int idx,
                                       int kernel_idx) const
  {
    auto it = _integrals.find({type, idx, kernel_idx});
    if (it == _integrals.end())
      throw std::out_of_range("Requested domain not found.");
    return it->second.entities;
  }

  std::span<const std::int32_t> domain_arg(IntegralType type, int rank, int idx,
                                           int kernel_idx) const
  {
    auto it = _edata.at(rank).find({type, idx, kernel_idx});
    if (it == _edata.at(rank).end())
      throw std::out_of_range("Requested domain for argument not found.");
 
    return std::visit([](const auto& v) -> std::span<const std::int32_t>
                      { return v; }, it->second);
  }

  std::span<const std::int32_t> domain_coeff(IntegralType type, int idx,
                                             int c) const
  {
    auto it = _cdata.find({type, idx, c});
    if (it == _cdata.end())
      throw std::out_of_range("No domain for requested integral.");
    return std::visit([](const auto& v) -> std::span<const std::int32_t>
                      { return v; }, it->second);
  }

  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};
    n.reserve(_coefficients.size() + 1);
    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 (integral type, idx, kernel_idx)
  std::map<std::tuple<IntegralType, int, int>,
           integral_data<scalar_type, geometry_type>>
      _integrals;
 
  // The mesh
  std::shared_ptr<const mesh::Mesh<geometry_type>> _mesh;
 
  // 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;
 
  // True if permutation data needs to be passed into these integrals
  bool _needs_facet_permutations;
 
  // Mapped domain index data for argument functions.
  //
  // Consider:
  //
  // entities  = this->domain(type, idx, kernel_idx);
  // entities0 = _edata[0][{type, idx, kernel_idx}];
  //
  // Then `entities[i]` is a mesh entity index (e.g., cell index) in
  // `_mesh`, and  `entities0[i]` is the index of the same entity but in
  // the mesh associated with the argument 0 (test function) space.
  std::vector<std::map<
      std::tuple<IntegralType, int, int>,
      std::variant<std::vector<std::int32_t>, std::span<const std::int32_t>>>>
      _edata;
 
  // Mapped domain index data for coefficient functions.
  //
  // Consider:
  //
  // entities  = this->domain(type, idx, kernel_idx);
  // entities0 = _cdata[{type, idx, c}];
  //
  // where `c` is the coefficient index.
  //
  // Then `entities[i]` is a mesh entity index (e.g., cell index) in
  // `_mesh`, and  `entities0[i]` is the index of the same entity but in
  // the mesh associated with the coefficient Function.
  std::map<
      std::tuple<IntegralType, int, int>,
      std::variant<std::vector<std::int32_t>, std::span<const std::int32_t>>>
      _cdata;
};
} // namespace dolfinx::fem