Hello!
I am trying to compute gradients of custom functions that depends on displacement (for topology optimization problems). when comparing serial and parallel runs, the “compliance = sum of strain energy in all elements” has same gradient values, but the “stress_sum = sum of von Mises stress in all elements” has different gradient values. I am using same class to compute function and gradients for both, so not sure why only mismatch occurs for stress. Below is a script to reproduce the results. Thank you for any help or suggestions! (This example runs with dolfinx v0.8.0)
import dolfinx as dfx
from dolfinx import fem
import dolfinx.fem.petsc
import numpy as np
from mpi4py import MPI
from petsc4py import PETSc
import ufl
import basix
comm = MPI.COMM_WORLD
# ------------- Solver Config.
class NonlinearPDE_SNESProblem:
def __init__(self, F, u, bc, comm):
V = u.function_space
du = ufl.TrialFunction(V)
self.L = fem.form(F)
self.a = fem.form(ufl.derivative(F, u, du))
self.bc = bc
self._F, self._J = None, None
self.u = u
self.comm = comm
def F(self, snes, x, F):
"""Assemble residual vector."""
x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
x.copy(self.u.x.petsc_vec)
self.u.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
with F.localForm() as f_local:
f_local.set(0.0)
fem.petsc.assemble_vector(F, self.L)
F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
F.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
fem.petsc.apply_lifting(F, [self.a], bcs=[self.bc], x0=[x], scale=-1.0)
F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
fem.petsc.set_bc(F, self.bc, x, -1.0)
F.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
def J(self, snes, x, J, P):
"""Assemble Jacobian matrix."""
x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
x.copy(self.u.x.petsc_vec)
self.u.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
J.zeroEntries()
fem.petsc.assemble_matrix(J, self.a, bcs=self.bc)
J.assemble()
# Regularize the Jacobian by adding εI
epsilon = 1e-10 # Small regularization parameter
J.shift(epsilon) # Add εI to the Jacobian
J.assemble()
def fnc_PrepSolver(Res, u, bcs, a, comm):
Vu = u.function_space
# Set up nonlinear problem
problem = NonlinearPDE_SNESProblem(Res, u, bcs, comm)
b = fem.petsc.create_petsc_vector(Vu.dofmap.index_map, Vu.dofmap.index_map_bs)
J = fem.petsc.create_matrix(problem.a)
# Create Newton solver and solve
snes = PETSc.SNES().create()
snes.setFunction(problem.F, b)
snes.setJacobian(problem.J, J)
snes.setTolerances(stol=1.0e-8, rtol=1e-7, atol=1e-7, max_it=200)
snes.getKSP().setType("preonly")
snes.getKSP().setTolerances(rtol=1e-8, atol=1e-10, max_it=1000)
snes.getKSP().getPC().setType("lu")
class _solver:
def __init__(self, snes_solver):
self.snes_solver = snes_solver
def solve(self, u_now):
u_temp = u_now.x.petsc_vec.copy()
u_temp.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
self.snes_solver.solve(None, u_temp)
iter = self.snes_solver.getIterationNumber()
converged = True if self.snes_solver.getConvergedReason() > 0 else False
residual_norm = self.snes_solver.getFunctionNorm()
return (iter, converged, residual_norm)
solver = _solver(snes)
return solver
#...................................................................................................
# # ----------------------------------------------- FEA setup
Lx, Ly, Lz = 3.0,1.0,1.0
nx, ny, nz = 6, 2, 1
domain = dfx.mesh.create_rectangle(MPI.COMM_WORLD, [[0.0,0.0], [Lx,Ly]],[nx,ny], dfx.mesh.CellType.quadrilateral)
metadata_={'quadrature_degree':4}
ds = ufl.Measure('ds', domain=domain, metadata=metadata_)
dx = ufl.Measure('dx', domain=domain, metadata={'quadrature_degree': metadata_['quadrature_degree'], 'quadrature_rule':'default'})
Vu_el = basix.ufl.element("CG", domain.basix_cell(), 1, shape=(domain.topology.dim,))
Vu = fem.functionspace(domain, Vu_el)
u = fem.Function(Vu)
u_test = ufl.TestFunction(Vu)
du = ufl.TrialFunction(Vu)
Vrho = fem.functionspace(domain, ("DG", 0))
Rho = fem.Function(Vrho)
Rho.vector.array[:] = 0.5
E = PETSc.ScalarType(7.1e10)
nu = PETSc.ScalarType(0.35)
mu = fem.Constant(domain, E / (2.0*(1.0 + nu)))
lmbda = fem.Constant(domain, E*nu / ((1.0 + nu)*(1.0 - 2.0*nu)))
gamma = PETSc.ScalarType(2700.0)
I = ufl.variable(ufl.Identity(len(u)))
F = ufl.variable(I + ufl.grad(u))
J = ufl.variable(ufl.det(F))
epsilon_u = ufl.variable(ufl.sym(F - I))
Psi = Rho * mu * ufl.inner(epsilon_u, epsilon_u) + lmbda/2 * ufl.tr(epsilon_u) * ufl.tr(epsilon_u)
PK1 = ufl.diff(Psi, F)
Res = ufl.inner(PK1, ufl.grad(u_test))*dx
Jac = ufl.derivative(Res, u, du)
#...................................................................................................
# ------------- Set BC
dimFE = domain.topology.dim
domain.topology.create_connectivity(dimFE, dimFE - 1)
domain.topology.create_connectivity(dimFE, 0)
domain.topology.create_connectivity(0, dimFE)
fdim = dimFE - 1
domain.topology.create_connectivity(fdim, fdim + 1)
locate_Left = lambda x: np.isclose(x[0], 0)
facets_left = dfx.mesh.locate_entities(domain, fdim, locate_Left)
Left_u1_dofs = fem.locate_dofs_topological(Vu.sub(0), fdim, facets_left)
Left_u2_dofs = fem.locate_dofs_topological(Vu.sub(1), fdim, facets_left)
bcs_1 = fem.dirichletbc(PETSc.ScalarType(0.0), Left_u1_dofs, Vu.sub(0)) # fix - xLeft
bcs_2 = fem.dirichletbc(PETSc.ScalarType(0.0), Left_u2_dofs, Vu.sub(1)) # fix - yLeft
locate_Right = lambda x: np.isclose(x[0], Lx)
facets_right = dfx.mesh.locate_entities(domain, fdim, locate_Right)
Right_u1_dofs = fem.locate_dofs_topological(Vu.sub(0), fdim, facets_right)
Right_u2_dofs = fem.locate_dofs_topological(Vu.sub(1), fdim, facets_right)
D_appl = fem.Constant(domain,PETSc.ScalarType(0.0))
bcs_3 = fem.dirichletbc(D_appl, Right_u1_dofs, Vu.sub(0)) # disp-loading - xRight
bcs_3adj = fem.dirichletbc(PETSc.ScalarType(0.0), Right_u1_dofs, Vu.sub(0)) # disp-loading - xRight
bcs = [bcs_1, bcs_2, bcs_3]
adj_bcs = [bcs_1, bcs_2, bcs_3adj]
# ------------- Solve
mySolver = fnc_PrepSolver(Res, u, bcs, Jac, comm)
D_appl.value = 0.02*Lx
(iter, converged, res_norm) = mySolver.solve(u)
u.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
#...................................................................................................
# # ----------------------------------------------- Properties of interest
class cls_F_dF():
''' class to compute function & gradient '''
def __init__(self, comm, adj_bcs, fx_field, Res, u, Rho):
self.comm = comm
self.adj_bcs = adj_bcs
self.lamAdj_vec = u.vector.duplicate()
self.fx_form = fem.form(fx_field)
self.fx_ini = np.abs(self.comm.allreduce(fem.assemble_scalar(self.fx_form), op=MPI.SUM))
self.dfdrho_form = fem.form(ufl.derivative(fx_field, Rho))
self.dfdrho_vec = fem.petsc.create_vector(self.dfdrho_form)
self.dfdrho_vec2 = self.dfdrho_vec.duplicate()
self.dfdu_form = fem.form(ufl.derivative(fx_field, u))
self.dfdu_vec = fem.petsc.create_vector(self.dfdu_form)
self.dRes_du_T = ufl.adjoint(ufl.derivative(Res, u))
self.dR_du_form = fem.form(self.dRes_du_T)
self.dR_du_mat = fem.petsc.create_matrix(self.dR_du_form)
self.dRes_drho_T = ufl.adjoint(ufl.derivative(Res, Rho))
self.dR_drho_form = fem.form(self.dRes_drho_T)
self.dR_drho_mat = fem.petsc.create_matrix(self.dR_drho_form)
self.solver = PETSc.KSP().create(self.comm)
self.solver.setOperators(self.dR_du_mat)
prefix = f"Adj_solver_{id(self)}"
self.solver.setOptionsPrefix(prefix)
opts = PETSc.Options()
opts.prefixPush(prefix)
opts.prefixPop()
self.solver.setFromOptions()
def fnc_Evaluate(self):
self.fx = self.comm.allreduce(fem.assemble_scalar(self.fx_form), op=MPI.SUM)
# ... dF/dRho + Lam*dR/du
with self.lamAdj_vec.localForm() as loc_vec:
loc_vec.set(0.0)
self.lamAdj_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
with self.dfdu_vec.localForm() as loc_vec:
loc_vec.set(0.0)
self.dfdu_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
with self.dfdrho_vec.localForm() as loc_vec:
loc_vec.set(0.0)
self.dfdrho_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
with self.dfdrho_vec2.localForm() as loc_vec:
loc_vec.set(0.0)
self.dfdrho_vec2.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
self.dR_du_mat.zeroEntries()
self.dR_drho_mat.zeroEntries()
fem.petsc.assemble_vector(self.dfdu_vec, self.dfdu_form)
self.dfdu_vec.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
self.dfdu_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
self.dfdu_vec.scale(-1)
fem.petsc.apply_lifting(self.dfdu_vec, [self.dR_du_form], [self.adj_bcs], x0=[self.lamAdj_vec], scale=1)
self.dfdu_vec.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
fem.petsc.set_bc(self.dfdu_vec, self.adj_bcs, self.lamAdj_vec, 1.0)
self.dfdu_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
fem.petsc.assemble_matrix(self.dR_du_mat, self.dR_du_form, bcs=self.adj_bcs)
self.dR_du_mat.assemble()
self.solver.setOperators(self.dR_du_mat)
self.solver.solve(self.dfdu_vec, self.lamAdj_vec)
self.lamAdj_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
fem.petsc.assemble_matrix(self.dR_drho_mat, self.dR_drho_form)
self.dR_drho_mat.assemble()
self.dR_drho_mat.mult(self.lamAdj_vec, self.dfdrho_vec2)
fem.petsc.assemble_vector(self.dfdrho_vec, self.dfdrho_form)
self.dfdrho_vec.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
self.dfdrho_vec.axpy(1.0, self.dfdrho_vec2)
def __del__(self):
self.dfdu_vec.destroy()
self.dfdrho_vec.destroy()
self.dfdrho_vec2.destroy()
self.lamAdj_vec.destroy()
self.dR_du_mat.destroy()
self.dR_drho_mat.destroy()
# .... number elements from origin --> for tracking as serial/parallel element orderings are different
ElemId = fem.Function(Vrho)
ElemId.interpolate(lambda x: (x[0]**2 + x[1]**2))
ElemId.x.scatter_forward()
ElemId_values = comm.gather(ElemId.vector.array.tolist(), root=0)
# ..... stress
Cauchy = PK1*F.T/J
dev_Cauchy = Cauchy - (1/3)*ufl.tr(Cauchy)*ufl.Identity(Cauchy.ufl_shape[0])
mises_sum = ufl.sqrt((3/2)*ufl.inner(dev_Cauchy, dev_Cauchy))*dx
fx_S = cls_F_dF(comm, adj_bcs, mises_sum, Res, u, Rho)
fx_S.fnc_Evaluate()
S, dS_dRho_vec = fx_S.fx, fx_S.dfdrho_vec
dS_dRho_values = comm.gather(dS_dRho_vec.array.tolist(), root=0)
if comm.rank == 0:
print(f"--------- stress_sum: parallel (4 process) -----------")
print(f"f: {S}")
global_values = [item for sublist in dS_dRho_values for item in sublist]
global_ElemId_values = [item for sublist in ElemId_values for item in sublist]
data_table = np.vstack([global_ElemId_values, global_values]).T
data_table = data_table[np.lexsort(data_table.T[:-1])]
print(f"[dist, df_dRho]:")
print(f"{data_table}")
print(f"--------------------------------")
# ..... compliance
compliance = (-ufl.inner(PK1, ufl.sym(ufl.grad(u))))*dx
fx_C = cls_F_dF(comm, adj_bcs, compliance, Res, u, Rho)
fx_C.fnc_Evaluate()
C, dC_dRho_vec = fx_C.fx, fx_C.dfdrho_vec
dC_dRho_values = comm.gather(dC_dRho_vec.array.tolist(), root=0)
if comm.rank == 0:
print(f"--------- compliance: parallel (4 process) -----------")
print(f"f: {C}")
global_values = [item for sublist in dC_dRho_values for item in sublist]
global_ElemId_values = [item for sublist in ElemId_values for item in sublist]
data_table = np.vstack([global_ElemId_values, global_values]).T
data_table = data_table[np.lexsort(data_table.T[:-1])]
print(f"[dist, df_dRho]:")
print(f"{data_table}")
print(f"--------------------------------")
