CUQIpy 1.3.0__py3-none-any.whl → 1.4.0.post0.dev61__py3-none-any.whl

cuqi/model/_model.py

@@ -5,48 +5,55 @@ 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.geometry import Geometry, _DefaultGeometry1D, _DefaultGeometry2D,\
+    _get_identity_geometries
 import cuqi
 import matplotlib.pyplot as plt
 from copy import copy
+from functools import partial
+from cuqi.utilities import force_ndarray
 
 class Model(object):
     """Generic model defined by a forward operator.
 
     Parameters
     -----------
-    forward : 2D ndarray or callable function.
-        Forward operator.
+    forward : callable function
+        Forward operator of the model. It takes one or more inputs and returns the model output.
 
-    range_geometry : integer or cuqi.geometry.Geometry
-        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
+    range_geometry : integer, a 1D or 2D tuple of integers, cuqi.geometry.Geometry 
+        If integer or 1D tuple of integers is given, a cuqi.geometry._DefaultGeometry1D is created with dimension of the integer.
+        If 2D tuple of integers is given, a cuqi.geometry._DefaultGeometry2D is created with dimensions of the tuple.
+        If cuqi.geometry.Geometry object is given, it is used as the range geometry of the model.
 
-    domain_geometry : integer or cuqi.geometry.Geometry
-        If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
+    domain_geometry : integer, a 1D or 2D tuple of integers, cuqi.geometry.Geometry or a tuple with items of any of the listed types
+        If integer or 1D tuple of integers is given, a cuqi.geometry._DefaultGeometry1D is created with dimension of the integer.
+        If 2D tuple of integers is given (and the forward model has one input only), a cuqi.geometry._DefaultGeometry2D is created with dimensions of the tuple.
+        If cuqi.geometry.Geometry is given, it is used as the domain geometry.
+        If tuple of the above types is given, a cuqi.geometry._ProductGeometry is created based on the tuple entries. This is used for models with multiple inputs where each entry in the tuple represents the geometry of each input.
 
-    gradient : callable function, optional
-        The direction-Jacobian product of the forward operator Jacobian with 
-        respect to the forward operator input, evaluated at a point (`wrt`).
-        The signature of the gradient function should be (`direction`, `wrt`),
-        where `direction` is the direction by which the Jacobian matrix is
-        multiplied and `wrt` is the point at which the Jacobian is computed.
+    gradient : callable function, a tuple of callable functions or None, optional
+        The direction-Jacobian product of the forward model Jacobian with respect to the model input, evaluated at the model input. For example, if the forward model inputs are `x` and `y`, the gradient callable signature should be (`direction`, `x`, `y`), in that order, where `direction` is the direction by which the Jacobian matrix is multiplied and `x` and `y` are the parameters at which the Jacobian is computed.
 
-    jacobian : callable function, optional
-        The Jacobian of the forward operator with respect to the forward operator input,
-        evaluated at a point (`wrt`). The signature of the Jacobian function should be (`wrt`).
-        The Jacobian function should return a 2D ndarray of shape (range_dim, domain_dim).
-        The Jacobian function is used to specify the gradient function by computing the vector-Jacobian
-        product (VJP), here we refer to the vector in the VJP as the `direction` since it is the direction at 
-        which the gradient is computed.
-        automatically and thus the gradient function should not be specified when the Jacobian
-        function is specified.
+        If the gradient function is a single callable function, it returns a 1D ndarray if the model has only one input. If the model has multiple inputs, this gradient function should return a tuple of 1D ndarrays, each representing the gradient with respect to each input.
+
+        If the gradient function is a tuple of callable functions, each callable function should return a 1D ndarray representing the gradient with respect to each input. The order of the callable functions in the tuple should match the order of the model inputs.
+
+    jacobian : callable function, a tuple of callable functions or None, optional
+        The Jacobian of the forward model with respect to the forward model input, evaluated at the model input. For example, if the forward model inputs are `x` and `y`, the jacobian signature should be (`x`, `y`), in that order, where `x` and `y` are the parameters at which the Jacobian is computed.
+
+        If the Jacobian function is a single callable function, it should return a 2D ndarray of shape (range_dim, domain_dim) if the model has only one input. If the model has multiple inputs, this Jacobian function should return a tuple of 2D ndarrays, each representing the Jacobian with respect to each input. 
+        
+        If the Jacobian function is a tuple of callable functions, each callable function should return a 2D ndarray representing the Jacobian with respect to each input. The order of the callable functions in the tuple should match the order of the model inputs.
+    
+        The Jacobian function is used to specify the gradient function by computing the vector-Jacobian product (VJP), here we refer to the vector in the VJP as the `direction` since it is the direction at which the gradient is computed. Either the gradient or the Jacobian can be specified, but not both.
 
 
     :ivar range_geometry: The geometry representing the range.
     :ivar domain_geometry: The geometry representing the domain.
 
-    Example
-    -------
+    Example 1
+    ----------
 
     Consider a forward model :math:`F: \mathbb{R}^2 \\rightarrow \mathbb{R}` defined by the following forward operator:
 
@@ -75,6 +82,9 @@ class Model(object):
 
         model = Model(forward, range_geometry=1, domain_geometry=2, jacobian=jacobian)
 
+        print(model(np.array([1, 1])))
+        print(model.gradient(np.array([1]), np.array([1, 1])))
+
     Alternatively, the gradient information in the forward model can be defined by direction-Jacobian product using the gradient keyword argument.
 
     This may be more efficient if forming the Jacobian matrix is expensive.
@@ -87,64 +97,162 @@ class Model(object):
         def forward(x):
             return 10*x[1] - 10*x[0]**3 + 5*x[0]**2 + 6*x[0]
 
-        def gradient(direction, wrt):
-            # Direction-Jacobian product direction@jacobian(wrt)
-            return direction@np.array([[-30*wrt[0]**2 + 10*wrt[0] + 6, 10]])
+        def gradient(direction, x):
+            # Direction-Jacobian product direction@jacobian(x)
+            return direction@np.array([[-30*x[0]**2 + 10*x[0] + 6, 10]])
 
         model = Model(forward, range_geometry=1, domain_geometry=2, gradient=gradient)
 
+        print(model(np.array([1, 1])))
+        print(model.gradient(np.array([1]), np.array([1, 1])))
+
+    Example 2
+    ----------
+    Alternatively, the example above can be defined as a model with multiple inputs: :math:`x` and :math:`y`:
+
+    .. code-block:: python
+
+        import numpy as np
+        from cuqi.model import Model
+        from cuqi.geometry import Discrete
+    
+        def forward(x, y):
+            return 10 * y - 10 * x**3 + 5 * x**2 + 6 * x
+    
+        def jacobian(x, y):
+            return (np.array([[-30 * x**2 + 10 * x + 6]]), np.array([[10]]))
+    
+        model = Model(
+            forward,
+            range_geometry=1,
+            domain_geometry=(Discrete(1), Discrete(1)),
+            jacobian=jacobian,
+        )
+    
+        print(model(1, 1))
+        print(model.gradient(np.array([1]), 1, 1))
     """
+
+    _supports_partial_eval = True
+    """Flag indicating that partial evaluation of Model objects is supported, i.e., calling the model object with only some of the inputs specified returns a model that can be called with the remaining inputs."""
+
     def __init__(self, forward, range_geometry, domain_geometry, gradient=None, jacobian=None):
 
-        #Check if input is callable
+        # Check if input is callable
         if callable(forward) is not True:
             raise TypeError("Forward needs to be callable function.")
-        
+
+        # Store forward func
+        self._forward_func = forward
+        self._stored_non_default_args = None
+
+        # Store range_geometry
+        self.range_geometry = range_geometry
+
+        # Store domain_geometry
+        self.domain_geometry = domain_geometry
+
+        # Additional checks for the forward operator
+        self._check_domain_geometry_consistent_with_forward()
+
         # Check if only one of gradient and jacobian is given
         if (gradient is not None) and (jacobian is not None):
             raise TypeError("Only one of gradient and jacobian should be specified")
-        
-        #Check if input is callable
-        if (gradient is not None) and (callable(gradient) is not True):
-            raise TypeError("Gradient needs to be callable function.")
-        
-        if (jacobian is not None) and (callable(jacobian) is not True):
-            raise TypeError("Jacobian needs to be callable function.")
-        
-        # Use jacobian function to specify gradient function (vector-Jacobian product)
+
+        # Check correct gradient form (check type, signature, etc.)
+        self._check_correct_gradient_jacobian_form(gradient, "gradient")
+
+        # Check correct jacobian form (check type, signature, etc.)
+        self._check_correct_gradient_jacobian_form(jacobian, "jacobian")
+
+        # If jacobian is provided, use it to specify gradient function
+        # (vector-Jacobian product)
         if jacobian is not None:
-            gradient = lambda direction, wrt: direction@jacobian(wrt)
- 
-        #Store forward func
-        self._forward_func = forward
+            gradient = self._use_jacobian_to_specify_gradient(jacobian)
+
         self._gradient_func = gradient
