Hi all, I am trying to compute the sensitivity with DOLFINx based on the adjoint method (lecture_12_sensitivities.pdf (fenicsproject.org)).
Given:
(1) PDE: \mathbf{R}(\mathbf{u})=\mathbf{0}
(2) Variable m
(3) Objective functional g(m,\mathbf{u}(m))
Compute the sensitivity:
\dfrac{\text{d}g(m,\mathbf{u}(m))}{\text{d}m} = \dfrac{\partial g}{\partial m}+\left(\dfrac{\partial \mathbf{R}}{\partial m}\right)^\text{T} \cdot \boldsymbol{\lambda}
where the adjoint vector \boldsymbol{\lambda} can be solved from
\left(\dfrac{\partial \mathbf{R}}{\partial \mathbf{u}}\right)^\text{T} \cdot \boldsymbol{\lambda} = -\dfrac{\partial g}{\partial \mathbf{u}}
I guess this is very similar to the dolfin-adjoint package. But it seems we do not have the corresponding package in DOLFINx.
I try to implement the adjoint method and it works well when running in serial. The MWE is as follows:
import dolfinx, ufl
from dolfinx.generation import RectangleMesh
from dolfinx.mesh import CellType
from dolfinx.fem import VectorFunctionSpace, FunctionSpace, Function, DirichletBC
from dolfinx.nls import NewtonSolver
from mpi4py import MPI
import numpy as np
import scipy.sparse
# mesh and function space
mesh = RectangleMesh(MPI.COMM_WORLD, [[0, 0, 0], [1, 1, 0]], \
[3, 3], CellType.quadrilateral)
V = VectorFunctionSpace(mesh, ("CG", 1))
S = FunctionSpace(mesh, ("DG", 0))
u = Function(V) # test function (displacment field)
v = ufl.TestFunction(V) # trial function
# kinematics
lambda_ = 0.5769
mu = Function(S) # in this case, we set the mu as a design variable
mu.vector.array = 0.3846
P = 2.0*mu*ufl.sym(ufl.grad(u)) + lambda_*ufl.tr(ufl.sym(ufl.grad(u)))*ufl.Identity(len(u))
dx = ufl.Measure("dx", metadata={"quadrature_degree": 4})
R = ufl.inner(ufl.grad(v), P)*dx
# boundary conditions
left = lambda x: np.isclose(x[0], 0)
right = lambda x: np.isclose(x[0], 1)
fdim = mesh.topology.dim - 1
u1 = Function(V)
boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, left)
with u1.vector.localForm() as loc:
loc.set(0)
bc1 = DirichletBC(u1, dolfinx.fem.locate_dofs_topological(V, fdim, boundary_facets))
u2 = Function(V)
boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, right)
with u2.vector.localForm() as loc:
loc.set(1)
bc2 = DirichletBC(u2, dolfinx.fem.locate_dofs_topological(V, fdim, boundary_facets))
bcs = [bc1, bc2]
# solve the problem
problem = dolfinx.fem.NonlinearProblem(R, u, bcs)
solver = NewtonSolver(MPI.COMM_WORLD, problem)
solver.solve(u)
# print(u.vector.array)
# set an arbitrary functinal
J = mu*ufl.dot(u, u)*dx
J_old = dolfinx.fem.assemble_scalar(J)
# partial derivative of J w.r.t. mu
dJdmu = dolfinx.fem.assemble_vector(ufl.derivative(J, mu))
dJdmu.assemble()
dJdmu = dJdmu.getArray()
# partial derivative of R w.r.t. mu
dRdmu = dolfinx.fem.assemble_matrix(ufl.derivative(R, mu)) # partial derivative
dRdmu.assemble()
ai, aj, av = dRdmu.getValuesCSR()
dRdmu = scipy.sparse.csr_matrix((av, aj, ai))
# reset the boundary condition
with u2.vector.localForm() as loc:
loc.set(0.0)
# solve the adjoint vector
lhs = ufl.adjoint(ufl.derivative(R, u))
rhs = -ufl.derivative(J, u)
problem = dolfinx.fem.LinearProblem(lhs, rhs, bcs=bcs)
lambda_vec = problem.solve().vector.array
# compute the sensitivity
dJdmu += dRdmu.T@lambda_vec
However, when running in parallel, I have to collect the displacement field to root 0 (based on Gather solutions in parallel in FEniCSX - FEniCS Project) and then evaluate the sensitivity in a global mesh.
So, my question is, is there a way to make the above code also support the parallel computation?