Repository: CUQI-DTU/CUQIpy

Issue:
Regularized Gaussian: TV with Gibbs and more
Roadmap for the planned features for the Regularized Gaussian distributions: (Each planned as their own PR)
- [x] 1: Add ADMM solver
- [x] 2: Add TV as option for RegularizedGaussian, see #575 
- [x] 3: Add Conjugacy for TV, see #587 
- [x] 4: Add nonnegativity constrained TV with Gibbs #612 

Add additional constraints:
- [x] 5: Add Monotonicity constraints, see #618 



Optional considerations:
- [x] Solver choice: Scipy linear least squares for nonnegativity and box constraints
- [ ] Solver choices: others?
- [ ] Write a how-to guide #616 
- [ ] Possible unification of the penalty terms, see #584

Hints from issue comments:
@chaozg: As suggested by @jeverink, this issue can be safely closed as 
- all core sub-issues are created and fixed
- optional sub-issues are being tracked separately.

Solution patch:
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 5654f5c11..0cf002c9c 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -22,6 +22,15 @@ jobs:
       - name: Sphinx build
         run: |
           sphinx-build docs ./public
+
+      # If PR-based build. Upload artifact of the documentation for download and review
+      - name: Upload artifact
+        if: ${{ github.event_name == 'pull_request' }}
+        uses: actions/upload-artifact@v3
+        with:
+          name: docs_build
+          path: ./public
+
       - name: Deploy
         uses: peaceiris/actions-gh-pages@v3
         if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
diff --git a/cuqi/experimental/mcmc/_langevin_algorithm.py b/cuqi/experimental/mcmc/_langevin_algorithm.py
index 9c1e9742c..5759a8548 100644
--- a/cuqi/experimental/mcmc/_langevin_algorithm.py
+++ b/cuqi/experimental/mcmc/_langevin_algorithm.py
@@ -3,7 +3,7 @@
 from cuqi.experimental.mcmc import Sampler
 from cuqi.implicitprior import RestorationPrior, MoreauYoshidaPrior
 from cuqi.array import CUQIarray
-from copy import deepcopy
+from copy import copy
 
 class ULA(Sampler): # Refactor to Proposal-based sampler?
     """Unadjusted Langevin algorithm (ULA) (Roberts and Tweedie, 1996)
@@ -316,11 +316,14 @@ def target(self, value):
 
     def _create_smoothed_target(self, value):
         """ Create a smoothed target using a Moreau-Yoshida envelope. """
-        copied_value = deepcopy(value)
+        copied_value = copy(value)
         if isinstance(copied_value.prior, RestorationPrior):
+            # Acceess the prior name
+            name = value.prior.name
             copied_value.prior = MoreauYoshidaPrior(
                 copied_value.prior,
-                self.smoothing_strength)
+                self.smoothing_strength,
+                name=name)
         return copied_value
 
     def validate_target(self):
diff --git a/cuqi/experimental/mcmc/_laplace_approximation.py b/cuqi/experimental/mcmc/_laplace_approximation.py
index 909ccc8b7..649b613fb 100644
--- a/cuqi/experimental/mcmc/_laplace_approximation.py
+++ b/cuqi/experimental/mcmc/_laplace_approximation.py
@@ -74,8 +74,8 @@ def model(self):
         return self.target.model     
     
     @property
-    def data(self):
-        return self.target.data
+    def _data(self):
+        return self.target.data - self.target.model._shift
 
     def _precompute(self):
 
@@ -89,7 +89,7 @@ def Lk_fun(x_k):
             return W.sqrt() @ D
         self.Lk_fun = Lk_fun
 
-        self._m = len(self.data)
+        self._m = len(self._data)
         self._L1 = self.likelihood.distribution.sqrtprec
 
         # If prior location is scalar, repeat it to match dimensions
@@ -101,17 +101,17 @@ def Lk_fun(x_k):
         # Initial Laplace approx
         self._L2 = Lk_fun(self.initial_point)
         self._L2mu = self._L2@self._priorloc
-        self._b_tild = np.hstack([self._L1@self.data, self._L2mu]) 
+        self._b_tild = np.hstack([self._L1@self._data, self._L2mu]) 
         
         # Least squares form
         def M(x, flag):
             if flag == 1:
-                out1 = self._L1 @ self.model.forward(x)
+                out1 = self._L1 @ self.model._forward_func_no_shift(x) # Use forward function which excludes shift
                 out2 = np.sqrt(1/self.prior.scale)*(self._L2 @ x)
                 out  = np.hstack([out1, out2])
             elif flag == 2:
                 idx = int(self._m)
-                out1 = self.model.adjoint(self._L1.T@x[:idx])
+                out1 = self.model._adjoint_func_no_shift(self._L1.T@x[:idx])
                 out2 = np.sqrt(1/self.prior.scale)*(self._L2.T @ x[idx:])
                 out  = out1 + out2                
             return out
@@ -121,7 +121,7 @@ def step(self):
         # Update Laplace approximation
         self._L2 = self.Lk_fun(self.current_point)
         self._L2mu = self._L2@self._priorloc
-        self._b_tild = np.hstack([self._L1@self.data, self._L2mu]) 
+        self._b_tild = np.hstack([self._L1@self._data, self._L2mu]) 
     
         # Sample from approximate posterior
         e = np.random.randn(len(self._b_tild))
@@ -139,9 +139,9 @@ def validate_target(self):
         if not isinstance(self.target, cuqi.distribution.Posterior):
             raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior'.")       
 
-        # Check Linear model
-        if not isinstance(self.likelihood.model, cuqi.model.LinearModel):
-            raise TypeError("Model needs to be linear")
+        # Check Affine model
+        if not isinstance(self.likelihood.model, cuqi.model.AffineModel):
+            raise TypeError("Model needs to be affine or linear")
 
         # Check Gaussian likelihood
         if not hasattr(self.likelihood.distribution, "sqrtprec"):
diff --git a/cuqi/experimental/mcmc/_rto.py b/cuqi/experimental/mcmc/_rto.py
index 55defd043..9858ed2f2 100644
--- a/cuqi/experimental/mcmc/_rto.py
+++ b/cuqi/experimental/mcmc/_rto.py
@@ -3,7 +3,7 @@
 from scipy.sparse.linalg import LinearOperator as scipyLinearOperator
 import numpy as np
 import cuqi
-from cuqi.solver import CGLS, FISTA
+from cuqi.solver import CGLS, FISTA, ADMM
 from cuqi.experimental.mcmc import Sampler
 
 
@@ -11,7 +11,7 @@ class LinearRTO(Sampler):
     """
     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.
+    Samples posterior related to the inverse problem with Gaussian likelihood and prior, and where the forward model is linear or more generally affine.
 
     Parameters
     ------------
@@ -22,7 +22,7 @@ class LinearRTO(Sampler):
         
         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.
+        model: is a m by n dimensional matrix, AffineModel 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.
@@ -71,12 +71,15 @@ def likelihoods(self):
 
     @property
     def model(self):
-        return self.target.model     
-    
+        return self.target.model 
+
     @property
-    def data(self):
-        return self.target.data
-    
+    def models(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return [self.target.model]
+        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior):
+            return self.target.models    
+
     def _precompute(self):
         L1 = [likelihood.distribution.sqrtprec for likelihood in self.likelihoods]
         L2 = self.prior.sqrtprec
@@ -84,8 +87,7 @@ def _precompute(self):
 
         # pre-computations
         self.n = self.prior.dim
-        self.b_tild = np.hstack([L@likelihood.data for (L, likelihood) in zip(L1, self.likelihoods)]+ [L2mu]) 
-
+        self.b_tild = np.hstack([L@(likelihood.data - model._shift) for (L, likelihood, model) in zip(L1, self.likelihoods, self.models)]+ [L2mu]) # With shift from AffineModel
         callability = [callable(likelihood.model) for likelihood in self.likelihoods]
         notcallability = [not c for c in callability]
         if all(notcallability):
@@ -94,7 +96,7 @@ def _precompute(self):
             # 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)]
+                    out1 = [L @ likelihood.model._forward_func_no_shift(x) for (L, likelihood) in zip(L1, self.likelihoods)] # Use forward function which excludes shift
                     out2 = L2 @ x
                     out  = np.hstack(out1 + [out2])
                 elif flag == 2:
@@ -103,7 +105,7 @@ def M(x, flag):
                     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])
+                        out1 += likelihood.model._adjoint_func_no_shift(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end])
                         idx_start = idx_end
                     out2 = L2.T @ x[idx_end:]
                     out  = out1 + out2                