-         
-        #Store range_geometry
-        if isinstance(range_geometry, tuple) and len(range_geometry) == 2:
-            self.range_geometry = _DefaultGeometry2D(range_geometry)
-        elif isinstance(range_geometry, int):
-            self.range_geometry = _DefaultGeometry1D(grid=range_geometry)
-        elif isinstance(range_geometry, Geometry):
-            self.range_geometry = range_geometry
-        elif range_geometry is None:
-            raise AttributeError("The parameter 'range_geometry' is not specified by the user and it connot be inferred from the attribute 'forward'.")
-        else:
-            raise TypeError("The parameter 'range_geometry' should be of type 'int', 2 dimensional 'tuple' or 'cuqi.geometry.Geometry'.")
-
-        #Store domain_geometry
-        if isinstance(domain_geometry, tuple) and len(domain_geometry) == 2:
-            self.domain_geometry = _DefaultGeometry2D(domain_geometry)
-        elif isinstance(domain_geometry, int):
-            self.domain_geometry = _DefaultGeometry1D(grid=domain_geometry)
-        elif isinstance(domain_geometry, Geometry):
-            self.domain_geometry = domain_geometry
-        elif domain_geometry is None:
-            raise AttributeError("The parameter 'domain_geometry' is not specified by the user and it connot be inferred from the attribute 'forward'.")
+
+        # Set gradient output stacked flag to False
+        self._gradient_output_stacked = False
+
+    @property
+    def _non_default_args(self):
+        if self._stored_non_default_args is None:
+            # Store non_default_args of the forward operator for faster caching
+            # when checking for those arguments.
+            self._stored_non_default_args =\
+                cuqi.utilities.get_non_default_args(self._forward_func)
+        return self._stored_non_default_args
+
+    @property
+    def number_of_inputs(self):
+        """ The number of inputs of the model. """
+        return len(self._non_default_args)
+
+    @property
+    def range_geometry(self):
+        """ The geometry representing the range of the model. """
+        return self._range_geometry
+
+    @range_geometry.setter
+    def range_geometry(self, value):
+        """ Update the range geometry of the model. """
+        if isinstance(value, Geometry):
+            self._range_geometry = value
+        elif isinstance(value, int):
+            self._range_geometry = self._create_default_geometry(value)
+        elif isinstance(value, tuple):
+            self._range_geometry = self._create_default_geometry(value)
+        elif value is None:
+            raise AttributeError(
+                "The parameter 'range_geometry' is not specified by the user and it cannot be inferred from the attribute 'forward'."
+            )
         else:
-            raise TypeError("The parameter 'domain_geometry' should be of type 'int', 2 dimensional 'tuple' or 'cuqi.geometry.Geometry'.")
+            raise TypeError(
+                " The allowed types for 'range_geometry' are: 'cuqi.geometry.Geometry', int, 1D tuple of int, or 2D tuple of int."
+            )
 
-        # Store non_default_args of the forward operator for faster caching when checking for those arguments.
-        self._non_default_args = cuqi.utilities.get_non_default_args(self._forward_func)
+    @property
+    def domain_geometry(self):
+        """ The geometry representing the domain of the model. """
+        return self._domain_geometry
+
+    @domain_geometry.setter
+    def domain_geometry(self, value):
+        """ Update the domain geometry of the model. """
+
+        if isinstance(value, Geometry):
+            self._domain_geometry = value
+        elif isinstance(value, int):
+            self._domain_geometry = self._create_default_geometry(value)
+        elif isinstance(value, tuple) and self.number_of_inputs == 1:
+            self._domain_geometry = self._create_default_geometry(value)
+        elif isinstance(value, tuple) and self.number_of_inputs > 1:
+            geometries = [item if isinstance(item, Geometry) else self._create_default_geometry(item) for item in value]
+            self._domain_geometry = cuqi.experimental.geometry._ProductGeometry(*geometries)
+        elif value is None:
+            raise AttributeError(
+                "The parameter 'domain_geometry' is not specified by the user and it cannot be inferred from the attribute 'forward'."
+            )
+        else:
+            raise TypeError(
+                "For forward model with 1 input, the allowed types for 'domain_geometry' are: 'cuqi.geometry.Geometry', int, 1D tuple of int, or 2D tuple of int. For forward model with multiple inputs, the 'domain_geometry' should be a tuple with items of any of the above types."
+            )
+
+    def _create_default_geometry(self, value):
+        """Private function that creates default geometries for the model."""
+        if isinstance(value, tuple) and len(value) == 1:
+            value = value[0]
+        if isinstance(value, Geometry):
+            return value
+        if isinstance(value, int):
+            return _DefaultGeometry1D(grid=value)
+        elif isinstance(value, tuple) and len(value) == 2:
+            return _DefaultGeometry2D(im_shape=value)
+        else:
+            raise ValueError(
+                "Default geometry creation can be specified by an integer or a 2D tuple of integers."
+            )
 
     @property
     def domain_dim(self):
@@ -160,341 +268,977 @@ class Model(object):
         """
         return self.range_geometry.par_dim
 
-    def _2fun(self, x, geometry, is_par):
-        """ Converts `x` to function values (if needed) using the appropriate 
-        geometry. For example, `x` can be the model input which need to be
-        converted to function value before being passed to 
-        :class:`~cuqi.model.Model` operators (e.g. _forward_func, _adjoint_func,
-        _gradient_func).
+    def _check_domain_geometry_consistent_with_forward(self):
+        """Private function that checks if the domain geometry of the model is
+        consistent with the forward operator."""
+        if (
+            not isinstance(
+                self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+            )
+            and self.number_of_inputs > 1
+        ):
+            raise ValueError(
+                "The forward operator input is specified by more than one argument. This is only supported for domain geometry of type tuple with items of type: cuqi.geometry.Geometry object, int, or 2D tuple of int."
+            )
+
+    def _check_correct_gradient_jacobian_form(self, func, func_type):
+        """Private function that checks if the gradient/jacobian parameter is
+        in the correct form. That is, check if the gradient/jacobian has the
+        correct type, signature, etc."""
+
+        if func is None:
+            return
+
+        # gradient/jacobian should be callable (for single input and multiple input case)
+        # or a tuple of callables (for multiple inputs case)
+        if isinstance(func, tuple):
+            # tuple length should be same as the number of inputs
+            if len(func) != self.number_of_inputs:
+                raise ValueError(
+                    f"The "
+                    + func_type.lower()
+                    + f" tuple length should be {self.number_of_inputs} for model with inputs {self._non_default_args}"
+                )
+            # tuple items should be callables or None
+            if not all([callable(func_i) or func_i is None for func_i in func]):
+                raise TypeError(
+                    func_type.capitalize()
+                    + " tuple should contain callable functions or None."
+                )
+
+        elif callable(func):
+            # temporarily convert gradient/jacobian to tuple for checking only
+            func = (func,)
+
+        else:
+            raise TypeError(
+                "Gradient needs to be callable function or tuple of callable functions."
+            )
+
+        expected_func_non_default_args = (
+            self._non_default_args
+            if not hasattr(self, "_original_non_default_args")
+            else self._original_non_default_args
+        )
+
+        if func_type.lower() == "gradient":
+            # prepend 'direction' to the expected gradient non default args
+            expected_func_non_default_args = [
+                "direction"
+            ] + expected_func_non_default_args
+
+        for func_i in func:
+            # make sure the signature of the gradient/jacobian function is correct
+            # that is, the same as the expected_func_non_default_args
+            if func_i is not None:
+                func_non_default_args = cuqi.utilities.get_non_default_args(func_i)
+
+                if list(func_non_default_args) != list(expected_func_non_default_args):
+                    raise ValueError(
+                        func_type.capitalize()
+                        + f" function signature should be {expected_func_non_default_args}"
+                    )
+
+    def _use_jacobian_to_specify_gradient(self, jacobian):
+        """Private function that uses the jacobian function to specify the
+        gradient function."""
+        # if jacobian is a single function and model has multiple inputs
+        if callable(jacobian) and self.number_of_inputs > 1:
+            gradient = self._create_gradient_lambda_function_from_jacobian_with_correct_signature(
+                jacobian, form='one_callable_multiple_inputs'
+            )
+        # Elif jacobian is a single function and model has only one input
+        elif callable(jacobian):
+            gradient = self._create_gradient_lambda_function_from_jacobian_with_correct_signature(
+                jacobian, form='one_callable_one_input'
+            )
+        # Else, jacobian is a tuple of jacobian functions
+        else:
+            gradient = []
+            for jac in jacobian:
+                if jac is not None:
+                    gradient.append(
+                        self._create_gradient_lambda_function_from_jacobian_with_correct_signature(
+                            jac, form='tuple_of_callables'
+                        )
+                    )
+                else:
+                    gradient.append(None)
+        return tuple(gradient) if isinstance(gradient, list) else gradient
+
+    def _create_gradient_lambda_function_from_jacobian_with_correct_signature(
+        self, jacobian, form
+    ):
+        """Private function that creates gradient lambda function from the
+        jacobian function, with the correct signature (based on the model
+        non_default_args).
+        """
+        # create the string representation of the lambda function
+        # for different forms of jacobian
+        if form=='one_callable_multiple_inputs':
+            grad_fun_str = (
+                "lambda direction, "
+                + ", ".join(self._non_default_args)
+                + ", jacobian: tuple([direction@jacobian("
+                + ", ".join(self._non_default_args)
+                + ")[i] for i in range("+str(self.number_of_inputs)+")])"
+            )
+        elif form=='tuple_of_callables' or form=='one_callable_one_input':
+            grad_fun_str = (
+                "lambda direction, "
+                + ", ".join(self._non_default_args)
+                + ", jacobian: direction@jacobian("
+                + ", ".join(self._non_default_args)
+                + ")"
+            )
+        else:
+            raise ValueError("form should be either 'one_callable' or 'tuple_of_callables'.")
+
+        # create the lambda function from the string
+        grad_func = eval(grad_fun_str)
+
+        # create partial function from the lambda function with jacobian as a
+        # fixed argument
+        grad_func = partial(grad_func, jacobian=jacobian)
+
+        return grad_func
+
+    def _2fun(self, geometry=None, is_par=True, **kwargs):
+        """ Converts `kwargs` to function values (if needed) using the geometry. For example, `kwargs` can be the model input which need to be converted to function value before being passed to :class:`~cuqi.model.Model` operators (e.g. _forward_func, _adjoint_func, _gradient_func).
 
         Parameters
         ----------
-        x : ndarray or cuqi.array.CUQIarray
-            The value to be converted.
-
         geometry : cuqi.geometry.Geometry
-            The geometry representing `x`.
+            The geometry representing the values in `kwargs`.
 
-        is_par : bool
-            If True, `x` is assumed to be parameters.
-            If False, `x` is assumed to be function values.
+        is_par : bool or a tuple of bools
+            If `is_par` is True, the values in `kwargs` are assumed to be parameters.
+            If `is_par` is False, the values in `kwargs` are assumed to be function values.
+            If `is_par` is a tuple of bools, the values in `kwargs` are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
+        
+        **kwargs : keyword arguments to be converted to function values.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
-            `x` represented as a function.
+        dict of the converted values
+        """
+        # Check kwargs and geometry are consistent and set up geometries list and
+        # is_par tuple
+        geometries, is_par = self._helper_pre_conversion_checks_and_processing(geometry, is_par, **kwargs)
+
+        # Convert to function values
+        for i, (k, v) in enumerate(kwargs.items()):
+            # Use CUQIarray funvals if geometry is consistent
+            if isinstance(v, CUQIarray) and v.geometry == geometries[i]:
+                kwargs[k] = v.funvals
+            # Else, if we still need to convert to function value (is_par[i] is True)
+            # we use the geometry par2fun method
+            elif is_par[i] and v is not None:
+                kwargs[k] = geometries[i].par2fun(v)
+            else:
+                # No need to convert
+                pass
+
+        return kwargs
+
+    def _helper_pre_conversion_checks_and_processing(self, geometry=None, is_par=True, **kwargs):
+        """ Helper function that checks if kwargs and geometry are consistent
+        and sets up geometries list and is_par tuple.
         """
