Loading...
Searching...
No Matches
assembler.h
Go to the documentation of this file.
// Copyright (C) 2018-2025 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once
#include "Function.h"
#include "FunctionSpace.h"
#include "assemble_expression_impl.h"
#include "assemble_matrix_impl.h"
#include "assemble_scalar_impl.h"
#include "assemble_vector_impl.h"
#include "pack.h"
#include "traits.h"
#include "utils.h"
#include <algorithm>
#include <basix/mdspan.hpp>
#include <cstdint>
#include <dolfinx/common/types.h>
#include <memory>
#include <optional>
#include <span>
#include <vector>
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class DirichletBC;
template <dolfinx::scalar T, std::floating_point U>
class Expression;
template <dolfinx::scalar T, std::floating_point U>
class Form;
template <std::floating_point T>
class FunctionSpace;
template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(
std::span<T> values, const fem::Expression<T, U>& e,
md::mdspan<const T, md::dextents<std::size_t, 2>> coeffs,
std::span<const T> constants, const mesh::Mesh<U>& mesh,
fem::MDSpan2 auto entities,
std::optional<
std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
element)
{
auto [X, Xshape] = e.X();
impl::tabulate_expression(values, e.kernel(), Xshape, e.value_size(), coeffs,
constants, mesh, entities, element);
}
template <dolfinx::scalar T, std::floating_point U>
void tabulate_expression(std::span<T> values, const fem::Expression<T, U>& e,
const mesh::Mesh<U>& mesh, fem::MDSpan2 auto entities)
{
std::optional<
std::pair<std::reference_wrapper<const FiniteElement<U>>, std::size_t>>
element = std::nullopt;
if (auto V = e.argument_space(); V)
{
std::size_t num_argument_dofs
= V->dofmap()->element_dof_layout().num_dofs() * V->dofmap()->bs();
assert(V->element());
element = {std::cref(*V->element()), num_argument_dofs};
}
std::vector<int> coffsets = e.coefficient_offsets();
const std::vector<std::shared_ptr<const Function<T, U>>>& coefficients
= e.coefficients();
std::vector<T> coeffs(entities.extent(0) * coffsets.back());
int cstride = coffsets.back();
{
std::vector<std::reference_wrapper<const Function<T, U>>> c;
std::ranges::transform(coefficients, std::back_inserter(c),
[](auto c) -> const Function<T, U>& { return *c; });
fem::pack_coefficients(c, coffsets, entities, std::span(coeffs));
}
std::vector<T> constants = fem::pack_constants(e);
tabulate_expression<T, U>(
values, e, md::mdspan(coeffs.data(), entities.size(), cstride),
std::span<const T>(constants), mesh, entities, element);
}
// -- 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 scalar_value_t<T>,
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();
if constexpr (std::is_same_v<U, scalar_value_t<T>>)
{
return impl::assemble_scalar(M, mesh->geometry().dofmap(),
mdspanx3_t(x.data(), x.size() / 3, 3),
constants, coefficients);
}
else
{
std::vector<scalar_value_t<T>> _x(x.begin(), x.end());
return impl::assemble_scalar(M, mesh->geometry().dofmap(),
mdspanx3_t(_x.data(), _x.size() / 3, 3),
constants, coefficients);
}
}
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));
}
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,
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)
{
// If all forms are null, there is nothing to do
if (std::ranges::all_of(a, [](auto ai) { return !ai; }))
return;
impl::apply_lifting(b, a, constants, coeffs, bcs1, 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,
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 (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); });
apply_lifting(b, a, _constants, _coeffs, bcs1, 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)
{
using mdspanx3_t
= md::mdspan<const scalar_value_t<T>,
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();
if constexpr (std::is_same_v<U, scalar_value_t<T>>)
{
impl::assemble_matrix(mat_add, a, mdspanx3_t(x.data(), x.size() / 3, 3),
constants, coefficients, dof_marker0, dof_marker1);
}
else
{
std::vector<scalar_value_t<T>> _x(x.begin(), x.end());
impl::assemble_matrix(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, std::span<const T> constants,
const std::map<std::pair<IntegralType, int>,
std::pair<std::span<const T>, int>>& coefficients,
const std::vector<std::reference_wrapper<const DirichletBC<T, U>>>& bcs)
{
// Index maps for dof ranges
// NOTE: For mixed-topology meshes, there will be multiple DOF maps,
// but the index maps are the same.
auto map0 = a.function_spaces().at(0)->dofmaps(0)->index_map;
auto map1 = a.function_spaces().at(1)->dofmaps(0)->index_map;
auto bs0 = a.function_spaces().at(0)->dofmaps(0)->index_map_bs();
auto bs1 = a.function_spaces().at(1)->dofmaps(0)->index_map_bs();
// Build dof markers
std::vector<std::int8_t> dof_marker0, dof_marker1;
assert(map0);
std::int32_t dim0 = bs0 * (map0->size_local() + map0->num_ghosts());
assert(map1);
std::int32_t dim1 = bs1 * (map1->size_local() + map1->num_ghosts());
for (std::size_t k = 0; k < bcs.size(); ++k)
{
assert(bcs[k].get().function_space());
if (a.function_spaces().at(0)->contains(*bcs[k].get().function_space()))
{
dof_marker0.resize(dim0, false);
bcs[k].get().mark_dofs(dof_marker0);
}
if (a.function_spaces().at(1)->contains(*bcs[k].get().function_space()))
{
dof_marker1.resize(dim1, false);
bcs[k].get().mark_dofs(dof_marker1);
}
}
// Assemble
assemble_matrix(mat_add, a, 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);
// Assemble
assemble_matrix(mat_add, a, std::span(constants),
make_coefficients_span(coefficients), bcs);
}
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(constants),
make_coefficients_span(coefficients), dof_marker0,
dof_marker1);
}
template <dolfinx::scalar T>
void set_diagonal(auto set_fn, std::span<const std::int32_t> rows,
T diagonal = 1.0)
{
for (std::size_t i = 0; i < rows.size(); ++i)
{
std::span diag_span(&diagonal, 1);
set_fn(rows.subspan(i, 1), rows.subspan(i, 1), diag_span);
}
}
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 = 1.0)
{
for (auto& bc : bcs)
{
if (V.contains(*bc.get().function_space()))
{
const auto [dofs, range] = bc.get().dof_indices();
set_diagonal(set_fn, dofs.first(range), diagonal);
}
}
}
} // namespace dolfinx::fem
dolfinx::fem::DirichletBC
Definition DirichletBC.h:255
dolfinx::fem::Expression
An Expression represents a mathematical expression evaluated at a pre-defined points on a reference c...
Definition Expression.h:41
dolfinx::fem::FiniteElement
Model of a finite element.
Definition FiniteElement.h:57
dolfinx::fem::Form
A representation of finite element variational forms.
Definition Form.h:117
dolfinx::fem::FunctionSpace
This class represents a finite element function space defined by a mesh, a finite element,...
Definition FunctionSpace.h:34
dolfinx::fem::Function
Definition Function.h:47
dolfinx::mesh::Mesh
A Mesh consists of a set of connected and numbered mesh topological entities, and geometry data.
Definition Mesh.h:23
dolfinx::fem::MDSpan2
Concept for mdspan of rank 1 or 2.
Definition traits.h:36
dolfinx::la::MatSet
Matrix accumulate/set concept for functions that can be used in assemblers to accumulate or set value...
Definition utils.h:28
dolfinx::scalar
Definition types.h:20
utils.h
Functions supporting finite element method operations.
dolfinx::fem
Finite element method functionality.
Definition assemble_expression_impl.h:23
dolfinx::fem::allocate_coefficient_storage
std::pair< std::vector< T >, int > allocate_coefficient_storage(const Form< T, U > &form, IntegralType integral_type, int id)
Allocate storage for coefficients of a pair (integral_type, id) from a Form.
Definition pack.h:170
dolfinx::fem::assemble_matrix
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)
Assemble bilinear form into a matrix. Matrix must already be initialised. Does not zero or finalise t...
Definition assembler.h:429
dolfinx::fem::assemble_scalar
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)
Assemble functional into scalar.
Definition assembler.h:158
dolfinx::fem::set_diagonal
void set_diagonal(auto set_fn, std::span< const std::int32_t > rows, T diagonal=1.0)
Sets a value to the diagonal of a matrix for specified rows.
Definition assembler.h:568
dolfinx::fem::tabulate_expression
void tabulate_expression(std::span< T > values, const fem::Expression< T, U > &e, md::mdspan< const T, md::dextents< std::size_t, 2 > > coeffs, std::span< const T > constants, const mesh::Mesh< U > &mesh, fem::MDSpan2 auto entities, std::optional< std::pair< std::reference_wrapper< const FiniteElement< U > >, std::size_t > > element)
Evaluate an Expression on cells or facets.
Definition assembler.h:64
dolfinx::fem::apply_lifting
void apply_lifting(V &&b, 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)
Modify the right-hand side vector to account for constraints (Dirichlet boundary condition constraint...
Definition assembler.h:326
dolfinx::fem::make_coefficients_span
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)
Create a map of std::spans from a map of std::vectors.
Definition assembler.h:132
dolfinx::fem::pack_coefficients
void pack_coefficients(const Form< T, U > &form, std::map< std::pair< IntegralType, int >, std::pair< std::vector< T >, int > > &coeffs)
Pack coefficients of a Form.
Definition pack.h:222
dolfinx::fem::assemble_vector
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)
Assemble linear form into a vector.
Definition assembler.h:218
dolfinx::fem::pack_constants
std::vector< T > pack_constants(std::vector< std::reference_wrapper< const fem::Constant< T > > > c)
Pack constants of an Expression or Form into a single array ready for assembly.
Definition pack.h:378
dolfinx::mesh
Mesh data structures and algorithms on meshes.
Definition DofMap.h:32
pack.h
Functions supporting the packing of coefficient data.