Repository: CUQI-DTU/CUQIpy

Issue:
Use Progress bar to print sampling progress and possibly other useful sampling information
Use Progress bar to print sampling progress and possibly provide other useful sampling information (e.g. display dynamic runtime convergence diagnostics)

Discussion:

Solution patch:
diff --git a/cuqi/experimental/mcmc/_hmc.py b/cuqi/experimental/mcmc/_hmc.py
index 8b4c89ea7..bee0343d0 100644
--- a/cuqi/experimental/mcmc/_hmc.py
+++ b/cuqi/experimental/mcmc/_hmc.py
@@ -109,10 +109,7 @@ def __init__(self, target=None, initial_point=None, max_depth=None,
 
     def _initialize(self):
 
-        # Arrays to store acceptance rate
-        self._acc = [None] # Overwrites acc from Sampler. TODO. Check if this is necessary
-
-        self._current_alpha_ratio = np.nan # Current alpha ratio is set to some
+        self._current_alpha_ratio = np.nan # Current alpha ratio will be set to some
                                            # value (other than np.nan) before 
                                            # being used
 
@@ -233,6 +230,7 @@ def step(self):
         r_minus, r_plus = np.copy(r_k), np.copy(r_k)
 
         # run NUTS
+        acc = 0
         while (s == 1) and (j <= self.max_depth):
             # sample a direction
             v = int(2*(np.random.rand() < 0.5)-1)
@@ -260,9 +258,8 @@ def step(self):
                 self.current_point = point_prime
                 self.current_target_logd = logd_prime
                 self.current_target_grad = np.copy(grad_prime)
-                self._acc.append(1)
-            else:
-                self._acc.append(0)
+                acc = 1
+
 
             # update number of particles, tree level, and stopping criterion
             n += n_prime
@@ -280,6 +277,8 @@ def step(self):
         if np.isnan(self.current_target_logd):
             raise NameError('NaN potential func')
 
+        return acc
+
     def tune(self, skip_len, update_count):
         """ adapt epsilon during burn-in using dual averaging"""
         k = update_count+1
diff --git a/cuqi/experimental/mcmc/_sampler.py b/cuqi/experimental/mcmc/_sampler.py
index 21a0fbb38..3a5111aed 100644
--- a/cuqi/experimental/mcmc/_sampler.py
+++ b/cuqi/experimental/mcmc/_sampler.py
@@ -220,7 +220,8 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
         if hasattr(self, "_pre_sample"): self._pre_sample()
 
         # Draw samples
-        for _ in tqdm( range(Ns), "Sample: "):
+        pbar = tqdm(range(Ns), "Sample: ")
+        for idx in pbar:
             
             # Perform one step of the sampler
             acc = self.step()
@@ -229,6 +230,9 @@ def sample(self, Ns, batch_size=0, sample_path='./CUQI_samples/') -> 'Sampler':
             self._acc.append(acc)
             self._samples.append(self.current_point)
 
+            # display acc rate at progress bar
+            pbar.set_postfix_str(f"acc rate: {np.mean(self._acc[-1-idx:]):.2%}")
+
             # Add sample to batch
             if batch_size > 0:
                 batch_handler.add_sample(self.current_point)
@@ -260,7 +264,8 @@ def warmup(self, Nb, tune_freq=0.1) -> 'Sampler':
         if hasattr(self, "_pre_warmup"): self._pre_warmup()
 
         # Draw warmup samples with tuning
-        for idx in tqdm(range(Nb), "Warmup: "):
+        pbar = tqdm(range(Nb), "Warmup: ")
+        for idx in pbar:
 
             # Perform one step of the sampler
             acc = self.step()
@@ -273,6 +278,9 @@ def warmup(self, Nb, tune_freq=0.1) -> 'Sampler':
             self._acc.append(acc)
             self._samples.append(self.current_point)
 
+            # display acc rate at progress bar
+            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)
 


Test patch:
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index 8c5564fe4..12147b13a 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -962,3 +962,32 @@ def test_if_invalid_sample_accepted(sampler: cuqi.experimental.mcmc.Sampler):
     assert (
         samples.min() > 0.0 - tol and samples.max() < 1.0 + tol
     ), f"Invalid samples accepted for sampler {sampler.__class__.__name__}."
+
+
+# Test NUTS acceptance rate
+@pytest.mark.parametrize(
+    "sampler",
+    [
+        cuqi.experimental.mcmc.NUTS(cuqi.distribution.Gaussian(0, 1)),
+        cuqi.experimental.mcmc.NUTS(cuqi.distribution.DistributionGallery('donut'))
+    ],
+)
+def test_nuts_acceptance_rate(sampler: cuqi.experimental.mcmc.Sampler):
+    """ Test that the NUTS sampler correctly updates the acceptance rate. """
+    # Fix random seed for reproducibility, but the test should be robust to seed
+    np.random.seed(0)
+
+    # Sample:
+    sampler.warmup(100).sample(100)
+
+    # Compute number of times samples were updated:
+    samples = sampler.get_samples().samples
+    counter = 0
+    for i in range(1, samples.shape[1]):
+        if np.any(samples[:, i] != samples[:, i - 1]):
+            counter += 1
+
+    # Compute the number of accepted samples according to the sampler
+    acc_rate_sum = sum(sampler._acc[2:])
+
+    assert np.isclose(counter, acc_rate_sum), "NUTS sampler does not update acceptance rate correctly: "+str(counter)+" != "+str(acc_rate_sum)
\ No newline at end of file