This tutorial is part of multiphenicsx.

Tutorial 06, case 5: Stokes problem with distributed control

In this tutorial we solve the optimal control problem

$$\min J(y, u) = \frac{1}{2} \int_{\Omega} |v - v_d|^2 dx + \frac{\alpha}{2} \int_{\Omega} |u|^2 dx$$ s.t. $$\begin{cases} - \Delta v + \nabla p = f + u & \text{in } \Omega\\ \text{div} v = 0 & \text{in } \Omega\\ v = 0 & \text{on } \partial\Omega \end{cases}$$

where $$\begin{align*} & \Omega & \text{unit square}\\ & u \in [L^2(\Omega)]^2 & \text{control variable}\\ & v \in [H^1_0(\Omega)]^2 & \text{state velocity variable}\\ & p \in L^2(\Omega) & \text{state pressure variable}\\ & \alpha > 0 & \text{penalization parameter}\\ & v_d & \text{desired state}\\ & f & \text{forcing term} \end{align*}$$ using an adjoint formulation solved by a one shot approach

import dolfinx.fem
import dolfinx.fem.petsc
import dolfinx.io
import dolfinx.mesh
import mpi4py.MPI
import numpy as np
import numpy.typing
import petsc4py.PETSc
import sympy
import ufl
import viskex

import multiphenicsx.fem
import multiphenicsx.fem.petsc

Mesh

mesh = dolfinx.mesh.create_unit_square(mpi4py.MPI.COMM_WORLD, 32, 32)

# Create connectivities required by the rest of the code
mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)