-        # Convert to function representation
-        # if x is CUQIarray and geometry are consistent, we obtain funvals
-        # directly
-        if isinstance(x, CUQIarray) and  x.geometry == geometry:
-            x = x.funvals
-        # Otherwise we use the geometry par2fun method
-        elif is_par:
-            x = geometry.par2fun(x)
-
-        return x
-
-    def _2par(self, val, geometry, to_CUQIarray=False, is_par=False):
-        """ Converts val, normally output of :class:~`cuqi.model.Model` 
-        operators (e.g. _forward_func, _adjoint_func, _gradient_func), to
-        parameters using the appropriate geometry.
+        # If len of kwargs is larger than 1, the geometry needs to be of type
+        # _ProductGeometry
+        if (
+            not isinstance(geometry, cuqi.experimental.geometry._ProductGeometry)
+            and len(kwargs) > 1
+        ):
+            raise ValueError(
+                "The input is specified by more than one argument. This is only "
+                + "supported for domain geometry of type "
+                + f"{cuqi.experimental.geometry._ProductGeometry.__name__}."
+            )
+
+        # If is_par is bool, make it a tuple of bools of the same length as
+        # kwargs
+        is_par = (is_par,) * len(kwargs) if isinstance(is_par, bool) else is_par
+
+        # Set up geometries list
+        geometries = (
+            geometry.geometries
+            if isinstance(geometry, cuqi.experimental.geometry._ProductGeometry)
+            else [geometry]
+        )
+
+        return geometries, is_par
+
+    def _2par(self, geometry=None, to_CUQIarray=False, is_par=False, **kwargs):    
+        """ Converts `kwargs` to parameters using the geometry. For example, `kwargs` can be the output of :class:`~cuqi.model.Model` operators (e.g. _forward_func, _adjoint_func, _gradient_func) which need to be converted to parameters before being returned.
 
         Parameters
         ----------
-        val : ndarray or cuqi.array.CUQIarray
-            The value to be converted to parameters.
-
         geometry : cuqi.geometry.Geometry
-            The geometry representing the argument `val`.
+            The geometry representing the values in `kwargs`.
 
-        to_CUQIarray : bool
-            If True, the returned value is wrapped as a cuqi.array.CUQIarray.
+        to_CUQIarray : bool or a tuple of bools
+            If `to_CUQIarray` is True, the values in `kwargs` will be wrapped in `CUQIarray`.
+            If `to_CUQIarray` is False, the values in `kwargs` will not be wrapped in `CUQIarray`.
+            If `to_CUQIarray` is a tuple of bools, the values in `kwargs` will be wrapped in `CUQIarray` or not based on the corresponding boolean value in the tuple.
         
-        is_par : bool
-            If True, `val` is assumed to be of parameter representation and
-            hence no conversion to parameters is performed.
+        is_par : bool or a tuple of bools
+            If `is_par` is True, the values in `kwargs` are assumed to be parameters.
+            If `is_par` is False, the values in `kwargs` are assumed to be function values.
+            If `is_par` is a tuple of bools, the values in `kwargs` are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
-            The value `val` represented as parameters.
+        dict of the converted values
         """
+        # Check kwargs and geometry are consistent and set up geometries list and
+        # is_par tuple
+        geometries, is_par = self._helper_pre_conversion_checks_and_processing(geometry, is_par, **kwargs)
+
+        # if to_CUQIarray is bool, make it a tuple of bools of the same length
+        # as kwargs
+        to_CUQIarray = (to_CUQIarray,) * len(kwargs) if isinstance(to_CUQIarray, bool) else to_CUQIarray
+
         # Convert to parameters
-        # if val is CUQIarray and geometry are consistent, we obtain parameters
-        # directly
-        if isinstance(val, CUQIarray) and val.geometry == geometry:
-            val = val.parameters
-        # Otherwise we use the geometry fun2par method
-        elif not is_par:
-            val = geometry.fun2par(val)
-
-        # Wrap val in CUQIarray if requested
-        if to_CUQIarray:
-            val = CUQIarray(val, is_par=True, geometry=geometry)
-
-        # Return val
-        return val
-        
+        for i , (k, v) in enumerate(kwargs.items()):
+            # Use CUQIarray parameters if geometry is consistent
+            if isinstance(v, CUQIarray) and v.geometry == geometries[i]:
+                v = v.parameters
+            # Else, if we still need to convert to parameter value (is_par[i] is False)
+            # we use the geometry fun2par method
+            elif not is_par[i] and v is not None:
+                v = geometries[i].fun2par(v)
+            else:
+                # No need to convert
+                pass
+
+            # Wrap the value v in CUQIarray if requested
+            if to_CUQIarray[i] and v is not None:
+                v = CUQIarray(v, is_par=True, geometry=geometries[i])
+
+            kwargs[k] = v
 
-    def _apply_func(self, func, func_range_geometry, func_domain_geometry, x, is_par, **kwargs):
-        """ Private function that applies the given function `func` to the input value `x`. It converts the input to function values (if needed) using the given `func_domain_geometry` and converts the output function values to parameters using the given `func_range_geometry`. It additionally handles the case of applying the function `func` to the cuqi.samples.Samples object.
+        return kwargs
 
-        kwargs are keyword arguments passed to the functions `func`.
+    def _apply_func(self, func=None, fwd=True, is_par=True, **kwargs):
+        """ Private function that applies the given function `func` to the input `kwargs`. It converts the input to function values (if needed) and converts the output to parameter values. It additionally handles the case of applying the function `func` to cuqi.samples.Samples objects.
         
         Parameters
         ----------
         func: function handler 
             The function to be applied.
 
-        func_range_geometry : cuqi.geometry.Geometry
-            The geometry representing the function `func` range.
-
-        func_domain_geometry : cuqi.geometry.Geometry
-            The geometry representing the function `func` domain.
-
-        x : ndarray or cuqi.array.CUQIarray
-            The input value to the operator.
+        fwd : bool
+            Flag indicating the direction of the operator to determine the range and domain geometries of the function.
+            If True the function is a forward operator.
+            If False the function is an adjoint operator.
 
-        is_par : bool
-            If True the input is assumed to be parameters.
-            If False the input is assumed to be function values.
+        is_par : bool or list of bool
+            If True, the inputs in `kwargs` are assumed to be parameters.
+            If False, the input in `kwargs` are assumed to be function values.
+            If `is_par` is a list of bools, the inputs are assumed to be parameters or function values based on the corresponding boolean value in the list.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
-            The output of the function `func` converted to parameters.
+        ndarray or cuqi.array.CUQIarray or cuqi.samples.Samples object
+            The output of the function.
         """ 
