Non-linear blocked newton solver with real spaces

[jorgensd]

Suggestion by @CastillonMiguel (and worked out with him)

# # Real function spaces
#
# Author: Jørgen S. Dokken
#
# License: MIT
#
# In this example we will show how to use the "real" function space to solve
# a singular Poisson problem.
# The problem at hand is:
# Find $u \in H^1(\Omega)$ such that
# \begin{align}
# -\Delta u &= f \quad \text{in } \Omega, \\
# \frac{\partial u}{\partial n} &= g \quad \text{on } \partial \Omega, \\
# \int_\Omega u &= h.
# \end{align}
#
# ## Lagrange multiplier
# We start by considering the equivalent optimization problem:
# Find $u \in H^1(\Omega)$ such that
# \begin{align}
# \min_{u \in H^1(\Omega)} J(u) = \min_{u \in H^1(\Omega)} \frac{1}{2}\int_\Omega \vert \nabla u \cdot \nabla u \vert \mathrm{d}x - \int_\Omega f u \mathrm{d}x - \int_{\partial \Omega} g u \mathrm{d}s,
# \end{align}
# such that
# \begin{align}
# \int_\Omega u = h.
# \end{align}
# We introduce a Lagrange multiplier $\lambda$ to enforce the constraint:
# \begin{align}
# \min_{u \in H^1(\Omega), \lambda\in \mathbb{R}} \mathcal{L}(u, \lambda) = \min_{u \in H^1(\Omega), \lambda\in \mathbb{R}} J(u) + \lambda (\int_\Omega u \mathrm{d}x-h).
# \end{align}
# We then compute the optimality conditions for the problem above
# \begin{align}
# \frac{\partial \mathcal{L}}{\partial u}[\delta u] &= \int_\Omega \nabla u \cdot \nabla \delta u \mathrm{d}x + \lambda\int \delta u \mathrm{d}x - \int_\Omega f \delta u ~\mathrm{d}x - \int_{\partial \Omega} g \delta u~\mathrm{d}s = 0, \\
# \frac{\partial \mathcal{L}}{\partial \lambda}[\delta \lambda] &=\delta \lambda (\int_\Omega u \mathrm{d}x -h)= 0.
# \end{align}
# We write the weak formulation:
#
# $$
# \begin{align}
# \int_\Omega \nabla u \cdot \nabla \delta u~\mathrm{d}x + \int_\Omega \lambda \delta u~\mathrm{d}x = \int_\Omega f \delta u~\mathrm{d}x + \int_{\partial \Omega} g v \mathrm{d}s\\
# \int_\Omega u \delta \lambda  \mathrm{d}x = h \delta \lambda .
# \end{align}
# $$
#
# where we have moved $\delta\lambda$ into the integral as it is a spatial constant.

# ## Implementation
# We start by creating the domain and derive the source terms $f$, $g$ and $h$ from our manufactured solution
# For this example we will use the following exact solution
# \begin{align}
# u_{exact}(x, y) = 0.3y^2 + \sin(2\pi x).
# \end{align}

from mpi4py import MPI
from petsc4py import PETSc
import dolfinx.fem.petsc

import numpy as np
from scifem import create_real_functionspace, assemble_scalar
import ufl

M = 20
mesh = dolfinx.mesh.create_unit_square(
    MPI.COMM_WORLD, M, M, dolfinx.mesh.CellType.triangle, dtype=np.float64
)
V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))


def u_exact(x):
    return 0.3 * x[1] ** 2 + ufl.sin(2 * ufl.pi * x[0])


x = ufl.SpatialCoordinate(mesh)
n = ufl.FacetNormal(mesh)
g = ufl.dot(ufl.grad(u_exact(x)), n)
f = -ufl.div(ufl.grad(u_exact(x)))
h = assemble_scalar(u_exact(x) * ufl.dx)

# We then create the Lagrange multiplier space

R = create_real_functionspace(mesh)

# Next, we can create a mixed-function space for our problem


