Fixing a mean value with an affine multi-point constraint

Fixing a mean value with an affine multi-point constraint#

Author Jørgen S. Dokken

This demo shows how an affine multi-point constraint can enforce a constraint on an integral of the solution,

\[ \int_\Omega u ~\mathrm{d}x = \gamma, \]

and how that relates to the more familiar way of imposing such a condition, a Lagrange multiplier in a real space (basix.ufl.real_element()), as seen in SciFEM real space demo. The two are algebraically equivalent; they differ sharply in cost, and this demo measures the difference.

Mathematical formulation#

We solve the pure Neumann Poisson problem

\[\begin{split} \begin{align*} -\Delta u &= f &&\text{in } \Omega,\\ \frac{\partial u}{\partial n} &= h &&\text{on } \partial\Omega, \end{align*} \end{split}\]

on \(\Omega=(0,1)^2\). The problem is singular: constants lie in the kernel, so \(u\) is determined only up to an additive constant, and the data must satisfy the compatibility condition \(\int_\Omega f~\mathrm{d}x + \int_{\partial\Omega} h~\mathrm{d}s = 0\). Prescribing the mean value \(\int_\Omega u~\mathrm{d}x = \gamma\) selects the unique solution.

We start by import the required modules:

Hide code cell source

 
from __future__ import annotations

import time
from pathlib import Path

from mpi4py import MPI

import basix.ufl
import numpy as np
import pandas
import pyvista
import ufl
from dolfinx import default_scalar_type, fem, mesh, plot

import dolfinx_mpc.utils
from dolfinx_mpc import LinearProblem, MultiPointConstraint

Problem data#

We use the manufactured solution \(u_{ex} = x^2 - x + C\). Everything else is derived from it with UFL rather than worked out by hand: the source is \(f=-\Delta u_{ex}\), the boundary flux is \(h=\nabla u_{ex}\cdot n\), and the target \(\gamma\) is the integral of \(u_{ex}\) itself. Deriving \(f\) and \(h\) from the same expression makes the compatibility condition hold by construction, since \(\int_\Omega -\Delta u_{ex}~\mathrm{d}x + \int_{\partial\Omega} \nabla u_{ex}\cdot n~\mathrm{d}s = 0\) is the divergence theorem. Since \(u_{ex}\) is a second order polynomial it lies in the second order Lagrange space, so the discrete solution should reproduce it up to solver accuracy.

C = 3.0
degree = 2

Transforming an integral condition into a multi-point constraint#

Expanding \(u_h=\sum_i u_i\phi_i\) turns the integral condition into a single linear equation between all the coefficients,

\[ \int_\Omega u_h~\mathrm{d}x = \sum_i u_i \int_\Omega \phi_i ~\mathrm{d}x = \sum_i w_i u_i = \gamma, \qquad w_i = \int_\Omega \phi_i~\mathrm{d}x. \]

The weights \(w\) are just the assembled vector of the linear form \(v\mapsto\int_\Omega v~\mathrm{d}x\). Picking one degree of freedom \(s\) with \(w_s\neq 0\) and solving for it gives exactly the affine multi-point constraint \(u = K\hat{u} + g\) of dolfinx_mpc.MultiPointConstraint,

\[ u_s = \sum_{i\neq s} \underbrace{\left(-\frac{w_i}{w_s}\right)}_{c_{si}} u_i + \underbrace{\frac{\gamma}{w_s}}_{g_s}. \]

So \(s\) is the single slave, every other degree of freedom in the support of the functional is a master, the coefficients are \(-w_i/w_s\).

There are two new concepts in this demo compared to the others:

  1. The inhomogeneity \(g_s=\frac{\gamma}{w_s}\) is computed and assigned to the constraint.

  2. The constraint involves an integral, so the weights \(w_i\) are computed by assembling a linear form rather than reading them off a mesh entity. We create a convenience function build_constraint that takes in:

  3. the function space V that our dolfinx_mpc.MultiPointConstraint will be built on

  4. the linear form weight_form that defines the weights

  5. The target value value, which is \(\gamma\)

  6. Any boundary conditions bcs that should be applied to the constraint. The function returns the finalized MPC as well as the number of masters in the MPC constraint.

