Repository: CUQI-DTU/CUQIpy

Issue:
Samplers accept invalid proposals
**Description**

With a bounded prior, invalid proposals are accepted after a MH step. This happens to all samplers involving a MH step.

As seen in the following plot, an invalid proposal, negative here, is accepted after a MH step, which is not supposed to happen, and then the sampler gets stuck there.

**Example to reproduce**

```python
import cuqi
beta = cuqi.distribution.Beta(0.5, 0.5)
simpler = cuqi.sampler.MALA(beta, x0=0.1, scale=0.1)
simples = simpler.sample(1000, 0)
plt.plot(simples.samples.T[:10])
```
![bug](https://github.com/user-attachments/assets/06045a3c-d212-4c49-ab36-3236f74dc2e2)

# DoD
- [x] Identity the reason behind:
  - [x]  the validity of proposals is not checked at a MH step
- [x] solution agreed by team and issue merged
  - [x] validate proposal with `isnan` and `isinf` at MH steps in relevant samplers, i.e., MH, CWMH, MALA, NUTS
- [x] proper test functions are created

Discussion:
@chaozg: @nabriis , I just want to draw your attention to this issue since you were not at the standup when I briefly reported it.@chaozg: Hi @jakobsj  could you review the DoD here?@jakobsj: DoD approved!

Solution patch:
diff --git a/cuqi/experimental/mcmc/__init__.py b/cuqi/experimental/mcmc/__init__.py
index 330b3c2d3..b2f8a5b0c 100644
--- a/cuqi/experimental/mcmc/__init__.py
+++ b/cuqi/experimental/mcmc/__init__.py
@@ -1,4 +1,112 @@
-""" Re-implementation of sampler module in a more object oriented way. """
+"""
+Re-implementation of sampler module in a more object-oriented way.
+
+Main changes for users
+----------------------
+
+1. Sampling API
+   ^^^^^^^^^^^^
+
+   Previously one would call the `.sample` or `sample_adapt` methods of a sampler instance at :py:mod:`cuqi.sampler` to sample from a target distribution and store the samples as the output as follows:
+
+   .. code-block:: python
+
+      from cuqi.sampler import MH
+      from cuqi.distribution import DistributionGallery
+
+      # Target distribution
+      target = DistributionGallery("donut")
+
+      # Set up sampler
+      sampler = MH(target)
+
+      # Sample from the target distribution (Alternatively calling sample with explicit scale parameter set in sampler)
+      samples = sampler.sample_adapt(Ns=100, Nb=100) # Burn-in (Nb) removed by default
+
+   This has now changed to to a more object-oriented API which provides more flexibility and control over the sampling process.
+
+   For example one can now more explicitly control when the sampler is tuned (warmup) and when it is sampling with fixed parameters.
+
+   .. code-block:: python
+
+      from cuqi.experimental.mcmc import MH
+      from cuqi.distribution import DistributionGallery
+
+      # Target distribution
+      target = DistributionGallery("donut")
+
+      # Set up sampler
+      sampler = MH(target)
+
+      # Sample from the target distribution
+      sampler.warmup(Nb=100)  # Explicit warmup (tuning) of sampler
+      sampler.sample(Ns=100)  # Sampling with fixed parameters
+      samples = sampler.get_samples().burnthin(Nb=100) # Getting samples and removing burn-in from warmup
+
+   Importantly, the removal of burn-in from e.g. warmup is now a separate step that is done after the sampling process is complete.
+
+2. Sampling API for BayesianProblem
+   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+   :py:class:`cuqi.problem.BayesianProblem` continues to have the same API for `sample_posterior` and the `UQ` method.
+
+   There is now a flag `experimental` that can be set to `True` to use the new MCMC samplers.
+   
+   By default, the flag is set to `False` and the old samplers are used.
+
+   For this more high-level interface, burn-in is automatically removed from the samples as was the case before.
+
+
+3. More options for Gibbs sampling
+   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+   There are now more options for Gibbs sampling. Previously it was only possible to sample with Gibbs for samplers  :py:class:`cuqi.sampler.LinearRTO`, :py:class:`cuqi.sampler.RegularizedLinearRTO`, :py:class:`cuqi.sampler.Conjugate`, and :py:class:`cuqi.sampler.ConjugateApprox`.
+
+   Now, it is possible to define a Gibbs sampling scheme using any sampler from the :py:mod:`cuqi.experimental.mcmc` module.
+
+   **Example using a NUTS-within-Gibbs scheme for a 1D deconvolution problem:**
+
+   .. code-block:: python
+
+      import cuqi
+      import numpy as np
+      from cuqi.distribution import Gamma, Gaussian, GMRF, JointDistribution
+      from cuqi.experimental.mcmc import NUTS, HybridGibbs, Conjugate
+      from cuqi.testproblem import Deconvolution1D
+
+      # Forward problem
+      A, y_data, info = Deconvolution1D(dim=128, phantom='sinc', noise_std=0.001).get_components()
+
+      # Bayesian Inverse Problem
+      s = Gamma(1, 1e-4)
+      x = GMRF(np.zeros(A.domain_dim), 50)
+      y = Gaussian(A @ x, lambda s: 1 / s)
+
+      # Posterior
+      target = JointDistribution(y, x, s)(y=y_data)
+
+      # Gibbs sampling strategy. Note we can define initial_points and various parameters for each sampler
+      sampling_strategy = {
+          "x": NUTS(max_depth=10, initial_point=np.zeros(A.domain_dim)),
+          "s": Conjugate()
+      }
+
+      # Here we do 10 internal steps with NUTS for each Gibbs step
+      num_sampling_steps = {
+          "x": 10,
+          "s": 1
+      }
+
+      sampler = HybridGibbs(target, sampling_strategy, num_sampling_steps)
+
+      sampler.warmup(50)
+      sampler.sample(200)
+      samples = sampler.get_samples().burnthin(Nb=50)
+
+      samples["x"].plot_ci(exact=info.exactSolution)
+"""
+
+
 
 from ._sampler import Sampler, ProposalBasedSampler
 from ._langevin_algorithm import ULA, MALA
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/_gibbs.py b/cuqi/experimental/mcmc/_gibbs.py
index d35a805d9..aa94eb873 100644
--- a/cuqi/experimental/mcmc/_gibbs.py
+++ b/cuqi/experimental/mcmc/_gibbs.py
@@ -1,6 +1,6 @@
 from cuqi.distribution import JointDistribution
 from cuqi.experimental.mcmc import Sampler
-from cuqi.samples import Samples
+from cuqi.samples import Samples, JointSamples
 from cuqi.experimental.mcmc import NUTS
 from typing import Dict
 import numpy as np
@@ -66,7 +66,7 @@ class HybridGibbs:
         import numpy as np
 
         # Model and data
-        A, y_obs, probinfo = cuqi.testproblem.Deconvolution1D(phantom='square').get_components()
+        A, y_obs, probinfo = cuqi.testproblem.Deconvolution1D(phantom='sinc').get_components()
         n = A.domain_dim
 
         # Define distributions
@@ -90,7 +90,11 @@ class HybridGibbs:
         sampler = cuqi.experimental.mcmc.HybridGibbs(posterior, sampling_strategy)
 
         # Run sampler
-        samples = sampler.sample(Ns=1000, Nb=200)
+        sampler.warmup(200)
+        sampler.sample(1000)
+
+        # Get samples removing burn-in
+        samples = sampler.get_samples().burnthin(200)
 
         # Plot results
         samples['x'].plot_ci(exact=probinfo.exactSolution)
@@ -161,7 +165,7 @@ def sample(self, Ns) -> 'HybridGibbs':
 
         """
 
-        for _ in tqdm(range(Ns)):
+        for _ in tqdm(range(Ns), "Sample: "):
 
             self.step()
 
@@ -184,7 +188,7 @@ def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
 
         tune_interval = max(int(tune_freq * Nb), 1)
 
-        for idx in tqdm(range(Nb)):
+        for idx in tqdm(range(Nb), "Warmup: "):
 
             self.step()
 
@@ -197,7 +201,7 @@ def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
         return self
 
     def get_samples(self) -> Dict[str, Samples]:
-        samples_object = {}
+        samples_object = JointSamples()
         for par_name in self.par_names:
             samples_array = np.array(self.samples[par_name]).T
             samples_object[par_name] = Samples(samples_array, self.target.get_density(par_name).geometry)
diff --git a/cuqi/experimental/mcmc/_hmc.py b/cuqi/experimental/mcmc/_hmc.py
index d2c9a4097..a5a9f14cc 100644
--- a/cuqi/experimental/mcmc/_hmc.py
+++ b/cuqi/experimental/mcmc/_hmc.py
@@ -256,7 +256,10 @@ 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)
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 3902d4c3f..21a0fbb38 100644
--- a/cuqi/experimental/mcmc/_sampler.py
+++ b/cuqi/experimental/mcmc/_sampler.py
@@ -220,7 +220,7 @@ 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) ):
+        for _ in tqdm( range(Ns), "Sample: "):
             
             # Perform one step of the sampler
             acc = self.step()
@@ -260,7 +260,7 @@ 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)):
+        for idx in tqdm(range(Nb), "Warmup: "):
 
             # Perform one step of the sampler
             acc = self.step()
diff --git a/cuqi/samples/__init__.py b/cuqi/samples/__init__.py
index 92873c0f3..a3d84d25d 100644
--- a/cuqi/samples/__init__.py
+++ b/cuqi/samples/__init__.py
@@ -1 +1 @@
-from ._samples import Samples
+from ._samples import Samples, JointSamples
diff --git a/cuqi/samples/_samples.py b/cuqi/samples/_samples.py
index fcdb0a12b..f7847c220 100644
--- a/cuqi/samples/_samples.py
+++ b/cuqi/samples/_samples.py
@@ -867,3 +867,23 @@ def __repr__(self) -> str:
                "Geometry:\n {}\n\n".format(self.geometry) + \
                "Shape:\n {}\n\n".format(self.shape) + \
                "Samples:\n {}\n\n".format(self.samples)
+
+class JointSamples(dict):
+    """ An object used to store samples from :class:`cuqi.distribution.JointDistribution`. 
+
+    This object is a simple overload of the dictionary class to allow easy access to certain methods 
+    of Samples objects without having to iterate over each key in the dictionary. 
+    
+    """
+
+    def burnthin(self, Nb, Nt=1):
+        """ Remove burn-in and thin samples for all samples in the dictionary. Returns a copy of the samples stored in the dictionary. """
+        return JointSamples({key: samples.burnthin(Nb, Nt) for key, samples in self.items()})
+
+    def __repr__(self) -> str: 
+        return "CUQIpy JointSamples Dict:\n" + \
+               "-------------------------\n\n" + \
+               "Keys:\n {}\n\n".format(list(self.keys())) + \
+               "Ns (number of samples):\n {}\n\n".format({key: samples.Ns for key, samples in self.items()}) + \
+               "Geometry:\n {}\n\n".format({key: samples.geometry for key, samples in self.items()}) + \
+               "Shape:\n {}\n\n".format({key: samples.shape for key, samples in self.items()})


Test patch:
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index 72a6c5547..ab3402b95 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -903,4 +903,59 @@ 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__}."