W = ufl.MixedFunctionSpace(V, R)
u = dolfinx.fem.Function(V)
lmbda = dolfinx.fem.Function(R)
w = [u, lmbda]
du, dl = ufl.TestFunctions(W)
dw = ufl.TrialFunctions(W)
# We can now define the variational problem

zero = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0))

F = ufl.inner(ufl.grad(u), ufl.grad(du)) * ufl.dx
F += ufl.inner(lmbda, du) * ufl.dx
F += ufl.inner(u, dl) * ufl.dx
F -= ufl.inner(f, du) * ufl.dx + ufl.inner(g, du) * ufl.ds
F += ufl.inner(zero, dl) * ufl.dx
F_blocked = ufl.extract_blocks(F)


J = sum(ufl.derivative(F, w[i], dw[i]) for i in range(len(w)))
J_blocked = ufl.extract_blocks(J)

import scifem
class RealSpaceNewtonSolver(scifem.BlockedNewtonSolver):

    def _assemble_residual(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
        """Assemble the residual F into the vector b.
        Args:
            x: The vector containing the latest solution
            b: Vector to assemble the residual into
        """
        # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
        with b.localForm() as b_local:
            b_local.set(0.0)
        dolfinx.fem.petsc.assemble_vector_block(
            b,
            self._F,
            self._a,
            bcs=self._bcs,
            x0=x,
            alpha=-1.0,
        )
        start_pos = 0
        for s in self._u:
            num_sub_dofs = (
                s.function_space.dofmap.index_map.size_local
                * s.function_space.dofmap.index_map_bs
            )            
            # FIXME: Add documentation here!
            if s.function_space.dofmap.index_map.size_global == 1:
                assert s.function_space.dofmap.index_map_bs == 1
                if s.function_space.dofmap.index_map.size_local > 0:
                    b.array_w[start_pos:start_pos+num_sub_dofs] -= h
            start_pos += num_sub_dofs
        b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)

# Note that we have defined the variational form in a block form, and
# that we have not included $h$ in the variational form. We will enforce this
# once we have assembled the right hand side vector.

solver = RealSpaceNewtonSolver(F_blocked, w,J=J_blocked, bcs=[],
                               petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})
solver.solve()

print(f"{dolfinx.fem.assemble_scalar(dolfinx.fem.form(w[0]*ufl.dx))}")
diff = w[0] - u_exact(x)
error = dolfinx.fem.form(ufl.inner(diff, diff) * ufl.dx)

print(f"L2 error: {np.sqrt(assemble_scalar(error)):.2e}")

with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solution.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(w[0])

[xinqi16]

Hi, I'm using this solver to solve a time evolution problem with the mesh topology changed frequently, so I need to create the solver many times. But I found that this would lead to

Other MPI error, error stack:
internal_Dist_graph_create_adjacent(125).: MPI_Dist_graph_create_adjacent(comm=0x84000005, indegree=0, sources=(nil), sourceweights=0x7feebf8c3e14, outdegree=0, destinations=(nil), destweights=0x7feebf8c3e14, MPI_INFO_NULL, reorder=0, comm_dist_graph=0x7ffe2e0070ec) failed
MPIR_Dist_graph_create_adjacent_impl(317): 
MPII_Comm_copy(937)......................: 
MPIR_Get_contextid_sparse_group(573).....: Too many communicators (0/2048 free on this process; ignore_id=0)
Abort(202996751) on node 0 (rank 0 in comm 416): application called MPI_Abort(comm=0x84000005, 202996751) - process 0

A minimum reproducible example can be done by simply adding

for i in range(2**18):
    print(f"Loop {i}")
    solver = RealSpaceNewtonSolver(F_blocked, w,J=J_blocked, bcs=[],
                                petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})

to this code. Could this problem be solved? I think this is possibly because we did not destroy the solver in the Newton solver.

[CastillonMiguel]

First of all, I've added the missing documentation for the FIXME comment. The code section now includes proper documentation explaining what's happening in that block

