Repository: CUQI-DTU/CUQIpy

Issue:
Enable sampler callback in Gibbs and change interface of callback to pass more information
**General description of task**

Enable sampler callback in Gibbs and enable callback to access more information, e.g. Ns, all samples, etc. This update is only in the experimental module and not breaking backward compatibility of the sampler module.

**Definition of Done (Feature/change)**
- [x] Change interface of callback to pass the sampler object itself and the sample_index and total number of samples instead of the current sample and the sample_index (note: we can also pass both sample object and current_sample)
- [x] Callback enabled in HybridGibbs
- [x] Update any documentation or demos that use callback (none use the experimental module callback)

**Definition of Done (Mandatory)**
- [x] Documentation added (docstrings on all public methods/classes)
- [x] Unit tests added/updated (and passing!)
- [ ] Code reviewed and approved by 2 CUQI-maintainers
- [x] Online documentation is rendered correctly (esp. math) 

**Defintion of Done (Optional)**
- [x] (I suggest no need for now) HowTo/tutorial added


**Example**
```python
def callback(sampler, sample_index, num_of_samples)
    if sample_index %  num_of_sample == 100:
        # do something. e.g. save samples, plot samples, look at some diagnostics
        # something like this
        samples = sampler.get_samples()
        print(samples.compute_ESS()) 

```

Hints from issue comments:
@amal-ghamdi: @jakobsj , ready for you DoD review!@jakobsj: @amal-ghamdi I would like to discuss this one to better understand the usage need for and implications of this potentially wide-ranging change. At a coming daily could you give a brief overview please?@amal-ghamdi: @jakobsj I think no need for me to update the DoD after our discussion today and you can approve if you agree to it. Just few comments:
- The demo we looked at for classical sampler is https://github.com/CUQI-DTU/CUQIpy/blob/main/demos/demo31_callback.py
- This PR also enables callback in Gibbs which was not supported.
- I suggest to skip howto/demo for this PR but let me know if you prefer otherwise.
- We talked about memory consumption and that the new approach is potentially saving memory, if you think profiling is needed I can look into that.

Thanks!@jakobsj: Thanks @amal-ghamdi .  Somehow I had understood passing the "samples" object, not the "sampler" object. Would one make more sense than the other (and why)?@amal-ghamdi: Thanks @jakobsj. Good point!

The logic I propose is that the sampler passes itself. 
- If the callback has access to the sampler, it can access the samples
- Also other diagnostics (like tree size in NUT)
- Additionally, the user can potentially do some of the customized exploration of sampler behaviour as @chaozg suggested if the user wants to explore such things (for example, explore step size adaptation, e.g. change steps size based on NUTS tree size).
- The call back can also access the forward model `sampler.target.model` which is something I needed in the brain efflux project to apply the model to samples mean, etc inside the call back because I was saving data and generating prediction/reconstruction plots every n_th sample 
- It is just a more flexible setup where you get access to everything you might or might not need
- Since I would think of callback as an advanced feature anyway, I do not mind giving the user more control by accessing the sampler object itself.

The second option,
If the sampler passes the samples object instead, the sampler will need to create a samples object every sampling step (which might or might not be a critical overhead, depending on how the samples array is created from the array list). It is a safer option in terms of not giving the user so much control, but there are limitations on the things the user can access, as explained above.  I can explore this option and inspect memory usage if we decide to go with the second option.

We can discuss more if there are still aspects you are concerned about.@jakobsj: Thank you @amal-ghamdi for the additional clarification. I now see that there are strong reasons to go with the proposed approach of the sampler passing itself for maximum flexibility rather than just the samples object.  DoD approved!

Solution patch:
diff --git a/cuqi/experimental/mcmc/_cwmh.py b/cuqi/experimental/mcmc/_cwmh.py
index b05e0c9ec..84ee1e106 100644
--- a/cuqi/experimental/mcmc/_cwmh.py
+++ b/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 
diff --git a/cuqi/experimental/mcmc/_gibbs.py b/cuqi/experimental/mcmc/_gibbs.py
index e7673cdd8..ba23c585b 100644
--- a/cuqi/experimental/mcmc/_gibbs.py
+++ b/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():
diff --git a/cuqi/experimental/mcmc/_hmc.py b/cuqi/experimental/mcmc/_hmc.py
index 288805755..4c56bbe17 100644
--- a/cuqi/experimental/mcmc/_hmc.py
+++ b/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
     -------
diff --git a/cuqi/experimental/mcmc/_langevin_algorithm.py b/cuqi/experimental/mcmc/_langevin_algorithm.py
index 5759a8548..021ac2bb2 100644
--- a/cuqi/experimental/mcmc/_langevin_algorithm.py
+++ b/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
     """
diff --git a/cuqi/experimental/mcmc/_laplace_approximation.py b/cuqi/experimental/mcmc/_laplace_approximation.py
index 649b613fb..b243609af 100644
--- a/cuqi/experimental/mcmc/_laplace_approximation.py
+++ b/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):
 
diff --git a/cuqi/experimental/mcmc/_rto.py b/cuqi/experimental/mcmc/_rto.py
index 22d51d693..bd2d6ea46 100644
--- a/cuqi/experimental/mcmc/_rto.py
+++ b/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):
diff --git a/cuqi/experimental/mcmc/_sampler.py b/cuqi/experimental/mcmc/_sampler.py
index b43b19ef7..90e11f2ee 100644
--- a/cuqi/experimental/mcmc/_sampler.py
+++ b/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. """
diff --git a/cuqi/geometry/_geometry.py b/cuqi/geometry/_geometry.py
index 3267ed89b..9245ee223 100644
--- a/cuqi/geometry/_geometry.py
+++ b/cuqi/geometry/_geometry.py
@@ -225,7 +225,12 @@ def __eq__(self, obj):
         return self._all_values_equal(obj)
 
     def __repr__(self) -> str:
