assembler.h

Go to the documentation of this file.

// Copyright (C) 2018-2026 Garth N. Wells 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 "assemble_matrix_impl.h"
#include "assemble_scalar_impl.h"
#include "assemble_vector_impl.h"
#include "pack.h"
#include "traits.h"
#include <algorithm>
#include <array>
#include <basix/mdspan.hpp>
#include <cstddef>
#include <cstdint>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/log.h>
#include <dolfinx/common/sort.h>
#include <dolfinx/common/types.h>
#include <dolfinx/la/utils.h>
#include <dolfinx/mesh/EntityMap.h>
#include <format>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <ranges>
#include <span>
#include <stdexcept>
#include <utility>
#include <vector>

 
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class DirichletBC;
template <dolfinx::scalar T, std::floating_point U>
class Form;
template <std::floating_point T>
class FunctionSpace;
 
// -- Helper functions -----------------------------------------------------

template <dolfinx::scalar T>
std::map<std::pair<IntegralType, int>, std::pair<std::span<const T>, int>>
make_coefficients_span(const std::map<std::pair<IntegralType, int>,
                                      std::pair<std::vector<T>, int>>& coeffs)
{
  using Key = typename std::remove_reference_t<decltype(coeffs)>::key_type;
  std::map<Key, std::pair<std::span<const T>, int>> c;
  std::ranges::transform(
      coeffs, std::inserter(c, c.end()),
      [](auto& e) -> typename decltype(c)::value_type
      { return {e.first, {e.second.first, e.second.second}}; });
  return c;
}
 
// -- Scalar ----------------------------------------------------------------

template <dolfinx::scalar T, std::floating_point U>
T assemble_scalar(
    const Form<T, U>& M, 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 = M.mesh();
  assert(mesh);
  std::span x = mesh->geometry().x();
 
  // Accumulate contributions from each cell type
  const int num_cell_types = mesh->topology()->cell_types().size();
  T val = 0;
  for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
  {
    // Geometry dofmap and data
    md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> x_dofmap
        = mesh->geometry().dofmaps().at(cell_type_idx);
    val += impl::assemble_scalar(M, x_dofmap,
                                 mdspanx3_t(x.data(), x.size() / 3, 3),
                                 constants, coefficients, cell_type_idx);
  }
  return val;
}

template <dolfinx::scalar T, std::floating_point U>
T assemble_scalar(const Form<T, U>& M)
{
  const std::vector<T> constants = pack_constants(M);
  auto coefficients = allocate_coefficient_storage(M);
  pack_coefficients(M, coefficients);
  return assemble_scalar(M, std::span(constants),
                         make_coefficients_span(coefficients));
}
 
// -- Vectors ----------------------------------------------------------------

// template <dolfinx::scalar T, std::floating_point U>
template <typename V, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires std::is_same_v<typename std::remove_cvref_t<V>::value_type, 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)
{
  impl::assemble_vector(b, L, constants, coefficients);
}

// template <dolfinx::scalar T, std::floating_point U>
// void assemble_vector(std::span<T> b, const Form<T, U>& L)
template <typename V, std::floating_point U,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires std::is_same_v<typename std::remove_cvref_t<V>::value_type, T>
void assemble_vector(V&& b, const Form<T, U>& L)
{
  auto coefficients = allocate_coefficient_storage(L);
  pack_coefficients(L, coefficients);
  const std::vector<T> constants = pack_constants(L);
  assemble_vector(b, L, std::span(constants),
                  make_coefficients_span(coefficients));
}
 
