assemble_vector_impl.h

// Copyright (C) 2018-2026 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 "DirichletBC.h"
#include "DofMap.h"
#include "Form.h"
#include "assemble_matrix_impl.h"
#include "traits.h"
#include <algorithm>
#include <array>
#include <basix/mdspan.hpp>
#include <concepts>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/mesh/Geometry.h>
#include <dolfinx/mesh/Mesh.h>
#include <dolfinx/mesh/Topology.h>
#include <functional>
#include <memory>
#include <optional>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <vector>
 
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class DirichletBC;
}
 
namespace dolfinx::fem::impl
{
using mdspan2_t = md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>;

template <typename V, MDSpan2Int32 XD, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires AssemblyVector<V, T>
void assemble_cells_vector(
    V&& b, GeometryPack<XD, U> geometry, const IndexList auto& cells,
    const FormArgumentCells<T> auto& arg0, const FEkernel<T, U> auto& kernel,
    std::span<const T> constants,
    md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
    ScratchBuffer<T> auto be_b, ScratchBuffer<U> auto cdofs_b)
{
  if (std::ranges::empty(cells))
    return;
 
  // By value: the sizes below fold only if not read through a
  // reference (see fem::DofMapPack). mdspan and span are two-word
  // copies.
  const auto dmap = arg0.dofmap.map;
  const auto bs = arg0.dofmap.bs;
  // By reference: a generated range (e.g. iota) does not convert to a
  // span, and a caller's std::vector must not be copied.
  const auto& cells0 = arg0.dofmap.entities;
  const auto& P0 = arg0.transform;
  std::span<const std::uint32_t> cell_info0 = arg0.cell_info;
 
  const auto ndofs_x = geometry.dofmap.extent(1);
  const auto ndofs = dmap.extent(1);
  assert(be_b.size() == static_cast<std::size_t>(bs) * ndofs);
  assert(cdofs_b.size() == 3 * static_cast<std::size_t>(ndofs_x));
 
  const std::int32_t* dmap_ptr = dmap.data_handle();
  const T* coeffs_data = coeffs.data_handle();
  const auto cstride = coeffs.extent(1);
 
  // P0 does not change across cells in this call, so whether it is a
  // set (non-null) transform is loop-invariant.
  const bool p0_set = is_transform_set(P0);
 
  const std::size_t num_cells = std::ranges::size(cells);
  assert(std::ranges::size(cells0) == num_cells);
 
  // The integration-domain and argument cell lists are usually the same
  // span, making the second lookup redundant. Loop-invariant, so the
  // branch predicts perfectly. Both concepts require a sized range.
  bool same_cells = false;
  if constexpr (std::ranges::contiguous_range<decltype(cells)>
                and std::ranges::contiguous_range<decltype(cells0)>)
  {
    same_cells = std::ranges::data(cells0) == std::ranges::data(cells)
                 and std::ranges::size(cells0) == num_cells;
  }
 
  // Iterate over active cells
  for (std::size_t index = 0; index < num_cells; ++index)
  {
    // Integration domain cell and test function cell
    const std::int32_t c = cells[index];
    const std::int32_t c0 = same_cells ? c : cells0[index];
 
    gather_cell_coordinates(geometry, c, cdofs_b.data());
 
    // Tabulate vector for cell
    std::ranges::fill(be_b, T(0));
    kernel(be_b.data(), coeffs_data + index * cstride, constants.data(),
           cdofs_b.data(), nullptr, nullptr, nullptr);
    if (p0_set)
      P0(std::span<T>(be_b), cell_info0, c0, 1);
 
    // Scatter cell vector to 'global' vector array
    const std::int32_t* dofs
        = dmap_ptr + static_cast<std::ptrdiff_t>(c0) * ndofs;
    for (std::size_t i = 0; i < ndofs; ++i)
    {
      const std::int32_t dof = bs * dofs[i];
      for (int k = 0; k < bs; ++k)
        b[dof + k] += be_b[bs * i + k];
    }
  }
}

template <typename V, MDSpan2Int32 XD, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires AssemblyVector<V, T>
void assemble_entities_vector(
    V&& b, GeometryPack<XD, U> geometry,
    md::mdspan<const std::int32_t,
               std::extents<std::size_t, md::dynamic_extent, 2>>
        entities,
    const FormArgumentEntities<T> auto& arg0, 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<T> auto be_b, ScratchBuffer<U> auto cdofs_b)
{
  if (entities.empty())
    return;
 
  // By value: the sizes below fold only if not read through a
  // reference (see fem::DofMapPack). mdspan and span are two-word
  // copies.
  const auto dmap = arg0.dofmap.map;
  const auto bs = arg0.dofmap.bs;
  const auto entities0 = arg0.dofmap.entities;
  const auto& P0 = arg0.transform;
  std::span<const std::uint32_t> cell_info0 = arg0.cell_info;
 
  const auto num_dofs = dmap.extent(1);
  const auto num_x_dofs_cell = geometry.dofmap.extent(1);
  assert(cdofs_b.size() == 3 * static_cast<std::size_t>(num_x_dofs_cell));
  assert(be_b.size() == static_cast<std::size_t>(bs) * num_dofs);
  assert(entities0.size() == entities.size());
 
  const std::int32_t* dmap_ptr = dmap.data_handle();
 
  const bool p0_set = is_transform_set(P0);
 
  const T* coeffs_data = coeffs.data_handle();
  const std::size_t cstride = coeffs.extent(1);
 
  for (std::size_t f = 0; f < entities.extent(0); ++f)
  {
    // Cell in the integration domain, local facet index relative to the
    // integration domain cell, and cell in the test function mesh
    std::int32_t cell = entities(f, 0);
    std::int32_t local_entity = entities(f, 1);
    std::int32_t cell0 = entities0(f, 0);
 
    gather_cell_coordinates(geometry, cell, cdofs_b.data());
 
    // Permutations
    std::uint8_t perm = perms.empty() ? 0 : perms(cell, local_entity);
 
    // Tabulate element vector
    std::ranges::fill(be_b, T(0));
    kernel(be_b.data(), coeffs_data + f * cstride, constants.data(),
           cdofs_b.data(), &local_entity, &perm, nullptr);
    if (p0_set)
      P0(std::span<T>(be_b), cell_info0, cell0, 1);
 
    // Add to global vector
    const std::int32_t* dofs
        = dmap_ptr + static_cast<std::ptrdiff_t>(cell0) * num_dofs;
    for (std::size_t i = 0; i < num_dofs; ++i)
      for (int k = 0; k < bs; ++k)
        b[bs * dofs[i] + k] += be_b[bs * i + k];
  }
}

template <typename V, MDSpan2Int32 XD, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires AssemblyVector<V, T>
void assemble_interior_facets_vector(
    V&& b, GeometryPack<XD, U> geometry,
    md::mdspan<const std::int32_t,
               std::extents<std::size_t, md::dynamic_extent, 2, 2>>
        facets,
    const FormArgumentFacets<T> auto& arg0, 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<T> auto be_b, ScratchBuffer<U> auto cdofs_b)
{
  if (facets.empty())
    return;
 
  // By value: the sizes below fold only if not read through a
  // reference (see fem::DofMapPack). mdspan and span are two-word
  // copies.
  const auto dmap = arg0.dofmap.map;
  const auto bs = arg0.dofmap.bs;
  const auto facets0 = arg0.dofmap.entities;
  const auto& P0 = arg0.transform;
  std::span<const std::uint32_t> cell_info0 = arg0.cell_info;
 
  const auto num_x_dofs_cell = geometry.dofmap.extent(1);
  const auto dmap_size = dmap.extent(1);
  assert(cdofs_b.size() == 2 * static_cast<std::size_t>(num_x_dofs_cell) * 3);
  assert(be_b.size() == static_cast<std::size_t>(bs) * 2 * dmap_size);
  assert(facets0.size() == facets.size());
  U* cdofs0 = cdofs_b.data();
  U* cdofs1 = cdofs_b.data() + num_x_dofs_cell * 3;
 
  const T* coeffs_data = coeffs.data_handle();
  const auto cstride = 2 * coeffs.extent(2);
 
  const bool p0_set = is_transform_set(P0);
 
  for (std::size_t f = 0; f < facets.extent(0); ++f)
  {
    // Cells in integration domain and test function domain meshes
    std::array<std::int32_t, 2> cells{facets(f, 0, 0), facets(f, 1, 0)};
    std::array<std::int32_t, 2> cells0{facets0(f, 0, 0), facets0(f, 1, 0)};
 
    // Local facet indices
    std::array<std::int32_t, 2> local_facet{facets(f, 0, 1), facets(f, 1, 1)};
 
    gather_cell_coordinates(geometry, cells[0], cdofs0);
    gather_cell_coordinates(geometry, cells[1], cdofs1);
 
    // Get dofmaps for cells. When integrating over interfaces between
    // two domains, the test function might only be defined on one side,
    // so we check which cells exist in the test function domain.
    std::span dmap0 = cells0[0] >= 0 ? std::span(&dmap(cells0[0], 0), dmap_size)
                                     : std::span<const std::int32_t>();
    std::span dmap1 = cells0[1] >= 0 ? std::span(&dmap(cells0[1], 0), dmap_size)
                                     : std::span<const std::int32_t>();
 
    // Tabulate element vector
    std::ranges::fill(be_b, T(0));
    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(be_b.data(), coeffs_data + f * cstride, constants.data(),
           cdofs_b.data(), local_facet.data(), perm.data(), nullptr);
 
    if (p0_set and cells0[0] >= 0)
      P0(std::span<T>(be_b), cell_info0, cells0[0], 1);
    if (p0_set and cells0[1] >= 0)
    {
      std::span sub_be(be_b.data() + bs * dmap_size, bs * dmap_size);
      P0(sub_be, cell_info0, cells0[1], 1);
    }
 
    // Add element vector to global vector
    for (std::size_t i = 0; i < dmap0.size(); ++i)
    {
      std::int32_t dof = bs * dmap0[i];
      std::int32_t offset = bs * i;
      for (int k = 0; k < bs; ++k)
        b[dof + k] += be_b[offset + k];
    }
    for (std::size_t i = 0; i < dmap1.size(); ++i)
    {
      std::int32_t dof = bs * dmap1[i];
      std::int32_t offset = bs * (i + dmap_size);
      for (int k = 0; k < bs; ++k)
        b[dof + k] += be_b[offset + k];
    }
  }
}

template <dolfinx::scalar T, std::floating_point U, typename V>
  requires AssemblyVector<V, T>
void lift_bc(V&& b, const Form<T, U>& a, auto bs0, auto bs1,
             std::span<const T> constants,
             const std::map<std::pair<IntegralType, int>,
                            std::pair<std::span<const T>, int>>& coefficients,
             std::span<const T> bc_values1,
             std::span<const std::int8_t> bc_markers1, std::span<const T> x0,
             T alpha)
{
  // Deduce runtime block sizes as fallback when compile-time sizes
  // not given. The block size of the dofmap and indexmap is the same
  // on all sub-topologies.
  assert(bs0 == a.function_spaces()[0]->dofmaps().front()->bs());
  assert(bs1 == a.function_spaces()[1]->dofmaps().front()->bs());
 
  auto lifting_fn = [bs0, bs1, alpha, &b, &bc_values1, &bc_markers1,
                     &x0](auto rows, auto cols, auto Ae)
  {
    const std::size_t nc = cols.size() * bs1;
    for (std::size_t i = 0; i < cols.size(); ++i)
    {
      for (int k = 0; k < bs1; ++k)
      {
        const std::int32_t ii = cols[i] * bs1 + k;
        if (bc_markers1[ii])
        {
          const T x_bc = bc_values1[ii];
          const T _x0 = x0.empty() ? 0 : x0[ii];
          for (std::size_t j = 0; j < rows.size(); ++j)
          {
            for (int m = 0; m < bs0; ++m)
            {
              const std::int32_t jj = rows[j] * bs0 + m;
              b[jj] -= Ae[(j * bs0 + m) * nc + (i * bs1 + k)] * alpha
                       * (x_bc - _x0);
            }
          }
        }
      }
    }
  };
 
  // Use dolfinx::fem::impl::assemble_matrix assembler to work on the
  // vector b. With LiftingMode=true, the kernel is only called on cells
  // that have BC-constrained DOFs in the column space.
  std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
  assert(mesh);
  std::span x = mesh->geometry().x();
  md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>> _x(
      x.data(), x.size() / 3, 3);
  impl::assemble_matrix<true>(lifting_fn, a, _x, constants, coefficients, {},
                              bc_markers1);
}

template <typename V, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires AssemblyVector<V, T>
void assemble_vector(
    V&& b, const Form<T, U>& L,
    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)
{
  // Integration domain mesh
  std::shared_ptr<const mesh::Mesh<U>> mesh = L.mesh();
  assert(mesh);
 
  // Test function mesh
  auto mesh0 = L.function_spaces().at(0)->mesh();
  assert(mesh0);
 
  const int num_cell_types = mesh->topology()->cell_types().size();
  for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
  {
    // Geometry dofmap and data
    mdspan2_t x_dofmap = mesh->geometry().dofmaps().at(cell_type_idx);
    GeometryPack geometry{x_dofmap, x};
 
    // Get dofmap data
    assert(L.function_spaces().at(0));
    auto element = L.function_spaces().at(0)->elements(cell_type_idx);
    assert(element);
    std::shared_ptr<const fem::DofMap> dofmap
        = L.function_spaces().at(0)->dofmaps().at(cell_type_idx);
    assert(dofmap);
    auto dofs = dofmap->map();
    const int bs = dofmap->bs();
 
    // Buffers reused across all integral kernels for this cell type,
    // 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<T> be_buffer(2 * bs * dofs.extent(1));
    std::vector<U> cdofs_buffer(2 * 3 * x_dofmap.extent(1));
    std::span be_b(be_buffer);
    std::span cdofs_b(cdofs_buffer);
    std::span be_b1 = be_b.first(bs * dofs.extent(1));
    std::span cdofs_b1 = cdofs_b.first(3 * x_dofmap.extent(1));
 
    const fem::DofTransformKernel<T> auto& P0
        = element->template dof_transformation_fn<T>(doftransform::standard);
 
    std::span<const std::uint32_t> cell_info0;
    if (element->needs_dof_transformations())
    {
      mesh0->topology_mutable()->create_cell_permutations();
      cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info());
    }
 
    for (int i = 0; i < L.num_integrals(IntegralType::cell, 0); ++i)
    {
      auto fn = L.kernel(IntegralType::cell, i, cell_type_idx);
      assert(fn);
      std::span cells = L.domain(IntegralType::cell, i, cell_type_idx);
      std::span cells0 = L.domain_arg(IntegralType::cell, 0, i, cell_type_idx);
      auto& [coeffs, cstride] = coefficients.at({IntegralType::cell, i});
      assert(cells.size() * cstride == coeffs.size());
      impl::dispatch_bs(
          bs,
          [&b, &geometry, &cells, &dofs, &cells0, &P0, &cell_info0, &fn,
           &constants, &coeffs, cstride, &be_b1, &cdofs_b1](auto bs)
          {
            impl::assemble_cells_vector(
                b, geometry, cells,
                FormArgument{DofMapPack{dofs, bs, cells0}, P0, cell_info0}, fn,
                constants, md::mdspan(coeffs.data(), cells.size(), cstride),
                be_b1, cdofs_b1);
          });
    }
 
    md::mdspan<const std::uint8_t, md::dextents<std::size_t, 2>> facet_perms;
    if (L.needs_facet_permutations())
    {
      facet_perms = impl::entity_permutations(
          *mesh->topology_mutable(), IntegralType::interior_facet,
          mesh->topology()->cell_types()[0]);
    }
 
    using mdspanx2_t
        = md::mdspan<const std::int32_t,
                     md::extents<std::size_t, md::dynamic_extent, 2>>;
    using mdspanx22_t
        = md::mdspan<const std::int32_t,
                     md::extents<std::size_t, md::dynamic_extent, 2, 2>>;
    using mdspanx2x_t
        = md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 2,
                                          md::dynamic_extent>>;
 
    for (int i = 0; i < L.num_integrals(IntegralType::interior_facet, 0); ++i)
    {
      auto fn = L.kernel(IntegralType::interior_facet, i, 0);
      assert(fn);
      auto& [coeffs, cstride]
          = coefficients.at({IntegralType::interior_facet, i});
      std::span facets = L.domain(IntegralType::interior_facet, i, 0);
      std::span facets1 = L.domain_arg(IntegralType::interior_facet, 0, i, 0);
      assert((facets.size() / 4) * 2 * cstride == coeffs.size());
 
      mdspanx22_t facets_mdspan(facets.data(), facets.size() / 4, 2, 2);
      mdspanx22_t facets1_mdspan(facets1.data(), facets1.size() / 4, 2, 2);
      impl::dispatch_bs(
          bs,
          [&b, &geometry, &facets_mdspan, &dofs, &facets1_mdspan, &P0,
           &cell_info0, &fn, &constants, &coeffs, &facets, cstride,
           &facet_perms, &be_b, &cdofs_b](auto bs)
          {
            impl::assemble_interior_facets_vector(
                b, geometry, facets_mdspan,
                FormArgument{DofMapPack{dofs, bs, facets1_mdspan}, P0,
                             cell_info0},
                fn, constants,
                mdspanx2x_t(coeffs.data(), facets.size() / 4, 2, cstride),
                facet_perms, be_b, cdofs_b);
          });
    }
 
    for (auto itg_type : {fem::IntegralType::exterior_facet,
                          fem::IntegralType::vertex, fem::IntegralType::ridge})
    {
      const int num_itg = L.num_integrals(itg_type, 0);
      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 (L.needs_facet_permutations())
      {
        perms = impl::entity_permutations(*mesh->topology_mutable(), itg_type,
                                          mesh->topology()->cell_types()[0]);
      }
 
      for (int i = 0; i < num_itg; ++i)
      {
        auto fn = L.kernel(itg_type, i, 0);
        assert(fn);
        auto& [coeffs, cstride] = coefficients.at({itg_type, i});
        std::span e = L.domain(itg_type, i, 0);
        mdspanx2_t entities(e.data(), e.size() / 2, 2);
        std::span e1 = L.domain_arg(itg_type, 0, i, 0);
        mdspanx2_t entities1(e1.data(), e1.size() / 2, 2);
        assert((entities.size() / 2) * cstride == coeffs.size());
        impl::dispatch_bs(
            bs,
            [&b, &geometry, &entities, &dofs, &entities1, &P0, &cell_info0, &fn,
             &constants, &coeffs, cstride, &perms, &be_b1, &cdofs_b1](auto bs)
            {
              impl::assemble_entities_vector(
                  b, geometry, entities,
                  FormArgument{DofMapPack{dofs, bs, entities1}, P0, cell_info0},
                  fn, constants,
                  md::mdspan(coeffs.data(), entities.extent(0), cstride), perms,
                  be_b1, cdofs_b1);
            });
      }
    }
  }
}

template <typename V, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires AssemblyVector<V, T>
void assemble_vector(
    V&& b, const Form<T, U>& L, std::span<const T> constants,
    const std::map<std::pair<IntegralType, int>,
                   std::pair<std::span<const T>, int>>& coefficients)
{
  using mdspanx3_t
      = md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>>;
 
  std::shared_ptr<const mesh::Mesh<U>> mesh = L.mesh();
  assert(mesh);
  auto x = mesh->geometry().x();
  impl::assemble_vector(b, L, mdspanx3_t(x.data(), x.size() / 3, 3), constants,
                        coefficients);
}
} // namespace dolfinx::fem::impl