Stokes problem with slip condition using NEST matrices#

Copyright (C) 2022 Nathan Sime

SPDX-License-Identifier: MIT

This demo illustrates how to apply a slip condition on an interface not aligned with the coordinate axis. The demos solves the Stokes problem using the nest functionality to avoid using mixed function spaces. The demo also illustrates how to use block preconditioners with PETSc. This demo uses the lower level API of dolfinx_mpc. See Mortar coupling: a Lagrange multiplier on a submesh for a higher level API for creating MPCs on NEST matrices.

Hide code cell source

from __future__ import annotations

from pathlib import Path

from mpi4py import MPI
from petsc4py import PETSc

import basix
import dolfinx.io
import gmsh
import numpy as np
import pyvista
import scipy.sparse.linalg
import ufl
from dolfinx import default_real_type, default_scalar_type, plot
from dolfinx.io import gmsh as gmshio
from ufl.core.expr import Expr

import dolfinx_mpc
import dolfinx_mpc.utils

We create a convenience function to create a mesh with gmsh.

def create_mesh_gmsh(
    L: int = 2,
    H: int = 1,
    res: float = 0.1,
    theta: float = np.pi / 5,
    wall_marker: int = 1,
    outlet_marker: int = 2,
    inlet_marker: int = 3,
):
    """
    Create a channel of length L, height H, rotated theta degrees
    around origin, with facet markers for inlet, outlet and walls.


    Parameters
    ----------
    L
        The length of the channel
    H
        Width of the channel
    res
        Mesh resolution (uniform)
    theta
        Rotation angle
    wall_marker
        Integer used to mark the walls of the channel
    outlet_marker
        Integer used to mark the outlet of the channel
    inlet_marker
        Integer used to mark the inlet of the channel
    """
    gmsh.initialize()
    if MPI.COMM_WORLD.rank == 0:
        gmsh.model.add("Square duct")

        # Create rectangular channel
        channel = gmsh.model.occ.addRectangle(0, 0, 0, L, H)
        gmsh.model.occ.synchronize()

        # Find entity markers before rotation
        surfaces = gmsh.model.occ.getEntities(dim=1)
        walls = []
        inlets = []
        outlets = []
        for surface in surfaces:
            com = gmsh.model.occ.getCenterOfMass(surface[0], surface[1])
            if np.allclose(com, [0, H / 2, 0]):
                inlets.append(surface[1])
            elif np.allclose(com, [L, H / 2, 0]):
                outlets.append(surface[1])
            elif np.isclose(com[1], 0) or np.isclose(com[1], H):
                walls.append(surface[1])
        # Rotate channel theta degrees in the xy-plane
        gmsh.model.occ.rotate([(2, channel)], 0, 0, 0, 0, 0, 1, theta)
        gmsh.model.occ.synchronize()

        # Add physical markers
        gmsh.model.addPhysicalGroup(2, [channel], 1)
        gmsh.model.setPhysicalName(2, 1, "Fluid volume")
        gmsh.model.addPhysicalGroup(1, walls, wall_marker)
        gmsh.model.setPhysicalName(1, wall_marker, "Walls")
        gmsh.model.addPhysicalGroup(1, inlets, inlet_marker)
        gmsh.model.setPhysicalName(1, inlet_marker, "Fluid inlet")
        gmsh.model.addPhysicalGroup(1, outlets, outlet_marker)
        gmsh.model.setPhysicalName(1, outlet_marker, "Fluid outlet")

        # Set uniform mesh size
        gmsh.option.setNumber("Mesh.CharacteristicLengthMin", res)
        gmsh.option.setNumber("Mesh.CharacteristicLengthMax", res)

        # Generate mesh
        gmsh.model.mesh.generate(2)
    # Convert gmsh model to DOLFINx Mesh and meshtags
    mesh_data = gmshio.model_to_mesh(gmsh.model, MPI.COMM_WORLD, 0, gdim=2)
    assert mesh_data.cell_tags is not None
    ft = mesh_data.facet_tags
    gmsh.finalize()
    return mesh_data.mesh, ft

Mesh and function space creation#

mesh, mt = create_mesh_gmsh(res=0.1)

fdim = mesh.topology.dim - 1

