CUQIpy 1.1.1.post0.dev36__py3-none-any.whl → 1.4.1.post0.dev124__py3-none-any.whl

cuqi/pde/__init__.py

@@ -4,3 +4,7 @@ from ._pde import (
     SteadyStateLinearPDE,
     TimeDependentLinearPDE
 )
+
+from ._observation_map import (
+    FD_spatial_gradient
+)

cuqi/pde/_observation_map.py

@@ -0,0 +1,36 @@
+import scipy
+import numpy as np
+"""
+This module contains observation map examples for PDE problems. The map can
+be passed to the `PDE` object initializer via the `observation_map` argument.
+
+For example on how to use set observation maps in time dependent PDEs, see
+`demos/howtos/time_dependent_linear_pde.py`.
+"""
+
+# 1. Steady State Observation Maps
+# --------------------------------
+
+# 2. Time-Dependent Observation Maps
+# -----------------------------------
+def FD_spatial_gradient(sol, grid, times):
+    """Time dependent observation map that computes the finite difference (FD) spatial gradient of a solution given at grid points (grid) and times (times). This map is supported for 1D spatial domains only.
+    
+    Parameters
+    ----------
+    sol : np.ndarray
+        The solution array of shape (number of grid points, number of time steps).
+
+    grid : np.ndarray
+        The spatial grid points of shape (number of grid points,).
+
+    times : np.ndarray
+        The discretized time steps of shape (number of time steps,)."""
+
+    if len(grid.shape) != 1:
+        raise ValueError("FD_spatial_gradient only supports 1D spatial domains.")
+    observed_quantity = np.zeros((len(grid)-1, len(times)))
+    for i in range(observed_quantity.shape[0]):
+        observed_quantity[i, :] = ((sol[i, :] - sol[i+1, :])/
+                                       (grid[i] - grid[i+1]))
+    return observed_quantity

cuqi/pde/_pde.py

@@ -3,6 +3,8 @@ 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):
     """
@@ -13,14 +15,15 @@ class PDE(ABC):
     PDE_form : callable function
         Callable function which returns a tuple of the needed PDE components (expected components are explained in the subclasses) 
 
-    observation_map: a function handle
-        A function that takes the PDE solution as input and the returns the observed solution. e.g. `observation_map=lambda u: u**2` or `observation_map=lambda u: u[0]`
-
     grid_sol: np.ndarray
         The grid on which solution is defined
 
     grid_obs: np.ndarray
-        The grid on which the observed solution should be interpolated (currently only supported for 1D problems).  
+        The grid on which the observed solution should be interpolated (currently only supported for 1D problems).
+
+    observation_map: a function handle
+        A function that takes the PDE solution, interpolated on `grid_obs`, as input and returns the observed solution. e.g., `observation_map=lambda u, grid_obs: u**2`.
+
     """
 
     def __init__(self, PDE_form, grid_sol=None, grid_obs=None, observation_map=None):
@@ -28,9 +31,10 @@ 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,parameter):
+    def assemble(self, *args, **kwargs):
         pass
 
     @abstractmethod
@@ -63,6 +67,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"):
@@ -93,6 +104,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):
     """
@@ -135,6 +188,10 @@ class LinearPDE(PDE):
             info = None
 
         return solution, info
+    
+    def interpolate_on_observed_domain(self, solution):
+        """Interpolate solution on observed space domain."""
+        raise NotImplementedError("interpolate_on_observed_domain method is not implemented for LinearPDE base class.")
 
 class SteadyStateLinearPDE(LinearPDE):
     """Linear steady state PDE.
@@ -142,22 +199,28 @@ 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.
+
+    observation_map: a function handle
+        A function that takes the PDE solution, interpolated on `grid_obs`, as input and returns the observed solution. e.g. `observation_map=lambda u, grid_obs: u**2`.
 
     kwargs: 
         See :class:`~cuqi.pde.LinearPDE` for the remaining keyword arguments. 
 
     Example
     -------- 
-    See demo demos/demo24_fwd_poisson.py for an illustration on how to use SteadyStateLinearPDE with varying solver choices. And demos demos/demo25_fwd_poisson_2D.py and demos/demo26_fwd_poisson_mixedBC.py for examples with mixed (Dirichlet and Neumann) boundary conditions problems. demos/demo25_fwd_poisson_2D.py also illustrates how to observe on a specific boundary, for example.
+    See demo demos/demo24_fwd_poisson.py for an illustration on how to use SteadyStateLinearPDE with varying solver choices. And demos demos/demo25_fwd_poisson_2d.py and demos/demo26_fwd_poisson_mixed_bc.py for examples with mixed (Dirichlet and Neumann) boundary conditions problems. demos/demo25_fwd_poisson_2d.py also illustrates how to observe on a specific boundary, for example.
     """
 
-    def __init__(self, PDE_form, **kwargs):
-        super().__init__(PDE_form, **kwargs)
+    def __init__(self, PDE_form, observation_map=None, **kwargs):
+        super().__init__(PDE_form, observation_map=observation_map, **kwargs)
 
-    def assemble(self, parameter):
+    def assemble(self, *args, **kwargs):
         """Assembles differential operator and rhs according to PDE_form"""
-        self.diff_op, self.rhs = self.PDE_form(parameter)
+        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."""
@@ -166,26 +229,34 @@ class SteadyStateLinearPDE(LinearPDE):
 
         return self._solve_linear_system(self.diff_op, self.rhs, self._linalg_solve, self._linalg_solve_kwargs)
 