@@ -129,16 +131,16 @@ def validate_target(self):
 
         # 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 isinstance(self.model, cuqi.model.AffineModel):
+                raise TypeError("Model needs to be linear or more generally affine")
 
             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 isinstance(likelihood.model, cuqi.model.AffineModel):
+                    raise TypeError("Model needs to be linear or more generally affine")
 
                 if not hasattr(likelihood.distribution, "sqrtprec"):
                     raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
@@ -159,6 +161,13 @@ class RegularizedLinearRTO(LinearRTO):
     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.
+    The sampler works by repeatedly solving regularized linear least squares problems for perturbed data.
+    The solver for these optimization problems is chosen based on how the regularized is provided in the implicit Gaussian prior.
+    Currently we use the following solvers:
+    FISTA: [1] Beck, Amir, and Marc Teboulle. "A fast iterative shrinkage-thresholding algorithm for linear inverse problems." SIAM journal on imaging sciences 2.1 (2009): 183-202.
+           Used when prior.proximal is callable.
+    ADMM:  [2] Boyd et al. "Distributed optimization and statistical learning via the alternating direction method of multipliers."Foundations and Trends® in Machine learning, 2011.
+           Used when prior.proximal is a list of penalty terms.
 
     Parameters
     ------------
@@ -169,12 +178,19 @@ class RegularizedLinearRTO(LinearRTO):
         Initial point for the sampler. *Optional*.
 
     maxit : int
-        Maximum number of iterations of the inner FISTA solver. *Optional*.
+        Maximum number of iterations of the inner FISTA/ADMM solver. *Optional*.
+
+    inner_max_it : int
+        Maximum number of iterations of the CGLS solver used within the ADMM 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.
 
+    penalty_parameter : int
+        Penalty parameter of the inner ADMM solver. *Optional*.
+        See [2] or `cuqi.solver.ADMM`
+
     abstol : float
         Absolute tolerance of the inner FISTA solver. *Optional*.
     
@@ -188,7 +204,7 @@ class RegularizedLinearRTO(LinearRTO):
         An example is shown in demos/demo31_callback.py.
         
     """
-    def __init__(self, target=None, initial_point=None, maxit=100, stepsize="automatic", abstol=1e-10, adaptive=True, **kwargs):
+    def __init__(self, target=None, initial_point=None, maxit=100, inner_max_it=10, stepsize="automatic", penalty_parameter=10, abstol=1e-10, adaptive=True, **kwargs):
         
         super().__init__(target=target, initial_point=initial_point, **kwargs)
 
@@ -197,10 +213,13 @@ def __init__(self, target=None, initial_point=None, maxit=100, stepsize="automat
         self.abstol = abstol   
         self.adaptive = adaptive
         self.maxit = maxit
+        self.inner_max_it = inner_max_it
+        self.penalty_parameter = penalty_parameter
 
     def _initialize(self):
         super()._initialize()
-        self._stepsize = self._choose_stepsize()
+        if self._inner_solver == "FISTA":
+            self._stepsize = self._choose_stepsize()
 
     @property
     def proximal(self):
@@ -210,8 +229,7 @@ def validate_target(self):
         super().validate_target()
         if not isinstance(self.target.prior, (cuqi.implicitprior.RegularizedGaussian, cuqi.implicitprior.RegularizedGMRF)):
             raise TypeError("Prior needs to be RegularizedGaussian or RegularizedGMRF")
-        if not callable(self.proximal):
-            raise TypeError("Proximal needs to be callable")
+        self._inner_solver = "FISTA" if callable(self.proximal) else "ADMM"
 
     def _choose_stepsize(self):
         if isinstance(self.stepsize, str):
@@ -235,8 +253,16 @@ def prior(self):
 
     def step(self):
         y = self.b_tild + np.random.randn(len(self.b_tild))
-        sim = FISTA(self.M, y, self.proximal,
-                    self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+
+        if self._inner_solver == "FISTA":
+            sim = FISTA(self.M, y, self.proximal,
+                        self.current_point, maxit = self.maxit, stepsize = self._stepsize, abstol = self.abstol, adaptive = self.adaptive)         
+        elif self._inner_solver == "ADMM":
+            sim = ADMM(self.M, y, self.proximal,
+                        self.current_point, self.penalty_parameter, maxit = self.maxit, inner_max_it = self.inner_max_it, adaptive = self.adaptive)  
+        else:
+            raise ValueError("Choice of solver not supported.")
+
         self.current_point, _ = sim.solve()
         acc = 1
         return acc
diff --git a/cuqi/experimental/mcmc/_sampler.py b/cuqi/experimental/mcmc/_sampler.py
index 6c28dd4c0..b43b19ef7 100644
--- a/cuqi/experimental/mcmc/_sampler.py
+++ b/cuqi/experimental/mcmc/_sampler.py
@@ -396,16 +396,19 @@ def __repr__(self):
         """ Return a string representation of the sampler. """
         if self.target is None:
             return f"Sampler: {self.__class__.__name__} \n Target: None"
-        self._ensure_initialized()
-        state = self.get_state()
-        msg = f" Sampler: \n\t {self.__class__.__name__} \n Target: \n \t {self.target} \n Current state: \n"
-        # Sort keys alphabetically
-        keys = sorted(state['state'].keys())
-        # Put _ in the end
-        keys = [key for key in keys if key[0] != '_'] + [key for key in keys if key[0] == '_']
-        for key in keys:
-            value = state['state'][key]
-            msg += f"\t {key}: {value} \n"
+        else:
+            msg = f"Sampler: {self.__class__.__name__} \n Target: \n \t {self.target} "
+            
+        if self._is_initialized:
+            state = self.get_state()
+            msg += f"\n Current state: \n"
+            # Sort keys alphabetically
+            keys = sorted(state['state'].keys())
+            # Put _ in the end
+            keys = [key for key in keys if key[0] != '_'] + [key for key in keys if key[0] == '_']
+            for key in keys:
+                value = state['state'][key]
+                msg += f"\t {key}: {value} \n"
         return msg
 
 class ProposalBasedSampler(Sampler, ABC):
diff --git a/cuqi/geometry/_geometry.py b/cuqi/geometry/_geometry.py
index 6361d357f..742ccf6ed 100644
--- a/cuqi/geometry/_geometry.py
+++ b/cuqi/geometry/_geometry.py
@@ -522,6 +522,9 @@ class Image2D(Geometry):
     Plotting is handled via matplotlib.pyplot.imshow.
     Colormap is defaulted to grayscale.
 
+    A How-To guide on the use of Image2D as the domain/range geometry of a CUQI 
+    :class:`Model` is available `here <https://cuqi-dtu.github.io/CUQIpy/user/_auto_howtos/myula.html>`_.
+
     Parameters
     -----------
     im_shape : tuple
diff --git a/cuqi/implicitprior/_regularizedGMRF.py b/cuqi/implicitprior/_regularizedGMRF.py
index c6096ed1e..a5c73d5d1 100644
--- a/cuqi/implicitprior/_regularizedGMRF.py
+++ b/cuqi/implicitprior/_regularizedGMRF.py
@@ -63,6 +63,7 @@ def __init__(self, mean=None, prec=None, bc_type='zero', order=1, proximal = Non
             
             # Underlying explicit GMRF
             self._gaussian = GMRF(mean, prec, bc_type=bc_type, order=order, **kwargs)
+            kwargs.pop("geometry", None)
             
             # Init from abstract distribution class
             super(Distribution, self).__init__(**kwargs)
diff --git a/cuqi/implicitprior/_regularizedGaussian.py b/cuqi/implicitprior/_regularizedGaussian.py
index b4d9d2909..f07b364cb 100644
--- a/cuqi/implicitprior/_regularizedGaussian.py
+++ b/cuqi/implicitprior/_regularizedGaussian.py
@@ -1,6 +1,8 @@
 from cuqi.utilities import get_non_default_args
 from cuqi.distribution import Distribution, Gaussian
 from cuqi.solver import ProjectNonnegative, ProjectBox, ProximalL1
+from cuqi.geometry import Continuous1D, Continuous2D, Image2D
+from cuqi.operator import FirstOrderFiniteDifference
 
 import numpy as np
 
@@ -39,17 +41,22 @@ class RegularizedGaussian(Distribution):
     sqrtprec
         See :class:`~cuqi.distribution.Gaussian` for details.
 
-    proximal : callable f(x, scale) or None
-        Euclidean proximal operator f of the regularization function g, that is, a solver for the optimization problem
-        min_z 0.5||x-z||_2^2+scale*g(x).
-
+    proximal : callable f(x, scale), list of tuples (callable proximal operator of f_i, linear operator L_i) or None
+        If callable:
+            Euclidean proximal operator f of the regularization function g, that is, a solver for the optimization problem
+            min_z 0.5||x-z||_2^2+scale*g(x).
+        If list of tuples (callable proximal operator of f_i, linear operator L_i):
+            Each callable proximal operator of f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
+            The corresponding regularization takes the form
+                sum_i f_i(L_i x),
+            where the sum ranges from 1 to an arbitrary n.
 
     projector : callable f(x) or None
         Euclidean projection onto the constraint C, that is, a solver for the optimization problem
         min_(z in C) 0.5||x-z||_2^2.
 
     constraint : string or None
-        Preset constraints. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
+        Preset constraints that generate the corresponding proximal parameter. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
         For "box", the following additional parameters can be passed:
             lower_bound : array_like or None
                 Lower bound of box, defaults to zero
@@ -57,10 +64,10 @@ class RegularizedGaussian(Distribution):
                 Upper bound of box, defaults to one
 
     regularization : string or None
-        Preset regularization. Can be set to "l1". Required for use in Gibbs in future update.
-        For "l1", the following additional parameters can be passed:
+        Preset regularization that generate the corresponding proximal parameter. Can be set to "l1" or 'tv'. Required for use in Gibbs in future update.
+        For "l1" or "tv", the following additional parameters can be passed:
             strength : scalar
-                Regularization parameter, i.e., strength*||x||_1 , defaults to one
+                Regularization parameter, i.e., strength*||Lx||_1, defaults to one
 
     """
         
