Coverage for python/src/dolfinx_mpc/assemble_matrix.py: 98%
176 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 09:58 +0000
1# Copyright (C) 2020-2021 Jørgen S. Dokken
2#
3# This file is part of DOLFINX_MPC
4#
5# SPDX-License-Identifier: MIT
6from __future__ import annotations
8from collections.abc import Sequence
9from typing import Optional, Union
11from mpi4py import MPI
12from petsc4py import PETSc as _PETSc
14import dolfinx.cpp as _cpp
15import dolfinx.fem as _fem
16import dolfinx.fem.petsc # noqa: F401
17import dolfinx.la
18import numpy as np
19import numpy.typing as npt
20from dolfinx import default_scalar_type
22from dolfinx_mpc import cpp
24from ._kind import Kind, blocked_layout, deprecated, single_type
25from .dirichletbc import BCData
26from .multipointconstraint import MultiPointConstraint
29def _assemble_form(
30 A: _PETSc.Mat, # type: ignore
31 form: _fem.Form,
32 constraint: Sequence[MultiPointConstraint],
33 bc_data: BCData,
34 num_threads: Optional[int] = 1,
35):
36 """
37 Assemble one compiled form into a matrix.
39 Additive: `A` is not zeroed. Does not finalise or add any diagonal entry —
40 those belong to the system rather than to a single form, see
41 :func:`_finalize_matrix`.
42 """
43 # Markers come from the cache rather than being rebuilt here, so a caller
44 # that reassembles the same form pays for them once; see `BCData`.
45 dof_marker0, dof_marker1 = bc_data.markers(*form.function_spaces)
46 cpp.mpc.assemble_matrix(
47 A,
48 form._cpp_object,
49 constraint[0]._cpp_object,
50 constraint[1]._cpp_object,
51 dof_marker0,
52 dof_marker1,
53 num_threads,
54 )
57def _add_diagonals(
58 slave_blocks: Sequence,
59 bc_blocks: Sequence[tuple[_PETSc.Mat, npt.NDArray[np.int32]]],
60 diagval: _PETSc.ScalarType = 1, # type: ignore
61):
62 """
63 Add the diagonal entries of the slave and Dirichlet rows. Does not assemble.
65 A diagonal entry belongs to a block of the system, not to a form: a block
66 may carry slaves without having a diagonal bilinear form to assemble, and a
67 block appearing in several forms must still receive exactly one entry. So
68 this runs once, after every form has been assembled.
70 Args:
71 slave_blocks: `(sub-matrix, constraint)` pairs whose slave rows get `diagval`
72 bc_blocks: `(sub-matrix, rows)` pairs whose Dirichlet rows get `diagval`
73 diagval: Value to place on the diagonal
74 """
75 for A_sub, mpc in slave_blocks:
76 cpp.mpc.insert_diagonal_slaves(A_sub, mpc._cpp_object, diagval)
78 # Both diagonals are added rather than inserted, so no flush is needed to
79 # take the matrix out of add mode. Adding is equivalent here because
80 # assembly leaves a Dirichlet row empty: the element matrix has its
81 # constrained rows and columns zeroed, slaves are rejected at finalize if
82 # they carry a condition, and a master that carries one is eliminated into
83 # the constraint offset rather than kept in the master list.
84 for A_sub, rows in bc_blocks:
85 _cpp.fem.petsc.set_diagonal(A_sub, rows, default_scalar_type(diagval), _PETSc.InsertMode.ADD_VALUES) # type: ignore
88def _finalize_matrix(
89 A: _PETSc.Mat, # type: ignore
90 slave_blocks: Sequence,
91 bc_blocks: Sequence[tuple[_PETSc.Mat, npt.NDArray[np.int32]]],
92 diagval: _PETSc.ScalarType = 1, # type: ignore
93):
94 """
95 Add the diagonal entries, see :func:`_add_diagonals`, and assemble `A`.
97 Args:
98 A: The matrix, with every form already assembled into it
99 slave_blocks: `(sub-matrix, constraint)` pairs whose slave rows get `diagval`
100 bc_blocks: `(sub-matrix, rows)` pairs whose Dirichlet rows get `diagval`
101 diagval: Value to place on the diagonal
102 """
103 _add_diagonals(slave_blocks, bc_blocks, diagval)
104 A.assemble()
107def assemble_matrix(
108 form: Union[_fem.Form, Sequence[Sequence[Optional[_fem.Form]]]],
109 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],
110 bcs: Optional[Sequence[_fem.DirichletBC]] = None,
111 diagval: _PETSc.ScalarType = 1, # type: ignore
112 A: Optional[_PETSc.Mat] = None, # type: ignore
113 num_threads: Optional[int] = 1,
114 bc_data: Optional[BCData] = None,
115 kind: Kind = None,
116) -> _PETSc.Mat: # type: ignore
117 """
118 Assemble a compiled DOLFINx bilinear form, or an array of them, into a PETSc matrix with
119 corresponding multi point constraints and Dirichlet boundary conditions.
121 As in :func:`dolfinx.fem.petsc.assemble_matrix`, the kind of matrix is selected by `kind`, or by
122 the type of `A` if it is supplied.
124 Args:
125 form: The compiled bilinear variational form, or a rank 2 list of them with `None` for a
126 block without a form
127 constraint: For a single form, its multi point constraint, or for a rectangular form a
128 list of 2 constraints on axis 0 & 1. For an array of forms, the constraint of each
129 block, which is used for the rows and the columns.
130 bcs: Sequence of Dirichlet boundary conditions
131 diagval: Value to set on the diagonal of the matrix
132 A: PETSc matrix to assemble into. Assembly is additive, so `A` is not
133 zeroed; call `A.zeroEntries()` first to discard its contents. If not
134 supplied a new matrix is created, which is already zeroed.
135 num_threads: The number of threads to use for certain operations
136 bc_data: A :class:`BCData` cache. Built from `bcs` when not supplied;
137 pass one to share it with the other assemblies of the same system.
138 kind: The kind of matrix to create when `A` is not supplied, see :func:`create_matrix`.
139 Returns:
140 _PETSc.Mat: The matrix with the assembled bi-linear form #type: ignore
141 """
142 if bc_data is None:
143 bc_data = BCData(bcs)
145 if isinstance(form, Sequence):
146 if not isinstance(constraint, Sequence):
147 raise ValueError("An array of forms needs one multi point constraint per block")
148 if A is None:
149 A = create_matrix(form, constraint, kind)
150 if A.getType() == "nest":
151 _assemble_matrix_nest(A, form, constraint, diagval=diagval, num_threads=num_threads, bc_data=bc_data)
152 else:
153 _assemble_matrix_block(A, form, constraint, diagval=diagval, num_threads=num_threads, bc_data=bc_data)
154 return A
156 if not isinstance(constraint, Sequence):
157 assert form.function_spaces[0] == form.function_spaces[1]
158 constraint = (constraint, constraint)
160 # Generate matrix with MPC sparsity pattern. A freshly created matrix is
161 # already zeroed; an `A` supplied by the caller is added into.
162 if A is None:
163 A = create_matrix(form, constraint, kind)
165 V0, V1 = form.function_spaces
166 for mpc, V in {id(c): (c, W) for c, W in zip(constraint, (V0, V1))}.values():
167 if not mpc._cpp_object.has_cross_block_masters:
168 _raise_if_constrained_masters([mpc], [V], bc_data)
169 _assemble_form(A, form, constraint, bc_data, num_threads)
171 slave_blocks = [(A, constraint[0])] if constraint[0] is constraint[1] else []
172 bc_blocks = [(A, bc_data.rows(V0))] if V0 is V1 else []
173 _finalize_matrix(A, slave_blocks, bc_blocks, diagval)
175 return A
178def create_matrix(
179 a: Union[_fem.Form, Sequence[Sequence[Optional[_fem.Form]]]],
180 constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],
181 kind: Kind = None,
182) -> _PETSc.Mat: # type: ignore
183 """
184 Create a PETSc matrix with the sparsity pattern of a bilinear form, or an array of them, under
185 multi point constraints.
187 As in :func:`dolfinx.fem.petsc.create_matrix`, three cases are supported:
189 1. A single form gives a matrix of the PETSc type `kind`, the default if `None`.
190 2. An array of forms with `kind` ``"nest"``, or a nested sequence of PETSc matrix types of the
191 same shape as the array, gives a matrix of type ``nest`` whose blocks have those types.
192 3. An array of forms with any other `kind` gives a single, monolithic matrix of PETSc type
193 `kind`, the default if `None` or ``"mpi"``, arranged as
194 :math:`A = [a_{ij}]` with the dofs of each block ordered ``[owned, ghosts]`` and the blocks
195 one after another. The ghosts include the masters added by the constraint of the block.
196 A diagonal block that has no form still reserves the diagonal entry of each of its slaves.
198 Args:
199 a: The compiled bilinear form, or a rank 2 list of them with `None` for a block without a form
200 constraint: As in :func:`assemble_matrix`
201 kind: The kind of matrix, as above
203 Returns:
204 The matrix, to assemble into with :func:`assemble_matrix`.
205 """
206 if not isinstance(a, Sequence):
207 if not isinstance(constraint, Sequence):
208 constraint = (constraint, constraint)
209 for mpc in constraint:
210 mpc._raise_if_not_finalized()
211 return cpp.mpc.create_matrix(
212 a._cpp_object, constraint[0]._cpp_object, constraint[1]._cpp_object, single_type(kind)
213 )
215 if not isinstance(constraint, Sequence):
216 raise ValueError("An array of forms needs one multi point constraint per block")
217 layout, types = blocked_layout(kind)
218 if layout == "nest":
219 return _create_matrix_nest(a, constraint, types)
220 return _create_matrix_block(a, constraint, matrix_type=types) # type: ignore[arg-type]
223def create_sparsity_pattern(form: _fem.Form, mpc: Union[MultiPointConstraint, Sequence[MultiPointConstraint]]):
224 """
225 Create sparsity-pattern for MPC given a compiled DOLFINx form
227 Args:
228 form: The form
229 mpc: For square forms, the MPC. For rectangular forms a list of 2 MPCs on
230 axis 0 & 1, respectively
231 """
232 if isinstance(mpc, Sequence):
233 assert len(mpc) == 2
234 for mpc_ in mpc:
235 mpc_._raise_if_not_finalized() # type: ignore
236 return cpp.mpc.create_sparsity_pattern(form._cpp_object, mpc[0]._cpp_object, mpc[1]._cpp_object)
237 else:
238 mpc._raise_if_not_finalized() # type: ignore
239 return cpp.mpc.create_sparsity_pattern(
240 form._cpp_object,
241 mpc._cpp_object, # type: ignore
242 mpc._cpp_object, # type: ignore
243 ) # type: ignore
246def _create_matrix_nest(
247 a: Sequence[Sequence[Optional[_fem.Form]]],
248 constraints: Sequence[MultiPointConstraint],
249 types: Optional[Sequence[Sequence[Optional[str]]]] = None,
250):
251 """
252 Create a PETSc matrix of type "nest" with the blocks of the types in `types`, if given.
254 A block has a matrix if it has a form or is a diagonal block, which holds the diagonal of its
255 slaves. If a constraint has masters in another block, every block has one, since those masters
256 put entries in blocks without a form.
257 """
258 assert len(constraints) == len(a)
259 for mpc in constraints:
260 mpc._raise_if_not_finalized()
261 forms = [[None if a_ij is None else a_ij._cpp_object for a_ij in a_i] for a_i in a]
262 mpcs = [mpc._cpp_object for mpc in constraints]
263 _types = None if types is None else [list(t) for t in types]
264 return cpp.mpc.create_matrix_nest(forms, mpcs, mpcs, _types)
267def create_matrix_nest(a: Sequence[Sequence[_fem.Form | None]], constraints: Sequence[MultiPointConstraint]):
268 """
269 Create a PETSc matrix of type "nest" with appropriate sparsity pattern
270 given the provided multi points constraints
272 .. deprecated::
273 Use :func:`create_matrix` with ``kind="nest"``.
275 Args:
276 a: The compiled bilinear variational form provided in a rank 2 list
277 constraints: An ordered list of multi point constraints
278 """
279 deprecated("create_matrix_nest", "create_matrix(a, constraints, kind='nest')")
280 return _create_matrix_nest(a, constraints)
283def _block_spaces(
284 a: Sequence[Sequence[Optional[_fem.Form]]], constraints: Sequence[MultiPointConstraint]
285) -> list[_fem.FunctionSpace]:
286 """The space of each diagonal block, from the forms, or from its constraint if no form has it.
288 A block without a form still has a space: its constraint's, which may carry Dirichlet
289 conditions, and whose dofs may be masters of other blocks.
290 """
291 spaces: list[Optional[_fem.FunctionSpace]] = [None] * len(constraints)
292 for i, a_row in enumerate(a):
293 for j, a_ij in enumerate(a_row):
294 if a_ij is None:
295 continue
296 if spaces[i] is None:
297 spaces[i] = a_ij.function_spaces[0]
298 if j < len(constraints) and spaces[j] is None:
299 spaces[j] = a_ij.function_spaces[1]
300 return [mpc.input_space if V is None else V for V, mpc in zip(spaces, constraints)]
303def _raise_if_constrained_masters(
304 constraints: Sequence[MultiPointConstraint], spaces: Sequence[_fem.FunctionSpace], bc_data: BCData
305):
306 """Raise if a Dirichlet condition of the assembly constrains a master the constraints kept.
308 Assembly adds the entries of a slave's row and column to those of its masters after the
309 constrained rows and columns are zeroed, so a constrained master must have been eliminated, by
310 giving the condition to the constraint of the master's block before finalizing it. The verdict
311 is cached in `bc_data`.
313 Args:
314 constraints: The constraint of each block, in the order they were finalized in if a
315 constraint has masters in another block
316 spaces: The space of each block, on which the conditions are stated
317 bc_data: The Dirichlet conditions of the assembly
319 Raises:
320 ValueError: On every process, if a kept master is constrained.
322 Note:
323 Collective.
324 """
325 key = tuple(id(mpc._cpp_object) for mpc in constraints)
326 if key in bc_data._checked:
327 return
328 # The markers of each block on its extended space: the owned dofs are numbered as in the
329 # space of the conditions, the ghosts are the extended space's own
330 markers: list[Optional[npt.NDArray[np.bool_]]] = []
331 for mpc, V in zip(constraints, spaces):
332 # Decided by the spaces of the conditions, the same on every process, as the scatter is
333 # collective; the markers of a process without dofs are empty
334 if not any(V.contains(bc.function_space) for bc in bc_data._bcs):
335 markers.append(None)
336 continue
337 owned_markers = bc_data.markers(V, V)[0]
338 dofmap = mpc.function_space.dofmap
339 marker = dolfinx.la.vector(dofmap.index_map, dofmap.index_map_bs, dtype=np.float64)
340 num_owned = dofmap.index_map.size_local * dofmap.index_map_bs
341 marker.array[:num_owned] = owned_markers[:num_owned]
342 marker.scatter_forward()
343 markers.append(marker.array > 0)
345 constrained = False
346 if any(marker is not None for marker in markers):
347 for k, mpc in enumerate(constraints):
348 masters = mpc._cpp_object.masters.array
349 blocks = np.asarray(mpc._cpp_object.master_blocks)
350 if blocks.size == 0:
351 blocks = np.full(masters.size, k, dtype=np.int32)
352 else:
353 # A master in the constraint's own block is in block `k` of the system, also when
354 # the constraint was finalized on its own; any other block is the master's
355 blocks = np.where(blocks == mpc._cpp_object.block, k, blocks)
356 for j, marker in enumerate(markers):
357 if marker is not None and marker[masters[blocks == j]].any():
358 constrained = True
359 if constraints[0].function_space.mesh.comm.allreduce(constrained, op=MPI.LOR):
360 raise ValueError(
361 "A Dirichlet condition of the assembly constrains a master of a multi point constraint. "
362 "Give the condition to the constraint of the master's space as well, "
363 "MultiPointConstraint(V, bcs=...), before finalizing, so that the master is eliminated."
364 )
365 bc_data._checked.add(key)
368def _assemble_blocks(
369 blocks: Sequence[Sequence[Optional[_PETSc.Mat]]], # type: ignore
370 a: Sequence[Sequence[Optional[_fem.Form]]],
371 constraints: Sequence[MultiPointConstraint],
372 bc_data: BCData,
373 diagval: _PETSc.ScalarType, # type: ignore
374 num_threads: Optional[int],
375):
376 """
377 Assemble an array of forms into the matrices of the blocks of a system, then add the diagonal
378 of the slave and Dirichlet rows of each diagonal block.
380 The entries of a master go to the block of the master, which may have no form.
382 Raises:
383 ValueError: On every process, if a Dirichlet condition constrains a master that the
384 constraints kept.
385 """
386 spaces = _block_spaces(a, constraints)
387 _raise_if_constrained_masters(constraints, spaces, bc_data)
388 for i, a_row in enumerate(a):
389 for j, a_ij in enumerate(a_row):
390 if a_ij is None:
391 continue
392 dof_marker0, dof_marker1 = bc_data.markers(*a_ij.function_spaces)
393 cpp.mpc.assemble_matrix_blocks(
394 blocks,
395 i,
396 j,
397 a_ij._cpp_object,
398 constraints[i]._cpp_object,
399 constraints[j]._cpp_object,
400 dof_marker0,
401 dof_marker1,
402 num_threads,
403 )
405 # The diagonal is a property of a block, so it is added once per diagonal block after
406 # every form has been assembled, including for a block that has no form of its own, whose
407 # pattern reserves its whole diagonal
408 for i, mpc in enumerate(constraints):
409 A_ii = blocks[i][i]
410 if A_ii is None:
411 continue
412 rows = bc_data.rows(spaces[i])
413 _add_diagonals([(A_ii, mpc)], [(A_ii, rows)] if rows.size > 0 else [], diagval)
416def _assemble_matrix_nest(
417 A: _PETSc.Mat, # type: ignore
418 a: Sequence[Sequence[Optional[_fem.Form]]],
419 constraints: Sequence[MultiPointConstraint],
420 bcs: Sequence[_fem.DirichletBC] = [],
421 diagval: _PETSc.ScalarType = 1, # type: ignore
422 num_threads: Optional[int] = 1,
423 bc_data: Optional[BCData] = None,
424):
425 """Assemble an array of forms into a PETSc matrix of type "nest"."""
426 if bc_data is None:
427 bc_data = BCData([bc for bc in bcs])
428 nr, nc = A.getNestSize()
429 blocks: list[list[Optional[_PETSc.Mat]]] = [] # type: ignore
430 for k in range(nr):
431 row = []
432 for col in range(nc):
433 A_kl = A.getNestSubMatrix(k, col)
434 row.append(None if A_kl.handle == 0 else A_kl)
435 blocks.append(row)
436 _assemble_blocks(blocks, a, constraints, bc_data, diagval, num_threads)
437 A.assemble()
440def assemble_matrix_nest(
441 A: _PETSc.Mat, # type: ignore
442 a: Sequence[Sequence[Optional[_fem.Form]]],
443 constraints: Sequence[MultiPointConstraint],
444 bcs: Sequence[_fem.DirichletBC] = [],
445 diagval: _PETSc.ScalarType = 1, # type: ignore
446 num_threads: Optional[int] = 1,
447 bc_data: Optional[BCData] = None,
448):
449 """
450 Assemble a compiled DOLFINx bilinear form into a PETSc matrix of type
451 "nest" with corresponding multi point constraints and Dirichlet boundary
452 conditions.
454 .. deprecated::
455 Use :func:`assemble_matrix`, which selects the layout from `A`.
457 Args:
458 A: PETSc matrix to assemble into. Assembly is additive, so `A` is not
459 zeroed; call `A.zeroEntries()` first to discard its contents.
460 a: The compiled bilinear variational form provided in a rank 2 list
461 constraints: An ordered list of multi point constraints
462 bcs: Sequence of Dirichlet boundary conditions
463 diagval: Value to set on the diagonal of the matrix (Default 1)
464 num_threads: The number of threads to use for certain operations
465 bc_data: A :class:`BCData` cache. Built from `bcs` when not supplied;
466 pass one to share it with the other assemblies of the same system.
467 """
468 deprecated("assemble_matrix_nest", "assemble_matrix(a, constraints, A=A)")
469 _assemble_matrix_nest(A, a, constraints, bcs, diagval, num_threads, bc_data)
472def _block_index_sets(constraints: Sequence[MultiPointConstraint]):
473 """Index sets selecting each block of a monolithic matrix, in the local numbering of the matrix."""
474 return _cpp.la.petsc.create_index_sets(
475 [
476 (mpc.function_space.dofmap.index_map._cpp_object, mpc.function_space.dofmap.index_map_bs)
477 for mpc in constraints
478 ]
479 )
482def _create_matrix_block(
483 a: Sequence[Sequence[Optional[_fem.Form]]],
484 constraints: Sequence[MultiPointConstraint],
485 constraints1: Optional[Sequence[MultiPointConstraint]] = None,
486 matrix_type: Optional[str] = None,
487):
488 """Create a monolithic PETSc matrix, see :func:`create_matrix`. The columns use `constraints1` if given."""
489 cols = constraints if constraints1 is None else constraints1
490 for mpc in (*constraints, *cols):
491 mpc._raise_if_not_finalized()
492 forms = [[None if a_ij is None else a_ij._cpp_object for a_ij in a_i] for a_i in a]
493 return cpp.mpc.create_matrix_block(
494 forms,
495 [mpc._cpp_object for mpc in constraints],
496 [mpc._cpp_object for mpc in cols],
497 matrix_type,
498 )
501def _assemble_matrix_block(
502 A: _PETSc.Mat, # type: ignore
503 a: Sequence[Sequence[Optional[_fem.Form]]],
504 constraints: Sequence[MultiPointConstraint],
505 bcs: Sequence[_fem.DirichletBC] = [],
506 diagval: _PETSc.ScalarType = 1, # type: ignore
507 num_threads: Optional[int] = 1,
508 bc_data: Optional[BCData] = None,
509):
510 """
511 Assemble an array of forms into a monolithic PETSc matrix made by :func:`create_matrix`.
513 Raises:
514 ValueError: On every process, if a Dirichlet condition constrains a master that the
515 constraints kept.
516 """
517 if bc_data is None:
518 bc_data = BCData(list(bcs))
519 is_ = _block_index_sets(constraints)
520 # The local sub-matrix of every block, as a master may put entries in any of them
521 blocks = [[A.getLocalSubMatrix(is_k, is_l) for is_l in is_] for is_k in is_]
522 try:
523 _assemble_blocks(blocks, a, constraints, bc_data, diagval, num_threads)
524 finally:
525 for is_k, row in zip(is_, blocks):
526 for is_l, A_kl in zip(is_, row):
527 A.restoreLocalSubMatrix(is_k, is_l, A_kl)
528 A.assemble()