From 063cba9fc390d13d512257db450ca9580dc2ef27 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Mon, 23 Sep 2024 09:07:17 +0000
Subject: [PATCH 1/9] Add blocked newton solver and appropriate test

---
 src/scifem/__init__.py              |   3 +-
 src/scifem/solvers.py               | 220 ++++++++++++++++++++++++++++
 tests/test_blocked_newton_solver.py |  56 +++++++
 3 files changed, 278 insertions(+), 1 deletion(-)
 create mode 100644 src/scifem/solvers.py
 create mode 100644 tests/test_blocked_newton_solver.py

diff --git a/src/scifem/__init__.py b/src/scifem/__init__.py
index 56cdd7f..a86febb 100644
--- a/src/scifem/__init__.py
+++ b/src/scifem/__init__.py
@@ -5,6 +5,7 @@
 from .point_source import PointSource
 from .assembly import assemble_scalar
 from . import xdmf
+from .solvers import NewtonSolver
 
 
 def create_real_functionspace(
@@ -37,4 +38,4 @@ def create_real_functionspace(
     return dolfinx.fem.FunctionSpace(mesh, ufl_e, cppV)
 
 
-__all__ = ["create_real_functionspace", "assemble_scalar", "PointSource", "xdmf"]
+__all__ = ["create_real_functionspace", "assemble_scalar", "PointSource", "xdmf", "NewtonSolver"]
diff --git a/src/scifem/solvers.py b/src/scifem/solvers.py
new file mode 100644
index 0000000..3da893c
--- /dev/null
+++ b/src/scifem/solvers.py
@@ -0,0 +1,220 @@
+import dolfinx
+from typing import Callable
+import ufl
+import numpy as np
+from petsc4py import PETSc
+
+__all__ = ["NewtonSolver"]
+
+
+class NewtonSolver:
+    max_iterations: int
+    _bcs: list[dolfinx.fem.DirichletBC]
+    A: PETSc.Mat
+    b: PETSc.Vec
+    _J: dolfinx.fem.Form | ufl.form.Form
+    _F: dolfinx.fem.Form | ufl.form.Form
+    _dx: PETSc.Vec
+    _error_on_convergence: bool
+    _pre_solve_callback: Callable[["NewtonSolver"], None] | None
+    _post_solve_callback: Callable[["NewtonSolver"], None] | None
+
+    def __init__(
+        self,
+        F: list[dolfinx.fem.form],
+        J: list[list[dolfinx.fem.form]],
+        w: list[dolfinx.fem.Function],
+        bcs: list[dolfinx.fem.DirichletBC] | None = None,
+        max_iterations: int = 5,
+        petsc_options: dict[str, str | float | int | None] | None = None,
+        error_on_nonconvergence: bool = True,
+    ):
+        """
+        Create a Newton solver for a block nonlinear problem ``F(u) = 0``.
+        Solved as ``J(u) du = -F(u)``, where ``J`` is the Jacobian of ``F``.
+
+        Args:
+            F: List of forms defining the residual
+            J: List of lists of forms defining the Jacobian
+            w: List of functions representing the solution
+            bcs: List of Dirichlet boundary conditions
+            max_iterations: Maximum number of iterations in Newton solver
+            petsc_options: PETSc options for Krylov subspace solver.
+            error_on_nonconvergence: Raise an error if the linear solver
+                does not converge or Newton solver doesn't converge
+        """
+        # Compile forms if not compiled. Will throw error it requires entity maps
+        self._F = dolfinx.fem.form(F)
+        self._J = dolfinx.fem.form(J)
+
+        # Store solution and accessible/modifiable properties
+        self._pre_solve_callback = None
+        self._post_solve_callback = None
+        self.max_iterations = max_iterations
+        self._error_on_convergence = error_on_nonconvergence
+        self.bcs = [] if bcs is None else bcs
+        self.w = w
+
+        # Create PETSc objects for block assembly
+        self.b = dolfinx.fem.petsc.create_vector_block(self._F)
+        self.A = dolfinx.fem.petsc.create_matrix_block(self._J)
+        self.dx = dolfinx.fem.petsc.create_vector_block(self._F)
+        self.x = dolfinx.fem.petsc.create_vector_block(self._F)
+
+        # Set PETSc options
+        opts = PETSc.Options()
+        if petsc_options is not None:
+            for k, v in petsc_options.items():
+                opts[k] = v
+
+        # Define KSP solver
+        self._solver = PETSc.KSP().create(self.b.getComm().tompi4py())
+        self._solver.setOperators(self.A)
+        self._solver.setFromOptions()
+
+        # Set matrix and vector PETSc options
+        self.A.setFromOptions()
+        self.b.setFromOptions()
+
+    def set_pre_solve_callback(self, callback: Callable[["NewtonSolver"], None]):
+        """Set a callback function that is called before each Newton iteration."""
+        self._pre_solve_callback = callback
+
+    def set_post_solve_callback(self, callback: Callable[["NewtonSolver"], None]):
+        """Set a callback function that is called after each Newton iteration."""
+        self._post_solve_callback = callback
+
+    def solve(self, tol=1e-6, beta=1.0) -> int:
+        """Solve the nonlinear problem using Newton's method.
+
+        Args:
+            tol: Tolerance for termination of Newton's method.
+            beta: Damping parameter for the update.
+
+        Returns:
+            The number of Newton iterations used to converge.
+
+        Note:
+            The tolerance is on the 0-norm of the update.
+        """
+        i = 1
+
+        while i <= self.max_iterations:
+            if self._pre_solve_callback is not None:
+                self._pre_solve_callback(self)
+
+            # Pack constants and coefficients
+            constants_L = [
+                form and dolfinx.cpp.fem.pack_constants(form._cpp_object) for form in self._F
+            ]
+            coeffs_L = [dolfinx.cpp.fem.pack_coefficients(form._cpp_object) for form in self._F]
+            constants_a = [
+                [
+                    dolfinx.cpp.fem.pack_constants(form._cpp_object)
+                    if form is not None
+                    else np.array([], dtype=PETSc.ScalarType)
+                    for form in forms
+                ]
+                for forms in self._J
+            ]
+            coeffs_a = [
+                [
+                    {} if form is None else dolfinx.cpp.fem.pack_coefficients(form._cpp_object)
+                    for form in forms
+                ]
+                for forms in self._J
+            ]
+
+            # Scatter previous solution `w` to `self.x`, the blocked version used for lifting
+            dolfinx.cpp.la.petsc.scatter_local_vectors(
+                self.x,
+                [si.x.petsc_vec.array_r for si in self.w],
+                [
+                    (
+                        si.function_space.dofmap.index_map,
+                        si.function_space.dofmap.index_map_bs,
+                    )
+                    for si in self.w
+                ],
+            )
+            self.x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+
+            # Assemble F(u_{i-1}) - J(u_D - u_{i-1}) and set du|_bc= u_D - u_{i-1}
+            with self.b.localForm() as b_local:
+                b_local.set(0.0)
+            dolfinx.fem.petsc.assemble_vector_block(
+                self.b,
+                self._F,
+                self._J,
+                bcs=self.bcs,
+                x0=self.x,
+                scale=-1.0,
+                coeffs_a=coeffs_a,
+                constants_a=constants_a,
+                coeffs_L=coeffs_L,
+                constants_L=constants_L,
+            )
+            self.b.ghostUpdate(PETSc.InsertMode.INSERT_VALUES, PETSc.ScatterMode.FORWARD)
+
+            # Assemble Jacobian
+            self.A.zeroEntries()
+            dolfinx.fem.petsc.assemble_matrix_block(
+                self.A, self._J, bcs=self.bcs, constants=constants_a, coeffs=coeffs_a
+            )
+            self.A.assemble()
+
+            self._solver.solve(self.b, self.dx)
+            if self._error_on_convergence:
+                if (status := self._solver.getConvergedReason()) <= 0:
+                    raise RuntimeError(f"Linear solver did not converge, got reason: {status}")
+
+            # Update solution
+            offset_start = 0
+            for s in self.w:
+                num_sub_dofs = (
+                    s.function_space.dofmap.index_map.size_local
+                    * s.function_space.dofmap.index_map_bs
+                )
+                s.x.petsc_vec.array_w[:num_sub_dofs] -= (
+                    beta * self.dx.array_r[offset_start : offset_start + num_sub_dofs]
+                )
+                s.x.petsc_vec.ghostUpdate(
+                    addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD
+                )
+                offset_start += num_sub_dofs
+
+            if self._post_solve_callback is not None:
+                self._post_solve_callback(self)
+
+            # Compute norm of update
+            correction_norm = self.dx.norm(0)
+            print(f"Iteration {i}: Correction norm {correction_norm}")
+            if correction_norm < tol:
+                return i
+            i += 1
+
+        if self._error_on_convergence:
+            raise RuntimeError("Newton solver did not converge")
+        else:
+            return self.max_iterations
+
+    @property
+    def F(self):
+        """The list of residuals where each entry is a ``dolfinx.fem.Form``."""
+        return self._F
+
+    @property
+    def J(self):
+        """
+        The Jacobian blocks represented as lists of lists where each entry
+        is a ``dolfinx.fem.Form``.
+        """
+        return self
+
+    def __del__(self):
+        """Clean up the solver by destroying PETSc objects."""
+        self.A.destroy()
+        self.b.destroy()
+        self.dx.destroy()
+        self._solver.destroy()
+        self.x.destroy()
diff --git a/tests/test_blocked_newton_solver.py b/tests/test_blocked_newton_solver.py
new file mode 100644
index 0000000..653add5
--- /dev/null
+++ b/tests/test_blocked_newton_solver.py
@@ -0,0 +1,56 @@
+from mpi4py import MPI
+
+import numpy as np
+import dolfinx
+from petsc4py import PETSc
+from scifem import NewtonSolver, assemble_scalar
+import ufl
+import pytest
+
+
+@pytest.mark.parametrize("factor", [1, -1])
+def test_NewtonSolver(factor):
+    dtype = PETSc.RealType
+    ftype = PETSc.ScalarType
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 2, 2, dtype=dtype)
+    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
+
+    W = ufl.MixedFunctionSpace(V, Q)
+    v, q = ufl.TestFunctions(W)
+    du, dp = ufl.TrialFunctions(W)
+
+    # Set nonzero initial guess
+    # Choose initial guess acording to the input factor
+    u = dolfinx.fem.Function(V, dtype=ftype)
+    u.x.array[:] = factor * 0.1
+    p = dolfinx.fem.Function(Q, dtype=ftype)
+    p.x.array[:] = factor * 0.02
+    x = ufl.SpatialCoordinate(mesh)
+    u_expr = 3 * x[0] ** 2
+    p_expr = 5 * x[1] ** 4
+    c0 = dolfinx.fem.Constant(mesh, ftype(0.3))
+    c1 = dolfinx.fem.Constant(mesh, ftype(0.82))
+    F0 = ufl.inner(c0 * u**2, v) * ufl.dx - ufl.inner(u_expr, v) * ufl.dx
+    F1 = ufl.inner(c1 * p**2, q) * ufl.dx - ufl.inner(p_expr, q) * ufl.dx
+    if dolfinx.__version__ == "0.8.0":
+        F = [F0, F1]
+        J = [
+            [ufl.derivative(F0, u, du), ufl.derivative(F0, p, du)],
+            [ufl.derivative(F1, u, dp), ufl.derivative(F1, p, dp)],
+        ]
+    elif dolfinx.__version__ == "0.9.0.0":
+        F_ = F0 + F1
+        F = list(ufl.extract_blocks(F_))
+        J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
+    petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
+    solver = NewtonSolver(F, J, [u, p], max_iterations=25, petsc_options=petsc_options)
+    solver.solve()
+
+    u_ex = factor * ufl.sqrt(u_expr / c0)
+    p_ex = factor * ufl.sqrt(p_expr / c1)
+    err_u = ufl.inner(u_ex - u, u_ex - u) * ufl.dx
+    err_p = ufl.inner(p_ex - p, p_ex - p) * ufl.dx
+    tol = np.finfo(dtype).eps * 1.0e3
+    assert assemble_scalar(err_u) < tol
+    assert assemble_scalar(err_p) < tol

From 7f3e30c83e7bb06b952941c52937dae283dd0b47 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Mon, 23 Sep 2024 09:12:24 +0000
Subject: [PATCH 2/9] Change grid size

---
 tests/test_blocked_newton_solver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_blocked_newton_solver.py b/tests/test_blocked_newton_solver.py
index 653add5..defe73f 100644
--- a/tests/test_blocked_newton_solver.py
+++ b/tests/test_blocked_newton_solver.py
@@ -12,7 +12,7 @@
 def test_NewtonSolver(factor):
     dtype = PETSc.RealType
     ftype = PETSc.ScalarType
-    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 2, 2, dtype=dtype)
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 12, dtype=dtype)
     V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
     Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
 

