This tutorial is part of multiphenicsx.

Tutorial 06, case 3a: advection diffusion reaction control problem with Neumann control

In this tutorial we solve the optimal control problem

$$\min J(y, u) = \frac{1}{2} \int_{\Omega} (y - y_d)^2 dx + \frac{\alpha}{2} \int_{\Gamma_2} u^2 ds$$ s.t. $$\begin{cases} - \epsilon \Delta y + \beta \cdot \nabla y + \sigma y = f & \text{in } \Omega\\ \epsilon \partial_n y = 0 & \text{on } \Gamma_1\\ \epsilon \partial_n y = u & \text{on } \Gamma_2\\ \epsilon \partial_n y = 0 & \text{on } \Gamma_3\\ y = 0 & \text{on } \Gamma_4 \end{cases}$$

where $$\begin{align*} & \Omega & \text{unit square}\\ & \Gamma_1 & \text{bottom boundary of the square}\\ & \Gamma_2 & \text{left boundary of the square}\\ & \Gamma_3 & \text{top boundary of the square}\\ & \Gamma_4 & \text{right boundary of the square}\\ & u \in L^2(\Gamma_2) & \text{control variable}\\ & y \in H^1(\Omega) & \text{state variable}\\ & \alpha > 0 & \text{penalization parameter}\\ & y_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 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_2 = boundaries.indices[boundaries.values == 2]
boundaries_4 = boundaries.indices[boundaries.values == 4]

# Define associated measures
ds = ufl.Measure("ds", subdomain_data=boundaries)

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 = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
U = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
Q = Y.clone()

Create element
Create element

Restrictions

dofs_Y = np.arange(0, Y.dofmap.index_map.size_local + Y.dofmap.index_map.num_ghosts)
dofs_U = dolfinx.fem.locate_dofs_topological(U, boundaries.dim, boundaries_2)
dofs_Q = dofs_Y
restriction_Y = multiphenicsx.fem.DofMapRestriction(Y.dofmap, dofs_Y)
restriction_U = multiphenicsx.fem.DofMapRestriction(U.dofmap, dofs_U)
restriction_Q = multiphenicsx.fem.DofMapRestriction(Q.dofmap, dofs_Q)
restriction = [restriction_Y, restriction_U, restriction_Q]

Trial and test functions

(y, u, p) = (ufl.TrialFunction(Y), ufl.TrialFunction(U), ufl.TrialFunction(Q))
(z, v, q) = (ufl.TestFunction(Y), ufl.TestFunction(U), ufl.TestFunction(Q))

Problem data

alpha = 1.e-5
y_d = 1.
epsilon = 1.e-1
beta = ufl.as_vector((-1., -2.))
sigma = 1.
ff = 1.
bc0 = petsc4py.PETSc.ScalarType(0)

Optimality conditions

state_operator = (epsilon * ufl.inner(ufl.grad(y), ufl.grad(q)) * ufl.dx
                  + ufl.inner(ufl.dot(beta, ufl.grad(y)), q) * ufl.dx + sigma * ufl.inner(y, q) * ufl.dx)
adjoint_operator = (epsilon * ufl.inner(ufl.grad(p), ufl.grad(z)) * ufl.dx
                    - ufl.inner(ufl.dot(beta, ufl.grad(p)), z) * ufl.dx + sigma * ufl.inner(p, z) * ufl.dx)
a = [[ufl.inner(y, z) * ufl.dx, None, adjoint_operator],
     [None, alpha * ufl.inner(u, v) * ds(2), - ufl.inner(p, v) * ds(2)],
     [state_operator, - ufl.inner(u, q) * ds(2), None]]
f = [ufl.inner(y_d, z) * ufl.dx,
     None,
     ufl.inner(ff, q) * ufl.dx]
a[2][2] = dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)) * ufl.inner(p, q) * ufl.dx
f[1] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), v) * ufl.dx
a_cpp = dolfinx.fem.form(a)
f_cpp = dolfinx.fem.form(f)
bdofs_Y_4 = dolfinx.fem.locate_dofs_topological(Y, mesh.topology.dim - 1, boundaries_4)
bdofs_Q_4 = dolfinx.fem.locate_dofs_topological(Q, mesh.topology.dim - 1, boundaries_4)
bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_4, Y),
      dolfinx.fem.dirichletbc(bc0, bdofs_Q_4, Q)]

Solution

(y, u, p) = (dolfinx.fem.Function(Y), dolfinx.fem.Function(U), dolfinx.fem.Function(Q))

Cost functional

J = 0.5 * ufl.inner(y - y_d, y - y_d) * ufl.dx + 0.5 * alpha * ufl.inner(u, u) * ds(2)
J_cpp = dolfinx.fem.form(J)

Uncontrolled functional value

# Extract state forms from the optimality conditions
a_state = ufl.replace(a[2][0], {q: z})
f_state = ufl.replace(f[2], {q: z})
a_state_cpp = dolfinx.fem.form(a_state)
f_state_cpp = dolfinx.fem.form(f_state)
bc_state = [bc[0]]

# Assemble the linear system for the state
A_state = dolfinx.fem.petsc.assemble_matrix(a_state_cpp, bcs=bc_state)
A_state.assemble()
F_state = dolfinx.fem.petsc.assemble_vector(f_state_cpp)
dolfinx.fem.petsc.apply_lifting(F_state, [a_state_cpp], [bc_state])
F_state.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)
dolfinx.fem.petsc.set_bc(F_state, bc_state)

# Solve
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, y.x.petsc_vec)
y.x.petsc_vec.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()

<petsc4py.PETSc.KSP at 0x7fbb68df2980>

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.23058804)

Uncontrolled J = 0.23058804871570932

viskex.dolfinx.plot_scalar_field(y, "uncontrolled state")

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 = multiphenicsx.fem.petsc.assemble_matrix_block(a_cpp, bcs=bc, restriction=(restriction, restriction))
A.assemble()
F = multiphenicsx.fem.petsc.assemble_vector_block(f_cpp, a_cpp, bcs=bc, restriction=restriction)

# Solve
yup = multiphenicsx.fem.petsc.create_vector_block(f_cpp, restriction=restriction)
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, yup)
yup.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()

<petsc4py.PETSc.KSP at 0x7fbb68e48ef0>

# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(yup, [Y.dofmap, U.dofmap, Q.dofmap], restriction) as yup_wrapper:
    for yup_wrapper_local, component in zip(yup_wrapper, (y, u, p)):
        with component.x.petsc_vec.localForm() as component_local:
            component_local[:] = yup_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.21175842)

Optimal J = 0.21175842908207035

viskex.dolfinx.plot_scalar_field(y, "state")

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(u, "control")

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, "adjoint")

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