Coverage for python/src/dolfinx_mpc/utils/test.py: 94%
142 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 09:58 +0000
1# Copyright (C) 2021-2022 Jørgen Schartum Dokken and Connor D. Pierce
2#
3# This file is part of DOLFINX_MPC
4#
5# SPDX-License-Identifier: MIT
6from __future__ import annotations
8__all__ = [
9 "gather_PETScVector",
10 "gather_PETScMatrix",
11 "compare_mpc_lhs",
12 "compare_mpc_rhs",
13 "gather_transformation_matrix",
14 "gather_constants",
15 "compare_CSR",
16]
18from typing import Any
20from mpi4py import MPI
21from petsc4py import PETSc
23import dolfinx.common
24import numpy as np
25import scipy.sparse
27import dolfinx_mpc
30def _gather_slaves_global(constraint):
31 """
32 Given a multi point constraint, return slaves for all processors with global dof numbering
33 """
34 imap = constraint.function_space.dofmap.index_map
35 num_local_slaves = constraint.num_local_slaves
36 block_size = constraint.function_space.dofmap.index_map_bs
37 _slaves = constraint.slaves
38 if num_local_slaves > 0:
39 slave_blocks = _slaves[:num_local_slaves] // block_size
40 slave_rems = _slaves[:num_local_slaves] % block_size
41 glob_slaves = np.asarray(imap.local_to_global(slave_blocks), dtype=np.int64) * block_size + slave_rems
42 else:
43 glob_slaves = np.array([], dtype=np.int64)
45 slaves = np.hstack(MPI.COMM_WORLD.allgather(glob_slaves))
46 return slaves
49def gather_constants(constraint, root=0):
50 """
51 Given a multi-point constraint, gather the constraint offsets :math:`g` of all
52 processes into a single global vector on `root`, indexed by global dof.
53 """
54 V = constraint.V
55 imap = V.dofmap.index_map
56 block_size = V.dofmap.index_map_bs
57 constants = constraint.constants
58 l_range = np.asarray(imap.local_range, dtype=np.int64) * block_size
59 ranges = MPI.COMM_WORLD.gather(l_range, root=root)
60 g_consts = MPI.COMM_WORLD.gather(constants[: l_range[1] - l_range[0]], root=root)
61 if MPI.COMM_WORLD.rank == root:
62 global_consts = np.zeros(imap.size_global * block_size, dtype=constants.dtype)
63 for r, vals in zip(ranges, g_consts):
64 global_consts[r[0] : r[1]] = vals
65 return global_consts
66 else:
67 return
70def gather_transformation_matrix(constraint, root=0):
71 """
72 Creates the transformation matrix K (dim x dim-len(slaves)) for a given MPC
73 and gathers it as a scipy CSR matrix on process 'root'.
75 Example:
77 For dim=3, where:
78 u_1 = alpha u_0 + beta u_2
80 Input:
81 slaves = [1]
82 masters = [0, 2]
83 coeffs = [alpha, beta]
84 offsets = [0, 1]
86 Output:
87 K = [[1,0], [alpha beta], [0,1]]
88 """
89 # Gather slaves from all procs
90 V = constraint.V
91 imap = constraint.function_space.dofmap.index_map
92 block_size = V.dofmap.index_map_bs
93 num_local_slaves = constraint.num_local_slaves
94 # Gather all global_slaves
95 slaves = constraint.slaves[:num_local_slaves]
96 if num_local_slaves > 0:
97 local_blocks = slaves // block_size
98 local_rems = slaves % block_size
99 glob_slaves = np.asarray(imap.local_to_global(local_blocks), dtype=np.int64) * block_size + local_rems
100 else:
101 glob_slaves = np.array([], dtype=np.int64)
102 all_slaves = np.hstack(MPI.COMM_WORLD.allgather(glob_slaves))
103 masters = constraint.masters.array
104 master_blocks = masters // block_size
105 master_rems = masters % block_size
106 coeffs = constraint.coefficients()[0]
107 offsets = constraint.masters.offsets
108 # Create sparse K matrix
109 K_val, rows, cols = [], [], []
111 # Add local contributions to K from local slaves
112 for slave, global_slave in zip(slaves, glob_slaves):
113 masters_index = (
114 np.asarray(
115 imap.local_to_global(master_blocks[offsets[slave] : offsets[slave + 1]]),
116 dtype=np.int64,
117 )
118 * block_size
119 + master_rems[offsets[slave] : offsets[slave + 1]]
120 )
121 coeffs_index = coeffs[offsets[slave] : offsets[slave + 1]]
122 # A slave with no masters is a pure offset u_s = g_s, so its row of K is
123 # zero and the value is carried by the constant vector instead
124 for master, coeff in zip(masters_index, coeffs_index):
125 count = sum(master > all_slaves)
126 K_val.append(coeff)
127 rows.append(global_slave)
128 cols.append(master - count)
130 # Add identity for all dofs on diagonal
131 l_range = V.dofmap.index_map.local_range
132 global_dofs = np.arange(l_range[0] * block_size, l_range[1] * block_size)
133 is_slave = np.isin(global_dofs, glob_slaves)
134 for i, dof in enumerate(global_dofs):
135 if not is_slave[i]:
136 K_val.append(1)
137 rows.append(dof)
138 cols.append(dof - sum(dof > all_slaves))
140 # Gather K to root
141 K_vals = MPI.COMM_WORLD.gather(np.asarray(K_val, dtype=coeffs.dtype), root=root)
142 rows_g = MPI.COMM_WORLD.gather(np.asarray(rows, dtype=np.int64), root=root)
143 cols_g = MPI.COMM_WORLD.gather(np.asarray(cols, dtype=np.int64), root=root)
145 if MPI.COMM_WORLD.rank == root:
146 K_sparse = scipy.sparse.coo_matrix((np.hstack(K_vals), (np.hstack(rows_g), np.hstack(cols_g)))).tocsr()
147 return K_sparse
150def petsc_to_local_CSR(A: PETSc.Mat, mpc: dolfinx_mpc.MultiPointConstraint): # type: ignore
151 """
152 Convert a PETSc matrix to a local CSR matrix (scipy) including ghost entries
153 """
154 global_indices = np.asarray(mpc.function_space.dofmap.index_map.global_indices(), dtype=PETSc.IntType) # type: ignore
155 sort_index = np.argsort(global_indices)
156 is_A = PETSc.IS().createGeneral(global_indices[sort_index]) # type: ignore
157 A_loc = A.createSubMatrices(is_A)[0]
158 ai, aj, av = A_loc.getValuesCSR()
159 A_csr = scipy.sparse.csr_matrix((av, aj, ai))
160 return A_csr[global_indices[:, None], global_indices]
163def gather_PETScMatrix(A: PETSc.Mat, root=0) -> scipy.sparse.csr_matrix: # type: ignore
164 """
165 Given a distributed PETSc matrix, gather in on process 'root' in
166 a scipy CSR matrix
167 """
168 ai, aj, av = A.getValuesCSR()
169 aj_all = MPI.COMM_WORLD.gather(aj, root=root) # type: ignore
170 av_all = MPI.COMM_WORLD.gather(av, root=root) # type: ignore
171 ai_all = MPI.COMM_WORLD.gather(ai, root=root) # type: ignore
172 if MPI.COMM_WORLD.rank == root:
173 ai_cum = [0]
174 for ai in ai_all: # type: ignore
175 offsets = ai[1:] + ai_cum[-1]
176 ai_cum.extend(offsets)
177 return scipy.sparse.csr_matrix((np.hstack(av_all), np.hstack(aj_all), ai_cum), shape=A.getSize()) # type: ignore
180def gather_PETScVector(vector: PETSc.Vec, root=0) -> np.ndarray: # type: ignore
181 """
182 Gather a PETScVector from different processors on
183 process 'root' as an numpy array
184 """
185 if vector.handle == 0:
186 raise RuntimeError("Vector has been destroyed prior to this call")
187 numpy_vec = np.zeros(vector.size, dtype=vector.array.dtype)
188 l_min = vector.owner_range[0]
189 l_max = vector.owner_range[1]
190 numpy_vec[l_min:l_max] += vector.array.astype(numpy_vec.dtype)
191 return np.asarray(sum(MPI.COMM_WORLD.allgather(numpy_vec)))
194def compare_CSR(A: scipy.sparse.csr_matrix, B: scipy.sparse.csr_matrix, atol=1e-10):
195 """Compare CSR matrices A and B"""
196 diff = np.abs(A - B)
197 assert diff.max() < atol
200def compare_mpc_lhs(
201 A_org: PETSc.Mat, # type: ignore
202 A_mpc: PETSc.Mat, # type: ignore
203 mpc: dolfinx_mpc.MultiPointConstraint,
204 root: int = 0,
205 atol: np.floating[Any] = 5e3 * np.finfo(dolfinx.default_scalar_type).resolution,
206):
207 """
208 Compare an unmodified matrix for the problem with the one assembled with a
209 multi point constraint.
211 The unmodified matrix is multiplied with K^T A K, where K is the global transformation matrix.
212 """
213 timer = dolfinx.common.Timer("~MPC: Compare matrices")
214 comm = mpc.V.mesh.comm
215 V = mpc.V
216 assert root < comm.size
217 is_complex = np.issubdtype(mpc.coefficients()[0].dtype, np.complexfloating) # type: ignore
218 scipy_dtype = np.complex128 if is_complex else np.float64
219 K = gather_transformation_matrix(mpc, root=root)
220 A_csr = gather_PETScMatrix(A_org, root=root)
221 # Get global slaves
222 glob_slaves = _gather_slaves_global(mpc)
223 A_mpc_csr = gather_PETScMatrix(A_mpc, root=root)
224 if MPI.COMM_WORLD.rank == root:
225 K = K.astype(scipy_dtype)
226 A_csr = A_csr.astype(scipy_dtype)
227 KTAK = np.conj(K.T) * A_csr * K
229 # Remove identity rows of MPC matrix
230 all_cols = np.arange(V.dofmap.index_map.size_global * V.dofmap.index_map_bs)
232 # Scipy >=1.17.0 requires special way of doing slicing to avoid memory overflow
233 # https://github.com/scipy/scipy/issues/24339
234 free_mask = np.flatnonzero(np.isin(all_cols, glob_slaves, invert=True).astype(np.int32))
235 mpc_without_slaves = A_mpc_csr[free_mask, :][:, free_mask]
237 # Compute difference
238 compare_CSR(KTAK, mpc_without_slaves, atol=atol)
240 timer.stop()
243def compare_mpc_rhs(
244 b_org: PETSc.Vec, # type: ignore
245 b: PETSc.Vec, # type: ignore
246 constraint: dolfinx_mpc.MultiPointConstraint,
247 root: int = 0,
248 A_org: PETSc.Mat = None, # type: ignore
249):
250 """
251 Compare an unconstrained RHS with an MPC rhs.
253 The reduced right hand side is :math:`K^T (b - A g)`, where :math:`g` is the constraint
254 offset. For a homogeneous constraint the :math:`A g` term vanishes and `A_org` is not
255 needed; for an affine constraint it must be supplied.
256 """
257 if constraint.has_inhomogeneity and A_org is None:
258 raise ValueError("A_org is required to compare the rhs of an inhomogeneous constraint")
259 glob_slaves = _gather_slaves_global(constraint)
260 b_org_np = gather_PETScVector(b_org, root=root)
261 b_np = gather_PETScVector(b, root=root)
262 K = gather_transformation_matrix(constraint, root=root)
263 constants = gather_constants(constraint, root=root)
264 A_csr = None if A_org is None else gather_PETScMatrix(A_org, root=root)
265 comm = constraint.V.mesh.comm
266 if comm.rank == root:
267 lifted = b_org_np if A_csr is None else b_org_np - A_csr @ constants
268 reduced_b = np.conj(K.T) @ lifted
269 all_cols = np.arange(constraint.V.dofmap.index_map.size_global * constraint.V.dofmap.index_map_bs)
270 cols_except_slaves = np.flatnonzero(np.isin(all_cols, glob_slaves, invert=True).astype(np.int32))
271 assert np.allclose(b_np[glob_slaves], 0)
272 assert np.allclose(b_np[cols_except_slaves], reduced_b)