+        # Specify the range and domain geometries of the function
+        # If forward operator, range geometry is the model range geometry and
+        # domain geometry is the model domain geometry
+        if fwd:
+            func_range_geometry = self.range_geometry
+            func_domain_geometry = self.domain_geometry
+        # If adjoint operator, range geometry is the model domain geometry and
+        # domain geometry is the model range geometry
+        else:
+            func_range_geometry = self.domain_geometry
+            func_domain_geometry = self.range_geometry
+
         # If input x is Samples we apply func for each sample
         # TODO: Check if this can be done all-at-once for computational speed-up
-        if isinstance(x,Samples):
-            out = np.zeros((func_range_geometry.par_dim, x.Ns))
-            # Recursively apply func to each sample
-            for idx, item in enumerate(x):
-                out[:,idx] = self._apply_func(func,
-                                              func_range_geometry,
-                                              func_domain_geometry,
-                                              item, is_par=True,
-                                              **kwargs)
-            return Samples(out, geometry=func_range_geometry)
-        
-        # store if input x is CUQIarray
-        is_CUQIarray = type(x) is CUQIarray
+        if any(isinstance(x, Samples) for x in kwargs.values()):
+            return self._handle_case_when_model_input_is_samples(func, fwd, **kwargs)
+
+        # store if any input x is CUQIarray
+        is_CUQIarray = any(isinstance(x, CUQIarray) for x in kwargs.values())
 
-        x = self._2fun(x, func_domain_geometry, is_par=is_par)
-        out = func(x, **kwargs)
+        # Convert input to function values
+        kwargs = self._2fun(geometry=func_domain_geometry, is_par=is_par, **kwargs)
 
-        # Return output as parameters 
-        # (and wrapped in CUQIarray if input was CUQIarray)
-        return self._2par(out, func_range_geometry, 
-                                    to_CUQIarray=is_CUQIarray)
+        # Apply the function
+        out = func(**kwargs)
 
-    def _parse_args_add_to_kwargs(self, *args, **kwargs):
-        """ Private function that parses the input arguments of the model and adds them as keyword arguments matching the non default arguments of the forward function. """
+        # Return output as parameters
+        # (wrapped in CUQIarray if any input was CUQIarray)
+        return self._2par(
+            geometry=func_range_geometry, to_CUQIarray=is_CUQIarray, **{"out": out}
+        )["out"]
 
+    def _handle_case_when_model_input_is_samples(self, func=None, fwd=True, **kwargs):
+        """Private function that calls apply_func for samples in the
+        Samples object(s).
+        """
+        # All kwargs should be Samples objects
+        if not all(isinstance(x, Samples) for x in kwargs.values()):
+            raise TypeError(
+                "If applying the function to Samples, all inputs should be Samples."
+            )
+
+        # All Samples objects should have the same number of samples
+        Ns = list(kwargs.values())[0].Ns
+        if not all(x.Ns == Ns for x in kwargs.values()):
+            raise ValueError(
+                "If applying the function to Samples, all inputs should have the same number of samples."
+            )
+
+        # Specify the range dimension of the function
+        range_dim = self.range_dim if fwd else self.domain_dim
+
+        # Create empty array to store the output
+        out = np.zeros((range_dim, Ns))
+
+        # Recursively apply func to each sample
+        for i in range(Ns):
+            kwargs_i = {
+                k: CUQIarray(v.samples[..., i], is_par=v.is_par, geometry=v.geometry)
+                for k, v in kwargs.items()
+            }
+            out[:, i] = self._apply_func(func=func, fwd=fwd, **kwargs_i)
+        # Specify the range geometries of the function
+        func_range_geometry = self.range_geometry if fwd else self.domain_geometry
+        return Samples(out, geometry=func_range_geometry)
+
+    def _parse_args_add_to_kwargs(
+        self, *args, is_par=True, non_default_args=None, map_name="model", **kwargs
+    ):
+        """ Private function that parses the input arguments and adds them as
+        keyword arguments matching (the order of) the non default arguments of
+        the forward function or other specified non_default_args list.
+        """
+        # If non_default_args is not specified, use the non_default_args of the
+        # model
+        if non_default_args is None:
+            non_default_args = self._non_default_args
+
+        # Either args or kwargs can be provided but not both
+        if len(args) > 0 and len(kwargs) > 0:
+            raise ValueError(
+                "The "
+                + map_name.lower()
+                + " input is specified both as positional and keyword arguments. This is not supported."
+            )
+
+        len_input = len(args) + len(kwargs)
+
+        # If partial evaluation, make sure input is not of type Samples
+        if len_input < len(non_default_args):
+            # If the argument is a Sample object, splitting or partial
+            # evaluation of the model is not supported
+            temp_args = args if len(args) > 0 else list(kwargs.values())
+            if any(isinstance(arg, Samples) for arg in temp_args):
+                raise ValueError(("When using Samples objects as input, the"
+                                +" user should provide a Samples object for"
+                                +f" each non_default_args {non_default_args}"
+                                +" of the model. That is, partial evaluation"
+                                +" or splitting is not supported for input"
+                                +" of type Samples."))
+
+        # If args are given, add them to kwargs
         if len(args) > 0:
 
-            if len(kwargs) > 0:
-                raise ValueError("The model input is specified both as positional and keyword arguments. This is not supported.")
-                
-            if len(args) != len(self._non_default_args):
-                raise ValueError("The number of positional arguments does not match the number of non-default arguments of the model.")
-            
+            # Check if the input is for multiple input case and is stacked,
+            # then split it
+            if len(args) < len(non_default_args):
+                args = self._split_in_case_of_stacked_args(*args, is_par=is_par)
+
             # Add args to kwargs following the order of non_default_args
             for idx, arg in enumerate(args):
-                kwargs[self._non_default_args[idx]] = arg
+                kwargs[non_default_args[idx]] = arg
+    
+        # Check kwargs matches non_default_args
+        if not (set(list(kwargs.keys())) <= set(non_default_args)):
+            if map_name == "gradient":
+                error_msg = f"The gradient input is specified by a direction and keywords arguments {list(kwargs.keys())} that does not match the non_default_args of the model {non_default_args}."
+            else:
+                error_msg = (
+                    "The "
+                    + map_name.lower()
+                    + f" input is specified by keywords arguments {list(kwargs.keys())} that does not match the non_default_args of the "
+                    + map_name
+                    + f" {non_default_args}."
+                )
+
+            raise ValueError(error_msg)
+
+        # Make sure order of kwargs is the same as non_default_args
+        kwargs = {k: kwargs[k] for k in non_default_args if k in kwargs}
 
         return kwargs
-        
+
+    def _split_in_case_of_stacked_args(self, *args, is_par=True):
+        """Private function that checks if the input args is a stacked
+        CUQIarray or numpy array and splits it into multiple arguments based on
+        the domain geometry of the model. Otherwise, it returns the input args
+        unchanged."""
+
+        # Check conditions for splitting and split if all conditions are met
+        is_CUQIarray = isinstance(args[0], CUQIarray)
+        is_numpy_array = isinstance(args[0], np.ndarray)
+
+        if ((is_CUQIarray or is_numpy_array) and
+           is_par and
+           len(args) == 1 and
+           args[0].shape == (self.domain_dim,) and
+           isinstance(self.domain_geometry, cuqi.experimental.geometry._ProductGeometry)):
+            # Split the stacked input
+            split_args = np.split(args[0], self.domain_geometry.stacked_par_split_indices)
+            # Convert split args to CUQIarray if input is CUQIarray
+            if is_CUQIarray:
+                split_args = [
+                    CUQIarray(arg, is_par=True, geometry=self.domain_geometry.geometries[i])
+                    for i, arg in enumerate(split_args)
+                ]
+            return split_args
+
+        else:
+            return args
+
     def forward(self, *args, is_par=True, **kwargs):
         """ Forward function of the model.
-        
-        Forward converts the input to function values (if needed) using the domain geometry of the model.
-        Forward converts the output function values to parameters using the range geometry of the model.
+
+        Forward converts the input to function values (if needed) using the domain geometry of the model. Then it applies the forward operator to the function values and converts the output to parameters using the range geometry of the model.
 
         Parameters
         ----------
-        *args : ndarray or cuqi.array.CUQIarray
-            The model input.
+        *args : ndarrays or cuqi.array.CUQIarray objects or cuqi.samples.Samples objects
+            Positional arguments for the forward operator. The forward operator input can be specified as either positional arguments or keyword arguments but not both.
 
-        is_par : bool
-            If True the input is assumed to be parameters.
-            If False the input is assumed to be function values.
-        
-        **kwargs : keyword arguments for model input.
-            Keywords must match the names of the non_default_args of the model.
+            If the input is specified as positional arguments, the order of the arguments should match the non_default_args of the model.
+
+        is_par : bool or a tuple of bools
+            If True, the inputs in `args` or `kwargs` are assumed to be parameters.
+            If False, the inputs in `args` or `kwargs` are assumed to be function values.
+            If `is_par` is a tuple of bools, the inputs are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
+
+        **kwargs : keyword arguments
+            keyword arguments for the forward operator. The forward operator input can be specified as either positional arguments or keyword arguments but not both.
+
+            If the input is specified as keyword arguments, the keys should match the non_default_args of the model.
 
         Returns
         -------
-        ndarray or cuqi.array.CUQIarray
+        ndarray or cuqi.array.CUQIarray or cuqi.samples.Samples object
             The model output. Always returned as parameters.
         """
 
