Add affine model nabr suggestions

cuqi/experimental/mcmc/_laplace_approximation.py

@@ -106,12 +106,12 @@ def Lk_fun(x_k):
         # Least squares form
         def M(x, flag):
             if flag == 1:
-                out1 = self._L1 @ self.model._forward_no_shift(x) # Use forward function which excludes shift
+                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_no_shift(self._L1.T@x[:idx]) # Use forward function which excludes shift
+                out1 = self.model.adjoint(self._L1.T@x[:idx])
                 out2 = np.sqrt(1/self.prior.scale)*(self._L2.T @ x[idx:])
                 out  = out1 + out2                
             return out

cuqi/experimental/mcmc/_rto.py

@@ -82,7 +82,7 @@ def models(self):
 
     @property
     def data(self):
-        return self.target.data - self.target.model._shift # Include shift from AffineModel here, self.data is never used?
+        return self.target.data - self.target.model._shift # Include shift from AffineModel here
 
     def _precompute(self):
         L1 = [likelihood.distribution.sqrtprec for likelihood in self.likelihoods]
@@ -91,7 +91,7 @@ def _precompute(self):
 
         # pre-computations
         self.n = self.prior.dim
-        self.b_tild = np.hstack([L@(likelihood.data - model._shift) for (L, likelihood, model) in zip(L1, self.likelihoods, self.models)]+ [L2mu]) # shift in data here, since self.data is not used
+        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):
@@ -100,7 +100,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_no_shift(x) for (L, likelihood) in zip(L1, self.likelihoods)] # Use forward function which excludes shift
+                    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:
@@ -109,7 +109,7 @@ def M(x, flag):
                     out1 = np.zeros(self.n)
                     for likelihood in self.likelihoods:
                         idx_end += len(likelihood.data)
-                        out1 += likelihood.model._adjoint_no_shift(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end]) # Use adjoint function which excludes shift
+                        out1 += likelihood.model.adjoint(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end])
                         idx_start = idx_end
                     out2 = L2.T @ x[idx_end:]
                     out  = out1 + out2                

cuqi/model/_model.py

@@ -1,17 +1,14 @@
 from __future__ import annotations
-from typing import List, Union
 import numpy as np
 from scipy.sparse import csc_matrix
 from scipy.sparse import hstack
 from scipy.linalg import solve
 from cuqi.samples import Samples
 from cuqi.array import CUQIarray
 from cuqi.geometry import Geometry, _DefaultGeometry1D, _DefaultGeometry2D, _get_identity_geometries
-from cuqi.utilities import infer_len
-import numbers
 import cuqi
 import matplotlib.pyplot as plt
-from copy import copy, deepcopy
+from copy import copy
 
 class Model(object):
     """Generic model defined by a forward operator.
@@ -475,7 +472,7 @@ def __repr__(self) -> str:
 
 
 class AffineModel(Model):
-    """ Model representing an affine operator, i.e. a linear operator with a fixed shift.
+    """ Model representing an affine operator, i.e. a linear operator with a fixed shift. For more details on linear operators, see :class:`~cuqi.model.LinearModel`.
 
     Parameters
     ----------
@@ -495,84 +492,87 @@ class AffineModel(Model):
     domain_geometry : cuqi.geometry.Geometry
         The geometry representing the domain.
 
-    Methods
-    -----------
-    :meth:`get_matrix` returns an ndarray with the matrix representing the forward operator.
-
     """
 
     def __init__(self, linear_operator, shift, linear_operator_adjoint=None, range_geometry=None, domain_geometry=None):
 
-        #Store shift as private attribute
-        self._shift = shift
-
-        if not callable(linear_operator):
-            forward_func = lambda x: self._matrix@x + self._shift
-            forward_func_noshift = lambda x: self._matrix@x
-            adjoint_func_noshift = lambda y: self._matrix.T@y
+        # If input represents a matrix, extract needed properties from it
+        if hasattr(linear_operator, '__matmul__') and hasattr(linear_operator, 'T'):
             matrix = linear_operator