# Create the function space
cellname = mesh.basix_cell()
Ve = basix.ufl.element(
    basix.ElementFamily.P,
    cellname,
    2,
    shape=(mesh.geometry.dim,),
    dtype=default_real_type,
)
Qe = basix.ufl.element(basix.ElementFamily.P, cellname, 1, dtype=default_real_type)

V = dolfinx.fem.functionspace(mesh, Ve)
Q = dolfinx.fem.functionspace(mesh, Qe)
Info    : Meshing 1D...
Info    : [  0%] Meshing curve 1 (Line)
Info    : [ 30%] Meshing curve 2 (Line)
Info    : [ 60%] Meshing curve 3 (Line)
Info    : [ 80%] Meshing curve 4 (Line)
Info    : Done meshing 1D (Wall 0.000367584s, CPU 0.000578s)
Info    : Meshing 2D...
Info    : Meshing surface 1 (Plane, Frontal-Delaunay)
Info    : Done meshing 2D (Wall 0.00864046s, CPU 0.008706s)
Info    : 275 nodes 552 elements

Defining boundary conditions#

Inlet velocity Dirichlet BC#

def inlet_velocity_expression(x):
    return np.stack(
        (
            np.sin(np.pi * np.sqrt(x[0] ** 2 + x[1] ** 2)),
            5 * x[1] * np.sin(np.pi * np.sqrt(x[0] ** 2 + x[1] ** 2)),
        )
    )


inlet_velocity = dolfinx.fem.Function(V)
inlet_velocity.interpolate(inlet_velocity_expression)
inlet_velocity.x.scatter_forward()
dofs = dolfinx.fem.locate_dofs_topological(V, 1, mt.find(3))
bc1 = dolfinx.fem.dirichletbc(inlet_velocity, dofs)
bcs: list[dolfinx.fem.DirichletBC] = [bc1]

Slip conditions for walls#

n = dolfinx_mpc.utils.create_normal_approximation(V, mt, 1)
with dolfinx.common.Timer("~Stokes: Create slip constraint"):
    mpc = dolfinx_mpc.MultiPointConstraint(V)
    mpc.create_slip_constraint(V, (mt, 1), n, bcs=bcs)
mpc.finalize()

mpc_q = dolfinx_mpc.MultiPointConstraint(Q)
mpc_q.finalize()

Variational problem#

def tangential_proj(u: Expr, n: Expr):
    """
    See for instance:
    https://link.springer.com/content/pdf/10.1023/A:1022235512626.pdf
    """
    return (ufl.Identity(u.ufl_shape[0]) - ufl.outer(n, n)) * u


def sym_grad(u: Expr):
    return ufl.sym(ufl.grad(u))


def T(u: Expr, p: Expr, mu: Expr):
    return 2 * mu * sym_grad(u) - p * ufl.Identity(u.ufl_shape[0])


# Traditional terms
mu = 1
f = dolfinx.fem.Constant(mesh, default_scalar_type((0, 0)))
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)
a00 = 2 * mu * ufl.inner(sym_grad(u), sym_grad(v)) * ufl.dx
a01 = -ufl.inner(p, ufl.div(v)) * ufl.dx
a10 = -ufl.inner(ufl.div(u), q) * ufl.dx
a11 = None

L0 = ufl.inner(f, v) * ufl.dx
L1 = ufl.inner(dolfinx.fem.Constant(mesh, default_scalar_type(0.0)), q) * ufl.dx

# No prescribed shear stress
n = ufl.FacetNormal(mesh)
g_tau = tangential_proj(dolfinx.fem.Constant(mesh, default_scalar_type(((0, 0), (0, 0)))) * n, n)
ds = ufl.Measure("ds", domain=mesh, subdomain_data=mt, subdomain_id=1)

# Terms due to slip condition
# Explained in for instance: https://arxiv.org/pdf/2001.10639.pdf
a00 = 2 * mu * ufl.inner(sym_grad(u), sym_grad(v)) * ufl.dx

a00 -= ufl.inner(ufl.outer(n, n) * ufl.dot(2 * mu * sym_grad(u), n), v) * ds
a01 -= ufl.inner(ufl.outer(n, n) * ufl.dot(-p * ufl.Identity(u.ufl_shape[0]), n), v) * ds
L0 += ufl.inner(g_tau, v) * ds

