Repository: CUQI-DTU/CUQIpy

Issue:
Acc rate not stored in HybridGibbs sampler
**Description**

Currently the acceptance rate is not maintained in the HybridGibbs sampler. This means that tuning is not happening using the full acc rate of the chain.

## DoD:
- [x] HybridGibbs ensures samplers store acc rate while sampling
- [x] HybridGibbs maintains state and history of samplers
- [x] unit tests are implemented to check the above mentioned behavior

Discussion:
@jakobsj: @nabriis can you add DoD here please?@nabriis: > @nabriis can you add DoD here please?

Done

Solution patch:
diff --git a/cuqi/experimental/mcmc/_gibbs.py b/cuqi/experimental/mcmc/_gibbs.py
index 8e627e30d..d35a805d9 100644
--- a/cuqi/experimental/mcmc/_gibbs.py
+++ b/cuqi/experimental/mcmc/_gibbs.py
@@ -1,6 +1,7 @@
 from cuqi.distribution import JointDistribution
 from cuqi.experimental.mcmc import Sampler
 from cuqi.samples import Samples
+from cuqi.experimental.mcmc import NUTS
 from typing import Dict
 import numpy as np
 import warnings
@@ -151,18 +152,50 @@ def validate_targets(self):
             sampler.validate_target()
 
     def sample(self, Ns) -> 'HybridGibbs':
-        """ Sample from the joint distribution using Gibbs sampling """
+        """ Sample from the joint distribution using Gibbs sampling
+
+        Parameters
+        ----------
+        Ns : int
+            The number of samples to draw.
+
+        """
+
         for _ in tqdm(range(Ns)):
+
             self.step()
+
             self._store_samples()
 
-    def warmup(self, Nb) -> 'HybridGibbs':
-        """ Warmup (tune) the Gibbs sampler """
+        return self
+
+    def warmup(self, Nb, tune_freq=0.1) -> 'HybridGibbs':
+        """ Warmup (tune) the samplers in the Gibbs sampling scheme
+
+        Parameters
+        ----------
+        Nb : int
+            The number of samples to draw during warmup.
+
+        tune_freq : float, optional
+            Frequency of tuning the samplers. Tuning is performed every tune_freq*Nb steps.
+
+        """
+
+        tune_interval = max(int(tune_freq * Nb), 1)
+
         for idx in tqdm(range(Nb)):
+
             self.step()