From e870d3719de811ee272661bd81bbd77010299bff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Mon, 23 Sep 2024 10:45:49 +0000
Subject: [PATCH 3/9] Revert fancy fixes from
 https://github.com/FEniCS/ufl/pull/308

---
 tests/test_blocked_newton_solver.py | 36 ++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/tests/test_blocked_newton_solver.py b/tests/test_blocked_newton_solver.py
index defe73f..d64df01 100644
--- a/tests/test_blocked_newton_solver.py
+++ b/tests/test_blocked_newton_solver.py
@@ -16,9 +16,18 @@ def test_NewtonSolver(factor):
     V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
     Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 2))
 
-    W = ufl.MixedFunctionSpace(V, Q)
-    v, q = ufl.TestFunctions(W)
-    du, dp = ufl.TrialFunctions(W)
+    # if dolfinx.__version__ == "0.8.0":
+    v = ufl.TestFunction(V)
+    q = ufl.TestFunction(Q)
+    du = ufl.TrialFunction(V)
+    dp = ufl.TrialFunction(Q)
+    # elif dolfinx.__version__ == "0.9.0.0":
+        # Relies on: https://github.com/FEniCS/ufl/pull/308
+        # W = ufl.MixedFunctionSpace(V, Q)
+        # v, q = ufl.TestFunctions(W)
+        # du, dp = ufl.TrialFunctions(W)
+    # else:
+    #     raise ValueError("Unsupported version of dolfinx")
 
     # Set nonzero initial guess
     # Choose initial guess acording to the input factor