+            linear_operator = lambda x: matrix@x
+            linear_operator_adjoint = lambda y: matrix.T@y
+            if range_geometry is None and hasattr(matrix, 'shape'):
+                range_geometry = _DefaultGeometry1D(grid=matrix.shape[0])
+            if domain_geometry is None and hasattr(matrix, 'shape'):
+                domain_geometry = _DefaultGeometry1D(grid=matrix.shape[1])
         else:
-            forward_func = lambda *args, **kwargs: self._2par(linear_operator(*args, **kwargs), self.range_geometry) + self._shift
-            forward_func_noshift = linear_operator
-            adjoint_func_noshift = linear_operator_adjoint
             matrix = None
 
-        #Check if input is callable
-        if not callable(forward_func_noshift):
+        # Check if input is callable
+        if not callable(linear_operator):
             raise TypeError("Linear operator must be defined as a matrix or a 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])  
+        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")
 
-        #Initialize Model class
-        super().__init__(forward_func, range_geometry, domain_geometry)
+        # Initialize Model class
+        super().__init__(linear_operator, range_geometry, domain_geometry)
 
-        #Store matrix privately
+        # Store matrix privately
         self._matrix = matrix
 
-        #Store linear operator privately
+        # Store shift as private attribute
+        self._shift = shift
+
+        # Store linear operator privately
         self._linear_operator = linear_operator
 
-        #Define gradient
-        self._gradient_func = lambda direction, wrt: adjoint_func_noshift(direction)
+        # Store adjoint function
+        self._linear_operator_adjoint = linear_operator_adjoint
+
+        # Define gradient
+        self._gradient_func = lambda direction, wrt: linear_operator_adjoint(direction)
 
-        #Add adjoint without shift
-        self._adjoint_func_noshift = adjoint_func_noshift
+        # 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)
+
+    def adjoint(self, y, is_par=True):
+        """ Adjoint of the model.
+
+        Adjoint converts the input to function values (if needed) using the range geometry of the model.
+        Adjoint converts the output function values to parameters using the range geometry of the model.
 
-        #Add forward without shift
-        self._forward_func_noshift = forward_func_noshift 
+        Parameters
+        ----------
+        y : ndarray or cuqi.array.CUQIarray
+            The adjoint model input.
 
-        # Use arguments from user's callable linear operator
-        if callable(linear_operator):
-            self._non_default_args = cuqi.utilities.get_non_default_args(linear_operator)
+        Returns
+        -------
+        ndarray or cuqi.array.CUQIarray
+            The adjoint model output. Always returned as parameters.
+        """
+        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, # Adjoint of x -> Ax+b is A^T y
+                                self.domain_geometry,
+                                self.range_geometry,
+                                y, is_par)
 
     def update_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
 
-        # update forward function with new shift
-        if not callable(self._linear_operator):
-            forward_func = lambda x: self._matrix@x + self._shift
-        else:
-            forward_func = lambda *args, **kwargs: self._2par(self._linear_operator(*args, **kwargs), self.range_geometry) + self._shift
-        self._forward_func = forward_func
-
-    def _forward_no_shift(self, x, is_par=True):
+    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._forward_func_noshift,
+        return self._apply_func(self._linear_operator,
                 self.range_geometry,
                 self.domain_geometry,
                 x, is_par)
-
-    def _adjoint_no_shift(self, y, is_par=True):
-        """ Helper function for computing the adjoint operator without the shift. """
-        return self._apply_func(self._adjoint_func_noshift,
-                        self.domain_geometry,
-                        self.range_geometry,
-                        y, is_par)
 
     def get_matrix(self):
         """
@@ -598,20 +598,6 @@ def get_matrix(self):
 
             return self._matrix   
 
-    # Related to model algebra
-    # def __add__(self, shift) -> Union[Model, LinearModel]:
-    #     """ Creates an AffineModel with a fixed shift added, i.e. model_shifted(x) = model(x) + shift. """
-    #     if isinstance(shift, numbers.Number) and shift == 0:
-    #         return copy(self)
-    #     if not isinstance(shift, numbers.Number) and infer_len(shift) != self.range_dim:
-    #         raise ValueError("Shift dimension does not match model range dimension.")
-    #     if isinstance(shift, Model):
-    #         return SumOfModels(self, shift)
-    #     if isinstance(shift, SumOfModels):
-    #         return SumOfModels(self, *shift._models)
-    #     new_model = copy(self)
-    #     new_model._shift = self._shift + shift
-    #     return new_model
 
 class LinearModel(AffineModel):
     """Model based on a Linear forward operator.
