Repository: CUQI-DTU/CUQIpy

Issue:
Documentation updates for training
realted to #463
Description:
-------------
Go over CUQIpy documentation and make light updates to help improve the documentation (e.g. fixes of typos, equations not rendered correctly) and identify larger documentation update tasks that are necessary

Definition of done:
-------------------
- [x] Light update (as described above) is done, reviewed, and merged
- [ ] Bigger and necessary documentation update tasks are identified and corresponding issues are created, assigned and linked from this issue. (No major issues found at this stage.)
- [x] The possibility of fixing errors related to documentation that appear when CUQIpy is installed on a Windows machine is investigated
- [x] #488
- [x] Add experimental mcmc to main API page


Notes on Windows installation:
------------------------------
- @chaozg went through the CUQIpy installation process on Windows with Tania by following the instructions listed at the newest CUQIpy [documentation](https://cuqi-dtu.github.io/CUQIpy/user/getting_started.html), and everything worked fine. Just some notes:
  - After installation, it might take some time, like half a minute, for vscode to update its kernel options, so the just installed environment appears as an available python kernel
  - For working with the terminal interface in vscode, we should choose "terminal prompt" as the terminal type, as other terminals don't know conda

Discussion:
@chaozg: Hi @jakobsj , the DoDs are ready for your review 😄 @jakobsj: @chaozg looks good, maybe with the addition in item no 2 that new issues created are linked from this issue.@chaozg: Thanks @jakobsj , just updated item 2 according to your suggestion to "...corresponding issues are created, assigned and linked from this issue."

Solution patch:
diff --git a/cuqi/array/_array.py b/cuqi/array/_array.py
index ec0c5bf00..0ceecb522 100644
--- a/cuqi/array/_array.py
+++ b/cuqi/array/_array.py
@@ -15,19 +15,7 @@ class CUQIarray(np.ndarray):
         Boolean flag whether input_array is to be interpreted as parameter (True) or function values (False).
 
     geometry : cuqi.geometry.Geometry, default None
-        Contains the geometry related of the data
-
-    Attributes
-    ----------
-    funvals : CUQIarray
-        Returns itself as function values.
-
-    parameters : CUQIarray
-        Returns itself as parameters.
-
-    Methods
-    ----------
-    :meth:`plot`: Plots the data as function or parameters.
+        Contains the geometry related of the data.
     """
 
     def __repr__(self) -> str: 
@@ -62,6 +50,7 @@ def __array_finalize__(self, obj):
 
     @property
     def funvals(self):
+        """ Returns itself as function values. """
         if self.is_par is True:
             vals = self.geometry.par2fun(self)
         else:
@@ -82,6 +71,7 @@ def funvals(self):
 
     @property
     def parameters(self):
+        """ Returns itself as parameters. """
         if self.is_par is False:
             if self.dtype == np.dtype('O'):
                 # If the current state if the CUQIarray is function values, and
@@ -109,6 +99,7 @@ def to_numpy(self):
                 f"Cannot convert {self.__class__.__name__} to numpy array")
 
     def plot(self, plot_par=False, **kwargs):
+        """ Plot the data as function or parameters. """
         if plot_par:
             kwargs["is_par"]=True
             return self.geometry.plot(self.parameters, plot_par=plot_par, **kwargs)
diff --git a/cuqi/distribution/_beta.py b/cuqi/distribution/_beta.py
index bf3947a7d..a57af1826 100644
--- a/cuqi/distribution/_beta.py
+++ b/cuqi/distribution/_beta.py
@@ -9,23 +9,19 @@ class Beta(Distribution):
     """
     Multivariate beta distribution of independent random variables x_i. Each is distributed according to the PDF function
 
-    f(x) = x^(alpha-1) * (1-x)^(beta-1) * Gamma(alpha+beta) / (Gamma(alpha)*Gamma(beta))
+    .. math::
 
-    where Gamma is the Gamma function.
+        f(x) = x^{(\\alpha-1)}(1-x)^{(\\beta-1)}\Gamma(\\alpha+\\beta) / (\Gamma(\\alpha)\Gamma(\\beta))
+
+    where :math:`\Gamma` is the Gamma function.
 
     Parameters
     ------------
     alpha: float or array_like
+           The shape parameter :math:`\\alpha` of the beta distribution.
 
     beta: float or array_like
-
-    Methods
-    -----------
-    sample: generate one or more random samples
-    pdf: evaluate probability density function
-    logpdf: evaluate log probability density function
-    cdf: evaluate cumulative probability function
-    gradient: evaluate the gradient of the logpdf
+          The shape parameter :math:`\\beta` of the beta distribution.
     
     Example
     -------
diff --git a/cuqi/distribution/_custom.py b/cuqi/distribution/_custom.py
index cbfac3b68..866dfbb32 100644
--- a/cuqi/distribution/_custom.py
+++ b/cuqi/distribution/_custom.py
@@ -12,12 +12,6 @@ class UserDefinedDistribution(Distribution):
     gradient_func: Function evaluating the gradient of the logpdf. Callable.
     sample_func: Function drawing samples from distribution. Callable.
     
-    Methods
-    -----------
-    sample: generate one or more random samples
-    logpdf: evaluate log probability density function
-    gradient: evaluate gradient of logpdf
-    
     Example
     -----------
     .. code-block:: python
diff --git a/cuqi/distribution/_distribution.py b/cuqi/distribution/_distribution.py
index 1867c5324..f4f5eaab1 100644
--- a/cuqi/distribution/_distribution.py
+++ b/cuqi/distribution/_distribution.py
@@ -28,37 +28,6 @@ class Distribution(Density, ABC):
     is_symmetric : bool, default None
         Indicator if distribution is symmetric.
 
-    Attributes
-    ----------
-    dim : int or None
-        Dimension of distribution.
-
-    name : str or None
-        Name of distribution.
-    
-    geometry : Geometry or None
-        Geometry of distribution.
-
-    is_cond : bool
-        Indicator if distribution is conditional.
-
-    Methods
-    -------
-    pdf():
-        Evaluate the probability density function.
-
-    logpdf():
-        Evaluate the log probability density function.
-
-    sample():
-        Generate one or more random samples.
-
-    get_conditioning_variables():
-        Return the conditioning variables of distribution.
-
-    get_mutable_variables():
-        Return the mutable variables (attributes and properties) of distribution.
-
     Notes
     -----
     A distribution can be conditional if one or more mutable variables are unspecified.
@@ -239,6 +208,7 @@ def _gradient(self, *args, **kwargs):
             "enable_FD().")
 
     def sample(self,N=1,*args,**kwargs):