@@ -33,16 +42,17 @@ def test_NewtonSolver(factor):
     c1 = dolfinx.fem.Constant(mesh, ftype(0.82))
     F0 = ufl.inner(c0 * u**2, v) * ufl.dx - ufl.inner(u_expr, v) * ufl.dx
     F1 = ufl.inner(c1 * p**2, q) * ufl.dx - ufl.inner(p_expr, q) * ufl.dx
-    if dolfinx.__version__ == "0.8.0":
-        F = [F0, F1]
-        J = [
-            [ufl.derivative(F0, u, du), ufl.derivative(F0, p, du)],
-            [ufl.derivative(F1, u, dp), ufl.derivative(F1, p, dp)],
-        ]
-    elif dolfinx.__version__ == "0.9.0.0":
-        F_ = F0 + F1
-        F = list(ufl.extract_blocks(F_))
-        J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
+    # if dolfinx.__version__ == "0.8.0":
+    F = [F0, F1]
+    J = [
+        [ufl.derivative(F0, u, du), ufl.derivative(F0, p, dp)],
+        [ufl.derivative(F1, u, du), ufl.derivative(F1, p, dp)],
+    ]
+    # elif dolfinx.__version__ == "0.9.0.0":
+        # Relies on: https://github.com/FEniCS/ufl/pull/308
+        # F_ = F0 + F1
+        # F = list(ufl.extract_blocks(F_))
+        # J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
     petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
     solver = NewtonSolver(F, J, [u, p], max_iterations=25, petsc_options=petsc_options)
     solver.solve()

