1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
|
# Copyright (C) 2025 Garth N. Wells, Jack S. Hale
#
# This file is part of DOLFINx (https://www.fenicsproject.org)
#
# SPDX-License-Identifier: LGPL-3.0-or-later
"""Functions for working with PETSc linear algebra objects.
Note:
Due to subtle issues in the interaction between petsc4py memory
management and the Python garbage collector, it is recommended that
the PETSc method ``destroy()`` is called on returned PETSc objects
once the object is no longer required. Note that ``destroy()`` is
collective over the object's MPI communicator.
"""
import functools
import itertools
import typing
from collections.abc import Iterable, Sequence
from petsc4py import PETSc
import numpy as np
import numpy.typing as npt
import dolfinx
from dolfinx.common import IndexMap
from dolfinx.la import Vector
assert dolfinx.has_petsc4py
__all__ = ["assign", "create_vector", "create_vector_wrap"]
def _ghost_update(x: PETSc.Vec, insert_mode: PETSc.InsertMode, scatter_mode: PETSc.ScatterMode): # type: ignore[name-defined]
"""Helper function for ghost updating PETSc vectors"""
if x.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined]
for x_sub in x.getNestSubVecs():
x_sub.ghostUpdate(addv=insert_mode, mode=scatter_mode)
x_sub.destroy()
else:
x.ghostUpdate(addv=insert_mode, mode=scatter_mode)
def _zero_vector(x: PETSc.Vec): # type: ignore[name-defined]
"""Helper function for zeroing out PETSc vectors"""
if x.getType() == PETSc.Vec.Type.NEST: # type: ignore[attr-defined]
for x_sub in x.getNestSubVecs():
with x_sub.localForm() as x_sub_local:
x_sub_local.set(0.0)
x_sub.destroy()
else:
with x.localForm() as x_local:
x_local.set(0.0)
def create_vector_wrap(x: Vector) -> PETSc.Vec: # type: ignore[name-defined]
"""Wrap a distributed DOLFINx vector as a PETSc vector.
Args:
x: The vector to wrap as a PETSc vector.
Returns:
A PETSc vector that shares data with ``x``.
"""
index_map = x.index_map
ghosts = index_map.ghosts.astype(PETSc.IntType) # type: ignore[attr-defined]
bs = x.block_size
size = (index_map.size_local * bs, index_map.size_global * bs)
return PETSc.Vec().createGhostWithArray( # type: ignore[attr-defined]
ghosts, x.array, size=size, bsize=bs, comm=index_map.comm
)
def create_vector(
maps: typing.Sequence[tuple[IndexMap, int]], kind: str | None = None
) -> PETSc.Vec: # type: ignore[name-defined]
"""Create a PETSc vector from a sequence of maps and blocksizes.
Three cases are supported:
1. If ``maps=[(im_0, bs_0), ..., (im_n, bs_n)]`` is a sequence of
indexmaps and blocksizes and ``kind`` is ``None``or is
``PETSc.Vec.Type.MPI``, a ghosted PETSc vector whith block structure
described by ``(im_i, bs_i)`` is created.
The created vector ``b`` is initialized such that on each MPI
process ``b = [b_0, b_1, ..., b_n, b_0g, b_1g, ..., b_ng]``, where
``b_i`` are the entries associated with the 'owned' degrees-of-
freedom for ``(im_i, bs_i)`` and ``b_ig`` are the 'unowned' (ghost)
entries.
If more than one tuple is supplied, the returned vector has an
attribute ``_blocks`` that holds the local offsets into ``b`` for
the (i) owned and (ii) ghost entries for each ``V[i]``. It can be
accessed by ``b.getAttr("_blocks")``. The offsets can be used to get
views into ``b`` for blocks, e.g.::
>>> offsets0, offsets1, = b.getAttr("_blocks")
>>> offsets0
(0, 12, 28)
>>> offsets1
(28, 32, 35)
>>> b0_owned = b.array[offsets0[0]:offsets0[1]]
>>> b0_ghost = b.array[offsets1[0]:offsets1[1]]
>>> b1_owned = b.array[offsets0[1]:offsets0[2]]
>>> b1_ghost = b.array[offsets1[1]:offsets1[2]]
2. If ``V=[(im_0, bs_0), ..., (im_n, bs_n)]`` is a sequence of function
space and ``kind`` is ``PETSc.Vec.Type.NEST``, a PETSc nested vector
(a 'nest' of ghosted PETSc vectors) is created.
Args:
map: Sequence of tuples of ``IndexMap`` and the associated block
size.
kind: PETSc vector type (``VecType``) to create.
Returns:
A PETSc vector with the prescribed layout. The vector is not
initialised to zero.
"""
if len(maps) == 1:
# Single space case
index_map, bs = maps[0]
ghosts = index_map.ghosts.astype(PETSc.IntType) # type: ignore[attr-defined]
size = (index_map.size_local * bs, index_map.size_global * bs)
b = PETSc.Vec().createGhost(ghosts, size=size, bsize=bs, comm=index_map.comm) # type: ignore
if kind == PETSc.Vec.Type.MPI: # type: ignore[attr-defined]
_assign_block_data(maps, b)
return b
if kind is None or kind == PETSc.Vec.Type.MPI: # type: ignore[attr-defined]
b = dolfinx.cpp.fem.petsc.create_vector_block(maps)
_assign_block_data(maps, b)
return b
elif kind == PETSc.Vec.Type.NEST: # type: ignore[attr-defined]
return dolfinx.cpp.fem.petsc.create_vector_nest(maps)
else:
raise NotImplementedError(
"Vector type must be specified for blocked/nested assembly."
f"Vector type '{kind}' not supported."
"Did you mean 'nest' or 'mpi'?"
)
@functools.singledispatch
def assign(
x0: npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]],
x1: PETSc.Vec, # type: ignore[name-defined]
):
"""Assign ``x0`` values to a PETSc vector ``x1``.
Values in ``x0``, which is possibly a stacked collection of arrays,
are assigned ``x1``. When ``x0`` holds a sequence of ``n``` arrays
and ``x1`` has type ``NEST``, the assignment is::
[x0[0]]
x1 = [x0[1]]
[.....]
[x0[n-1]]
When ``x0`` holds a sequence of ``n`` arrays and ``x1`` **does
not** have type ``NEST``, the assignment is::
[x0_owned[0]]
x1 = [.....]
[x0_owned[n-1]]
[x0_ghost[0]]
[.....]
[x0_ghost[n-1]]
Args:
x0: An array or list of arrays that will be assigned to ``x1``.
x1: Vector to assign values to.
"""
if x1.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined]
x1_nest = x1.getNestSubVecs()
for _x0, _x1 in zip(x0, x1_nest):
with _x1.localForm() as x:
x.array_w[:] = _x0
else:
with x1.localForm() as _x:
if isinstance(x0, Sequence):
start = 0
for _x0 in x0:
end = start + _x0.shape[0]
_x.array_w[start:end] = _x0
start = end
else:
_x.array_w[:] = x0
@assign.register
def _(
x0: PETSc.Vec, # type: ignore[name-defined]
x1: npt.NDArray[np.inexact] | Sequence[npt.NDArray[np.inexact]],
):
"""Assign PETSc vector ``x0`` values to (blocked) array(s) ``x1``.
This function performs the reverse of the assignment performed by
the version of :func:`.assign(x0: (npt.NDArray[np.inexact] |
list[npt.NDArray[np.inexact]]), x1: PETSc.Vec)`.
Args:
x0: Vector that will have its values assigned to ``x1``.
x1: An array or list of arrays to assign to.
"""
if x0.getType() == PETSc.Vec.Type().NEST: # type: ignore[attr-defined]
x0_nest = x0.getNestSubVecs()
for _x0, _x1 in zip(x0_nest, x1):
with _x0.localForm() as x:
_x1[:] = x.array_r[:] # type: ignore[index]
else:
with x0.localForm() as _x0:
if isinstance(x1, Sequence):
start = 0
for _x1 in x1:
end = start + _x1.shape[0]
_x1[:] = _x0.array_r[start:end]
start = end
else:
x1[:] = _x0.array_r[:]
def _assign_block_data(maps: Iterable[tuple[IndexMap, int]], vec: PETSc.Vec): # type: ignore[name-defined]
"""Assign block data to a PETSc vector.
Args:
maps: Iterable of tuples each containing an ``IndexMap`` and the
associated block size ``bs``.
vec: PETSc vector to assign block data to.
"""
# Early exit if the vector already has block data or is a nest vector
if vec.getAttr("_blocks") is not None or vec.getType() == "nest":
return
maps = [(index_map, bs) for index_map, bs in maps]
off_owned = tuple(
itertools.accumulate(maps, lambda off, m: off + m[0].size_local * m[1], initial=0)
)
off_ghost = tuple(
itertools.accumulate(
maps, lambda off, m: off + m[0].num_ghosts * m[1], initial=off_owned[-1]
)
)
vec.setAttr("_blocks", (off_owned, off_ghost))
|