+        """ Sample from the distribution. """
 
         if self.is_cond:
             raise ValueError(f"Cannot sample from conditional distribution. Missing conditioning variables: {self.get_conditioning_variables()}")
@@ -263,6 +233,7 @@ def _sample(self,N):
         pass
 
     def pdf(self,x):
+        """ Evaluate the log probability density function of the distribution. """
         return np.exp(self.logpdf(x))
 
     def _condition(self, *args, **kwargs):
diff --git a/cuqi/distribution/_gamma.py b/cuqi/distribution/_gamma.py
index a1f5b2b90..c8a408f88 100644
--- a/cuqi/distribution/_gamma.py
+++ b/cuqi/distribution/_gamma.py
@@ -6,14 +6,18 @@
 class Gamma(Distribution):
     """
     Represents a multivariate Gamma distribution characterized by shape and rate parameters of independent random variables x_i. Each is distributed according to the PDF function
+
+    .. math::
     
-    f(x_i; shape, rate) = rate^shape * x_i^(shape-1) * exp(-rate * x_i) / Gamma(shape)
+        f(x_i; \\alpha, \\beta) = \\beta^\\alpha x_i^{\\alpha-1} \\exp(-\\beta x_i) / \Gamma(\\alpha)
 
-    where `shape` and `rate` are the parameters of the distribution, and Gamma is the Gamma function.
+    where shape :math:`\\alpha` and rate :math:`\\beta` are the parameters of the distribution, and :math:`\Gamma` is the Gamma function.
 
     In case shape and/or rate are arrays, the pdf looks like
 
-    f(x_i; shape_i, rate_i) = rate_i^shape_i * x_i^(shape_i-1) * exp(-rate_i * x_i) / Gamma(shape_i)
+    .. math::
+    
+        f(x_i; \\alpha_i, \\beta_i) = \\beta_i^{\\alpha_i} x_i^{\\alpha_i-1} \\exp(-\\beta_i x_i) / \Gamma(\\alpha_i)
 
     Parameters
     ----------
diff --git a/cuqi/distribution/_gaussian.py b/cuqi/distribution/_gaussian.py
index e3523f75f..3708acb67 100644
--- a/cuqi/distribution/_gaussian.py
+++ b/cuqi/distribution/_gaussian.py
@@ -730,18 +730,6 @@ class JointGaussianSqrtPrec(Distribution):
     ------------
     means: List of means for each Gaussian distribution.
     sqrtprecs: List of sqrt precision matricies for each Gaussian distribution.
-
-    Attributes
-    ------------
-    sqrtprec: Returns the sqrt precision matrix of the joined gaussian in stacked form.
-    sqrtprecTimesMean: Returns the sqrt precision matrix times the mean of the distribution.
-    
-    Methods
-    -----------
-    sample: generate one or more random samples (NotImplemented)
-    pdf: evaluate probability density function (NotImplemented)
-    logpdf: evaluate log probability density function (NotImplemented)
-    cdf: evaluate cumulative probability function (NotImplemented)
     """    
     def __init__(self,means=None,sqrtprecs=None,is_symmetric=True,**kwargs):
 