From 56467f1a528408fad868e63e2d2640d78783efcb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= <dokken@simula.no>
Date: Mon, 23 Sep 2024 12:08:59 +0000
Subject: [PATCH 4/9] pre-commit

---
 tests/test_blocked_newton_solver.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/tests/test_blocked_newton_solver.py b/tests/test_blocked_newton_solver.py
index d64df01..b38e48d 100644
--- a/tests/test_blocked_newton_solver.py
+++ b/tests/test_blocked_newton_solver.py
@@ -22,10 +22,10 @@ def test_NewtonSolver(factor):
     du = ufl.TrialFunction(V)
     dp = ufl.TrialFunction(Q)
     # elif dolfinx.__version__ == "0.9.0.0":
-        # Relies on: https://github.com/FEniCS/ufl/pull/308
-        # W = ufl.MixedFunctionSpace(V, Q)
-        # v, q = ufl.TestFunctions(W)
-        # du, dp = ufl.TrialFunctions(W)
+    # Relies on: https://github.com/FEniCS/ufl/pull/308
+    # W = ufl.MixedFunctionSpace(V, Q)
+    # v, q = ufl.TestFunctions(W)
+    # du, dp = ufl.TrialFunctions(W)
     # else:
     #     raise ValueError("Unsupported version of dolfinx")
 
@@ -49,10 +49,10 @@ def test_NewtonSolver(factor):
         [ufl.derivative(F1, u, du), ufl.derivative(F1, p, dp)],
     ]
     # elif dolfinx.__version__ == "0.9.0.0":