def build_constraint(V, weight_form, value, bcs=()):
    """Constrain ``L(u) = value`` and report how many masters survived."""
    mpc = MultiPointConstraint(V, bcs=list(bcs))
    mpc.add_integral_constraint(weight_form, value, bcs=list(bcs))
    mpc.finalize()
    kept = sum(len(mpc.masters.links(s)) for s in mpc.slaves[: mpc.num_local_slaves])
    return mpc, V.mesh.comm.allreduce(kept, op=MPI.SUM)

Variational problem#

The forms are the standard pure Neumann Poisson ones; nothing in them knows about the constraint.

def build_poisson_forms(V: fem.FunctionSpace, u_ex: ufl.core.expr.Expr) -> tuple[ufl.Form, ufl.Form]:
    """Bilinear and linear form of the pure Neumann Poisson problem.

    The source and the boundary flux are both differentiated out of ``u_ex``, so
    the data is compatible by construction.
    """
    u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
    n = ufl.FacetNormal(V.mesh)
    f = -ufl.div(ufl.grad(u_ex))
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx
    L = ufl.inner(f, v) * ufl.dx + ufl.inner(ufl.dot(ufl.grad(u_ex), n), v) * ufl.ds
    return a, L


def exact_mean(domain: mesh.Mesh, u_ex):
    """The target value, integrated from the manufactured solution."""
    compiled_mean = fem.form(u_ex * ufl.dx(domain=domain))
    return compiled_mean.mesh.comm.allreduce(fem.assemble_scalar(compiled_mean), op=MPI.SUM)

Setting up the problem#

The reduced operator of a cell integral constraint is dense, so the mesh is deliberately coarse; the cost table at the end of the demo shows why.

N = 16
domain = mesh.create_unit_square(MPI.COMM_WORLD, N, N)
comm = domain.comm
V = fem.functionspace(domain, ("Lagrange", degree))
x = ufl.SpatialCoordinate(domain)
u_ex = x[0] ** 2 - x[0] + C
a, L = build_poisson_forms(V, u_ex)
gamma = exact_mean(domain, u_ex)

We use the form ufl.TestFunction(V) * ufl.dx to compute the weights \(w_i\) of the integral constraint.

comm.Barrier()
_t0 = time.perf_counter()
mpc, num_masters = build_constraint(V, ufl.conj(ufl.TestFunction(V)) * ufl.dx, gamma)
comm.Barrier()
_t1 = time.perf_counter()
t_constraint = _t1 - _t0

Solving#

dolfinx_mpc.LinearProblem handles the affine constraint for us: it refreshes the offset with update_constants and applies dolfinx_mpc.apply_mpc_lifting(), which contributes the \(-K^TAg\) term. Note that no null space has to be attached even though \(A\) is singular: the constraint removes the constant kernel, because \(w\cdot\mathbf{1} = |\Omega| \neq 0\).

petsc_options = {
    "ksp_type": "preonly",
    "pc_type": "lu",
    "pc_factor_mat_solver_type": "mumps",
    "ksp_error_if_not_converged": True,
}
problem = LinearProblem(a, L, mpc, bcs=[], petsc_options=petsc_options)
uh = problem.solve()
comm.Barrier()
t_mpc = time.perf_counter() - _t1

def operator_stats(A, comm, singular=False, root=0):
    """Global nnz and 2-norm condition number of an assembled operator.

    The condition number is computed from a dense SVD on ``root``, so this is a
    diagnostic for demo sized problems only. Used just below to make the
    verification tolerance conditioning-aware; see "Cost and conditioning"
    further down for why $K^TAK$ needs this at all.
    """
    # petsc4py defaults to MatInfoType.GLOBAL_SUM, so this is already reduced
    # over the communicator and must not be summed again.
    nnz = int(A.getInfo()["nz_used"])
    A_csr = dolfinx_mpc.utils.gather_PETScMatrix(A, root=root)
    cond = None
    if comm.rank == root:
        sv = np.linalg.svd(A_csr.toarray(), compute_uv=False)
        # A is singular for the pure Neumann problem, so compare against the
        # smallest *nonzero* singular value.
        cond = sv[0] / (sv[-2] if singular else sv[-1])
    return nnz, comm.bcast(cond, root=root)