@@ -783,6 +771,7 @@ def dim(self):
 
     @property
     def sqrtprec(self):
+        """ Returns the sqrt precision matrix of the joined gaussian in stacked form. """
         if spa.issparse(self._sqrtprecs[0]):
             return spa.vstack((self._sqrtprecs))
         else:
@@ -790,6 +779,7 @@ def sqrtprec(self):
 
     @property
     def sqrtprecTimesMean(self):
+        """ Returns the sqrt precision matrix times the mean of the distribution."""
         result = []
         for i in range(len(self._means)):
             result.append((self._sqrtprecs[i]@self._means[i]).flatten())
diff --git a/cuqi/distribution/_inverse_gamma.py b/cuqi/distribution/_inverse_gamma.py
index 1cf1bb15a..e8537d435 100644
--- a/cuqi/distribution/_inverse_gamma.py
+++ b/cuqi/distribution/_inverse_gamma.py
@@ -8,9 +8,11 @@ class InverseGamma(Distribution):
     """
     Multivariate inverse gamma distribution of independent random variables x_i. Each is distributed according to the PDF function
 
-    f(x) = (x-location)^(-shape-1) * exp(-scale/(x-location)) / (scale^(-shape)*Gamma(shape))
+    .. math::
 
-    where shape, location and scale are the shape, location and scale of x_i, respectively. And Gamma is the Gamma function.
+        f(x) = (x-\\beta)^{(-\\alpha-1)} * \exp(-\\gamma/(x-\\beta)) / (\\gamma^{(-\\alpha)}*\Gamma(\\alpha))
+
+    where shape :math:`\\alpha`, location :math:`\\beta` and scale :math:`\\gamma` are the shape, location and scale of x_i, respectively. And :math:`\Gamma` is the Gamma function.
 
     Parameters
     ------------
@@ -23,14 +25,6 @@ class InverseGamma(Distribution):
     scale: float or array_like
         The scale of the inverse gamma distribution (non-negative)
 
-    
-    Methods
-    -----------
-    sample: generate one or more random samples
-    pdf: evaluate probability density function
-    logpdf: evaluate log probability density function
-    cdf: evaluate cumulative probability function
-    
     Example
     -------
     .. code-block:: python
diff --git a/cuqi/distribution/_lognormal.py b/cuqi/distribution/_lognormal.py
index 62a4fec83..3e4b657ab 100644
--- a/cuqi/distribution/_lognormal.py
+++ b/cuqi/distribution/_lognormal.py
@@ -16,13 +16,6 @@ class Lognormal(Distribution):
     cov: np.ndarray
         Covariance matrix of the normal distribution used to define the lognormal distribution 
     
-    Methods
-    -----------
-    sample: generate one or more random samples
-    pdf: evaluate probability density function
-    logpdf: evaluate log probability density function
-    cdf: evaluate cumulative probability function
-    
     Example
     -------
     .. code-block:: python
diff --git a/cuqi/distribution/_modifiedhalfnormal.py b/cuqi/distribution/_modifiedhalfnormal.py
index 327aebf60..9702dea77 100644
--- a/cuqi/distribution/_modifiedhalfnormal.py
+++ b/cuqi/distribution/_modifiedhalfnormal.py
@@ -7,14 +7,17 @@
 class ModifiedHalfNormal(Distribution):
     """
     Represents a modified half-normal (MHN) distribution, a three-parameter family of distributions generalizing the Gamma distribution.