a: list[list[dolfinx.fem.Form]] = dolfinx.fem.form(
    [
        [a00, a01],
        [a10, a11],
    ]
)
L: list[dolfinx.fem.Form] = dolfinx.fem.form([L0, L1])

# Assemble LHS matrix and RHS vector
with dolfinx.common.Timer("~Stokes: Assemble LHS and RHS"):
    A = dolfinx_mpc.assemble_matrix(a, [mpc, mpc_q], bcs, kind="nest")
    b = dolfinx_mpc.assemble_vector(L, [mpc, mpc_q], kind="nest")

# Set Dirichlet boundary condition values in the RHS
bcs1 = dolfinx.fem.bcs_by_block(dolfinx.fem.extract_function_spaces(a, 1), bcs)
dolfinx_mpc.apply_lifting(b, a, bcs1, constraint=[mpc, mpc_q])
for b_sub in b.getNestSubVecs():
    b_sub.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)  # type: ignore
bcs0 = dolfinx.fem.bcs_by_block(dolfinx.fem.extract_function_spaces(L), bcs)
dolfinx.fem.petsc.set_bc(b, bcs0)

# Preconditioner
P11 = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(ufl.inner(p, q) * ufl.dx))
P = PETSc.Mat().createNest([[A.getNestSubMatrix(0, 0), None], [None, P11]])  # type: ignore
P.assemble()

Solve variational problem#

ksp = PETSc.KSP().create(mesh.comm)  # type: ignore
ksp.setOperators(A, P)
ksp.setMonitor(
    lambda ctx, it, r: PETSc.Sys.Print(  # type: ignore
        f"Iteration: {it:>4d}, |r| = {r:.3e}"
    )
)
ksp.setType("minres")
tol = 2 * np.finfo(mesh.geometry.x.dtype).eps
ksp.setTolerances(rtol=tol, atol=tol)
ksp.getPC().setType("fieldsplit")
ksp.getPC().setFieldSplitType(PETSc.PC.CompositeType.ADDITIVE)  # type: ignore

nested_IS = P.getNestISs()
ksp.getPC().setFieldSplitIS(("u", nested_IS[0][0]), ("p", nested_IS[0][1]))

ksp_u, ksp_p = ksp.getPC().getFieldSplitSubKSP()
ksp_u.setType("preonly")
ksp_u.setTolerances(rtol=tol, atol=tol)
ksp_u.getPC().setType("gamg")
ksp_p.setType("preonly")
ksp_p.getPC().setType("jacobi")
ksp_p.setTolerances(rtol=tol, atol=tol)
ksp.setFromOptions()

Uh = b.copy()
ksp.solve(b, Uh)

for Uh_sub in Uh.getNestSubVecs():
    Uh_sub.ghostUpdate(
        addv=PETSc.InsertMode.INSERT,  # type: ignore
        mode=PETSc.ScatterMode.FORWARD,  # type: ignore
    )  # type: ignore

Hide code cell output