-        kwargs = self._parse_args_add_to_kwargs(*args, **kwargs)
-
-        # Check kwargs matches non_default_args
-        if set(list(kwargs.keys())) != set(self._non_default_args):
-            raise ValueError(f"The model input is specified by a keywords arguments {kwargs.keys()} that does not match the non_default_args of the model {self._non_default_args}.")
-
-        # For now only support one input
-        if len(kwargs) > 1:
-            raise ValueError("The model input is specified by more than one argument. This is not supported.")
-
-        # Get input matching the non_default_args
-        x = kwargs[self._non_default_args[0]]
-
-        # If input is a distribution, we simply change the parameter name of model to match the distribution name
-        if isinstance(x, cuqi.distribution.Distribution):
-            if x.dim != self.domain_dim:
-                raise ValueError("Attempting to match parameter name of Model with given distribution, but distribution dimension does not match model domain dimension.")
-            new_model = copy(self)
-            new_model._non_default_args = [x.name] # Defaults to x if distribution had no name
-            return new_model
+        # Add args to kwargs and ensure the order of the arguments matches the
+        # non_default_args of the forward function
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, **kwargs, is_par=is_par, map_name="model"
+        )
+        # Extract args from kwargs
+        args = list(kwargs.values())
+
+        if len(kwargs) == 0:
+            return self
+
+        partial_arguments = len(kwargs) < len(self._non_default_args)
+
+        # If input is a distribution, we simply change the parameter name of
+        # model to match the distribution name
+        if all(isinstance(x, cuqi.distribution.Distribution)
+               for x in kwargs.values()):
+            if partial_arguments:
+                raise ValueError(
+                    "Partial evaluation of the model is not supported for distributions."
+                )
+            return self._handle_case_when_model_input_is_distributions(kwargs)
 
         # If input is a random variable, we handle it separately
-        if isinstance(x, cuqi.experimental.algebra.RandomVariable):
-            return self._handle_random_variable(x)
-        
+        elif all(isinstance(x, cuqi.experimental.algebra.RandomVariable)
+               for x in kwargs.values()):
+            if partial_arguments:
+                raise ValueError(
+                    "Partial evaluation of the model is not supported for random variables."
+                )
+            return self._handle_case_when_model_input_is_random_variables(kwargs)
+
         # If input is a Node from internal abstract syntax tree, we let the Node handle the operation
         # We use NotImplemented to indicate that the operation is not supported from the Model class
         # in case of operations such as "@" that can be interpreted as both __matmul__ and __rmatmul__
         # the operation may be delegated to the Node class.
-        if isinstance(x, cuqi.experimental.algebra.Node):
+        elif any(isinstance(args_i, cuqi.experimental.algebra.Node) for args_i in args):
             return NotImplemented
 
+        # if input is partial, we create a new model with the partial input
+        if partial_arguments:
+            # Create is_par_partial from the is_par to contain only the relevant parts
+            if isinstance(is_par, (list, tuple)):
+                is_par_partial = tuple(
+                    is_par[i]
+                    for i in range(self.number_of_inputs)
+                    if self._non_default_args[i] in kwargs.keys()
+                )
+            else:
+                is_par_partial = is_par
+            # Build a partial model with the given kwargs
+            partial_model = self._build_partial_model(kwargs, is_par_partial)
+            return partial_model
+
         # Else we apply the forward operator
-        return self._apply_func(self._forward_func,
-                                self.range_geometry,
-                                self.domain_geometry,
-                                x, is_par)
+        # if model has _original_non_default_args, we use it to replace the
+        # kwargs keys so that it matches self._forward_func signature
+        if hasattr(self, '_original_non_default_args'):
+            kwargs = {k:v for k,v in zip(self._original_non_default_args, args)}
+        return self._apply_func(func=self._forward_func,
+                                fwd=True,
+                                is_par=is_par,
+                                **kwargs)
+
+    def _correct_distribution_dimension(self, distributions):
+        """Private function that checks if the dimension of the
+        distributions matches the domain dimension of the model."""
+        if len(distributions) == 1:
+            return list(distributions)[0].dim == self.domain_dim
+        elif len(distributions) > 1 and isinstance(
+            self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+        ):
+            return all(
+                d.dim == self.domain_geometry.par_dim_list[i]
+                for i, d in enumerate(distributions)
+            )
+        else:
+            return False
+
+    def _build_partial_model(self, kwargs, is_par):
+        """Private function that builds a partial model substituting the given
+        keyword arguments with their values. The created partial model will have
+        as inputs the non-default arguments that are not in the kwargs."""
+
+        # Extract args from kwargs
+        args = list(kwargs.values())
+
+        # Define original_non_default_args which represents the complete list of
+        # non-default arguments of the forward function.
+        original_non_default_args = (
+            self._original_non_default_args
+            if hasattr(self, "_original_non_default_args")
+            else self._non_default_args
+        )
+
+        if hasattr(self, "_original_non_default_args"):
+            # Split the _original_non_default_args into two lists:
+            # 1. reduced_original_non_default_args: the _original_non_default_args
+            # corresponding to the _non_default_args that are not in kwargs
+            # 2. substituted_non_default_args: the _original_non_default_args
+            # corresponding to the _non_default_args that are in kwargs
+            reduced_original_non_default_args = [
+                original_non_default_args[i]
+                for i in range(self.number_of_inputs)
+                if self._non_default_args[i] not in kwargs.keys()
+            ]
+            substituted_non_default_args = [
+                original_non_default_args[i]
+                for i in range(self.number_of_inputs)
+                if self._non_default_args[i] in kwargs.keys()
+            ]
+            # Replace the keys in kwargs with the substituted_non_default_args
+            # so that the kwargs match the signature of the _forward_func
+            kwargs = {k: v for k, v in zip(substituted_non_default_args, args)}
+
+        # Create a partial domain geometry with the geometries corresponding
+        # to the non-default arguments that are not in kwargs (remaining
+        # unspecified inputs)
+        partial_domain_geometry = cuqi.experimental.geometry._ProductGeometry(
+            *[
+                self.domain_geometry.geometries[i]
+                for i in range(self.number_of_inputs)
+                if original_non_default_args[i] not in kwargs.keys()
+            ]
+        )
+
+        if len(partial_domain_geometry.geometries) == 1:
+            partial_domain_geometry = partial_domain_geometry.geometries[0]
+
+        # Create a domain geometry with the geometries corresponding to the
+        # non-default arguments that are specified
+        substituted_domain_geometry = cuqi.experimental.geometry._ProductGeometry(
+            *[
+                self.domain_geometry.geometries[i]
+                for i in range(self.number_of_inputs)
+                if original_non_default_args[i] in kwargs.keys()
+            ]
+        )
+
+        if len(substituted_domain_geometry.geometries) == 1:
+            substituted_domain_geometry = substituted_domain_geometry.geometries[0]
+
+        # Create new model with partial input
+        # First, we convert the input to function values
+        kwargs = self._2fun(geometry=substituted_domain_geometry, is_par=is_par, **kwargs)
+
+        # Second, we create a partial function for the forward operator
+        partial_forward = partial(self._forward_func, **kwargs)
+
+        # Third, if applicable, we create a partial function for the gradient
+        if isinstance(self._gradient_func, tuple):
+            # If gradient is a tuple, we create a partial function for each
+            # gradient function in the tuple
+            partial_gradient = tuple(
+                (
+                    partial(self._gradient_func[i], **kwargs)
+                    if self._gradient_func[i] is not None
+                    else None
+                )
+                for i in range(self.number_of_inputs)
+                if original_non_default_args[i] not in kwargs.keys()
+            )
+            if len(partial_gradient) == 1:
+                partial_gradient = partial_gradient[0]
+
+        elif callable(self._gradient_func):
+            raise NotImplementedError(
+                "Partial forward model is only supported for gradient/jacobian functions that are tuples of callable functions."
+            )
 
