Repository: CUQI-DTU/CUQIpy

Issue:
In MYULA, use shallow copy instead of deepcopy when smoothing the target
Description:
-------------
In MYULA, we copy the target when it is not smoothed to create a smoothed target (basically changing the prior from `RestoratorationPrior` type to `MoreauYoshidaPrior` type). Deepcopy is not needed and it causes and error when the target posterior contains, for example, an underlying FEniCS model. See the error msg below. The error is due to the inability of `pickle`, which is used by deep copy, to pickle some fenics objects.

DoD:
------
- [x] (2 line change in the code) replace copy with deepcopy
- [x] Document the error message that arise from using MYULA with target containing fenics based model
- [x] Get the change to be approved

**note** that MYULA tests passes in both cases (with copy and with deepcopy) and no need for additional testing.

Error
-----
```
{
	"name": "TypeError",
	"message": "cannot pickle 'dolfin.cpp.generation.IntervalMesh' object",
	"stack": "---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In [26], line 1
----> 1 sampler =MYULA(BP.posterior, )

File ~/Documents/research_code/CUQI-DTU/CUQIpy/cuqi/experimental/mcmc/_langevin_algorithm.py:303, in MYULA.__init__(self, target, scale, smoothing_strength, **kwargs)
    301 def __init__(self, target=None, scale=1.0, smoothing_strength=0.1, **kwargs):
    302     self.smoothing_strength = smoothing_strength
--> 303     super().__init__(target=target, scale=scale, **kwargs)

File ~/Documents/research_code/CUQI-DTU/CUQIpy/cuqi/experimental/mcmc/_langevin_algorithm.py:73, in ULA.__init__(self, target, scale, **kwargs)
     71 def __init__(self, target=None, scale=1.0, **kwargs):
---> 73     super().__init__(target, **kwargs)
     74     self.initial_scale = scale

File ~/Documents/research_code/CUQI-DTU/CUQIpy/cuqi/experimental/mcmc/_sampler.py:67, in Sampler.__init__(self, target, initial_point, callback)
     46 def __init__(self, target:cuqi.density.Density=None, initial_point=None, callback=None):
     47     \"\"\" Initializer for abstract base class for all samplers.
     48 
     49     Any subclassing samplers should simply store input parameters as part of the __init__ method. 
   (...)
     64 
     65     \"\"\"
---> 67     self.target = target
     68     self.initial_point = initial_point
     69     self.callback = callback

File ~/Documents/research_code/CUQI-DTU/CUQIpy/cuqi/experimental/mcmc/_langevin_algorithm.py:312, in MYULA.target(self, value)
    308 self._target = value
    310 if self._target is not None:
    311     # Create a smoothed target
--> 312     self._smoothed_target = self._create_smoothed_target(value)
    314     # Validate the target
    315     self.validate_target()

File ~/Documents/research_code/CUQI-DTU/CUQIpy/cuqi/experimental/mcmc/_langevin_algorithm.py:319, in MYULA._create_smoothed_target(self, value)
    317 def _create_smoothed_target(self, value):
    318     \"\"\" Create a smoothed target using a Moreau-Yoshida envelope. \"\"\"
--> 319     copied_value = deepcopy(value)
    320     if isinstance(copied_value.prior, RestorationPrior):
    321         copied_value.prior = MoreauYoshidaPrior(
    322             copied_value.prior,
    323             self.smoothing_strength)

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269 if state is not None:
    270     if deep:
--> 271         state = deepcopy(state, memo)
    272     if hasattr(y, '__setstate__'):
    273         y.__setstate__(state)

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
    229 memo[id(x)] = y
    230 for key, value in x.items():
--> 231     y[deepcopy(key, memo)] = deepcopy(value, memo)
    232 return y

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269 if state is not None:
    270     if deep:
--> 271         state = deepcopy(state, memo)
    272     if hasattr(y, '__setstate__'):
    273         y.__setstate__(state)

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
    229 memo[id(x)] = y
    230 for key, value in x.items():
--> 231     y[deepcopy(key, memo)] = deepcopy(value, memo)
    232 return y

    [... skipping similar frames: _deepcopy_dict at line 231 (5 times), _reconstruct at line 271 (5 times), deepcopy at line 172 (5 times), deepcopy at line 146 (5 times)]

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:172, in deepcopy(x, memo, _nil)
    170                 y = x
    171             else:
--> 172                 y = _reconstruct(x, memo, *rv)
    174 # If is its own copy, don't memoize.
    175 if y is not x:

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:271, in _reconstruct(x, memo, func, args, state, listiter, dictiter, deepcopy)
    269 if state is not None:
    270     if deep:
--> 271         state = deepcopy(state, memo)
    272     if hasattr(y, '__setstate__'):
    273         y.__setstate__(state)

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:146, in deepcopy(x, memo, _nil)
    144 copier = _deepcopy_dispatch.get(cls)
    145 if copier is not None:
--> 146     y = copier(x, memo)
    147 else:
    148     if issubclass(cls, type):

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:231, in _deepcopy_dict(x, memo, deepcopy)
    229 memo[id(x)] = y
    230 for key, value in x.items():
--> 231     y[deepcopy(key, memo)] = deepcopy(value, memo)
    232 return y

File ~/opt/miniconda3/envs/fenicsproject/lib/python3.10/copy.py:161, in deepcopy(x, memo, _nil)
    159 reductor = getattr(x, \"__reduce_ex__\", None)
    160 if reductor is not None:
--> 161     rv = reductor(4)
    162 else:
    163     reductor = getattr(x, \"__reduce__\", None)

TypeError: cannot pickle 'dolfin.cpp.generation.IntervalMesh' object"
}
```

Discussion:
@amal-ghamdi: @jakobsj , can you review the DoD. Thank you.@jakobsj: @amal-ghamdi DoD approved!

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..9ab5fb90a 100644
--- a/cuqi/experimental/mcmc/_rto.py
+++ b/cuqi/experimental/mcmc/_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.
@@ -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")
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/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/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_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)