Create a separate decorator for the progress

docs/api/neurodamus.core.rst

@@ -20,7 +20,7 @@ neurodamus.core package
    .. rubric:: Decorators
 
    .. autosummary::
-      return_neuron_timings_and_progress
+      return_neuron_timings
       mpi_no_errors
       run_only_rank0
 
@@ -67,7 +67,7 @@ Module API
 
 **Decorators**
 
-.. autofunction:: return_neuron_timings_and_progress
+.. autofunction:: return_neuron_timings
 
 .. autofunction:: mpi_no_errors
 

neurodamus/core/_utils.py

@@ -69,15 +69,23 @@ def rank0_wrapper(*args, **kw):
 
 
 class SimulationProgress:
-    def __init__(self):
+    def __init__(self, f):
         """
-        Utility class which will set up a timer to perioducally check the amount of time lapsed
+        Decorator class which will set up a timer to perioducally check the amount of time lapsed
         in the simulation compared to the final tstop value.  This is converted into a percentage
         of the job complete which is then printed to the console.
         """
+        self.f = f
+        wraps(f)(self)
+
+    def __call__(self, *args, **kwargs):
+        # Start / update progress tracking
         self.last_time_check = time.time()
         self.sim_start = self.last_time_check
         self.update_progress()
+        # Execute the decorated function
+        result = self.f(*args, **kwargs)
+        return result
 
     @run_only_rank0
     def update_progress(self):
@@ -99,7 +107,7 @@ def update_progress(self):
         Nd.cvode.event(sim_t + 1, self.update_progress)
 
 
-def return_neuron_timings_and_progress(f):
+def return_neuron_timings(f):
     """Decorator to collect, return timings and show the progress on a neuron run
     """
     @wraps(f)
@@ -110,7 +118,6 @@ def timings_wrapper(*args, **kw):
         pc = MPI.pc
         wait_base = pc.wait_time()
 
-        SimulationProgress()
         f(*args, **kw)  # Discard return values
 
         tdat[0] = pc.wait_time() - wait_base

neurodamus/node.py

@@ -13,7 +13,7 @@
 from collections import namedtuple, defaultdict
 from contextlib import contextmanager
 
-from .core import MPI, mpi_no_errors, return_neuron_timings_and_progress, run_only_rank0
+from .core import MPI, mpi_no_errors, return_neuron_timings, run_only_rank0, SimulationProgress
 from .core import NeurodamusCore as Nd
 from .core.configuration import CircuitConfig, Feature, GlobalConfig, SimConfig
 from .core._engine import EngineBase
@@ -1342,7 +1342,8 @@ def run_all(self):
         return timings
 
     # -
-    @return_neuron_timings_and_progress
+    @return_neuron_timings
+    @SimulationProgress
     def _run_neuron(self):
         self.solve()
         logging.info("Simulation finished.")