diff --git a/src/dxtb/__init__.py b/src/dxtb/__init__.py index 87deeceb..b27ace98 100644 --- a/src/dxtb/__init__.py +++ b/src/dxtb/__init__.py @@ -22,7 +22,7 @@ """ # import timer first to get correct total time -from dxtb._src.timing import timer +from dxtb._src.timing import timer, kill_timer timer.start("Import") timer.start("PyTorch", parent_uid="Import") @@ -54,7 +54,7 @@ ############################################################################### -# stop timers and remove from global namespace +# stop timers and remove PyTorch from global namespace for cleaner API del torch timer.stop("dxtb") timer.stop("Import") @@ -62,12 +62,17 @@ ############################################################################### __all__ = [ - "calculators", "components", + # + "calculators", "Calculator", "GFN1_XTB", "GFN2_XTB", + # "IndexHelper", + # + "kill_timer", "timer", + # "__version__", ] diff --git a/src/dxtb/_src/timing/timer.py b/src/dxtb/_src/timing/timer.py index e56392d3..925eb04e 100644 --- a/src/dxtb/_src/timing/timer.py +++ b/src/dxtb/_src/timing/timer.py @@ -29,7 +29,7 @@ import time -__all__ = ["timer"] +__all__ = ["timer", "create_timer", "kill_timer"] class TimerError(Exception): @@ -188,8 +188,8 @@ def __init__( ) -> None: self.label = label self.timers = {} - self._enabled = True - self._subtimer_parent_map = {} + self._enabled: bool = True + self._subtimer_parent_map: dict[str, str] = {} self._autostart = autostart self._cuda_sync = cuda_sync self._only_parents = only_parents @@ -368,6 +368,9 @@ def kill(self) -> None: self.reset() self.stop_all() + self.timers.clear() + self._subtimer_parent_map.clear() + def get_time(self, uid: str) -> float: """ Get the elapsed time of a timer. @@ -463,6 +466,66 @@ def print(self, v: int = 5, precision: int = 3) -> None: # pragma: no cover precision=precision, ) + def __str__(self) -> str: + """Return a string representation of the :class:`._Timers` instance.""" + timers_repr = ", ".join( + f"'{label}': {timer.elapsed_time:.3f}s" + for label, timer in self.timers.items() + ) + + return ( + f"{self.__class__.__name__}(" + f"label={self.label!r}, " + f"enabled={self._enabled}, " + f"cuda_sync={self._cuda_sync}, " + f"only_parents={self._only_parents}, " + f"timers={{{timers_repr}}}" + f")" + ) + + def __repr__(self) -> str: + """Return a string representation of the :class:`._Timers` instance.""" + return str(self) + + +def create_timer(autostart: bool = True, cuda_sync: bool = False) -> _Timers: + """ + Create a new timer instance. + + Parameters + ---------- + autostart : bool, optional + Whether to start the total timer automatically. Defaults to ``True``. + cuda_sync : bool, optional + Whether to call :func:`torch.cuda.synchronize` after CUDA operations. + Defaults to ``False``. + + Returns + ------- + _Timers + Instance of the timer class. + + Note + ---- + Delete the timer instance with :func:`.kill_timer` when it is no longer + needed or throws errors when reusing it. + """ + global timer + timer = _Timers(autostart=autostart, cuda_sync=cuda_sync) + return timer + + +def kill_timer() -> None: + """Delete the global timer instance.""" + global timer + if "timer" not in globals(): + raise TimerError( + "Cannot delete timer instance; timer was never initialized." + ) + + timer.kill() + del timer + -timer = _Timers(autostart=True, cuda_sync=False) +timer = create_timer(autostart=True, cuda_sync=False) """Global instance of the timer class."""