Verification#

Three things are checked: that the constraint is satisfied and that the manufactured solution is recovered.

mean_value = comm.allreduce(fem.assemble_scalar(fem.form(uh * ufl.dx)), op=MPI.SUM)
error = np.sqrt(comm.allreduce(fem.assemble_scalar(fem.form((uh - u_ex) ** 2 * ufl.dx)), op=MPI.SUM))
_, cond_mpc = operator_stats(problem.A, comm)

Hide code cell source

 
if comm.rank == 0:
    print("----Verification----")
    print(f"  dofs                  {V.dofmap.index_map.size_global * V.dofmap.index_map_bs}")
    print(f"  masters               {num_masters}")
    print(f"  mean(u_h)             {mean_value:.15f} (target {gamma:.15f})")
    print(f"  |mean(u_h) - gamma|   {abs(mean_value - gamma):.3e}")
    print(f"  L2(u_h - u_ex)        {error:.3e}")
    print(f"  cond(K^TAK)           {cond_mpc:.3e}")

----Verification----
  dofs                  1089
  masters               799
  mean(u_h)             2.833333333333336 (target 2.833333333333333)
  |mean(u_h) - gamma|   3.109e-15
  L2(u_h - u_ex)        4.748e-11
  cond(K^TAK)           4.576e+05

tol = 100 * np.finfo(default_scalar_type()).eps
assert abs(mean_value - gamma) < tol

u_ex lies in the discrete space, so the discretization is exact. The remaining error is the conditioning of the dense reduced operator: a direct solve loses roughly \(\log_{10}\kappa\) digits, so the bound below scales with the measured \(\kappa(K^TAK)\) rather than a fixed constant – a flat bound tight enough for float64 would be far too strict for a lower-precision build, where tol itself is orders of magnitude larger to begin with.

assert error < 10 * cond_mpc * tol

Relation to a real space#

The constrained problem has the saddle point form

\[\begin{split} \begin{pmatrix} A & w \\ w^T & 0\end{pmatrix} \begin{pmatrix} u \\ \lambda \end{pmatrix} = \begin{pmatrix} b \\ \gamma \end{pmatrix}, \end{split}\]

where \(\lambda\) is the scalar Lagrange multiplier carried by a real space. Since \(\mathrm{range}(K)=\{v: w^Tv=0\}\), the first block row says \(b-Au \in \mathrm{span}(w)\), i.e. \(K^T(b-Au)=0\); together with \(w^Tu=\gamma\) that is the reduced system \(K^TAK\hat{u}=K^T(b-Ag)\).

The same problem, solved as a saddle point system with a scalar multiplier in a real space. Because the real basis function is identically one, inner(lam, v) * dx assembles the dense row \(w\) directly. The right hand side entry of the constraint row must be \(\gamma\) itself, so the constant we integrate against the real test function is \(\gamma/|\Omega|\).

