Form.h

// Copyright (C) 2019-2025 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 <basix/mdspan.hpp>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/types.h>
#include <dolfinx/mesh/Mesh.h>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#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          
};
 
namespace impl
{
std::vector<std::int32_t> compute_domain(
    auto entities, std::span<const std::int32_t> entity_map,
    std::optional<int> codim = std::nullopt,
    std::optional<
        std::reference_wrapper<const graph::AdjacencyList<std::int32_t>>>
        c_to_f
    = std::nullopt)
{
  static_assert(entities.rank() == 1 or entities.rank() == 2);
  std::vector<std::int32_t> mapped_entities;
  mapped_entities.reserve(entities.size());
  if constexpr (entities.rank() == 1)
  {
    std::span ents(entities.data_handle(), entities.size());
    std::ranges::transform(ents, std::back_inserter(mapped_entities),
                           [&entity_map](auto e) { return entity_map[e]; });
  }
  else if (entities.rank() == 2)
  {
    assert(codim.value() >= 0);
    if (codim.value() == 0)
    {
      for (std::size_t i = 0; i < entities.extent(0); ++i)
      {
        // Add cell and the local facet index
        mapped_entities.insert(mapped_entities.end(),
                               {entity_map[entities(i, 0)], entities(i, 1)});
      }
    }
    else if (codim.value() == 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.
      for (std::size_t i = 0; i < entities.extent(0); ++i)
      {
        // Get the facet index, and add cell and the local facet index
        std::int32_t facet
            = c_to_f->get().links(entities(i, 0))[entities(i, 1)];
        mapped_entities.insert(mapped_entities.end(),
                               {entity_map[facet], entities(i, 1)});
      }
    }
    else
      throw std::runtime_error("Codimension > 1 not supported.");
  }
  else
    throw std::runtime_error("Integral type not supported.");
 
  return mapped_entities;
}
} // 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::map<std::shared_ptr<const mesh::Mesh<geometry_type>>,
                     std::span<const std::int32_t>>& 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.");
 
    // Check consistency of mesh(es)
    {
      // Integration domain mesh is passed, so check that it is (1)
      // common for spaces and coefficients (2) or an entity_map is
      // available
      for (auto& space : _function_spaces)
      {
        if (auto mesh0 = space->mesh();
            mesh0 != _mesh and !entity_maps.contains(mesh0))
        {
          throw std::runtime_error(
              "Incompatible mesh. argument entity_maps must be provided.");
        }
      }
      for (auto& c : coefficients)
      {
        if (auto mesh0 = c->function_space()->mesh();
            mesh0 != _mesh and !entity_maps.contains(mesh0))
        {
          throw std::runtime_error(
              "Incompatible mesh. coefficient entity_maps must be provided.");
        }
      }
    }
 
    for (auto& space : _function_spaces)
    {
      // Working map: [integral type, domain ID, 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
      {
        auto it = entity_maps.find(mesh0);
        assert(it != entity_maps.end());
        std::span<const std::int32_t> entity_map = it->second;
        for (auto& [key, itg] : _integrals)
        {
          auto [type, id, kernel_idx] = key;
          std::vector<std::int32_t> e;
          if (type == IntegralType::cell)
          {
            e = impl::compute_domain(
                md::mdspan(itg.entities.data(), itg.entities.size()),
                entity_map);
          }
          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 = impl::compute_domain(
                md::mdspan<const std::int32_t,
                           md::extents<std::size_t, md::dynamic_extent, 2>>(
                    itg.entities.data(), itg.entities.size() / 2, 2),
                entity_map, codim, *c_to_f);
          }
          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, id, kernel_idx] = key;
      for (int c : integral.coeffs)
      {
        if (auto mesh0 = coefficients.at(c)->function_space()->mesh();
            mesh0 == _mesh)
        {
          _cdata.insert({{type, id, c}, std::span(integral.entities)});
        }
        else
        {
          auto it = entity_maps.find(mesh0);
          assert(it != entity_maps.end());
          std::span<const std::int32_t> entity_map = it->second;
          std::vector<std::int32_t> e;
          if (type == IntegralType::cell)
          {
            e = impl::compute_domain(
                md::mdspan(integral.entities.data(), integral.entities.size()),
                entity_map);
          }
          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 = impl::compute_domain(
                md::mdspan<const std::int32_t,
                           md::extents<std::size_t, md::dynamic_extent, 2>>(
                    integral.entities.data(), integral.entities.size() / 2, 2),
                entity_map, codim, *c_to_f);
          }
          else
            throw std::runtime_error("Integral type not supported.");
          _cdata.insert({{type, id, 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::set<IntegralType> set;
    std::ranges::for_each(_integrals, [&set](auto& x)
                          { set.insert(std::get<0>(x.first)); });
    return set;
  }

  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;
  }

  std::vector<int> integral_ids(IntegralType type) const
  {
    std::vector<int> ids;
    for (auto& [key, integral] : _integrals)
    {
      auto [t, id, kernel_idx] = key;
      if (t == type)
        ids.push_back(id);
    }
 
    // IDs may be repeated in mixed-topology meshes, so remove
    // duplicates
    std::sort(ids.begin(), ids.end());
    auto it = std::unique(ids.begin(), ids.end());
    ids.erase(it, ids.end());
    return ids;
  }

  std::span<const std::int32_t> domain(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 domain not found.");
    return it->second.entities;
  }

  std::span<const std::int32_t> domain_arg(IntegralType type, int rank, int id,
                                           int kernel_idx) const
  {
    auto it = _edata.at(rank).find({type, id, kernel_idx});
    if (it == _edata.at(rank).end())
      throw std::runtime_error("Requested domain for argument not found.");
    try
    {
      return std::get<std::span<const std::int32_t>>(it->second);
    }
    catch (std::bad_variant_access& e)
    {
      return std::get<std::vector<std::int32_t>>(it->second);
    }
  }

  std::span<const std::int32_t> domain_coeff(IntegralType type, int id,
                                             int c) const
  {
    auto it = _cdata.find({type, id, c});
    if (it == _cdata.end())
      throw std::runtime_error("No domain for requested integral.");
    try
    {
      return std::get<std::span<const std::int32_t>>(it->second);
    }
    catch (std::bad_variant_access& e)
    {
      return std::get<std::vector<std::int32_t>>(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