@@ -677,53 +663,19 @@ 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):
+    def __init__(self, forward, adjoint=None, range_geometry=None, domain_geometry=None):
 
         #Initialize as AffineModel with shift=0
         super().__init__(forward, 0, adjoint, range_geometry, domain_geometry)
 
-        if not callable(forward):
-            adjoint_func = lambda y: self._matrix.T@y
-        else:
-            adjoint_func = adjoint
-
-        #Check if input is callable
-        if not callable(adjoint_func):
-            raise TypeError("Adjoint of linear operator must be defined as a callable function of some kind")
-
-        #Add adjoint
-        self._adjoint_func = adjoint_func
-
-    def adjoint(self, y, is_par=True):
-        """ Adjoint of the model.
-
-        Adjoint converts the input to function values (if needed) using the range geometry of the model.
-        Adjoint converts the output function values to parameters using the range geometry of the model.
-
-        Parameters
-        ----------
-        y : ndarray or cuqi.array.CUQIarray
-            The adjoint model input.
-
-        Returns
-        -------
-        ndarray or cuqi.array.CUQIarray
-            The adjoint model output. Always returned as parameters.
-        """
-        return self._apply_func(self._adjoint_func,
-                                self.domain_geometry,
-                                self.range_geometry,
-                                y, is_par)
-
     def __matmul__(self, x):
         return self.forward(x)
 
     @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

cuqi/sampler/_laplace_approximation.py

@@ -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_no_shift(x) # Use forward function which excludes shift
+                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_no_shift(self._L1.T@x[:idx]) # Use forward function which excludes shift
+                out1 = self._model.adjoint(self._L1.T@x[:idx])
                 out2 = np.sqrt(1/self.target.prior.scale)*(self._L2.T @ x[idx:])
                 out  = out1 + out2                
             return out 

cuqi/sampler/_rto.py

@@ -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_no_shift(x) for (L, likelihood) in zip(L1, self.likelihoods)] # Use forward function which excludes shift
+                    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_no_shift(likelihood.distribution.sqrtprec.T@x[idx_start:idx_end]) # Use adjoint function which excludes shift
+                        out1 += likelihood.model.adjoint(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                
@@ -151,7 +151,7 @@ def models(self):
 
     @property
     def data(self):
-        return self.target.data - self.target.model._shift # Include shift from AffineModel here, self.data never used
+        return self.target.data - self.target.model._shift # Include shift from AffineModel here
 
     def _sample(self, N, Nb):   
         Ns = N+Nb   # number of simulations        

demos/try17_AffineModel.py

@@ -1,7 +1,7 @@
+# %%
 import numpy as np
 import matplotlib.pyplot as plt
 import sys
-sys.path.append("..") 
 
 from cuqi.model import AffineModel, LinearModel
 from cuqi.experimental.mcmc import MH, LinearRTO
@@ -86,7 +86,7 @@
 posterior = JointDistribution(x, y)(y=y_obs)
 
 print(posterior.likelihood.model.forward(np.array([1,1])))
-print(posterior.likelihood.model._forward_no_shift(np.array([1,1])))
+print(posterior.likelihood.model._forward_func_no_shift(np.array([1,1])))
 
 # Sample posterior with MH
 np.random.seed(1000000)
@@ -111,4 +111,4 @@
 plt.legend()
 
 fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(5,4), squeeze=False)
-cs = samplesAffine_burnin.plot_pair(marginals = True)
+cs = samplesAffine_burnin.plot_pair(marginals = True)