-        # Relies on: https://github.com/FEniCS/ufl/pull/308
-        # F_ = F0 + F1
-        # F = list(ufl.extract_blocks(F_))
-        # J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
+    # Relies on: https://github.com/FEniCS/ufl/pull/308
+    # F_ = F0 + F1
+    # F = list(ufl.extract_blocks(F_))
+    # J = ufl.extract_blocks(ufl.derivative(F_, u, du) + ufl.derivative(F_, p, dp))
     petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
     solver = NewtonSolver(F, J, [u, p], max_iterations=25, petsc_options=petsc_options)
     solver.solve()

From 3cc503df1f0e42b8961099b71c87356b742e54b2 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Wed, 25 Sep 2024 09:01:08 +0200
Subject: [PATCH 5/9] Add example for blocked newton solver

---
 _toc.yml                   |   1 +
 examples/blocked_solver.py | 166 +++++++++++++++++++++++++++++++++++++
 2 files changed, 167 insertions(+)
 create mode 100644 examples/blocked_solver.py

diff --git a/_toc.yml b/_toc.yml
index 03f1315..45b0ca5 100644
--- a/_toc.yml
+++ b/_toc.yml
@@ -10,6 +10,7 @@ parts:
     - file: "examples/point_source.py"
     - file: "examples/xdmf_point_cloud.py"
     - file: "examples/mark_entities.py"
+    - file: "examples/blocked_solver.py"
   - caption: Reference
     chapters:
     - file: "docs/api.rst"