Iteration:    0, |r| = 1.429e+01
Iteration:    1, |r| = 5.283e+00
Iteration:    2, |r| = 5.262e+00
Iteration:    3, |r| = 2.000e+00
Iteration:    4, |r| = 1.491e+00
Iteration:    5, |r| = 1.307e+00
Iteration:    6, |r| = 7.062e-01
Iteration:    7, |r| = 5.817e-01
Iteration:    8, |r| = 5.281e-01
Iteration:    9, |r| = 3.995e-01
Iteration:   10, |r| = 3.950e-01
Iteration:   11, |r| = 3.090e-01
Iteration:   12, |r| = 2.403e-01
Iteration:   13, |r| = 2.343e-01
Iteration:   14, |r| = 1.549e-01
Iteration:   15, |r| = 1.462e-01
Iteration:   16, |r| = 1.198e-01
Iteration:   17, |r| = 9.283e-02
Iteration:   18, |r| = 9.268e-02
Iteration:   19, |r| = 7.094e-02
Iteration:   20, |r| = 5.685e-02
Iteration:   21, |r| = 5.627e-02
Iteration:   22, |r| = 4.370e-02
Iteration:   23, |r| = 4.200e-02
Iteration:   24, |r| = 3.575e-02
Iteration:   25, |r| = 2.664e-02
Iteration:   26, |r| = 2.661e-02
Iteration:   27, |r| = 2.111e-02
Iteration:   28, |r| = 1.916e-02
Iteration:   29, |r| = 1.808e-02
Iteration:   30, |r| = 1.391e-02
Iteration:   31, |r| = 1.276e-02
Iteration:   32, |r| = 1.195e-02
Iteration:   33, |r| = 8.872e-03
Iteration:   34, |r| = 8.443e-03
Iteration:   35, |r| = 7.097e-03
Iteration:   36, |r| = 5.463e-03
Iteration:   37, |r| = 5.440e-03
Iteration:   38, |r| = 4.080e-03
Iteration:   39, |r| = 3.549e-03
Iteration:   40, |r| = 3.459e-03
Iteration:   41, |r| = 2.637e-03
Iteration:   42, |r| = 2.517e-03
Iteration:   43, |r| = 2.125e-03
Iteration:   44, |r| = 1.635e-03
Iteration:   45, |r| = 1.634e-03
Iteration:   46, |r| = 1.189e-03
Iteration:   47, |r| = 9.547e-04
Iteration:   48, |r| = 9.306e-04
Iteration:   49, |r| = 6.346e-04
Iteration:   50, |r| = 5.522e-04
Iteration:   51, |r| = 5.058e-04
Iteration:   52, |r| = 3.324e-04
Iteration:   53, |r| = 3.170e-04
Iteration:   54, |r| = 2.697e-04
Iteration:   55, |r| = 1.934e-04
Iteration:   56, |r| = 1.892e-04
Iteration:   57, |r| = 1.549e-04
Iteration:   58, |r| = 1.152e-04
Iteration:   59, |r| = 1.152e-04
Iteration:   60, |r| = 8.718e-05
Iteration:   61, |r| = 6.428e-05
Iteration:   62, |r| = 6.424e-05
Iteration:   63, |r| = 4.624e-05
Iteration:   64, |r| = 3.659e-05
Iteration:   65, |r| = 3.572e-05
Iteration:   66, |r| = 2.350e-05
Iteration:   67, |r| = 2.054e-05
Iteration:   68, |r| = 1.863e-05
Iteration:   69, |r| = 1.228e-05
Iteration:   70, |r| = 1.171e-05
Iteration:   71, |r| = 9.059e-06
Iteration:   72, |r| = 6.164e-06
Iteration:   73, |r| = 6.162e-06
Iteration:   74, |r| = 4.096e-06
Iteration:   75, |r| = 3.061e-06
Iteration:   76, |r| = 2.977e-06
Iteration:   77, |r| = 1.852e-06
Iteration:   78, |r| = 1.572e-06
Iteration:   79, |r| = 1.456e-06
Iteration:   80, |r| = 9.679e-07
Iteration:   81, |r| = 8.637e-07
Iteration:   82, |r| = 7.667e-07
Iteration:   83, |r| = 5.196e-07
Iteration:   84, |r| = 5.152e-07
Iteration:   85, |r| = 3.863e-07
Iteration:   86, |r| = 3.602e-07
Iteration:   87, |r| = 2.895e-07
Iteration:   88, |r| = 1.771e-07
Iteration:   89, |r| = 1.765e-07
Iteration:   90, |r| = 1.197e-07
Iteration:   91, |r| = 9.171e-08
Iteration:   92, |r| = 9.018e-08
Iteration:   93, |r| = 5.844e-08
Iteration:   94, |r| = 4.590e-08
Iteration:   95, |r| = 4.297e-08
Iteration:   96, |r| = 2.643e-08
Iteration:   97, |r| = 2.486e-08
Iteration:   98, |r| = 1.959e-08
Iteration:   99, |r| = 1.616e-08
Iteration:  100, |r| = 1.456e-08
Iteration:  101, |r| = 9.692e-09
Iteration:  102, |r| = 9.298e-09
Iteration:  103, |r| = 7.971e-09
Iteration:  104, |r| = 5.392e-09
Iteration:  105, |r| = 4.994e-09
Iteration:  106, |r| = 4.340e-09
Iteration:  107, |r| = 2.795e-09
Iteration:  108, |r| = 2.678e-09
Iteration:  109, |r| = 2.087e-09
Iteration:  110, |r| = 1.739e-09
Iteration:  111, |r| = 1.539e-09
Iteration:  112, |r| = 1.042e-09
Iteration:  113, |r| = 1.042e-09
Iteration:  114, |r| = 7.192e-10
Iteration:  115, |r| = 5.466e-10
Iteration:  116, |r| = 5.369e-10
Iteration:  117, |r| = 3.679e-10
Iteration:  118, |r| = 3.027e-10
Iteration:  119, |r| = 2.892e-10
Iteration:  120, |r| = 1.909e-10
Iteration:  121, |r| = 1.626e-10
Iteration:  122, |r| = 1.546e-10
Iteration:  123, |r| = 9.862e-11
Iteration:  124, |r| = 8.149e-11
Iteration:  125, |r| = 7.503e-11
Iteration:  126, |r| = 4.660e-11
Iteration:  127, |r| = 4.418e-11
Iteration:  128, |r| = 3.469e-11
Iteration:  129, |r| = 2.749e-11
Iteration:  130, |r| = 2.562e-11
Iteration:  131, |r| = 1.599e-11
Iteration:  132, |r| = 1.501e-11
Iteration:  133, |r| = 1.258e-11
Iteration:  134, |r| = 8.796e-12
Iteration:  135, |r| = 8.751e-12
Iteration:  136, |r| = 6.568e-12
Iteration:  137, |r| = 5.360e-12
Iteration:  138, |r| = 5.072e-12
Iteration:  139, |r| = 3.094e-12
Iteration:  140, |r| = 2.680e-12
Iteration:  141, |r| = 2.313e-12
Iteration:  142, |r| = 1.375e-12
Iteration:  143, |r| = 1.290e-12
Iteration:  144, |r| = 9.976e-13
Iteration:  145, |r| = 5.916e-13
Iteration:  146, |r| = 5.766e-13
Iteration:  147, |r| = 4.291e-13
Iteration:  148, |r| = 3.000e-13
Iteration:  149, |r| = 3.000e-13
Iteration:  150, |r| = 2.113e-13
Iteration:  151, |r| = 1.888e-13
Iteration:  152, |r| = 1.562e-13
Iteration:  153, |r| = 1.161e-13
Iteration:  154, |r| = 1.150e-13
Iteration:  155, |r| = 8.070e-14
Iteration:  156, |r| = 7.176e-14
Iteration:  157, |r| = 6.351e-14
Iteration:  158, |r| = 4.027e-14
Iteration:  159, |r| = 3.866e-14
Iteration:  160, |r| = 3.006e-14
Iteration:  161, |r| = 1.892e-14
Iteration:  162, |r| = 1.843e-14
Iteration:  163, |r| = 1.399e-14
Iteration:  164, |r| = 9.139e-15
Iteration:  165, |r| = 9.128e-15
Iteration:  166, |r| = 6.027e-15

