Coverage for python/src/dolfinx_mpc/dictcondition.py: 93%
136 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 09:58 +0000
1# Copyright (C) 2020-2026 Jørgen S. Dokken
2#
3# This file is part of DOLFINX_MPC
4#
5# SPDX-License-Identifier: MIT
6"""A multi-point constraint given as a dictionary from the coordinates of slaves to those of their masters."""
8from __future__ import annotations
10import hashlib
11import typing
13from mpi4py import MPI
15import dolfinx
16import dolfinx.fem as fem
17import numpy as np
18import numpy.typing as npt
19from dolfinx import default_scalar_type
21from .container import _deprecated, _tolerance
24def close_to(
25 point: np.typing.NDArray[np.float64 | np.float32],
26 atol: float | None = None,
27 *,
28 distance_tol: float | None = None,
29):
30 """
31 Convenience function for locating a point [x,y,z]
32 within an array x [[x0,...,xN],[y0,...,yN], [z0,...,zN]].
34 Args:
35 point: The point should be padded to 3D
36 atol: Deprecated, use `distance_tol`.
37 distance_tol: Every component of a located point is within ``distance_tol + 1e-5 |point|``
38 of the point's. Defaults to `500` machine epsilon of the type of `point`, or of
39 ``dolfinx.default_real_type`` if it is not a floating point type.
40 """
41 if atol is not None:
42 _deprecated("atol", "`distance_tol`")
43 distance_tol = atol if distance_tol is None else distance_tol
44 point = np.asarray(point)
45 real_type = point.dtype if np.issubdtype(point.dtype, np.floating) else dolfinx.default_real_type
46 tol = _tolerance(distance_tol, real_type)
47 return lambda x: np.isclose(x, point, atol=tol).all(axis=0)
50def _close_pairs(
51 points: npt.NDArray[np.float64], candidate_points: npt.NDArray[np.float64], atol: float, rtol: float = 1e-5
52) -> tuple[npt.NDArray[np.int64], npt.NDArray[np.int64]]:
53 """The pairs `(i, k)` for which row `k` of `candidate_points` is close to point `i`, as in
54 :func:`close_to`: every component within ``atol + rtol * |point|``.
56 The rows are sorted by their projection on a generic direction. The rows that can be close to a
57 point are then a window of the sorted rows, found for all points at once by a binary search,
58 and only those are compared with the point.
59 """
60 direction = np.array([1.0, 1.0 / np.pi, 1.0 / np.e])
61 projection = candidate_points @ direction
62 order = np.argsort(projection, kind="stable")
63 sorted_projection = projection[order]
64 tol = atol + rtol * np.abs(points)
65 # A row with every component within tol of the point projects within tol @ |direction| of it,
66 # widened by the rounding of the projections
67 centre = points @ direction
68 width = tol @ np.abs(direction) + 1e-12 * (1.0 + np.abs(centre))
69 lo = np.searchsorted(sorted_projection, centre - width, side="left")
70 hi = np.searchsorted(sorted_projection, centre + width, side="right")
71 counts = hi - lo
72 # Expand the windows into one candidate (point, row) per slot, the slots of all points one after
73 # the other. For counts = [2, 0, 3] and lo = [5, 9, 1]: point = [0, 0, 2, 2, 2], first = [0, 2, 2],
74 # within = [0, 1, 0, 1, 2] and position = [5, 6, 1, 2, 3].
75 # The point of each slot
76 point = np.repeat(np.arange(len(points)), counts)
77 # The first slot of each point
78 first = np.cumsum(counts) - counts
79 # The place of each slot in the window of its point: 0, 1, ...
80 within = np.arange(counts.sum()) - np.repeat(first, counts)
81 # The position of each slot in the sorted rows, and the row there
82 position = np.repeat(lo, counts) + within
83 row = order[position]
84 close = (np.abs(candidate_points[row] - points[point]) <= tol[point]).all(axis=1)
85 return point[close], row[close]
88def _digest(
89 slave_master_dict: dict[bytes, dict[bytes, typing.Any]],
90 subspace_slave: int | None,
91 subspace_master: int | None,
92) -> bytes | None:
93 """A fingerprint of the input, equal on processes given the same input, or None if the keys are not bytes."""
94 h = hashlib.sha256(repr((subspace_slave, subspace_master)).encode())
95 try:
96 for slave, masters in slave_master_dict.items():
97 h.update(len(slave).to_bytes(8, "little") + slave)
98 h.update(len(masters).to_bytes(8, "little"))
99 for master, coeff in masters.items():
100 h.update(len(master).to_bytes(8, "little") + master)
101 h.update(np.complex128(coeff).tobytes())
102 except (TypeError, ValueError):
103 return None
104 return h.digest()
107def _check_input(
108 comm: MPI.Comm,
109 slave_master_dict: dict[bytes, dict[bytes, typing.Any]],
110 subspace_slave: int | None,
111 subspace_master: int | None,
112 real_type: np.dtype,
113 scalar_type: np.dtype,
114) -> None:
115 """Raise on every process unless the input is the same on every process, and valid.
117 The slaves and masters are matched across processes by their position in the dictionary, and
118 every process reads the coefficients from its own copy, so all must have the same.
119 """
120 digest = _digest(slave_master_dict, subspace_slave, subspace_master)
121 if not comm.allreduce(digest is not None, op=MPI.LAND):
122 raise TypeError("The coordinates of the slaves and masters must be given as bytes")
123 if not comm.allreduce(digest == comm.bcast(digest, root=0), op=MPI.LAND):
124 raise ValueError("Every process must pass the same dictionary, in the same order, with the same sub spaces")
125 # The input is the same everywhere, so the checks below give the same verdict on every process
126 itemsize = real_type.itemsize
127 for key in (k for slave, masters in slave_master_dict.items() for k in (slave, *masters)):
128 if len(key) % itemsize != 0 or not 0 < len(key) // itemsize <= 3:
129 raise ValueError(f"A coordinate must be 1 to 3 values of the mesh's coordinate type, {real_type}")
130 if not np.issubdtype(scalar_type, np.complexfloating):
131 for masters in slave_master_dict.values():
132 if any(np.imag(c) != 0 for c in masters.values()):
133 raise ValueError(f"A complex coefficient cannot be used for a constraint of scalar type {scalar_type}")
136@typing.no_type_check
137def create_dictionary_constraint(
138 V: fem.functionspace,
139 slave_master_dict: dict[bytes, dict[bytes, float]],
140 subspace_slave: int | None = None,
141 subspace_master: int | None = None,
142 dtype: npt.DTypeLike | None = None,
143 distance_tol: float | None = None,
144):
145 """
146 Returns a multi point constraint for a given function space
147 and dictionary constraint.
149 Every process must pass the same dictionary, in the same order, and the same sub spaces: the
150 slaves and masters are matched across processes by their position in it. This is checked.
152 Args:
153 V: The function space
154 slave_master_dict: The dictionary
155 subspace_slave: If using mixed or vector space, and only want to use dofs from
156 a sub space as slave add index here.
157 subspace_master: Subspace index for mixed or vector spaces
158 dtype: The scalar type of the coefficients. Defaults to the default scalar type.
159 distance_tol: Every component of the coordinate of a dof is within
160 ``distance_tol + 1e-5 |x|`` of that of its key. Defaults to `500` machine epsilon of the
161 coordinate type of the mesh.
163 Returns:
164 The slaves on this process, owned first and then ghosts, in the order of the dictionary,
165 and their masters, coefficients, the owners of the masters and the offsets.
167 Raises:
168 ValueError: On every process, if the processes are not given the same input, or if no
169 process has a degree of freedom at a master of a slave.
171 Examples:
172 If the dof `D` located at `[d0,d1]` should be constrained to the dofs `E` and
173 F at `[e0,e1]` and `[f0,f1]` as :math:`D = \\alpha E + \\beta F`
174 the dictionary should be:
176 .. highlight:: python
177 .. code-block:: python
179 {np.array([d0, d1], dtype=mesh.geometry.x.dtype).tobytes():
180 {numpy.array([e0, e1], dtype=mesh.geometry.x.dtype).tobytes(): alpha,
181 numpy.array([f0, f1], dtype=mesh.geometry.x.dtype).tobytes(): beta}}
183 Note:
184 Collective.
185 """
186 comm = V.mesh.comm
187 real_type = np.dtype(V.mesh.geometry.x.dtype)
188 scalar_type = np.dtype(default_scalar_type if dtype is None else dtype)
189 _check_input(comm, slave_master_dict, subspace_slave, subspace_master, real_type, scalar_type)
191 bs = V.dofmap.index_map_bs
192 index_map = V.dofmap.index_map
193 local_size = index_map.size_local * bs
194 atol = _tolerance(distance_tol, real_type)
196 def dof_table(subspace):
197 """The coordinates of the dofs of `V`, or of its sub space, and their indices in `V`, local
198 to the process, ghosts included. A blocked space has a row per component, all at the
199 coordinate of the block, so a point finds all of them."""
200 if subspace is None:
201 coordinates, table_bs = V.tabulate_dof_coordinates(), bs
202 dofs = np.arange(len(coordinates) * bs, dtype=np.int64)
203 else:
204 # One map per dofmap, that is per cell type of the mesh
205 V_sub, sub_to_V = V.sub(subspace).collapse()
206 if len(sub_to_V) != 1:
207 raise NotImplementedError("A dictionary constraint on a mesh of several cell types")
208 coordinates, table_bs = V_sub.tabulate_dof_coordinates(), V_sub.dofmap.index_map_bs
209 dofs = np.asarray(sub_to_V[0], dtype=np.int64)
210 return np.repeat(coordinates.astype(np.float64), table_bs, axis=0), dofs
212 def points(keys):
213 """The coordinates in `keys`, padded to 3D."""
214 out = np.zeros((len(keys), 3), dtype=np.float64)
215 for k, key in enumerate(keys):
216 coordinates = np.frombuffer(key, dtype=real_type)
217 out[k, : len(coordinates)] = coordinates
218 return out
220 # Master j of slave i is entry starts[i] + j of the flat layout, the same on every process
221 slave_keys = list(slave_master_dict.keys())
222 master_keys = [master for key in slave_keys for master in slave_master_dict[key]]
223 num_slaves = len(slave_keys)
224 starts = np.zeros(num_slaves + 1, dtype=np.int64)
225 starts[1:] = np.cumsum([len(slave_master_dict[key]) for key in slave_keys])
226 num_entries = int(starts[-1])
228 # One reduction carries everything: the global index and the owner of each master, filled by
229 # the process owning it, whether each slave is on some process, and the errors of locating
230 reduced = np.full(2 * num_entries + num_slaves + 2, -1, dtype=np.int64)
231 masters_all = reduced[:num_entries]
232 owners_all = reduced[num_entries : 2 * num_entries]
233 found = reduced[2 * num_entries : 2 * num_entries + num_slaves]
234 errors = reduced[2 * num_entries + num_slaves :]
236 # The slaves, owned or ghost, at their points
237 slave_coordinates, slave_table = dof_table(subspace_slave)
238 i, k = _close_pairs(points(slave_keys), slave_coordinates, atol)
239 count = np.bincount(i, minlength=num_slaves)
240 single = count[i] == 1
241 slave_dofs = np.full(num_slaves, -1, dtype=np.int64)
242 slave_dofs[i[single]] = slave_table[k[single]]
243 found[:] = slave_dofs >= 0
244 errors[0] = (count > 1).any()
246 # The masters this process owns, as only the owner of a master fills it in
247 master_coordinates, master_table = dof_table(subspace_master)
248 owned = master_table < local_size
249 e, k = _close_pairs(points(master_keys), master_coordinates[owned], atol)
250 count = np.bincount(e, minlength=num_entries)
251 single = count[e] == 1
252 blocks, components = np.divmod(master_table[owned][k[single]], bs)
253 masters_all[e[single]] = index_map.local_to_global(blocks.astype(np.int32)) * bs + components
254 owners_all[e[single]] = comm.rank
255 errors[1] = (count > 1).any()
256 comm.Allreduce(MPI.IN_PLACE, reduced, op=MPI.MAX)
258 # The reduced data is the same on every process, and so is every verdict on it
259 if errors[0]:
260 raise RuntimeError("Multiple slaves found at same point. You should use sub-space locators.")
261 if errors[1]:
262 raise RuntimeError("Multiple masters found at same point. You should use sub-space locators.")
263 unresolved = [i for i in np.flatnonzero(found) if (masters_all[starts[i] : starts[i + 1]] < 0).any()]
264 if len(unresolved) > 0:
265 point = np.frombuffer(slave_keys[unresolved[0]], dtype=real_type)
266 raise ValueError(
267 f"No process has a degree of freedom at a master of {len(unresolved)} slave(s), the first at {point}"
268 )
270 coeffs_all = np.array([c for key in slave_keys for c in slave_master_dict[key].values()])
271 if not np.issubdtype(scalar_type, np.complexfloating):
272 coeffs_all = coeffs_all.real
273 coeffs_all = coeffs_all.astype(scalar_type)
274 # The slaves on this process, owned first, then ghosts, each in the order of the dictionary
275 held = np.flatnonzero(slave_dofs >= 0)
276 order = np.concatenate([held[slave_dofs[held] < local_size], held[slave_dofs[held] >= local_size]])
277 counts = starts[order + 1] - starts[order]
278 entries = np.repeat(starts[order] - (np.cumsum(counts) - counts), counts) + np.arange(counts.sum())
279 offsets = np.zeros(len(order) + 1, dtype=np.int32)
280 offsets[1:] = np.cumsum(counts)
281 return (
282 slave_dofs[order].astype(np.int32),
283 masters_all[entries].copy(),
284 coeffs_all[entries],
285 owners_all[entries].astype(np.int32),
286 offsets,
287 )