Coverage for python/src/dolfinx_mpc/integralcondition.py: 98%
82 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-07 09:58 +0000
1# Copyright (C) 2026 Jørgen S. Dokken
2#
3# This file is part of DOLFINX_MPC
4#
5# SPDX-License-Identifier: MIT
6"""Express a scalar integral condition as a multi point constraint."""
8from __future__ import annotations
10import typing
12from mpi4py import MPI
14import dolfinx.fem as _fem
15import numpy as np
16import numpy.typing as npt
17import ufl
18from dolfinx import default_scalar_type, la
20from .container import _deprecated, _tolerance
23def create_integral_constraint(
24 V: _fem.FunctionSpace,
25 weight_form: ufl.Form,
26 value: np.floating | np.complexfloating | float | complex,
27 bcs: typing.Optional[typing.Sequence[_fem.DirichletBC]] = None,
28 rtol: np.floating | float | None = None,
29 *,
30 coefficient_tol: np.floating | float | None = None,
31) -> tuple[
32 npt.NDArray[np.int32],
33 npt.NDArray[np.int64],
34 npt.NDArray,
35 npt.NDArray[np.int32],
36 npt.NDArray[np.int32],
37 _fem.Function,
38]:
39 r"""Build the constraint data enforcing a scalar integral condition.
41 A linear functional :math:`L` applied to :math:`u_h=\sum_i u_i\phi_i` is a
42 single linear equation between all the coefficients,
44 .. math::
45 L(u_h) = \sum_i w_i u_i = \gamma, \qquad w_i = L(\phi_i),
47 where :math:`w` is the assembled vector of ``weight_form``. Solving it for
48 one degree of freedom :math:`s` gives the affine relation
50 .. math::
51 u_s = \sum_{i \neq s} \left(-\frac{w_i}{w_s}\right) u_i
52 + \frac{\gamma}{w_s},
54 that is a constraint with a single slave, every other degree of freedom in
55 the support of the functional as a master, and the inhomogeneity
56 :math:`\gamma/w_s`. Typical uses are fixing the mean value of a solution,
57 :math:`\int_\Omega u~\mathrm{d}x = \gamma`, which makes a pure Neumann
58 problem non-singular, or prescribing a boundary average or a flow rate,
59 :math:`\int_\Gamma u\cdot n~\mathrm{d}s = \gamma`, without prescribing the
60 profile that carries it.
62 Args:
63 V: The function space the constraint acts on.
64 weight_form: A linear form in ``ufl.TestFunction(V)`` defining the
65 functional, for instance ``v * ufl.dx`` or
66 ``ufl.dot(v, n) * ds(tag)``.
67 value: The prescribed value :math:`\gamma` of the functional.
68 bcs: Dirichlet conditions on ``V``. A constrained degree of freedom is
69 never chosen as the slave, but may still appear as a master, in
70 which case passing the same conditions to
71 :class:`MultiPointConstraint` folds its contribution into the
72 constraint offset.
73 rtol: Deprecated, use `coefficient_tol`.
74 coefficient_tol: A master whose coefficient is below `coefficient_tol`
75 times the largest one is dropped. The weights of degrees of freedom
76 outside the support of the functional are returned by quadrature as
77 roundoff rather than as exact zeros, so the threshold is relative
78 rather than a test against zero. Defaults to `500` machine epsilon
79 of the real type of the scalar type: a threshold tight enough for
80 `float64` roundoff is far below `float32` roundoff, and would then
81 discard nothing, inflating the master count and the conditioning of
82 the reduced operator.
84 Returns:
85 The ``slaves``, ``masters``, ``coeffs``, ``owners`` and ``offsets``
86 arrays accepted by :meth:`MultiPointConstraint.add_constraint`, and a
87 :class:`dolfinx.fem.Function` holding the inhomogeneity, to be passed to
88 :class:`MultiPointConstraint` as ``rhs_coeffs``.
90 Note:
91 Collective. Every process contributes its owned weights to the slave's
92 owner, which is the only one that needs the whole functional; the owner
93 then reduces them to the (generally far smaller) list of significant
94 masters and forwards that, not the raw weights, to any process ghosting
95 the slave. Those are the only processes that declare the constraint.
96 """
97 bcs = [] if bcs is None else list(bcs)
98 comm = V.mesh.comm
99 imap = V.dofmap.index_map
100 bs = V.dofmap.index_map_bs
101 num_owned = imap.size_local * bs
102 dtype = np.dtype(default_scalar_type)
103 if rtol is not None:
104 _deprecated("rtol", "`coefficient_tol`")
105 coefficient_tol = rtol if coefficient_tol is None else coefficient_tol
106 rtol = _tolerance(coefficient_tol, dtype)
107 mpi_scalar = MPI._typedict[dtype.char]
109 arguments = ufl.algorithms.extract_arguments(weight_form)
110 if len(arguments) != 1 or arguments[0].number() != 0:
111 raise ValueError(f"weight_form must be linear in a single test function, got {len(arguments)} argument(s)")
112 if arguments[0].ufl_function_space() != V:
113 raise ValueError("The test function of weight_form must be in the function space of the constraint")
115 # Assemble the functional and accumulate ghost contributions onto the owner
116 w = _fem.assemble_vector(_fem.form(weight_form, dtype=dtype))
117 w.scatter_reverse(la.InsertMode.add)
118 w_owned = w.array[:num_owned]
120 # A Dirichlet degree of freedom may be a master, but must not be the slave.
121 # Zero those weights in a scratch copy: `set` writes alpha * (value - x0), so
122 # alpha=0 masks them whatever the condition prescribes.
123 masked = w.array.copy()
124 for bc in bcs:
125 bc.set(masked, alpha=0.0)
127 # The largest weight makes the best slave, since it bounds every coefficient
128 # by one in magnitude. That matters, because cond(K^H A K) grows with the
129 # square of the coefficient norm and the weights can span many orders of
130 # magnitude. Reducing over the masked weights picks it without communicating
131 # either the weights or the Dirichlet markers.
132 local_best = int(np.argmax(np.abs(masked[:num_owned]))) if num_owned > 0 else 0
133 local_magnitude = float(abs(masked[local_best])) if num_owned > 0 else -1.0
134 best_magnitude, slave_global = comm.allreduce(
135 (local_magnitude, int(imap.local_range[0] * bs) + local_best if num_owned > 0 else -1),
136 op=MPI.MAXLOC,
137 )
138 # `best_magnitude` is |w_s|, never the signed weight, so it is negative only
139 # for the sentinel a process with no owned dofs contributes. Comparing it
140 # against zero would be too weak: a weight outside the support of the
141 # functional comes back from quadrature as roundoff rather than as an exact
142 # zero, so a slave whose weight is negligible *relative to the functional*
143 # would otherwise be accepted and give coefficients of order 1/rtol.
144 scale = comm.allreduce(float(np.abs(w_owned).max(initial=0.0)), op=MPI.MAX)
145 if best_magnitude <= rtol * scale:
146 raise RuntimeError(
147 "No admissible slave: the functional vanishes, to within rtol, on "
148 "every degree of freedom that is not constrained by a Dirichlet condition"
149 )
151 # Only the processes holding the slave declare the constraint, so only they
152 # need the weights: gather to one of them and forward to the few others.
153 #
154 # `global_to_local` gives the slave's local block on every rank that owns
155 # or ghosts it; `imap.local_range`/`size_local` then say which of the two
156 # for free, no communication, since both are purely local properties of
157 # the index map. A rank that only ghosts the slave also already knows who
158 # owns it: `owners` is aligned with the ghost list, so the position of the
159 # slave's block within it (past `size_local`) is the owning rank directly,
160 # again with no communication -- the whole point of ghost ownership info
161 # existing on the index map in the first place.
162 slave_block = int(imap.global_to_local(np.array([slave_global // bs], dtype=np.int64))[0])
163 holds_slave = slave_block != -1
164 is_owner = holds_slave and slave_block < imap.size_local
165 owner_hint = comm.rank if is_owner else (int(imap.owners[slave_block - imap.size_local]) if holds_slave else -1)
166 # Every other process still needs to know root, if only to call the
167 # Gatherv below with a value of `root` that agrees with everyone else's;
168 # propagating the one fact a holding rank already has to the rest of the
169 # communicator is a single scalar reduction (there is exactly one owner,
170 # so MAX recovers it) rather than an allgather of one integer per rank.
171 root = comm.allreduce(owner_hint, op=MPI.MAX)
172 assert root != -1, "No process holds the slave, but one must"
174 # The owned degrees of freedom of a process form a contiguous global
175 # range, and the processes concatenate in rank order, so a gathered
176 # entry's position is already its global index; `counts` is needed for
177 # the Gatherv below regardless, to size and place each process's part.
178 counts = np.array(comm.allgather(num_owned), dtype=np.int32)
179 displ = np.concatenate(([0], np.cumsum(counts)[:-1])).astype(np.int32)
180 total = int(counts.sum())
182 # Root needs to know who else to Send to, i.e. who ghosts the slave. Rather
183 # than every process announcing whether it holds the slave (an allgather
184 # over the whole communicator), only the ranks that ghost it declare an
185 # edge, to the single rank they already know is the owner; root discovers
186 # them by building the graph and reading back its incoming edges. Every
187 # process must still call this (it is collective), but only a handful ever
188 # declare a nonzero degree, so the underlying exchange is between just
189 # those ranks and root instead of all of them.
190 destinations = [root] if (holds_slave and not is_owner) else []
191 graph = comm.Create_dist_graph([comm.rank], [len(destinations)], destinations, reorder=False)
192 ghost_ranks, _, _ = graph.Get_dist_neighbors()
193 graph.Free()
195 all_weights = np.empty(total, dtype=dtype) if is_owner else None
196 # mpi4py reads a three-entry tuple as (buffer, counts, datatype), so the
197 # datatype has to be spelled out whenever displacements are given
198 recvbuf = (all_weights, counts, displ, mpi_scalar) if is_owner else None
199 comm.Gatherv(np.ascontiguousarray(w_owned, dtype=dtype), recvbuf, root=root)
201 rhs_coeffs = _fem.Function(V, dtype=dtype)
202 rhs_coeffs.x.array[:] = 0
203 slaves = np.zeros(0, dtype=np.int32)
204 masters = np.zeros(0, dtype=np.int64)
205 coeffs = np.zeros(0, dtype=dtype)
206 owners = np.zeros(0, dtype=np.int32)
207 offsets = np.zeros(1, dtype=np.int32)
208 if is_owner:
209 assert all_weights is not None
210 w_slave = all_weights[slave_global]
211 # A global master index is its position in `all_weights` (established
212 # above), so building an index array and deleting the slave's entry from
213 # it -- and separately from the coefficients and owners -- is three
214 # total-sized allocate-and-copy passes for something one boolean mask
215 # does in one. The slave's own would-be coefficient is exactly ±1 (it
216 # has the largest |w| by construction), i.e. the largest a coefficient
217 # can ever be, so it must be zeroed before the max below, not just
218 # excluded from the final selection, or it would silently set the
219 # threshold instead of being subject to it.
220 all_coeffs = -all_weights / w_slave
221 all_coeffs[slave_global] = 0.0
222 # Drop negligible coefficients; they change nothing in the constraint but
223 # each costs a ghost, a row of the sparsity pattern and an entry in every
224 # element matrix modification. Filtering here, once, rather than after
225 # forwarding the raw weights to each ghost, means every ghost gets the
226 # already-reduced answer instead of redoing this reduction on its own
227 # copy of the (generally much larger) dense weight vector.
228 significant = np.abs(all_coeffs) >= rtol * np.abs(all_coeffs).max(initial=0.0)
229 significant[slave_global] = False # never a master of itself, even if threshold == 0
230 all_owners_full = np.repeat(np.arange(comm.size, dtype=np.int32), counts)
231 masters = np.flatnonzero(significant).astype(np.int64)
232 coeffs = np.ascontiguousarray(all_coeffs[significant], dtype=dtype)
233 owners = np.ascontiguousarray(all_owners_full[significant], dtype=np.int32)
234 # Non-blocking: issuing every send before waiting on any of them means
235 # root pays for the slowest destination once, not the sum of all of
236 # them. ghost_ranks stays bounded by the slave's local mesh valence
237 # (how many subdomains meet at that dof), not by the size of comm, but
238 # there is no reason to pay sequential latency for it regardless.
239 requests = [comm.isend((masters, coeffs, owners, w_slave), dest=int(other), tag=0) for other in ghost_ranks]
240 MPI.Request.waitall(requests)
241 elif holds_slave:
242 masters, coeffs, owners, w_slave = comm.recv(source=root, tag=0)
244 if holds_slave:
245 slaves = np.array([slave_block * bs + slave_global % bs], dtype=np.int32)
246 offsets = np.array([0, masters.size], dtype=np.int32)
247 rhs_coeffs.x.array[slaves[0]] = dtype.type(value) / w_slave
248 rhs_coeffs.x.scatter_forward()
249 return slaves, masters, coeffs, owners, offsets, rhs_coeffs