CUQIpy 1.0.0.post0.dev127__py3-none-any.whl → 1.0.0.post0.dev180__py3-none-any.whl

{CUQIpy-1.0.0.post0.dev127.dist-info → CUQIpy-1.0.0.post0.dev180.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: CUQIpy
-Version: 1.0.0.post0.dev127
+Version: 1.0.0.post0.dev180
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License: Apache License

{CUQIpy-1.0.0.post0.dev127.dist-info → CUQIpy-1.0.0.post0.dev180.dist-info}/RECORD

@@ -1,6 +1,6 @@
 cuqi/__init__.py,sha256=LsGilhl-hBLEn6Glt8S_l0OJzAA1sKit_rui8h-D-p0,488
 cuqi/_messages.py,sha256=fzEBrZT2kbmfecBBPm7spVu7yHdxGARQB4QzXhJbCJ0,415
-cuqi/_version.py,sha256=2GhcRA-eAhvzsZ7hWqzmLdIV7UEkRWfrTcDXqFBbrWE,510
+cuqi/_version.py,sha256=nZ9Rz10NtUnQZt09kW9YMygjxV0G-qsipgAL3xBFeAU,510
 cuqi/config.py,sha256=wcYvz19wkeKW2EKCGIKJiTpWt5kdaxyt4imyRkvtTRA,526
 cuqi/diagnostics.py,sha256=5OrbJeqpynqRXOe5MtOKKhe7EAVdOEpHIqHnlMW9G_c,3029
 cuqi/array/__init__.py,sha256=-EeiaiWGNsE3twRS4dD814BIlfxEsNkTCZUc5gjOXb0,30
@@ -32,13 +32,14 @@ cuqi/distribution/_normal.py,sha256=UeoTtGDT7YSf4ZNo2amlVF9K-YQpYbf8q76jcRJTVFw,
 cuqi/distribution/_posterior.py,sha256=zAfL0GECxekZ2lBt1W6_LN0U_xskMwK4VNce5xAF7ig,5018
 cuqi/distribution/_uniform.py,sha256=7xJmCZH_LPhuGkwEDGh-_CTtzcWKrXMOxtTJUFb7Ydo,1607
 cuqi/experimental/__init__.py,sha256=vhZvyMX6rl8Y0haqCzGLPz6PSUKyu75XMQbeDHqTTrw,83
-cuqi/experimental/mcmc/__init__.py,sha256=vVnohcm4EIUwbp1sr3LbB0BkXO8jyZsbiKMJmIgetYY,314
-cuqi/experimental/mcmc/_cwmh.py,sha256=G-8YjMqPraZm1Pm3n6scFkpa65gdtI1WTQxlL21etEI,8066
-cuqi/experimental/mcmc/_langevin_algorithm.py,sha256=ckVHDXLaw8hsUaOAFAEs7bL2Ny7W1QBKSc4AAC-TCis,9986
-cuqi/experimental/mcmc/_mh.py,sha256=AslackZJ3hPUNQfy70Fh9WoRPtcsHGqulI0LQgGHBus,3477
-cuqi/experimental/mcmc/_pcn.py,sha256=ma3pFqFgOmE7woZ41B5CGccKEuaacJPTmKvSEQhvtzs,3981
-cuqi/experimental/mcmc/_rto.py,sha256=nQHpSnUlE65TXBnGFk88JLawR44af2Rtle0oFCGYgaQ,11540
-cuqi/experimental/mcmc/_sampler.py,sha256=s-15bElbSZFHJlaV9gmiwk-UAneQEU9W-y5tLy_NMCU,11197
+cuqi/experimental/mcmc/__init__.py,sha256=S4aXYpnO75HQcwDYfr1-ki8UvlenPDXxshES5avtBF0,340
+cuqi/experimental/mcmc/_cwmh.py,sha256=yRlTk5a1QYfH3JyCecfOOTeDf-4-tmJ3Tl2Bc3pyp1Y,7336
+cuqi/experimental/mcmc/_hmc.py,sha256=qqAyoAajLE_JenYMgAbD3tknuEf75AJu-ufF69GKGk4,19384
+cuqi/experimental/mcmc/_langevin_algorithm.py,sha256=MX48u3GYgCckB6Q5h5kXr_qdIaLQH2toOG5u29OY7gk,8245
+cuqi/experimental/mcmc/_mh.py,sha256=aIV1Ntq0EAq3QJ1_X-DbP7eDAL-d_Or7d3RUO-R48I4,3090
+cuqi/experimental/mcmc/_pcn.py,sha256=3M8zhQGQa53Gz04AkC8wJM61_5rIjGVnhPefi8m4dbY,3531
+cuqi/experimental/mcmc/_rto.py,sha256=jSPznr34XPfWM6LysWIiN4hE-vtyti3cHyvzy9ruykg,11349
+cuqi/experimental/mcmc/_sampler.py,sha256=_5Uo2F-Mta46w3lo7WBVNwvTLYhES_BzMTJrKxA00c8,14861
 cuqi/geometry/__init__.py,sha256=Tz1WGzZBY-QGH3c0GiyKm9XHN8MGGcnU6TUHLZkzB3o,842
 cuqi/geometry/_geometry.py,sha256=WYFC-4_VBTW73b2ldsnfGYKvdSiCE8plr89xTSmkadg,46804
 cuqi/implicitprior/__init__.py,sha256=ZRZ9fgxgEl5n0A9F7WCl1_jid-GUiC8ZLkyTmGQmFlY,100
@@ -59,7 +60,7 @@ cuqi/sampler/_conjugate.py,sha256=ztmUR3V3qZk9zelKx48ULnmMs_zKTDUfohc256VOIe8,27
 cuqi/sampler/_conjugate_approx.py,sha256=xX-X71EgxGnZooOY6CIBhuJTs3dhcKfoLnoFxX3CO2g,1938
 cuqi/sampler/_cwmh.py,sha256=VlAVT1SXQU0yD5ZeR-_ckWvX-ifJrMweFFdFbxdfB_k,7775
 cuqi/sampler/_gibbs.py,sha256=N7qcePwMkRtxINN5JF0FaMIdDCXZGqsfKjfha_KHCck,8627
-cuqi/sampler/_hmc.py,sha256=76nPkvNU0wLSg4qvm-1s048MzQasl5Qk94sHpyeJ5hM,14819
+cuqi/sampler/_hmc.py,sha256=EUTefZir-wapoZ7OZFb5M5vayL8z6XksZRMY1BpbuXc,15027
 cuqi/sampler/_langevin_algorithm.py,sha256=o5EyvaR6QGAD7LKwXVRC3WwAP5IYJf5GoMVWl9DrfOA,7861
 cuqi/sampler/_laplace_approximation.py,sha256=u018Z5eqlcq_cIwD9yNOaA15dLQE_vUWaee5Xp8bcjg,6454
 cuqi/sampler/_mh.py,sha256=V5tIdn-KdfWo4J_Nbf-AH6XwKWblWUyc4BeuSikUHsE,7062
@@ -75,8 +76,8 @@ cuqi/testproblem/_testproblem.py,sha256=x769LwwRdJdzIiZkcQUGb_5-vynNTNALXWKato7s
 cuqi/utilities/__init__.py,sha256=EfxHLdsyDNugbmbzs43nV_AeKcycM9sVBjG9WZydagA,351
 cuqi/utilities/_get_python_variable_name.py,sha256=QwlBVj2koJRA8s8pWd554p7-ElcI7HUwY32HknaR92E,1827
 cuqi/utilities/_utilities.py,sha256=At3DOXRdF3GwLkVcM2FXooGyjAGfPkIM0bRzhTfLmWk,8046
-CUQIpy-1.0.0.post0.dev127.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
-CUQIpy-1.0.0.post0.dev127.dist-info/METADATA,sha256=bHtJEgkhpP_C50jwE7SwIpOY0Hm48V2qMKy_v6k0SiU,18393
-CUQIpy-1.0.0.post0.dev127.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-CUQIpy-1.0.0.post0.dev127.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
-CUQIpy-1.0.0.post0.dev127.dist-info/RECORD,,
+CUQIpy-1.0.0.post0.dev180.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
+CUQIpy-1.0.0.post0.dev180.dist-info/METADATA,sha256=YSVheKM46Y6Al3zXEu11wlCkJxAasZw2onXB0nTLN08,18393
+CUQIpy-1.0.0.post0.dev180.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+CUQIpy-1.0.0.post0.dev180.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
+CUQIpy-1.0.0.post0.dev180.dist-info/RECORD,,

cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-04-05T11:26:24+0200",
+ "date": "2024-04-22T18:01:04+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "476cc08200034336b46144fac4c819f8298fa587",
- "version": "1.0.0.post0.dev127"
+ "full-revisionid": "be4b485322b1be52f78dfe6d03694d6ba2000b11",
+ "version": "1.0.0.post0.dev180"
 }
 '''  # END VERSION_JSON
 

cuqi/experimental/mcmc/__init__.py

@@ -6,3 +6,4 @@ from ._mh import MHNew
 from ._pcn import pCNNew
 from ._rto import LinearRTONew, RegularizedLinearRTONew
 from ._cwmh import CWMHNew
+from ._hmc import NUTSNew

cuqi/experimental/mcmc/_cwmh.py

@@ -144,7 +144,7 @@ class CWMHNew(ProposalBasedSamplerNew):
 
         # Propose a sample x_all_components from the proposal distribution
         # for all the components
-        target_eval_t = self.current_target
+        target_eval_t = self.current_target_logd
         if isinstance(self.proposal,cuqi.distribution.Distribution):
             x_all_components = self.proposal(
                 location= self.current_point, scale=self.scale).sample()
@@ -175,7 +175,7 @@ class CWMHNew(ProposalBasedSamplerNew):
 
             x_star = x_t.copy()
 
-        self.current_target = target_eval_t
+        self.current_target_logd = target_eval_t
         self.current_point = x_t
 
         return acc
@@ -199,23 +199,3 @@ class CWMHNew(ProposalBasedSamplerNew):
         # Update the scale parameter
         self.scale = np.minimum(scale_temp, np.ones(self.dim))
         self._scale_temp = scale_temp
-
-    def get_state(self):
-        current_point = self.current_point
-        if isinstance(current_point, CUQIarray):
-            current_point = current_point.to_numpy()
-
-        return {'sampler_type': 'CWMH',
-                'current_point': current_point,
-                'current_target': self.current_target,
-                'scale': self.scale}
-
-    def set_state(self, state):
-        current_point = state['current_point']
-        if not isinstance(current_point, CUQIarray):
-            current_point = CUQIarray(current_point,
-                                      geometry=self.target.geometry)
-
-        self.current_point = current_point
-        self.current_target = state['current_target']
-        self.scale = state['scale']

cuqi/experimental/mcmc/_hmc.py

@@ -0,0 +1,470 @@
+import numpy as np
+import numpy as np
+from cuqi.experimental.mcmc import SamplerNew
+from cuqi.array import CUQIarray
+from numbers import Number
+
+class NUTSNew(SamplerNew):
+    """No-U-Turn Sampler (Hoffman and Gelman, 2014).
+
+    Samples a distribution given its logpdf and gradient using a Hamiltonian
+    Monte Carlo (HMC) algorithm with automatic parameter tuning.
+
+    For more details see: See Hoffman, M. D., & Gelman, A. (2014). The no-U-turn
+    sampler: Adaptively setting path lengths in Hamiltonian Monte Carlo. Journal
+    of Machine Learning Research, 15, 1593-1623.
+
+    Parameters
+    ----------
+    target : `cuqi.distribution.Distribution`
+        The target distribution to sample. Must have logpdf and gradient method.
+        Custom logpdfs and gradients are supported by using a
+        :class:`cuqi.distribution.UserDefinedDistribution`.
+    
+    initial_point : ndarray
+        Initial parameters. *Optional*. If not provided, the initial point is 
+        an array of ones.
+
+    max_depth : int
+        Maximum depth of the tree >=0 and the default is 15.
+
+    step_size : None or float
+        If step_size is provided (as positive float), it will be used as initial
+        step size. If None, the step size will be estimated by the sampler.
+
+    opt_acc_rate : float
+        The optimal acceptance rate to reach if using adaptive step size.
+        Suggested values are 0.6 (default) or 0.8 (as in stan). In principle,
+        opt_acc_rate should be in (0, 1), however, choosing a value that is very
+        close to 1 or 0 might lead to poor performance of the sampler.
+
+    callback : callable, *Optional*
+        If set this function will be called after every sample.
+        The signature of the callback function is
+        `callback(sample, sample_index)`,
+        where `sample` is the current sample and `sample_index` is the index of
+        the sample.
+        An example is shown in demos/demo31_callback.py.
+
+    Example
+    -------
+    .. code-block:: python
+
+        # Import cuqi
+        import cuqi
+
+        # Define a target distribution
+        tp = cuqi.testproblem.WangCubic()
+        target = tp.posterior
+
+        # Set up sampler
+        sampler = cuqi.experimental.mcmc.NUTSNew(target)
+
+        # Sample
+        sampler.warmup(5000)
+        sampler.sample(10000)
+
+        # Get samples
+        samples = sampler.get_samples()
+
+        # Plot samples
+        samples.plot_pair()
+
+    After running the NUTS sampler, run diagnostics can be accessed via the 
+    following attributes:
+
+    .. code-block:: python
+
+        # Number of tree nodes created each NUTS iteration
+        sampler.num_tree_node_list
+
+        # Step size used in each NUTS iteration
+        sampler.epsilon_list
+
+        # Suggested step size during adaptation (the value of this step size is
+        # only used after adaptation).
+        sampler.epsilon_bar_list
+
+    """
+
+    _STATE_KEYS = SamplerNew._STATE_KEYS.union({'_epsilon', '_epsilon_bar',
+                                                '_H_bar', '_mu',
+                                                '_alpha', '_n_alpha'})
+
+    _HISTORY_KEYS = SamplerNew._HISTORY_KEYS.union({'num_tree_node_list',
+                                                    'epsilon_list',
+                                                    'epsilon_bar_list'})
+
+    def __init__(self, target, initial_point=None, max_depth=15,
+                 step_size=None, opt_acc_rate=0.6, **kwargs):
+        super().__init__(target, initial_point=initial_point, **kwargs)
+
+        # Assign parameters as attributes
+        self.max_depth = max_depth
+        self.step_size = step_size
+        self.opt_acc_rate = opt_acc_rate
+        
+        # Set current point 
+        self.current_point = self.initial_point
+
+        # 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
+
+        # Arrays to store acceptance rate
+        self._acc = [None]
+
+        # NUTS run diagnostic:
+        # number of tree nodes created each NUTS iteration
+        self._num_tree_node = 0
+        # Create lists to store NUTS run diagnostics
+        self._create_run_diagnostic_attributes()
+
+    #=========================================================================
+    #============================== Properties ===============================
+    #=========================================================================
+    @property
+    def max_depth(self):
+        return self._max_depth
+
+    @max_depth.setter
+    def max_depth(self, value):
+        if not isinstance(value, int):
+            raise TypeError('max_depth must be an integer.')
+        if value < 0:
+            raise ValueError('max_depth must be >= 0.')
+        self._max_depth = value
+
+    @property
+    def step_size(self):
+        return self._step_size
+    
+    @step_size.setter
+    def step_size(self, value):
+        if value is None:
+            pass # NUTS will adapt the step size
+
+        # step_size must be a positive float, raise error otherwise
+        elif isinstance(value, bool)\
+            or not isinstance(value, Number)\
+            or value <= 0:
+            raise TypeError('step_size must be a positive float or None.')
+        self._step_size = value
+
+    @property
+    def opt_acc_rate(self):
+        return self._opt_acc_rate
+    
+    @opt_acc_rate.setter
+    def opt_acc_rate(self, value):
+        if not isinstance(value, Number) or value <= 0 or value >= 1:
+            raise ValueError('opt_acc_rate must be a float in (0, 1).')
+        self._opt_acc_rate = value
+
+    #=========================================================================
+    #================== Implement methods required by SamplerNew =============
+    #=========================================================================
+    def validate_target(self):
+        # Check if the target has logd and gradient methods
+        try:
+            current_target_logd, current_target_grad =\
+            self._nuts_target(np.ones(self.dim))
+        except:
+            raise ValueError('Target must have logd and gradient methods.')
+
+    def reset(self):
+        # Call the parent reset method
+        super().reset()
+        # Reset NUTS run diagnostic attributes
+        self._reset_run_diagnostic_attributes()
+
+    def step(self):
+        # Convert current_point, logd, and grad to numpy arrays
+        # if they are CUQIarray objects
+        if isinstance(self.current_point, CUQIarray):
+            self.current_point = self.current_point.to_numpy() 
+        if isinstance(self.current_target_logd, CUQIarray):
+            self.current_target_logd = self.current_target_logd.to_numpy()
+        if isinstance(self.current_target_grad, CUQIarray):
+            self.current_target_grad = self.current_target_grad.to_numpy()
+
+        # reset number of tree nodes for each iteration
+        self._num_tree_node = 0
+
+        # copy current point, logd, and grad in local variables
+        point_k = self.current_point.copy() # initial position (parameters)
+        logd_k = self.current_target_logd
+        grad_k = self.current_target_grad.copy() # initial gradient
+        
+        # compute r_k and Hamiltonian
+        r_k = self._Kfun(1, 'sample') # resample momentum vector
+        Ham = logd_k - self._Kfun(r_k, 'eval') # Hamiltonian
+
+        # slice variable
+        log_u = Ham - np.random.exponential(1, size=1)
+
+        # initialization
+        j, s, n = 0, 1, 1
+        point_minus, point_plus = np.copy(point_k), np.copy(point_k)
+        grad_minus, grad_plus = np.copy(grad_k), np.copy(grad_k)
+        r_minus, r_plus = np.copy(r_k), np.copy(r_k)
+
+        # run NUTS
+        while (s == 1) and (j <= self.max_depth):
+            # sample a direction
+            v = int(2*(np.random.rand() < 0.5)-1)
+
+            # build tree: doubling procedure
+            if (v == -1):
+                point_minus, r_minus, grad_minus, _, _, _, \
+                    point_prime, logd_prime, grad_prime,\
+                        n_prime, s_prime, alpha, n_alpha = \
+                            self._BuildTree(point_minus, r_minus, grad_minus,
+                                            Ham, log_u, v, j, self._epsilon)
+            else:
+                _, _, _, point_plus, r_plus, grad_plus, \
+                    point_prime, logd_prime, grad_prime,\
+                        n_prime, s_prime, alpha, n_alpha = \
+                            self._BuildTree(point_plus, r_plus, grad_plus,
+                                            Ham, log_u, v, j, self._epsilon)
+
+            # 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):
+                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)
+
+            # update number of particles, tree level, and stopping criterion
+            n += n_prime
+            dpoints = point_plus - point_minus
+            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
+
+        # update run diagnostic attributes
+        self._update_run_diagnostic_attributes(
+            self._num_tree_node, self._epsilon, self._epsilon_bar)
+        
+        self._epsilon = self._epsilon_bar 
+        if np.isnan(self.current_target_logd):
+            raise NameError('NaN potential func')
+
+    def tune(self, skip_len, update_count):
+        """ adapt epsilon during burn-in using dual averaging"""
+        k = update_count+1
+
+        # Fixed parameters that do not change during the run
+        gamma, t_0, kappa = 0.05, 10, 0.75 # kappa in (0.5, 1]
+
+        eta1 = 1/(k + t_0)
+        self._H_bar = (1-eta1)*self._H_bar +\
+            eta1*(self.opt_acc_rate - (self._alpha/self._n_alpha))
+        self._epsilon = np.exp(self._mu - (np.sqrt(k)/gamma)*self._H_bar)
+        eta = k**(-kappa)
+        self._epsilon_bar =\
+            np.exp(eta*np.log(self._epsilon) +(1-eta)*np.log(self._epsilon_bar))
+
+    def _pre_warmup(self):
+        super()._pre_warmup()
+
+        self.current_target_logd, self.current_target_grad =\
+            self._nuts_target(self.current_point)
+
+        # Set up tuning parameters (only first time tuning is called)
+        # Note:
+        #  Parameters changes during the tune run
+        #    self._epsilon_bar
+        #    self._H_bar
+        #    self._epsilon
+        #  Parameters that does not change during the run
+        #    self._mu
+
+        if self._epsilon is None:
+            # parameters dual averaging
+            self._epsilon = self._FindGoodEpsilon()
+            # Parameter mu, does not change during the run
+            self._mu = np.log(10*self._epsilon)
+
+        if self._epsilon_bar is None: # Initial value of epsilon_bar
+            self._epsilon_bar = 1
+
+        if self._H_bar is None: # Initial value of H_bar
+            self._H_bar = 0
+
+    def _pre_sample(self):
+        super()._pre_sample()
+
+        self.current_target_logd, self.current_target_grad =\
+            self._nuts_target(self.current_point)
+        
+        # Set up epsilon and epsilon_bar if not set
+        if self._epsilon is None:
+            if self.step_size is None:
+                step_size = self._FindGoodEpsilon()
+            else:
+                step_size = self.step_size
+            self._epsilon = step_size
+            self._epsilon_bar = step_size
+
+    #=========================================================================
+    def _nuts_target(self, x): # returns logposterior tuple evaluation-gradient
+        return self.target.logd(x), self.target.gradient(x)
+
+    #=========================================================================
+    # auxiliary standard Gaussian PDF: kinetic energy function
+    # d_log_2pi = d*np.log(2*np.pi)
+    def _Kfun(self, r, flag):
+        if flag == 'eval': # evaluate
+            return 0.5*(r.T @ r) #+ d_log_2pi 
+        if flag == 'sample': # sample
+            return np.random.standard_normal(size=self.dim)
+
+    #=========================================================================
+    def _FindGoodEpsilon(self, epsilon=1):
+        point_k = self.current_point
+        self.current_target_logd, self.current_target_grad = self._nuts_target(
+            point_k)
+        logd = self.current_target_logd
+        grad = self.current_target_grad
+
+        r = self._Kfun(1, 'sample')    # resample a momentum
+        Ham = logd - self._Kfun(r, 'eval')     # initial Hamiltonian
+        _, r_prime, logd_prime, grad_prime = self._Leapfrog(
+            point_k, r, grad, epsilon)
+
+        # trick to make sure the step is not huge, leading to infinite values of
+        # the likelihood
+        k = 1
+        while np.isinf(logd_prime) or np.isinf(grad_prime).any():
+            k *= 0.5
+            _, r_prime, logd_prime, grad_prime = self._Leapfrog(
+                point_k, r, grad, epsilon*k)
+        epsilon = 0.5*k*epsilon
+
+        # doubles/halves the value of epsilon until the accprob of the Langevin
+        # proposal crosses 0.5
+        Ham_prime = logd_prime - self._Kfun(r_prime, 'eval')
+        log_ratio = Ham_prime - Ham
+        a = 1 if log_ratio > np.log(0.5) else -1
+        while (a*log_ratio > -a*np.log(2)):
+            epsilon = (2**a)*epsilon
+            _, r_prime, logd_prime, _ = self._Leapfrog(
+                point_k, r, grad, epsilon)
+            Ham_prime = logd_prime - self._Kfun(r_prime, 'eval')
+            log_ratio = Ham_prime - Ham
+        return epsilon
+
+    #=========================================================================
+    def _Leapfrog(self, point_old, r_old, grad_old, epsilon):
+        # symplectic integrator: trajectories preserve phase space volumen
+        r_new = r_old + 0.5*epsilon*grad_old     # half-step
+        point_new = point_old + epsilon*r_new     # full-step
+        logd_new, grad_new = self._nuts_target(point_new)     # new gradient
+        r_new += 0.5*epsilon*grad_new     # half-step
+        return point_new, r_new, logd_new, grad_new
+
+    #=========================================================================
+    def _BuildTree(
+            self, point_k, r, grad, Ham, log_u, v, j, epsilon, Delta_max=1000):
+        # Increment the number of tree nodes counter
+        self._num_tree_node += 1
+
+        if (j == 0):     # base case
+            # single leapfrog step in the direction v
+            point_prime, r_prime, logd_prime, grad_prime = self._Leapfrog(
+                point_k, r, grad, v*epsilon)
+            Ham_prime = logd_prime - self._Kfun(r_prime, 'eval') # Hamiltonian
+                                                                 # eval
+            n_prime = int(log_u <= Ham_prime)     # if particle is in the slice
+            s_prime = int(log_u < Delta_max + Ham_prime)     # check U-turn
+            #
+            diff_Ham = Ham_prime - Ham
+
+            # Compute the acceptance probability
+            # alpha_prime = min(1, np.exp(diff_Ham))
+            # written in a stable way to avoid overflow when computing
+            # exp(diff_Ham) for large values of diff_Ham
+            alpha_prime = 1 if diff_Ham > 0 else np.exp(diff_Ham)
+            n_alpha_prime = 1
+            #
+            point_minus, point_plus = point_prime, point_prime
+            r_minus, r_plus = r_prime, r_prime
+            grad_minus, grad_plus = grad_prime, grad_prime
+        else: 
+            # recursion: build the left/right subtrees
+            point_minus, r_minus, grad_minus, point_plus, r_plus, grad_plus, \
+                point_prime, logd_prime, grad_prime,\
+                    n_prime, s_prime, alpha_prime, n_alpha_prime = \
+                        self._BuildTree(point_k, r, grad,
+                                        Ham, log_u, v, j-1, epsilon)
+            if (s_prime == 1): # do only if the stopping criteria does not 
+                               # verify at the first subtree
+                if (v == -1):
+                    point_minus, r_minus, grad_minus, _, _, _, \
+                        point_2prime, logd_2prime, grad_2prime,\
+                            n_2prime, s_2prime, alpha_2prime, n_alpha_2prime = \
+                                self._BuildTree(point_minus, r_minus, grad_minus,
+                                                Ham, log_u, v, j-1, epsilon)
+                else:
+                    _, _, _, point_plus, r_plus, grad_plus, \
+                        point_2prime, logd_2prime, grad_2prime,\
+                            n_2prime, s_2prime, alpha_2prime, n_alpha_2prime = \
+                                self._BuildTree(point_plus, r_plus, grad_plus,
+                                                Ham, log_u, v, j-1, epsilon)
+
+                # Metropolis step
+                alpha2 = n_2prime / max(1, (n_prime + n_2prime))
+                if (np.random.rand() <= alpha2):
+                    point_prime = np.copy(point_2prime)
+                    logd_prime = np.copy(logd_2prime)
+                    grad_prime = np.copy(grad_2prime)
+
+                # update number of particles and stopping criterion
+                alpha_prime += alpha_2prime
+                n_alpha_prime += n_alpha_2prime
+                dpoints = point_plus - point_minus
+                s_prime = s_2prime *\
+                    int((dpoints@r_minus.T)>=0) * int((dpoints@r_plus.T)>=0)
+                n_prime += n_2prime
+
+        return point_minus, r_minus, grad_minus, point_plus, r_plus, grad_plus,\
+            point_prime, logd_prime, grad_prime,\
+                n_prime, s_prime, alpha_prime, n_alpha_prime
+
+    #=========================================================================
+    #======================== Diagnostic methods =============================
+    #=========================================================================
+
+    def _create_run_diagnostic_attributes(self):
+        """A method to create attributes to store NUTS run diagnostic."""
+        self._reset_run_diagnostic_attributes()
+
+    def _reset_run_diagnostic_attributes(self):
+        """A method to reset attributes to store NUTS run diagnostic."""
+        # List to store number of tree nodes created each NUTS iteration
+        self.num_tree_node_list = []
+        # List of step size used in each NUTS iteration 
+        self.epsilon_list = []
+        # List of burn-in step size suggestion during adaptation 
+        # only used when adaptation is done
+        # remains fixed after adaptation (after burn-in)
+        self.epsilon_bar_list = []
+
+    def _update_run_diagnostic_attributes(self, n_tree, eps, eps_bar):
+        """A method to update attributes to store NUTS run diagnostic."""
+        # Store the number of tree nodes created in iteration k
+        self.num_tree_node_list.append(n_tree)
+        # Store the step size used in iteration k
+        self.epsilon_list.append(eps)
+        # Store the step size suggestion during adaptation in iteration k
+        self.epsilon_bar_list.append(eps_bar)

cuqi/experimental/mcmc/_langevin_algorithm.py

@@ -60,14 +60,17 @@ class ULANew(SamplerNew): # Refactor to Proposal-based sampler?
     A Deblur example can be found in demos/demo27_ULA.py
     # TODO: update demo once sampler merged
     """
+
+    _STATE_KEYS = SamplerNew._STATE_KEYS.union({'current_target_logd', 'scale', 'current_target_grad'})
+
     def __init__(self, target, scale=1.0, **kwargs):
 
         super().__init__(target, **kwargs)
 
         self.scale = scale
         self.current_point = self.initial_point
-        self.current_target_eval = self.target.logd(self.current_point)
-        self.current_target_grad_eval = self.target.gradient(self.current_point)
+        self.current_target_logd = self.target.logd(self.current_point)
+        self.current_target_grad = self.target.gradient(self.current_point)
         self._acc = [1] # TODO. Check if we need this
 
     def validate_target(self):
@@ -99,15 +102,15 @@ class ULANew(SamplerNew): # Refactor to Proposal-based sampler?
             1 (accepted)
         """
         self.current_point = x_star
-        self.current_target_eval = target_eval_star
-        self.current_target_grad_eval = target_grad_star
+        self.current_target_logd = target_eval_star
+        self.current_target_grad = target_grad_star
         acc = 1
         return acc
 
     def step(self):
         # propose state
         xi = cuqi.distribution.Normal(mean=np.zeros(self.dim), std=np.sqrt(self.scale)).sample()
-        x_star = self.current_point + 0.5*self.scale*self.current_target_grad_eval + xi
+        x_star = self.current_point + 0.5*self.scale*self.current_target_grad + xi
 
         # evaluate target
         target_eval_star, target_grad_star = self.target.logd(x_star), self.target.gradient(x_star)
@@ -120,26 +123,6 @@ class ULANew(SamplerNew): # Refactor to Proposal-based sampler?
     def tune(self, skip_len, update_count):
         pass
 
-    def get_state(self):
-        if isinstance(self.current_point, CUQIarray):
-            self.current_point = self.current_point.to_numpy()
-        if isinstance(self.current_target_eval, CUQIarray):
-            self.current_target_eval = self.current_target_eval.to_numpy()
-        if isinstance(self.current_target_grad_eval, CUQIarray):
-            self.current_target_grad_eval = self.current_target_grad_eval.to_numpy()
-        return {'sampler_type': 'ULA', 'current_point': self.current_point, \
-                'current_target_eval': self.current_target_eval, \
-                'current_target_grad_eval': self.current_target_grad_eval, \
-                'scale': self.scale}
-
-    def set_state(self, state):
-        temp = CUQIarray(state['current_point'] , geometry=self.target.geometry)
-        self.current_point = temp
-        temp = CUQIarray(state['current_target_eval'] , geometry=self.target.geometry)
-        self.current_target_eval = temp
-        temp = CUQIarray(state['current_target_grad_eval'] , geometry=self.target.geometry)
-        self.current_target_grad_eval = temp
-        self.scale = state['scale']
 
 class MALANew(ULANew): # Refactor to Proposal-based sampler?
     """  Metropolis-adjusted Langevin algorithm (MALA) (Roberts and Tweedie, 1996)
@@ -219,9 +202,9 @@ class MALANew(ULANew): # Refactor to Proposal-based sampler?
         scaler
             1 if accepted, 0 otherwise
         """
-        log_target_ratio = target_eval_star - self.current_target_eval
+        log_target_ratio = target_eval_star - self.current_target_logd
         log_prop_ratio = self._log_proposal(self.current_point, x_star, target_grad_star) \
-            - self._log_proposal(x_star, self.current_point,  self.current_target_grad_eval)
+            - self._log_proposal(x_star, self.current_point,  self.current_target_grad)
         log_alpha = min(0, log_target_ratio + log_prop_ratio)
 
         # accept/reject with Metropolis
@@ -229,8 +212,8 @@ class MALANew(ULANew): # Refactor to Proposal-based sampler?
         log_u = np.log(np.random.rand())
         if (log_u <= log_alpha) and (np.isnan(target_eval_star) == False):
             self.current_point = x_star
-            self.current_target_eval = target_eval_star
-            self.current_target_grad_eval = target_grad_star
+            self.current_target_logd = target_eval_star
+            self.current_target_grad = target_grad_star
             acc = 1
         return acc
 
@@ -241,15 +224,3 @@ class MALANew(ULANew): # Refactor to Proposal-based sampler?
         mu = theta_k + ((self.scale)/2)*g_logpi_k
         misfit = theta_star - mu
         return -0.5*((1/(self.scale))*(misfit.T @ misfit))
-
-    def get_state(self):
-        if isinstance(self.current_point, CUQIarray):
-            self.current_point = self.current_point.to_numpy()
-        if isinstance(self.current_target_eval, CUQIarray):
-            self.current_target_eval = self.current_target_eval.to_numpy()
-        if isinstance(self.current_target_grad_eval, CUQIarray):
-            self.current_target_grad_eval = self.current_target_grad_eval.to_numpy()
-        return {'sampler_type': 'MALA', 'current_point': self.current_point, \
-                'current_target_eval': self.current_target_eval, \
-                'current_target_grad_eval': self.current_target_grad_eval, \
-                'scale': self.scale}

cuqi/experimental/mcmc/_mh.py

@@ -23,6 +23,8 @@ class MHNew(ProposalBasedSamplerNew):
 
     """
 
+    _STATE_KEYS = ProposalBasedSamplerNew._STATE_KEYS.union({'scale', '_scale_temp'})
+
     def __init__(self, target, proposal=None, scale=1, **kwargs):
         super().__init__(target, proposal=proposal, scale=scale, **kwargs)
         # Due to a bug? in old MH, we must keep track of this extra variable to match behavior.
@@ -54,7 +56,7 @@ class MHNew(ProposalBasedSamplerNew):
         target_eval_star = self.target.logd(x_star)
 
         # ratio and acceptance probability
-        ratio = target_eval_star - self.current_target # proposal is symmetric
+        ratio = target_eval_star - self.current_target_logd # proposal is symmetric
         alpha = min(0, ratio)
 
         # accept/reject
@@ -62,7 +64,7 @@ class MHNew(ProposalBasedSamplerNew):
         acc = 0
         if (u_theta <= alpha):
             self.current_point = x_star
-            self.current_target = target_eval_star
+            self.current_target_logd = target_eval_star
             acc = 1
         
         return acc
@@ -79,13 +81,3 @@ class MHNew(ProposalBasedSamplerNew):
 
         # update parameters
         self.scale = min(self._scale_temp, 1)
-
-    def get_state(self):
-        return {'sampler_type': 'MH', 'current_point': self.current_point.to_numpy(), 'current_target': self.current_target.to_numpy(), 'scale': self.scale}
-
-    def set_state(self, state):
-        temp = CUQIarray(state['current_point'] , geometry=self.target.geometry)
-        self.current_point = temp
-        temp = CUQIarray(state['current_target'] , geometry=self.target.geometry)
-        self.current_target = temp
-        self.scale = state['scale']

cuqi/experimental/mcmc/_pcn.py

@@ -5,13 +5,15 @@ from cuqi.array import CUQIarray
 
 class pCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
 
+    _STATE_KEYS = SamplerNew._STATE_KEYS.union({'scale', 'current_likelihood_logd'})
+
     def __init__(self, target, scale=1.0, **kwargs):
 
         super().__init__(target, **kwargs)
 
         self.scale = scale
         self.current_point = self.initial_point
-        self.current_loglike_eval = self._loglikelihood(self.current_point)
+        self.current_likelihood_logd = self._loglikelihood(self.current_point)
 
         self._acc = [1] # TODO. Check if we need this
 
@@ -33,7 +35,7 @@ class pCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
         loglike_eval_star =  self._loglikelihood(x_star) 
 
         # ratio and acceptance probability
-        ratio = loglike_eval_star - self.current_loglike_eval  # proposal is symmetric
+        ratio = loglike_eval_star - self.current_likelihood_logd  # proposal is symmetric
         alpha = min(0, ratio)
 
         # accept/reject
@@ -41,7 +43,7 @@ class pCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
         u_theta = np.log(np.random.rand())
         if (u_theta <= alpha):
             self.current_point = x_star
-            self.current_loglike_eval = loglike_eval_star
+            self.current_likelihood_logd = loglike_eval_star
             acc = 1
         
         return acc
@@ -87,15 +89,3 @@ class pCNNew(SamplerNew):  # Refactor to Proposal-based sampler?
 
     def tune(self, skip_len, update_count):
         pass
-
-    def get_state(self):
-        return {'sampler_type': 'PCN', 'current_point': self.current_point.to_numpy(), \
-                'current_loglike_eval': self.current_loglike_eval.to_numpy(), \
-                'scale': self.scale}
-
-    def set_state(self, state):
-        temp = CUQIarray(state['current_point'] , geometry=self.target.geometry)
-        self.current_point = temp
-        temp = CUQIarray(state['current_loglike_eval'] , geometry=self.target.geometry)
-        self.current_loglike_eval = temp
-        self.scale = state['scale']

cuqi/experimental/mcmc/_rto.py

@@ -50,6 +50,7 @@ class LinearRTONew(SamplerNew):
 
         if initial_point is None: #TODO: Replace later with a getter
             self.initial_point = np.zeros(self.dim)
+            self._samples = [self.initial_point]
 
         self.current_point = self.initial_point
         self._acc = [1] # TODO. Check if we need this
@@ -188,12 +189,6 @@ class LinearRTONew(SamplerNew):
         if not hasattr(self.prior, "sqrtprecTimesMean"):
             raise TypeError("Prior must contain a sqrtprecTimesMean attribute")
 
-    def get_state(self): #TODO: LinearRTO only need initial_point for reproducibility?
-        return {'sampler_type': 'LinearRTO'}
-
-    def set_state(self, state): #TODO: LinearRTO only need initial_point for reproducibility?
-        pass
-
 class RegularizedLinearRTONew(LinearRTONew):
     """
     Regularized Linear RTO (Randomize-Then-Optimize) sampler.
@@ -272,4 +267,4 @@ class RegularizedLinearRTONew(LinearRTONew):
                     maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
         self.current_point, _ = sim.solve()
         acc = 1
-        return acc
+        return acc

cuqi/experimental/mcmc/_sampler.py

@@ -21,6 +21,11 @@ class SamplerNew(ABC):
     Samples are stored in a list to allow for dynamic growth of the sample set. Returning samples is done by creating a new Samples object from the list of samples.
 
     """
+    _STATE_KEYS = {'current_point'}
+    """ Set of keys for the state dictionary. """
+
+    _HISTORY_KEYS = {'_samples', '_acc'}
+    """ Set of keys for the history dictionary. """
 
     def __init__(self, target: cuqi.density.Density, initial_point=None, callback=None):
         """ Initializer for abstract base class for all samplers.
@@ -45,8 +50,10 @@ class SamplerNew(ABC):
         # Choose initial point if not given
         if initial_point is None:
             initial_point = np.ones(self.dim)
+
+        self.initial_point = initial_point
         
-        self._samples = [initial_point]
+        self._samples = [initial_point] # Remove. See #324.
 
     # ------------ Abstract methods to be implemented by subclasses ------------
     
@@ -65,29 +72,16 @@ class SamplerNew(ABC):
         """ Validate the target is compatible with the sampler. Called when the target is set. Should raise an error if the target is not compatible. """
         pass
 
-    @abstractmethod
-    def get_state(self):
-        """ Return the state of the sampler. """
+    # -- _pre_sample and _pre_warmup methods: can be overridden by subclasses --
+    def _pre_sample(self):
+        """ Any code that needs to be run before sampling. """
         pass
 
-    @abstractmethod
-    def set_state(self, state):
-        """ Set the state of the sampler. """
+    def _pre_warmup(self):
+        """ Any code that needs to be run before warmup. """
         pass
 
-
     # ------------ Public attributes ------------
-
-    @property
-    def initial_point(self):
-        """ Return the initial point of the sampler. This is always the first sample. """
-        return self._samples[0]
-    
-    @initial_point.setter
-    def initial_point(self, value):
-        """ Set the initial point of the sampler. """
-        self._samples[0] = value
-    
     @property
     def dim(self):
         """ Dimension of the target density. """
@@ -109,6 +103,15 @@ class SamplerNew(ABC):
         self._target = value
         self.validate_target()
 
+    @property
+    def current_point(self):
+        """ The current point of the sampler. """
+        return self._current_point
+    
+    @current_point.setter
+    def current_point(self, value):
+        """ Set the current point of the sampler. """
+        self._current_point = value
 
     # ------------ Public methods ------------
 
@@ -156,6 +159,9 @@ class SamplerNew(ABC):
         if batch_size > 0:
             batch_handler = _BatchHandler(batch_size, sample_path)
 
+        # Any code that needs to be run before sampling
+        self._pre_sample()
+
         # Draw samples
         for _ in progressbar( range(Ns) ):
             
@@ -191,6 +197,9 @@ class SamplerNew(ABC):
 
         tune_interval = max(int(tune_freq * Nb), 1)
 
+        # Any code that needs to be run before warmup
+        self._pre_warmup()
+
         # Draw warmup samples with tuning
         for idx in progressbar(range(Nb)):
 
@@ -209,6 +218,94 @@ class SamplerNew(ABC):
             self._call_callback(self.current_point, len(self._samples)-1)
 
         return self
+    
+    def get_state(self) -> dict:
+        """ Return the state of the sampler. 
+
+        The state is used when checkpointing the sampler.
+
+        The state of the sampler is a dictionary with keys 'metadata' and 'state'.
+        The 'metadata' key contains information about the sampler type.
+        The 'state' key contains the state of the sampler.
+
+        For example, the state of a "MH" sampler could be:
+
+        state = {
+            'metadata': {
+                'sampler_type': 'MH'
+            },
+            'state': {
+                'current_point': np.array([...]),
+                'current_target_logd': -123.45,
+                'scale': 1.0,
+                ...
+            }
+        }
+        """   
+        state = {
+            'metadata': {
+                'sampler_type': self.__class__.__name__
+            },
+            'state': {
+                key: getattr(self, key) for key in self._STATE_KEYS
+            }
+        }
+        return state
+
+    def set_state(self, state: dict):
+        """ Set the state of the sampler.
+
+        The state is used when loading the sampler from a checkpoint.
+
+        The state of the sampler is a dictionary with keys 'metadata' and 'state'.
+
+        For example, the state of a "MH" sampler could be:
+
+        state = {
+            'metadata': {
+                'sampler_type': 'MH'
+            },
+            'state': {
+                'current_point': np.array([...]),
+                'current_target_logd': -123.45,
+                'scale': 1.0,
+                ...
+            }
+        }
+        """
+        if state['metadata']['sampler_type'] != self.__class__.__name__:
+            raise ValueError(f"Sampler type in state dictionary ({state['metadata']['sampler_type']}) does not match the type of the sampler ({self.__class__.__name__}).")
+        
+        for key, value in state['state'].items():
+            if key in self._STATE_KEYS:
+                setattr(self, key, value)
+            else:
+                raise ValueError(f"Key {key} not recognized in state dictionary of sampler {self.__class__.__name__}.")
+            
+    def get_history(self) -> dict:
+        """ Return the history of the sampler. """
+        history = {
+            'metadata': {
+                'sampler_type': self.__class__.__name__
+            },
+            'history': {
+                key: getattr(self, key) for key in self._HISTORY_KEYS
+            }
+        }
+        return history
+    
+    def set_history(self, history: dict):
+        """ Set the history of the sampler. """
+        if history['metadata']['sampler_type'] != self.__class__.__name__:
+            raise ValueError(f"Sampler type in history dictionary ({history['metadata']['sampler_type']}) does not match the type of the sampler ({self.__class__.__name__}).")
+        
+        for key, value in history['history'].items():
+            if key in self._HISTORY_KEYS:
+                setattr(self, key, value)
+            else:
+                raise ValueError(f"Key {key} not recognized in history dictionary of sampler {self.__class__.__name__}.")
+
+    # ------------ Private methods ------------
 
     def _call_callback(self, sample, sample_index):
         """ Calls the callback function. Assumes input is sample and sample index"""
@@ -218,6 +315,9 @@ class SamplerNew(ABC):
 
 class ProposalBasedSamplerNew(SamplerNew, ABC):
     """ Abstract base class for samplers that use a proposal distribution. """
+
+    _STATE_KEYS = SamplerNew._STATE_KEYS.union({'current_target_logd', 'scale'})
+
     def __init__(self, target, proposal=None, scale=1, **kwargs):
         """ Initializer for proposal based samplers. 
 
@@ -240,7 +340,7 @@ class ProposalBasedSamplerNew(SamplerNew, ABC):
         super().__init__(target, **kwargs)
 
         self.current_point = self.initial_point
-        self.current_target = self.target.logd(self.current_point)
+        self.current_target_logd = self.target.logd(self.current_point)
         self.proposal = proposal
         self.scale = scale
 

cuqi/sampler/_hmc.py

@@ -83,6 +83,9 @@ class NUTS(Sampler):
         self.max_depth = max_depth
         self.adapt_step_size = adapt_step_size
         self.opt_acc_rate = opt_acc_rate
+        # if this flag is True, the samples and the burn-in will be returned
+        # otherwise, the burn-in will be truncated
+        self._return_burnin = False
 
         # NUTS run diagnostic
         # number of tree nodes created each NUTS iteration
@@ -226,9 +229,10 @@ class NUTS(Sampler):
             if np.isnan(joint_eval[k]):
                 raise NameError('NaN potential func')
             
-        # apply burn-in 
-        theta = theta[:, Nb:]
-        joint_eval = joint_eval[Nb:]
+        # apply burn-in
+        if not self._return_burnin: 
+            theta = theta[:, Nb:]
+            joint_eval = joint_eval[Nb:]
         return theta, joint_eval, step_sizes
 
     #=========================================================================