Merge pull request #623 from CUQI-DTU/callback_update

Enable callback in Gibbs and update callback interface to pass more information to the callback method

Author: amal-ghamdi

cuqi/experimental/mcmc/_cwmh.py

@@ -31,12 +31,9 @@ class CWMH(ProposalBasedSampler):
     initial_point : ndarray
         Initial parameters. *Optional*
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is
-        `callback(sample, sample_index)`, where `sample` is the current sample
-        and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     kwargs : dict
         Additional keyword arguments to be passed to the base class 

cuqi/experimental/mcmc/_gibbs.py

@@ -58,6 +58,10 @@ class HybridGibbs:
         will call its step method in each Gibbs step.
         Default is 1 for all variables.
 
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
+
     Example
     -------
     .. code-block:: python
@@ -103,7 +107,7 @@ class HybridGibbs:
             
     """
 
-    def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampler], num_sampling_steps: Dict[str, int] = None):
+    def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampler], num_sampling_steps: Dict[str, int] = None, callback=None):
 
         # Store target and allow conditioning to reduce to a single density
         self.target = target() # Create a copy of target distribution (to avoid modifying the original)
@@ -120,6 +124,9 @@ def __init__(self, target: JointDistribution, sampling_strategy: Dict[str, Sampl
         # Initialize sampler (after target is set)
         self._initialize()
 
+        # Set the callback function
+        self.callback = callback
+
     def _initialize(self):
         """ Initialize sampler """
 
@@ -158,13 +165,15 @@ def sample(self, Ns) -> 'HybridGibbs':
             The number of samples to draw.
 
         """
-
-        for _ in tqdm(range(Ns), "Sample: "):
+        for idx in tqdm(range(Ns), "Sample: "):
 
             self.step()
 
             self._store_samples()
 
+            # Call callback function if specified
+            self._call_callback(idx, Ns)
+
         return self
 
     def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
@@ -192,6 +201,9 @@ def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
 
             self._store_samples()
 
+            # Call callback function if specified
+            self._call_callback(idx, Nb)
+
         return self
 
     def get_samples(self) -> Dict[str, Samples]:
@@ -263,6 +275,11 @@ def tune(self, skip_len, update_count):
             self.samplers[par_name].tune(skip_len=skip_len, update_count=update_count)
 
     # ------------ Private methods ------------
+    def _call_callback(self, sample_index, num_of_samples):
+        """ Calls the callback function. Assumes input is sampler, sample index, and total number of samples """
+        if self.callback is not None:
+            self.callback(self, sample_index, num_of_samples)
+
     def _initialize_samplers(self):
         """ Initialize samplers """
         for sampler in self.samplers.values():

cuqi/experimental/mcmc/_hmc.py

@@ -38,13 +38,9 @@ class NUTS(Sampler):
         opt_acc_rate should be in (0, 1), however, choosing a value that is very
         close to 1 or 0 might lead to poor performance of the sampler.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is