diff --git a/examples/blocked_solver.py b/examples/blocked_solver.py
new file mode 100644
index 0000000..810ded0
--- /dev/null
+++ b/examples/blocked_solver.py
@@ -0,0 +1,166 @@
+# # Nonlinear elasticity with blocked Newton solver
+#
+# Author: Henrik N. T. Finsberg
+#
+# SPDX-License-Identifier: MIT
+
+# In this example we will solve a nonlinear elasticity problem using a blocked Newton solver.
+# We consider a unit cube domain $\Omega = [0, 1]^3$ with Dirichlet boundary conditions on the left face and traction force on the right face, and we seek a displacement field $\mathbf{u}: \Omega \to \mathbb{R}^3$ that solves the momentum balance equation
+#
+# $$
+# \begin{align}
+# \nabla \cdot \mathbf{P} = 0, \quad \mathbf{X} \in \Omega, \\
+# \end{align}
+# $$
+#
+# where $\mathbf{P}$ is the first Piola-Kirchhoff stress tensor, and $\mathbf{X}$ is the reference configuration. We consider a Neo-Hookean material model, where the strain energy density is given by
+#
+# $$
+# \begin{align}
+# \psi = \frac{\mu}{2}(\text{tr}(\mathbf{C}) - 3),
+# \end{align}
+# $$
+#
+# and the first Piola-Kirchhoff stress tensor is given by
+#
+# $$
+# \begin{align}
+# \mathbf{P} = \frac{\partial \psi}{\partial \mathbf{F}}
+# \end{align}
+# $$
+#
+# where $\mathbf{F} = \nabla \mathbf{u} + \mathbf{I}$ is the deformation gradient, $\mathbf{C} = \mathbf{F}^T \mathbf{F}$ is the right Cauchy-Green tensor, $\mu$ is the shear modulus, and $p$ is the pressure.
+# We also enforce the incompressibility constraint
+#
+# $$
+# \begin{align}
+# J = \det(\mathbf{F}) = 1,
+# \end{align}
+# $$
+#
+# so that the total Lagrangian is given by
+#
+# $$
+# \begin{align}
+# \mathcal{L}(\mathbf{u}, p) = \int_{\Omega} \psi \, dx - \int_{\partial \Omega} t \cdot \mathbf{u} \, ds + \int_{\Omega} p(J - 1) \, dx.
+# \end{align}
+# $$
+#
+# The Euler-Lagrange equations for this problem are given by: Find $\mathbf{u} \in V$ and $p \in Q$ such that
+#
+# $$
+# \begin{align}
+# D_{\delta \mathbf{u} } \mathcal{L}(\mathbf{u}, p) = 0, \quad \forall \delta \mathbf{u} \in V, \\
+# D_{\delta p} \mathcal{L}(\mathbf{u}, p) = 0, \quad \forall \delta p \in Q,
+# \end{align}
+# $$
+#
+# where $V$ is the displacement space and $Q$ is the pressure space. For this we select Taylor-Hood finite elements for the displacement and pressure spaces, with second order Lagrange elements for $\mathbf{u}$ and first order Lagrange elements for $p$.
+#
+#
+
+from mpi4py import MPI
+import numpy as np
+import ufl
+import dolfinx
+import scifem
+
+# We create the mesh and the function spaces
+
+mesh = dolfinx.mesh.create_unit_cube(
+    MPI.COMM_WORLD, 3,3,3, dolfinx.mesh.CellType.hexahedron, dtype=np.float64
+)
+V = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, (3,)))
+Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+
+# And the test and trial functions
+#
+
+v = ufl.TestFunction(V)
+q = ufl.TestFunction(Q)
+du = ufl.TrialFunction(V)
+dp = ufl.TrialFunction(Q)
+u = dolfinx.fem.Function(V)
+p = dolfinx.fem.Function(Q)
+
+
+# Next we create the facet tags for the left and right faces
+
+def left(x):
+    return np.isclose(x[0], 0)
+
+
+def right(x):
+    return np.isclose(x[0], 1)
+
+
+facet_tags = scifem.create_entity_markers(
+    mesh, mesh.topology.dim - 1, [(1, left, True), (2, right, True)]
+)
+
+# We create the Dirichlet boundary conditions on the left face
+
+facets_left = facet_tags.find(1)
+dofs_left = dolfinx.fem.locate_dofs_topological(V, 2, facets_left)
+u_bc_left = dolfinx.fem.Function(V)
+u_bc_left.x.array[:] = 0
+bc = dolfinx.fem.dirichletbc(u_bc_left, dofs_left)
+
+# Define the deformation gradient, right Cauchy-Green tensor, and invariants of the deformation tensors
+
+d = len(u)
+I = ufl.Identity(d)             # Identity tensor
+F = I + ufl.grad(u)             # Deformation gradient
+C = F.T*F                       # Right Cauchy-Green tensor
+I1 = ufl.tr(C)                  # First invariant of C
+J  = ufl.det(F)                 # Jacobian of F
+
+# Traction for to be applied on the right face
+
+t = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(10.0))
+
+N = ufl.FacetNormal(mesh)
+
+# Material parameters and strain energy density
+mu = dolfinx.fem.Constant(mesh, 10.0)
+psi = (mu / 2)*(I1 - 3)
+
+# We for the total Lagrangian
+
+L = psi*ufl.dx - ufl.dot(t * N, u)*ufl.ds(subdomain_data=facet_tags, subdomain_id=2)  + p * (J - 1) * ufl.dx
+
+# and take the first variation of the total Lagrangian to obtain the residual
+
+r_u = ufl.derivative(L, u, v)
+r_p = ufl.derivative(L, p, q)
+R = [r_u, r_p]
+
+# We do the same for the second variation to obtain the Jacobian
+
+K = [
+    [ufl.derivative(r_u, u, du), ufl.derivative(r_u, p, dp)],
+    [ufl.derivative(r_p, u, du), ufl.derivative(r_p, p, dp)],
+]
+
+# Now we can create the Newton solver and solve the problem
+
+petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
+solver = scifem.NewtonSolver(R, K, [u, p], max_iterations=25, bcs=[bc], petsc_options=petsc_options)
+solver.solve()
+
+# Finally, we can visualize the solution using `pyvista`
+
+import pyvista
+pyvista.start_xvfb()
+p = pyvista.Plotter()
+topology, cell_types, geometry = dolfinx.plot.vtk_mesh(V)
+grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
+grid["u"] = u.x.array.reshape((geometry.shape[0], 3))
+p.add_mesh(grid, style="wireframe", color="k")
+warped = grid.warp_by_vector("u", factor=1.5)
+p.add_mesh(warped, show_edges=True)
+p.show_axes()
+if not pyvista.OFF_SCREEN:
+    p.show()
+else:
+    figure_as_array = p.screenshot("displacement.png")

From 311a2017c7045be463274e92b02d9197604d5935 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Wed, 25 Sep 2024 09:17:53 +0200
Subject: [PATCH 6/9] Use inner instead of dot

---
 examples/blocked_solver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/blocked_solver.py b/examples/blocked_solver.py
index 810ded0..d257e09 100644
--- a/examples/blocked_solver.py
+++ b/examples/blocked_solver.py
@@ -127,7 +127,7 @@ def right(x):
 
 # We for the total Lagrangian
 
