DOLFINx
DOLFINx C++ interface
assembler.h
1// Copyright (C) 2018-2021 Garth N. Wells
2//
3// This file is part of DOLFINx (https://www.fenicsproject.org)
4//
5// SPDX-License-Identifier: LGPL-3.0-or-later
6
7#pragma once
8
9#include "assemble_matrix_impl.h"
10#include "assemble_scalar_impl.h"
11#include "assemble_vector_impl.h"
12#include <cstdint>
13#include <memory>
14#include <span>
15#include <vector>
16
17namespace dolfinx::fem
18{
19template <typename T>
20class DirichletBC;
21template <typename T>
22class Form;
23class FunctionSpace;
24
25// -- Helper functions -----------------------------------------------------
26
28template <typename T>
29std::map<std::pair<dolfinx::fem::IntegralType, int>,
30 std::pair<std::span<const T>, int>>
31make_coefficients_span(const std::map<std::pair<IntegralType, int>,
32 std::pair<std::vector<T>, int>>& coeffs)
33{
34 using Key = typename std::remove_reference_t<decltype(coeffs)>::key_type;
35 std::map<Key, std::pair<std::span<const T>, int>> c;
36 std::transform(coeffs.cbegin(), coeffs.cend(), std::inserter(c, c.end()),
37 [](auto& e) -> typename decltype(c)::value_type {
38 return {e.first, {e.second.first, e.second.second}};
39 });
40 return c;
41}
42
43// -- Scalar ----------------------------------------------------------------
44
54template <typename T>
55T assemble_scalar(
56 const Form<T>& M, const std::span<const T>& constants,
57 const std::map<std::pair<IntegralType, int>,
58 std::pair<std::span<const T>, int>>& coefficients)
59{
60 return impl::assemble_scalar(M, constants, coefficients);
61}
62
68template <typename T>
69T assemble_scalar(const Form<T>& M)
70{
71 const std::vector<T> constants = pack_constants(M);
72 auto coefficients = allocate_coefficient_storage(M);
73 pack_coefficients(M, coefficients);
74 return assemble_scalar(M, std::span(constants),
75 make_coefficients_span(coefficients));
76}
77
78// -- Vectors ----------------------------------------------------------------
79
88template <typename T>
89void assemble_vector(
90 std::span<T> b, const Form<T>& L, const std::span<const T>& constants,
91 const std::map<std::pair<IntegralType, int>,
92 std::pair<std::span<const T>, int>>& coefficients)
93{
94 impl::assemble_vector(b, L, constants, coefficients);
95}
96
101template <typename T>
102void assemble_vector(std::span<T> b, const Form<T>& L)
103{
104 auto coefficients = allocate_coefficient_storage(L);
105 pack_coefficients(L, coefficients);
106 const std::vector<T> constants = pack_constants(L);
107 assemble_vector(b, L, std::span(constants),
108 make_coefficients_span(coefficients));
109}
110
111// FIXME: clarify how x0 is used
112// FIXME: if bcs entries are set
113
114// FIXME: need to pass an array of Vec for x0?
115// FIXME: clarify zeroing of vector
116
129template <typename T>
130void apply_lifting(
131 std::span<T> b, const std::vector<std::shared_ptr<const Form<T>>>& a,
132 const std::vector<std::span<const T>>& constants,
133 const std::vector<std::map<std::pair<IntegralType, int>,
134 std::pair<std::span<const T>, int>>>& coeffs,
135 const std::vector<std::vector<std::shared_ptr<const DirichletBC<T>>>>& bcs1,
136 const std::vector<std::span<const T>>& x0, double scale)
137{
138 impl::apply_lifting(b, a, constants, coeffs, bcs1, x0, scale);
139}
140
153template <typename T>
154void apply_lifting(
155 std::span<T> b, const std::vector<std::shared_ptr<const Form<T>>>& a,
156 const std::vector<std::vector<std::shared_ptr<const DirichletBC<T>>>>& bcs1,
157 const std::vector<std::span<const T>>& x0, double scale)
158{
159 std::vector<
160 std::map<std::pair<IntegralType, int>, std::pair<std::vector<T>, int>>>
161 coeffs;
162 std::vector<std::vector<T>> constants;
163 for (auto _a : a)
164 {
165 if (_a)
166 {
167 auto coefficients = allocate_coefficient_storage(*_a);
168 pack_coefficients(*_a, coefficients);
169 coeffs.push_back(coefficients);
170 constants.push_back(pack_constants(*_a));
171 }
172 else
173 {
174 coeffs.push_back(std::map<std::pair<IntegralType, int>,
175 std::pair<std::vector<T>, int>>());
176 constants.push_back({});
177 }
178 }
179
180 std::vector<std::span<const T>> _constants(constants.begin(),
181 constants.end());
182 std::vector<std::map<std::pair<IntegralType, int>,
183 std::pair<std::span<const T>, int>>>
184 _coeffs;
185 std::transform(coeffs.cbegin(), coeffs.cend(), std::back_inserter(_coeffs),
186 [](auto& c) { return make_coefficients_span(c); });
187 apply_lifting(b, a, _constants, _coeffs, bcs1, x0, scale);
188}
189
190// -- Matrices ---------------------------------------------------------------
191
199template <typename T, typename U>
200void assemble_matrix(
201 U mat_add, const Form<T>& a, const std::span<const T>& constants,
202 const std::map<std::pair<IntegralType, int>,
203 std::pair<std::span<const T>, int>>& coefficients,
204 const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs)
205{
206 // Index maps for dof ranges
207 auto map0 = a.function_spaces().at(0)->dofmap()->index_map;
208 auto map1 = a.function_spaces().at(1)->dofmap()->index_map;
209 auto bs0 = a.function_spaces().at(0)->dofmap()->index_map_bs();
210 auto bs1 = a.function_spaces().at(1)->dofmap()->index_map_bs();
211
212 // Build dof markers
213 std::vector<std::int8_t> dof_marker0, dof_marker1;
214 assert(map0);
215 std::int32_t dim0 = bs0 * (map0->size_local() + map0->num_ghosts());
216 assert(map1);
217 std::int32_t dim1 = bs1 * (map1->size_local() + map1->num_ghosts());
218 for (std::size_t k = 0; k < bcs.size(); ++k)
219 {
220 assert(bcs[k]);
221 assert(bcs[k]->function_space());
222 if (a.function_spaces().at(0)->contains(*bcs[k]->function_space()))
223 {
224 dof_marker0.resize(dim0, false);
225 bcs[k]->mark_dofs(dof_marker0);
226 }
227
228 if (a.function_spaces().at(1)->contains(*bcs[k]->function_space()))
229 {
230 dof_marker1.resize(dim1, false);
231 bcs[k]->mark_dofs(dof_marker1);
232 }
233 }
234
235 // Assemble
236 impl::assemble_matrix(mat_add, a, constants, coefficients, dof_marker0,
237 dof_marker1);
238}
239
245template <typename T, typename U>
246void assemble_matrix(
247 U mat_add, const Form<T>& a,
248 const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs)
249{
250 // Prepare constants and coefficients
251 const std::vector<T> constants = pack_constants(a);
252 auto coefficients = allocate_coefficient_storage(a);
253 pack_coefficients(a, coefficients);
254
255 // Assemble
256 assemble_matrix(mat_add, a, std::span(constants),
257 make_coefficients_span(coefficients), bcs);
258}
259
272template <typename T, typename U>
273void assemble_matrix(
274 U mat_add, const Form<T>& a, const std::span<const T>& constants,
275 const std::map<std::pair<IntegralType, int>,
276 std::pair<std::span<const T>, int>>& coefficients,
277 const std::span<const std::int8_t>& dof_marker0,
278 const std::span<const std::int8_t>& dof_marker1)
279
280{
281 impl::assemble_matrix(mat_add, a, constants, coefficients, dof_marker0,
282 dof_marker1);
283}
284
295template <typename T, typename U>
296void assemble_matrix(U mat_add, const Form<T>& a,
297 const std::span<const std::int8_t>& dof_marker0,
298 const std::span<const std::int8_t>& dof_marker1)
299
300{
301 // Prepare constants and coefficients
302 const std::vector<T> constants = pack_constants(a);
303 auto coefficients = allocate_coefficient_storage(a);
304 pack_coefficients(a, coefficients);
305
306 // Assemble
307 assemble_matrix(mat_add, a, std::span(constants),
308 make_coefficients_span(coefficients), dof_marker0,
309 dof_marker1);
310}
311
322template <typename T, typename U>
323void set_diagonal(U set_fn, const std::span<const std::int32_t>& rows,
324 T diagonal = 1.0)
325{
326 for (std::size_t i = 0; i < rows.size(); ++i)
327 {
328 std::span diag_span(&diagonal, 1);
329 set_fn(rows.subspan(i, 1), rows.subspan(i, 1), diag_span);
330 }
331}
332
348template <typename T, typename U>
349void set_diagonal(U set_fn, const fem::FunctionSpace& V,
350 const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
351 T diagonal = 1.0)
352{
353 for (const auto& bc : bcs)
354 {
355 assert(bc);
356 if (V.contains(*bc->function_space()))
357 {
358 const auto [dofs, range] = bc->dof_indices();
359 set_diagonal(set_fn, dofs.first(range), diagonal);
360 }
361 }
362}
363
364// -- Setting bcs ------------------------------------------------------------
365
366// FIXME: Move these function elsewhere?
367
368// FIXME: clarify x0
369// FIXME: clarify what happens with ghosts
370
374template <typename T>
375void set_bc(std::span<T> b,
376 const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
377 const std::span<const T>& x0, double scale = 1.0)
378{
379 if (b.size() > x0.size())
380 throw std::runtime_error("Size mismatch between b and x0 vectors.");
381 for (const auto& bc : bcs)
382 {
383 assert(bc);
384 bc->set(b, x0, scale);
385 }
386}
387
391template <typename T>
392void set_bc(std::span<T> b,
393 const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
394 double scale = 1.0)
395{
396 for (const auto& bc : bcs)
397 {
398 assert(bc);
399 bc->set(b, scale);
400 }
401}
402
403} // namespace dolfinx::fem
Object for setting (strong) Dirichlet boundary conditions.
Definition: DirichletBC.h:125
A representation of finite element variational forms.
Definition: Form.h:63
const std::vector< std::shared_ptr< const FunctionSpace > > & function_spaces() const
Return function spaces for all arguments.
Definition: Form.h:181
This class represents a finite element function space defined by a mesh, a finite element,...
Definition: FunctionSpace.h:31
bool contains(const FunctionSpace &V) const
Check whether V is subspace of this, or this itself.
Definition: FunctionSpace.cpp:69
Finite element method functionality.
Definition: assemble_matrix_impl.h:25
std::vector< typename U::scalar_type > pack_constants(const U &u)
Pack constants of u of generic type U ready for assembly.
Definition: utils.h:970
void set_bc(std::span< T > b, const std::vector< std::shared_ptr< const DirichletBC< T > > > &bcs, const std::span< const T > &x0, double scale=1.0)
Set bc values in owned (local) part of the vector, multiplied by 'scale'. The vectors b and x0 must h...
Definition: assembler.h:375
std::map< std::pair< dolfinx::fem::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 std::span from a map of std::vector.
Definition: assembler.h:31
std::pair< std::vector< T >, int > allocate_coefficient_storage(const Form< T > &form, IntegralType integral_type, int id)
Allocate storage for coefficients of a pair (integral_type, id) from a fem::Form form.
Definition: utils.h:675
void pack_coefficients(const Form< T > &form, IntegralType integral_type, int id, const std::span< T > &c, int cstride)
Pack coefficients of a Form for a given integral type and domain id.
Definition: utils.h:738
void set_diagonal(U set_fn, const std::span< const std::int32_t > &rows, T diagonal=1.0)
Sets a value to the diagonal of a matrix for specified rows. It is typically called after assembly....
Definition: assembler.h:323