# # Real function spaces
#
# Author: Jørgen S. Dokken
#
# License: MIT
#
# In this example we will show how to use the "real" function space to solve
# a singular Poisson problem.
# The problem at hand is:
# Find $u \in H^1(\Omega)$ such that
# \begin{align}
# -\Delta u &= f \quad \text{in } \Omega, \\
# \frac{\partial u}{\partial n} &= g \quad \text{on } \partial \Omega, \\
# \int_\Omega u &= h.
# \end{align}
#
# ## Lagrange multiplier
# We start by considering the equivalent optimization problem:
# Find $u \in H^1(\Omega)$ such that
# \begin{align}
# \min_{u \in H^1(\Omega)} J(u) = \min_{u \in H^1(\Omega)} \frac{1}{2}\int_\Omega \vert \nabla u \cdot \nabla u \vert \mathrm{d}x - \int_\Omega f u \mathrm{d}x - \int_{\partial \Omega} g u \mathrm{d}s,
# \end{align}
# such that
# \begin{align}
# \int_\Omega u = h.
# \end{align}
# We introduce a Lagrange multiplier $\lambda$ to enforce the constraint:
# \begin{align}
# \min_{u \in H^1(\Omega), \lambda\in \mathbb{R}} \mathcal{L}(u, \lambda) = \min_{u \in H^1(\Omega), \lambda\in \mathbb{R}} J(u) + \lambda (\int_\Omega u \mathrm{d}x-h).
# \end{align}
# We then compute the optimality conditions for the problem above
# \begin{align}
# \frac{\partial \mathcal{L}}{\partial u}[\delta u] &= \int_\Omega \nabla u \cdot \nabla \delta u \mathrm{d}x + \lambda\int \delta u \mathrm{d}x - \int_\Omega f \delta u ~\mathrm{d}x - \int_{\partial \Omega} g \delta u~\mathrm{d}s = 0, \\
# \frac{\partial \mathcal{L}}{\partial \lambda}[\delta \lambda] &=\delta \lambda (\int_\Omega u \mathrm{d}x -h)= 0.
# \end{align}
# We write the weak formulation:
#
# $$
# \begin{align}
# \int_\Omega \nabla u \cdot \nabla \delta u~\mathrm{d}x + \int_\Omega \lambda \delta u~\mathrm{d}x = \int_\Omega f \delta u~\mathrm{d}x + \int_{\partial \Omega} g v \mathrm{d}s\\
# \int_\Omega u \delta \lambda  \mathrm{d}x = h \delta \lambda .
# \end{align}
# $$
#
# where we have moved $\delta\lambda$ into the integral as it is a spatial constant.

# ## Implementation
# We start by creating the domain and derive the source terms $f$, $g$ and $h$ from our manufactured solution
# For this example we will use the following exact solution
# \begin{align}
# u_{exact}(x, y) = 0.3y^2 + \sin(2\pi x).
# \end{align}

from mpi4py import MPI
from petsc4py import PETSc
import dolfinx.fem.petsc

import numpy as np
from scifem import create_real_functionspace, assemble_scalar
import ufl

M = 20
mesh = dolfinx.mesh.create_unit_square(
    MPI.COMM_WORLD, M, M, dolfinx.mesh.CellType.triangle, dtype=np.float64
)
V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))


def u_exact(x):
    return 0.3 * x[1] ** 2 + ufl.sin(2 * ufl.pi * x[0])


x = ufl.SpatialCoordinate(mesh)
n = ufl.FacetNormal(mesh)
g = ufl.dot(ufl.grad(u_exact(x)), n)
f = -ufl.div(ufl.grad(u_exact(x)))
h = assemble_scalar(u_exact(x) * ufl.dx)

# We then create the Lagrange multiplier space

R = create_real_functionspace(mesh)

# Next, we can create a mixed-function space for our problem


W = ufl.MixedFunctionSpace(V, R)
u = dolfinx.fem.Function(V)
lmbda = dolfinx.fem.Function(R)
w = [u, lmbda]
du, dl = ufl.TestFunctions(W)
dw = ufl.TrialFunctions(W)
# We can now define the variational problem