-L = psi*ufl.dx - ufl.dot(t * N, u)*ufl.ds(subdomain_data=facet_tags, subdomain_id=2)  + p * (J - 1) * ufl.dx
+L = psi*ufl.dx - ufl.inner(t * N, u)*ufl.ds(subdomain_data=facet_tags, subdomain_id=2)  + p * (J - 1) * ufl.dx
 
 # and take the first variation of the total Lagrangian to obtain the residual
 

From a3ab14010fd4d0d451e23da8d029c33754159a06 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Wed, 25 Sep 2024 10:49:55 +0200
Subject: [PATCH 7/9] Use DPc elements for the pressure

---
 _config.yml                |  3 +++
 examples/blocked_solver.py | 19 ++++++++++++++-----
 pyproject.toml             |  2 +-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/_config.yml b/_config.yml
index a801317..c53e94b 100644
--- a/_config.yml
+++ b/_config.yml
@@ -28,11 +28,13 @@ parse:
     - linkify
     - html_admonition
 
+
 sphinx:
 
   config:
     html_last_updated_fmt: "%b %d, %Y"
     nb_execution_show_tb: True
+    bibtex_bibfiles: ["docs/refs.bib"]
     suppress_warnings: ["mystnb.unknown_mime_type"]
     nb_custom_formats:  # https://jupyterbook.org/en/stable/file-types/jupytext.html#file-types-custom
         .py:
@@ -43,5 +45,6 @@ sphinx:
   - 'sphinx.ext.autodoc'
   - 'sphinx.ext.napoleon'
   - 'sphinx.ext.viewcode'
+  - 'sphinxcontrib.bibtex'
 
 exclude_patterns: [".pytest_cache/*" ,"tests", venv, .vcode, .ruff_cache, .github, .git]
diff --git a/examples/blocked_solver.py b/examples/blocked_solver.py
index d257e09..2613044 100644
--- a/examples/blocked_solver.py
+++ b/examples/blocked_solver.py
@@ -46,6 +46,7 @@
 # \end{align}
 # $$
 #
+# Here $t$ is the traction force which is set to $10$ on the right face of the cube and $0$ elsewhere.
 # The Euler-Lagrange equations for this problem are given by: Find $\mathbf{u} \in V$ and $p \in Q$ such that
 #
 # $$
@@ -55,7 +56,8 @@
 # \end{align}
 # $$
 #
-# where $V$ is the displacement space and $Q$ is the pressure space. For this we select Taylor-Hood finite elements for the displacement and pressure spaces, with second order Lagrange elements for $\mathbf{u}$ and first order Lagrange elements for $p$.
+# where $V$ is the displacement space and $Q$ is the pressure space. For this we select $Q_2/P_1$ elements i.e second order Lagrange elements for $\mathbf{u}$ and [first order discontinuous polynomial cubical elements](https://defelement.com/elements/examples/quadrilateral-dpc-1.html) for $p$, which is a stable element for incompressible elasticity {cite}`auricchio2013approximation`.
+# Note also that the Euler-Lagrange equations can be derived automatically using `ufl`.
 #
 #
 
@@ -68,10 +70,10 @@
 # We create the mesh and the function spaces
 
 mesh = dolfinx.mesh.create_unit_cube(
-    MPI.COMM_WORLD, 3,3,3, dolfinx.mesh.CellType.hexahedron, dtype=np.float64
+    MPI.COMM_WORLD, 3, 3, 3, dolfinx.mesh.CellType.hexahedron, dtype=np.float64
 )
 V = dolfinx.fem.functionspace(mesh, ("Lagrange", 2, (3,)))
-Q = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+Q = dolfinx.fem.functionspace(mesh, ("DPC", 1,))
 
 # And the test and trial functions
 #
@@ -155,12 +157,19 @@ def right(x):
 p = pyvista.Plotter()
 topology, cell_types, geometry = dolfinx.plot.vtk_mesh(V)
 grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
+linear_grid = pyvista.UnstructuredGrid(*dolfinx.plot.vtk_mesh(mesh))
 grid["u"] = u.x.array.reshape((geometry.shape[0], 3))
-p.add_mesh(grid, style="wireframe", color="k")
+p.add_mesh(linear_grid, style="wireframe", color="k")
 warped = grid.warp_by_vector("u", factor=1.5)
