Skip to content

Commit e764036

Browse files
jorgensdjorgensd
committed
UPdate newton solver to be safe in parallel (#305)
1 parent 15c02f8 commit e764036

2 files changed

Lines changed: 65 additions & 27 deletions

File tree

‎chapter4/newton-solver.ipynb‎

Lines changed: 37 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -184,11 +184,24 @@
184184
{
185185
"cell_type": "code",
186186
"execution_count": null,
187-
"id": "14",
187+
"id": "0d5ce361",
188188
"metadata": {},
189189
"outputs": [],
190190
"source": [
191191
"solver = PETSc.KSP().create(mesh.comm)\n",
192+
"solver.setType(\"preonly\")\n",
193+
"solver.getPC().setType(\"lu\")\n",
194+
"solver.getPC().setFactorSolverType(\"mumps\")\n",
195+
"solver.setErrorIfNotConverged(True)"
196+
]
197+
},
198+
{
199+
"cell_type": "code",
200+
"execution_count": null,
201+
"id": "14",
202+
"metadata": {},
203+
"outputs": [],
204+
"source": [
192205
"solver.setOperators(A)\n",
193206
"du = dolfinx.fem.Function(V)"
194207
]
@@ -257,7 +270,7 @@
257270
"\n",
258271
" # Compute norm of update\n",
259272
" correction_norm = du.x.petsc_vec.norm(0)\n",
260-
" print(f\"Iteration {i}: Correction norm {correction_norm}\")\n",
273+
" PETSc.Sys.Print(f\"Iteration {i}: Correction norm {correction_norm}\")\n",
261274
" if correction_norm < 1e-10:\n",
262275
" break\n",
263276
" solutions[i, :] = uh.x.array[sort_order]"
@@ -279,7 +292,7 @@
279292
"outputs": [],
280293
"source": [
281294
"dolfinx.fem.petsc.assemble_vector(L, residual)\n",
282-
"print(f\"Final residual {L.norm(0)}\")\n",
295+
"PETSc.Sys.Print(f\"Final residual {L.norm(0)}\")\n",
283296
"A.destroy()\n",
284297
"L.destroy()\n",
285298
"solver.destroy()"
@@ -314,7 +327,7 @@
314327
" u_ex = root(x)\n",
315328
" L2_error = dolfinx.fem.form(ufl.inner(uh - u_ex, uh - u_ex) * ufl.dx)\n",
316329
" global_L2 = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(L2_error), op=MPI.SUM)\n",
317-
" print(f\"L2-error (root {j}) {np.sqrt(global_L2)}\")\n",
330+
" PETSc.Sys.Print(f\"L2-error (root {j}) {np.sqrt(global_L2)}\")\n",
318331
"\n",
319332
" kwargs = {} if j == 0 else {\"label\": \"u_exact\"}\n",
320333
" plt.plot(x_spacing, root(x_spacing.reshape(1, -1)), *args, **kwargs)\n",
@@ -410,7 +423,11 @@
410423
"A = dolfinx.fem.petsc.create_matrix(jacobian)\n",
411424
"L = dolfinx.fem.petsc.create_vector(dolfinx.fem.extract_function_spaces(residual))\n",
412425
"solver = PETSc.KSP().create(mesh.comm)\n",
413-
"solver.setOperators(A)"
426+
"solver.setOperators(A)\n",
427+
"solver.setType(\"preonly\")\n",
428+
"solver.getPC().setType(\"lu\")\n",
429+
"solver.getPC().setFactorSolverType(\"mumps\")\n",
430+
"solver.setErrorIfNotConverged(True)"
414431
]
415432
},
416433
{
@@ -461,15 +478,20 @@
461478
" dolfinx.fem.petsc.assemble_matrix(A, jacobian, bcs=[bc])\n",
462479
" A.assemble()\n",
463480
" dolfinx.fem.petsc.assemble_vector(L, residual)\n",
481+
"\n",
482+
" # Compute b - alpha * J(u_D-u_(i-1))\n",
483+
" dolfinx.fem.petsc.apply_lifting(\n",
484+
" L, [jacobian], [[bc]], x0=[uh.x.petsc_vec], alpha=-1.0\n",
485+
" )\n",
464486
" L.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)\n",
465-
" L.scale(-1)\n",
466487
"\n",
467-
" # Compute b - J(u_D-u_(i-1))\n",
468-
" dolfinx.fem.petsc.apply_lifting(L, [jacobian], [[bc]], x0=[uh.x.petsc_vec], alpha=1)\n",
469-
" # Set du|_bc = u_{i-1}-u_D\n",
470-
" dolfinx.fem.petsc.set_bc(L, [bc], uh.x.petsc_vec, 1.0)\n",
488+
" # Set du|_bc = - (u_{i-1}-u_D)\n",
489+
" dolfinx.fem.petsc.set_bc(L, [bc], uh.x.petsc_vec, -1.0)\n",
471490
" L.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)\n",
472491
"\n",
492+
" # Compute negative residual\n",
493+
" L.scale(-1)\n",
494+
"\n",
473495
" # Solve linear problem\n",
474496
" solver.solve(L, du.x.petsc_vec)\n",
475497
" du.x.scatter_forward()\n",
@@ -486,8 +508,10 @@
486508
" np.sqrt(mesh.comm.allreduce(dolfinx.fem.assemble_scalar(error), op=MPI.SUM))\n",
487509
" )\n",
488510
" du_norm.append(correction_norm)\n",
489-
"\n",
490-
" print(f\"Iteration {i}: Correction norm {correction_norm}, L2 error: {L2_error[-1]}\")\n",
511+
" PETSc.Sys.Print(\n",
512+
" f\"Iteration {i}: Correction norm {correction_norm}, L2 error: {L2_error[-1]}\",\n",
513+
" flush=True,\n",
514+
" )\n",
491515
" if correction_norm < 1e-10:\n",
492516
" break"
493517
]
@@ -542,8 +566,7 @@
542566
"outputs": [],
543567
"source": [
544568
"error_max = domain.comm.allreduce(np.max(np.abs(uh.x.array - u_D.x.array)), op=MPI.MAX)\n",
545-
"if domain.comm.rank == 0:\n",
546-
" print(f\"Error_max: {error_max:.2e}\")"
569+
"PETSc.Sys.Print(f\"Error_max: {error_max:.2e}\")"
547570
]
548571
},
549572
{

‎chapter4/newton-solver.py‎

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
# extension: .py
77
# format_name: light
88
# format_version: '1.5'
9-
# jupytext_version: 1.18.1
9+
# jupytext_version: 1.19.1
1010
# kernelspec:
1111
# display_name: Python 3 (ipykernel)
1212
# language: python
@@ -99,6 +99,11 @@ def root_1(x):
9999
# Next, we create the linear solver and the vector to hold `du`.
100100

101101
solver = PETSc.KSP().create(mesh.comm)
102+
solver.setType("preonly")
103+
solver.getPC().setType("lu")
104+
solver.getPC().setFactorSolverType("mumps")
105+
solver.setErrorIfNotConverged(True)
106+
102107
solver.setOperators(A)
103108
du = dolfinx.fem.Function(V)
104109

@@ -139,15 +144,15 @@ def root_1(x):
139144

140145
# Compute norm of update
141146
correction_norm = du.x.petsc_vec.norm(0)
142-
print(f"Iteration {i}: Correction norm {correction_norm}")
147+
PETSc.Sys.Print(f"Iteration {i}: Correction norm {correction_norm}")
143148
if correction_norm < 1e-10:
144149
break
145150
solutions[i, :] = uh.x.array[sort_order]
146151

147152
# We now compute the magnitude of the residual.
148153

149154
dolfinx.fem.petsc.assemble_vector(L, residual)
150-
print(f"Final residual {L.norm(0)}")
155+
PETSc.Sys.Print(f"Final residual {L.norm(0)}")
151156
A.destroy()
152157
L.destroy()
153158
solver.destroy()
@@ -167,7 +172,7 @@ def root_1(x):
167172
u_ex = root(x)
168173
L2_error = dolfinx.fem.form(ufl.inner(uh - u_ex, uh - u_ex) * ufl.dx)
169174
global_L2 = mesh.comm.allreduce(dolfinx.fem.assemble_scalar(L2_error), op=MPI.SUM)
170-
print(f"L2-error (root {j}) {np.sqrt(global_L2)}")
175+
PETSc.Sys.Print(f"L2-error (root {j}) {np.sqrt(global_L2)}")
171176

172177
kwargs = {} if j == 0 else {"label": "u_exact"}
173178
plt.plot(x_spacing, root(x_spacing.reshape(1, -1)), *args, **kwargs)
@@ -227,6 +232,10 @@ def u_exact(x):
227232
L = dolfinx.fem.petsc.create_vector(dolfinx.fem.extract_function_spaces(residual))
228233
solver = PETSc.KSP().create(mesh.comm)
229234
solver.setOperators(A)
235+
solver.setType("preonly")
236+
solver.getPC().setType("lu")
237+
solver.getPC().setFactorSolverType("mumps")
238+
solver.setErrorIfNotConverged(True)
230239

231240
# Since this problem has strong Dirichlet conditions, we need to apply lifting to the right hand side of our Newton problem.
232241
# We previously had that we wanted to solve the system:
@@ -262,15 +271,20 @@ def u_exact(x):
262271
dolfinx.fem.petsc.assemble_matrix(A, jacobian, bcs=[bc])
263272
A.assemble()
264273
dolfinx.fem.petsc.assemble_vector(L, residual)
274+
275+
# Compute b - alpha * J(u_D-u_(i-1))
276+
dolfinx.fem.petsc.apply_lifting(
277+
L, [jacobian], [[bc]], x0=[uh.x.petsc_vec], alpha=-1.0
278+
)
265279
L.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
266-
L.scale(-1)
267280

268-
# Compute b - J(u_D-u_(i-1))
269-
dolfinx.fem.petsc.apply_lifting(L, [jacobian], [[bc]], x0=[uh.x.petsc_vec], alpha=1)
270-
# Set du|_bc = u_{i-1}-u_D
271-
dolfinx.fem.petsc.set_bc(L, [bc], uh.x.petsc_vec, 1.0)
281+
# Set du|_bc = - (u_{i-1}-u_D)
282+
dolfinx.fem.petsc.set_bc(L, [bc], uh.x.petsc_vec, -1.0)
272283
L.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)
273284