def bottom(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the bottom boundary."""
    return abs(x[1] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]


def left(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the left boundary."""
    return abs(x[0] - 0.) < np.finfo(float).eps  # type: ignore[no-any-return]


def top(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the top boundary."""
    return abs(x[1] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]


def right(x: np.typing.NDArray[np.float64]) -> np.typing.NDArray[np.bool_]:
    """Condition that defines the right boundary."""
    return abs(x[0] - 1.) < np.finfo(float).eps  # type: ignore[no-any-return]


boundaries_entities = dict()
boundaries_values = dict()
for (boundary, boundary_id) in zip((bottom, left, top, right), (1, 2, 3, 4)):
    boundaries_entities[boundary_id] = dolfinx.mesh.locate_entities_boundary(
        mesh, mesh.topology.dim - 1, boundary)
    boundaries_values[boundary_id] = np.full(
        boundaries_entities[boundary_id].shape, boundary_id, dtype=np.int32)
boundaries_entities_unsorted = np.hstack(list(boundaries_entities.values()))
boundaries_values_unsorted = np.hstack(list(boundaries_values.values()))
boundaries_entities_argsort = np.argsort(boundaries_entities_unsorted)
boundaries_entities_sorted = boundaries_entities_unsorted[boundaries_entities_argsort]
boundaries_values_sorted = boundaries_values_unsorted[boundaries_entities_argsort]
boundaries = dolfinx.mesh.meshtags(
    mesh, mesh.topology.dim - 1,
    boundaries_entities_sorted, boundaries_values_sorted)
boundaries.name = "boundaries"

boundaries_1234 = boundaries.indices[np.isin(boundaries.values, (1, 2, 3, 4))]

viskex.dolfinx.plot_mesh(mesh)

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

viskex.dolfinx.plot_mesh_tags(mesh, boundaries, "boundaries")

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

Function spaces

Y_velocity = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, (mesh.geometry.dim, )))
Y_pressure = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
U = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, (mesh.geometry.dim, )))
Q_velocity = Y_velocity.clone()
Q_pressure = Y_pressure.clone()

Trial and test functions

(v, p) = (ufl.TrialFunction(Y_velocity), ufl.TrialFunction(Y_pressure))
(w, q) = (ufl.TestFunction(Y_velocity), ufl.TestFunction(Y_pressure))
u = ufl.TrialFunction(U)
r = ufl.TestFunction(U)
(z, b) = (ufl.TrialFunction(Q_velocity), ufl.TrialFunction(Q_pressure))
(s, d) = (ufl.TestFunction(Q_velocity), ufl.TestFunction(Q_pressure))

Problem data

alpha = 1.e-5
epsilon = 1.e-5
x, y = sympy.symbols("x[0], x[1]")
psi_d = 10 * (1 - sympy.cos(0.8 * np.pi * x)) * (1 - sympy.cos(0.8 * np.pi * y)) * (1 - x)**2 * (1 - y)**2
v_d_x = sympy.lambdify([x, y], psi_d.diff(y, 1))
v_d_y = sympy.lambdify([x, y], - psi_d.diff(x, 1))
v_d = dolfinx.fem.Function(Y_velocity)
v_d.interpolate(lambda x: np.stack((v_d_x(x[0], x[1]), v_d_y(x[0], x[1])), axis=0))
ff = dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2)))
bc0 = np.zeros((2, ), dtype=petsc4py.PETSc.ScalarType)

Optimality conditions

a = [[ufl.inner(v, w) * ufl.dx, None, None, ufl.inner(ufl.grad(z), ufl.grad(w)) * ufl.dx,
      - ufl.inner(b, ufl.div(w)) * ufl.dx],
     [None, None, None, - ufl.inner(ufl.div(z), q) * ufl.dx, epsilon * ufl.inner(b, q) * ufl.dx],
     [None, None, alpha * ufl.inner(u, r) * ufl.dx, - ufl.inner(z, r) * ufl.dx, None],
     [ufl.inner(ufl.grad(v), ufl.grad(s)) * ufl.dx, - ufl.inner(p, ufl.div(s)) * ufl.dx,
      - ufl.inner(u, s) * ufl.dx, None, None],
     [- ufl.inner(ufl.div(v), d) * ufl.dx, epsilon * ufl.inner(p, d) * ufl.dx, None, None, None]]
f = [ufl.inner(v_d, w) * ufl.dx,
     None,
     None,
     ufl.inner(ff, s) * ufl.dx,
     None]
a[3][3] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(z, s) * ufl.dx
f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), q) * ufl.dx
f[2] = ufl.inner(dolfinx.fem.Constant(mesh, tuple(petsc4py.PETSc.ScalarType(0) for _ in range(2))), r) * ufl.dx
f[4] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), d) * ufl.dx
a_cpp = dolfinx.fem.form(a)
f_cpp = dolfinx.fem.form(f)
bdofs_Y_velocity_1234 = dolfinx.fem.locate_dofs_topological(
    Y_velocity, mesh.topology.dim - 1, boundaries_1234)
bdofs_Q_velocity_1234 = dolfinx.fem.locate_dofs_topological(
    Q_velocity, mesh.topology.dim - 1, boundaries_1234)
bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_velocity_1234, Y_velocity),
      dolfinx.fem.dirichletbc(bc0, bdofs_Q_velocity_1234, Q_velocity)]

Solution

(v, p) = (dolfinx.fem.Function(Y_velocity), dolfinx.fem.Function(Y_pressure))
u = dolfinx.fem.Function(U)
(z, b) = (dolfinx.fem.Function(Q_velocity), dolfinx.fem.Function(Q_pressure))

Cost functional

J = 0.5 * ufl.inner(v - v_d, v - v_d) * ufl.dx + 0.5 * alpha * ufl.inner(u, u) * ufl.dx
J_cpp = dolfinx.fem.form(J)

Uncontrolled functional value

# Extract state forms from the optimality conditions
a_state = [[ufl.replace(a[i][j], {s: w, d: q}) if a[i][j] is not None else None
            for j in (0, 1)] for i in (3, 4)]
f_state = [ufl.replace(f[i], {s: w, d: q}) for i in (3, 4)]
a_state_cpp = dolfinx.fem.form(a_state)
f_state_cpp = dolfinx.fem.form(f_state)
bc_state = [bc[0]]

# Assemble the block linear system for the state
A_state = dolfinx.fem.petsc.assemble_matrix_block(a_state_cpp, bcs=bc_state)
A_state.assemble()
F_state = dolfinx.fem.petsc.assemble_vector_block(f_state_cpp, a_state_cpp, bcs=bc_state)

# Solve
vp = dolfinx.fem.petsc.create_vector_block([f_cpp[j] for j in (0, 1)])
ksp = petsc4py.PETSc.KSP()
ksp.create(mesh.comm)
ksp.setOperators(A_state)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(F_state, vp)
vp.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()

<petsc4py.PETSc.KSP at 0x7f94991bf150>

# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(vp, [c.function_space.dofmap for c in (v, p)]) as vp_wrapper:
    for vp_wrapper_local, component in zip(vp_wrapper, (v, p)):
        with component.x.petsc_vec.localForm() as component_local:
            component_local[:] = vp_wrapper_local

J_uncontrolled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)
print("Uncontrolled J =", J_uncontrolled)
assert np.isclose(J_uncontrolled, 0.1784536)

Uncontrolled J = 0.17845361493420545

viskex.dolfinx.plot_vector_field(v, "uncontrolled state velocity", glyph_factor=1e-1)

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

viskex.dolfinx.plot_scalar_field(p, "uncontrolled state pressure")

error: XDG_RUNTIME_DIR is invalid or not set in the environment.

MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

Optimal control

# Assemble the block linear system for the optimality conditions
A = dolfinx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc)
A.assemble()
F = dolfinx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc)

# Solve
vpuzb = dolfinx.fem.petsc.create_vector_block(f_cpp)
ksp = petsc4py.PETSc.KSP()
ksp.create(mesh.comm)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setFromOptions()
ksp.solve(F, vpuzb)
vpuzb.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()

<petsc4py.PETSc.KSP at 0x7f9484970860>

# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(
        vpuzb, [c.function_space.dofmap for c in (v, p, u, z, b)]) as vpuzb_wrapper:
    for vpuzb_wrapper_local, component in zip(vpuzb_wrapper, (v, p, u, z, b)):
        with component.x.petsc_vec.localForm() as component_local:
            component_local[:] = vpuzb_wrapper_local

J_controlled = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(J_cpp), op=mpi4py.MPI.SUM)
print("Optimal J =", J_controlled)
assert np.isclose(J_controlled, 0.0052941)

Optimal J = 0.00529410631025776

viskex.dolfinx.plot_vector_field(v, "state velocity", glyph_factor=1e-1)

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

viskex.dolfinx.plot_scalar_field(p, "state pressure")

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

viskex.dolfinx.plot_vector_field(u, "control", glyph_factor=1e-3)

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

viskex.dolfinx.plot_vector_field(z, "adjoint velocity", glyph_factor=1e2)

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen

viskex.dolfinx.plot_scalar_field(b, "adjoint pressure")

error: XDG_RUNTIME_DIR is invalid or not set in the environment.
MESA: error: ZINK: failed to choose pdev
glx: failed to create drisw screen