Repository: CUQI-DTU/CUQIpy

Issue:
update readme
# DoD
- [x] Decide if we remove the link to the training repo since the notebooks there are retired?
- [x] Add links to CUQI book website and showcase repo?

Hints from issue comments:

Solution patch:
diff --git a/README.md b/README.md
index e609ba40a..53274b7f5 100644
--- a/README.md
+++ b/README.md
@@ -23,8 +23,9 @@ CUQIpy is part of the [CUQI project](https://www.compute.dtu.dk/english/cuqi) su
 
 - **Papers:** CUQIpy on ArXiv - [Part I](https://arxiv.org/abs/2305.16949) & [Part II](https://arxiv.org/abs/2305.16951)
 
-- **Training:** [Demo & training notebooks](https://github.com/CUQI-DTU/CUQIpy-demos)
+- **CUQI book:** [CUQI book website](https://cuqi-dtu.github.io/CUQI-Book/)
 
+- **User showcases:** [Showcase repository](https://github.com/CUQI-DTU/CUQIpy-User-Showcase/)
 
 ## 🚀 Quickstart
 Install CUQIpy using pip:
diff --git a/cuqi/distribution/_uniform.py b/cuqi/distribution/_uniform.py
index 550b296da..a61ae71e6 100644
--- a/cuqi/distribution/_uniform.py
+++ b/cuqi/distribution/_uniform.py
@@ -21,6 +21,9 @@ def __init__(self, low=None, high=None, is_symmetric=True, **kwargs):
         self.high = high      
 
     def logpdf(self, x):
+        """
+        Evaluate the logarithm of the PDF at the given values of x.
+        """
         # First check whether x is outside bounds.
         # It is outside if any coordinate is outside the interval.
         if np.any(x < self.low) or np.any(x > self.high):
@@ -38,6 +41,15 @@ def logpdf(self, x):
             return_val = np.log(1.0/v)
         return return_val
 
+    def gradient(self, x):
+        """
+        Computes the gradient of logpdf at the given values of x.
+        """
+        if np.any(x < self.low) or np.any(x > self.high):
+            return np.NaN*np.ones_like(x)
+        else:
+            return np.zeros_like(x)
+
     def _sample(self,N=1, rng=None):
 
         if rng is not None:
diff --git a/cuqi/experimental/mcmc/_cwmh.py b/cuqi/experimental/mcmc/_cwmh.py
index 3ef9f780c..b05e0c9ec 100644
--- a/cuqi/experimental/mcmc/_cwmh.py
+++ b/cuqi/experimental/mcmc/_cwmh.py
@@ -158,7 +158,9 @@ def step(self):
 
             # accept/reject
             u_theta = np.log(np.random.rand())
-            if (u_theta <= alpha): # accept
+            if (u_theta <= alpha) and \
+               (not np.isnan(target_eval_star)) and \
+               (not np.isinf(target_eval_star)):
                 x_t[j] = x_all_components[j]
                 target_eval_t = target_eval_star
                 acc[j] = 1
diff --git a/cuqi/experimental/mcmc/_hmc.py b/cuqi/experimental/mcmc/_hmc.py
index d2c9a4097..bee0343d0 100644
--- a/cuqi/experimental/mcmc/_hmc.py
+++ b/cuqi/experimental/mcmc/_hmc.py
@@ -88,14 +88,16 @@ class NUTS(Sampler):
     """
 
     _STATE_KEYS = Sampler._STATE_KEYS.union({'_epsilon', '_epsilon_bar',
-                                                '_H_bar', '_mu',
-                                                '_alpha', '_n_alpha'})
+                                                '_H_bar',
+                                                'current_target_logd',
+                                                'current_target_grad',
+                                                'max_depth'})
 
     _HISTORY_KEYS = Sampler._HISTORY_KEYS.union({'num_tree_node_list',
                                                     'epsilon_list',
                                                     'epsilon_bar_list'})
 
-    def __init__(self, target=None, initial_point=None, max_depth=15,
+    def __init__(self, target=None, initial_point=None, max_depth=None,
                  step_size=None, opt_acc_rate=0.6, **kwargs):
         super().__init__(target, initial_point=initial_point, **kwargs)
 
@@ -104,30 +106,20 @@ def __init__(self, target=None, initial_point=None, max_depth=15,
         self.step_size = step_size
         self.opt_acc_rate = opt_acc_rate
 
-        # Initialize epsilon and epsilon_bar
-        # epsilon is the step size used in the current iteration
-        # after warm up and one sampling step, epsilon is updated
-        # to epsilon_bar for the remaining sampling steps.
-        self._epsilon = None
-        self._epsilon_bar = None
-        self._H_bar = None
-
-        # Extra parameters for tuning
-        self._n_alpha = None
-        self._alpha = None
-
 
     def _initialize(self):
 
-        # Arrays to store acceptance rate
-        self._acc = [None] # Overwrites acc from Sampler. TODO. Check if this is necessary
-
-        self._alpha = 0 # check if meaningful value
-        self._n_alpha = 0 # check if meaningful value
+        self._current_alpha_ratio = np.nan # Current alpha ratio will be set to some
+                                           # value (other than np.nan) before 
+                                           # being used
 
         self.current_target_logd, self.current_target_grad = self._nuts_target(self.current_point)
 
-        # parameters dual averaging
+        # Parameters dual averaging
+        # Initialize epsilon and epsilon_bar
+        # epsilon is the step size used in the current iteration
+        # after warm up and one sampling step, epsilon is updated
+        # to epsilon_bar for the remaining sampling steps.
         if self.step_size is None:
             self._epsilon = self._FindGoodEpsilon()
         else:
@@ -155,6 +147,8 @@ def max_depth(self):
 
     @max_depth.setter
     def max_depth(self, value):
+        if value is None:
+            value = 15 # default value
         if not isinstance(value, int):
             raise TypeError('max_depth must be an integer.')
         if value < 0:
@@ -236,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)
@@ -256,13 +251,15 @@ def step(self):
 
             # Metropolis step
             alpha2 = min(1, (n_prime/n)) #min(0, np.log(n_p) - np.log(n))
-            if (s_prime == 1) and (np.random.rand() <= alpha2):
+            if (s_prime == 1) and \
+                (np.random.rand() <= alpha2) and \
+                (not np.isnan(logd_prime)) and \
+                (not np.isinf(logd_prime)):
                 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
@@ -270,8 +267,7 @@ def step(self):
             s = s_prime *\
                 int((dpoints @ r_minus.T) >= 0) * int((dpoints @ r_plus.T) >= 0)
             j += 1
-            self._alpha = alpha
-            self._n_alpha = n_alpha
+            self._current_alpha_ratio = alpha/n_alpha
 
         # update run diagnostic attributes
         self._update_run_diagnostic_attributes(
@@ -281,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
@@ -290,7 +288,7 @@ def tune(self, skip_len, update_count):
 
         eta1 = 1/(k + t_0)
         self._H_bar = (1-eta1)*self._H_bar +\
-            eta1*(self.opt_acc_rate - (self._alpha/self._n_alpha))
+            eta1*(self.opt_acc_rate - (self._current_alpha_ratio))
         self._epsilon = np.exp(self._mu - (np.sqrt(k)/gamma)*self._H_bar)
         eta = k**(-kappa)
         self._epsilon_bar =\
@@ -306,12 +304,12 @@ def _pre_warmup(self):
         #    self._epsilon
         #  Parameters that does not change during the run
         #    self._mu
-
+        self._ensure_initialized()
         if self._epsilon_bar == "unset": # Initial value of epsilon_bar for tuning
             self._epsilon_bar = 1
 
     def _pre_sample(self):
-
+        self._ensure_initialized()
         if self._epsilon_bar == "unset": # Initial value of epsilon_bar for sampling
             self._epsilon_bar = self._epsilon
             
diff --git a/cuqi/experimental/mcmc/_langevin_algorithm.py b/cuqi/experimental/mcmc/_langevin_algorithm.py
index 2331507a3..d4af60114 100644
--- a/cuqi/experimental/mcmc/_langevin_algorithm.py
+++ b/cuqi/experimental/mcmc/_langevin_algorithm.py
@@ -102,10 +102,14 @@ def _accept_or_reject(self, x_star, target_eval_star, target_grad_star):
         scalar
             1 (accepted)
         """
-        self.current_point = x_star
-        self.current_target_logd = target_eval_star
-        self.current_target_grad = target_grad_star
-        acc = 1
+        acc = 0
+        if (not np.isnan(target_eval_star)) and \
+           (not np.isinf(target_eval_star)):
+            self.current_point = x_star
+            self.current_target_logd = target_eval_star
+            self.current_target_grad = target_grad_star
+            acc = 1
+            
         return acc
 
     def step(self):
@@ -211,7 +215,9 @@ def _accept_or_reject(self, x_star, target_eval_star, target_grad_star):
         # accept/reject with Metropolis
         acc = 0
         log_u = np.log(np.random.rand())
-        if (log_u <= log_alpha) and (np.isnan(target_eval_star) == False):
+        if (log_u <= log_alpha) and \
+           (not np.isnan(target_eval_star)) and \
+           (not np.isinf(target_eval_star)):
             self.current_point = x_star
             self.current_target_logd = target_eval_star
             self.current_target_grad = target_grad_star
diff --git a/cuqi/experimental/mcmc/_mh.py b/cuqi/experimental/mcmc/_mh.py
index 435679023..1322d9b7d 100644
--- a/cuqi/experimental/mcmc/_mh.py
+++ b/cuqi/experimental/mcmc/_mh.py
@@ -57,7 +57,9 @@ def step(self):
         # accept/reject
         u_theta = np.log(np.random.rand())
         acc = 0
-        if (u_theta <= alpha):
+        if (u_theta <= alpha) and \
+           (not np.isnan(target_eval_star)) and \
+           (not np.isinf(target_eval_star)):
             self.current_point = x_star
             self.current_target_logd = target_eval_star
             acc = 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)
 
diff --git a/cuqi/problem/_problem.py b/cuqi/problem/_problem.py
index 169a5e8d4..824a2d878 100644
--- a/cuqi/problem/_problem.py
+++ b/cuqi/problem/_problem.py
@@ -614,7 +614,7 @@ def _samplepCN(self, Ns, Nb, callback=None, experimental=False):
 
             scale = 0.02
 
-            sampler = cuqi.experimental.mcmc.pCN(self.posterior, scale, callback=callback)
+            sampler = cuqi.experimental.mcmc.PCN(self.posterior, scale, callback=callback)
 
             ti = time.time()
 


Test patch:
diff --git a/tests/test_distributions_shape.py b/tests/test_distributions_shape.py
index 718e398ed..387f12fc7 100644
--- a/tests/test_distributions_shape.py
+++ b/tests/test_distributions_shape.py
@@ -95,7 +95,8 @@ def test_multivariate_scalar_vars_gradient(dist):
     try:
         assert np.allclose(
             dist_from_vec.gradient(val),
-            dist_from_dim.gradient(val)
+            dist_from_dim.gradient(val),
+            equal_nan=True
         )
     except NotImplementedError:
         pass  # Pass the test if NotImplementedError is raised
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index 72a6c5547..12147b13a 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -348,7 +348,8 @@ def create_conjugate_target(type:str):
     cuqi.experimental.mcmc.UGLA(create_lmrf_prior_target(dim=16)),
     cuqi.experimental.mcmc.Direct(cuqi.distribution.Gaussian(np.zeros(10), 1)),
     cuqi.experimental.mcmc.Conjugate(create_conjugate_target("Gaussian-Gamma")),
-    cuqi.experimental.mcmc.ConjugateApprox(create_conjugate_target("LMRF-Gamma"))
+    cuqi.experimental.mcmc.ConjugateApprox(create_conjugate_target("LMRF-Gamma")),
+    cuqi.experimental.mcmc.NUTS(cuqi.testproblem.Deconvolution1D(dim=10).posterior, max_depth=4)
 ]
     
 # List of samplers from cuqi.experimental.mcmc that should be skipped for checkpoint testing
@@ -359,7 +360,6 @@ def create_conjugate_target(type:str):
     cuqi.experimental.mcmc.PCN,
     cuqi.experimental.mcmc.CWMH,
     cuqi.experimental.mcmc.RegularizedLinearRTO, # Due to the _choose_stepsize method
-    cuqi.experimental.mcmc.NUTS,
     cuqi.experimental.mcmc.HybridGibbs
 ]
 
@@ -431,6 +431,7 @@ def test_checkpointing(sampler: cuqi.experimental.mcmc.Sampler):
     cuqi.experimental.mcmc.LinearRTO(cuqi.testproblem.Deconvolution1D(dim=10).posterior),
     cuqi.experimental.mcmc.RegularizedLinearRTO(create_regularized_target(dim=16)),
     cuqi.experimental.mcmc.UGLA(create_lmrf_prior_target(dim=32)),
+    cuqi.experimental.mcmc.NUTS(cuqi.testproblem.Deconvolution1D(dim=10).posterior, max_depth=4)
 ]
 
 
@@ -477,6 +478,7 @@ def test_history_keys(sampler: cuqi.experimental.mcmc.Sampler):
 state_exception_keys = {
     cuqi.experimental.mcmc.ULA: 'scale',
     cuqi.experimental.mcmc.MALA: 'scale',
+    cuqi.experimental.mcmc.NUTS: 'max_depth'
 }
 
 @pytest.mark.parametrize("sampler", state_history_targets)
@@ -593,7 +595,8 @@ def compare_attributes(attr1, attr2, key=''):
             for k in attr1.keys():
                 compare_attributes(attr1[k], attr2[k], f"{key}['{k}']")
         elif isinstance(attr1, Number):
-            assert np.allclose(attr1, attr2), f"Number: Attribute '{key}' differs after initialization"
+            assert (np.isnan(attr1) and np.isnan(attr2)) \
+                or np.allclose(attr1, attr2), f"Number: Attribute '{key}' differs after initialization"
         elif isinstance(attr1, str):
             assert attr1 == attr2, f"String: Attribute '{key}' differs after initialization"
         else:
@@ -903,4 +906,88 @@ def test_HybridGibbs_updates_state_only_after_accepting_sample():
         # In case sampler has accepted samples, the state should have been updated
         else:
             assert sampler_states[key] != new_state, f"Sampler {key} state was erroneously not updated in Gibbs scheme, even when new samples were accepted. State: \n {new_state}"
-            
+
+sampler_instances_for_bounded_distribution = [
+    cuqi.experimental.mcmc.MH(
+        target=cuqi.distribution.Beta(0.5, 0.5),
+        initial_point=np.array([0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.ULA(
+        cuqi.distribution.Beta(0.5, 0.5),
+        initial_point=np.array([0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.MALA(
+        cuqi.distribution.Beta(0.5, 0.5),
+        initial_point=np.array([0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.NUTS(
+        cuqi.distribution.Beta(0.5, 0.5), initial_point=np.array([0.1])
+    ),
+    cuqi.experimental.mcmc.MH(
+        target=cuqi.distribution.Beta(np.array([0.5, 0.5]), np.array([0.5, 0.5])),
+        initial_point=np.array([0.1, 0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.CWMH(
+        target=cuqi.distribution.Beta(
+            np.array([0.5, 0.5]), np.array([0.5, 0.5])
+        ),
+        initial_point=np.array([0.1, 0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.ULA(
+        cuqi.distribution.Beta(np.array([0.5, 0.5]), np.array([0.5, 0.5])),
+        initial_point=np.array([0.1, 0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.MALA(
+        cuqi.distribution.Beta(np.array([0.5, 0.5]), np.array([0.5, 0.5])),
+        initial_point=np.array([0.1, 0.1]),
+        scale=0.1,
+    ),
+    cuqi.experimental.mcmc.NUTS(
+        cuqi.distribution.Beta(np.array([0.5, 0.5]), np.array([0.5, 0.5])),
+        initial_point=np.array([0.1, 0.1]),
+    ),
+]
+
+@pytest.mark.parametrize("sampler", sampler_instances_for_bounded_distribution)
+def test_if_invalid_sample_accepted(sampler: cuqi.experimental.mcmc.Sampler):
+    sampler.sample(50)
+    samples = sampler.get_samples().samples
+    tol = 1e-8
+    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