Loading...
Searching...
No Matches
petsc.h
// Copyright (C) 2018-2021 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier: LGPL-3.0-or-later
#pragma once
#ifdef HAS_PETSC
#include "Form.h"
#include "assembler.h"
#include "utils.h"
#include <concepts>
#include <dolfinx/la/petsc.h>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <petscmat.h>
#include <petscvec.h>
#include <span>
#include <utility>
#include <vector>
namespace dolfinx::common
{
class IndexMap;
}
namespace dolfinx::fem
{
template <dolfinx::scalar T, std::floating_point U>
class DirichletBC;
namespace petsc
{
template <std::floating_point T>
Mat create_matrix(const Form<PetscScalar, T>& a,
std::optional<std::string> type = std::nullopt)
{
la::SparsityPattern pattern = fem::create_sparsity_pattern(a);
pattern.finalize();
return la::petsc::create_matrix(a.mesh()->comm(), pattern, type);
}
template <std::floating_point T>
Mat create_matrix_block(
const std::vector<std::vector<const Form<PetscScalar, T>*>>& a,
std::optional<std::string> type = std::nullopt)
{
// Extract and check row/column ranges
std::array<std::vector<std::shared_ptr<const FunctionSpace<T>>>, 2> V
= fem::common_function_spaces(extract_function_spaces(a));
std::array<std::vector<int>, 2> bs_dofs;
for (std::size_t i = 0; i < 2; ++i)
{
for (auto& _V : V[i])
bs_dofs[i].push_back(_V->dofmap()->bs());
}
// Build sparsity pattern for each block
std::shared_ptr<const mesh::Mesh<T>> mesh;
std::vector<std::vector<std::unique_ptr<la::SparsityPattern>>> patterns(
V[0].size());
for (std::size_t row = 0; row < V[0].size(); ++row)
{
for (std::size_t col = 0; col < V[1].size(); ++col)
{
if (const Form<PetscScalar, T>* form = a[row][col]; form)
{
patterns[row].push_back(std::make_unique<la::SparsityPattern>(
create_sparsity_pattern(*form)));
if (!mesh)
mesh = form->mesh();
}
else
patterns[row].push_back(nullptr);
}
}
if (!mesh)
throw std::runtime_error("Could not find a Mesh.");
// Compute offsets for the fields
std::array<std::vector<std::pair<
std::reference_wrapper<const common::IndexMap>, int>>,
2>
maps;
for (std::size_t d = 0; d < 2; ++d)
{
for (auto& space : V[d])
{
maps[d].emplace_back(*space->dofmap()->index_map,
space->dofmap()->index_map_bs());
}
}
// Create merged sparsity pattern
std::vector<std::vector<const la::SparsityPattern*>> p(V[0].size());
for (std::size_t row = 0; row < V[0].size(); ++row)
for (std::size_t col = 0; col < V[1].size(); ++col)
p[row].push_back(patterns[row][col].get());
la::SparsityPattern pattern(mesh->comm(), p, maps, bs_dofs);
pattern.finalize();
// FIXME: Add option to pass customised local-to-global map to PETSc
// Mat constructor
// Initialise matrix
Mat A = la::petsc::create_matrix(mesh->comm(), pattern, type);
// Create row and column local-to-global maps (field0, field1, field2,
// etc), i.e. ghosts of field0 appear before owned indices of field1
std::array<std::vector<PetscInt>, 2> _maps;
for (int d = 0; d < 2; ++d)
{
// TODO: Index map concatenation has already been computed inside
// the SparsityPattern constructor, but we also need it here to
// build the PETSc local-to-global map. Compute outside and pass
// into SparsityPattern constructor.
// TODO: avoid concatenating the same maps twice in case that V[0]
// == V[1].
const std::vector<
std::pair<std::reference_wrapper<const common::IndexMap>, int>>& map
= maps[d];
std::vector<PetscInt>& _map = _maps[d];
// Concatenate the block index map in the row and column directions
const auto [rank_offset, local_offset, ghosts, _]
= common::stack_index_maps(map);
for (std::size_t f = 0; f < map.size(); ++f)
{
auto offset = local_offset[f];
const common::IndexMap& imap = map[f].first.get();
int bs = map[f].second;
for (std::int32_t i = 0; i < bs * imap.size_local(); ++i)
_map.push_back(i + rank_offset + offset);
for (std::int32_t i = 0; i < bs * imap.num_ghosts(); ++i)
_map.push_back(ghosts[f][i]);
}
}
// Create PETSc local-to-global map/index sets and attach to matrix
ISLocalToGlobalMapping petsc_local_to_global0;
ISLocalToGlobalMappingCreate(MPI_COMM_SELF, 1, _maps[0].size(),
_maps[0].data(), PETSC_COPY_VALUES,
&petsc_local_to_global0);
if (V[0] == V[1])
{
MatSetLocalToGlobalMapping(A, petsc_local_to_global0,
petsc_local_to_global0);
ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
}
else
{
ISLocalToGlobalMapping petsc_local_to_global1;
ISLocalToGlobalMappingCreate(MPI_COMM_SELF, 1, _maps[1].size(),
_maps[1].data(), PETSC_COPY_VALUES,
&petsc_local_to_global1);
MatSetLocalToGlobalMapping(A, petsc_local_to_global0,
petsc_local_to_global1);
ISLocalToGlobalMappingDestroy(&petsc_local_to_global0);
ISLocalToGlobalMappingDestroy(&petsc_local_to_global1);
}
return A;
}
template <std::floating_point T>
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)
{
// Extract and check row/column ranges
auto V = fem::common_function_spaces(extract_function_spaces(a));
// Loop over each form and create matrix
int rows = a.size();
int cols = a.front().size();
std::vector<Mat> mats(rows * cols, nullptr);
std::shared_ptr<const mesh::Mesh<T>> mesh;
for (int i = 0; i < rows; ++i)
{
for (int j = 0; j < cols; ++j)
{
if (const Form<PetscScalar, T>* form = a[i][j]; form)
{
if (types)
mats[i * cols + j] = create_matrix(*form, types->at(i).at(j));
else
mats[i * cols + j] = create_matrix(*form, std::nullopt);
mesh = form->mesh();
}
}
}
if (!mesh)
throw std::runtime_error("Could not find a Mesh.");
// Initialise block (MatNest) matrix
Mat A;
MatCreate(mesh->comm(), &A);
MatSetType(A, MATNEST);
MatNestSetSubMats(A, rows, nullptr, cols, nullptr, mats.data());
MatSetUp(A);
// De-reference Mat objects
for (std::size_t i = 0; i < mats.size(); ++i)
{
if (mats[i])
MatDestroy(&mats[i]);
}
return A;
}
Vec create_vector_block(
const std::vector<
std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
Vec create_vector_nest(
const std::vector<
std::pair<std::reference_wrapper<const common::IndexMap>, int>>& maps);
// -- Vectors ----------------------------------------------------------------
template <std::floating_point T>
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)
{
Vec b_local;
VecGhostGetLocalForm(b, &b_local);
PetscInt n = 0;
VecGetSize(b_local, &n);
PetscScalar* array = nullptr;
VecGetArray(b_local, &array);
std::span<PetscScalar> _b(array, n);
fem::assemble_vector(_b, L, constants, coeffs);
VecRestoreArray(b_local, &array);
VecGhostRestoreLocalForm(b, &b_local);
}
template <std::floating_point T>
void assemble_vector(Vec b, const Form<PetscScalar, T>& L)
{
Vec b_local;
VecGhostGetLocalForm(b, &b_local);
PetscInt n = 0;
VecGetSize(b_local, &n);
PetscScalar* array = nullptr;
VecGetArray(b_local, &array);
std::span<PetscScalar> _b(array, n);
fem::assemble_vector(_b, L);
VecRestoreArray(b_local, &array);
VecGhostRestoreLocalForm(b, &b_local);
}
// FIXME: clarify how x0 is used
// FIXME: if bcs entries are set
// FIXME: need to pass an array of Vec for x0?
// FIXME: clarify zeroing of vector
template <std::floating_point T>
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)
{
Vec b_local;
VecGhostGetLocalForm(b, &b_local);
PetscInt n = 0;
VecGetSize(b_local, &n);
PetscScalar* array = nullptr;
VecGetArray(b_local, &array);
std::span<PetscScalar> _b(array, n);
if (x0.empty())
fem::apply_lifting(_b, a, constants, coeffs, bcs1, {}, alpha);
else
{
std::vector<std::span<const PetscScalar>> x0_ref;
std::vector<Vec> x0_local(a.size());
std::vector<const PetscScalar*> x0_array(a.size());
for (std::size_t i = 0; i < a.size(); ++i)
{
assert(x0[i]);
VecGhostGetLocalForm(x0[i], &x0_local[i]);
PetscInt n = 0;
VecGetSize(x0_local[i], &n);
VecGetArrayRead(x0_local[i], &x0_array[i]);
x0_ref.emplace_back(x0_array[i], n);
}
std::vector x0_tmp(x0_ref.begin(), x0_ref.end());
fem::apply_lifting(_b, a, constants, coeffs, bcs1, x0_tmp, alpha);
for (std::size_t i = 0; i < x0_local.size(); ++i)
{
VecRestoreArrayRead(x0_local[i], &x0_array[i]);
VecGhostRestoreLocalForm(x0[i], &x0_local[i]);
}
}
VecRestoreArray(b_local, &array);
VecGhostRestoreLocalForm(b, &b_local);
}
// FIXME: clarify how x0 is used
// FIXME: if bcs entries are set
// FIXME: need to pass an array of Vec for x0?
// FIXME: clarify zeroing of vector
template <std::floating_point T>
void apply_lifting(
Vec b,
const std::vector<
std::optional<std::reference_wrapper<const Form<PetscScalar, double>>>>&
a,
const std::vector<std::vector<
std::reference_wrapper<const DirichletBC<PetscScalar, double>>>>& bcs1,
const std::vector<Vec>& x0, PetscScalar alpha)
{
Vec b_local;
VecGhostGetLocalForm(b, &b_local);
PetscInt n = 0;
VecGetSize(b_local, &n);
PetscScalar* array = nullptr;
VecGetArray(b_local, &array);
std::span<PetscScalar> _b(array, n);
if (x0.empty())
fem::apply_lifting(_b, a, bcs1, {}, alpha);
else
{
std::vector<std::span<const PetscScalar>> x0_ref;
std::vector<Vec> x0_local(a.size());
std::vector<const PetscScalar*> x0_array(a.size());
for (std::size_t i = 0; i < a.size(); ++i)
{
assert(x0[i]);
VecGhostGetLocalForm(x0[i], &x0_local[i]);
PetscInt n = 0;
VecGetSize(x0_local[i], &n);
VecGetArrayRead(x0_local[i], &x0_array[i]);
x0_ref.emplace_back(x0_array[i], n);
}
std::vector x0_tmp(x0_ref.begin(), x0_ref.end());
fem::apply_lifting(_b, a, bcs1, x0_tmp, alpha);
for (std::size_t i = 0; i < x0_local.size(); ++i)
{
VecRestoreArrayRead(x0_local[i], &x0_array[i]);
VecGhostRestoreLocalForm(x0[i], &x0_local[i]);
}
}
VecRestoreArray(b_local, &array);
VecGhostRestoreLocalForm(b, &b_local);
}
// -- Setting bcs ------------------------------------------------------------
// FIXME: Move these function elsewhere?
// FIXME: clarify x0
// FIXME: clarify what happens with ghosts
template <std::floating_point T>
void set_bc(
Vec b,
const std::vector<std::reference_wrapper<const DirichletBC<PetscScalar, T>>>
bcs,
const Vec x0, PetscScalar alpha = 1)
{
PetscInt n = 0;
VecGetLocalSize(b, &n);
PetscScalar* array = nullptr;
VecGetArray(b, &array);
std::span<PetscScalar> _b(array, n);
if (x0)
{
Vec x0_local;
VecGhostGetLocalForm(x0, &x0_local);
PetscInt n = 0;
VecGetSize(x0_local, &n);
const PetscScalar* array = nullptr;
VecGetArrayRead(x0_local, &array);
std::span<const PetscScalar> _x0(array, n);
for (auto& bc : bcs)
bc.set(_b, _x0, alpha);
VecRestoreArrayRead(x0_local, &array);
VecGhostRestoreLocalForm(x0, &x0_local);
}
else
{
for (auto& bc : bcs)
bc->set(_b, std::nullopt, alpha);
}
VecRestoreArray(b, &array);
}
} // namespace petsc
} // namespace dolfinx::fem
#endif
assembler.h
Functions supporting assembly of finite element fem::Form and fem::Expression.
dolfinx::common::IndexMap
Definition IndexMap.h:97
dolfinx::fem::DirichletBC
Definition DirichletBC.h:255
utils.h
Functions supporting finite element method operations.
dolfinx::MPI::size
int size(MPI_Comm comm)
Definition MPI.cpp:72
dolfinx::common
Miscellaneous classes, functions and types.
Definition dolfinx_common.h:8
dolfinx::fem
Finite element method functionality.
Definition assemble_expression_impl.h:23
dolfinx::fem::extract_function_spaces
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:172
dolfinx::fem::create_sparsity_pattern
la::SparsityPattern create_sparsity_pattern(const Form< T, U > &a)
Create a sparsity pattern for a given form.
Definition utils.h:197