Put NestVec into DOLFINx Function#

uh = dolfinx.fem.Function(mpc.function_space)
ph = dolfinx.fem.Function(mpc_q.function_space)
dolfinx.fem.petsc.assign(Uh, [uh, ph])
uh.x.scatter_forward()
ph.x.scatter_forward()

# Backsubstitute to update slave dofs in solution vector
mpc.backsubstitution(uh)
mpc_q.backsubstitution(ph)

Output#

uh.name = "u"
ph.name = "p"
outdir = Path("results")
outdir.mkdir(exist_ok=True, parents=True)

with dolfinx.io.XDMFFile(mesh.comm, outdir / "demo_stokes_nest.xdmf", "w") as outfile:
    outfile.write_mesh(mesh)
    outfile.write_meshtags(mt, mesh.geometry)
    outfile.write_function(ph)

with dolfinx.io.VTXWriter(mesh.comm, outdir / "stokes_nest_uh.bp", uh, engine="BP4") as vtx:
    vtx.write(0.0)

The velocity, drawn as arrows coloured by the speed. Along the tilted walls the flow runs parallel to the boundary, which is the slip condition. Each process draws the cells it owns, which are gathered on the first process. uh lives in the extended space of the constraint, whose leading entries are those of V.

pyvista.global_theme.allow_empty_mesh = True
owned_cells = np.arange(mesh.topology.index_map(mesh.topology.dim).size_local, dtype=np.int32)
grid = pyvista.UnstructuredGrid(*plot.vtk_mesh(V, entities=owned_cells))
velocity = np.zeros((grid.n_points, 3))
velocity[:, : mesh.geometry.dim] = uh.x.array.real[: grid.n_points * mesh.geometry.dim].reshape(grid.n_points, -1)
grid.point_data["u"] = velocity
grid.point_data["|u|"] = np.linalg.norm(velocity, axis=1)
pieces = mesh.comm.gather(grid, root=0)
if pieces is not None:  # only the first process received the grids
    pieces = [piece for piece in pieces if piece.n_points > 0]
    clim = [min(p["|u|"].min(initial=np.inf) for p in pieces), max(p["|u|"].max(initial=0) for p in pieces)]
    plotter = pyvista.Plotter(window_size=[800, 500])
    for piece in pieces:
        plotter.add_mesh(piece, color="lightgray", opacity=0.4)
        plotter.add_mesh(
            piece.glyph(orient="u", scale="|u|", factor=0.08),
            scalars="|u|",
            cmap="viridis",
            clim=clim,
            scalar_bar_args={"title": "speed"},
        )
    plotter.view_xy()
    plotter.camera.tight(padding=0.15, view="xy", adjust_render_window=False)
    # The figure is named after the demo, as the gallery of the documentation expects
    figure = Path("demo_stokes_nest.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_3965/3751583190.py:30: 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"))