-        `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of
-        the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     Example
     -------

cuqi/experimental/mcmc/_langevin_algorithm.py

@@ -32,11 +32,9 @@ class ULA(Sampler): # Refactor to Proposal-based sampler?
         be smaller than 1/L, where L is the Lipschitz of the gradient of the log
         target density, logd).
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
 
     Example
@@ -164,11 +162,9 @@ class MALA(ULA): # Refactor to Proposal-based sampler?
         be smaller than 1/L, where L is the Lipschitz of the gradient of the log
         target density, logd).
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
 
     Example
@@ -288,12 +284,9 @@ class MYULA(ULA):
     smoothing_strength : float
         This parameter controls the smoothing strength of MYULA.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of
-        the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     A Deblur example can be found in demos/howtos/myula.py
     # TODO: update demo once sampler merged
@@ -378,12 +371,9 @@ class PnPULA(MYULA):
         This parameter controls the smoothing strength of PnP-ULA.
 
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of
-        the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
 
     # TODO: update demo once sampler merged
     """

cuqi/experimental/mcmc/_laplace_approximation.py

@@ -43,11 +43,9 @@ class UGLA(Sampler):
         sampling easier but results in a worse approximation. See details in Section 3.3 of the paper.
         If not provided, it defaults to 1e-5.
 
-    callback : callable, *Optional*
-        If set, this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
     """
     def __init__(self, target=None, initial_point=None, maxit=50, tol=1e-4, beta=1e-5, **kwargs):
 

cuqi/experimental/mcmc/_rto.py

@@ -36,11 +36,9 @@ class LinearRTO(Sampler):
     tol : float
         Tolerance of the inner CGLS solver. *Optional*.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         
     """
     def __init__(self, target=None, initial_point=None, maxit=10, tol=1e-6, **kwargs):
@@ -204,11 +202,9 @@ class RegularizedLinearRTO(LinearRTO):
     solver : string
         If set to "ScipyLinearLSQ", solver is set to cuqi.solver.ScipyLinearLSQ, otherwise FISTA/ISTA or ADMM is used. Note "ScipyLinearLSQ" can only be used with `RegularizedGaussian` of `box` or `nonnegativity` constraint. *Optional*.
 
-    callback : callable, *Optional*
-        If set this function will be called after every sample.
-        The signature of the callback function is `callback(sample, sample_index)`,
-        where `sample` is the current sample and `sample_index` is the index of the sample.
-        An example is shown in demos/demo31_callback.py.
+    callback : callable, optional
+        A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+        The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         
     """
     def __init__(self, target=None, initial_point=None, maxit=100, inner_max_it=10, stepsize="automatic", penalty_parameter=10, abstol=1e-10, adaptive=True, solver=None, inner_abstol=None, **kwargs):

cuqi/experimental/mcmc/_sampler.py

@@ -59,9 +59,8 @@ def __init__(self, target:cuqi.density.Density=None, initial_point=None, callbac
             The initial point for the sampler. If not given, the sampler will choose an initial point.
 
         callback : callable, optional
-            A function that will be called after each sample is drawn. The function should take two arguments: the sample and the index of the sample.
-            The sample is a 1D numpy array and the index is an integer. The callback function is useful for monitoring the sampler during sampling.
-
+            A function that will be called after each sampling step. It can be useful for monitoring the sampler during sampling.
+            The function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature is: `callback(sampler, sample_index, num_of_samples)`.
         """
 
         self.target = target
@@ -209,7 +208,6 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
             The path to save the samples. If not specified, the samples are saved to the current working directory under a folder called 'CUQI_samples'.
 
         """
-
         self._ensure_initialized()
 
         # Initialize batch handler
@@ -235,7 +233,7 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
                 batch_handler.add_sample(self.current_point)
 
             # Call callback function if specified            
-            self._call_callback(self.current_point, len(self._samples)-1)
+            self._call_callback(idx, Ns)
 
         return self
 
@@ -276,7 +274,7 @@ def warmup(self, Nb, tune_freq=0.1) -> 'Sampler':
             pbar.set_postfix_str(f"acc rate: {np.mean(self._acc[-1-idx:]):.2%}")
 
             # Call callback function if specified
-            self._call_callback(self.current_point, len(self._samples)-1)
+            self._call_callback(idx, Nb)
 
         return self
 
@@ -367,10 +365,10 @@ def set_history(self, history: dict):
                 raise ValueError(f"Key {key} not recognized in history dictionary of sampler {self.__class__.__name__}.")
 
     # ------------ Private methods ------------
-    def _call_callback(self, sample, sample_index):
-        """ Calls the callback function. Assumes input is sample and sample index"""
+    def _call_callback(self, sample_index, num_of_samples):
+        """ Calls the callback function. Assumes input is sampler, sample index, and total number of samples """
         if self.callback is not None:
-            self.callback(sample, sample_index)
+            self.callback(self, sample_index, num_of_samples)
 
     def _validate_initialization(self):
         """ Validate the initialization of the sampler by checking all state and history keys are set. """

cuqi/problem/_problem.py

@@ -304,6 +304,7 @@ def sample_posterior(self, Ns, Nb=None, callback=None, experimental=False) -> cu
             The signature of the callback function is `callback(sample, sample_index)`,
             where `sample` is the current sample and `sample_index` is the index of the sample.
             An example is shown in demos/demo31_callback.py.
+            Note: if the parameter `experimental` is set to True, the callback function should take three arguments: the sampler object, the index of the current sampling step, the total number of requested samples. The last two arguments are integers. An example of the callback function signature in the case is: `callback(sampler, sample_index, num_of_samples)`.
 
         experimental : bool, *Optional*
             If set to True, the sampler selection will use the samplers from the :mod:`cuqi.experimental.mcmc` module.
@@ -848,16 +849,14 @@ def _sampleGibbs(self, Ns, Nb, callback=None, experimental=False):
             print(f"burn-in: {Nb/Ns*100:g}%")
             print("")
 
-            if callback is not None:
-                raise NotImplementedError("Callback not implemented for Gibbs sampler")
-
             # Start timing
             ti = time.time()
 
             # Sampling strategy
             sampling_strategy = self._determine_sampling_strategy(experimental=True)
 
-            sampler = cuqi.experimental.mcmc.HybridGibbs(self._target, sampling_strategy)
+            sampler = cuqi.experimental.mcmc.HybridGibbs(
+                self._target, sampling_strategy, callback=callback)
             sampler.warmup(Nb)
             sampler.sample(Ns)
             samples = sampler.get_samples()
@@ -876,7 +875,7 @@ def _sampleGibbs(self, Ns, Nb, callback=None, experimental=False):
             print("")
 
             if callback is not None:
-                raise NotImplementedError("Callback not implemented for Gibbs sampler")
+                raise NotImplementedError("Callback not implemented for Gibbs sampler. It is only implemented for experimental Gibbs sampler.")
 
             # Start timing
             ti = time.time()