-            self.tune(idx)
+
+            # Tune the sampler at tuning intervals (matching behavior of Sampler class)
+            if (idx + 1) % tune_interval == 0:
+                self.tune(tune_interval, idx // tune_interval) 
+                
             self._store_samples()
 
+        return self
+
     def get_samples(self) -> Dict[str, Samples]:
         samples_object = {}
         for par_name in self.par_names:
@@ -182,38 +215,65 @@ def step(self):
             # Get sampler
             sampler = self.samplers[par_name]
 
-            # Set initial parameters using current point and scale (subset of state)
-            # This does not store the full state from e.g. NUTS sampler
-            # But works on samplers like MH, PCN, ULA, MALA, LinearRTO, UGLA, CWMH
-            # that only use initial_point and initial_scale
-            sampler.initial_point = self.current_samples[par_name]
-            if hasattr(sampler, 'initial_scale'): sampler.initial_scale = sampler.scale
+            # Instead of simply changing the target of the sampler, we reinitialize it.
+            # This is to ensure that all internal variables are set to match the new target.
+            # To return the sampler to the old state and history, we first extract the state and history
+            # before reinitializing the sampler and then set the state and history back to the sampler
+
+            # Extract state and history from sampler
+            if isinstance(sampler, NUTS): # Special case for NUTS as it is not playing nice with get_state and get_history
+                sampler.initial_point = sampler.current_point
+            else:
+                sampler_state = sampler.get_state()
+                sampler_history = sampler.get_history()
 
             # Reinitialize sampler
-            # This makes the sampler lose all of its state.
-            # This is only OK because we set the initial values above from the previous state
             sampler.reinitialize()
 
+            # Set state and history back to sampler
+            if not isinstance(sampler, NUTS): # Again, special case for NUTS.
+                sampler.set_state(sampler_state)
+                sampler.set_history(sampler_history)
+
             # Run pre_warmup and pre_sample methods for sampler
             # TODO. Some samplers (NUTS) seem to require to run _pre_warmup before _pre_sample
             self._pre_warmup_and_pre_sample_sampler(sampler)
 
-            # Take MCMC steps
+            # Allow for multiple sampling steps in each Gibbs step
             for _ in range(self.num_sampling_steps[par_name]):
-                sampler.step()
+                # Sampling step
+                acc = sampler.step()
+
+                # Store acceptance rate in sampler (matching behavior of Sampler class Sample method)
+                sampler._acc.append(acc)
 
             # Extract samples (Ensure even 1-dimensional samples are 1D arrays)
-            self.current_samples[par_name] = sampler.current_point.reshape(-1)
+            if isinstance(sampler.current_point, np.ndarray):
+                self.current_samples[par_name] = sampler.current_point.reshape(-1)
+            else:
+                self.current_samples[par_name] = sampler.current_point
+
+    def tune(self, skip_len, update_count):
+        """ Run a single tuning step on each of the samplers in the Gibbs sampling scheme
+
+        Parameters
+        ----------
+        skip_len : int
+            Defines the number of steps in between tuning (i.e. the tuning interval).
+
+        update_count : int
+            The number of times tuning has been performed. Can be used for internal bookkeeping.
 
-    def tune(self, idx):
-        """ Tune each of the samplers """
+        """
         for par_name in self.par_names:
-            self.samplers[par_name].tune(skip_len=1, update_count=idx)
+            self.samplers[par_name].tune(skip_len=skip_len, update_count=update_count)
 
     # ------------ Private methods ------------
     def _initialize_samplers(self):
         """ Initialize samplers """
         for sampler in self.samplers.values():
+            if isinstance(sampler, NUTS):
+                print(f'Warning: NUTS sampler is not fully stateful in HybridGibbs. Sampler will be reinitialized in each Gibbs step.')
             sampler.initialize()
 
     def _initialize_num_sampling_steps(self):
diff --git a/requirements-dev.txt b/requirements-dev.txt
index b10b7c98a..10cb5012c 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -7,4 +7,3 @@ sphinx-panels
 sphinx-copybutton
 sphinx-gallery
 versioneer==0.26
-tqdm # For experimental mcmc module
diff --git a/requirements.txt b/requirements.txt
index 6c1415bf6..3ae2538d6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,4 @@ matplotlib
 numpy>=1.17.0
 scipy<1.13
 arviz
+tqdm


Test patch:
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index 575c23976..72a6c5547 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -823,3 +823,84 @@ def test_HybridGibbs_initial_point_setting():
     assert sampler.current_samples["d"] == 3
     assert sampler.current_samples["s"] == 1 # Default initial point of MH is 1.
     assert np.allclose(sampler.current_samples["x"], 0.5*np.ones(10))
+
+def test_HybridGibbs_stores_acc_rate():
+    """ Test that the HybridGibbs sampler stores the acceptance rate of the samplers correctly. Also ensures that history is maintained during warmup and sampling. """
+
+    # Number of warmup and sampling steps
+    Ns = 25
+    Nb = 25
+    
+    # Bayesian Problem
+    s = cuqi.distribution.Gaussian(1, 1)
+    d = cuqi.distribution.Uniform(1, 100)
+    x = cuqi.distribution.Gaussian(lambda s: s, lambda d: 1/d, geometry=1)
+
+    # Joint distribution
+    joint = cuqi.distribution.JointDistribution(x, d, s)
+
+    # Sampling strategy
+    sampling_strategy = {
+        "d" : cuqi.experimental.mcmc.MH(initial_point=3),
+        "s" : cuqi.experimental.mcmc.PCN(initial_point=3),
+        "x" : cuqi.experimental.mcmc.MALA(initial_point=0)
+    }
+
+    # Hybrid Gibbs sampler
+    sampler = cuqi.experimental.mcmc.HybridGibbs(joint, sampling_strategy=sampling_strategy)
+
+    # Run the sampler for a few steps
+    sampler.warmup(Nb).sample(Ns)
+
+    # Check that the samplers have acc rate correctly updated and maintained
+    assert len(sampler.samplers["d"].get_history()["history"]["_acc"]) == Nb + Ns + 1
+    assert len(sampler.samplers["s"].get_history()["history"]["_acc"]) == Nb + Ns + 1
+    assert len(sampler.samplers["x"].get_history()["history"]["_acc"]) == Nb + Ns + 1
+
+def test_HybridGibbs_updates_state_only_after_accepting_sample():
+    """ Test that the HybridGibbs sampler updates the state only after accepting a sample. """
+
+    # Number of warmup and sampling steps
+    Ns = 25
+    Nb = 25
+    
+    # Bayesian Problem
+    s = cuqi.distribution.Gaussian(1, 1)
+    d = cuqi.distribution.Uniform(1, 100)
+    x = cuqi.distribution.Gaussian(lambda s: s, lambda d: 1/d, geometry=1)
+
+    # Joint distribution
+    joint = cuqi.distribution.JointDistribution(x, d, s)
+
+    # Sampling strategy
+    sampling_strategy = {
+        "d" : cuqi.experimental.mcmc.MH(initial_point=3),
+        "s" : cuqi.experimental.mcmc.PCN(initial_point=3),
+        "x" : cuqi.experimental.mcmc.MALA(initial_point=0, scale=1e-1) # Relatively high scale may lead to no accepted samples
+    }
+
+    # Hybrid Gibbs sampler
+    sampler = cuqi.experimental.mcmc.HybridGibbs(joint, sampling_strategy=sampling_strategy)
+
+    # Run the sampler for a few steps
+    sampler.warmup(Nb).sample(Ns)
+
+    # Store states of samplers
+    sampler_states = {key: sampler.samplers[key].get_state() for key in sampler.samplers.keys()}
+
+    # Run warmup and sampling again
+    sampler.warmup(Nb).sample(Ns)
+
+    # Verify that the state is only updated if a sample is accepted.
+    for key in sampler.samplers.keys():
+
+        # Get new state for specific sampler
+        new_state = sampler.samplers[key].get_state()
+
+        # In case sampler has no accepted samples, the state should not have been updated
+        if (np.sum(sampler.samplers[key].get_history()["history"]["_acc"][Ns+Nb+1:]) == 0):
+            assert sampler_states[key] == new_state, f"Sampler {key} state was erroneously updated in Gibbs scheme, but no accepted samples. State: \n {new_state}"
+        # 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}"
+