def solve_real_space(domain, degree, u_ex, value):
    """Reference solve using a Lagrange multiplier in a real space."""
    V = fem.functionspace(domain, ("Lagrange", degree))
    R = fem.functionspace(domain, basix.ufl.real_element(domain.basix_cell(), dtype=domain.geometry.x.dtype))
    W = ufl.MixedFunctionSpace(V, R)
    u, lam = ufl.TrialFunctions(W)
    v, mu = ufl.TestFunctions(W)
    n = ufl.FacetNormal(domain)
    f = -ufl.div(ufl.grad(u_ex))

    volume = domain.comm.allreduce(
        fem.assemble_scalar(fem.form(fem.Constant(domain, default_scalar_type(1.0)) * ufl.dx)), op=MPI.SUM
    )
    a = ufl.inner(ufl.grad(u), ufl.grad(v)) * ufl.dx + ufl.inner(lam, v) * ufl.dx + ufl.inner(u, mu) * ufl.dx
    L = (
        ufl.inner(f, v) * ufl.dx
        + ufl.inner(ufl.dot(ufl.grad(u_ex), n), v) * ufl.ds
        + ufl.inner(fem.Constant(domain, default_scalar_type(value / volume)), mu) * ufl.dx
    )
    problem = fem.petsc.LinearProblem(
        ufl.extract_blocks(a),
        ufl.extract_blocks(L),
        bcs=[],
        kind="mpi",
        petsc_options_prefix="demo_mean_value_real_",
        petsc_options={
            "ksp_type": "preonly",
            "pc_type": "lu",
            "pc_factor_mat_solver_type": "mumps",
            "ksp_error_if_not_converged": True,
        },
    )
    uh, lamh = problem.solve()
    # The problem is returned alongside the solution: it owns the PETSc matrix,
    # which is destroyed with it.
    return uh, lamh, problem

We compare the solution of the Lagrange multiplier problem to the one with the multi-point constraint.

_t2 = time.perf_counter()
u_real, lam_real, real_problem = solve_real_space(domain, degree, u_ex, gamma)
comm.Barrier()
t_real = time.perf_counter() - _t2

num_owned = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
assert isinstance(uh, fem.Function)
_diff = np.max(np.abs(uh.x.array[:num_owned] - u_real.x.array[:num_owned])) if num_owned else 0.0
mpc_vs_real = comm.allreduce(_diff, op=MPI.MAX)
# The real-space (saddle point) solve keeps A's own, much better conditioning,
# so this difference is dominated by uh's error and needs the same
# conditioning-aware bound as the check against u_ex above.
assert mpc_vs_real < 10 * cond_mpc * tol
print(f"  max|u_mpc - u_real|   {mpc_vs_real:.3e}")

  max|u_mpc - u_real|   1.148e-10

Cost and conditioning#

Eliminating the slave is not free. Order the degrees of freedom with the masters \(m\) first and the slave \(s\) last, so that the constraint above is \(u=K\hat{u}+g\) and both \(A\) and \(K\) split as

\[\begin{split} A = \begin{pmatrix} A_{mm} & A_{ms} \\ A_{sm} & A_{ss}\end{pmatrix}, \qquad K = \begin{pmatrix} I \\ c^T \end{pmatrix}, \qquad c_i = -\frac{w_i}{w_s}, \end{split}\]

where \(c\) collects the coefficients of the one slave and \(A_{ss}\) is a scalar. Multiplying out,

\[ K^TAK = A_{mm} + A_{ms}c^T + cA_{sm} + \left(cc^T\right)A_{ss}. \]

The final term is a rank one matrix that is dense over every pair of masters that survive filtering, since \(c\) has no near-zero entries left once they are dropped. For a cell integral nearly every degree of freedom is a master – only the P2 vertex dofs are filtered out, since a P2 vertex basis function integrates to exactly zero on simplices – so \(K^TAK\) is essentially full, and its norm is inflated by \(\lVert c\rVert^2\), which costs roughly a factor \(M\) in the condition number. The saddle point system pays neither price: it keeps \(A\) intact and appends one row and column. That row is itself dense – every \(w_i\) is stored, so the system grows by exactly \(2N\) entries – but \(2N\) is linear in the problem size, against the \(M^2\) of the elimination. operator_stats, defined and used above to make the verification tolerance conditioning-aware, measures both effects; the sweep below repeats it over a refinement range.

Cost and conditioning, measured#

The helper below repeats the solve over a small refinement sweep so that the growth is visible rather than asserted.