-p.add_mesh(warped, show_edges=True)
+p.add_mesh(warped, show_edges=False)
 p.show_axes()
 if not pyvista.OFF_SCREEN:
     p.show()
 else:
     figure_as_array = p.screenshot("displacement.png")
+
+
+# # References
+# ```{bibliography}
+# ```
+#
diff --git a/pyproject.toml b/pyproject.toml
index 316ce90..2164d8c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,7 +21,7 @@ repository = "https://github.com/scientificcomputing/scifem.git"
 [project.optional-dependencies]
 h5py = ["h5py"]
 adios2 = ["adios2"]
-docs = ["jupyter-book", "jupytext", "pyvista[trame]"]
+docs = ["jupyter-book", "jupytext", "pyvista[trame]", "sphinxcontrib-bibtex"]
 dev = ["ruff", "mypy", "bump-my-version", "pre-commit"]
 test = ["pytest", "petsc4py", "h5py"]
 all = ["scifem[docs,dev,test,audio2,h5py]"]

From 590d1a10e7a0dae6c5b43a62387df7780aedaf67 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Wed, 25 Sep 2024 12:32:11 +0200
Subject: [PATCH 8/9] Add missing bib file

---
 docs/refs.bib | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 docs/refs.bib

diff --git a/docs/refs.bib b/docs/refs.bib
new file mode 100644
index 0000000..218a050
--- /dev/null
+++ b/docs/refs.bib
@@ -0,0 +1,9 @@
+@article{auricchio2013approximation,
+  title={Approximation of incompressible large deformation elastic problems: some unresolved issues},
+  author={Auricchio, Ferdinando and Beir{\~a}o da Veiga, Louren{\c{c}}o and Lovadina, Carlo and Reali, Alessandro and Taylor, Robert L and Wriggers, Peter},
+  journal={Computational Mechanics},
+  volume={52},
+  pages={1153--1167},
+  year={2013},
+  publisher={Springer}
+}

From 2d124f2be633ce0c536947988478d713e4453028 Mon Sep 17 00:00:00 2001
From: Henrik Finsberg <henrikfinsberg@hotmail.com>
Date: Wed, 25 Sep 2024 12:49:49 +0200
Subject: [PATCH 9/9] Add logger and pass callback functions to the solver

---
 examples/blocked_solver.py | 16 ++++++++++++++++
 src/scifem/solvers.py      | 10 +++++++---
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/examples/blocked_solver.py b/examples/blocked_solver.py
index 2613044..d0e367a 100644
--- a/examples/blocked_solver.py
+++ b/examples/blocked_solver.py
@@ -144,10 +144,26 @@ def right(x):
     [ufl.derivative(r_p, u, du), ufl.derivative(r_p, p, dp)],
 ]
 
+
 # Now we can create the Newton solver and solve the problem
 
 petsc_options = {"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"}
 solver = scifem.NewtonSolver(R, K, [u, p], max_iterations=25, bcs=[bc], petsc_options=petsc_options)
+
+# We can also set a callback function that is called before and after the solve, which takes the solver object as input
+
+
+def pre_solve(solver: scifem.NewtonSolver):
+    print(f"Starting solve with {solver.max_iterations} iterations")
+
+
+def post_solve(solver: scifem.NewtonSolver):
+    print(f"Solve completed in with correction norm {solver.dx.norm(0)}")
+
+
+solver.set_pre_solve_callback(pre_solve)
+solver.set_post_solve_callback(post_solve)
+
 solver.solve()
 
 # Finally, we can visualize the solution using `pyvista`
diff --git a/src/scifem/solvers.py b/src/scifem/solvers.py
index 3da893c..93907e3 100644
--- a/src/scifem/solvers.py
+++ b/src/scifem/solvers.py
@@ -1,11 +1,15 @@
-import dolfinx
 from typing import Callable
-import ufl
+import logging
+
 import numpy as np
 from petsc4py import PETSc
+import ufl
+import dolfinx
 
 __all__ = ["NewtonSolver"]
 
+logger = logging.getLogger(__name__)
+
 
 class NewtonSolver:
     max_iterations: int
@@ -188,7 +192,7 @@ def solve(self, tol=1e-6, beta=1.0) -> int:
 
             # Compute norm of update
             correction_norm = self.dx.norm(0)
-            print(f"Iteration {i}: Correction norm {correction_norm}")
+            logger.info(f"Iteration {i}: Correction norm {correction_norm}")
             if correction_norm < tol:
                 return i
             i += 1