diff --git a/docs/api/neurodamus.core.rst b/docs/api/neurodamus.core.rst index 0796ba764..e8943e9ec 100644 --- a/docs/api/neurodamus.core.rst +++ b/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 diff --git a/neurodamus/core/_utils.py b/neurodamus/core/_utils.py index 8f82ddc78..ff41ce5e0 100644 --- a/neurodamus/core/_utils.py +++ b/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 diff --git a/neurodamus/node.py b/neurodamus/node.py index 685a9b864..414309419 100644 --- a/neurodamus/node.py +++ b/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.")