-    The distribution is continuous with pdf 
-    f(x; alpha, beta, gamma) propto x^(alpha-1) * exp(-beta * x^2 + gamma * x)
+    The distribution is continuous with pdf
+
+    .. math::
+    
+        f(x; \\alpha, \\beta, \\gamma) \propto x^{(\\alpha-1)} * \exp(-\\beta * x^2 + \\gamma * x)
 
     The MHN generalizes the half-normal distribution, because
-    f(x; 1, beta, 0) propto exp(-beta * x^2)
+    :math:`f(x; 1, \\beta, 0) \propto \exp(-\\beta * x^2)`
 
     The MHN generalizes the gamma distribution because
-    f(x; alpha, 0, -gamma) propto x^(alpha-1) * exp(- gamma * x)
+    :math:`f(x; \\alpha, 0, -\\gamma) \propto x^{(\\alpha-1)} * \exp(- \\gamma * x)`
 
     Reference:
     [1] Sun, et al. "The Modified-Half-Normal distribution: Properties and an efficient sampling scheme." Communications in Statistics-Theory and Methods
@@ -22,13 +25,13 @@ class ModifiedHalfNormal(Distribution):
     Parameters
     ----------
     alpha : float
-        The polynomial exponent parameter of the MHN distribution. Must be positive.
+        The polynomial exponent parameter :math:`\\alpha` of the MHN distribution. Must be positive.
 
     beta : float
-        The quadratic exponential parameter of the MHN distribution. Must be positive.
+        The quadratic exponential parameter :math:`\\beta` of the MHN distribution. Must be positive.
 
     gamma : float
