Merge pull request #533 from CUQI-DTU/add_info_to_progress_bar

Display acc rate on progress bar

Author: chaozg

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

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)
 

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)