From 2d8f6994c1e664d70b3429858da196a53a4c2c93 Mon Sep 17 00:00:00 2001
From: Sebastian Blauth <sebastian.blauth@itwm.fraunhofer.de>
Date: Mon, 21 Oct 2024 17:19:58 +0200
Subject: [PATCH] Add debug logs for Newton solver and snes solver

---
 cashocs/nonlinear_solvers/newton_solver.py | 5 +++++
 cashocs/nonlinear_solvers/snes.py          | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/cashocs/nonlinear_solvers/newton_solver.py b/cashocs/nonlinear_solvers/newton_solver.py
index 247fd6ab..11bcba55 100644
--- a/cashocs/nonlinear_solvers/newton_solver.py
+++ b/cashocs/nonlinear_solvers/newton_solver.py
@@ -218,6 +218,7 @@ def _print_output(self) -> None:
 
     def _assemble_matrix(self) -> None:
         """Assembles the matrix for solving the linear problem."""
+        log.begin("Assembling the Jacobian for Newton's method.", level=log.DEBUG)
         self.assembler.assemble(self.A_fenics)
         self.A_fenics.ident_zeros()
         self.A_matrix = fenics.as_backend_type(  # pylint: disable=invalid-name
@@ -233,6 +234,8 @@ def _assemble_matrix(self) -> None:
         else:
             self.P_matrix = None
 
+        log.end()
+
     def _compute_eta_inexact(self) -> None:
         """Computes the parameter ``eta`` for the inexact Newton method."""
         if self.inexact and isinstance(self.eta, float):
@@ -373,6 +376,7 @@ def _check_for_divergence(self) -> None:
 
     def _compute_residual(self) -> None:
         """Computes the residual of the nonlinear system."""
+        log.begin("Assembling the residual for Newton's method.", level=log.DEBUG)
         self.residual = fenics.PETScVector(self.comm)
         self.assembler.assemble(self.residual, self.u.vector())
         if (
@@ -384,6 +388,7 @@ def _compute_residual(self) -> None:
             self.residual[:] -= self.residual_shift[:]
 
         self.b = fenics.as_backend_type(self.residual).vec()
+        log.end()
 
     def _compute_convergence_tolerance(self) -> float:
         """Computes the tolerance for the Newton solver.
diff --git a/cashocs/nonlinear_solvers/snes.py b/cashocs/nonlinear_solvers/snes.py
index f11bdb23..ab2241c8 100644
--- a/cashocs/nonlinear_solvers/snes.py
+++ b/cashocs/nonlinear_solvers/snes.py
@@ -32,6 +32,7 @@
 
 from cashocs import _exceptions
 from cashocs import _utils
+from cashocs import log
 
 if TYPE_CHECKING:
     from cashocs import _typing
@@ -172,6 +173,7 @@ def assemble_function(
             f: The vector in which the function evaluation is stored.
 
         """
+        log.begin("Assembling the residual for Newton's method.", level=log.DEBUG)
         self.u.vector().vec().setArray(x)
         self.u.vector().apply("")
         f = fenics.PETScVector(f)
@@ -184,6 +186,7 @@ def assemble_function(
         ):
             self.assembler_shift.assemble(self.residual_shift, self.u.vector())
             f[:] -= self.residual_shift[:]
+        log.end()
 
     def assemble_jacobian(
         self,
@@ -202,6 +205,7 @@ def assemble_jacobian(
 
         """
         if not self.is_preassembled:
+            log.begin("Assembling the Jacobian for Newton's method.", level=log.DEBUG)
             self.u.vector().vec().setArray(x)
             self.u.vector().apply("")
 
@@ -213,6 +217,7 @@ def assemble_jacobian(
                 P = fenics.PETScMatrix(P)  # pylint: disable=invalid-name
                 self.assembler_pc.assemble(P)
                 P.ident_zeros()
+            log.end()
         else:
             self.is_preassembled = False