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 <concepts>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/EntityMap.h>
#include <dolfinx/mesh/Mesh.h>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <ranges>
#include <span>
#include <tuple>
#include <utility>
#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           
};

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::runtime_error("Form Mesh is null.");
 
    // 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::runtime_error(
            "Incompatible mesh. argument entity_maps must be provided.");
      }
      return *it;
    };
 
    // A helper function to compute the (cell, local_facet) pairs in the
    // argument/coefficient domain from the (cell, local_facet) pairs in
    // `this->mesh()`.
    auto compute_facet_domains
        = [&](const auto& int_ents_mesh, int codim, const auto& c_to_f,
              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 if (codim == 1)
      {
        // In the codim 1 case, we need to map facets in `this->mesh()`
        // to cells in the argument/coefficient mesh, so here we extract
        // the facet index using the cell-to-facet connectivity.
        for (std::size_t i = 0; i < int_ents_mesh.size(); i += 2)
        {
          entities.push_back(
              c_to_f->links(int_ents_mesh[i])[int_ents_mesh[i + 1]]);
        }
      }
      else
        throw std::runtime_error("Codimension > 1 not supported.");
 
      // 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_facet_index) pairs in the
      // argument/coefficient domain. Since `create_submesh`preserves
      // the local facet index (with respect to the cell), we can use
      // the local facet indices from the input integration entities
      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;
    };
 
    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;
          if (type == IntegralType::cell)
            e = emap.sub_topology_to_topology(itg.entities, inverse);
          else if (type == IntegralType::exterior_facet
                   or type == IntegralType::interior_facet)
          {
            const mesh::Topology topology = *_mesh->topology();
            int tdim = topology.dim();
            assert(mesh0);
            int codim = tdim - mesh0->topology()->dim();
            assert(codim >= 0);
            auto c_to_f = topology.connectivity(tdim, tdim - 1);
            assert(c_to_f);
 
            e = compute_facet_domains(itg.entities, codim, c_to_f, emap,
                                      inverse);
          }
          else
            throw std::runtime_error("Integral type not supported.");
 
          vdata.insert({key, std::move(e)});
        }
      }
 
      _edata.push_back(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;
          if (type == IntegralType::cell)
            e = emap.sub_topology_to_topology(integral.entities, inverse);
          else if (type == IntegralType::exterior_facet
                   or type == IntegralType::interior_facet)
          {
            const mesh::Topology topology = *_mesh->topology();
            int tdim = topology.dim();
            assert(mesh0);
            int codim = tdim - mesh0->topology()->dim();
            auto c_to_f = topology.connectivity(tdim, tdim - 1);
            assert(c_to_f);
 
            e = compute_facet_domains(integral.entities, codim, c_to_f, emap,
                                      inverse);
          }
          else
            throw std::runtime_error("Integral type not supported.");
          _cdata.insert({{type, idx, c}, std::move(e)});
        }
      }
    }
  }

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

  std::set<IntegralType> integral_types() const
  {
    std::vector<IntegralType> set_data;
    std::ranges::transform(_integrals, std::back_inserter(set_data),
                           [](auto& x) { return std::get<0>(x.first); });
    return std::set<IntegralType>(set_data.begin(), set_data.end());
  }

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

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

  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::runtime_error("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::runtime_error("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::runtime_error("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};
    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, id, celltype)
  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(IntegralType, integral(id), kernel_idx];
  // entities0 = _edata[0][IntegralType, integral(id), 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 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(IntegralType, integral(id), kernel_idx];
  // entities0 = _cdata[IntegralType, integral(id), 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