Coverage for python/src/dolfinx_mpc/utils/mpc_utils.py: 87%
271 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 mpi4py import MPI
9from petsc4py import PETSc
11import dolfinx.common as _common
12import dolfinx.cpp as _cpp
13import dolfinx.fem as _fem
14import dolfinx.geometry as _geometry
15import dolfinx.la as _la
16import dolfinx.log as _log
17import dolfinx.mesh as _mesh
18import numpy as np
19import ufl
20from dolfinx import default_scalar_type as _dt
22import dolfinx_mpc.cpp
24__all__ = [
25 "rotation_matrix",
26 "facet_normal_approximation",
27 "log_info",
28 "rigid_motions_nullspace",
29 "determine_closest_block",
30 "create_normal_approximation",
31 "create_point_to_point_constraint",
32]
35def rotation_matrix(axis, angle):
36 # See https://en.wikipedia.org/wiki/Rotation_matrix,
37 # Subsection: Rotation_matrix_from_axis_and_angle.
38 if np.isclose(np.inner(axis, axis), 1):
39 n_axis = axis
40 else:
41 # Normalize axis
42 n_axis = axis / np.sqrt(np.inner(axis, axis))
44 # Define cross product matrix of axis
45 axis_x = np.array([[0, -n_axis[2], n_axis[1]], [n_axis[2], 0, -n_axis[0]], [-n_axis[1], n_axis[0], 0]])
46 identity = np.cos(angle) * np.eye(3)
47 outer = (1 - np.cos(angle)) * np.outer(n_axis, n_axis)
48 return np.sin(angle) * axis_x + identity + outer
51def facet_normal_approximation(
52 V,
53 mt: _mesh.MeshTags,
54 mt_id: int,
55 tangent=False,
56 jit_options: dict = {},
57 form_compiler_options: dict = {},
58):
59 """
60 Approximate the facet normal by projecting it into the function space for a set of facets
62 Args:
63 V: The function space to project into
64 mt: The `dolfinx.mesh.MeshTagsMetaClass` containing facet markers
65 mt_id: The id for the facets in `mt` we want to represent the normal at
66 tangent: To approximate the tangent to the facet set this flag to `True`
67 jit_options: Parameters used in CFFI JIT compilation of C code generated by FFCx.
68 See https://github.com/FEniCS/dolfinx/blob/main/python/dolfinx/jit.py#L22-L37
69 for all available parameters. Takes priority over all other parameter values.
70 form_compiler_options: Parameters used in FFCx compilation of this form. Run `ffcx - -help` at
71 the commandline to see all available options. Takes priority over all
72 other parameter values, except for `scalar_type` which is determined by
73 DOLFINx.
74 """
75 timer = _common.Timer("~MPC: Facet normal projection")
76 comm = V.mesh.comm
77 n = ufl.FacetNormal(V.mesh)
78 nh = _fem.Function(V)
79 u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
80 ds = ufl.ds(domain=V.mesh, subdomain_data=mt, subdomain_id=mt_id)
81 if tangent:
82 if V.mesh.geometry.dim == 1:
83 raise ValueError("Tangent not defined for 1D problem")
84 elif V.mesh.geometry.dim == 2:
85 a = ufl.inner(u, v) * ds
86 L = ufl.inner(ufl.as_vector([-n[1], n[0]]), v) * ds
87 else:
89 def tangential_proj(u, n):
90 """
91 See for instance:
92 https://link.springer.com/content/pdf/10.1023/A:1022235512626.pdf
93 """
94 return (ufl.Identity(u.ufl_shape[0]) - ufl.outer(n, n)) * u
96 c = _fem.Constant(V.mesh, [1, 1, 1])
97 a = ufl.inner(u, v) * ds
98 L = ufl.inner(tangential_proj(c, n), v) * ds
99 else:
100 a = ufl.inner(u, v) * ds
101 L = ufl.inner(n, v) * ds
103 # Find all dofs that are not boundary dofs
104 imap = V.dofmap.index_map
105 all_blocks = np.arange(imap.size_local, dtype=np.int32)
106 top_blocks = _fem.locate_dofs_topological(V, V.mesh.topology.dim - 1, mt.find(mt_id))
107 deac_blocks = all_blocks[np.isin(all_blocks, top_blocks, invert=True)]
109 # Note there should be a better way to do this
110 # Create sparsity pattern only for constraint + bc
111 bilinear_form = _fem.form(a, jit_options=jit_options, form_compiler_options=form_compiler_options)
112 pattern = _fem.create_sparsity_pattern(bilinear_form)
113 pattern.insert_diagonal(deac_blocks)
114 pattern.finalize()
115 u_0 = _fem.Function(V)
116 u_0.x.petsc_vec.set(0)
118 bc_deac = _fem.dirichletbc(u_0, deac_blocks)
119 A = _cpp.la.petsc.create_matrix(comm, pattern._cpp_object, None)
120 A.zeroEntries()
122 # Assemble the matrix with all entries
123 form_coeffs = _cpp.fem.pack_coefficients(bilinear_form._cpp_object)
124 form_consts = _cpp.fem.pack_constants(bilinear_form._cpp_object)
125 markers = _fem.petsc._matrix_bc_markers(bilinear_form, [bc_deac])
126 _cpp.fem.petsc.assemble_matrix(A, bilinear_form._cpp_object, form_consts, form_coeffs, *markers, False)
127 if bilinear_form.function_spaces[0] is bilinear_form.function_spaces[1]:
128 A.assemblyBegin(PETSc.Mat.AssemblyType.FLUSH) # type: ignore
129 A.assemblyEnd(PETSc.Mat.AssemblyType.FLUSH) # type: ignore
130 dofs = np.empty(0, dtype=np.int32)
131 if bilinear_form.function_spaces[0].contains(bc_deac.function_space):
132 dofs, owned = bc_deac.dof_indices()
133 rows = np.array(dofs, dtype=np.int32)
134 _cpp.fem.petsc.set_diagonal(A, rows, 1.0, PETSc.InsertMode.INSERT_VALUES) # type: ignore
135 A.assemble()
136 linear_form = _fem.form(L, jit_options=jit_options, form_compiler_options=form_compiler_options)
137 b = _fem.petsc.assemble_vector(linear_form)
139 _fem.petsc.apply_lifting(b, [bilinear_form], [[bc_deac]])
140 b.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE) # type: ignore
141 _fem.petsc.set_bc(b, [bc_deac])
143 # Solve Linear problem
144 solver = PETSc.KSP().create(V.mesh.comm) # type: ignore
145 solver.setType("cg")
146 solver.rtol = 1e-8
147 solver.setOperators(A)
148 solver.solve(b, nh.x.petsc_vec)
149 nh.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD) # type: ignore
150 timer.stop()
151 solver.destroy()
152 b.destroy()
153 return nh
156def log_info(message):
157 """
158 Wrapper for logging a simple string on the zeroth communicator
159 Reverting the log level
160 """
161 old_level = _log.get_log_level()
162 if MPI.COMM_WORLD.rank == 0:
163 _log.set_log_level(_log.LogLevel.INFO)
164 _log.log(_log.LogLevel.INFO, message)
165 _log.set_log_level(old_level)
168def rigid_motions_nullspace(V: _fem.FunctionSpace):
169 """
170 Function to build nullspace for 2D/3D elasticity.
172 Args:
173 V: The function space
174 """
175 _x = _fem.Function(V)
176 # Get geometric dim
177 gdim = V.mesh.geometry.dim
178 assert gdim == 2 or gdim == 3
180 # Set dimension of nullspace
181 dim = 3 if gdim == 2 else 6
183 # Create list of vectors for null space
184 nullspace_basis = [
185 _la.vector(V.dofmap.index_map, bs=V.dofmap.index_map_bs, dtype=PETSc.ScalarType) # type: ignore
186 for i in range(dim)
187 ]
189 basis = [b.array for b in nullspace_basis]
190 dofs = [V.sub(i).dofmap.list.reshape(-1) for i in range(gdim)]
192 # Build translational null space basis
193 for i in range(gdim):
194 basis[i][dofs[i]] = 1.0
195 # Build rotational null space basis
196 x = V.tabulate_dof_coordinates()
197 dofs_block = V.dofmap.list.reshape(-1)
198 x0, x1, x2 = x[dofs_block, 0], x[dofs_block, 1], x[dofs_block, 2]
199 if gdim == 2:
200 basis[2][dofs[0]] = -x1
201 basis[2][dofs[1]] = x0
202 elif gdim == 3:
203 basis[3][dofs[0]] = -x1
204 basis[3][dofs[1]] = x0
206 basis[4][dofs[0]] = x2
207 basis[4][dofs[2]] = -x0
208 basis[5][dofs[2]] = x1
209 basis[5][dofs[1]] = -x2
210 for b in nullspace_basis:
211 b.scatter_forward()
213 _la.orthonormalize(nullspace_basis)
214 assert _la.is_orthonormal(nullspace_basis, float(np.finfo(_x.x.array.dtype).eps))
215 local_size = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
216 basis_petsc = [
217 PETSc.Vec().createWithArray(x[:local_size], bsize=gdim, comm=V.mesh.comm) # type: ignore
218 for x in basis
219 ]
220 return PETSc.NullSpace().create(comm=V.mesh.comm, vectors=basis_petsc) # type: ignore
223def determine_closest_block(V, point):
224 """
225 Determine the closest dofs (in a single block) to a point and the distance
226 """
227 # Create boundingboxtree of cells connected to boundary facets
228 tdim = V.mesh.topology.dim
229 boundary_facets = _mesh.exterior_facet_indices(V.mesh.topology)
230 V.mesh.topology.create_connectivity(tdim - 1, tdim)
231 f_to_c = V.mesh.topology.connectivity(tdim - 1, tdim)
232 boundary_cells = []
233 for facet in boundary_facets:
234 boundary_cells.extend(f_to_c.links(facet))
235 cell_imap = V.mesh.topology.index_map(tdim)
236 boundary_cells = np.array(np.unique(boundary_cells), dtype=np.int32)
237 boundary_cells = boundary_cells[boundary_cells < cell_imap.size_local]
238 bb_tree = _geometry.bb_tree(V.mesh, tdim, entities=boundary_cells, padding=0.0)
239 midpoint_tree = _geometry.create_midpoint_tree(V.mesh, tdim, boundary_cells)
241 # Find facet closest
242 point = np.reshape(point, (1, 3)).astype(V.mesh.geometry.x.dtype)
243 closest_cell = _geometry.compute_closest_entity(bb_tree, midpoint_tree, V.mesh, point)[0]
245 # Set distance high if cell is not owned
246 if cell_imap.size_local < closest_cell or closest_cell == -1:
247 R = 1e5
248 else:
249 # Get cell geometry
250 p = V.mesh.geometry.x
251 V.mesh.topology.create_connectivity(tdim, tdim)
252 entities = _mesh.entities_to_geometry(V.mesh, tdim, np.array([closest_cell], dtype=np.int32), False)
253 R = np.linalg.norm(_geometry.compute_distance_gjk(point, p[entities[0]]))
255 # Find processor with cell closest to point
256 global_distances = MPI.COMM_WORLD.allgather(R)
257 owning_processor = np.argmin(global_distances)
259 dofmap = V.dofmap
260 imap = dofmap.index_map
261 ghost_owner = imap.owners
262 local_max = imap.size_local
263 # Determine which block of dofs is closest
264 min_distance = max(R, 1e5)
265 minimal_distance_block = None
266 min_dof_owner = owning_processor
267 if MPI.COMM_WORLD.rank == owning_processor:
268 x = V.tabulate_dof_coordinates()
269 cell_blocks = dofmap.cell_dofs(closest_cell)
270 for block in cell_blocks:
271 distance = np.linalg.norm(_geometry.compute_distance_gjk(point, x[block]))
272 if distance < min_distance:
273 # If cell owned by processor, but not the closest dof
274 if block < local_max:
275 min_dof_owner = MPI.COMM_WORLD.rank
276 else:
277 min_dof_owner = ghost_owner[block - local_max]
278 minimal_distance_block = block
279 min_distance = distance
280 min_dof_owner = MPI.COMM_WORLD.bcast(min_dof_owner, root=owning_processor)
281 # If dofs not owned by cell
282 if owning_processor != min_dof_owner:
283 owning_processor = min_dof_owner
285 if MPI.COMM_WORLD.rank == min_dof_owner:
286 # Re-search using the closest cell
287 x = V.tabulate_dof_coordinates()
288 cell_blocks = dofmap.cell_dofs(closest_cell)
289 for block in cell_blocks:
290 distance = np.linalg.norm(_geometry.compute_distance_gjk(point, x[block]))
291 if distance < min_distance:
292 # If cell owned by processor, but not the closest dof
293 if block < local_max:
294 min_dof_owner = MPI.COMM_WORLD.rank
295 else:
296 min_dof_owner = ghost_owner[block - local_max]
297 minimal_distance_block = block
298 min_distance = distance
299 assert min_dof_owner == owning_processor
300 return owning_processor, [minimal_distance_block]
301 else:
302 return owning_processor, []
305def create_point_to_point_constraint(V, slave_point, master_point, vector=None):
306 # Determine which processor owns the dof closest to the slave and master point
307 slave_proc, slave_block = determine_closest_block(V, slave_point)
308 master_proc, master_block = determine_closest_block(V, master_point)
309 is_master_proc = MPI.COMM_WORLD.rank == master_proc
310 is_slave_proc = MPI.COMM_WORLD.rank == slave_proc
312 block_size = V.dofmap.index_map_bs
313 imap = V.dofmap.index_map
314 # Output structures
315 slaves, masters, coeffs, owners, offsets = [], [], [], [], []
316 # Information required to handle vector as input
317 zero_indices, slave_index = None, None
318 if vector is not None:
319 zero_indices = np.argwhere(np.isclose(vector, 0)).T[0]
320 slave_index = np.argmax(np.abs(vector))
321 if is_slave_proc:
322 assert len(slave_block) == 1
323 slave_block_g = imap.local_to_global(np.asarray(slave_block, dtype=np.int32))[0]
324 if vector is None:
325 slaves = np.arange(
326 slave_block[0] * block_size,
327 slave_block[0] * block_size + block_size,
328 dtype=np.int32,
329 )
330 else:
331 assert len(vector) == block_size
332 # Check for input vector (Should be of same length as number of slaves)
333 # All entries should not be zero
334 assert not np.isin(slave_index, zero_indices)
335 # Check vector for zero contributions
336 slaves = np.array([slave_block[0] * block_size + slave_index], dtype=np.int32)
337 for i in range(block_size):
338 if i != slave_index and not np.isin(i, zero_indices):
339 masters.append(slave_block_g * block_size + i)
340 owners.append(slave_proc)
341 coeffs.append(-vector[i] / vector[slave_index])
343 global_masters = None
344 if is_master_proc:
345 assert len(master_block) == 1
346 master_block_g = imap.local_to_global(np.asarray(master_block, dtype=np.int32))[0]
347 masters_as_glob = np.arange(
348 master_block_g * block_size, master_block_g * block_size + block_size, dtype=np.int64
349 )
350 else:
351 masters_as_glob = np.array([], dtype=np.int64)
353 ghost_processors = []
354 shared_indices = dolfinx_mpc.cpp.mpc.compute_shared_indices(V._cpp_object)
356 if is_master_proc and is_slave_proc:
357 # If slaves and masters are on the same processor finalize local work
358 if vector is None:
359 masters = masters_as_glob
360 owners = np.full(len(masters), master_proc, dtype=np.int32)
361 coeffs = np.ones(len(masters), dtype=_dt)
362 offsets = np.arange(0, len(masters) + 1, dtype=np.int32)
363 else:
364 for i in range(len(masters_as_glob)):
365 if not np.isin(i, zero_indices):
366 masters.append(masters_as_glob[i])
367 owners.append(master_proc)
368 coeffs.append(vector[i] / vector[slave_index])
369 offsets = [0, len(masters)]
370 else:
371 # Send/Recv masters from other processor
372 if is_master_proc:
373 MPI.COMM_WORLD.send(masters_as_glob, dest=slave_proc, tag=10)
374 if is_slave_proc:
375 global_masters = MPI.COMM_WORLD.recv(source=master_proc, tag=10)
376 for i, master in enumerate(global_masters):
377 if not np.isin(i, zero_indices):
378 masters.append(master)
379 owners.append(master_proc)
380 if vector is None:
381 coeffs.append(1)
382 else:
383 coeffs.append(vector[i] / vector[slave_index])
384 if vector is None:
385 offsets = np.arange(0, len(slaves) + 1, dtype=np.int32)
386 else:
387 offsets = np.array([0, len(masters)], dtype=np.int32)
388 ghost_processors = shared_indices.links(slave_block[0])
390 # Broadcast processors containg slave
391 ghost_processors = MPI.COMM_WORLD.bcast(ghost_processors, root=slave_proc)
392 if is_slave_proc:
393 for proc in ghost_processors:
394 MPI.COMM_WORLD.send(slave_block_g * block_size + slaves % block_size, dest=proc, tag=20 + proc)
395 MPI.COMM_WORLD.send(coeffs, dest=proc, tag=30 + proc)
396 MPI.COMM_WORLD.send(owners, dest=proc, tag=40 + proc)
397 MPI.COMM_WORLD.send(masters, dest=proc, tag=50 + proc)
398 MPI.COMM_WORLD.send(offsets, dest=proc, tag=60 + proc)
400 # Receive data for ghost slaves
401 ghost_slaves, ghost_masters, ghost_coeffs, ghost_owners, ghost_offsets = [], [], [], [], []
402 if np.isin(MPI.COMM_WORLD.rank, ghost_processors):
403 # Convert recieved slaves to the corresponding ghost index
404 recv_slaves = MPI.COMM_WORLD.recv(source=slave_proc, tag=20 + MPI.COMM_WORLD.rank)
405 ghost_coeffs = MPI.COMM_WORLD.recv(source=slave_proc, tag=30 + MPI.COMM_WORLD.rank)
406 ghost_owners = MPI.COMM_WORLD.recv(source=slave_proc, tag=40 + MPI.COMM_WORLD.rank)
407 ghost_masters = MPI.COMM_WORLD.recv(source=slave_proc, tag=50 + MPI.COMM_WORLD.rank)
408 ghost_offsets = MPI.COMM_WORLD.recv(source=slave_proc, tag=60 + MPI.COMM_WORLD.rank)
410 # Unroll ghost blocks
411 ghosts = imap.ghosts
412 ghost_dofs = [g * block_size + i for g in ghosts for i in range(block_size)]
414 ghost_slaves = np.zeros(len(recv_slaves), dtype=np.int32)
415 local_size = imap.size_local
416 for i, slave in enumerate(recv_slaves):
417 idx = np.argwhere(ghost_dofs == slave)[0, 0]
418 ghost_slaves[i] = local_size * block_size + idx
419 slaves = np.asarray(np.append(slaves, ghost_slaves), dtype=np.int32)
420 masters = np.asarray(np.append(masters, ghost_masters), dtype=np.int64)
421 coeffs = np.asarray(np.append(coeffs, ghost_coeffs), dtype=_dt)
422 owners = np.asarray(np.append(owners, ghost_owners), dtype=np.int32)
423 offsets = np.asarray(np.append(offsets, ghost_offsets), dtype=np.int32)
424 return slaves, masters, coeffs, owners, offsets
427def create_normal_approximation(V: _fem.FunctionSpace, mt: _cpp.mesh.MeshTags_int32, value: int):
428 """
429 Creates a normal approximation for the dofs in the closure of the attached entities.
430 Where a dof is attached to multiple entities, an average is computed.
432 Args:
433 V: The function space
434 mt: The meshtag containing the indices
435 value: Value for the entities in the mesh tag to compute normal on
437 Returns:
438 nh: The normal vector
439 """
440 nh = _fem.Function(V)
441 n_cpp = dolfinx_mpc.cpp.mpc.create_normal_approximation(V._cpp_object, mt.dim, mt.find(value))
442 nh._cpp_object = n_cpp
443 return nh