-    def __call__(self, *args, **kwargs):
-        return self.forward(*args, **kwargs)
+        else:
+            partial_gradient = None
+
+        # Lastly, we create the partial model with the partial forward
+        # operator (we set the gradient function later)
+        partial_model = Model(
+            forward=partial_forward,
+            range_geometry=self.range_geometry,
+            domain_geometry=partial_domain_geometry,
+        )
+
+        # Set the _original_non_default_args (if applicable) and
+        # _stored_non_default_args of the partial model
+        if hasattr(self, "_original_non_default_args"):
+            partial_model._original_non_default_args = reduced_original_non_default_args
+        partial_model._stored_non_default_args = [
+            self._non_default_args[i]
+            for i in range(self.number_of_inputs)
+            if original_non_default_args[i] not in kwargs.keys()
+        ]
+
+        # Set the gradient function of the partial model
+        partial_model._check_correct_gradient_jacobian_form(
+            partial_gradient, "gradient"
+        )
+        partial_model._gradient_func = partial_gradient
+
+        return partial_model
+
+    def _handle_case_when_model_input_is_distributions(self, kwargs):
+        """Private function that handles the case of the input being a
+        distribution or multiple distributions."""
+
+        if not self._correct_distribution_dimension(kwargs.values()):
+            raise ValueError(
+                "Attempting to match parameter name of Model with given distribution(s), but distribution(s) dimension(s) does not match model input dimension(s)."
+            )
+        new_model = copy(self)
+
+        # Store the original non_default_args of the model
+        new_model._original_non_default_args = (
+            self._original_non_default_args
+            if hasattr(self, "_original_non_default_args")
+            else self._non_default_args
+        )
+
+        # Update the non_default_args of the model to match the distribution
+        # names. Defaults to x in the case of only one distribution that has no
+        # name
+        new_model._stored_non_default_args = [x.name for x in kwargs.values()]
+
+        # If there is a repeated name, raise an error
+        if len(set(new_model._stored_non_default_args)) != len(
+            new_model._stored_non_default_args
+        ):
+            raise ValueError(
+                "Attempting to match parameter name of Model with given distributions, but distribution names are not unique. Please provide unique names for the distributions."
+            )
+
+        return new_model
+
+    def _handle_case_when_model_input_is_random_variables(self, kwargs):
+        """ Private function that handles the case of the input being a random variable. """
+        # If random variable is not a leaf-type node (e.g. internal node) we return NotImplemented
+        if any(not isinstance(x.tree, cuqi.experimental.algebra.VariableNode) for x in kwargs.values()):
+            return NotImplemented        
+
+        # Extract the random variable distributions and check dimensions consistency with domain geometry
+        distributions = [value.distribution for value in kwargs.values()]
+        if not self._correct_distribution_dimension(distributions):
+            raise ValueError("Attempting to match parameter name of Model with given random variable(s), but random variable dimension(s) does not match model input dimension(s).")
+
+        new_model = copy(self)
 
-    def gradient(self, direction, wrt, is_direction_par=True, is_wrt_par=True):
-        """ Gradient of the forward operator (Direction-Jacobian product)
+        # Store the original non_default_args of the model
+        new_model._original_non_default_args = self._non_default_args
 
-        For non-linear models the gradient is computed using the
-        forward operator and the Jacobian of the forward operator.
+        # Update the non_default_args of the model to match the random variable
+        # names. Defaults to x in the case of only one random variable that has
+        # no name
+        new_model._stored_non_default_args = [x.name for x in distributions]
+
+        # If there is a repeated name, raise an error
+        if len(set(new_model._stored_non_default_args)) != len(
+            new_model._stored_non_default_args
+        ):
+            raise ValueError(
+                "Attempting to match parameter name of Model with given random variables, but random variables names are not unique. Please provide unique names for the random variables."
+            )
+
+        return new_model
+
+    def gradient(
+        self, direction, *args, is_direction_par=True, is_var_par=True, **kwargs
+    ):
+        """Gradient of the forward operator (Direction-Jacobian product)
+
+        The gradient computes the Vector-Jacobian product (VJP) of the forward operator evaluated at the given model input and the given vector (direction).
 
         Parameters
         ----------
-        direction : ndarray
-            The direction to compute the gradient. The Jacobian is applied to this direction.
+        direction : ndarray or cuqi.array.CUQIarray
+            The direction at which to compute the gradient.
+
+        *args : ndarrays or cuqi.array.CUQIarray objects
+            Positional arguments for the values at which to compute the gradient. The gradient operator input can be specified as either positional arguments or keyword arguments but not both.
 
-        wrt : ndarray
-            The point to compute the Jacobian at. This is only used for non-linear models.
+            If the input is specified as positional arguments, the order of the arguments should match the non_default_args of the model.
 
         is_direction_par : bool
             If True, `direction` is assumed to be parameters.
             If False, `direction` is assumed to be function values.
 
-        is_wrt_par : bool
-            If True, `wrt` is assumed to be parameters.
-            If False, `wrt` is assumed to be function values.
-        
+        is_var_par : bool or a tuple of bools
+            If True, the inputs in `args` or `kwargs` are assumed to be parameters.
+            If False, the inputs in `args` or `kwargs` are assumed to be function values.
+            If `is_var_par` is a tuple of bools, the inputs in `args` or `kwargs` are assumed to be parameters or function values based on the corresponding boolean value in the tuple.
         """
-        # Obtain the parameters representation of wrt and raise an error if it
-        # cannot be obtained
-        error_message = \
-            "For the gradient to be computed, is_wrt_par needs " +\
-            "to be True and wrt needs to be parameter value, not function " +\
-            "value. Alternatively, the model domain_geometry: "+\
-            f"{self.domain_geometry} " +\
-            "should have an implementation of the method fun2par"
+        # Add args to kwargs and ensure the order of the arguments matches the
+        # non_default_args of the forward function
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, **kwargs, is_par=is_var_par, map_name="gradient"
+        )
+
+        # Obtain the parameters representation of the variables and raise an
+        # error if it cannot be obtained
+        error_message = (
+            "For the gradient to be computed, is_var_par needs to be True and the variables in kwargs needs to be parameter value, not function value. Alternatively, the model domain_geometry:"
+            + f" {self.domain_geometry} "
+            + "should have an implementation of the method fun2par"
+        )
         try:
-            wrt_par = self._2par(wrt, 
-                                 geometry=self.domain_geometry,
-                                 is_par=is_wrt_par,
-                                 to_CUQIarray=False,
-                                 )
+            kwargs_par = self._2par(
+                geometry=self.domain_geometry,
+                is_par=is_var_par,
+                to_CUQIarray=False,
+                **kwargs,
+            )
         # NotImplementedError will be raised if fun2par of the geometry is not
         # implemented and ValueError will be raised when imap is not set in
         # MappedGeometry
         except ValueError as e:
-            raise ValueError(error_message +
-                             " ,including an implementation of imap for " +
-                             "MappedGeometry")
+            raise ValueError(
+                error_message
+                + " ,including an implementation of imap for MappedGeometry"
+            )
         except NotImplementedError as e:
             raise NotImplementedError(error_message)
-        
-        # Check for other errors that may prevent computing the gradient
-        self._check_gradient_can_be_computed(direction, wrt)
-
-        wrt = self._2fun(wrt, self.domain_geometry, is_par=is_wrt_par)
-
-        # Store if the input direction is CUQIarray
-        is_direction_CUQIarray = type(direction) is CUQIarray
-
-        direction = self._2fun(direction,
-                               self.range_geometry,
-                               is_par=is_direction_par)
 
-        grad = self._gradient_func(direction, wrt)
-        grad_is_par = False # Assume gradient is function values
-        
-        # If domain_geometry has gradient attribute, we apply it to the gradient
-        # The gradient returned by the domain_geometry.gradient is assumed to be
-        # parameters
-        if hasattr(self.domain_geometry, 'gradient'):
-            grad = self.domain_geometry.gradient(grad, wrt_par)
-            grad_is_par = True # Gradient is parameters
-
-        # we convert the computed gradient to parameters
-        grad = self._2par(grad,
-                          self.domain_geometry,
-                          to_CUQIarray=is_direction_CUQIarray,
-                          is_par=grad_is_par)
+        # Check for other errors that may prevent computing the gradient
+        self._check_gradient_can_be_computed(direction, kwargs)
+
+        # Also obtain the function values representation of the variables
+        kwargs_fun = self._2fun(
+            geometry=self.domain_geometry, is_par=is_var_par, **kwargs
+        )
+
+        # Store if any of the inputs is a CUQIarray
+        to_CUQIarray = isinstance(direction, CUQIarray) or any(
+            isinstance(x, CUQIarray) for x in kwargs_fun.values()
+        )
+
+        # Turn to_CUQIarray to a tuple of bools of the same length as kwargs_fun
+        to_CUQIarray = tuple([to_CUQIarray] * len(kwargs_fun))
+
+        # Convert direction to function value
+        direction_fun = self._2fun(
+            direction=direction, geometry=self.range_geometry, is_par=is_direction_par
+        )
+
+        # If model has _original_non_default_args, we use it to replace the
+        # kwargs keys so that it matches self._gradient_func signature
+        if hasattr(self, '_original_non_default_args'):
+            args_fun = list(kwargs_fun.values())
+            kwargs_fun = {
+                k: v for k, v in zip(self._original_non_default_args, args_fun)
+            }
+        # Append the direction to the kwargs_fun as first input
+        kwargs_fun_grad_input = {**direction_fun, **kwargs_fun}
+
+        # Form 1 of gradient (callable)
+        if callable(self._gradient_func):
+            grad = self._gradient_func(**kwargs_fun_grad_input)
+            grad_is_par = False  # Assume gradient is function value
+
+        # Form 2 of gradient (tuple of callables)
+        elif isinstance(self._gradient_func, tuple):
+            grad = []
+            for i, grad_func in enumerate(self._gradient_func):
+                if grad_func is not None:
+                    grad.append(grad_func(**kwargs_fun_grad_input))
+                else:
+                    grad.append(None)
+                    # set the ith item of to_CUQIarray tuple to False
+                    # because the ith gradient is None
+                    to_CUQIarray = to_CUQIarray[:i] + (False,) + to_CUQIarray[i + 1 :]
+            grad_is_par = False  # Assume gradient is function value
+
+        grad = self._apply_chain_rule_to_account_for_domain_geometry_gradient(
+            kwargs_par, grad, grad_is_par, to_CUQIarray
+        )
+
+        if len(grad) == 1:
+            return list(grad.values())[0]
+        elif self._gradient_output_stacked:
+            return np.hstack(
+                [
+                    (
+                        v.to_numpy()
+                        if isinstance(v, CUQIarray)
+                        else force_ndarray(v, flatten=True)
+                    )
+                    for v in list(grad.values())
+                ]
+            )
 
         return grad
-    
-    def _check_gradient_can_be_computed(self, direction, wrt):
-        """ Private function that checks if the gradient can be computed. By
+
+    def _check_gradient_can_be_computed(self, direction, kwargs_dict):
+        """Private function that checks if the gradient can be computed. By
         raising an error for the cases where the gradient cannot be computed."""
 
         # Raise an error if _gradient_func function is not set
         if self._gradient_func is None:
             raise NotImplementedError("Gradient is not implemented for this model.")