def measure(N: int) -> dict:
    """Solve at resolution ``N`` and collect cost, conditioning and timings."""
    domain = mesh.create_unit_square(MPI.COMM_WORLD, N, N)
    comm = domain.comm
    V = fem.functionspace(domain, ("Lagrange", degree))
    x = ufl.SpatialCoordinate(domain)
    u_ex = x[0] ** 2 - x[0] + C
    a, L = build_poisson_forms(V, u_ex)
    gamma = exact_mean(domain, u_ex)

    comm.Barrier()
    t0 = time.perf_counter()
    mpc, num_masters = build_constraint(V, ufl.conj(ufl.TestFunction(V)) * ufl.dx, gamma)
    comm.Barrier()
    t1 = time.perf_counter()
    problem = LinearProblem(a, L, mpc, bcs=[], petsc_options=petsc_options)
    problem.solve()
    comm.Barrier()
    t2 = time.perf_counter()

    t3 = time.perf_counter()
    _, _, real_problem = solve_real_space(domain, degree, u_ex, gamma)
    comm.Barrier()
    t_real = time.perf_counter() - t3

    A_plain = fem.petsc.assemble_matrix(fem.form(a), bcs=[])
    A_plain.assemble()
    nnz_A, cond_A = operator_stats(A_plain, comm, singular=True)
    nnz_mpc, cond_mpc = operator_stats(problem.A, comm)
    nnz_real, cond_real = operator_stats(real_problem.A, comm)
    A_plain.destroy()
    return {
        "N": V.dofmap.index_map.size_global * V.dofmap.index_map_bs,
        "M": num_masters,
        "norm_c": float(np.linalg.norm(mpc.coefficients()[0])),
        "nnz_A": nnz_A,
        "nnz_mpc": nnz_mpc,
        "nnz_real": nnz_real,
        "cond_A": cond_A,
        "cond_mpc": cond_mpc,
        "cond_real": cond_real,
        "t_constraint": t1 - t0,
        "t_mpc": t2 - t1,
        "t_real": t_real,
    }

rows = [measure(N) for N in (8, 12, 16)]

COLUMNS = {
    "N": "dofs",
    "M": "masters",
    "norm_c": "||c||",
    "nnz_A": "nnz(A)",
    "nnz_mpc": "nnz(KtAK)",
    "nnz_real": "nnz(saddle)",
    "cond_A": "cond(A)",
    "cond_mpc": "cond(KtAK)",
    "cond_real": "cond(saddle)",
    "t_constraint": "build [s]",
    "t_mpc": "mpc [s]",
    "t_real": "real [s]",
    "mpc/real": "mpc/real",
}
FORMATS = {
    "||c||": "{:.2f}",
    "cond(A)": "{:.3e}",
    "cond(KtAK)": "{:.3e}",
    "cond(saddle)": "{:.3e}",
    "build [s]": "{:.4f}",
    "mpc [s]": "{:.4f}",
    "real [s]": "{:.4f}",
    "mpc/real": "{:.2f}",
}

table = pandas.DataFrame(rows)
table["mpc/real"] = table["t_mpc"] / table["t_real"]
table = table[list(COLUMNS)].rename(columns=COLUMNS).set_index("dofs")
table.style.format(FORMATS)
  masters ||c|| nnz(A) nnz(KtAK) nnz(saddle) cond(A) cond(KtAK) cond(saddle) build [s] mpc [s] real [s] mpc/real
dofs                        
289 207 13.53 3073 46967 3651 3.300e+02 3.088e+04 3.300e+02 0.0014 0.0112 0.0065 1.72
625 455 20.47 6817 216141 8067 7.018e+02 1.486e+05 7.018e+02 0.0011 0.0472 0.0073 6.46
1089 799 27.40 12033 654503 14211 1.212e+03 4.576e+05 1.212e+03 0.0012 0.1494 0.0091 16.42

The table shows nnz(KtAK) growing like \(M^2\) and cond(KtAK) like \(\lVert c\rVert^2\sim M\), because the rank one term \(cc^TA_{ss}\) couples every master to every other. Here the functional is supported on the whole mesh, so \(M\) is the number of degrees of freedom (bar the vertex dofs, whose weights vanish for P2 on simplices) and the reduced operator is essentially full. A real space instead appends a single row and column – dense, but only \(2N\) entries, so nnz(saddle) grows linearly where nnz(KtAK) grows quadratically – and leaves the conditioning of \(A\) untouched.