namespace impl
{
template <dolfinx::scalar T, std::floating_point U>
std::vector<std::int8_t> bc_dof_markers(
    const FunctionSpace<U>& V,
    const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs)
{
  std::vector<std::int8_t> markers;
  for (auto& bc : bcs)
  {
    assert(bc.get().function_space());
    if (V.contains(*bc.get().function_space()))
    {
      if (markers.empty())
      {
        std::shared_ptr<const DofMap> dofmap = V.dofmaps().front();
        std::shared_ptr<const common::IndexMap> map = dofmap->index_map;
        assert(map);
        markers.resize(dofmap->index_map_bs()
                           * (map->size_local() + map->num_ghosts()),
                       0);
      }
      bc.get().mark_dofs(markers);
    }
  }
  return markers;
}

template <dolfinx::scalar T, std::floating_point U>
std::array<std::span<const std::int8_t>, 2> bc_dof_markers_pair(
    const Form<T, U>& a,
    const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs,
    std::array<std::vector<std::int8_t>, 2>& storage)
{
  storage[0] = bc_dof_markers(*a.function_spaces().at(0), bcs);
  std::span<const std::int8_t> marker0(storage[0]);
  if (a.function_spaces().at(0) == a.function_spaces().at(1))
    return {marker0, marker0};
 
  storage[1] = bc_dof_markers(*a.function_spaces().at(1), bcs);
  return {marker0, std::span<const std::int8_t>(storage[1])};
}

template <dolfinx::scalar T, std::floating_point U>
std::pair<std::vector<std::vector<std::int8_t>>, std::vector<std::vector<T>>>
bc_lifting_data(
    const std::vector<std::optional<std::reference_wrapper<const Form<T, U>>>>&
        a,
    const std::vector<
        std::vector<std::reference_wrapper<const DirichletBC<T, U>>>>& bcs1)
{
  if (a.size() != bcs1.size())
  {
    throw std::invalid_argument(
        "Mismatch in size between a and bcs in assembler.");
  }
 
  std::vector<std::vector<std::int8_t>> markers(a.size());
  std::vector<std::vector<T>> values(a.size());
  for (std::size_t j = 0; j < a.size(); ++j)
  {
    if (a[j] and !bcs1[j].empty())
    {
      std::shared_ptr<const DofMap> dofmap
          = a[j]->get().function_spaces().at(1)->dofmaps().front();
      std::shared_ptr<const common::IndexMap> map1 = dofmap->index_map;
      assert(map1);
      const std::int32_t crange
          = dofmap->index_map_bs() * (map1->size_local() + map1->num_ghosts());
      markers[j].assign(crange, 0);
      values[j].assign(crange, 0);
      for (auto& bc : bcs1[j])
      {
        bc.get().mark_dofs(markers[j]);
        bc.get().set(values[j], std::nullopt, 1);
      }
    }
  }
 
  return {std::move(markers), std::move(values)};
}
} // namespace impl

template <typename V,
          std::floating_point U
          = scalar_value_t<typename std::remove_cvref_t<V>::value_type>,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires std::is_same_v<typename std::remove_cvref_t<V>::value_type, T>