-        
-        # Raise error if either the direction or wrt are Samples object
-        if isinstance(direction, Samples) or isinstance(wrt, Samples):
-            raise ValueError("cuqi.samples.Samples input values for arguments `direction` and `wrt` are not supported")
-        
+
+        # Raise an error if either the direction or kwargs are Samples objects
+        if isinstance(direction, Samples) or any(
+            isinstance(x, Samples) for x in kwargs_dict.values()
+        ):
+            raise NotImplementedError(
+                "Gradient is not implemented for input of type Samples."
+            )
+
         # Raise an error if range_geometry is not in the list returned by
-        # `_get_identity_geometries()`. i.e. The Jacobian of its 
-        # par2fun map is not identity.  
-        #TODO: Add range geometry gradient to the chain rule 
+        # `_get_identity_geometries()`. i.e. The Jacobian of its
+        # par2fun map is not identity.
+        # TODO: Add range geometry gradient to the chain rule
         if not type(self.range_geometry) in _get_identity_geometries():
-            raise NotImplementedError("Gradient not implemented for model {} with range geometry {}".format(self,self.range_geometry)) 
-        
-        # Raise an error if domain_geometry does not have gradient attribute and
-        # is not in the list returned by `_get_identity_geometries()`. i.e. the
-        # Jacobian of its par2fun map is not identity.  
-        if not hasattr(self.domain_geometry, 'gradient') and \
-            not type(self.domain_geometry) in _get_identity_geometries():
-            raise NotImplementedError("Gradient not implemented for model {} with domain geometry {}".format(self,self.domain_geometry))
-
-    def _handle_random_variable(self, x):
-        """ Private function that handles the case of the input being a random variable. """
-        # If random variable is not a leaf-type node (e.g. internal node) we return NotImplemented
-        if not isinstance(x.tree, cuqi.experimental.algebra.VariableNode):
-            return NotImplemented        
-        
-        # In leaf-type node case we simply change the parameter name of model to match the random variable name
-        dist = x.distribution
-        if dist.dim != self.domain_dim:
-            raise ValueError("Attempting to match parameter name of Model with given random variable, but random variable dimension does not match model domain dimension.")
-        
-        new_model = copy(self)
-        new_model._non_default_args = [dist.name]
-        return new_model
+            raise NotImplementedError(
+                "Gradient is not implemented for model {} with range geometry {}. You can use one of the geometries in the list {}.".format(
+                    self,
+                    self.range_geometry,
+                    [i_g.__name__ for i_g in _get_identity_geometries()],
+                )
+            )
+
+        # Raise an error if domain_geometry (or its components in case of
+        # _ProductGeometry) does not have gradient attribute and is not in the
+        # list returned by `_get_identity_geometries()`. i.e. The Jacobian of its
+        # par2fun map is not identity.
+        domain_geometries = (
+            self.domain_geometry.geometries
+            if isinstance(
+                self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+            )
+            else [self.domain_geometry]
+        )
+        for domain_geometry in domain_geometries:
+            if (
+                not hasattr(domain_geometry, "gradient")
+                and not type(domain_geometry) in _get_identity_geometries()
+            ):
+                raise NotImplementedError(
+                    "Gradient is not implemented for model \n{}\nwith domain geometry (or domain geometry component) {}. The domain geometries should have gradient method or be from the geometries in the list {}.".format(
+                        self,
+                        domain_geometry,
+                        [i_g.__name__ for i_g in _get_identity_geometries()],
+                    )
+                )
+
+    def _apply_chain_rule_to_account_for_domain_geometry_gradient(self,
+                                                                  kwargs_par,
+                                                                  grad,
+                                                                  grad_is_par,
+                                                                  to_CUQIarray):
+        """ Private function that applies the chain rule to account for the
+        gradient of the domain geometry. That is, it computes the gradient of
+        the function values with respect to the parameters values."""
+        # Create list of domain geometries
+        geometries = (
+            self.domain_geometry.geometries
+            if isinstance(self.domain_geometry, cuqi.experimental.geometry._ProductGeometry)
+            else [self.domain_geometry]
+        )
+
+        # turn grad_is_par to a tuple of bools if it is not already
+        if isinstance(grad_is_par, bool):
+            grad_is_par = tuple([grad_is_par]*self.number_of_inputs)
+
+        # If the domain geometry is a _ProductGeometry and the gradient is
+        # stacked, split it
+        if (
+            isinstance(
+                self.domain_geometry, cuqi.experimental.geometry._ProductGeometry
+            )
+            and not isinstance(grad, (list, tuple))
+            and isinstance(grad, np.ndarray)
+        ):
+            grad = np.split(grad, self.domain_geometry.stacked_par_split_indices)
+
+        # If the domain geometry is not a _ProductGeometry, turn grad into a
+        # list of length 1, so that we can iterate over it
+        if not isinstance(self.domain_geometry, cuqi.experimental.geometry._ProductGeometry):
+            grad = [grad]
+
+        # apply the gradient of each geometry component
+        grad_kwargs = {}
+        for i, (k, v_par) in enumerate(kwargs_par.items()):
+            if hasattr(geometries[i], 'gradient') and grad[i] is not None:
+                grad_kwargs[k] = geometries[i].gradient(grad[i], v_par)
+                # update the ith component of grad_is_par to True
+                grad_is_par = grad_is_par[:i] + (True,) + grad_is_par[i+1:]
+            else:
+                grad_kwargs[k] = grad[i]
+
+        # convert the computed gradient to parameters
+        grad = self._2par(geometry=self.domain_geometry,
+                          to_CUQIarray=to_CUQIarray,
+                          is_par=grad_is_par,
+                          **grad_kwargs)
+
+        return grad
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
 
     def __len__(self):
         return self.range_dim
 
     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))