Hide code cell source

# Verification
# Transfer data from the MPC problem to numpy arrays for comparison
with dolfinx.common.Timer("~Stokes: Verification of problem by global matrix reduction"):
    W = dolfinx.fem.functionspace(mesh, basix.ufl.mixed_element([Ve, Qe]))
    V, V_to_W = W.sub(0).collapse()
    _, Q_to_W = W.sub(1).collapse()

    # Inlet velocity Dirichlet BC
    inlet_velocity = dolfinx.fem.Function(V)
    inlet_velocity.interpolate(inlet_velocity_expression)
    inlet_velocity.x.scatter_forward()
    W0 = W.sub(0)
    dofs = dolfinx.fem.locate_dofs_topological((W0, V), 1, mt.find(3))
    bc1 = dolfinx.fem.dirichletbc(inlet_velocity, dofs, W0)

    # Collect Dirichlet boundary conditions
    bcs = [bc1]

    # Slip conditions for walls
    n = dolfinx_mpc.utils.create_normal_approximation(V, mt, 1)
    with dolfinx.common.Timer("~Stokes: Create slip constraint"):
        mpc = dolfinx_mpc.MultiPointConstraint(W)
        mpc.create_slip_constraint(W.sub(0), (mt, 1), n, bcs=bcs)
    mpc.finalize()

    (u, p) = ufl.TrialFunctions(W)
    (v, q) = ufl.TestFunctions(W)
    a = (2 * mu * ufl.inner(sym_grad(u), sym_grad(v)) - ufl.inner(p, ufl.div(v)) - ufl.inner(ufl.div(u), q)) * ufl.dx
    L = ufl.inner(f, v) * ufl.dx

    # Terms due to slip condition
    # Explained in for instance: https://arxiv.org/pdf/2001.10639.pdf
    a -= ufl.inner(ufl.outer(n, n) * ufl.dot(T(u, p, mu), n), v) * ds
    L += ufl.inner(g_tau, v) * ds

    af = dolfinx.fem.form(a)
    Lf = dolfinx.fem.form(L)

    # Solve the MPC problem using a global transformation matrix
    # and numpy solvers to get reference values
    # Generate reference matrices and unconstrained solution
    A_org = dolfinx.fem.petsc.assemble_matrix(af, bcs)
    A_org.assemble()
    L_org = dolfinx.fem.petsc.assemble_vector(Lf)

    dolfinx.fem.petsc.apply_lifting(L_org, [af], [bcs])
    L_org.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)  # type: ignore
    dolfinx.fem.petsc.set_bc(L_org, bcs)
    root = 0

    # Gather LHS, RHS and solution on one process
    A_csr = dolfinx_mpc.utils.gather_PETScMatrix(A_org, root=root)
    K = dolfinx_mpc.utils.gather_transformation_matrix(mpc, root=root)
    L_np = dolfinx_mpc.utils.gather_PETScVector(L_org, root=root)

    u_mpc = dolfinx_mpc.utils.gather_PETScVector(uh.x.petsc_vec, root=root)
    p_mpc = dolfinx_mpc.utils.gather_PETScVector(ph.x.petsc_vec, root=root)
    up_mpc = np.hstack([u_mpc, p_mpc])
    if MPI.COMM_WORLD.rank == root:
        KTAK = K.T * A_csr * K
        reduced_L = K.T @ L_np
        # Solve linear system
        d = scipy.sparse.linalg.spsolve(KTAK, reduced_L)
        # Back substitution to full solution vector
        uh_numpy = K @ d
        np.testing.assert_allclose(
            np.linalg.norm(uh_numpy, 2), np.linalg.norm(up_mpc, 2), atol=1e3 * tol, rtol=1e3 * tol
        )
