Add dolfin timers

smart/model.py

@@ -214,6 +214,7 @@ def initialize(self, initialize_solver=True):
         """Main model initialization function, split into 5 subfunctions"""
 
         # Solver related parameters
+        timer = d.Timer("Initialize Model")
         self._base_t = Decimal("0." + (self.config.solver["time_precision"] - 1) * "0" + "1")
         self.t = self.rounded_decimal(0.0)
         self.dt = self.rounded_decimal(self.config.solver["initial_dt"])
@@ -238,9 +239,11 @@ def initialize(self, initialize_solver=True):
             self.cc.print()
             self.print_meshes()
             self.rc.print()
+        timer.stop()
 
     def _init_1(self):
         """Checking validity of model"""
+        timer = d.Timer("Initialize model step 1")
         logger.debug("Checking validity of model (step 1 of ZZ)", extra=dict(format_type="title"))
         self._init_1_1_check_mesh_dimensionality()
         self._init_1_2_check_namespace_conflicts()
@@ -249,10 +252,12 @@ def _init_1(self):
             "Step 1 of initialization completed successfully!",
             extra=dict(text_color="magenta"),
         )
+        timer.stop()
 
     def _init_2(self):
         """Cross-container dependent initializations
         (requires information from multiple containers)"""
+        timer = d.Timer("Initialize model step 2")
         logger.debug(
             "Cross-Container Dependent Initializations (step 2 of ZZ)",
             extra=dict(
@@ -270,9 +275,11 @@ def _init_2(self):
             "Step 2 of initialization completed " "successfully!",
             extra=dict(text_color="magenta"),
         )
+        timer.stop()
 
     def _init_3(self):
         """Mesh-related initializations"""
+        timer = d.Timer("Initialize model step 3")
         logger.debug(
             "Mesh-related Initializations (step 3 of ZZ)",
             extra=dict(format_type="title"),
@@ -289,9 +296,11 @@ def _init_3(self):
             "Step 3 of initialization completed successfully!",
             extra=dict(format_type="log_important"),
         )
+        timer.stop()
 
     def _init_4(self):
         """Dolfin function initializations"""
+        timer = d.Timer("Initialize model step 4")
         logger.debug("Dolfin Initializations (step 4 of ZZ)", extra=dict(format_type="title"))
         self._init_4_0_initialize_dolfin_parameters()
         self._init_4_1_get_active_compartments()
@@ -301,11 +310,13 @@ def _init_4(self):
         self._init_4_5_name_functions()
         self._init_4_6_check_dolfin_function_validity()
         self._init_4_7_set_initial_conditions()
+        timer.stop()
 
     def _init_5(self, initialize_solver):
         """Convert reactions to fluxes and define variational form.
         If initialize_solver is true, also initialize solver.
         """
+        timer = d.Timer("Initialize model step 5")
         logger.debug(
             "Dolfin fluxes, forms, and problems+solvers (step 5 of ZZ)",
             extra=dict(
@@ -317,6 +328,8 @@ def _init_5(self, initialize_solver):
         if initialize_solver:
             self.initialize_discrete_variational_problem_and_solver()
 
+        timer.stop()
+
     # Step 1 - Checking model validity
 
     def _init_1_1_check_mesh_dimensionality(self):

smart/solvers.py

@@ -49,6 +49,7 @@ def __init__(
         all_compartments: List[Compartment],
         stopwatches: Dict[str, Stopwatch],
     ):
+        timer = d.Timer("Initialize smartSNESProblem")
         self.u = u
         self.Fforms = Fforms
         self.Jforms_all = Jforms_all
@@ -112,8 +113,10 @@ def __init__(
                 f"Forms {self.empty_forms} are empty. Skipping assembly.",
                 extra=dict(format_type="data"),
             )
+        timer.stop()
 
     def init_petsc_matnest(self):
+        timer = d.Timer("Initialize PETSc Nested Matrix")
         Jforms = self.Jforms_all
         dim = self.dim
         Jpetsc = []
@@ -191,11 +194,13 @@ def init_petsc_matnest(self):
             # self.Jpetsc_nest = self.d_to_p(d.PETScNestMatrix(Jpetsc))
         self.Jpetsc_nest.assemble()
         logger.info(f"Jpetsc_nest assembled, size = {self.Jpetsc_nest.size}")
+        timer.stop()
 
     def d_to_p(self, dolfin_matrix):
         return d.as_backend_type(dolfin_matrix).mat()
 
     def init_petsc_vecnest(self):
+        timer = d.Timer("Initialize PETSc Nested Vector")
         dim = self.dim
         logger.info("Initializing block residual vector", extra=dict(format_type="assembly"))
 
@@ -236,6 +241,7 @@ def init_petsc_vecnest(self):
         else:
             self.Fpetsc_nest = p.Vec().createNest(Fpetsc, comm=self.comm)
         self.Fpetsc_nest.assemble()
+        timer.stop()
 
     def assemble_Jnest(self, Jnest):
         """Assemble Jacobian nest matrix.
@@ -246,6 +252,7 @@ def assemble_Jnest(self, Jnest):
 
 
         """
+        timer = d.Timer("SNES Assemble Jacobian Nested Matrix")
         logger.debug("Assembling block Jacobian", extra=dict(format_type="assembly"))
         self.stopwatches["snes jacobian assemble"].start()
         dim = self.dim
@@ -315,6 +322,7 @@ def assemble_Jnest(self, Jnest):
         Jnest.assemble()
 
         self.stopwatches["snes jacobian assemble"].pause()
+        timer.stop()
 
     def assemble_Fnest(self, Fnest):
         """
@@ -324,6 +332,7 @@ def assemble_Fnest(self, Fnest):
             Fnest : petsc4py.Vec
                 PETSc nest vector representing the residual
         """
+        timer = d.Timer("SNES Assemble Residual Nest Vector")
         dim = self.dim
         logger.debug("Assembling block residual vector", extra=dict(format_type="assembly"))
         self.stopwatches["snes residual assemble"].start()
@@ -347,6 +356,7 @@ def assemble_Fnest(self, Fnest):
 
         Fnest.assemble()
         self.stopwatches["snes residual assemble"].pause()
+        timer.stop()
 
     def copy_u(self, unest):
         if self.is_single_domain:
@@ -375,6 +385,7 @@ def init_petsc_matrix(self, i, j, nnz_guess=None, set_lgmap=False, assemble=Fals
             nnz_guess : number of non-zeros (per row) to guess for the matrix
             assemble : whether to assemble the matrix or not (Boolean)
         """
+        timer = d.Timer("SNES Initialize Zero Matrices")
         self.stopwatches["snes initialize zero matrices"].start()
 
         M = p.Mat().create(comm=self.comm)
@@ -403,6 +414,7 @@ def init_petsc_matrix(self, i, j, nnz_guess=None, set_lgmap=False, assemble=Fals
         if assemble:
             M.assemble()
         self.stopwatches["snes initialize zero matrices"].pause()
+        timer.stop()
 
         return M
 
@@ -413,13 +425,16 @@ def init_petsc_vector(self, j, assemble=False):
             j : index
             assemble : whether to assemble the vector or not (Boolean)
         """
+        timer = d.Timer("SNES Initialize Zero Vectors")
         V = p.Vec().create(comm=self.comm)
         V.setSizes((self.local_sizes[j], self.global_sizes[j]))
         V.setUp()
         # V.setLGMap(self.lgmaps_petsc[j])
 
         if assemble:
             V.assemble()
+        timer.stop()
+
         return V
 
     def Jijk_name(self, i: int, j: int, k: Optional[int] = None):