The timings say the same in wall clock terms. Building the constraint is cheap and mesh independent; what costs is solving with the operator it produces, and the ratio to the real space solve grows with the mesh. For a functional supported on a facet \(M\) is far smaller and both penalties shrink with it, see Prescribing a boundary average with an affine multi-point constraint.

Visualization#

Each process builds a PyVista grid over the cells it owns, so a shared cell is not drawn twice, and the grids are gathered onto one process and drawn into a single figure with common colour limits.

pyvista.global_theme.allow_empty_mesh = True


def gather_grids(u: fem.Function, V: fem.FunctionSpace, name: str, root: int = 0):
    """Owned-cell PyVista grids with ``u`` attached, gathered on ``root``.

    Returns the list of grids on rank 0 (``None`` elsewhere) and the global
    value range, so that every piece can be drawn with the same colour limits.
    """
    comm = V.mesh.comm
    bs = V.dofmap.index_map_bs
    tdim = V.mesh.topology.dim
    owned_cells = np.arange(V.mesh.topology.index_map(tdim).size_local, dtype=np.int32)
    grid = pyvista.UnstructuredGrid(*plot.vtk_mesh(V, entities=owned_cells))
    # vtk_mesh emits one point per dof block, in local numbering
    values = u.x.array.real[: grid.n_points * bs]
    grid.point_data[name] = values if bs == 1 else values.reshape(-1, bs)
    local = np.linalg.norm(values.reshape(-1, bs), axis=1) if bs > 1 else values
    lo = comm.allreduce(float(local.min()) if local.size else np.inf, op=MPI.MIN)
    hi = comm.allreduce(float(local.max()) if local.size else -np.inf, op=MPI.MAX)
    # gather returns the list on `root` and None everywhere else, so the caller
    # can test the result instead of comparing ranks itself
    return comm.gather(grid, root=root), [lo, hi]


# `uh` lives in the constraint's extended space, which carries the master dofs as
# extra ghosts, so its array is longer than the original space in parallel. The
# extended index map keeps the original dofs first, so the leading entries are
# exactly the values of the original space.

u_plot = fem.Function(V)
u_plot.x.array[:] = uh.x.array[: u_plot.x.array.size]
pieces, clim = gather_grids(u_plot, V, "u")

# The solution varies only with $x$, and the constraint has picked the vertical
# offset $C-1/6$ out of the one parameter family the pure Neumann problem admits.
# Warping by the value makes that offset visible as the height above zero.

if pieces is not None:  # only the root process received the grids
    plotter = pyvista.Plotter(window_size=[700, 500])
    plotter.add_text(f"mean(u) = {mean_value:.6f}", font_size=10)
    for piece in pieces:
        plotter.add_mesh(
            piece.warp_by_scalar("u", factor=0.15),
            scalars="u",
            cmap="viridis",
            clim=clim,
            show_edges=False,
            scalar_bar_args={"vertical": True},
        )
    plotter.view_isometric()  # type: ignore[call-arg]
    plotter.camera.zoom(1.4)
    # The figure is named after the demo, as the gallery of the documentation expects
    figure = Path("demo_mean_value_constraint.py")
    if pyvista.OFF_SCREEN:
        plotter.screenshot(figure.with_suffix(".png"))
    else:
        # The interactive scene, for the gallery
        plotter.export_vtksz(figure.with_suffix(".vtksz"))
        plotter.show(screenshot=figure.with_suffix(".png"))

/tmp/ipykernel_2376/1621171634.py:59: PyVistaDeprecationWarning: Plotter.export_vtksz is deprecated. Use plotter.trame.export_vtksz(...) instead. Install via `pip install trame-pyvista`.
  plotter.export_vtksz(figure.with_suffix(".vtksz"))