Merge pull request #656 from CUQI-DTU/multiple_inputs_pde

Multiple inputs support for pde gradient

Author: amal-ghamdi

cuqi/model/_model.py

@@ -179,7 +179,7 @@ def _non_default_args(self):
             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 @@ def _2fun(self, geometry=None, is_par=True, **kwargs):
             # 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 @@ def _2par(self, geometry=None, to_CUQIarray=False, is_par=False, **kwargs):
             # 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 @@ def _parse_args_add_to_kwargs(
                 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 @@ def _handle_case_when_model_input_is_distributions(self, kwargs):
         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 @@ def _apply_chain_rule_to_account_for_domain_geometry_gradient(self,
 
         # 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 @@ def __init__(self, PDE: cuqi.pde.PDE, range_geometry, domain_geometry):
         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 @@ def _forward_func(self, **kwargs):
 
         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:

cuqi/pde/_pde.py

@@ -3,6 +3,7 @@
 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 @@ def __init__(self, PDE_form, grid_sol=None, grid_obs=None, observation_map=None)
         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 @@ def _compare_grid(grid1,grid2):
 
         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 @@ def grid_obs(self,value):
     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 @@ def __init__(self, PDE_form, **kwargs):
 
     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 @@ def __init__(self, PDE_form, time_steps, time_obs='final', method='forward_euler
     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 @@ def method(self, value):
 
     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):

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