-
+        kwargs = {}
+        if self.number_of_inputs > 1:
+            pad = " " * len("CUQI {}: ".format(self.__class__.__name__))
+            kwargs["pad"]=pad  
+        return "CUQI {}: {} -> {}.\n    Forward parameters: {}.".format(self.__class__.__name__,self.domain_geometry.__repr__(**kwargs),self.range_geometry,self._non_default_args)
 
 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`.
@@ -533,7 +1277,7 @@ class AffineModel(Model):
         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
@@ -559,11 +1303,50 @@ class AffineModel(Model):
         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")
 
+        # If linear operator is of type Model, it needs to be a LinearModel
+        if isinstance(linear_operator, Model) and not isinstance(
+            linear_operator, LinearModel
+        ):
+            raise TypeError(
+                "The linear operator should be a LinearModel object, a callable function or a matrix."
+            )
+
+        # If the adjoint operator is of type Model, it needs to be a LinearModel
+        if isinstance(linear_operator_adjoint, Model) and not isinstance(
+            linear_operator_adjoint, LinearModel
+        ):
+            raise TypeError(
+                "The adjoint linear operator should be a LinearModel object, a callable function or a matrix."
+            )
+
+        # Additional checks if the linear_operator is not a LinearModel:
+        if not isinstance(linear_operator, LinearModel):
+            # Ensure the linear operator has exactly one input argument
+            if len(cuqi.utilities.get_non_default_args(linear_operator)) != 1:
+                raise ValueError(
+                    "The linear operator should have exactly one input argument."
+                )
+            # Ensure the adjoint linear operator has exactly one input argument
+            if (
+                linear_operator_adjoint is not None
+                and len(cuqi.utilities.get_non_default_args(linear_operator_adjoint))
+                != 1
+            ):
+                raise ValueError(
+                    "The adjoint linear operator should have exactly one input argument."
+                )
+
         # 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.")
 
+        # Store linear operator privately
+        # Note: we need to set the _linear_operator before calling the
+        # super().__init__() because it is needed when calling the property
+        # _non_default_args within the super().__init__()
+        self._linear_operator = linear_operator
+
         # Initialize Model class
         super().__init__(linear_operator, range_geometry, domain_geometry)
 
@@ -573,20 +1356,27 @@ class AffineModel(Model):
         # 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)
+        self._gradient_func = lambda direction, *args, **kwargs: 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)
+        # Set stored_non_default_args to None
+        self._stored_non_default_args = None
+
+    @property
+    def _non_default_args(self):
+        if self._stored_non_default_args is None:
+            # Use arguments from user's callable linear operator
+            self._stored_non_default_args = cuqi.utilities.get_non_default_args(
+                self._linear_operator
+            )
+        return self._stored_non_default_args
 
     @property
     def shift(self):
@@ -599,19 +1389,35 @@ class AffineModel(Model):
         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 _forward_func_no_shift(self, *args, is_par=True, **kwargs):
+        """Helper function for computing the forward operator without the shift."""
+        # convert args to kwargs
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, **kwargs, map_name="model", is_par=is_par
+        )
+        args = list(kwargs.values())
+        # if model has _original_non_default_args, we use it to replace the
+        # kwargs keys so that it matches self._linear_operator signature
+        if hasattr(self, '_original_non_default_args'):
+            kwargs = {k:v for k,v in zip(self._original_non_default_args, args)}
+        return self._apply_func(self._linear_operator, **kwargs, is_par=is_par)
+
+    def _adjoint_func_no_shift(self, *args, is_par=True, **kwargs):
+        """Helper function for computing the adjoint operator without the shift."""
+        # convert args to kwargs
+        kwargs = self._parse_args_add_to_kwargs(
+            *args,
+            **kwargs,
+            map_name='adjoint',
+            is_par=is_par,
+            non_default_args=cuqi.utilities.get_non_default_args(
+                self._linear_operator_adjoint
+            ),
+        )
+        return self._apply_func(
+            self._linear_operator_adjoint, **kwargs, is_par=is_par, fwd=False
+        )
 
-    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.
@@ -677,50 +1483,67 @@ class LinearModel(AffineModel):
     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.
     """
-    
+
     def __init__(self, forward, adjoint=None, range_geometry=None, domain_geometry=None):
 
-        #Initialize as AffineModel with shift=0
+        # Initialize as AffineModel with shift=0
         super().__init__(forward, 0, adjoint, range_geometry, domain_geometry)
 
-    def adjoint(self, y, is_par=True):
+    def adjoint(self, *args, is_par=True, **kwargs):
         """ 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.
+        Adjoint converts the input to function values (if needed) using the range geometry of the model then applies the adjoint operator to the function values and converts the output function values to parameters using the domain geometry of the model.
 
         Parameters
         ----------
-        y : ndarray or cuqi.array.CUQIarray
-            The adjoint model input.
+        *args : ndarrays or cuqi.array.CUQIarray object
+            Positional arguments for the adjoint operator ( maximum one argument). The adjoint operator input can be specified as either positional arguments or keyword arguments but not both.
+
+        **kwargs : keyword arguments
+            keyword arguments for the adjoint operator (maximum one argument). The adjoint operator input can be specified as either positional arguments or keyword arguments but not both.
+
+            If the input is specified as keyword arguments, the keys should match the non_default_args of the model.
 
         Returns
         -------
         ndarray or cuqi.array.CUQIarray
             The adjoint model output. Always returned as parameters.
         """
+        kwargs = self._parse_args_add_to_kwargs(
+            *args,
+            **kwargs,
+            map_name='adjoint',
+            is_par=is_par,
+            non_default_args=cuqi.utilities.get_non_default_args(
+                self._linear_operator_adjoint
+            ),
+        )
+
+        # length of kwargs should be 1
+        if len(kwargs) > 1:
+            raise ValueError(
+                "The adjoint operator input is specified by more than one argument. This is not supported."
+            )
         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)
+        return self._apply_func(
+            self._linear_operator_adjoint, **kwargs, is_par=is_par, fwd=False
+        )
+
+    def __matmul__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
 
-    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:
-            #TODO: Can we compute this faster while still in sparse format?
+            # TODO: Can we compute this faster while still in sparse format?
             mat = csc_matrix((self.range_dim,0)) #Sparse (m x 1 matrix)
             e = np.zeros(self.domain_dim)
-            
+
             # Stacks sparse matrices on csc matrix
             for i in range(self.domain_dim):
                 e[i] = 1
@@ -728,7 +1551,7 @@ class LinearModel(AffineModel):
                 mat = hstack((mat,col_vec[:,None])) #mat[:,i] = self.forward(e)
                 e[i] = 0
 
-            #Store matrix for future use
+            # Store matrix for future use
             self._matrix = mat
 
             return self._matrix   
@@ -736,61 +1559,129 @@ class LinearModel(AffineModel):
     @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._linear_operator_adjoint,
+            self._linear_operator,
+            self.domain_geometry,
+            self.range_geometry,
+        )
         if self._matrix is not None:
             transpose._matrix = self._matrix.T
         return transpose
-        
+
 
 class PDEModel(Model):
     """
     Model based on an underlying cuqi.pde.PDE.
-    In the forward operation the PDE is assembled, solved and observed.
+    In the forward method the PDE is assembled, solved and observed.
     
     Parameters
     -----------
-    forward : 2D ndarray or callable function.
-        Forward operator assembling, solving and observing the pde.
+    PDE : cuqi.pde.PDE
+        The PDE that specifies the forward operator.
 
-    range_geometry : integer or cuqi.geometry.Geometry (optional)
+    range_geometry : integer or cuqi.geometry.Geometry, optional
         If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
-    domain_geometry : integer or cuqi.geometry.Geometry (optional)
+    domain_geometry : integer or cuqi.geometry.Geometry, optional
         If integer is given, a cuqi.geometry._DefaultGeometry is created with dimension of the integer.
 
 
     :ivar range_geometry: The geometry representing the range.
     :ivar domain_geometry: The geometry representing the domain.
     """
-    def __init__(self, PDE: cuqi.pde.PDE, range_geometry, domain_geometry):
+    def __init__(self, PDE: cuqi.pde.PDE, range_geometry, domain_geometry, **kwargs):
 
         if not isinstance(PDE, cuqi.pde.PDE):
             raise ValueError("PDE needs to be a cuqi PDE.")
+        # PDE needs to be set before calling super().__init__
+        # for the property _non_default_args to work
+        self.pde = PDE
+        self._stored_non_default_args = None
+
+        # If gradient or jacobian is not provided, we create it from the PDE
+        if not np.any([k in kwargs.keys() for k in ["gradient", "jacobian"]]):
+            # Create gradient or jacobian function to pass to the Model based on
+            # the PDE object. The dictionary derivative_kwarg contains the
+            # created function along with the function type (either "gradient"
+            # or "jacobian")
+            derivative_kwarg = self._create_derivative_function()
+            # append derivative_kwarg to kwargs
+            kwargs.update(derivative_kwarg)
+
+        super().__init__(forward=self._forward_func_pde,
+                         range_geometry=range_geometry,
+                         domain_geometry=domain_geometry,
+                         **kwargs)
+
+    @property
+    def _non_default_args(self):
+        if self._stored_non_default_args is None:
+            # extract the non-default arguments of the PDE
+            self._stored_non_default_args = self.pde._non_default_args
 
-        super().__init__(self._forward_func, range_geometry, domain_geometry, gradient=self._gradient_func)
+        return self._stored_non_default_args
 
-        self.pde = PDE
+    def _forward_func_pde(self, **kwargs):
 
-    def _forward_func(self, x):
-        
-        self.pde.assemble(parameter=x)
+        self.pde.assemble(**kwargs)
 
         sol, info = self.pde.solve()
 
         obs = self.pde.observe(sol)
 
         return obs
-    
-    def _gradient_func(self, direction, wrt):
-        """ Compute direction-Jacobian product (gradient) of the model. """
+
+    def _create_derivative_function(self):
+        """Private function that creates the derivative function (gradient or
+        jacobian) based on the PDE object. The derivative function is created as
+        a lambda function that takes the direction and the parameters as input 
+        and returns the gradient or jacobian of the PDE. This private function
+        returns a dictionary with the created function and the function type
+        (either "gradient" or "jacobian")."""
+
         if hasattr(self.pde, "gradient_wrt_parameter"):
-            return self.pde.gradient_wrt_parameter(direction, wrt)
+            # Build the string that will be used to create the lambda function
+            function_str = (
+                "lambda direction, "
+                + ", ".join(self._non_default_args)
+                + ", pde_func: pde_func(direction, "
+                + ", ".join(self._non_default_args)
+                + ")"
+            )
+
+            # create the lambda function from the string
+            function = eval(function_str)
+
+            # create partial function from the lambda function with gradient_wrt_parameter
+            # as the first argument
+            grad_func = partial(function, pde_func=self.pde.gradient_wrt_parameter)
+
+            # Return the gradient function
+            return {"gradient": grad_func}
+
         elif hasattr(self.pde, "jacobian_wrt_parameter"):
-            return direction@self.pde.jacobian_wrt_parameter(wrt)
+            # Build the string that will be used to create the lambda function
+            function_str = (
+                "lambda "
+                + ", ".join(self._non_default_args)
+                + ", pde_func: pde_func( "
+                + ", ".join(self._non_default_args)
+                + ")"
+            )
+
+            # create the lambda function from the string
+            function = eval(function_str)
+
+            # create partial function from the lambda function with jacobian_wrt_parameter
+            # as the first argument
+            jacobian_func = partial(function, pde_func=self.pde.jacobian_wrt_parameter)
+
+            # Return the jacobian function
+            return {"jacobian": jacobian_func}
         else:
-            raise NotImplementedError("Gradient is not implemented for this model.")
+            return {} # empty dictionary if no gradient or jacobian is found
 
     # Add the underlying PDE class name to the repr.
     def __repr__(self) -> str:
         return super().__repr__()+"\n    PDE: {}.".format(self.pde.__class__.__name__)
-