void apply_lifting(
    V&& b,
    const std::vector<std::optional<std::reference_wrapper<const Form<T, U>>>>&
        a,
    const std::vector<std::span<const T>>& constants,
    const std::vector<std::map<std::pair<IntegralType, int>,
                               std::pair<std::span<const T>, int>>>& coeffs,
    const std::vector<std::span<const std::int8_t>>& bc_markers1,
    const std::vector<std::span<const T>>& bc_values1,
    const std::vector<std::span<const T>>& x0, T alpha)
{
  // If all forms are null, there is nothing to do
  if (std::ranges::all_of(a, [](auto ai) { return !ai; }))
    return;
 
  common::Timer t("[Apply lifting]");
 
  if (!x0.empty() and x0.size() != a.size())
  {
    throw std::invalid_argument(
        "Mismatch in size between x0 and bilinear form in assembler.");
  }
 
  if (bc_markers1.size() != a.size() or bc_values1.size() != a.size())
  {
    throw std::invalid_argument(
        "Mismatch in size between a and bc markers/values in assembler.");
  }
 
  for (std::size_t j = 0; j < a.size(); ++j)
  {
    if (!a[j] or bc_markers1[j].empty())
      continue;
 
    assert(a[j]->get().function_spaces().at(0));
    auto V1 = a[j]->get().function_spaces()[1];
    assert(V1);
 
    std::shared_ptr<const DofMap> dofmap = V1->dofmaps().front();
    auto map1 = dofmap->index_map;
    assert(map1);
    const std::size_t crange
        = dofmap->index_map_bs() * (map1->size_local() + map1->num_ghosts());
    if (bc_markers1[j].size() != crange or bc_values1[j].size() != crange)
    {
      throw std::invalid_argument(std::format(
          "bc markers/values for block {} have length {}/{}, expected {}.", j,
          bc_markers1[j].size(), bc_values1[j].size(), crange));
    }
 
    const int bs0 = a[j]->get().function_spaces()[0]->dofmaps().front()->bs();
    const int bs1 = dofmap->bs();
 
    std::span<const T> _x0;
    if (!x0.empty())
      _x0 = x0[j];
 
    impl::dispatch_bs(bs0, bs1,
                      [&b, &a, j, &constants, &coeffs, &bc_values1,
                       &bc_markers1, &_x0, alpha](auto bs0, auto bs1)
                      {
                        impl::lift_bc(b, a[j]->get(), bs0, bs1, constants[j],
                                      coeffs[j], bc_values1[j], bc_markers1[j],
                                      _x0, alpha);
                      });
  }
}

template <typename V,
          std::floating_point U
          = scalar_value_t<typename std::remove_cvref_t<V>::value_type>,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires std::is_same_v<typename std::remove_cvref_t<V>::value_type, T>
void apply_lifting(
    V&& b,
    const std::vector<std::optional<std::reference_wrapper<const Form<T, U>>>>&
        a,
    const std::vector<std::span<const T>>& constants,
    const std::vector<std::map<std::pair<IntegralType, int>,
                               std::pair<std::span<const T>, int>>>& coeffs,
    const std::vector<
        std::vector<std::reference_wrapper<const DirichletBC<T, U>>>>& bcs1,
    const std::vector<std::span<const T>>& x0, T alpha)
{
  auto [bc_markers1, bc_values1] = impl::bc_lifting_data(a, bcs1);
  apply_lifting(
      b, a, constants, coeffs,
      std::vector<std::span<const std::int8_t>>(bc_markers1.begin(),
                                                bc_markers1.end()),
      std::vector<std::span<const T>>(bc_values1.begin(), bc_values1.end()), x0,
      alpha);
}

template <typename V,
          std::floating_point U
          = scalar_value_t<typename std::remove_cvref_t<V>::value_type>,
          dolfinx::scalar T = typename std::remove_cvref_t<V>::value_type>
  requires std::is_same_v<typename std::remove_cvref_t<V>::value_type, T>
void apply_lifting(
    V&& b,
    const std::vector<std::optional<std::reference_wrapper<const Form<T, U>>>>&
        a,
    const std::vector<
        std::vector<std::reference_wrapper<const DirichletBC<T, U>>>>& bcs1,
    const std::vector<std::span<const T>>& x0, T alpha)
{
  std::vector<
      std::map<std::pair<IntegralType, int>, std::pair<std::vector<T>, int>>>
      coeffs;
  std::vector<std::vector<T>> constants;
  for (const auto& _a : a)
  {
    if (_a)
    {
      auto coefficients = allocate_coefficient_storage(_a->get());
      pack_coefficients(_a->get(), coefficients);
      coeffs.push_back(coefficients);
      constants.push_back(pack_constants(_a->get()));
    }
    else
    {
      coeffs.emplace_back();
      constants.emplace_back();
    }
  }
 
  std::vector<std::span<const T>> _constants(constants.begin(),
                                             constants.end());
  std::vector<std::map<std::pair<IntegralType, int>,
                       std::pair<std::span<const T>, int>>>
      _coeffs;
  std::ranges::transform(coeffs, std::back_inserter(_coeffs),
                         [](auto& c) { return make_coefficients_span(c); });
 
  auto [bc_markers1, bc_values1] = impl::bc_lifting_data(a, bcs1);
  apply_lifting(
      b, a, _constants, _coeffs,
      std::vector<std::span<const std::int8_t>>(bc_markers1.begin(),
                                                bc_markers1.end()),
      std::vector<std::span<const T>>(bc_values1.begin(), bc_values1.end()), x0,
      alpha);
}
 