-        The linear exponential parameter of the MHN distribution.
+        The linear exponential parameter :math:`\\gamma` of the MHN distribution.
 
     """
     def __init__(self, alpha=None, beta=None, gamma=None, is_symmetric=False, **kwargs):
diff --git a/cuqi/distribution/_normal.py b/cuqi/distribution/_normal.py
index ebb940747..b101c4788 100644
--- a/cuqi/distribution/_normal.py
+++ b/cuqi/distribution/_normal.py
@@ -12,13 +12,6 @@ class Normal(Distribution):
     mean: mean of distribution
     std: standard deviation
     
-    Methods
-    -----------
-    sample: generate one or more random samples
-    pdf: evaluate probability density function
-    logpdf: evaluate log probability density function
-    cdf: evaluate cumulative probability function
-    
     Example
     -----------
     .. code-block:: python
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/cuqi/samples/_samples.py b/cuqi/samples/_samples.py
index 993846f8e..fcdb0a12b 100644
--- a/cuqi/samples/_samples.py
+++ b/cuqi/samples/_samples.py
@@ -36,24 +36,6 @@ class Samples(object):
     geometry : cuqi.geometry.Geometry, default None
         Contains the geometry related of the samples
 
-    Attributes
-    ----------
-    shape : tuple
-        Returns the shape of samples.
-
-    Ns : int
-        Returns the number of samples
-
-    Methods
-    ----------
-    :meth:`plot`: Plots one or more samples.
-    :meth:`plot_ci`: Plots a credibility interval for the samples.
-    :meth:`plot_mean`: Plots the mean of the samples.
-    :meth:`plot_std`: Plots the std of the samples.
-    :meth:`plot_chain`: Plots all samples of one or more variables (MCMC chain).
-    :meth:`hist_chain`: Plots histogram of all samples of a single variable (MCMC chain).
-    :meth:`burnthin`: Removes burn-in and thins samples.
-    :meth:`diagnostics`: Conducts diagnostics on the chain.
     """
     def __init__(self, samples, geometry=None, is_par=True, is_vec=True):
         self.geometry = geometry
@@ -83,6 +65,7 @@ def __iter__(self):
 
     @property
     def shape(self):
+        """Returns the shape of samples."""
         return self.samples.shape
 
     @property
@@ -408,6 +391,7 @@ def plot_std(self,*args,**kwargs):
         return ax
 
     def plot(self,sample_indices=None,*args,**kwargs):
+        """ Plots one or more samples. """
         Ns = self.Ns
         Np = 5 # Number of samples to plot if Ns > 5
 
@@ -447,6 +431,7 @@ def plot_chain(self, variable_indices=None, *args, **kwargs):
         return lines
     
     def hist_chain(self,variable_indices,*args,**kwargs):
+        """ Plots samples histogram of variables with indices specified in variable_indices. """
 
         self._raise_error_if_not_vec(self.hist_chain.__name__)
 
@@ -580,6 +565,7 @@ def plot_ci(self, percent=95, exact=None, *args, plot_envelope_kwargs=None, **kw
 
 
     def diagnostics(self):
+        """ Conducts diagnostics on the chain (Geweke test). """
         # Geweke test
         Geweke(self.samples.T)
 
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index c90ba5283..cb64f4a05 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -174,9 +174,12 @@ class LS(object):
     """Wrapper for :meth:`scipy.optimize.least_squares`.
 
     Solve nonlinear least-squares problems with bounds:
+
+    .. math::
     
-    minimize F(x) = 0.5 * sum(rho(f_i(x)**2), i = 0, ..., m-1)
-    subject to lb <= x <= ub
+        \min F(x) = 0.5 * \sum(\\rho(f_i(x)^2), i = 0, ..., m-1)
+
+    subject to :math:`lb <= x <= ub`.
     
     Parameters
     ----------
diff --git a/docs/api/index.rst b/docs/api/index.rst
index 40d37de3f..e4039bdcb 100644
--- a/docs/api/index.rst
+++ b/docs/api/index.rst
@@ -19,6 +19,9 @@ Some useful auxiliary helper modules:
    - :doc:`cuqi.samples <_autosummary/cuqi.samples>` contains tools for storing and manipulating MCMC samples.
    - :doc:`cuqi.solver <_autosummary/cuqi.solver>` contains tools point estimation of posteriors.
 
+Experimental modules:
+   - :doc:`cuqi.experimental <_autosummary/cuqi.experimental>` contains experimental modules that are not yet stable.
+
 Full API overview
    A complete auto generated overview of the API can be found below (or by navigating using the sidebar).
 
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}"
+