def __init__(self, comm: MPI.Intracomm, problem: NonlinearProblem):
"""A Newton solver for non-linear problems."""
super().__init__(comm)
# Create matrix and vector to be used for assembly
# of the non-linear problem
self._A = fem.create_matrix(problem.a)
self.setJ(problem.J, self._A)
self._b = fem.create_vector(problem.L)
self.setF(problem.F, self._b)
self.set_form(problem.form)
def monolithic_solve():
"""Monolithic version"""
E = P * P
W = FunctionSpace(mesh, E)
U = Function(W)
dU = ufl.TrialFunction(W)
u0, u1 = ufl.split(U)
v0, v1 = ufl.TestFunctions(W)
F = inner((u0**2 + 1) * ufl.grad(u0), ufl.grad(v0)) * dx \
+ inner((u1**2 + 1) * ufl.grad(u1), ufl.grad(v1)) * dx \
- inner(f, v0) * ufl.dx - inner(g, v1) * dx
J = derivative(F, U, dU)
F, J = form(F), form(J)
u0_bc = Function(V0)
u0_bc.interpolate(bc_val_0)
u1_bc = Function(V1)
u1_bc.interpolate(bc_val_1)
bdofsW0_V0 = locate_dofs_topological((W.sub(0), V0), facetdim,
bndry_facets)
bdofsW1_V1 = locate_dofs_topological((W.sub(1), V1), facetdim,
bndry_facets)
bcs = [
dirichletbc(u0_bc, bdofsW0_V0, W.sub(0)),
dirichletbc(u1_bc, bdofsW1_V1, W.sub(1))
]
Jmat = create_matrix(J)
Fvec = create_vector(F)
snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.setTolerances(rtol=1.0e-15, max_it=10)
snes.getKSP().setType("preonly")
snes.getKSP().getPC().setType("lu")
problem = NonlinearPDE_SNESProblem(F, J, U, bcs)
snes.setFunction(problem.F_mono, Fvec)
snes.setJacobian(problem.J_mono, J=Jmat, P=None)
U.sub(0).interpolate(initial_guess_u)
U.sub(1).interpolate(initial_guess_p)
x = create_vector(F)
x.array = U.vector.array_r
snes.solve(None, x)
assert snes.getKSP().getConvergedReason() > 0
assert snes.getConvergedReason() > 0
return x.norm()
def __init__(
self,
F_form: ufl.Form,
u: dolfinx.fem.Function,
bcs=[],
J_form: ufl.Form = None,
bounds=None,
petsc_options={},
form_compiler_parameters={},
jit_parameters={},
monitor=None,
prefix=None,
):
self.u = u
self.bcs = bcs
self.bounds = bounds
# Give PETSc solver options a unique prefix
if prefix is None:
prefix = "snes_{}".format(str(id(self))[0:4])
self.prefix = prefix
if self.bounds is not None:
self.lb = bounds[0]
self.ub = bounds[1]
V = self.u.function_space
self.comm = V.mesh.comm
self.F_form = dolfinx.fem.form(F_form)
if J_form is None:
J_form = ufl.derivative(F_form, self.u, ufl.TrialFunction(V))
self.J_form = dolfinx.fem.form(J_form)
self.petsc_options = petsc_options
self.b = create_vector(self.F_form)
self.a = create_matrix(self.J_form)
self.monitor = monitor
self.solver = self.solver_setup()
def test_overlapping_bcs():
"""Test that, when boundaries condition overlap, the last provided
boundary condition is applied"""
n = 23
mesh = create_unit_square(MPI.COMM_WORLD, n, n)
V = FunctionSpace(mesh, ("Lagrange", 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)
a = form(inner(u, v) * dx)
L = form(inner(1, v) * dx)
dofs_left = locate_dofs_geometrical(V, lambda x: x[0] < 1.0 / (2.0 * n))
dofs_top = locate_dofs_geometrical(V, lambda x: x[1] > 1.0 - 1.0 /
(2.0 * n))
dof_corner = np.array(list(set(dofs_left).intersection(set(dofs_top))),
dtype=np.int64)
# Check only one dof pair is found globally
assert len(set(np.concatenate(MPI.COMM_WORLD.allgather(dof_corner)))) == 1
bcs = [
dirichletbc(PETSc.ScalarType(0), dofs_left, V),
dirichletbc(PETSc.ScalarType(123.456), dofs_top, V)
]
A, b = create_matrix(a), create_vector(L)
assemble_matrix(A, a, bcs=bcs)
A.assemble()
# Check the diagonal (only on the rank that owns the row)
d = A.getDiagonal()
if len(dof_corner) > 0 and dof_corner[0] < V.dofmap.index_map.size_local:
assert np.isclose(d.array_r[dof_corner[0]], 1.0)
with b.localForm() as b_loc:
b_loc.set(0)
assemble_vector(b, L)
apply_lifting(b, [a], [bcs])
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
set_bc(b, bcs)
b.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
if len(dof_corner) > 0:
with b.localForm() as b_loc:
assert b_loc[dof_corner[0]] == 123.456
def test_nonlinear_pde_snes():
"""Test Newton solver for a simple nonlinear PDE"""
# Create mesh and function space
mesh = create_unit_square(MPI.COMM_WORLD, 12, 15)
V = FunctionSpace(mesh, ("Lagrange", 1))
u = Function(V)
v = TestFunction(V)
F = inner(5.0, v) * dx - ufl.sqrt(u * u) * inner(
grad(u), grad(v)) * dx - inner(u, v) * dx
u_bc = Function(V)
u_bc.x.array[:] = 1.0
bc = dirichletbc(u_bc, locate_dofs_geometrical(V, lambda x: np.logical_or(np.isclose(x[0], 0.0),
np.isclose(x[0], 1.0))))
# Create nonlinear problem
problem = NonlinearPDE_SNESProblem(F, u, bc)
u.x.array[:] = 0.9
b = la.create_petsc_vector(V.dofmap.index_map, V.dofmap.index_map_bs)
J = fem.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(rtol=1.0e-9, max_it=10)
snes.getKSP().setType("preonly")
snes.getKSP().setTolerances(rtol=1.0e-9)
snes.getKSP().getPC().setType("lu")
snes.solve(None, u.vector)
assert snes.getConvergedReason() > 0
assert snes.getIterationNumber() < 6
# Modify boundary condition and solve again
u_bc.x.array[:] = 0.6
snes.solve(None, u.vector)
assert snes.getConvergedReason() > 0
assert snes.getIterationNumber() < 6
def matrix(self):
return fem.create_matrix(self.a)
def matrix(self):
return create_matrix(self.a)
def solve_displ_system(J,
F,
intern_var0,
intern_var1,
expr_compiled,
w0,
w1,
bcs,
df,
dt,
t,
k,
io_callback=None,
refinement_callback=None,
rtol=1.e-6,
atol=1.e-16,
max_its=20):
"""Solve system for displacement"""
rank = MPI.COMM_WORLD.rank
t0 = time()
A = create_matrix(J)
if rank == 0:
logger.info(f"[Timer] Preallocation matrix time {time() - t0:.3f}")
t0 = time()
t0 = time()
b = create_vector(F)
if rank == 0:
logger.info(f"[Timer] Preallocation vector time {time() - t0:.3f}")
with b.localForm() as local:
local.set(0.0)
assemble_vector(b, F)
apply_lifting(b, [J], [bcs])
b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
set_bc(b, bcs)
resid_norm0 = b.norm()
if rank == 0:
print(f"Initial DISPL residual: {resid_norm0:.2e}")
iter = 0
# Total number of NR iterations including refinement attempts
iter_total = 0
converged = False
refine = 0
scale = 1.0
rel_resid_norm = 1.0
bcs_init = []
for bc in bcs:
bcs_init += [bc.value.vector.duplicate()]
with bcs_init[-1].localForm() as loc0, bc.value.vector.localForm(
) as loc1:
loc1.copy(loc0)
df0 = np.copy(df.value)
while converged is False:
if iter > max_its or rel_resid_norm > 1.e+2:
refine += 1
iter = 0
if rank == 0:
logger.info(79 * "!")
logger.info(
f"Restarting NR with rel. stepsize: {1.0 / (2 ** refine)}")
with w0["displ"].vector.localForm(
) as w_local, w1["displ"].vector.localForm() as w1_local:
w_local.copy(w1_local)
df.value = df0
scale = 1.0 / 2**refine
if refine > 10:
raise ConvergenceError(
"Inner adaptivity reqiures > 10 refinements.")
# Reset and scale boundary condition
for i, bc in enumerate(bcs_init):
with bcs[i].value.vector.localForm(
) as bcsi_local, bc.localForm() as bc_local:
bc_local.copy(bcsi_local)
bcsi_local.scale(scale)
df.value *= scale
dt.value *= scale
if refinement_callback is not None:
refinement_callback(scale)
if rank == 0:
logger.info("Newton iteration for displ {}".format(iter))
if iter > 0:
for bc in bcs:
with bc.value.vector.localForm() as locvec:
locvec.set(0.0)
A.zeroEntries()
with b.localForm() as local:
local.set(0.0)
size = A.getSize()[0]
local_size = A.getLocalSize()[0]
t0 = time()
assemble_matrix(A, J, bcs)
A.assemble()
_time = time() - t0
Anorm = A.norm()
if rank == 0:
logger.info(f"[Timer] A0 size: {size}, local size: {local_size}")
logger.info(f"[Timer] A0 assembly: {_time:.4f}")
logger.info(f"[Timer] A0 assembly dofs/s: {size / _time:.1f}")
logger.info(f"A0 norm: {Anorm:.4f}")
t0 = time()
assemble_vector(b, F)
apply_lifting(b, [J], [bcs])
b.ghostUpdate(addv=PETSc.InsertMode.ADD,
mode=PETSc.ScatterMode.REVERSE)
set_bc(b, bcs)
bnorm = b.norm()
if rank == 0:
logger.info(f"[Timer] b0 assembly {time() - t0:.4f}")
logger.info(f"b norm: {bnorm:.4f}")
nsp = build_nullspace(w1["displ"].function_space)
ksp = PETSc.KSP()
ksp.create(MPI.COMM_WORLD)
ksp.setOptionsPrefix("disp")
opts = PETSc.Options()
A.setNearNullSpace(nsp)
A.setBlockSize(3)
ksp.setOperators(A)
x = A.createVecRight()
ksp.setFromOptions()
t0 = time()
ksp.solve(b, x)
t1 = time() - t0
opts.view()
xnorm = x.norm()
if rank == 0:
its = ksp.its
t1 = time() - t0
dofsps = int(size / t1)
logger.info(f"[Timer] A0 converged in: {its}")
logger.info(f"[Timer] A0 solve: {t1:.4f}")
logger.info(f"[Timer] A0 solve dofs/s: {dofsps:.1f}")
logger.info(f"Increment norm: {xnorm}")
# TODO: Local axpy segfaults, could ghostUpdate be avoided?
w1["displ"].vector.axpy(1.0, x)
w1["displ"].vector.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
#
# Evaluate true residual and check
#
with b.localForm() as local:
local.set(0.0)
assemble_vector(b, F)
apply_lifting(b, [J], [bcs])
b.ghostUpdate(addv=PETSc.InsertMode.ADD,
mode=PETSc.ScatterMode.REVERSE)
set_bc(b, bcs)
norm = b.norm()
rel_resid_norm = norm / resid_norm0
rel_dx_norm = x.norm() / w1["displ"].vector.norm()
if rank == 0:
logger.info("---")
logger.info(f"Abs. resid norm: {norm:.2e}")
logger.info(f"Rel. dx norm: {rel_dx_norm:.2e}")
logger.info(f"Rel. resid norm: {rel_resid_norm:.2e}")
logger.info("---")
iter += 1
iter_total += 1
if rel_resid_norm < rtol or norm < atol:
if rank == 0:
logger.info(
f"Newton converged in: {iter}, total: {iter_total}")
if io_callback is not None:
io_callback(intern_var1, w1, t.value + dt.value)
return scale
def __init__(self,
a: ufl.Form,
L: ufl.Form,
bcs: typing.List[fem.DirichletBC] = [],
u: fem.Function = None,
petsc_options={},
form_compiler_parameters={},
jit_parameters={}):
"""Initialize solver for a linear variational problem.
Parameters
----------
a
A bilinear UFL form, the left hand side of the variational problem.
L
A linear UFL form, the right hand side of the variational problem.
bcs
A list of Dirichlet boundary conditions.
u
The solution function. It will be created if not provided.
petsc_options
Parameters that is passed to the linear algebra backend PETSc.
For available choices for the 'petsc_options' kwarg, see the
`PETSc-documentation <https://www.mcs.anl.gov/petsc/documentation/index.html>`.
form_compiler_parameters
Parameters used in FFCx compilation of this form. Run `ffcx --help` at
the commandline to see all available options. Takes priority over all
other parameter values, except for `scalar_type` which is determined by
DOLFINx.
jit_parameters
Parameters used in CFFI JIT compilation of C code generated by FFCx.
See `python/dolfinx/jit.py` for all available parameters.
Takes priority over all other parameter values.
.. code-block:: python
problem = LinearProblem(a, L, [bc0, bc1], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
"""
self._a = fem.Form(a,
form_compiler_parameters=form_compiler_parameters,
jit_parameters=jit_parameters)
self._A = fem.create_matrix(self._a)
self._L = fem.Form(L,
form_compiler_parameters=form_compiler_parameters,
jit_parameters=jit_parameters)
self._b = fem.create_vector(self._L)
if u is None:
# Extract function space from TrialFunction (which is at the
# end of the argument list as it is numbered as 1, while the
# Test function is numbered as 0)
self.u = fem.Function(a.arguments()[-1].ufl_function_space())
else:
self.u = u
self.bcs = bcs
self._solver = PETSc.KSP().create(
self.u.function_space.mesh.mpi_comm())
self._solver.setOperators(self._A)
# Give PETSc solver options a unique prefix
solver_prefix = "dolfinx_solve_{}".format(id(self))
self._solver.setOptionsPrefix(solver_prefix)
# Set PETSc options
opts = PETSc.Options()
opts.prefixPush(solver_prefix)
for k, v in petsc_options.items():
opts[k] = v
opts.prefixPop()
self._solver.setFromOptions()
def test_assembly_solve_taylor_hood_nl(mesh):
"""Assemble Stokes problem with Taylor-Hood elements and solve."""
gdim = mesh.geometry.dim
P2 = VectorFunctionSpace(mesh, ("Lagrange", 2))
P1 = FunctionSpace(mesh, ("Lagrange", 1))
def boundary0(x):
"""Define boundary x = 0"""
return np.isclose(x[0], 0.0)
def boundary1(x):
"""Define boundary x = 1"""
return np.isclose(x[0], 1.0)
def initial_guess_u(x):
u_init = np.row_stack(
(np.sin(x[0]) * np.sin(x[1]), np.cos(x[0]) * np.cos(x[1])))
if gdim == 3:
u_init = np.row_stack((u_init, np.cos(x[2])))
return u_init
def initial_guess_p(x):
return -x[0]**2 - x[1]**3
u_bc_0 = Function(P2)
u_bc_0.interpolate(
lambda x: np.row_stack(tuple(x[j] + float(j) for j in range(gdim))))
u_bc_1 = Function(P2)
u_bc_1.interpolate(
lambda x: np.row_stack(tuple(np.sin(x[j]) for j in range(gdim))))
facetdim = mesh.topology.dim - 1
bndry_facets0 = locate_entities_boundary(mesh, facetdim, boundary0)
bndry_facets1 = locate_entities_boundary(mesh, facetdim, boundary1)
bdofs0 = locate_dofs_topological(P2, facetdim, bndry_facets0)
bdofs1 = locate_dofs_topological(P2, facetdim, bndry_facets1)
bcs = [dirichletbc(u_bc_0, bdofs0), dirichletbc(u_bc_1, bdofs1)]
u, p = Function(P2), Function(P1)
du, dp = ufl.TrialFunction(P2), ufl.TrialFunction(P1)
v, q = ufl.TestFunction(P2), ufl.TestFunction(P1)
F = [
inner(ufl.grad(u), ufl.grad(v)) * dx + inner(p, ufl.div(v)) * dx,
inner(ufl.div(u), q) * dx
]
J = [[derivative(F[0], u, du),
derivative(F[0], p, dp)],
[derivative(F[1], u, du),
derivative(F[1], p, dp)]]
P = [[J[0][0], None], [None, inner(dp, q) * dx]]
F, J, P = form(F), form(J), form(P)
# -- Blocked and monolithic
Jmat0 = create_matrix_block(J)
Pmat0 = create_matrix_block(P)
Fvec0 = create_vector_block(F)
snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.setTolerances(rtol=1.0e-15, max_it=10)
snes.getKSP().setType("minres")
snes.getKSP().getPC().setType("lu")
problem = NonlinearPDE_SNESProblem(F, J, [u, p], bcs, P=P)
snes.setFunction(problem.F_block, Fvec0)
snes.setJacobian(problem.J_block, J=Jmat0, P=Pmat0)
u.interpolate(initial_guess_u)
p.interpolate(initial_guess_p)
x0 = create_vector_block(F)
with u.vector.localForm() as _u, p.vector.localForm() as _p:
scatter_local_vectors(x0, [_u.array_r, _p.array_r],
[(u.function_space.dofmap.index_map,
u.function_space.dofmap.index_map_bs),
(p.function_space.dofmap.index_map,
p.function_space.dofmap.index_map_bs)])
x0.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
snes.solve(None, x0)
assert snes.getConvergedReason() > 0
# -- Blocked and nested
Jmat1 = create_matrix_nest(J)
Pmat1 = create_matrix_nest(P)
Fvec1 = create_vector_nest(F)
snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.setTolerances(rtol=1.0e-15, max_it=10)
nested_IS = Jmat1.getNestISs()
snes.getKSP().setType("minres")
snes.getKSP().setTolerances(rtol=1e-12)
snes.getKSP().getPC().setType("fieldsplit")
snes.getKSP().getPC().setFieldSplitIS(["u", nested_IS[0][0]],
["p", nested_IS[1][1]])
ksp_u, ksp_p = snes.getKSP().getPC().getFieldSplitSubKSP()
ksp_u.setType("preonly")
ksp_u.getPC().setType('lu')
ksp_p.setType("preonly")
ksp_p.getPC().setType('lu')
problem = NonlinearPDE_SNESProblem(F, J, [u, p], bcs, P=P)
snes.setFunction(problem.F_nest, Fvec1)
snes.setJacobian(problem.J_nest, J=Jmat1, P=Pmat1)
u.interpolate(initial_guess_u)
p.interpolate(initial_guess_p)
x1 = create_vector_nest(F)
for x1_soln_pair in zip(x1.getNestSubVecs(), (u, p)):
x1_sub, soln_sub = x1_soln_pair
soln_sub.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
soln_sub.vector.copy(result=x1_sub)
x1_sub.ghostUpdate(addv=PETSc.InsertMode.INSERT,
mode=PETSc.ScatterMode.FORWARD)
x1.set(0.0)
snes.solve(None, x1)
assert snes.getConvergedReason() > 0
assert nest_matrix_norm(Jmat1) == pytest.approx(Jmat0.norm(), 1.0e-12)
assert Fvec1.norm() == pytest.approx(Fvec0.norm(), 1.0e-12)
assert x1.norm() == pytest.approx(x0.norm(), 1.0e-12)
# -- Monolithic
P2_el = ufl.VectorElement("Lagrange", mesh.ufl_cell(), 2)
P1_el = ufl.FiniteElement("Lagrange", mesh.ufl_cell(), 1)
TH = P2_el * P1_el
W = FunctionSpace(mesh, TH)
U = Function(W)
dU = ufl.TrialFunction(W)
u, p = ufl.split(U)
du, dp = ufl.split(dU)
v, q = ufl.TestFunctions(W)
F = inner(ufl.grad(u), ufl.grad(v)) * dx + inner(p, ufl.div(v)) * dx \
+ inner(ufl.div(u), q) * dx
J = derivative(F, U, dU)
P = inner(ufl.grad(du), ufl.grad(v)) * dx + inner(dp, q) * dx
F, J, P = form(F), form(J), form(P)
bdofsW0_P2_0 = locate_dofs_topological((W.sub(0), P2), facetdim,
bndry_facets0)
bdofsW0_P2_1 = locate_dofs_topological((W.sub(0), P2), facetdim,
bndry_facets1)
bcs = [
dirichletbc(u_bc_0, bdofsW0_P2_0, W.sub(0)),
dirichletbc(u_bc_1, bdofsW0_P2_1, W.sub(0))
]
Jmat2 = create_matrix(J)
Pmat2 = create_matrix(P)
Fvec2 = create_vector(F)
snes = PETSc.SNES().create(MPI.COMM_WORLD)
snes.setTolerances(rtol=1.0e-15, max_it=10)
snes.getKSP().setType("minres")
snes.getKSP().getPC().setType("lu")
problem = NonlinearPDE_SNESProblem(F, J, U, bcs, P=P)
snes.setFunction(problem.F_mono, Fvec2)
snes.setJacobian(problem.J_mono, J=Jmat2, P=Pmat2)
U.sub(0).interpolate(initial_guess_u)
U.sub(1).interpolate(initial_guess_p)
x2 = create_vector(F)
x2.array = U.vector.array_r
snes.solve(None, x2)
assert snes.getConvergedReason() > 0
assert Jmat2.norm() == pytest.approx(Jmat0.norm(), 1.0e-12)
assert Fvec2.norm() == pytest.approx(Fvec0.norm(), 1.0e-12)
assert x2.norm() == pytest.approx(x0.norm(), 1.0e-12)