Merge pull request #486 from CUQI-DTU/add_AffineModel

Add AffineModel

Author: nabriis

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"):

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,21 +71,23 @@ 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
         L2mu = self.prior.sqrtprecTimesMean
 
         # pre-computations
         self.n = self.prior.dim
-        self.b_tild = np.hstack([L@likelihood.data for (L, likelihood) in zip(L1, self.likelihoods)]+ [L2mu]) 
-
+        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")

cuqi/model/__init__.py

@@ -1 +1 @@
-from ._model import Model, LinearModel, PDEModel
+from ._model import Model, LinearModel, PDEModel, AffineModel

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