285+
# Compute negative residual
286+
L.scale(-1)
287+
274288
# Solve linear problem
275289
solver.solve(L, du.x.petsc_vec)
276290
du.x.scatter_forward()
@@ -287,8 +301,10 @@ def u_exact(x):
287301
np.sqrt(mesh.comm.allreduce(dolfinx.fem.assemble_scalar(error), op=MPI.SUM))
288302
)
289303
du_norm.append(correction_norm)
290-
291-
print(f"Iteration {i}: Correction norm {correction_norm}, L2 error: {L2_error[-1]}")
304+
PETSc.Sys.Print(
305+
f"Iteration {i}: Correction norm {correction_norm}, L2 error: {L2_error[-1]}",
306+
flush=True,
307+
)
292308
if correction_norm < 1e-10:
293309
break
294310

@@ -315,8 +331,7 @@ def u_exact(x):
315331
# We compute the max error and plot the solution
316332

317333
error_max = domain.comm.allreduce(np.max(np.abs(uh.x.array - u_D.x.array)), op=MPI.MAX)
318-
if domain.comm.rank == 0:
319-
print(f"Error_max: {error_max:.2e}")
334+
PETSc.Sys.Print(f"Error_max: {error_max:.2e}")
320335

321336
u_topology, u_cell_types, u_geometry = dolfinx.plot.vtk_mesh(V)
322337
u_grid = pyvista.UnstructuredGrid(u_topology, u_cell_types, u_geometry)

0 commit comments

Comments
 (0)