Parallel custom Newton solver

General
1

Hello, the following custom newton solver works fine in serial but breaks down in parallel and I really cannot tell why… any help would be really appreciated.

import basix
import dolfinx
import numpy as np
import ufl
from dolfinx import fem, io, mesh
from dolfinx.fem.petsc import apply_lifting, assemble_matrix, assemble_vector, set_bc
from mpi4py import MPI
from petsc4py import PETSc


def main():
    comm = MPI.COMM_WORLD
    rank = comm.rank


    domain = mesh.create_rectangle(
        comm,
        points=[np.array([0.0, 0.0]), np.array([1.0, 1.0])],
        n=[64, 64],
        cell_type=mesh.CellType.triangle,
    )
    fdim = domain.topology.dim - 1

    ve = basix.ufl.element("CG", domain.basix_cell(), 2, shape=(domain.geometry.dim,))
    pe = basix.ufl.element("CG", domain.basix_cell(), 1)

    V = fem.functionspace(domain, ve)
    Q = fem.functionspace(domain, pe)
    VP = fem.functionspace(domain, basix.ufl.mixed_element([ve, pe]))

    vp_n = fem.Function(VP)
    u, p = ufl.split(vp_n)
    v, q = ufl.TestFunctions(VP)

    def top_boundary(x):
        return np.isclose(x[1], 1.0)

    def walls_except_top(x):
        on_bottom = np.isclose(x[1], 0.0)
        on_left = np.isclose(x[0], 0.0)
        on_right = np.isclose(x[0], 1.0)
        return np.logical_or(np.logical_or(on_bottom, on_left), on_right)

    top_facets = mesh.locate_entities_boundary(domain, fdim, top_boundary)
    wall_facets = mesh.locate_entities_boundary(domain, fdim, walls_except_top)

    u_top = fem.Function(V)
    u_top.interpolate(
        lambda x: np.vstack(
            (
                np.full(x.shape[1], 1.0, dtype=PETSc.ScalarType),
                np.zeros(x.shape[1], dtype=PETSc.ScalarType),
            )
        )
    )

    u_zero = fem.Function(V)
    p_zero = fem.Function(Q)

    top_dofs_u = fem.locate_dofs_topological((VP.sub(0), V), fdim, top_facets)
    wall_dofs_u = fem.locate_dofs_topological((VP.sub(0), V), fdim, wall_facets)
    p_pin_dofs = fem.locate_dofs_geometrical(
        (VP.sub(1), Q),
        lambda x: np.logical_and(np.isclose(x[0], 0.0), np.isclose(x[1], 0.0)),
    )

    bc_top = fem.dirichletbc(u_top, top_dofs_u, VP.sub(0))
    bc_walls = fem.dirichletbc(u_zero, wall_dofs_u, VP.sub(0))
    bc_p_pin = fem.dirichletbc(p_zero, p_pin_dofs, VP.sub(1))
    bcs = [bc_top, bc_walls, bc_p_pin]

    dx = ufl.Measure("dx", domain=domain)
    rho = fem.Constant(domain, PETSc.ScalarType(1.0))
    mu = fem.Constant(domain, PETSc.ScalarType(0.01))

    F_lin = (
        2.0 * mu * ufl.inner(ufl.sym(ufl.grad(u)), ufl.sym(ufl.grad(v))) * dx
        - p * ufl.div(v) * dx
        + q * ufl.div(u) * dx
    )
    F_nonlin = rho * ufl.inner(ufl.grad(u) * u, v) * dx

    dvp_ = {"n": vp_n}
    lmbda = 1.0
    atol = 1e-8
    max_it = 10

    ksp = PETSc.KSP().create(comm=MPI.COMM_WORLD)
    ksp.setType("preonly")
    ksp.getPC().setType("lu")
    ksp.getPC().setFactorSolverType("mumps")

    F_total = F_lin + F_nonlin
    J = ufl.derivative(F_total, dvp_["n"])
    J_form = fem.form(J)
    F_form = fem.form(F_total)

    dvp_res = fem.Function(dvp_["n"].function_space)

    if rank == 0:
        print("Starting continuation + Newton iterations...", flush=True)

    residual = 1.0e8

    while residual > atol and iter_count < max_it:

        b = assemble_vector(F_form)
        b.scale(-1.0)


        apply_lifting(b, [J_form], [bcs], x0=[dvp_["n"].vector], scale=1.0)
        set_bc(b, bcs, dvp_["n"].vector, 1.0)
        b.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD
        )

        A = assemble_matrix(J_form, bcs=bcs)
        A.assemble()

        ksp.setOperators(A)
        ksp.solve(b, dvp_res.vector)
        dvp_res.x.scatter_forward()
        dvp_["n"].vector.axpy(lmbda, dvp_res.vector)

        residual = b.norm()
        iter_count += 1

        if rank == 0:
            print(
                f"  iter={iter_count:03d} residual={residual:.6e}",
                flush=True,
            )

        A.destroy()
        b.destroy()
        dvp_res.vector.zeroEntries()

if __name__ == "__main__":
    main()
2

Ghost updates is not done correctly, as for instance covered in:

I’ll update the tutorial:

to reflect this.

3

Fix made at:

4

Thank you very much for the prompt reply as usual and for having fixed the tutorials @dokken.

Cheers!