Add debug logs for Newton solver and snes solver

sblauth · Oct 21, 2024 · 2d8f699 · 2d8f699
1 parent d95073e
commit 2d8f699
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git 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
@@ -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