// -- Matrices ---------------------------------------------------------------

template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix(
    la::MatSet<T> auto mat_add, const Form<T, U>& a,
    std::span<const T> constants,
    const std::map<std::pair<IntegralType, int>,
                   std::pair<std::span<const T>, int>>& coefficients,
    std::span<const std::int8_t> dof_marker0,
    std::span<const std::int8_t> dof_marker1)
 
{
  common::Timer t_assm("[Assemble Matrix]");
  using mdspanx3_t
      = md::mdspan<const U, md::extents<std::size_t, md::dynamic_extent, 3>>;
 
  std::shared_ptr<const mesh::Mesh<U>> mesh = a.mesh();
  assert(mesh);
  std::span x = mesh->geometry().x();
  impl::assemble_matrix<false>(mat_add, a,
                               mdspanx3_t(x.data(), x.size() / 3, 3), constants,
                               coefficients, dof_marker0, dof_marker1);
}

template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix(
    auto mat_add, const Form<T, U>& a,
    const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs)
{
  // Prepare constants and coefficients
  const std::vector<T> constants = pack_constants(a);
  auto coefficients = allocate_coefficient_storage(a);
  pack_coefficients(a, coefficients);
 
  std::array<std::vector<std::int8_t>, 2> markers;
  auto [dof_marker0, dof_marker1] = impl::bc_dof_markers_pair(a, bcs, markers);
 
  // Assemble
  assemble_matrix(mat_add, a, std::span<const T>(constants),
                  make_coefficients_span(coefficients), dof_marker0,
                  dof_marker1);
}

template <dolfinx::scalar T, std::floating_point U>
void assemble_matrix(auto mat_add, const Form<T, U>& a,
                     std::span<const std::int8_t> dof_marker0,
                     std::span<const std::int8_t> dof_marker1)
 
{
  // Prepare constants and coefficients
  const std::vector<T> constants = pack_constants(a);
  auto coefficients = allocate_coefficient_storage(a);
  pack_coefficients(a, coefficients);
 
  // Assemble
  assemble_matrix(mat_add, a, std::span<const T>(constants),
                  make_coefficients_span(coefficients), dof_marker0,
                  dof_marker1);
}

template <dolfinx::scalar T, std::floating_point U>
void set_diagonal(
    auto set_fn, const FunctionSpace<U>& V,
    const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs,
    T diagonal = T(1))
{
  logger().debug("Set diagonal");
 
  std::vector<std::int32_t> rows;
  for (auto& bc : bcs)
  {
    if (V.contains(*bc.get().function_space()))
    {
      const auto [dofs, range] = bc.get().dof_indices();
      std::span<const std::int32_t> owned = dofs.first(range);
      rows.insert(rows.end(), owned.begin(), owned.end());
    }
  }
 
  // A condition's dofs are strictly increasing (a DirichletBC
  // precondition), so one condition needs no sort. Several give sorted
  // runs, which a comparison sort handles poorly and which are often
  // already in order, hence the check before radix sorting.
  if (!std::ranges::is_sorted(rows))
    dolfinx::radix_sort(rows);
 
  // Overlapping conditions can repeat a row
  rows.erase(std::ranges::unique(rows).begin(), rows.end());
  la::set_diagonal(set_fn, rows, diagonal);
}
 
} // namespace dolfinx::fem