Callback improvements: timing, csv, interval

CHANGELOG.md

@@ -1,6 +1,9 @@
 * XX.X.X
   - New features:
     - LSQR algorithm added to the CIL algorithm class (#1975)
+    - Callback improvements:
+      - Added `Callback.interval` and `Callback.skip_iteration(Algorithm)` convenience methods (#1909)
+      - Added `TimingCallback` and `CSVCallback`
   - Bug fixes:
     - `CentreOfRotationCorrector.image_sharpness` data is now correctly smoothed to reduce aliasing artefacts and improve robustness. (#2202)
     - `PaganinProcessor` now correctly applies scaling with magnification for cone-beam geometry (#2225)

Wrappers/Python/cil/optimisation/utilities/callbacks.py

@@ -17,12 +17,16 @@
 # - CIL Developers, listed at: https://github.com/TomographicImaging/CIL/blob/master/NOTICE.txt
 
 
+import csv
 from abc import ABC, abstractmethod
 from functools import partialmethod
+from pathlib import Path
+from time import time
+from typing import List, Optional
 
+import numpy as np
 from tqdm.auto import tqdm as tqdm_auto
 from tqdm.std import tqdm as tqdm_std
-import numpy as np
 
 tqdm_std.monitor_interval = 0  # disable background monitoring thread
 
@@ -34,14 +38,20 @@ class Callback(ABC):
     ----------
     verbose: int, choice of 0,1,2, default 1
         0=quiet, 1=info, 2=debug.
+    interval: int, default 1
+        Number of algorithm iterations between callback calls.
     '''
-    def __init__(self, verbose=1):
+    def __init__(self, verbose=1, interval=1):
         self.verbose = verbose
+        self.interval = interval
 
     @abstractmethod
     def __call__(self, algorithm):
         pass
 
+    def skip_iteration(self, algorithm) -> bool:
+        return algorithm.iteration % min(self.interval, max(1, algorithm.update_objective_interval)) != 0 and algorithm.iteration != algorithm.max_iteration
+
 
 class _OldCallback(Callback):
     '''Converts an old-style :code:`def callback` to a new-style :code:`class Callback`.
@@ -68,8 +78,8 @@ class ProgressCallback(Callback):
     **tqdm_kwargs:
         Passed to :code:`tqdm_class`.
     '''
-    def __init__(self, verbose=1, tqdm_class=tqdm_auto, **tqdm_kwargs):
-        super().__init__(verbose=verbose)
+    def __init__(self, verbose=1, interval=1, tqdm_class=tqdm_auto, **tqdm_kwargs):
+        super().__init__(verbose=verbose, interval=interval)
         self.tqdm_class = tqdm_class
         self.tqdm_kwargs = tqdm_kwargs
         self._obj_len = 0  # number of objective updates
@@ -80,6 +90,7 @@ def __call__(self, algorithm):
             tqdm_kwargs.setdefault('total', algorithm.max_iteration)
             tqdm_kwargs.setdefault('disable', not self.verbose)
             tqdm_kwargs.setdefault('initial', max(0, algorithm.iteration))
+            tqdm_kwargs.setdefault('miniters', min(self.interval, max(1, algorithm.update_objective_interval)))
             self.pbar = self.tqdm_class(**tqdm_kwargs)
         if (obj_len := len(algorithm.objective)) != self._obj_len:
             self.pbar.set_postfix(algorithm.objective_to_dict(self.verbose>=2), refresh=False)
@@ -130,18 +141,11 @@ class TextProgressCallback(ProgressCallback):
 
     Parameters
     ----------
-    miniters: int, default :code:`Algorithm.update_objective_interval`
+    miniters: int, default :code:`min(self.interval, max(1, Algorithm.update_objective_interval))`
         Number of algorithm iterations between screen prints.
     '''
     __init__ = partialmethod(ProgressCallback.__init__, tqdm_class=_TqdmText)
 
-    def __call__(self, algorithm):
-        if not hasattr(self, 'pbar'):
-            self.tqdm_kwargs['miniters'] = min((
-                self.tqdm_kwargs.get('miniters', algorithm.update_objective_interval),
-                algorithm.update_objective_interval))
-        return super().__call__(algorithm)
-
 
 class LogfileCallback(TextProgressCallback):
     ''':code:`TextProgressCallback` but to a file instead of screen.
@@ -157,6 +161,46 @@ def __init__(self, log_file, mode='a', **kwargs):
         self.fd = open(log_file, mode=mode)
         super().__init__(file=self.fd, **kwargs)
 
+
+class CSVCallback(Callback):
+    """Saves :code:`algo.loss` in :code:`csv_file`"""
+    def __init__(self, csv_file='objectives.csv', append=False, **kwargs):
+        super().__init__(**kwargs)
+        csv_file = Path(csv_file)
+        csv_file.parent.mkdir(parents=True, exist_ok=True)
+        if csv_file.exists() and append:
+            self.csv = csv.writer(csv_file.open('a', buffering=1))
+        else:
+            self.csv = csv.writer(csv_file.open('w', buffering=1))
+            self.csv.writerow(("iteration", "objective"))
+
+    def __call__(self, algorithm):
+        if not self.skip_iteration(algorithm):
+            self.csv.writerow((algorithm.iteration, algorithm.get_last_loss()))
+
+
+class TimingCallback(Callback):
+    """
+    Measures time taken by each iteration in :code:`self.times`,
+    excluding time taken by other specified (nested) :code:`callbacks`.
+    """
+    def __init__(self, callbacks: Optional[List[Callback]]=None, **kwargs):
+        super().__init__(**kwargs)
+        self.times = []
+        self.callbacks = callbacks or []
+        self.reset()
+
+    def reset(self):
+        self.offset = time()
+
+    def __call__(self, algorithm):
+        time_excluding_callbacks = (now := time()) - self.offset
+        self.times.append(time_excluding_callbacks)
+        for c in self.callbacks:
+            c(algorithm)
+        self.offset += time() - now
+
+
 class EarlyStoppingObjectiveValue(Callback):
     '''Callback that stops iterations if the change in the objective value is less than a provided threshold value.
 
@@ -172,12 +216,12 @@ class EarlyStoppingObjectiveValue(Callback):
     def __init__(self, threshold=1e-6):
         self.threshold=threshold
 
-
     def __call__(self, algorithm):
         if len(algorithm.loss)>=2:
             if np.abs(algorithm.loss[-1]-algorithm.loss[-2])<self.threshold:
                 raise StopIteration
 
+
 class CGLSEarlyStopping(Callback):
     '''Callback to work with CGLS. It causes the algorithm to terminate if  :math:`||A^T(Ax-b)||_2 < \epsilon||A^T(Ax_0-b)||_2` where `epsilon` is set to default as '1e-6', :math:`x` is the current iterate and :math:`x_0` is the initial value.
     It will also terminate if the algorithm begins to diverge i.e. if :math:`||x||_2> \omega`, where `omega` is set to default as 1e6.
@@ -197,9 +241,7 @@ def __init__(self, epsilon=1e-6, omega=1e6):
         self.epsilon=epsilon
         self.omega=omega
 
-
     def __call__(self, algorithm):
-
         if (algorithm.norms <= algorithm.norms0 * self.epsilon):
             print('The norm of the residual is less than {} times the norm of the initial residual and so the algorithm is terminated'.format(self.epsilon))
             raise StopIteration

Wrappers/Python/test/test_algorithms.py

@@ -1665,11 +1665,10 @@ def old_callback(iteration, objective, solution):
 
         log = NamedTemporaryFile(delete=False)
         log.close()
-        algo.run(20, callbacks=[callbacks.LogfileCallback(
-            log.name)], callback=old_callback)
+        algo.run(20, callbacks=[callbacks.LogfileCallback(log.name, interval=5)], callback=old_callback)
         with open(log.name, 'r') as fd:
             self.assertListEqual(
-                ["64/83", "74/83", "83/83", ""],
+                ['64/83', '69/83', '74/83', '79/83', '83/83', ''],
                 [line.lstrip().split(" ", 1)[0] for line in fd.readlines()])
         unlink(log.name)