CUQIpy 1.3.0.post0.dev277__tar.gz → 1.3.0.post0.dev292__tar.gz

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/CUQIpy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: CUQIpy
-Version: 1.3.0.post0.dev277
+Version: 1.3.0.post0.dev292
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License:                                  Apache License

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: CUQIpy
-Version: 1.3.0.post0.dev277
+Version: 1.3.0.post0.dev292
 Summary: Computational Uncertainty Quantification for Inverse problems in Python
 Maintainer-email: "Nicolai A. B. Riis" <nabr@dtu.dk>, "Jakob S. Jørgensen" <jakj@dtu.dk>, "Amal M. Alghamdi" <amaal@dtu.dk>, Chao Zhang <chaz@dtu.dk>
 License:                                  Apache License

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/cuqi/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-06-26T10:10:00+0300",
+ "date": "2025-07-04T09:41:59+0300",
  "dirty": false,
  "error": null,
- "full-revisionid": "52c8f8ffc38956c4aee969ccc70cd3db24ee774e",
- "version": "1.3.0.post0.dev277"
+ "full-revisionid": "7536b0d728d704a7ab46bca7ff4756075ba0e573",
+ "version": "1.3.0.post0.dev292"
 }
 '''  # END VERSION_JSON
 

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/cuqi/model/_model.py

@@ -179,7 +179,7 @@ class Model(object):
             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. """
@@ -422,7 +422,7 @@ class Model(object):
             # 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) 
+            # 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)
@@ -496,7 +496,7 @@ class Model(object):
             # 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) 
+            # 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)
@@ -665,7 +665,7 @@ class Model(object):
                 error_msg = (
                     "The "
                     + map_name.lower()
-                    + f" input is specified by a keywords arguments {list(kwargs.keys())} that does not match the non_default_args of the "
+                    + 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}."
                 )
@@ -808,7 +808,11 @@ class Model(object):
         new_model = copy(self)
 
         # Store the original non_default_args of the model
-        new_model._original_non_default_args = self._non_default_args
+        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
@@ -1052,7 +1056,7 @@ class Model(object):
 
         # 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]*len(grad))
+            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
@@ -1451,7 +1455,7 @@ class PDEModel(Model):
     :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.")
@@ -1460,23 +1464,30 @@ class PDEModel(Model):
         self.pde = PDE
         self._stored_non_default_args = None
 
-        super().__init__(self._forward_func, range_geometry, domain_geometry)
+        # 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 = cuqi.utilities.get_non_default_args(
-                self.pde.PDE_form
-            )
-            # remove t from the non-default arguments
-            self._stored_non_default_args = self._non_default_args
-            if "t" in self._non_default_args:
-                self._stored_non_default_args.remove("t")
+            self._stored_non_default_args = self.pde._non_default_args
 
         return self._stored_non_default_args
 
-    def _forward_func(self, **kwargs):
+    def _forward_func_pde(self, **kwargs):
 
         self.pde.assemble(**kwargs)
 
@@ -1486,14 +1497,55 @@ class PDEModel(Model):
 
         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:

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/cuqi/pde/_pde.py

@@ -3,6 +3,7 @@ import scipy
 from inspect import getsource
 from scipy.interpolate import interp1d
 import numpy as np
+from cuqi.utilities import get_non_default_args
 
 
 class PDE(ABC):
@@ -29,6 +30,7 @@ class PDE(ABC):
         self.grid_sol = grid_sol
         self.grid_obs = grid_obs
         self.observation_map = observation_map
+        self._stored_non_default_args = None
 
     @abstractmethod
     def assemble(self, *args, **kwargs):
@@ -64,6 +66,13 @@ class PDE(ABC):
 
         return equal_arrays
 
+    @property
+    def _non_default_args(self):
+        """Returns the non-default arguments of the PDE_form function"""
+        if self._stored_non_default_args is None:
+            self._stored_non_default_args = get_non_default_args(self.PDE_form)
+        return self._stored_non_default_args
+
     @property
     def grid_sol(self):
         if hasattr(self,"_grid_sol"):
@@ -94,6 +103,48 @@ class PDE(ABC):
     def grids_equal(self):
         return self._grids_equal
 
+    def _parse_args_add_to_kwargs(
+        self, *args,  map_name, **kwargs):
+        """ Private function that parses the input arguments and adds them as
+        keyword arguments matching (the order of) the non default arguments of
+        the pde class.
+        """
+
+        # If any args are given, add them to kwargs
+        if len(args) > 0:
+            if len(kwargs) > 0:
+                raise ValueError(
+                    + map_name.lower()
+                    + " input is specified both as positional and keyword arguments. This is not supported."
+                )
+
+            # Check if the number of args does not match the number of
+            # non_default_args of the model
+            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 "
+                    + map_name.lower()
+                    + "."
+                )
+
+            # Add args to kwargs following the order of non_default_args
+            for idx, arg in enumerate(args):
+                kwargs[self._non_default_args[idx]] = arg
+
+        # Check kwargs matches non_default_args
+        if set(list(kwargs.keys())) != set(self._non_default_args):
+            error_msg = (
+                map_name.lower()
+                + f" input is specified by keywords arguments {list(kwargs.keys())} that does not match the non_default_args of "
+                + map_name
+                + f" {self._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 self._non_default_args}
+
+        return kwargs
 
 class LinearPDE(PDE):
     """
@@ -143,7 +194,7 @@ class SteadyStateLinearPDE(LinearPDE):
     Parameters
     -----------   
     PDE_form : callable function
-        Callable function with signature `PDE_form(parameter)` where `parameter` is the Bayesian parameter. The function returns a tuple with the discretized differential operator A and right-hand-side b. The types of A and b are determined by what the method :meth:`linalg_solve` accepts as first and second parameters, respectively. 
+        Callable function with signature `PDE_form(parameter1, parameter2, ...)` where `parameter1`, `parameter2`, etc. are the Bayesian unknown parameters (the user can choose any names for these parameters, e.g. `a`, `b`, etc.). The function returns a tuple with the discretized differential operator A and right-hand-side b. The types of A and b are determined by what the method :meth:`linalg_solve` accepts as first and second parameters, respectively. 
 
     kwargs: 
         See :class:`~cuqi.pde.LinearPDE` for the remaining keyword arguments. 
@@ -158,7 +209,10 @@ class SteadyStateLinearPDE(LinearPDE):
 
     def assemble(self, *args, **kwargs):
         """Assembles differential operator and rhs according to PDE_form"""
-        self.diff_op, self.rhs = self.PDE_form(*args, **kwargs)
+        kwargs = self._parse_args_add_to_kwargs(
+            *args, map_name="assemble", **kwargs
+        )
+        self.diff_op, self.rhs = self.PDE_form(**kwargs)
 
     def solve(self):
         """Solve the PDE and returns the solution and an information variable `info` which is a tuple of all variables returned by the function `linalg_solve` after the solution."""
@@ -186,7 +240,7 @@ class TimeDependentLinearPDE(LinearPDE):
     Parameters
     -----------   
     PDE_form : callable function
-        Callable function with signature `PDE_form(parameter, t)` where `parameter` is the Bayesian parameter and `t` is the time at which the PDE form is evaluated. The function returns a tuple of (`differential_operator`, `source_term`, `initial_condition`) where `differential_operator` is the linear operator at time `t`, `source_term` is the source term at time `t`, and `initial_condition` is the initial condition. The types of `differential_operator` and `source_term` are determined by what the method :meth:`linalg_solve` accepts as linear operator and right-hand side, respectively. The type of `initial_condition` should be the same type as the solution returned by :meth:`linalg_solve`.
+        Callable function with signature `PDE_form(parameter1, parameter2, ..., t)` where `parameter1`, `parameter2`, etc. are the Bayesian unknown parameters (the user can choose any names for these parameters, e.g. `a`, `b`, etc.) and `t` is the time at which the PDE form is evaluated. The function returns a tuple of (`differential_operator`, `source_term`, `initial_condition`) where `differential_operator` is the linear operator at time `t`, `source_term` is the source term at time `t`, and `initial_condition` is the initial condition. The types of `differential_operator` and `source_term` are determined by what the method :meth:`linalg_solve` accepts as linear operator and right-hand side, respectively. The type of `initial_condition` should be the same type as the solution returned by :meth:`linalg_solve`.
 
     time_steps : ndarray 
         An array of the discretized times corresponding to the time steps that starts with the initial time and ends with the final time
@@ -228,6 +282,18 @@ class TimeDependentLinearPDE(LinearPDE):
     def method(self):
         return self._method
 
+    @property
+    def _non_default_args(self):
+        """Returns the non-default arguments of the PDE_form function"""
+        if self._stored_non_default_args is None:
+            self._stored_non_default_args = get_non_default_args(self.PDE_form)
+            # Remove the time argument from the non-default arguments
+            # since it is provided automatically by `solve` method and is not
+            # an argument to be inferred in Bayesian inference setting.
+            if 't' in self._stored_non_default_args:
+                self._stored_non_default_args.remove('t')
+        return self._stored_non_default_args
+
     @method.setter
     def method(self, value):
         if value.lower() != 'forward_euler' and value.lower() != 'backward_euler':
@@ -237,13 +303,13 @@ class TimeDependentLinearPDE(LinearPDE):
 
     def assemble(self, *args, **kwargs):
         """Assemble PDE"""
+        kwargs = self._parse_args_add_to_kwargs(*args, map_name="assemble", **kwargs)
         self._parameter_kwargs = kwargs
-        self._parameter_args = args
 
     def assemble_step(self, t):
         """Assemble time step at time t"""
         self.diff_op, self.rhs, self.initial_condition = self.PDE_form(
-            *self._parameter_args, **self._parameter_kwargs, t=t
+            **self._parameter_kwargs, t=t
         )
 
     def solve(self):

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/cuqi/utilities/_utilities.py

@@ -188,7 +188,7 @@ def approx_derivative(func, wrt, direction=None, epsilon=np.sqrt(np.finfo(float)
     # We compute the Jacobian matrix of func using forward differences.
     # If the function is scalar-valued, we compute the gradient instead.
     # If the direction is provided, we compute the direction-Jacobian product.
-    wrt = np.asarray(wrt)
+    wrt = force_ndarray(wrt, flatten=True)
     f0 = func(wrt)
     Matr = np.zeros([infer_len(wrt), infer_len(f0)])
     dx = np.zeros(len(wrt))

{cuqipy-1.3.0.post0.dev277 → cuqipy-1.3.0.post0.dev292}/tests/test_model.py

@@ -7,15 +7,15 @@ import cuqi
 import pytest
 from scipy import optimize
 from copy import copy, deepcopy
-from cuqi.geometry import _identity_geometries, _DefaultGeometry1D, _DefaultGeometry2D, Geometry, Discrete, Image2D
+from cuqi.geometry import _identity_geometries, _DefaultGeometry1D, _DefaultGeometry2D, Geometry, Discrete, Image2D, KLExpansion
 from cuqi.utilities import force_ndarray
 from cuqi.experimental.geometry import _ProductGeometry
 
-def test_PDE_model_multiple_input():
-    """ Test that the PDE model can accept multiple inputs specified as positional arguments or keyword arguments """
+def test_steady_state_PDE_model_multiple_input():
+    """ Test that the steady state PDE model and gradient can accept multiple inputs specified as positional arguments or keyword arguments """
     pde_test_model = MultipleInputTestModel.helper_build_steady_state_PDE_test_model()
     pde_test_model.populate_model_variations()
-    CUQI_pde = pde_test_model.model_variations[0] # PDE model with multiple inputs
+    CUQI_pde = pde_test_model.model_variations[1] # PDE model with multiple inputs
 
     # Check that the model has correct parameter name
     assert CUQI_pde._non_default_args == ['mag', 'kappa_scale']
@@ -29,6 +29,54 @@ def test_PDE_model_multiple_input():
     # Check that the two outputs are the same
     assert np.allclose(output1, output2)
 
+    # Assert evaluating gradient works
+    direction = np.random.randn(CUQI_pde.range_dim)
+
+    # Make sure gradient can be computed with positional or keyword arguments
+    grad1 = CUQI_pde.gradient(direction, mag=2, kappa_scale=2)
+    grad2 = CUQI_pde.gradient(direction, 2, 2)
+
+    # Passing wrong kwargs should raise an error
+    with pytest.raises(
+        ValueError,
+        match=r"The gradient input is specified by a direction and keywords arguments \['mag', 'kappa'\] that does not match the non_default_args of the model \['mag', 'kappa_scale'\].",
+    ):
+        CUQI_pde.gradient(direction, mag=2, kappa=2)
+
+def test_time_dependent_PDE_model_multiple_input():
+    """ Test that the time dependent PDE model and gradient can accept multiple inputs specified as positional arguments or keyword arguments """
+    pde_test_model = MultipleInputTestModel.helper_build_time_dependent_PDE_test_model()
+    pde_test_model.populate_model_variations()
+    CUQI_pde = pde_test_model.model_variations[1] # PDE model with multiple inputs
+
+    # Check that the model has correct parameter name
+    assert CUQI_pde._non_default_args == ['mag', 'IC']
+
+    # Check that we can provide parameter names when evaluating the model
+    mag = 2
+    IC_ = np.random.randn(CUQI_pde.domain_geometry.geometries[1].par_dim)
+    output1 = CUQI_pde(mag=mag, IC=IC_)
+
+    # And check that we can provide positional arguments
+    output2 = CUQI_pde(mag, IC_)
+
+    # Check that the two outputs are the same
+    assert np.allclose(output1, output2)
+
+    # Assert evaluating gradient works
+    direction = np.random.randn(CUQI_pde.range_dim)
+
+    # Make sure gradient can be computed with positional or keyword arguments
+    grad1 = CUQI_pde.gradient(direction, mag=mag, IC=IC_)
+    grad2 = CUQI_pde.gradient(direction, mag, IC_)
+
+    # Passing wrong kwargs should raise an error
+    with pytest.raises(
+        ValueError,
+        match=r"The gradient input is specified by a direction and keywords arguments \['mag', 'IC_value'\] that does not match the non_default_args of the model \['mag', 'IC'\].",
+    ):
+        CUQI_pde.gradient(direction, mag=mag, IC_value=IC_)
+
 def test_constructing_gradient_from_jacobian():
     """ Test that the gradient is correctly constructed from the
     jacobian when only the jacobian is specified """
@@ -392,18 +440,17 @@ class MultipleInputTestModel:
 
     def populate_model_variations(self):
         """Populate the `model_variations` list with different variations of the model that share the same forward map but differ in some other aspect like gradient, jacobian, etc."""
+
         if self.pde is not None:
-            self.populate_pde_model_variations()
-        elif self.forward_map is not None:
-            self.populate_general_model_variations()
+            first_kwarg = {"PDE": self.pde}
+        else:
+            first_kwarg = {"forward": self.forward_map}
 
-    def populate_general_model_variations(self):
-        """Populate the `model_variations` for general models that are not PDE models."""
         # Model with forward only
         model = self.model_class(
-            forward=self.forward_map,
+            **first_kwarg,
             domain_geometry=self.domain_geometry,
-            range_geometry=self.range_geometry,
+            range_geometry=self.range_geometry
         )
         model._do_test_gradient = False  # do not test this model for gradient
         self.model_variations.append(model)
@@ -411,7 +458,7 @@ class MultipleInputTestModel:
         # Model with gradient of from 1 (callable)
         if self.gradient_form1 is not None:
             model = self.model_class(
-                forward=self.forward_map,
+                **first_kwarg,
                 gradient=self.gradient_form1,
                 domain_geometry=self.domain_geometry,
                 range_geometry=self.range_geometry,
@@ -421,7 +468,7 @@ class MultipleInputTestModel:
         # Model with gradient of from 2 (tuple of callables)
         if self.gradient_form2 is not None:
             model = self.model_class(
-                forward=self.forward_map,
+                **first_kwarg,
                 gradient=self.gradient_form2,
                 domain_geometry=self.domain_geometry,
                 range_geometry=self.range_geometry,
@@ -431,7 +478,7 @@ class MultipleInputTestModel:
         # Model with gradient of from 2 incomplete (tuple of callables with some None elements)
         if self.gradient_form2_incomplete is not None:
             model = self.model_class(
-                forward=self.forward_map,
+                **first_kwarg,
                 gradient=self.gradient_form2_incomplete,
                 domain_geometry=self.domain_geometry,
                 range_geometry=self.range_geometry,
@@ -441,7 +488,7 @@ class MultipleInputTestModel:
         # Model with jacobian of from 1 (callable)
         if self.jacobian_form1 is not None:
             model = self.model_class(
-                forward=self.forward_map,
+                **first_kwarg,
                 jacobian=self.jacobian_form1,
                 domain_geometry=self.domain_geometry,
                 range_geometry=self.range_geometry,
@@ -451,7 +498,7 @@ class MultipleInputTestModel:
         # Model with jacobian of from 2 (tuple of callables)
         if self.jacobian_form2 is not None:
             model = self.model_class(
-                forward=self.forward_map,
+                **first_kwarg,
                 jacobian=self.jacobian_form2,
                 domain_geometry=self.domain_geometry,
                 range_geometry=self.range_geometry,
@@ -461,24 +508,13 @@ class MultipleInputTestModel:
         # Model with jacobian of from 2 incomplete (tuple of callables with some None elements)
         if self.jacobian_form2_incomplete is not None:
             model = self.model_class(
-                forward=self.forward_map,
+                **first_kwarg,
                 jacobian=self.jacobian_form2_incomplete,
                 domain_geometry=self.domain_geometry,
                 range_geometry=self.range_geometry,
             )
             self.model_variations.append(model)
 
-    def populate_pde_model_variations(self):
-        """Populate the `model_variations` for PDE models."""
-        # Model with PDE (no gradient)
-        model = self.model_class(
-            self.pde,
-            domain_geometry=self.domain_geometry,
-            range_geometry=self.range_geometry,
-        )
-        model._do_test_gradient = False  # do not test this model for gradient
-        self.model_variations.append(model)
-
     @staticmethod
     def create_model_test_case_combinations():
         """Create all combinations of test model variations and test cases
@@ -510,6 +546,16 @@ class MultipleInputTestModel:
         TestCase.create_test_cases_for_test_model(test_model)
         test_model_list.append(test_model)
 
+        # Model 4
+        test_model = MultipleInputTestModel.helper_build_time_dependent_PDE_test_model()
+        test_model.populate_model_variations()
+        test_model.input_bounds = [
+            0.1,
+            4,
+        ]  # choose input from uniform distribution in [0.1, 4]
+        TestCase.create_test_cases_for_test_model(test_model)
+        test_model_list.append(test_model)
+
         # Append all combinations of test model variations and test cases
         # to model_test_case_combinations
         for test_model in test_model_list:
@@ -589,7 +635,7 @@ class MultipleInputTestModel:
 
     @staticmethod
     def helper_build_steady_state_PDE_test_model():
-        """Build a PDE model with a steady state Poisson equation and two inputs: mag and kappa_scale. This model does not have gradient or jacobian functions."""
+        """Build a PDE model with a steady state Poisson equation and two inputs: mag and kappa_scale."""
 
         # Poisson equation setup
         dim = 20  # Number of nodes
@@ -635,6 +681,68 @@ class MultipleInputTestModel:
         test_model.domain_geometry = (Discrete(["mag"]), Discrete(["kappa_scale"]))
         test_model.range_geometry = Continuous1D(len(grid_obs))
 
+        # Gradient with respect to mag
+        def gradient_mag(direction, mag, kappa_scale):
+            def fwd_mag(mag_):
+                CUQI_pde.assemble(mag_, kappa_scale)
+                u, _ = CUQI_pde.solve()
+                obs_u = CUQI_pde.observe(u)
+                return obs_u
+            mag = mag.to_numpy() if isinstance(mag, CUQIarray) else mag 
+            return direction @ cuqi.utilities.approx_derivative(fwd_mag, mag)
+
+        # Gradient with respect to kappa_scale
+        def gradient_kappa_scale(direction, mag, kappa_scale):
+            def fwd_kappa_scale(kappa_scale_):
+                CUQI_pde.assemble(mag, kappa_scale_)
+                u, _ = CUQI_pde.solve()
+                obs_u = CUQI_pde.observe(u)
+                return obs_u
+            kappa_scale = kappa_scale.to_numpy() if isinstance(kappa_scale, CUQIarray) else kappa_scale
+            return direction @ cuqi.utilities.approx_derivative(fwd_kappa_scale, kappa_scale)
+
+        # Gradient with respect to all inputs (form 1, callable)
+        def gradient_form1(direction, mag, kappa_scale):
+            grad_mag = gradient_mag(direction, mag, kappa_scale)
+            grad_kappa_scale = gradient_kappa_scale(direction, mag, kappa_scale)
+            return (grad_mag, grad_kappa_scale)
+
+        # Assign the gradient functions to the test model
+        test_model.gradient_form1 = gradient_form1
+        test_model.gradient_form2 = (gradient_mag, gradient_kappa_scale)
+        test_model.gradient_form2_incomplete = (gradient_mag, None)
+
+        # Jacobian with respect to mag
+        def jacobian_mag(mag, kappa_scale):
+            def fwd_mag(mag_):
+                CUQI_pde.assemble(mag_, kappa_scale)
+                u, _ = CUQI_pde.solve()
+                obs_u = CUQI_pde.observe(u)
+                return obs_u
+            mag = mag.to_numpy() if isinstance(mag, CUQIarray) else mag
+            return cuqi.utilities.approx_derivative(fwd_mag, mag).reshape(-1, 1)
+
+        # Jacobian with respect to kappa_scale
+        def jacobian_kappa_scale(mag, kappa_scale):
+            def fwd_kappa_scale(kappa_scale_):
+                CUQI_pde.assemble(mag, kappa_scale_)
+                u, _ = CUQI_pde.solve()
+                obs_u = CUQI_pde.observe(u)
+                return obs_u
+            kappa_scale = kappa_scale.to_numpy() if isinstance(kappa_scale, CUQIarray) else kappa_scale
+            return cuqi.utilities.approx_derivative(fwd_kappa_scale, kappa_scale).reshape(-1, 1)
+
+        # Jacobian with respect to all inputs (form 1, callable)
+        def jacobian_form1(mag, kappa_scale):
+            jac_mag = jacobian_mag(mag, kappa_scale)
+            jac_kappa_scale = jacobian_kappa_scale(mag, kappa_scale)
+            return (jac_mag, jac_kappa_scale)
+
+        # Assign the jacobian functions to the test model
+        test_model.jacobian_form1 = jacobian_form1
+        test_model.jacobian_form2 = (jacobian_mag, jacobian_kappa_scale)
+        test_model.jacobian_form2_incomplete = (jacobian_mag, None)
+
         return test_model
 
     @staticmethod
@@ -699,6 +807,103 @@ class MultipleInputTestModel:
         test_model.model_class = cuqi.model.Model
         return test_model
 
+    @staticmethod
+    def helper_build_time_dependent_PDE_test_model():
+        """Build a PDE model with a time-dependent PDE and two inputs: mag, and IC."""
+
+        # Prepare PDE form
+        N = 20   # Number of solution nodes
+        endpoint = 1.0 # Length of the domain
+        max_time = 0.1 # Maximum time
+        dx = endpoint/(N+1)   # space step size
+        cfl = 5/11 # the cfl condition to have a stable solution
+        dt_approx = cfl*dx**2 # defining approximate time step size
+        max_iter = int(max_time/dt_approx) # number of time steps
+        Dxx_matr = (np.diag( -2*np.ones(N) ) + np.diag(np.ones(N-1),-1) + np.diag(np.ones(N-1),1))/dx**2
+        Dxx = lambda mag: mag * Dxx_matr # FD diffusion operator
+
+        # Grids for model
+        grid_domain = np.linspace(dx, endpoint, N, endpoint=False)
+        grid_range = np.linspace(dx, endpoint, N, endpoint=False) 
+        time_steps = np.linspace(0,max_time,max_iter+1,endpoint=True)
+
+        # PDE form (mag, IC, time)
+        def PDE_form(mag, IC, t): return (Dxx(mag), np.zeros(N), IC)
+        PDE = cuqi.pde.TimeDependentLinearPDE(
+            PDE_form, time_steps, grid_sol=grid_domain, grid_obs=grid_range, method='backward_euler')
+
+        # Build the test model
+        test_model = MultipleInputTestModel()
+        test_model.model_class = cuqi.model.PDEModel
+        test_model.pde = PDE
+        test_model.domain_geometry = (Discrete(["mag"]), Continuous1D(grid_domain))
+        test_model.range_geometry = Continuous1D(grid_range)
+        
+        test_model.model_class = cuqi.model.PDEModel
+
+        # Gradient with respect to mag
+        def gradient_mag(direction, mag, IC):
+            def fwd_mag(mag_):
+                PDE.assemble(mag_, IC)
+                u, _ = PDE.solve()
+                obs_u = PDE.observe(u)
+                return obs_u
+            mag = mag.to_numpy() if isinstance(mag, CUQIarray) else mag 
+            return direction @ cuqi.utilities.approx_derivative(fwd_mag, mag)
+        
+        # Gradient with respect to IC
+        def gradient_IC(direction, mag, IC):
+            def fwd_IC(IC_):
+                PDE.assemble(mag, IC_)
+                u, _ = PDE.solve()
+                obs_u = PDE.observe(u)
+                return obs_u
+            IC = IC.to_numpy() if isinstance(IC, CUQIarray) else IC 
+            return direction @ cuqi.utilities.approx_derivative(fwd_IC, IC)
+        
+        # Gradient with respect to all inputs (form 1, callable)
+        def gradient_form1(direction, mag, IC):
+            grad_mag = gradient_mag(direction, mag, IC)
+            grad_IC = gradient_IC(direction, mag, IC)
+            return (grad_mag, grad_IC)
+        
+        # Assign the gradient functions to the test model
+        test_model.gradient_form1 = gradient_form1
+        test_model.gradient_form2 = (gradient_mag, gradient_IC)
+        test_model.gradient_form2_incomplete = (gradient_mag, None)
+
+        # Jacobian with respect to mag
+        def jacobian_mag(mag, IC):
+            def fwd_mag(mag_):
+                PDE.assemble(mag_, IC)
+                u, _ = PDE.solve()
+                obs_u = PDE.observe(u)
+                return obs_u
+            mag = mag.to_numpy() if isinstance(mag, CUQIarray) else mag 
+            return cuqi.utilities.approx_derivative(fwd_mag, mag)
+        
+        # Jacobian with respect to IC
+        def jacobian_IC(mag, IC):
+            def fwd_IC(IC_):
+                PDE.assemble(mag, IC_)
+                u, _ = PDE.solve()
+                obs_u = PDE.observe(u)
+                return obs_u
+            IC = IC.to_numpy() if isinstance(IC, CUQIarray) else IC 
+            return cuqi.utilities.approx_derivative(fwd_IC, IC)
+        
+        # Jacobian with respect to all inputs (form 1, callable)
+        def jacobian_form1(mag, IC):
+            jac_mag = jacobian_mag(mag, IC)
+            jac_IC = jacobian_IC(mag, IC)
+            return (jac_mag, jac_IC)
+        
+        # Assign the jacobian functions to the test model
+        test_model.jacobian_form1 = jacobian_form1
+        test_model.jacobian_form2 = (jacobian_mag, jacobian_IC)
+        test_model.jacobian_form2_incomplete = (jacobian_mag, None)
+
+        return test_model
 
 class TestCase:
     """Class representing a test case for a test model. A test case consists of the input values, the expected output values, and the expected output types and error messages."""
@@ -1754,7 +1959,7 @@ def test_linear_model_allow_other_parameter_names():
     # Check providing the wrong parameter name raises an error
     with pytest.raises(
         ValueError,
-        match=r"The model input is specified by a keywords arguments \['y'\] that does not match the non_default_args of the model \['x'\].",
+        match=r"The model input is specified by keywords arguments \['y'\] that does not match the non_default_args of the model \['x'\].",
     ):
         model_x_w(y=1)
 
@@ -1779,7 +1984,7 @@ def test_linear_model_allow_other_parameter_names():
     # Check providing the wrong parameter name raises an error
     with pytest.raises(
         ValueError,
-        match=r"The adjoint input is specified by a keywords arguments \['v'\] that does not match the non_default_args of the adjoint \['w'\].",
+        match=r"The adjoint input is specified by keywords arguments \['v'\] that does not match the non_default_args of the adjoint \['w'\].",
     ):
         model_x_w.adjoint(v=1)
 
@@ -1938,4 +2143,4 @@ def test_setting_domain_and_range_geometry(domain_geometry, range_geometry, num_
 
     # Check the model domain and range geometries are of the expected type
     assert model.domain_geometry == expected_domain_geometry
-    assert model.range_geometry == expected_range_geometry
+    assert model.range_geometry == expected_range_geometry