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} - \Delta y = f & \text{in } \Omega\\ \partial_n y = 0 & \text{on } \Gamma_1\\ y = u & \text{on } \Gamma_2\\ \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 = 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]
boundaries_24 = boundaries.indices[np.isin(boundaries.values, (2, 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
Y = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
U = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
L = U.clone()
Q = Y.clone()
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_L = dofs_U
dofs_Q = dofs_Y
restriction_Y = multiphenicsx.fem.DofMapRestriction(Y.dofmap, dofs_Y)
restriction_U = multiphenicsx.fem.DofMapRestriction(U.dofmap, dofs_U)
restriction_L = multiphenicsx.fem.DofMapRestriction(L.dofmap, dofs_L)
restriction_Q = multiphenicsx.fem.DofMapRestriction(Q.dofmap, dofs_Q)
restriction = [restriction_Y, restriction_U, restriction_L, restriction_Q]
(y, u, l, p) = (ufl.TrialFunction(Y), ufl.TrialFunction(U), ufl.TrialFunction(L), ufl.TrialFunction(Q))
(z, v, m, q) = (ufl.TestFunction(Y), ufl.TestFunction(U), ufl.TestFunction(L), ufl.TestFunction(Q))
alpha = 1.e-5
y_d = 1.
x = ufl.SpatialCoordinate(mesh)
ff = 10 * ufl.sin(2 * ufl.pi * x[0]) * ufl.sin(2 * ufl.pi * x[1])
bc0 = petsc4py.PETSc.ScalarType(0)
a = [[ufl.inner(y, z) * ufl.dx, None, ufl.inner(l, z) * ds(2), ufl.inner(ufl.grad(p), ufl.grad(z)) * ufl.dx],
[None, alpha * ufl.inner(u, v) * ds(2), - ufl.inner(l, v) * ds(2), None],
[ufl.inner(y, m) * ds(2), - ufl.inner(u, m) * ds(2), None, None],
[ufl.inner(ufl.grad(y), ufl.grad(q)) * ufl.dx, None, None, None]]
f = [ufl.inner(y_d, z) * ufl.dx,
None,
None,
ufl.inner(ff, q) * ufl.dx]
a[3][3] = 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
f[2] = ufl.inner(dolfinx.fem.Constant(mesh, petsc4py.PETSc.ScalarType(0)), m) * 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_24 = dolfinx.fem.locate_dofs_topological(Q, mesh.topology.dim - 1, boundaries_24)
bc = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_4, Y),
dolfinx.fem.dirichletbc(bc0, bdofs_Q_24, Q)]
(y, u, l, p) = (dolfinx.fem.Function(Y), dolfinx.fem.Function(U), dolfinx.fem.Function(L), dolfinx.fem.Function(Q))
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)
# Extract state forms from the optimality conditions
a_state = ufl.replace(a[3][0], {q: z})
f_state = ufl.replace(f[3], {q: z})
a_state_cpp = dolfinx.fem.form(a_state)
f_state_cpp = dolfinx.fem.form(f_state)
bdofs_Y_24 = dolfinx.fem.locate_dofs_topological(Y, mesh.topology.dim - 1, boundaries_24)
bc_state = [dolfinx.fem.dirichletbc(bc0, bdofs_Y_24, Y)]
# 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 0x7fc5248f8860>
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.5038977)
Uncontrolled J = 0.5038976704195648
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
# 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
yulp = 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, yulp)
yulp.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
ksp.destroy()
<petsc4py.PETSc.KSP at 0x7fc5249385e0>
# Split the block solution in components
with multiphenicsx.fem.petsc.BlockVecSubVectorWrapper(
yulp, [Y.dofmap, U.dofmap, L.dofmap, Q.dofmap], restriction) as yulp_wrapper:
for yulp_wrapper_local, component in zip(yulp_wrapper, (y, u, l, p)):
with component.x.petsc_vec.localForm() as component_local:
component_local[:] = yulp_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.1281224)
Optimal J = 0.12812239588290494
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(l, "lambda")
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