@@ -75,6 +82,7 @@ def __init__(self, mean=None, cov=None, prec=None, sqrtcov=None, sqrtprec=None,
         
         # We init the underlying Gaussian first for geometry and dimensionality handling
         self._gaussian = Gaussian(mean=mean, cov=cov, prec=prec, sqrtcov=sqrtcov, sqrtprec=sqrtprec, **kwargs)
+        kwargs.pop("geometry", None)
 
         # Init from abstract distribution class
         super().__init__(**kwargs)
@@ -88,12 +96,6 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         if (proximal is not None) + (projector is not None) + (constraint is not None) + (regularization is not None) != 1:
             raise ValueError("Precisely one of proximal, projector, constraint or regularization needs to be provided.")
 
-        if proximal is not None:
-            if not callable(proximal):
-                raise ValueError("Proximal needs to be callable.")
-            if len(get_non_default_args(proximal)) != 2:
-                raise ValueError("Proximal should take 2 arguments.")
-            
         if projector is not None:
             if not callable(projector):
                 raise ValueError("Projector needs to be callable.")
@@ -104,7 +106,8 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         self._preset = None
         
         if proximal is not None:
-            self._proximal = proximal
+            # No need to generate the proximal and associated information
+            self.proximal = proximal
         elif projector is not None:
             self._proximal = lambda z, gamma: projector(z)
         elif (isinstance(constraint, str) and constraint.lower() == "nonnegativity"):
@@ -113,15 +116,48 @@ def _parse_regularization_input_arguments(self, proximal, projector, constraint,
         elif (isinstance(constraint, str) and constraint.lower() == "box"):
             lower = optional_regularization_parameters["lower_bound"]
             upper = optional_regularization_parameters["upper_bound"]
-            self._proximal = lambda z, gamma: ProjectBox(z, lower, upper)
+            self._proximal = lambda z, _: ProjectBox(z, lower, upper)
             self._preset = "box" # Not supported in Gibbs
         elif (isinstance(regularization, str) and regularization.lower() in ["l1"]):
-            strength = optional_regularization_parameters["strength"]
-            self._proximal = lambda z, gamma: ProximalL1(z, gamma*strength)
+            self._strength = optional_regularization_parameters["strength"]
+            self._proximal = lambda z, gamma: ProximalL1(z, gamma*self._strength)
             self._preset = "l1"
+        elif (isinstance(regularization, str) and regularization.lower() in ["tv"]):
+            self._strength = optional_regularization_parameters["strength"]
+            if isinstance(self.geometry, (Continuous1D, Continuous2D, Image2D)):
+                self._transformation = FirstOrderFiniteDifference(self.geometry.fun_shape, bc_type='neumann')
+            else:
+                raise ValueError("Geometry not supported for total variation")
+            
+            self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+            self._regularization_oper = self._transformation
+
+            self._proximal = [(self._regularization_prox, self._regularization_oper)]
+            self._preset = "tv"
         else:
             raise ValueError("Regularization not supported")
 
+    
+    @property
+    def transformation(self):
+        return self._transformation
+    
+    @property
+    def strength(self):
+        return self._strength
+        
+    @strength.setter
+    def strength(self, value):
+        if self._preset not in self.regularization_options():
+            raise TypeError("Strength is only used when the regularization is set to l1 or TV.")
+
+        self._strength = value
+        if self._preset == "tv":
+            self._regularization_prox = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+            self._proximal = [(self._regularization_prox, self._regularization_oper)]
+        elif self._preset == "l1":
+            self._proximal = lambda z, gamma: ProximalL1(z, gamma*self._strength)
+
     # This is a getter only attribute for the underlying Gaussian
     # It also ensures that the name of the underlying Gaussian
     # matches the name of the implicit regularized Gaussian
@@ -135,6 +171,25 @@ def gaussian(self):
     def proximal(self):
         return self._proximal
     
+    @proximal.setter
+    def proximal(self, value):
+        if callable(value):
+            if len(get_non_default_args(value)) != 2:
+                raise ValueError("Proximal should take 2 arguments.")
+        elif isinstance(value, list):
+            for (prox, op) in value:
+                if len(get_non_default_args(prox)) != 2:
+                    raise ValueError("Proximal should take 2 arguments.")
+                if op.shape[1] != self.geometry.par_dim:
+                    raise ValueError("Incorrect shape of linear operator in proximal list.")
+        else:
+            raise ValueError("Proximal needs to be callable or a list. See documentation.")
+        
+        self._proximal = value
+
+        # For all the presets, self._proximal is set directly, 
+        self._preset = None
+            
     @property
     def preset(self):
         return self._preset
@@ -154,7 +209,7 @@ def constraint_options():
 
     @staticmethod
     def regularization_options():
-        return ["l1"]
+        return ["l1", "tv"]
 
 
     # --- Defer behavior of the underlying Gaussian --- #
@@ -206,16 +261,18 @@ def sqrtcov(self):
     def sqrtcov(self, value):
         self.gaussian.sqrtcov = value     
     
-    def get_conditioning_variables(self):
-        return self.gaussian.get_conditioning_variables()
-    
     def get_mutable_variables(self):
-        return self.gaussian.get_mutable_variables()
+        mutable_vars = self.gaussian.get_mutable_variables().copy()
+        if self.preset in self.regularization_options():
+            mutable_vars += ["strength"]
+        return mutable_vars
     
     # Overwrite the condition method such that the underlying Gaussian is conditioned in general, except when conditioning on self.name
     # which means we convert Distribution to Likelihood or EvaluatedDensity.
     def _condition(self, *args, **kwargs):
-
+        if self.preset in self.regularization_options():
+            return super()._condition(*args, **kwargs)
+        
         # Handle positional arguments (similar code as in Distribution._condition)
         cond_vars = self.get_conditioning_variables()
         kwargs = self._parse_args_add_to_kwargs(cond_vars, *args, **kwargs)
@@ -275,7 +332,7 @@ class ConstrainedGaussian(RegularizedGaussian):
         min_(z in C) 0.5||x-z||_2^2.
 
     constraint : string or None
-        Preset constraints. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
+        Preset constraints that generate the corresponding proximal parameter. Can be set to "nonnegativity" and "box". Required for use in Gibbs.
         For "box", the following additional parameters can be passed:
             lower_bound : array_like or None
                 Lower bound of box, defaults to zero
diff --git a/cuqi/model/__init__.py b/cuqi/model/__init__.py
index 030c393e9..48fd02670 100644
--- a/cuqi/model/__init__.py
+++ b/cuqi/model/__init__.py
@@ -1 +1 @@
-from ._model import Model, LinearModel, PDEModel
+from ._model import Model, LinearModel, PDEModel, AffineModel
diff --git a/cuqi/model/_model.py b/cuqi/model/_model.py
index 365b72f1f..0b9abf265 100644
--- a/cuqi/model/_model.py
+++ b/cuqi/model/_model.py
@@ -469,8 +469,126 @@ def __len__(self):
 
     def __repr__(self) -> str:
         return "CUQI {}: {} -> {}.\n    Forward parameters: {}.".format(self.__class__.__name__,self.domain_geometry,self.range_geometry,cuqi.utilities.get_non_default_args(self))
-    
-class LinearModel(Model):
+
+
+class AffineModel(Model):
+    """ Model class representing an affine model, i.e. a linear operator with a fixed shift. For linear models, represented by a linear operator only, see :class:`~cuqi.model.LinearModel`.
+
+    The affine model is defined as:
+
+    .. math::
+
+        x \\mapsto Ax + shift
+
+    where :math:`A` is the linear operator and :math:`shift` is the shift.
+
+    Parameters
+    ----------
+
+    linear_operator : 2d ndarray, callable function or cuqi.model.LinearModel
+        The linear operator. If ndarray is given, the operator is assumed to be a matrix.
+
+    shift : scalar or array_like
+        The shift to be added to the forward operator.
+
+    linear_operator_adjoint : callable function, optional
+        The adjoint of the linear operator. Also used for computing gradients.
+
+    range_geometry : cuqi.geometry.Geometry
+        The geometry representing the range.
+
+    domain_geometry : cuqi.geometry.Geometry
+        The geometry representing the domain.
+
+    """
+
+    def __init__(self, linear_operator, shift, linear_operator_adjoint=None, range_geometry=None, domain_geometry=None):
+
+        # If input represents a matrix, extract needed properties from it
+        if hasattr(linear_operator, '__matmul__') and hasattr(linear_operator, 'T'):
+            if linear_operator_adjoint is not None:
+                raise ValueError("Adjoint of linear operator should not be provided when linear operator is a matrix. If you want to provide an adjoint, use a callable function for the linear operator.")
+            
+            matrix = linear_operator
+
+            linear_operator = lambda x: matrix@x
+            linear_operator_adjoint = lambda y: matrix.T@y
+
+            if range_geometry is None:
+                if hasattr(matrix, 'shape'):
+                    range_geometry = _DefaultGeometry1D(grid=matrix.shape[0])
+                elif isinstance(matrix, LinearModel):
+                    range_geometry = matrix.range_geometry
+
+            if domain_geometry is None:
+                if hasattr(matrix, 'shape'):
+                    domain_geometry = _DefaultGeometry1D(grid=matrix.shape[1])
+                elif isinstance(matrix, LinearModel):
+                    domain_geometry = matrix.domain_geometry
+        else:
+            matrix = None
+
+        # Ensure that the operators are a callable functions (either provided or created from matrix)
+        if not callable(linear_operator):
+            raise TypeError("Linear operator must be defined as a matrix or a callable function of some kind")
+        if linear_operator_adjoint is not None and not callable(linear_operator_adjoint):
+            raise TypeError("Linear operator adjoint must be defined as a callable function of some kind")
+
+        # Check size of shift and match against range_geometry
+        if not np.isscalar(shift):
+            if len(shift) != range_geometry.par_dim:
+                raise ValueError("The shift should have the same dimension as the range geometry.")
+
+        # Initialize Model class
+        super().__init__(linear_operator, range_geometry, domain_geometry)
+
+        # Store matrix privately
+        self._matrix = matrix
+
+        # Store shift as private attribute
+        self._shift = shift
+
+        # Store linear operator privately
+        self._linear_operator = linear_operator
+
+        # Store adjoint function
+        self._linear_operator_adjoint = linear_operator_adjoint
+
+        # Define gradient
+        self._gradient_func = lambda direction, wrt: linear_operator_adjoint(direction)
+
+        # Update forward function to include shift (overwriting the one from Model class)
+        self._forward_func = lambda *args, **kwargs: linear_operator(*args, **kwargs) + shift
+
+        # Use arguments from user's callable linear operator (overwriting those found by Model class)
+        self._non_default_args = cuqi.utilities.get_non_default_args(linear_operator)
+
+    @property
+    def shift(self):
+        """ The shift of the affine model. """
+        return self._shift
+
+    @shift.setter
+    def shift(self, value):
+        """ Update the shift of the affine model. Updates both the shift value and the underlying forward function. """
+        self._shift = value
+        self._forward_func = lambda *args, **kwargs: self._linear_operator(*args, **kwargs) + value
+
+    def _forward_func_no_shift(self, x, is_par=True):
+        """ Helper function for computing the forward operator without the shift. """
+        return self._apply_func(self._linear_operator,
+                self.range_geometry,
+                self.domain_geometry,
+                x, is_par)
+
+    def _adjoint_func_no_shift(self, y, is_par=True):
+        """ Helper function for computing the adjoint operator without the shift. """
+        return self._apply_func(self._linear_operator_adjoint,
+                self.domain_geometry,
+                self.range_geometry,
+                y, is_par)
+
+class LinearModel(AffineModel):
     """Model based on a Linear forward operator.
 
     Parameters
@@ -534,45 +652,11 @@ def adjoint(y):
     Note that you would need to specify the range and domain geometries in this
     case as they cannot be inferred from the forward and adjoint functions.
     """
-    # Linear forward model with forward and adjoint (transpose).
     
-    def __init__(self,forward,adjoint=None,range_geometry=None,domain_geometry=None):
-        #Assume forward is matrix if not callable (TODO: add more checks)
-        if not callable(forward):      
-            forward_func = lambda x: self._matrix@x
-            adjoint_func = lambda y: self._matrix.T@y
-            matrix = forward
-        else:
-            forward_func = forward
-            adjoint_func = adjoint
-            matrix = None
-
-        #Check if input is callable
-        if callable(adjoint_func) is not True:
-            raise TypeError("Adjoint needs to be callable function of some kind")
-
-        # Use matrix to derive range_geometry and domain_geometry
-        if matrix is not None:
-            if range_geometry is None:
-                range_geometry = _DefaultGeometry1D(grid=matrix.shape[0])
-            if domain_geometry is None:
-                domain_geometry = _DefaultGeometry1D(grid=matrix.shape[1])  
-
-        #Initialize Model class
-        super().__init__(forward_func,range_geometry,domain_geometry)
-
-        #Add adjoint
-        self._adjoint_func = adjoint_func
-
-        #Store matrix privately
-        self._matrix = matrix
-
-        #Add gradient
-        self._gradient_func = lambda direction, wrt: self._adjoint_func(direction)
+    def __init__(self, forward, adjoint=None, range_geometry=None, domain_geometry=None):
 
-        # if matrix is not None: 
-        #     assert(self.range_dim  == matrix.shape[0]), "The parameter 'forward' dimensions are inconsistent with the parameter 'range_geometry'"
-        #     assert(self.domain_dim == matrix.shape[1]), "The parameter 'forward' dimensions are inconsistent with parameter 'domain_geometry'"
+        #Initialize as AffineModel with shift=0
+        super().__init__(forward, 0, adjoint, range_geometry, domain_geometry)
 
     def adjoint(self, y, is_par=True):
         """ Adjoint of the model.
@@ -590,16 +674,21 @@ def adjoint(self, y, is_par=True):
         ndarray or cuqi.array.CUQIarray
             The adjoint model output. Always returned as parameters.
         """
-        return self._apply_func(self._adjoint_func,
+        if self._linear_operator_adjoint is None:
+            raise ValueError("No adjoint operator was provided for this model.")
+        return self._apply_func(self._linear_operator_adjoint,
                                 self.domain_geometry,
                                 self.range_geometry,
                                 y, is_par)
 
-
+    def __matmul__(self, x):
+        return self.forward(x)
+        
     def get_matrix(self):
         """
         Returns an ndarray with the matrix representing the forward operator.
         """
+
         if self._matrix is not None: #Matrix exists so return it
             return self._matrix
         else:
@@ -617,15 +706,12 @@ def get_matrix(self):
             #Store matrix for future use
             self._matrix = mat
 
-            return self._matrix
-
-    def __matmul__(self, x):
-        return self.forward(x)
+            return self._matrix   
 
     @property
     def T(self):
         """Transpose of linear model. Returns a new linear model acting as the transpose."""
-        transpose = LinearModel(self.adjoint,self.forward,self.domain_geometry,self.range_geometry)
+        transpose = LinearModel(self.adjoint, self.forward, self.domain_geometry, self.range_geometry)
         if self._matrix is not None:
             transpose._matrix = self._matrix.T
         return transpose
diff --git a/cuqi/sampler/_laplace_approximation.py b/cuqi/sampler/_laplace_approximation.py
index 5495bc802..031aa0430 100644
--- a/cuqi/sampler/_laplace_approximation.py
+++ b/cuqi/sampler/_laplace_approximation.py
@@ -60,9 +60,9 @@ def __init__(self, target, x0=None, maxit=50, tol=1e-4, beta=1e-5, rng=None, **k
         if not isinstance(self.target, cuqi.distribution.Posterior):
             raise ValueError(f"To initialize an object of type {self.__class__}, 'target' need to be of type 'cuqi.distribution.Posterior'.")       
 
-        # Check Linear model
-        if not isinstance(self.target.likelihood.model, cuqi.model.LinearModel):
-            raise TypeError("Model needs to be linear")
+        # Check Affine model
+        if not isinstance(self.target.likelihood.model, cuqi.model.AffineModel):
+            raise TypeError("Model needs to be affine or linear")
 
         # Check Gaussian likelihood
         if not hasattr(self.target.likelihood.distribution, "sqrtprec"):
@@ -126,7 +126,7 @@ def Lk_fun(x_k):
 
         # Pre-computations
         self._model = self.target.likelihood.model   
-        self._data = self.target.likelihood.data
+        self._data = self.target.likelihood.data - self.target.model._shift
         self._m = len(self._data)
         self._L1 = self.target.likelihood.distribution.sqrtprec
 
@@ -146,12 +146,12 @@ def Lk_fun(x_k):
         # Least squares form
         def M(x, flag):
             if flag == 1:
-                out1 = self._L1 @ self._model.forward(x)
+                out1 = self._L1 @ self._model._forward_func_no_shift(x) # Use forward function which excludes shift
                 out2 = np.sqrt(1/self.target.prior.scale)*(self._L2 @ x)
                 out  = np.hstack([out1, out2])
             elif flag == 2:
                 idx = int(self._m)
-                out1 = self._model.adjoint(self._L1.T@x[:idx])
+                out1 = self._model._adjoint_func_no_shift(self._L1.T@x[:idx])
                 out2 = np.sqrt(1/self.target.prior.scale)*(self._L2.T @ x[idx:])
                 out  = out1 + out2                
             return out 
diff --git a/cuqi/sampler/_rto.py b/cuqi/sampler/_rto.py
index 6cb549c8b..343cdde31 100644
--- a/cuqi/sampler/_rto.py
+++ b/cuqi/sampler/_rto.py
@@ -11,7 +11,7 @@ class LinearRTO(Sampler):
     """
     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.
+    Samples posterior related to the inverse problem with Gaussian likelihood and prior, and where the forward model is linear or more generally affine.
 
     Parameters
     ------------
@@ -22,7 +22,7 @@ class LinearRTO(Sampler):
         
         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.
+        model: is a m by n dimensional matrix, AffineModel 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.
@@ -59,8 +59,8 @@ def __init__(self, target, x0=None, maxit=10, tol=1e-6, shift=0, **kwargs):
                 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")
+            if not isinstance(model, cuqi.model.AffineModel):
+                raise TypeError("Model needs to be cuqi.model.AffineModel or matrix")
 
             # Likelihood
             L = cuqi.distribution.Gaussian(model, sqrtprec=L_sqrtprec).to_likelihood(data)
@@ -95,7 +95,7 @@ def __init__(self, target, x0=None, maxit=10, tol=1e-6, shift=0, **kwargs):
 
         # pre-computations
         self.n = len(self.x0)
-        self.b_tild = np.hstack([L@likelihood.data for (L, likelihood) in zip(L1, self.likelihoods)]+ [L2mu]) 
+        self.b_tild = np.hstack([L@(likelihood.data - model._shift) for (L, likelihood, model) in zip(L1, self.likelihoods, self.models)]+ [L2mu])
 
         callability = [callable(likelihood.model) for likelihood in self.likelihoods]
         notcallability = [not c for c in callability]
@@ -105,7 +105,7 @@ def __init__(self, target, x0=None, maxit=10, tol=1e-6, shift=0, **kwargs):
             # 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)]
+                    out1 = [L @ likelihood.model._forward_func_no_shift(x) for (L, likelihood) in zip(L1, self.likelihoods)] # Use forward function which excludes shift
                     out2 = L2 @ x
                     out  = np.hstack(out1 + [out2])
                 elif flag == 2:
@@ -114,7 +114,7 @@ def M(x, flag):
                     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])
+                        out1 += likelihood.model._adjoint_func_no_shift(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end]) # Use adjoint function which excludes shift
                         idx_start = idx_end
                     out2 = L2.T @ x[idx_end:]
                     out  = out1 + out2                
@@ -143,8 +143,11 @@ def model(self):
         return self.target.model     
     
     @property
-    def data(self):
-        return self.target.data
+    def models(self):
+        if isinstance(self.target, cuqi.distribution.Posterior):
+            return [self.target.model]
+        elif isinstance(self.target, cuqi.distribution.MultipleLikelihoodPosterior):
+            return self.target.models    
 
     def _sample(self, N, Nb):   
         Ns = N+Nb   # number of simulations        
@@ -175,8 +178,8 @@ def _check_posterior(self):
 
         # 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 isinstance(self.model, cuqi.model.AffineModel):
+                raise TypeError("Model needs to be linear or affine")
 
             if not hasattr(self.likelihood.distribution, "sqrtprec"):
                 raise TypeError("Distribution in Likelihood must contain a sqrtprec attribute")
diff --git a/cuqi/solver/_solver.py b/cuqi/solver/_solver.py
index 6fd66f6aa..98a52a389 100644
--- a/cuqi/solver/_solver.py
+++ b/cuqi/solver/_solver.py
@@ -669,7 +669,7 @@ class ADMM(object):
         - flag=2 indicates multiplication of the transpose of A with vector x, that is  A.T @ x.
     b : ndarray.
     penalty_terms : List of tuples (callable proximal operator of f_i, linear operator L_i)
-        Each callable proximal operator f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
+        Each callable proximal operator of f_i accepts two arguments (x, p) and should return the minimizer of p/2||x-z||^2 + f(x) over z for some f.
     x0 : ndarray. Initial guess.
     penalty_parameter : Trade-off between linear least squares and regularization term in the solver iterates. Denoted as "rho" in [1].
     maxit : The maximum number of iterations.
diff --git a/demos/howtos/Image2D.py b/demos/howtos/Image2D.py
new file mode 100644
index 000000000..0a1ac893a
--- /dev/null
+++ b/demos/howtos/Image2D.py
@@ -0,0 +1,70 @@
+"""
+The use of Image2D
+==================
+
+This script demonstrates the use of :class:`~cuqi.geometry.Image2D` for domain and range geometries in CUQIpy forward models. The settings of this geometry, in particular the so-called `visual_only` flag, depend on how the forward model input is operated on in the forward model. `visual_only` flag tells the geometry to interpret the variable as a 2D image only when visualizing the variable but keep it as a vector when being operated on in the forward model. Generally, if the forward operator expects an image input, set `visual_only` to `False`, which is the default in CUQIpy. Otherwise, if it expects a vector input, set `visual_only` to `True` in the domain geometry. In the same way, the range geometry `visual_only` setting should be consistent with your forward operator output.
+"""
+
+# %%
+# Import necessary libraries
+import cuqi
+import numpy as np
+import matplotlib.pyplot as plt
+
+# %%
+# Define dimensions
+dim_x, dim_y = 2, 4
+
+# %%
+# Example illustrating the default behavior: `visual_only=False`
+# With `visual_only=False`, the underlying structure of a CUQI array or sample will be reshaped to the
+# shape of the specified geometry, i.e., an image of size (dim_x, dim_y), before being passed to the model.
+# Here we intend to define a forward operator that downsamples the input image by a factor of 2. It's easier
+# to define such a forward operator in terms of an image. We do not need to explicitly set `visual_only` to 
+#`False` because it is the default behavior.
+
+def forward_func_a(image):
+    assert image.shape == (dim_x, dim_y)
+    return image[::2, ::2]
+
+model_A = cuqi.model.LinearModel(forward_func_a,
+                                 domain_geometry=cuqi.geometry.Image2D((dim_x, dim_y)),
+                                 range_geometry=cuqi.geometry.Image2D((dim_x // 2, dim_y // 2)))
+
+print(model_A)
+
+x_sample = cuqi.array.CUQIarray(np.linspace(0, 1, dim_x * dim_y),
+                                geometry=cuqi.geometry.Image2D((dim_x, dim_y)))
+
+# Plot the original sample
+plt.figure()
+x_sample.plot(cmap='Greens')
+plt.title("input image x")
+
+# Apply the model and plot the result
+plt.figure()
+y_sample = model_A @ x_sample
+y_sample.plot(cmap='Greens')
+plt.title("output y=model_A(x)")
+
+# %%
+# Example with `visual_only=True`
+# With `visual_only=True`, the underlying structure of a cuqi array or sample will not be changed before
+# being passed to the model. This is useful when the model expects a vector input. Here we define a 
+# forward operator that reverses the input vector, which can be easily defined in terms of a vector.
+def forward_func_b(x):
+    assert x.shape == (dim_x * dim_y, )
+    return x[::-1]
+
+model_B = cuqi.model.LinearModel(forward_func_b,
+                                 domain_geometry=cuqi.geometry.Image2D((dim_x, dim_y), visual_only=True),
+                                 range_geometry=cuqi.geometry.Image2D((dim_x, dim_y), visual_only=True))
+
+# Note we explicitely set the `visual_only` flag to be `True` here.
+print(model_B)
+
+# Apply the model and plot the result
+plt.figure()
+y_sample = model_B @ x_sample
+y_sample.plot(cmap='Greens')
+plt.title("output y=model_B(x)")
\ No newline at end of file
diff --git a/demos/try17_AffineModel.py b/demos/try17_AffineModel.py
new file mode 100644
index 000000000..2e7233ec4
--- /dev/null
+++ b/demos/try17_AffineModel.py
@@ -0,0 +1,134 @@
+# %%
+import numpy as np
+import matplotlib.pyplot as plt
+import sys
+
+from cuqi.model import AffineModel, LinearModel
+from cuqi.experimental.mcmc import MH, LinearRTO
+from cuqi.distribution import Gaussian, LMRF, GMRF, JointDistribution
+from cuqi.geometry import Discrete, Continuous1D
+from cuqi.array import CUQIarray
+
+# %% Test problem
+### y = Ax + b
+### Infer x, given the model A and b, and noisy observations of y
+### There are two ways to solve this problem: 
+### 1) either you use the linear model only y=Ax realizing that you can subtract the bias from the data or
+### 2) define an affine model and let CUQIpy handle it.
+### This script demonstrates that both methods provide the same results.
+
+# Dimensions
+n = 2
+m = 5
+
+# Unknown parameter
+x_truevec = np.array([3, 2])
+x_true = CUQIarray(x_truevec, geometry=Discrete(["x1", "x2"]))
+
+# Observation times
+t = np.linspace(0, m, m, endpoint=False)
+
+# Forward model
+Amat = np.ones((m, n))
+Amat[:, 1] = t
+
+# shift/bias
+b = np.ones(m) * 3
+
+# Geometries
+domain_geometry = Discrete(n)
+range_geometry = Continuous1D(t)
+
+# Prior
+x_mean = np.array([2.0, 1.0])
+x_std = np.ones(n)
+x = Gaussian(x_mean, sqrtcov=x_std, geometry=domain_geometry)
+
+# Noisy data
+np.random.seed(10)
+y_std = np.ones(m)
+y_obs = Gaussian(Amat @ x_truevec + b, sqrtcov=y_std, geometry=range_geometry).sample()
+
+# %% Linear model, include bias (b) from model in data term such that data = y-b
+
+# CUQI Model
+A = LinearModel(Amat, domain_geometry=domain_geometry, range_geometry=range_geometry)
+
+# Likelihood
+y = Gaussian(A @ x, sqrtcov=y_std, geometry=range_geometry)
+
+# Posterior
+posterior = JointDistribution(x, y)(y=y_obs - b)
+
+# Sample posterior with linear RTO
+np.random.seed(1000000)
+Linearsampler = MH(posterior, initial_point=x_mean)  # Init point must be float (NOT AFFINE MODEL SPECIFIC)
+Linearsampler.warmup(200)
+Linearsampler.sample(10000)
+samplesLinear = Linearsampler.get_samples()
+samplesLinear_burnin = samplesLinear.burnthin(200)
+
+# Compute posterior mean and variance and plot
+meanLinear = samplesLinear_burnin.mean()
+varLinear = samplesLinear_burnin.variance()
+
+print("Posterior mean and variance using Linear model: ")
+print("mean: {}".format(meanLinear))
+print("cov: {}".format(varLinear))
+
+fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5, 4))
+ax.plot(t, A(samplesLinear_burnin).mean(), linestyle="--", label="Sample mean using Linear model")
+ax.plot(t, y_obs - b, "o", label="Noisy data")
+ax.set_xlabel("t")
+ax.set_ylabel("y")
+plt.legend()
+
+cs = samplesLinear_burnin.plot_pair(marginals=True)
+plt.suptitle("Posterior samples using Linear model")
+
+
+# %% Affine model, include bias (b) in affine model such that data = y
+
+# CUQI Model
+affine_model = AffineModel(Amat, b, domain_geometry=domain_geometry, range_geometry=range_geometry)
+
+# Likelihood
+y = Gaussian(affine_model(x), sqrtcov=y_std, geometry=range_geometry)
+
+# Posterior
+posterior = JointDistribution(x, y)(y=y_obs)
+
+print(posterior.likelihood.model.forward(np.array([1, 1])))
+print(posterior.likelihood.model._forward_func_no_shift(np.array([1, 1])))
+
+# Sample posterior with MH
+np.random.seed(1000000)
+Affinesampler = MH(posterior, initial_point=x_mean)  # Init point must be float (NOT AFFINE MODEL SPECIFIC)
+Affinesampler.warmup(200)
+Affinesampler.sample(10000)
+samplesAffine = Affinesampler.get_samples()
+samplesAffine_burnin = samplesAffine.burnthin(200)
+
+# Compute posterior mean and variance and plot
+meanAffine = samplesAffine_burnin.mean()
+varAffine = samplesAffine_burnin.variance()
+
+print("Posterior mean and variance from Affine model: ")
+print("mean: {}".format(meanAffine))
+print("cov: {}".format(varAffine))
+
+fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(5, 4))
+ax.plot(t, affine_model(samplesAffine_burnin).mean(), linestyle="--", label="Sample mean using Affine model")
+ax.plot(t, y_obs, "o", label="Noisy data")
+ax.set_xlabel("t")
+ax.set_ylabel("y")
+plt.legend()
+
+
+cs = samplesAffine_burnin.plot_pair(marginals=True)
+plt.suptitle("Posterior samples using Affine model")
+
+# %%
+# Compare the results from the two models
+# They should provide exactly the same results as the seed is the same
+assert np.allclose(samplesLinear.samples, samplesAffine.samples)
diff --git a/test.ipynb b/test.ipynb
deleted file mode 100644
index cde176af6..000000000
--- a/test.ipynb
+++ /dev/null
@@ -1,153 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[-3.99513417e-03 -1.32339656e-01 -4.52822633e-02 -7.44973888e-02\n",
-      " -3.35005208e-11]\n"
-     ]
-    }
-   ],
-   "source": [
-    "import numpy as np\n",
-    "import scipy as sp\n",
-    "\n",
-    "from cuqi.solver import CGLS, LM, FISTA, ADMM, ProximalL1, ProjectNonnegative\n",
-    "from scipy.optimize import lsq_linear\n",
-    "\n",
-    "\n",
-    "def test_ADMM_matrix_form():\n",
-    "    # Parameters\n",
-    "    rng = np.random.default_rng(seed = 42)\n",
-    "    m, n = 10, 5\n",
-    "    A = rng.standard_normal((m, n))\n",
-    "    b = rng.standard_normal(m)\n",
-    "    \n",
-    "    k = 4\n",
-    "    L = rng.standard_normal((k, n))\n",
-    "\n",
-    "    x0 = np.zeros(n)\n",
-    "    sol, _ = ADMM(A, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],\n",
-    "                   x0, 10, maxit = 100, adaptive = True).solve()\n",
-    "\n",
-    "    print(sol)\n",
-    "    ref_sol = np.array([-3.99513417e-03, -1.32339656e-01, -4.52822633e-02, -7.44973888e-02, -3.35005208e-11])\n",
-    "    # Compare\n",
-    "    assert np.allclose(sol, ref_sol, atol=1e-4)\n",
-    "\n",
-    "test_ADMM_matrix_form()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[-3.99513417e-03 -1.32339656e-01 -4.52822633e-02 -7.44973888e-02\n",
-      " -3.35005152e-11]\n"
-     ]
-    }
-   ],
-   "source": [
-    "\n",
-    "\n",
-    "def test_ADMM_function_form():\n",
-    "    # Parameters\n",
-    "    rng = np.random.default_rng(seed = 42)\n",
-    "    m, n = 10, 5\n",
-    "    A = rng.standard_normal((m, n))\n",
-    "    def A_fun(x, flag):\n",
-    "        if flag == 1:\n",
-    "            return A@x\n",
-    "        if flag == 2:\n",
-    "            return A.T@x\n",
-    "        \n",
-    "    b = rng.standard_normal(m)\n",
-    "    \n",
-    "    k = 4\n",
-    "    L = rng.standard_normal((k, n))\n",
-    "\n",
-    "    x0 = np.zeros(n)\n",
-    "    sol, _ = ADMM(A_fun, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],\n",
-    "                   x0, 10, maxit = 100, adaptive = True).solve()\n",
-    "\n",
-    "    print(sol)\n",
-    "    ref_sol = np.array([-3.99513417e-03, -1.32339656e-01, -4.52822633e-02, -7.44973888e-02, -3.35005208e-11])\n",
-    "    # Compare\n",
-    "    assert np.allclose(sol, ref_sol, atol=1e-4)\n",
-    "\n",
-    "test_ADMM_function_form()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[-3.99513417e-03 -1.32339656e-01 -4.52822633e-02 -7.44973888e-02\n",
-      " -3.35005208e-11]\n"
-     ]
-    }
-   ],
-   "source": [
-    "\n",
-    "# Parameters\n",
-    "rng = np.random.default_rng(seed = 42)\n",
-    "m, n = 10, 5\n",
-    "A = rng.standard_normal((m, n))\n",
-    "b = rng.standard_normal(m)\n",
-    "    \n",
-    "k = 4\n",
-    "L = rng.standard_normal((k, n))\n",
-    "\n",
-    "x0 = np.zeros(n)\n",
-    "sol, _ = ADMM(A, b, [(ProximalL1, np.eye(n)), (lambda z, _ : ProjectNonnegative(z), L)],\n",
-    "                   x0, 10, maxit = 100, adaptive = False).solve()\n",
-    "\n",
-    "print(sol)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "base",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.9"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}


Test patch:
diff --git a/tests/test_implicit_priors.py b/tests/test_implicit_priors.py
index c4fd363c4..ad617a3c1 100644
--- a/tests/test_implicit_priors.py
+++ b/tests/test_implicit_priors.py
@@ -16,7 +16,7 @@ def test_RegularizedGaussian_guarding_statements():
         cuqi.implicitprior.RegularizedGaussian(np.zeros(5), 1, proximal=lambda s,z: s, constraint="nonnegativity")
 
     # Proximal
-    with pytest.raises(ValueError, match="Proximal needs to be callable"):
+    with pytest.raises(ValueError, match="Proximal needs to be callable or a list. See documentation."):
         cuqi.implicitprior.RegularizedGaussian(np.zeros(5), 1, proximal=1)
 
     with pytest.raises(ValueError, match="Proximal should take 2 arguments"):
@@ -104,3 +104,37 @@ def test_RegularizedUnboundedUniform_is_RegularizedGaussian():
     x = cuqi.implicitprior.RegularizedUnboundedUniform(cuqi.geometry.Continuous1D(5), regularization="l1", strength = 5.0)
     
     assert np.allclose(x.gaussian.sqrtprec, 0.0)
+
+def test_RegularizedGaussian_conditioning_constrained():
+    """ Test that conditioning the implicit regularized Gaussian works as expected """
+    
+    x  = cuqi.implicitprior.RegularizedGMRF(lambda a:a*np.ones(2**2),
+                     prec = lambda b:5*b,
+                     constraint = "nonnegativity",
+                     geometry = cuqi.geometry.Image2D((2,2)))
+    
+    assert x.get_mutable_variables() == ['mean', 'prec']
+    assert x.get_conditioning_variables() == ['a', 'b']
+
+    x = x(a=1, b=2)
+
+    assert np.allclose(x.mean, [1, 1, 1, 1])
+    assert np.allclose(x.prec, 10)
+
+def test_RegularizedGaussian_conditioning_strength():
+    """ Test that conditioning the implicit regularized Gaussian works as expected """
+    
+    x  = cuqi.implicitprior.RegularizedGMRF(lambda a:a*np.ones(2**2),
+                     prec = lambda b:5*b,
+                     regularization = "tv",
+                     strength = lambda c:c*2,
+                     geometry = cuqi.geometry.Image2D((2,2)))
+    
+    assert x.get_mutable_variables() == ['mean', 'prec', 'strength']
+    assert x.get_conditioning_variables() == ['a', 'b', 'c']
+
+    x = x(a=1, b=2, c=3)
+
+    assert np.allclose(x.mean, [1, 1, 1, 1])
+    assert np.allclose(x.prec, 10)
+    assert np.allclose(x.strength, 6)
diff --git a/tests/test_model.py b/tests/test_model.py
index 84377a8f4..0bb44beb4 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -5,6 +5,7 @@
 import cuqi
 import pytest
 from scipy import optimize
+from copy import copy
 
 
 @pytest.mark.parametrize("seed",[(0),(1),(2)])
@@ -471,3 +472,45 @@ def gradient(dir, wrt):
 
     assert np.allclose(model_grad.gradient(dir, wrt), model_jac.gradient(dir, wrt))
 
+# Parametrize over models
+@pytest.mark.parametrize("model", [cuqi.testproblem.Deconvolution1D().model,
+                                   cuqi.testproblem.Heat1D().model])
+def test_AffineModel_Correct_result(model):
+    """ Test creating a shifted linear model from a linear model """
+
+    # Random vectors
+    x = np.random.randn(model.domain_dim)
+    b = np.random.randn(model.range_dim)
+
+    A_affine = cuqi.model.AffineModel(model, b, range_geometry = model.range_geometry, domain_geometry = model.domain_geometry)
+
+    # Dimension check
+    assert A_affine.range_dim == model.range_dim
+
+    # Check that the shifted linear model is correct
+    assert np.allclose(A_affine(x), model(x) + b)
+
+def test_AffineModel_update_shift():
+
+    A = np.eye(2)
+    b = np.array([1, 2])
+    x = np.array([1, 1])
+    new_shift = np.array([2,-1])
+    model = cuqi.model.AffineModel(A, b)
+    model_copy = copy(model)
+
+    # check model output
+    assert np.all(model(x) == np.array([2,3]))
+
+    # check model output with updated shift
+    model.shift = new_shift
+    assert np.all(model(x) == np.array([3,0]))
+
+    # check model output of copied model
+    assert np.all(model_copy(x) == np.array([2,3]))
+
+    # check model output of copied model with updated shift
+    model_copy.shift = new_shift
+    assert np.all(model_copy(x) == np.array([3,0]))
+
+
diff --git a/tests/test_sampler.py b/tests/test_sampler.py
index 70bc5eb86..5126638cc 100644
--- a/tests/test_sampler.py
+++ b/tests/test_sampler.py
@@ -510,3 +510,115 @@ def test_Gibbs_geometry_matches():
     assert samples["d"].geometry == target.get_density("d").geometry
     assert samples["l"].geometry == target.get_density("l").geometry
     assert samples["x"].geometry == target.get_density("x").geometry
+
+def test_RTO_with_AffineModel_is_equivalent_to_LinearModel_and_shifted_data():
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Shift
+    shift = np.random.rand(A.domain_dim)
+
+    # Define Bayesian Problem
+    x = cuqi.distribution.GMRF(np.zeros(A.domain_dim), 100)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs-shift)
+
+    # Set up LinearRTO with both models
+    sampler_linear = cuqi.sampler.LinearRTO(posterior)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_linear = sampler_linear.sample(10, 2)
+
+    # Define AffineModel
+    affine_model = cuqi.model.AffineModel(A, shift)
+
+    # Set up LinearRTO with AffineModel
+    y = cuqi.distribution.Gaussian(affine_model, 0.01**2)
+    posterior_affine = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up LinearRTO with AffineModel
+    sampler_affine = cuqi.sampler.LinearRTO(posterior_affine)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_affine = sampler_affine.sample(10, 2)
+
+    # Check that the samples are the same
+    assert np.allclose(samples_linear.samples, samples_affine.samples)
+
+def test_RegularizedRTO_with_AffieModel_is_equivalent_to_LinearModel_and_shifted_data():
+    """ Test that sampling with RegularizedRTO with an AffineModel is equivalent to sampling with LinearModel and shifted data. """
+
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Shift
+    shift = np.random.rand(A.domain_dim)
+
+    # Define Bayesian Problem
+    x = cuqi.implicitprior.NonnegativeGMRF(np.zeros(A.domain_dim), 100)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs-shift)
+
+    # Set up RegularizedRTO with both models
+    sampler_linear = cuqi.sampler.RegularizedLinearRTO(posterior)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_linear = sampler_linear.sample(10, 2)
+
+    # Define AffineModel
+    affine_model = cuqi.model.AffineModel(A, shift)
+
+    # Set up RegularizedRTO with AffineModel
+    y = cuqi.distribution.Gaussian(affine_model, 0.01**2)
+    posterior_affine = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up RegularizedRTO with AffineModel
+    sampler_affine = cuqi.sampler.RegularizedLinearRTO(posterior_affine)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_affine = sampler_affine.sample(10, 2)
+
+    # Check that the samples are the same
+    assert np.allclose(samples_linear.samples, samples_affine.samples)
+        
+def test_UGLA_with_AffineModel_is_equivalent_to_LinearModel_and_shifted_data():
+    """ Test that sampling with UGLA with an AffineModel is equivalent to sampling with LinearModel and shifted data. """
+
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Shift
+    shift = np.random.rand(A.domain_dim)
+
+    # Define Bayesian Problem
+    x = cuqi.distribution.LMRF(np.zeros(A.domain_dim), 0.01)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs-shift)
+
+    # Set up UGLA with both models
+    sampler_linear = cuqi.sampler.UGLA(posterior)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_linear = sampler_linear.sample(10, 2)
+
+    # Define AffineModel
+    affine_model = cuqi.model.AffineModel(A, shift)
+
+    # Set up UGLA with AffineModel
+    y = cuqi.distribution.Gaussian(affine_model, 0.01**2)
+    posterior_affine = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up UGLA with AffineModel
+    sampler_affine = cuqi.sampler.UGLA(posterior_affine)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_affine = sampler_affine.sample(10, 2)
+
+    # Check that the samples are the same
+    assert np.allclose(samples_linear.samples, samples_affine.samples)
diff --git a/tests/zexperimental/test_mcmc.py b/tests/zexperimental/test_mcmc.py
index fc9bd6ecf..177962dcd 100644
--- a/tests/zexperimental/test_mcmc.py
+++ b/tests/zexperimental/test_mcmc.py
@@ -330,12 +330,15 @@ def create_myula_target(dim=16):
     """Create a target for MYULA."""
     def func(x, restoration_strength=1):
         return x, True
-    likelihood = cuqi.testproblem.Deconvolution1D(
+    y = cuqi.testproblem.Deconvolution1D(
         dim=dim).posterior.likelihood 
-    restoration_prior = cuqi.implicitprior.RestorationPrior(
-        func, geometry=likelihood.model.domain_geometry)
+    x = cuqi.implicitprior.RestorationPrior(
+        func, geometry=y.model.domain_geometry)
+    # access name property to ensure RestorationPrior
+    # name is set 
+    x.name
     posterior = cuqi.distribution.Posterior(
-        likelihood, restoration_prior)
+        y, x)
     return posterior
 
 def create_myula_smoothed_target(dim=16):
@@ -368,6 +371,10 @@ def test_myula():
     assert samples_ula.Ns == 10
     assert samples_myula.Ns == 10
     assert np.allclose(samples_myula.samples, samples_ula.samples)
+    # Assert the name of the MoreauYoshidaPrior created by MYULA is
+    # the same name as the corresponding RestorationPrior (both are named x)
+    assert myula._smoothed_target.prior.name == 'x'
+    assert myula.target.prior.name == 'x'
 
 def test_myula_object_creation_fails_with_target_without_restore_method():
     """ Test that MYULA object creation fails with target that does not
@@ -1134,4 +1141,120 @@ def test_nuts_acceptance_rate(sampler: cuqi.experimental.mcmc.Sampler):
     # Compute the number of accepted samples according to the sampler
     acc_rate_sum = sum(sampler._acc[2:])
 
-    assert np.isclose(counter, acc_rate_sum), "NUTS sampler does not update acceptance rate correctly: "+str(counter)+" != "+str(acc_rate_sum)
\ No newline at end of file
+    assert np.isclose(counter, acc_rate_sum), "NUTS sampler does not update acceptance rate correctly: "+str(counter)+" != "+str(acc_rate_sum)
+    
+# ============ Testing of AffineModel with RTO-type samplers ============
+
+def test_LinearRTO_with_AffineModel_is_equivalent_to_LinearModel_and_shifted_data():
+    """ Test that LinearRTO with AffineModel is equivalent to LinearRTO with LinearModel and shifted data. """
+
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Shift
+    shift = np.random.rand(A.domain_dim)
+
+    # Define Bayesian Problem
+    x = cuqi.distribution.GMRF(np.zeros(A.domain_dim), 100)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs-shift)
+
+    # Set up LinearRTO with both models
+    sampler_linear = cuqi.experimental.mcmc.LinearRTO(posterior)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_linear = sampler_linear.warmup(2).sample(10).get_samples()
+
+    # Define AffineModel
+    affine_model = cuqi.model.AffineModel(A, shift)
+
+    # Set up LinearRTO with AffineModel
+    y = cuqi.distribution.Gaussian(affine_model, 0.01**2)
+    posterior_affine = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up LinearRTO with AffineModel
+    sampler_affine = cuqi.experimental.mcmc.LinearRTO(posterior_affine)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_affine = sampler_affine.warmup(2).sample(10).get_samples()
+
+    # Check that the samples are the same
+    assert np.allclose(samples_linear.samples, samples_affine.samples)
+
+def test_RegularizedLinearRTO_with_AffineModel_is_equivalent_to_RegularizedLinearModel_and_shifted_data():
+    """ Test that RegularizedLinearRTO with AffineModel is equivalent to RegularizedLinearRTO with RegularizedLinearModel and shifted data. """
+
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Shift
+    shift = np.random.rand(A.domain_dim)
+
+    # Define Bayesian Problem
+    x = cuqi.implicitprior.NonnegativeGMRF(np.zeros(A.domain_dim), 100)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs-shift)
+
+    # Set up LinearRTO with both models
+    sampler_linear = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_linear = sampler_linear.warmup(2).sample(10).get_samples()
+
+    # Define AffineModel
+    affine_model = cuqi.model.AffineModel(A, shift)
+
+    # Set up LinearRTO with AffineModel
+    y = cuqi.distribution.Gaussian(affine_model, 0.01**2)
+    posterior_affine = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up LinearRTO with AffineModel
+    sampler_affine = cuqi.experimental.mcmc.RegularizedLinearRTO(posterior_affine)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_affine = sampler_affine.warmup(2).sample(10).get_samples()
+
+    # Check that the samples are the same
+    assert np.allclose(samples_linear.samples, samples_affine.samples)
+
+def test_UGLA_with_AffineModel_is_equivalent_to_LinearModel_and_shifted_data():
+    """ Test that UGLA with AffineModel is equivalent to LinearModel and shifted data. """
+    
+    # Define LinearModel and data
+    A, y_obs, _ = cuqi.testproblem.Deconvolution1D().get_components()
+
+    # Define Shift
+    shift = np.random.rand(A.domain_dim)
+
+    # Define Bayesian Problem
+    x = cuqi.distribution.LMRF(np.zeros(A.domain_dim), 0.01)
+    y = cuqi.distribution.Gaussian(A@x, 0.01**2)
+    posterior = cuqi.distribution.JointDistribution(x, y)(y=y_obs-shift)
+
+    # Set up LinearRTO with both models
+    sampler_linear = cuqi.experimental.mcmc.UGLA(posterior)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_linear = sampler_linear.warmup(2).sample(10).get_samples()
+
+    # Define AffineModel
+    affine_model = cuqi.model.AffineModel(A, shift)
+
+    # Set up LinearRTO with AffineModel
+    y = cuqi.distribution.Gaussian(affine_model, 0.01**2)
+    posterior_affine = cuqi.distribution.JointDistribution(x, y)(y=y_obs)
+
+    # Set up LinearRTO with AffineModel
+    sampler_affine = cuqi.experimental.mcmc.UGLA(posterior_affine)
+
+    # Sample with fixes seed
+    np.random.seed(0)
+    samples_affine = sampler_affine.warmup(2).sample(10).get_samples()
+
+    # Check that the samples are the same
+    assert np.allclose(samples_linear.samples, samples_affine.samples)