A.destroy()
b.destroy()
for Uh_sub in Uh.getNestSubVecs():
    Uh_sub.destroy()
Uh.destroy()
ksp.destroy()
<petsc4py.PETSc.KSP at 0x7f2edfbd6980>

Expand the cell below to inspect timings of various functions

Hide code cell output

[MPI_MAX] Summary of timings (s)                                            |  reps       avg       tot
-------------------------------------------------------------------------------------------------------
Compute connectivity 1-0                                                    |     1  0.000005  0.000005
Compute connectivity 2-0                                                    |     1  0.000001  0.000001
Compute entities by key matching: number entities                           |     1  0.000085  0.000085
Compute entities of dim = 1                                                 |     1  0.000175  0.000175
Compute ghost indices                                                       |     1  0.000032  0.000032
Compute local part of mesh dual graph                                       |     1  0.000772  0.000772
Compute local part of mesh dual graph: 0                                    |     1  0.000007  0.000007
Compute local part of mesh dual graph: 1                                    |     1  0.000038  0.000038
Compute local part of mesh dual graph: 3                                    |     1  0.000704  0.000704
Compute local part of mesh dual graph: 4                                    |     1  0.000012  0.000012
Compute local part of mesh dual graph: 5                                    |     1  0.000007  0.000007
Compute local-to-local map                                                  |     1  0.000003  0.000003
Compute-local-to-global links for global/local adjacency list               |     1  0.000007  0.000007
Distribute row-wise data (scalable)                                         |     1  0.000040  0.000040
Dofmap builder: build dofmap data                                           |     6  0.000166  0.000995
Dofmap builder: compute dof reordering map                                  |     6  0.000083  0.000497
Dofmap builder: get dofmap global indices                                   |     6  0.000054  0.000325
Dofmap builder: init dofmap from element dofmap                             |     6  0.000018  0.000107
Entity local indexing                                                       |     1  0.000056  0.000056
Entity local indexing: build candidates                                     |     1  0.000012  0.000012
Entity local indexing: candidate exchange                                   |     1  0.000005  0.000005
Entity local indexing: global index exchange                                |     1  0.000002  0.000002
Entity local indexing: neighbourhood setup                                  |     1  0.000023  0.000023
Entity local indexing: ownership                                            |     1  0.000003  0.000003
Entity local indexing: renumber                                             |     1  0.000001  0.000001
Graph: create_level_structure                                               |     7  0.000011  0.000076
Reverse Cuthill-McKee ordering                                              |     3  0.000045  0.000135
SparsityPattern::finalize                                                   |     6  0.000160  0.000962
Threaded part                                                               |     1  0.000017  0.000017
Topology: 0                                                                 |     1  0.000004  0.000004
Topology: 1                                                                 |     1  0.000011  0.000011
Topology: 2                                                                 |     1  0.000004  0.000004
Topology: 4                                                                 |     1  0.000043  0.000043
Topology: 5                                                                 |     1  0.000010  0.000010
Topology: 6                                                                 |     1  0.000016  0.000016
Topology: 7                                                                 |     1  0.000002  0.000002
Topology: build and sort transposed index list                              |     1  0.000003  0.000003
Topology: create