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])