-        return "{}{}".format(self.__class__.__name__,self.par_shape)
+        if self.par_shape == self.fun_shape:
+            return "{}[{}]".format(self.__class__.__name__,
+                                   self.par_shape if len(self.par_shape) != 1 else self.par_shape[0])
+        return "{}[{}: {}]".format(self.__class__.__name__,
+                                   self.par_shape if len(self.par_shape) != 1 else self.par_shape[0],
+                                   self.fun_shape if len(self.fun_shape) != 1 else self.fun_shape[0])
 
     def _all_values_equal(self, obj):
         """Returns true of all values of the object and self are equal"""
diff --git a/cuqi/problem/_problem.py b/cuqi/problem/_problem.py
index 3323214c6..9d52bc4d2 100644
--- a/cuqi/problem/_problem.py
+++ b/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()


Test patch:
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index eb60fd557..83245706e 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -1510,3 +1510,101 @@ def test_RegularizedLinearRTO_ScipyLinearLSQ_option_invalid():
 
     with pytest.raises(ValueError, match="ScipyLinearLSQ"):
         sampler = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior, solver = "ScipyLinearLSQ")
+
+# ============ Start testing sampler callback ============
+# Samplers that should be tested for callback
+callback_testing_sampler_classes = [
+    cls
+    for _, cls in inspect.getmembers(cuqi.experimental.mcmc, inspect.isclass)
+    if cls not in [cuqi.experimental.mcmc.Sampler, cuqi.experimental.mcmc.ProposalBasedSampler]
+]
+
+# Instances of samplers that should be tested for callback
+callback_testing_sampler_instances = [
+    cuqi.experimental.mcmc.MH(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.PCN(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.CWMH(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.ULA(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.MALA(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.NUTS(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.LinearRTO(
+        target=cuqi.testproblem.Deconvolution1D(dim=10).posterior
+    ),
+    cuqi.experimental.mcmc.RegularizedLinearRTO(
+        target=create_regularized_target(dim=16)
+    ),
+    cuqi.experimental.mcmc.UGLA(target=create_lmrf_prior_target(dim=16)),
+    cuqi.experimental.mcmc.Direct(target=cuqi.distribution.Gaussian(np.zeros(10), 1)),
+    cuqi.experimental.mcmc.Conjugate(target=create_conjugate_target("Gaussian-Gamma")),
+    cuqi.experimental.mcmc.ConjugateApprox(
+        target=create_conjugate_target("LMRF-Gamma")
+    ),
+    cuqi.experimental.mcmc.MYULA(target=create_myula_target(dim=16)),
+    cuqi.experimental.mcmc.PnPULA(target=create_myula_target(dim=16)),
+    cuqi.experimental.mcmc.HybridGibbs(
+        target=HybridGibbs_target_1(),
+        sampling_strategy={
+            "x" : cuqi.experimental.mcmc.NUTS(max_depth=7),
+            "s" : cuqi.experimental.mcmc.Conjugate()
+        },
+    ),
+]
+
+@pytest.mark.parametrize("sampler", callback_testing_sampler_instances)
+def test_passing_callback_to_sampler(sampler):
+    """ Test that the callback function is called by the sampler and
+    the sampler passes the correct arguments to the callback function. """
+
+    Ns = 10
+    indices = [] # list to store the indices of the samples
+
+    # Define the callback function
+    def callback_function(callback_sampler, sample_index, num_samples):
+        assert isinstance(callback_sampler, sampler.__class__)
+        assert num_samples == Ns
+        indices.append(sample_index)
+
+    # Create a sampler with callback function
+    if not isinstance(sampler, cuqi.experimental.mcmc.HybridGibbs):
+        my_sampler = sampler.__class__(sampler.target, callback=callback_function)
+    else:
+        my_sampler = cuqi.experimental.mcmc.HybridGibbs(
+            sampler.target,
+            sampling_strategy={
+                "x": cuqi.experimental.mcmc.NUTS(max_depth=7),
+                "s": cuqi.experimental.mcmc.Conjugate(),
+            },
+            callback=callback_function,
+        )
+
+    # Sample Ns warmup samples
+    my_sampler.warmup(Ns)
+    assert np.allclose(indices, np.arange(Ns))
+
+    # reset the indices
+    indices = []
+
+    # Sample Ns samples
+    my_sampler.sample(Ns)
+    assert np.allclose(indices, np.arange(Ns))
+
+def test_all_samplers_that_should_be_tested_for_callback_are_in_the_tested_list():
+    """ Test that all samplers that should be tested for callback are in the callback_testing_sampler_instances. """
+    # The classes of the tested samplers:
+    tested_classes = [sampler.__class__ for sampler in callback_testing_sampler_instances]
+    for cls in callback_testing_sampler_classes:
+        assert cls in tested_classes, f"Sampler {cls} is not tested for callback."
+
+# ============= End testing sampler callback =============
\ No newline at end of file