Compute the sensitivity in parallel with DOLFINx

dolfinx
9

Use dolfinx.fem.petsc.assemble_ instead of dolfinx.fem.assemble_ resolves the issue.
MWE:

import dolfinx
import dolfinx.fem.petsc
import ufl
from dolfinx.fem import form
from dolfinx.mesh import CellType
from dolfinx.fem import functionspace, Function, DirichletBC
from dolfinx.nls.petsc import NewtonSolver
from dolfinx.fem.petsc import NonlinearProblem, LinearProblem
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np

# mesh and function space
mesh = dolfinx.mesh.create_rectangle(MPI.COMM_WORLD, np.array([[0, 0], [1, 1]]), [3, 3], cell_type=CellType.triangle)

V = functionspace(mesh, ("CG", 1, (2,)))
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(P, ufl.grad(v))*dx

# boundary conditions
def left(x): return np.isclose(x[0], 0)
def right(x): return np.isclose(x[0], 1)

fdim = mesh.topology.dim - 1

u1 = Function(V)
with u1.vector.localForm() as loc:
    loc.set(0)

boundary_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, left)
bc1 = dolfinx.fem.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 = dolfinx.fem.dirichletbc(u2, dolfinx.fem.locate_dofs_topological(
    V, fdim, boundary_facets))
bcs = [bc1, bc2]

# solve the problem
problem = 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(form(J))

# partial derivative of J w.r.t. mu
dJdmu = dolfinx.fem.petsc.assemble_vector(form(ufl.derivative(J, mu)))
# dJdmu.assemble()

# partial derivative of R w.r.t. mu
dRdmu_form = form(ufl.adjoint(ufl.derivative(R, mu)))
dRdmu = dolfinx.fem.petsc.assemble_matrix(dRdmu_form)  # partial derivative
dRdmu.assemble()

# 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 = LinearProblem(lhs, rhs, bcs=bcs)
lmbda = problem.solve()

dJdmu += dRdmu*lmbda.vector
dJdmu.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
print(dJdmu.array)
2 Likes