CUQIpy 1.0.0.post0.dev91__py3-none-any.whl → 1.0.0.post0.dev145__py3-none-any.whl

{CUQIpy-1.0.0.post0.dev91.dist-info → CUQIpy-1.0.0.post0.dev145.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: CUQIpy
-Version: 1.0.0.post0.dev91
+Version: 1.0.0.post0.dev145
 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.dev91.dist-info → CUQIpy-1.0.0.post0.dev145.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=wjkb-AlAnGV3_a5XYx4H4xHteXak-c9qRWEs8JAhr3k,509
+cuqi/_version.py,sha256=Qc2vgfclx0EbaHAQZJpzQEPFGq1LO7Yj1wMffmN8Llo,510
 cuqi/config.py,sha256=wcYvz19wkeKW2EKCGIKJiTpWt5kdaxyt4imyRkvtTRA,526
 cuqi/diagnostics.py,sha256=5OrbJeqpynqRXOe5MtOKKhe7EAVdOEpHIqHnlMW9G_c,3029
 cuqi/array/__init__.py,sha256=-EeiaiWGNsE3twRS4dD814BIlfxEsNkTCZUc5gjOXb0,30
@@ -32,12 +32,13 @@ 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=MV1lHnRJTBsJMYEywnEx4iTEEAvN1SLIXhpRU_eiX-8,258
-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/_sampler.py,sha256=hIW_CIyv1e2jAJiHAZ8kbguZp3u24W7OnQKqcFDuzgQ,11043
+cuqi/experimental/mcmc/__init__.py,sha256=vVnohcm4EIUwbp1sr3LbB0BkXO8jyZsbiKMJmIgetYY,314
+cuqi/experimental/mcmc/_cwmh.py,sha256=yRlTk5a1QYfH3JyCecfOOTeDf-4-tmJ3Tl2Bc3pyp1Y,7336
+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=hbZwUHcEZFSSVd2tICcp9FdcK9UKB_-izdM7w4xwijs,14408
 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
@@ -74,8 +75,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.dev91.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
-CUQIpy-1.0.0.post0.dev91.dist-info/METADATA,sha256=KG8D1oXbdnsjU-xCHTLbwmzUCXkndGBNZ9-xGP6XguI,18392
-CUQIpy-1.0.0.post0.dev91.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-CUQIpy-1.0.0.post0.dev91.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
-CUQIpy-1.0.0.post0.dev91.dist-info/RECORD,,
+CUQIpy-1.0.0.post0.dev145.dist-info/LICENSE,sha256=kJWRPrtRoQoZGXyyvu50Uc91X6_0XRaVfT0YZssicys,10799
+CUQIpy-1.0.0.post0.dev145.dist-info/METADATA,sha256=REboBBeNZ3O4SrH9hNrNAYgvQPU9XP49ezl4bHyCigQ,18393
+CUQIpy-1.0.0.post0.dev145.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+CUQIpy-1.0.0.post0.dev145.dist-info/top_level.txt,sha256=AgmgMc6TKfPPqbjV0kvAoCBN334i_Lwwojc7HE3ZwD0,5
+CUQIpy-1.0.0.post0.dev145.dist-info/RECORD,,

cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2024-04-04T16:40:53+0200",
+ "date": "2024-04-15T12:29:37+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "91a1edb3f1ae9412a27489f3faa9a4a9c8ff8a4a",
- "version": "1.0.0.post0.dev91"
+ "full-revisionid": "744dccfaf90c7e50a11b7c5a4a59d5438c37525d",
+ "version": "1.0.0.post0.dev145"
 }
 '''  # END VERSION_JSON
 

cuqi/experimental/mcmc/__init__.py

@@ -4,4 +4,5 @@ from ._sampler import SamplerNew, ProposalBasedSamplerNew
 from ._langevin_algorithm import ULANew, MALANew
 from ._mh import MHNew
 from ._pcn import pCNNew
+from ._rto import LinearRTONew, RegularizedLinearRTONew
 from ._cwmh import CWMHNew

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/_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

@@ -0,0 +1,270 @@
+import scipy as sp
+from scipy.linalg.interpolative import estimate_spectral_norm
+from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
+import numpy as np
+import cuqi
+from cuqi.solver import CGLS, FISTA
+from cuqi.experimental.mcmc import SamplerNew
+from cuqi.array import CUQIarray
+
+
+class LinearRTONew(SamplerNew):
+    """
+    Linear RTO (Randomize-Then-Optimize) sampler.
+
+    Samples posterior related to the inverse problem with Gaussian likelihood and prior, and where the forward model is Linear.
+
+    Parameters
+    ------------
+    target : `cuqi.distribution.Posterior`, `cuqi.distribution.MultipleLikelihoodPosterior` or 5-dimensional tuple.
+        If target is of type cuqi.distribution.Posterior or cuqi.distribution.MultipleLikelihoodPosterior, it represents the posterior distribution.
+        If target is a 5-dimensional tuple, it assumes the following structure:
+        (data, model, L_sqrtprec, P_mean, P_sqrtrec)
+        
+        Here:
+        data: is a m-dimensional numpy array containing the measured data.
+        model: is a m by n dimensional matrix or LinearModel representing the forward model.
+        L_sqrtprec: is the squareroot of the precision matrix of the Gaussian likelihood.
+        P_mean: is the prior mean.
+        P_sqrtprec: is the squareroot of the precision matrix of the Gaussian mean.
+
+    initial_point : `np.ndarray` 
+        Initial point for the sampler. *Optional*.
+
+    maxit : int
+        Maximum number of iterations of the inner CGLS solver. *Optional*.
+
+    tol : float
+        Tolerance of the inner CGLS solver. *Optional*.
+
+    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.
+        
+    """
+    def __init__(self, target, initial_point=None, maxit=10, tol=1e-6, **kwargs):
+
+        super().__init__(target=target, initial_point=initial_point, **kwargs)
+
+        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
+
+        # Other parameters
+        self.maxit = maxit
+        self.tol = tol
+
+    @property
+    def prior(self):
+        return self.target.prior
+
+    @property
+    def likelihood(self):
+        return self.target.likelihood
+    
+    @property
+    def likelihoods(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return [self.target.likelihood]
+        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior):
+            return self.target.likelihoods
+
+    @property
+    def model(self):
+        return self.target.model     
+    
+    @property
+    def data(self):
+        return self.target.data
+
+    @SamplerNew.target.setter
+    def target(self, value):
+        """ Set the target density. Runs validation of the target. """
+        # Accept tuple of inputs and construct posterior
+        if isinstance(value, tuple) and len(value) == 5:
+            # Structure (data, model, L_sqrtprec, P_mean, P_sqrtprec)
+            data = value[0]
+            model = value[1]
+            L_sqrtprec = value[2]
+            P_mean = value[3]
+            P_sqrtprec = value[4]
+
+            # If numpy matrix convert to CUQI model
+            if isinstance(model, np.ndarray) and len(model.shape) == 2:
+                model = cuqi.model.LinearModel(model)
+
+            # Check model input
+            if not isinstance(model, cuqi.model.LinearModel):
+                raise TypeError("Model needs to be cuqi.model.LinearModel or matrix")
+
+            # Likelihood
+            L = cuqi.distribution.Gaussian(model, sqrtprec=L_sqrtprec).to_likelihood(data)
+
+            # Prior TODO: allow multiple priors stacked
+            #if isinstance(P_mean, list) and isinstance(P_sqrtprec, list):
+            #    P = cuqi.distribution.JointGaussianSqrtPrec(P_mean, P_sqrtprec)
+            #else:
+            P = cuqi.distribution.Gaussian(P_mean, sqrtprec=P_sqrtprec)
+
+            # Construct posterior
+            value = cuqi.distribution.Posterior(L, P)
+        super(LinearRTONew, type(self)).target.fset(self, value)
+        self._precompute()
+
+    def _precompute(self):
+        L1 = [likelihood.distribution.sqrtprec for likelihood in self.likelihoods]
+        L2 = self.prior.sqrtprec
+        L2mu = self.prior.sqrtprecTimesMean
+
+        # pre-computations
+        self.n = self.prior.dim
+        self.b_tild = np.hstack([L@likelihood.data for (L, likelihood) in zip(L1, self.likelihoods)]+ [L2mu]) 
+
+        callability = [callable(likelihood.model) for likelihood in self.likelihoods]
+        notcallability = [not c for c in callability]
+        if all(notcallability):
+            self.M = sp.sparse.vstack([L@likelihood.model for (L, likelihood) in zip(L1, self.likelihoods)] + [L2])
+        elif all(callability):
+            # in this case, model is a function doing forward and backward operations
+            def M(x, flag):
+                if flag == 1:
+                    out1 = [L @ likelihood.model.forward(x) for (L, likelihood) in zip(L1, self.likelihoods)]
+                    out2 = L2 @ x
+                    out  = np.hstack(out1 + [out2])
+                elif flag == 2:
+                    idx_start = 0
+                    idx_end = 0
+                    out1 = np.zeros(self.n)
+                    for likelihood in self.likelihoods:
+                        idx_end += len(likelihood.data)
+                        out1 += likelihood.model.adjoint(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end])
+                        idx_start = idx_end
+                    out2 = L2.T @ x[idx_end:]
+                    out  = out1 + out2                
+                return out   
+            self.M = M  
+        else:
+            raise TypeError("All likelihoods need to be callable or none need to be callable.")
+
+    def step(self):
+        y = self.b_tild + np.random.randn(len(self.b_tild))
+        sim = CGLS(self.M, y, self.current_point, self.maxit, self.tol)            
+        self.current_point, _ = sim.solve()
+        acc = 1
+        return acc
+
+    def tune(self, skip_len, update_count):
+        pass
+    
+    def validate_target(self):
+        # Check target type
+        if not isinstance(self.target, (cuqi.distribution.Posterior, cuqi.distribution.MultipleLikelihoodPosterior)):
+            raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior' or 'cuqi.distribution.MultipleLikelihoodPosterior'.")       
+
+        # Check Linear model and Gaussian likelihood(s)
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            if not isinstance(self.model, cuqi.model.LinearModel):
+                raise TypeError("Model needs to be linear")
+
+            if not hasattr(self.likelihood.distribution, "sqrtprec"):
+                raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
+            
+        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior): # Elif used for further alternatives, e.g., stacked posterior
+            for likelihood in self.likelihoods:
+                if not isinstance(likelihood.model, cuqi.model.LinearModel):
+                    raise TypeError("Model needs to be linear")
+
+                if not hasattr(likelihood.distribution, "sqrtprec"):
+                    raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
+        
+        # Check Gaussian prior
+        if not hasattr(self.prior, "sqrtprec"):
+            raise TypeError("prior must contain a sqrtprec attribute")
+
+        if not hasattr(self.prior, "sqrtprecTimesMean"):
+            raise TypeError("Prior must contain a sqrtprecTimesMean attribute")
+
+class RegularizedLinearRTONew(LinearRTONew):
+    """
+    Regularized Linear RTO (Randomize-Then-Optimize) sampler.
+
+    Samples posterior related to the inverse problem with Gaussian likelihood and implicit Gaussian prior, and where the forward model is Linear.
+
+    Parameters
+    ------------
+    target : `cuqi.distribution.Posterior`
+        See `cuqi.sampler.LinearRTO`
+
+    initial_point : `np.ndarray` 
+        Initial point for the sampler. *Optional*.
+
+    maxit : int
+        Maximum number of iterations of the inner FISTA solver. *Optional*.
+        
+    stepsize : string or float
+        If stepsize is a string and equals either "automatic", then the stepsize is automatically estimated based on the spectral norm.
+        If stepsize is a float, then this stepsize is used.
+
+    abstol : float
+        Absolute tolerance of the inner FISTA solver. *Optional*.
+    
+    adaptive : bool
+        If True, FISTA is used as inner solver, otherwise ISTA is used. *Optional*.
+
+    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.
+        
+    """
+    def __init__(self, target, initial_point=None, maxit=100, stepsize="automatic", abstol=1e-10, adaptive=True, **kwargs):
+        
+        super().__init__(target=target, initial_point=initial_point, **kwargs)
+
+        # Other parameters
+        self.stepsize = stepsize
+        self.abstol = abstol   
+        self.adaptive = adaptive
+        self.proximal = target.prior.proximal
+        self._stepsize = self._choose_stepsize()
+        self.maxit = maxit
+
+    @LinearRTONew.target.setter
+    def target(self, value):
+        if not callable(value.prior.proximal):
+            raise TypeError("Projector needs to be callable")
+        return super(RegularizedLinearRTONew, type(self)).target.fset(self, value)
+
+    def _choose_stepsize(self):
+        if isinstance(self.stepsize, str):
+            if self.stepsize in ["automatic"]:
+                if not callable(self.M):
+                    M_op = scipyLinearOperator(self.M.shape, matvec = lambda v: self.M@v, rmatvec = lambda w: self.M.T@w)
+                else:
+                    M_op = scipyLinearOperator((len(self.b_tild), self.n), matvec = lambda v: self.M(v,1), rmatvec = lambda w: self.M(w,2))
+                    
+                _stepsize = 0.99/(estimate_spectral_norm(M_op)**2)
+                # print(f"Estimated stepsize for regularized Linear RTO: {_stepsize}")
+            else:
+                raise ValueError("Stepsize choice not supported")
+        else:
+            _stepsize = self.stepsize
+        return _stepsize
+
+    @property
+    def prior(self):
+        return self.target.prior.gaussian
+
+    def step(self):
+        y = self.b_tild + np.random.randn(len(self.b_tild))
+        sim = FISTA(self.M, y, self.current_point, self.proximal,
+                    maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+        self.current_point, _ = sim.solve()
+        acc = 1
+        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,24 +72,7 @@ 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. """
-        pass
-
-    @abstractmethod
-    def set_state(self, state):
-        """ Set the state of the sampler. """
-        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]
-    
     @property
     def dim(self):
         """ Dimension of the target density. """
@@ -104,6 +94,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 ------------
 
@@ -204,6 +203,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"""
@@ -213,6 +300,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. 
 
@@ -235,7 +325,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