assemble_scalar_impl.h

// Copyright (C) 2019-2025 Garth N. Wells and Paul T. Kühner
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include "Constant.h"
#include "Form.h"
#include "FunctionSpace.h"
#include "traits.h"
#include <algorithm>
#include <basix/mdspan.hpp>
#include <concepts>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <memory>
#include <vector>
 
namespace dolfinx::fem::impl
{
template <dolfinx::scalar T, MDSpan2Int32 XD, std::floating_point U>
T assemble_cells_scalar(
    GeometryPack<XD, U> geometry, const IndexList auto& cells,
    const FEkernel<T, U> auto& kernel, std::span<const T> constants,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    ScratchBuffer<U> auto cdofs_b)
{
  T value(0);
  if (std::ranges::empty(cells))
    return value;
 
  const auto x_dofmap = geometry.dofmap;
  const auto ndofs_x = x_dofmap.extent(1);
  assert(cdofs_b.size() == 3 * static_cast<std::size_t>(ndofs_x));
 
  const T* coeffs_data = coeffs.data_handle();
  const auto cstride = coeffs.extent(1);
 
  // Iterate over all cells
  const std::size_t num_cells = std::ranges::size(cells);
  for (std::size_t index = 0; index < num_cells; ++index)
  {
    std::int32_t c = cells[index];
 
    gather_cell_coordinates(geometry, c, cdofs_b.data());
    kernel(&value, coeffs_data + index * cstride, constants.data(),
           cdofs_b.data(), nullptr, nullptr, nullptr);
  }
 
  return value;
}

template <dolfinx::scalar T, MDSpan2Int32 XD, std::floating_point U>
T assemble_entities_scalar(
    GeometryPack<XD, U> geometry,
    md::mdspan<const std::int32_t,
               md::extents<std::size_t, md::dynamic_extent, 2>>
        entities,
    const FEkernel<T, U> auto& kernel, std::span<const T> constants,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms,
    ScratchBuffer<U> auto cdofs_b)
{
  T value(0);
  if (entities.empty())
    return value;
 
  const auto x_dofmap = geometry.dofmap;
  const auto ndofs_x = x_dofmap.extent(1);
  assert(cdofs_b.size() == 3 * static_cast<std::size_t>(ndofs_x));
 
  const T* coeffs_data = coeffs.data_handle();
  const auto cstride = coeffs.extent(1);
 
  // Iterate over all facets
  for (std::size_t f = 0; f < entities.extent(0); ++f)
  {
    std::int32_t cell = entities(f, 0);
    std::int32_t local_entity = entities(f, 1);
 
    gather_cell_coordinates(geometry, cell, cdofs_b.data());
 
    // Permutations
    std::uint8_t perm = perms.empty() ? 0 : perms(cell, local_entity);
    kernel(&value, coeffs_data + f * cstride, constants.data(), cdofs_b.data(),
           &local_entity, &perm, nullptr);
  }
 
  return value;
}

template <dolfinx::scalar T, MDSpan2Int32 XD, std::floating_point U>
T assemble_interior_facets_scalar(
    GeometryPack<XD, U> geometry,
    md::mdspan<const std::int32_t,
               md::extents<std::size_t, md::dynamic_extent, 2, 2>>
        facets,
    const FEkernel<T, U> auto& kernel, std::span<const T> constants,
    md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                    md::dynamic_extent>>
        coeffs,
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms,
    ScratchBuffer<U> auto cdofs_b)
{
  T value(0);
  if (facets.empty())
    return value;
 
  // Create data structures used in assembly
  const auto x_dofmap = geometry.dofmap;
  const auto ndofs_x = x_dofmap.extent(1);
  assert(cdofs_b.size() == 2 * 3 * static_cast<std::size_t>(ndofs_x));
  U* cdofs0 = cdofs_b.data();
  U* cdofs1 = cdofs_b.data() + 3 * ndofs_x;
 
  const T* coeffs_data = coeffs.data_handle();
  const auto cstride = 2 * coeffs.extent(2);
 
  // Iterate over all facets
  for (std::size_t f = 0; f < facets.extent(0); ++f)
  {
    std::array cells = {facets(f, 0, 0), facets(f, 1, 0)};
    std::array local_facet = {facets(f, 0, 1), facets(f, 1, 1)};
 
    gather_cell_coordinates(geometry, cells[0], cdofs0);
    gather_cell_coordinates(geometry, cells[1], cdofs1);
 
    std::array perm = perms.empty()
                          ? std::array<std::uint8_t, 2>{0, 0}
                          : std::array{perms(cells[0], local_facet[0]),
                                       perms(cells[1], local_facet[1])};
    kernel(&value, coeffs_data + f * cstride, constants.data(), cdofs_b.data(),
           local_facet.data(), perm.data(), nullptr);
  }
 
  return value;
}

template <dolfinx::scalar T, std::floating_point U>
T assemble_scalar(
    const fem::Form<T, U>& M, mdspan2_t x_dofmap,
    md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>> x,
    std::span<const T> constants,
    const std::map<std::pair<IntegralType, int>,
                   std::pair<std::span<const T>, int>>& coefficients,
    std::size_t cell_type_idx)
{
  std::shared_ptr<const mesh::Mesh<U>> mesh = M.mesh();
  assert(mesh);
 
  // Sized for the worst case (interior facets, which touch two cells).
  // The kernels require an exactly-sized buffer, so the one-cell
  // integrals get the leading half.
  std::vector<U> cdofs_b(2 * 3 * x_dofmap.extent(1));
  std::span cdofs_b1 = std::span(cdofs_b).first(3 * x_dofmap.extent(1));
  GeometryPack geometry{x_dofmap, x};
 
  T value = 0;
  for (int i = 0; i < M.num_integrals(IntegralType::cell, cell_type_idx); ++i)
  {
    auto fn = M.kernel(IntegralType::cell, i, cell_type_idx);
    assert(fn);
    auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
    std::span<const std::int32_t> cells
        = M.domain(IntegralType::cell, i, cell_type_idx);
    assert(cells.size() * cstride == coeffs.size());
    value += impl::assemble_cells_scalar(
        geometry, cells, fn, constants,
        md::mdspan(coeffs.data(), cells.size(), cstride), cdofs_b1);
  }
 
  md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
  if (M.needs_facet_permutations())
  {
    facet_perms = impl::entity_permutations(
        *mesh->topology_mutable(), IntegralType::interior_facet,
        mesh->topology()->cell_types()[cell_type_idx]);
  }
 
  for (int i = 0;
       i < M.num_integrals(IntegralType::interior_facet, cell_type_idx); ++i)
  {
    auto fn = M.kernel(IntegralType::interior_facet, i, cell_type_idx);
    assert(fn);
    auto& [coeffs, cstride]
        = coefficients.at({IntegralType::interior_facet, i});
    std::span facets = M.domain(IntegralType::interior_facet, i, cell_type_idx);
 
    constexpr std::size_t num_adjacent_cells = 2;
    // Two values per each adj. cell (cell index and local facet index).
    constexpr std::size_t shape1 = 2 * num_adjacent_cells;
 
    assert((facets.size() / shape1) * 2 * cstride == coeffs.size());
    value += impl::assemble_interior_facets_scalar(
        geometry,
        md::mdspan<const std::int32_t,
                   md::extents<std::size_t, md::dynamic_extent, 2, 2>>(
            facets.data(), facets.size() / shape1, 2, 2),
        fn, constants,
        md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                        md::dynamic_extent>>(
            coeffs.data(), facets.size() / shape1, 2, cstride),
        facet_perms, std::span(cdofs_b));
  }
 
  for (auto itg_type : {fem::IntegralType::exterior_facet,
                        fem::IntegralType::vertex, fem::IntegralType::ridge})
  {
    const int num_itg = M.num_integrals(itg_type, cell_type_idx);
    if (num_itg == 0)
      continue;
 
    // Each integral type is over entities of a different
    // codimension, so only the permutations this form actually
    // integrates over are computed.
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> perms;
    if (M.needs_facet_permutations())
    {
      perms = impl::entity_permutations(
          *mesh->topology_mutable(), itg_type,
          mesh->topology()->cell_types()[cell_type_idx]);
    }
 
    for (int i = 0; i < num_itg; ++i)
    {
      auto fn = M.kernel(itg_type, i, cell_type_idx);
      assert(fn);
      auto& [coeffs, cstride] = coefficients.at({itg_type, i});
 
      std::span entities = M.domain(itg_type, i, cell_type_idx);
 
      // Two values per each adj. cell (cell index and local entity index).
      assert((entities.size() / 2) * cstride == coeffs.size());
      value += impl::assemble_entities_scalar(
          geometry,
          md::mdspan<const std::int32_t,
                     md::extents<std::size_t, md::dynamic_extent, 2>>(
              entities.data(), entities.size() / 2, 2),
          fn, constants,
          md::mdspan(coeffs.data(), entities.size() / 2, cstride), perms,
          cdofs_b1);
    }
  }
 
  return value;
}
 
} // namespace dolfinx::fem::impl