zero = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(0.0))

F = ufl.inner(ufl.grad(u), ufl.grad(du)) * ufl.dx
F += ufl.inner(lmbda, du) * ufl.dx
F += ufl.inner(u, dl) * ufl.dx
F -= ufl.inner(f, du) * ufl.dx + ufl.inner(g, du) * ufl.ds
F += ufl.inner(zero, dl) * ufl.dx
F_blocked = ufl.extract_blocks(F)


J = sum(ufl.derivative(F, w[i], dw[i]) for i in range(len(w)))
J_blocked = ufl.extract_blocks(J)

import scifem
class RealSpaceNewtonSolver(scifem.BlockedNewtonSolver):

    def _assemble_residual(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
        """Assemble the residual F into the vector b.
        Args:
            x: The vector containing the latest solution
            b: Vector to assemble the residual into
        """
        # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
        with b.localForm() as b_local:
            b_local.set(0.0)
        dolfinx.fem.petsc.assemble_vector_block(
            b,
            self._F,
            self._a,
            bcs=self._bcs,
            x0=x,
            alpha=-1.0,
        )
        start_pos = 0
        for s in self._u:
            num_sub_dofs = (
                s.function_space.dofmap.index_map.size_local
                * s.function_space.dofmap.index_map_bs
            )            
            # Enforce constraint h for the Lagrange multiplier (real function space)
            # The constraint equation is: ∫_Ω u dx = h
            # This modifies the residual for the real space DOF by subtracting h
            if s.function_space.dofmap.index_map.size_global == 1:
                assert s.function_space.dofmap.index_map_bs == 1
                if s.function_space.dofmap.index_map.size_local > 0:
                    b.array_w[start_pos:start_pos+num_sub_dofs] -= h
            start_pos += num_sub_dofs
        b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)

# Note that we have defined the variational form in a block form, and
# that we have not included $h$ in the variational form. We will enforce this
# once we have assembled the right hand side vector.

solver = RealSpaceNewtonSolver(F_blocked, w,J=J_blocked, bcs=[],
                               petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})
solver.solve()

print(f"{dolfinx.fem.assemble_scalar(dolfinx.fem.form(w[0]*ufl.dx))}")
diff = w[0] - u_exact(x)
error = dolfinx.fem.form(ufl.inner(diff, diff) * ufl.dx)

print(f"L2 error: {np.sqrt(assemble_scalar(error)):.2e}")

with dolfinx.io.XDMFFile(MPI.COMM_WORLD, "solution.xdmf", "w") as xdmf:
    xdmf.write_mesh(mesh)
    xdmf.write_function(w[0])

[CastillonMiguel]

Hi, I'm using this solver to solve a time evolution problem with the mesh topology changed frequently, so I need to create the solver many times. But I found that this would lead to

Other MPI error, error stack:
internal_Dist_graph_create_adjacent(125).: MPI_Dist_graph_create_adjacent(comm=0x84000005, indegree=0, sources=(nil), sourceweights=0x7feebf8c3e14, outdegree=0, destinations=(nil), destweights=0x7feebf8c3e14, MPI_INFO_NULL, reorder=0, comm_dist_graph=0x7ffe2e0070ec) failed
MPIR_Dist_graph_create_adjacent_impl(317): 
MPII_Comm_copy(937)......................: 
MPIR_Get_contextid_sparse_group(573).....: Too many communicators (0/2048 free on this process; ignore_id=0)
Abort(202996751) on node 0 (rank 0 in comm 416): application called MPI_Abort(comm=0x84000005, 202996751) - process 0

A minimum reproducible example can be done by simply adding

for i in range(2**18):
print(f"Loop {i}")
solver = RealSpaceNewtonSolver(F_blocked, w,J=J_blocked, bcs=[],
petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})
to this code. Could this problem be solved? I think this is possibly because we did not destroy the solver in the Newton solver.