-
-    def observe(self, solution):
-            
+    def interpolate_on_observed_domain(self, solution):
+        """Interpolate solution on observed space grid."""
         if self.grids_equal:
             solution_obs = solution
         else:
             solution_obs = interp1d(self.grid_sol, solution, kind='quadratic')(self.grid_obs)
+        return solution_obs
+
+    def observe(self, solution):
+        """Apply observation operator to the solution. This includes
+        interpolation to observation points (if different from the
+        solution grid) then applying the observation map (if provided)."""
+ 
+        # Interpolate solution on observed domain
+        solution_obs = self.interpolate_on_observed_domain(solution)
 
         if self.observation_map is not None:
-            solution_obs = self.observation_map(solution_obs)
-                
+            solution_obs = self.observation_map(solution_obs, self.grid_obs)
+
         return solution_obs
-        
+
 class TimeDependentLinearPDE(LinearPDE):
     """Time Dependent Linear PDE with fixed time stepping using Euler method (backward or forward).
     
     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
@@ -196,16 +267,20 @@ class TimeDependentLinearPDE(LinearPDE):
     method: str
         Time stepping method. Currently two options are available `forward_euler` and  `backward_euler`.
 
+    observation_map: a function handle
+        A function that takes the PDE solution, interpolated on `grid_obs` and `time_obs`, as input and returns the observed solution. e.g. `observation_map=lambda u, grid_obs, time_obs: u**2`.
+
     kwargs: 
         See :class:`~cuqi.pde.LinearPDE` for the remaining keyword arguments 
  
     Example
     -----------  
-    See demos/demo34_TimeDependentLinearPDE.py for 1D heat and 1D wave equations.
+    See demos/howtos/time_dependent_linear_pde.py for 1D heat and 1D wave equations examples. It demonstrates setting up `TimeDependentLinearPDE` objects, including the choice of time stepping methods, observation domain, and observation map.
     """
 
-    def __init__(self, PDE_form, time_steps, time_obs='final', method='forward_euler', **kwargs):
-        super().__init__(PDE_form, **kwargs)
+    def __init__(self, PDE_form, time_steps, time_obs='final',
+                 method='forward_euler', observation_map=None, **kwargs):
+        super().__init__(PDE_form, observation_map=observation_map, **kwargs)
 
         self.time_steps = time_steps
         self.method = method
@@ -227,6 +302,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':
@@ -234,13 +321,16 @@ class TimeDependentLinearPDE(LinearPDE):
                 "method can be set to either `forward_euler` or `backward_euler`")
         self._method = value
 
-    def assemble(self, parameter):
+    def assemble(self, *args, **kwargs):
         """Assemble PDE"""
-        self._parameter = parameter
+        kwargs = self._parse_args_add_to_kwargs(*args, map_name="assemble", **kwargs)
+        self._parameter_kwargs = kwargs
 
     def assemble_step(self, t):
         """Assemble time step at time t"""
-        self.diff_op, self.rhs, self.initial_condition = self.PDE_form(self._parameter, t)
+        self.diff_op, self.rhs, self.initial_condition = self.PDE_form(
+            **self._parameter_kwargs, t=t
+        )
 
     def solve(self):
         """Solve PDE by time-stepping"""
@@ -269,8 +359,8 @@ class TimeDependentLinearPDE(LinearPDE):
 
         return u, info
 
-    def observe(self, solution):
-
+    def interpolate_on_observed_domain(self, solution):
+        """Interpolate solution on observed time and space points."""
         # If observation grid is the same as solution grid and observation time
         # is the final time step then no need to interpolate
         if self.grids_equal and np.all(self.time_steps[-1:] == self._time_obs):
@@ -279,7 +369,7 @@ class TimeDependentLinearPDE(LinearPDE):
         # Interpolate solution in time and space to the observation
         # time and space
         else:
-            # Raise error if solution is 2D or 3D in space 
+            # Raise error if solution is 2D or 3D in space
             if len(solution.shape) > 2:
                 raise ValueError("Interpolation of solutions of 2D and 3D "+ 
                                  "space dimensions based on the provided "+
@@ -287,19 +377,30 @@ class TimeDependentLinearPDE(LinearPDE):
                                  "You can, instead, pass a custom "+
                                  "observation_map and pass grid_obs and "+
                                  "time_obs as None.")
-            
+
             # Interpolate solution in space and time to the observation
             # time and space
             solution_obs = scipy.interpolate.RectBivariateSpline(
-                self.grid_sol, self.time_steps, solution)(self.grid_obs,
-                                                          self._time_obs)
+                self.grid_sol, self.time_steps, solution
+            )(self.grid_obs, self._time_obs)
+
+        return solution_obs
+
+    def observe(self, solution):
+        """Apply observation operator to the solution. This includes
+        interpolation to observation points (if different from the
+        solution grid) then applying the observation map (if provided)."""
 
+        # Interpolate solution on observed domain
+        solution_obs = self.interpolate_on_observed_domain(solution)
+        
         # Apply observation map
         if self.observation_map is not None:
-            solution_obs = self.observation_map(solution_obs)
-        
+            solution_obs = self.observation_map(solution_obs, self.grid_obs, 
+                                                self._time_obs)
+
         # squeeze if only one time observation
         if len(self._time_obs) == 1:
             solution_obs = solution_obs.squeeze()
 
-        return solution_obs
+        return solution_obs