DOLFINx 0.12.0.0
DOLFINx C++
Loading...
Searching...
No Matches
petsc.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#ifdef HAS_PETSC
10
11#include "Form.h"
12#include "assembler.h"
13#include "utils.h"
14#include <cassert>
15#include <concepts>
16#include <cstdint>
17#include <dolfinx/la/petsc.h>
18#include <functional>
19#include <map>
20#include <memory>
21#include <numeric>
22#include <optional>
23#include <petscmat.h>
24#include <petscvec.h>
25#include <ranges>
26#include <span>
27#include <stdexcept>
28#include <utility>
29#include <vector>
30
31namespace dolfinx::common
32{
33class IndexMap;
34}
35
36namespace dolfinx::fem
37{
38template <dolfinx::scalar T, std::floating_point U>
39class DirichletBC;
40
42namespace petsc
43{
44// Check a PETSc error code and throw a descriptive exception if it is
45// non-zero. Expects a local `PetscErrorCode ierr` in scope.
46#define CHECK_ERROR(NAME) \
47 do \
48 { \
49 if (ierr != 0) \
50 la::petsc::error(ierr, __FILE__, NAME); \
51 } while (0)
52
59template <std::floating_point T>
61 std::optional<std::string> type = std::nullopt)
62{
64 pattern.finalize();
65 return la::petsc::create_matrix(a.mesh()->comm(), pattern, type);
66}
67
78template <std::floating_point T>
80 const std::vector<std::vector<const Form<PetscScalar, T>*>>& a,
81 std::optional<std::string> type = std::nullopt)
82{
83 // Extract and check row/column ranges
84 std::array<std::vector<std::shared_ptr<const FunctionSpace<T>>>, 2> V
86 std::array<std::vector<int>, 2> bs_dofs;
87 for (std::size_t i = 0; i < 2; ++i)
88 {
89 for (auto& _V : V[i])
90 bs_dofs[i].push_back(_V->dofmap()->bs());
91 }
92
93 // Build sparsity pattern for each block
94 std::shared_ptr<const mesh::Mesh<T>> mesh;
95 std::vector<std::vector<std::unique_ptr<la::SparsityPattern>>> patterns(
96 V[0].size());
97 for (std::size_t row = 0; row < V[0].size(); ++row)
98 {
99 for (std::size_t col = 0; col < V[1].size(); ++col)
100 {
101 if (const Form<PetscScalar, T>* form = a[row][col]; form)
102 {
103 patterns[row].push_back(std::make_unique<la::SparsityPattern>(
105 if (!mesh)
106 mesh = form->mesh();
107 }
108 else
109 patterns[row].push_back(nullptr);
110 }
111 }
112
113 if (!mesh)
114 throw std::runtime_error("Could not find a Mesh.");
115
116 // Compute offsets for the fields
117 std::array<std::vector<std::pair<
118 std::reference_wrapper<const common::IndexMap>, int>>,
119 2>
120 maps;
121 for (std::size_t d = 0; d < 2; ++d)
122 {
123 for (auto& space : V[d])
124 {
125 maps[d].emplace_back(*space->dofmap()->index_map,
126 space->dofmap()->index_map_bs());
127 }
128 }
129
130 // Create merged sparsity pattern
131 std::vector<std::vector<const la::SparsityPattern*>> p(V[0].size());
132 for (std::size_t row = 0; row < V[0].size(); ++row)
133 for (std::size_t col = 0; col < V[1].size(); ++col)
134 p[row].push_back(patterns[row][col].get());
135
136 la::SparsityPattern pattern(mesh->comm(), p, maps, bs_dofs);
137 pattern.finalize();
138
139 // FIXME: Add option to pass customised local-to-global map to PETSc
140 // Mat constructor
141
142 // TODO: Index map concatenation has already been computed inside
143 // the SparsityPattern constructor, but we also need it here to
144 // build the PETSc local-to-global map. Compute outside and pass
145 // into SparsityPattern constructor.
146
147 // Initialise matrix
148 Mat A = la::petsc::create_matrix(mesh->comm(), pattern, type);
149
150 // Create row and column local-to-global maps (field0, field1, field2,
151 // etc), i.e. ghosts of field0 appear before owned indices of field1
152 std::array<std::vector<PetscInt>, 2> _maps;
153 for (int d = 0; d < 2; ++d)
154 {
155 if (d == 1 and V[0] == V[1])
156 {
157 // Row and column spaces are identical, so the concatenated
158 // index map for d=1 is identical to the one already computed
159 // for d=0 -- reuse it rather than paying for a second,
160 // communication-heavy call to stack_index_maps.
161 _maps[1] = _maps[0];
162 continue;
163 }
164
165 const std::vector<
166 std::pair<std::reference_wrapper<const common::IndexMap>, int>>& map
167 = maps[d];
168 std::vector<PetscInt>& _map = _maps[d];
169
170 // Concatenate the block index map in the row and column directions
171 const auto [rank_offset, local_offset, ghosts, _]
173 const std::size_t num_ghosts
174 = std::accumulate(ghosts.begin(), ghosts.end(), std::size_t(0),
175 [](std::size_t n, auto& g) { return n + g.size(); });
176 _map.reserve(local_offset.back() + num_ghosts);
177 for (std::size_t f = 0; f < map.size(); ++f)
178 {
179 auto offset = local_offset[f];
180 const common::IndexMap& imap = map[f].first.get();
181 int bs = map[f].second;
182 auto owned
183 = std::views::iota(std::int32_t(0), bs * imap.size_local())
184 | std::views::transform([offset, rank_offset](std::int32_t i)
185 { return i + rank_offset + offset; });
186 _map.insert(_map.end(), owned.begin(), owned.end());
187 _map.insert(_map.end(), ghosts[f].begin(), ghosts[f].end());
188 }
189 }
190
191 // Create PETSc local-to-global map/index sets and attach to matrix
192 ISLocalToGlobalMapping petsc_local_to_global0;
193 PetscErrorCode ierr = ISLocalToGlobalMappingCreate(
194 MPI_COMM_SELF, 1, _maps[0].size(), _maps[0].data(), PETSC_COPY_VALUES,
195 &petsc_local_to_global0);
196 CHECK_ERROR("ISLocalToGlobalMappingCreate");
197 if (V[0] == V[1])
198 {
199 ierr = MatSetLocalToGlobalMapping(A, petsc_local_to_global0,
200 petsc_local_to_global0);
201 CHECK_ERROR("MatSetLocalToGlobalMapping");
202 ierr = ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
203 CHECK_ERROR("ISLocalToGlobalMappingDestroy");
204 }
205 else
206 {
207 ISLocalToGlobalMapping petsc_local_to_global1;
208 ierr = ISLocalToGlobalMappingCreate(MPI_COMM_SELF, 1, _maps[1].size(),
209 _maps[1].data(), PETSC_COPY_VALUES,
210 &petsc_local_to_global1);
211 CHECK_ERROR("ISLocalToGlobalMappingCreate");
212 ierr = MatSetLocalToGlobalMapping(A, petsc_local_to_global0,
213 petsc_local_to_global1);
214 CHECK_ERROR("MatSetLocalToGlobalMapping");
215 ierr = ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
216 CHECK_ERROR("ISLocalToGlobalMappingDestroy");
217 ierr = ISLocalToGlobalMappingDestroy(&petsc_local_to_global1);
218 CHECK_ERROR("ISLocalToGlobalMappingDestroy");
219 }
220
221 return A;
222}
223
227template <std::floating_point T>
229 const std::vector<std::vector<const Form<PetscScalar, T>*>>& a,
230 std::optional<std::vector<std::vector<std::optional<std::string>>>> types)
231{
232 if (a.empty())
233 throw std::runtime_error("Rectangular array of forms must be non-empty.");
234
235 // Extract and check row/column ranges
237
238 // Loop over each form and create matrix
239 int rows = a.size();
240 int cols = a.front().size();
241 std::vector<Mat> mats(rows * cols, nullptr);
242 std::shared_ptr<const mesh::Mesh<T>> mesh;
243 for (int i = 0; i < rows; ++i)
244 {
245 for (int j = 0; j < cols; ++j)
246 {
247 if (const Form<PetscScalar, T>* form = a[i][j]; form)
248 {
249 if (types)
250 mats[i * cols + j] = create_matrix(*form, types->at(i).at(j));
251 else
252 mats[i * cols + j] = create_matrix(*form, std::nullopt);
253 mesh = form->mesh();
254 }
255 }
256 }
257
258 if (!mesh)
259 throw std::runtime_error("Could not find a Mesh.");
260
261 // Initialise block (MatNest) matrix. On error, destroy the
262 // already-created sub-matrices in `mats` before propagating, since
263 // the nest (which would otherwise take joint ownership of them) was
264 // never successfully assembled.
265 Mat A;
266 try
267 {
268 PetscErrorCode ierr = MatCreate(mesh->comm(), &A);
269 CHECK_ERROR("MatCreate");
270 ierr = MatSetType(A, MATNEST);
271 CHECK_ERROR("MatSetType");
272 ierr = MatNestSetSubMats(A, rows, nullptr, cols, nullptr, mats.data());
273 CHECK_ERROR("MatNestSetSubMats");
274 ierr = MatSetUp(A);
275 CHECK_ERROR("MatSetUp");
276 }
277 catch (...)
278 {
279 for (Mat& m : mats)
280 if (m)
281 MatDestroy(&m);
282 throw;
283 }
284
285 // De-reference Mat objects
286 for (Mat& m : mats)
287 if (m)
288 MatDestroy(&m);
289
290 return A;
291}
292
297 const std::vector<
298 std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
299
302 const std::vector<
303 std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
304
305// -- Vectors ----------------------------------------------------------------
306
319template <std::floating_point T>
321 Vec b, const Form<PetscScalar, T>& L,
322 std::span<const PetscScalar> constants,
323 const std::map<std::pair<IntegralType, int>,
324 std::pair<std::span<const PetscScalar>, int>>& coeffs)
325{
326 Vec b_local;
327 PetscErrorCode ierr = VecGhostGetLocalForm(b, &b_local);
328 CHECK_ERROR("VecGhostGetLocalForm");
329 PetscInt n = 0;
330 ierr = VecGetSize(b_local, &n);
331 CHECK_ERROR("VecGetSize");
332 PetscScalar* array = nullptr;
333 ierr = VecGetArray(b_local, &array);
334 CHECK_ERROR("VecGetArray");
335 std::span<PetscScalar> _b(array, n);
336 fem::assemble_vector(_b, L, constants, coeffs);
337 ierr = VecRestoreArray(b_local, &array);
338 CHECK_ERROR("VecRestoreArray");
339 ierr = VecGhostRestoreLocalForm(b, &b_local);
340 CHECK_ERROR("VecGhostRestoreLocalForm");
341}
342
353template <std::floating_point T>
355{
356 Vec b_local;
357 PetscErrorCode ierr = VecGhostGetLocalForm(b, &b_local);
358 CHECK_ERROR("VecGhostGetLocalForm");
359 PetscInt n = 0;
360 ierr = VecGetSize(b_local, &n);
361 CHECK_ERROR("VecGetSize");
362 PetscScalar* array = nullptr;
363 ierr = VecGetArray(b_local, &array);
364 CHECK_ERROR("VecGetArray");
365 std::span<PetscScalar> _b(array, n);
367 ierr = VecRestoreArray(b_local, &array);
368 CHECK_ERROR("VecRestoreArray");
369 ierr = VecGhostRestoreLocalForm(b, &b_local);
370 CHECK_ERROR("VecGhostRestoreLocalForm");
371}
372
373// FIXME: clarify zeroing of vector
374
404template <std::floating_point T>
406 Vec b,
407 std::vector<
408 std::optional<std::reference_wrapper<const Form<PetscScalar, T>>>>
409 a,
410 const std::vector<std::span<const PetscScalar>>& constants,
411 const std::vector<std::map<std::pair<IntegralType, int>,
412 std::pair<std::span<const PetscScalar>, int>>>&
413 coeffs,
414 const std::vector<
415 std::vector<std::reference_wrapper<const DirichletBC<PetscScalar, T>>>>&
416 bcs1,
417 const std::vector<Vec>& x0, PetscScalar alpha)
418{
419 if (!x0.empty() and x0.size() != a.size())
420 throw std::runtime_error("Mismatch between x0 and a in apply_lifting.");
421
422 Vec b_local;
423 PetscErrorCode ierr = VecGhostGetLocalForm(b, &b_local);
424 CHECK_ERROR("VecGhostGetLocalForm");
425 PetscInt n = 0;
426 ierr = VecGetSize(b_local, &n);
427 CHECK_ERROR("VecGetSize");
428 PetscScalar* array = nullptr;
429 ierr = VecGetArray(b_local, &array);
430 CHECK_ERROR("VecGetArray");
431 std::span<PetscScalar> _b(array, n);
432
433 if (x0.empty())
434 fem::apply_lifting(_b, a, constants, coeffs, bcs1, {}, alpha);
435 else
436 {
437 std::vector<std::span<const PetscScalar>> x0_ref;
438 std::vector<Vec> x0_local(a.size());
439 std::vector<const PetscScalar*> x0_array(a.size());
440 for (std::size_t i = 0; i < a.size(); ++i)
441 {
442 assert(x0[i]);
443 ierr = VecGhostGetLocalForm(x0[i], &x0_local[i]);
444 CHECK_ERROR("VecGhostGetLocalForm");
445 PetscInt n0 = 0;
446 ierr = VecGetSize(x0_local[i], &n0);
447 CHECK_ERROR("VecGetSize");
448 ierr = VecGetArrayRead(x0_local[i], &x0_array[i]);
449 CHECK_ERROR("VecGetArrayRead");
450 x0_ref.emplace_back(x0_array[i], n0);
451 }
452
453 fem::apply_lifting(_b, a, constants, coeffs, bcs1, x0_ref, alpha);
454
455 for (std::size_t i = 0; i < x0_local.size(); ++i)
456 {
457 ierr = VecRestoreArrayRead(x0_local[i], &x0_array[i]);
458 CHECK_ERROR("VecRestoreArrayRead");
459 ierr = VecGhostRestoreLocalForm(x0[i], &x0_local[i]);
460 CHECK_ERROR("VecGhostRestoreLocalForm");
461 }
462 }
463
464 ierr = VecRestoreArray(b_local, &array);
465 CHECK_ERROR("VecRestoreArray");
466 ierr = VecGhostRestoreLocalForm(b, &b_local);
467 CHECK_ERROR("VecGhostRestoreLocalForm");
468}
469
470// FIXME: clarify zeroing of vector
471
497template <std::floating_point T>
499 Vec b,
500 const std::vector<
501 std::optional<std::reference_wrapper<const Form<PetscScalar, T>>>>& a,
502 const std::vector<
503 std::vector<std::reference_wrapper<const DirichletBC<PetscScalar, T>>>>&
504 bcs1,
505 const std::vector<Vec>& x0, PetscScalar alpha)
506{
507 if (!x0.empty() and x0.size() != a.size())
508 throw std::runtime_error("Mismatch between x0 and a in apply_lifting.");
509
510 Vec b_local;
511 PetscErrorCode ierr = VecGhostGetLocalForm(b, &b_local);
512 CHECK_ERROR("VecGhostGetLocalForm");
513 PetscInt n = 0;
514 ierr = VecGetSize(b_local, &n);
515 CHECK_ERROR("VecGetSize");
516 PetscScalar* array = nullptr;
517 ierr = VecGetArray(b_local, &array);
518 CHECK_ERROR("VecGetArray");
519 std::span<PetscScalar> _b(array, n);
520
521 if (x0.empty())
522 fem::apply_lifting(_b, a, bcs1, {}, alpha);
523 else
524 {
525 std::vector<std::span<const PetscScalar>> x0_ref;
526 std::vector<Vec> x0_local(a.size());
527 std::vector<const PetscScalar*> x0_array(a.size());
528 for (std::size_t i = 0; i < a.size(); ++i)
529 {
530 assert(x0[i]);
531 ierr = VecGhostGetLocalForm(x0[i], &x0_local[i]);
532 CHECK_ERROR("VecGhostGetLocalForm");
533 PetscInt n0 = 0;
534 ierr = VecGetSize(x0_local[i], &n0);
535 CHECK_ERROR("VecGetSize");
536 ierr = VecGetArrayRead(x0_local[i], &x0_array[i]);
537 CHECK_ERROR("VecGetArrayRead");
538 x0_ref.emplace_back(x0_array[i], n0);
539 }
540
541 fem::apply_lifting(_b, a, bcs1, x0_ref, alpha);
542
543 for (std::size_t i = 0; i < x0_local.size(); ++i)
544 {
545 ierr = VecRestoreArrayRead(x0_local[i], &x0_array[i]);
546 CHECK_ERROR("VecRestoreArrayRead");
547 ierr = VecGhostRestoreLocalForm(x0[i], &x0_local[i]);
548 CHECK_ERROR("VecGhostRestoreLocalForm");
549 }
550 }
551
552 ierr = VecRestoreArray(b_local, &array);
553 CHECK_ERROR("VecRestoreArray");
554 ierr = VecGhostRestoreLocalForm(b, &b_local);
555 CHECK_ERROR("VecGhostRestoreLocalForm");
556}
557
558// -- Setting bcs ------------------------------------------------------------
559
560// FIXME: Move these function elsewhere?
561
574template <std::floating_point T>
575void set_bc(Vec b,
576 const std::vector<
577 std::reference_wrapper<const DirichletBC<PetscScalar, T>>>& bcs,
578 std::optional<const Vec> x0, PetscScalar alpha = 1)
579{
580 PetscInt n = 0;
581 PetscErrorCode ierr = VecGetLocalSize(b, &n);
582 CHECK_ERROR("VecGetLocalSize");
583 PetscScalar* array = nullptr;
584 ierr = VecGetArray(b, &array);
585 CHECK_ERROR("VecGetArray");
586 std::span<PetscScalar> _b(array, n);
587 if (x0.has_value())
588 {
589 Vec x0_local;
590 ierr = VecGhostGetLocalForm(x0.value(), &x0_local);
591 CHECK_ERROR("VecGhostGetLocalForm");
592 PetscInt n0 = 0;
593 ierr = VecGetSize(x0_local, &n0);
594 CHECK_ERROR("VecGetSize");
595 const PetscScalar* x0_array = nullptr;
596 ierr = VecGetArrayRead(x0_local, &x0_array);
597 CHECK_ERROR("VecGetArrayRead");
598 std::span<const PetscScalar> _x0(x0_array, n0);
599 for (auto& bc : bcs)
600 bc.get().set(_b, _x0, alpha);
601 ierr = VecRestoreArrayRead(x0_local, &x0_array);
602 CHECK_ERROR("VecRestoreArrayRead");
603 ierr = VecGhostRestoreLocalForm(x0.value(), &x0_local);
604 CHECK_ERROR("VecGhostRestoreLocalForm");
605 }
606 else
607 {
608 for (auto& bc : bcs)
609 bc.get().set(_b, std::nullopt, alpha);
610 }
611 ierr = VecRestoreArray(b, &array);
612 CHECK_ERROR("VecRestoreArray");
613}
614
615#undef CHECK_ERROR
616
617} // namespace petsc
618} // namespace dolfinx::fem
619
620#endif
Functions supporting assembly of finite element fem::Form and fem::Expression.
Definition IndexMap.h:95
std::int32_t size_local() const noexcept
Number of indices owned by this process.
Definition IndexMap.cpp:945
Definition DirichletBC.h:258
A representation of finite element variational forms.
Definition Form.h:118
std::shared_ptr< const mesh::Mesh< geometry_type > > mesh() const
Common mesh for the form (the 'integration domain').
Definition Form.h:367
Definition SparsityPattern.h:26
void finalize()
Finalize sparsity pattern and communicate off-process entries.
Definition SparsityPattern.cpp:264
Functions supporting finite element method operations.
Miscellaneous classes, functions and types.
Definition dolfinx_common.h:8
std::tuple< std::int64_t, std::vector< std::int32_t >, std::vector< std::vector< std::int64_t > >, std::vector< std::vector< int > > > stack_index_maps(const std::vector< std::pair< std::reference_wrapper< const IndexMap >, int > > &maps)
Compute layout data and ghost indices for a stacked (concatenated) index map, i.e....
Definition IndexMap.cpp:659
Helper functions for assembly into PETSc data structures.
Definition petsc.h:43
Mat create_matrix(const Form< PetscScalar, T > &a, std::optional< std::string > type=std::nullopt)
Create a matrix.
Definition petsc.h:60
Mat create_matrix_block(const std::vector< std::vector< const Form< PetscScalar, T > * > > &a, std::optional< std::string > type=std::nullopt)
Initialise a monolithic matrix for an array of bilinear forms.
Definition petsc.h:79
Mat create_matrix_nest(const std::vector< std::vector< const Form< PetscScalar, T > * > > &a, std::optional< std::vector< std::vector< std::optional< std::string > > > > types)
Create nested (MatNest) matrix.
Definition petsc.h:228
void assemble_vector(Vec b, const Form< PetscScalar, T > &L, std::span< const PetscScalar > constants, const std::map< std::pair< IntegralType, int >, std::pair< std::span< const PetscScalar >, int > > &coeffs)
Assemble linear form into an already allocated PETSc vector.
Definition petsc.h:320
void set_bc(Vec b, const std::vector< std::reference_wrapper< const DirichletBC< PetscScalar, T > > > &bcs, std::optional< const Vec > x0, PetscScalar alpha=1)
Definition petsc.h:575
Vec create_vector_block(const std::vector< std::pair< std::reference_wrapper< const common::IndexMap >, int > > &maps)
Initialise monolithic vector. Vector is not zeroed.
Definition petsc.cpp:21
void apply_lifting(Vec b, std::vector< std::optional< std::reference_wrapper< const Form< PetscScalar, T > > > > a, const std::vector< std::span< const PetscScalar > > &constants, const std::vector< std::map< std::pair< IntegralType, int >, std::pair< std::span< const PetscScalar >, int > > > &coeffs, const std::vector< std::vector< std::reference_wrapper< const DirichletBC< PetscScalar, T > > > > &bcs1, const std::vector< Vec > &x0, PetscScalar alpha)
Modify RHS vector to account for Dirichlet boundary conditions.
Definition petsc.h:405
Vec create_vector_nest(const std::vector< std::pair< std::reference_wrapper< const common::IndexMap >, int > > &maps)
Create nested (VecNest) vector. Vector is not zeroed.
Definition petsc.cpp:66
Finite element method functionality.
Definition assemble_expression_impl.h:23
std::vector< std::vector< std::array< std::shared_ptr< const FunctionSpace< U > >, 2 > > > extract_function_spaces(const std::vector< std::vector< const Form< T, U > * > > &a)
Extract test (0) and trial (1) function spaces pairs for each bilinear form for a rectangular array o...
Definition utils.h:173
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)
Modify the right-hand side vector to account for constraints (Dirichlet boundary condition constraint...
Definition assembler.h:341
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:233
la::SparsityPattern create_sparsity_pattern(const Form< T, U > &a)
Create a sparsity pattern for a given form.
Definition utils.h:198
std::array< std::vector< std::shared_ptr< const FunctionSpace< T > > >, 2 > common_function_spaces(const std::vector< std::vector< std::array< std::shared_ptr< const FunctionSpace< T > >, 2 > > > &V)
Extract FunctionSpaces for (0) rows blocks and (1) columns blocks from a rectangular array of (test,...
Definition FunctionSpace.h:433
Mat create_matrix(MPI_Comm comm, const SparsityPattern &sp, std::optional< std::string_view > type=std::nullopt)
Definition petsc.cpp:260
Mesh data structures and algorithms on meshes.
Definition DofMap.h:32