Hi, I'm not sure about the root cause of this issue, but it might be possible that this approach could help.

Reproduction attempt: I tried to reproduce this issue using your minimal example, but I cannot reproduce the error on my system. This could be due to different MPI implementations or system configurations.

Root cause: Each RealSpaceNewtonSolver instance likely creates PETSc objects that hold MPI communicators, but these may not be properly cleaned up when the solver is destroyed.

Quick workaround to test: Try explicitly deleting the solver:

for i in range(2**18):
    print(f"Loop {i}")
    solver = RealSpaceNewtonSolver(F_blocked, w, J=J_blocked, bcs=[],
                                   petsc_options={"ksp_type": "preonly", "pc_type": "lu", 
                                                  "pc_factor_mat_solver_type": "mumps"})
    solver.solve()
    del solver  # Explicitly delete

[xinqi16]

Hi, I'm using this solver to solve a time evolution problem with the mesh topology changed frequently, so I need to create the solver many times. But I found that this would lead to

Other MPI error, error stack:
internal_Dist_graph_create_adjacent(125).: MPI_Dist_graph_create_adjacent(comm=0x84000005, indegree=0, sources=(nil), sourceweights=0x7feebf8c3e14, outdegree=0, destinations=(nil), destweights=0x7feebf8c3e14, MPI_INFO_NULL, reorder=0, comm_dist_graph=0x7ffe2e0070ec) failed
MPIR_Dist_graph_create_adjacent_impl(317): 
MPII_Comm_copy(937)......................: 
MPIR_Get_contextid_sparse_group(573).....: Too many communicators (0/2048 free on this process; ignore_id=0)
Abort(202996751) on node 0 (rank 0 in comm 416): application called MPI_Abort(comm=0x84000005, 202996751) - process 0

A minimum reproducible example can be done by simply adding
for i in range(2**18):
print(f"Loop {i}")
solver = RealSpaceNewtonSolver(F_blocked, w,J=J_blocked, bcs=[],
petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})
to this code. Could this problem be solved? I think this is possibly because we did not destroy the solver in the Newton solver.

Hi, I'm not sure about the root cause of this issue, but it might be possible that this approach could help.

Reproduction attempt: I tried to reproduce this issue using your minimal example, but I cannot reproduce the error on my system. This could be due to different MPI implementations or system configurations.

Root cause: Each RealSpaceNewtonSolver instance likely creates PETSc objects that hold MPI communicators, but these may not be properly cleaned up when the solver is destroyed.

Quick workaround to test: Try explicitly deleting the solver:

for i in range(2**18):
print(f"Loop {i}")
solver = RealSpaceNewtonSolver(F_blocked, w, J=J_blocked, bcs=[],
petsc_options={"ksp_type": "preonly", "pc_type": "lu",
"pc_factor_mat_solver_type": "mumps"})
solver.solve()
del solver # Explicitly delete

It seems that this doesn't work. It pumps out with

...
Loop 2025
Other MPI error, error stack:
internal_Dist_graph_create_adjacent(125).: MPI_Dist_graph_create_adjacent(comm=0x84000005, indegree=0, sources=(nil), sourceweights=0x7f1b72ec3e14, outdegree=0, destinations=(nil), destweights=0x7f1b72ec3e14, MPI_INFO_NULL, reorder=0, comm_dist_graph=0x7ffc69d85e2c) failed
MPIR_Dist_graph_create_adjacent_impl(317): 
MPII_Comm_copy(937)......................: 
MPIR_Get_contextid_sparse_group(573).....: Too many communicators (0/2048 free on this process; ignore_id=0)
Abort(202996751) on node 0 (rank 0 in comm 416): application called MPI_Abort(comm=0x84000005, 202996751) - process 0

with my code as

for i in range(2**18):
    print(f"Loop {i}")
    solver = RealSpaceNewtonSolver(F_blocked, w,J=J_blocked, bcs=[],
                                petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})
    del solver

Anyway, I decided to use pctools to write my own newton solver, and it works. Thanks a lot for your help!