Enable general observation function for PDE, and accessing original solution (or other observable quantities)

cuqi/pde/__init__.py

@@ -4,3 +4,6 @@
     SteadyStateLinearPDE,
     TimeDependentLinearPDE
 )
+
+from ._observation_map import (_extract_spatial_temporal_obs,
+                               spatial_gradient)

cuqi/pde/_observation_map.py

@@ -0,0 +1,20 @@
+import scipy
+import numpy as np
+
+def spatial_gradient(sol, grid, times):
+    """Observation map that computes the spatial gradient of the solution given
+    at grid points (grid) and times (times)."""
+    solution_obs = np.zeros((len(grid)-1, len(times)))
+    for i in range(solution_obs.shape[0]):
+        solution_obs[i, :] = ((sol[i, :] - sol[i+1, :])/
+                                       (grid[i] - grid[i+1]))
+    return solution_obs
+
+def _extract_spatial_temporal_obs(sol, grid, times, grid_obs, time_obs):
+    """Private function to extract solution at observation points in space and
+    time."""
+    # Interpolate solution in space and time to the observation
+    # time and space
+    solution_obs = scipy.interpolate.RectBivariateSpline(
+        grid, times, sol)(grid_obs, time_obs)
+    return solution_obs

cuqi/pde/_pde.py

@@ -4,6 +4,7 @@
 from scipy.interpolate import interp1d
 import numpy as np
 from cuqi.utilities import get_non_default_args
+from ._observation_map import _extract_spatial_temporal_obs
 
 
 class PDE(ABC):
@@ -16,7 +17,7 @@ class PDE(ABC):
         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]`
+        A function that takes the PDE solution as input and returns the observed solution. e.g. `observation_map=lambda u: u**2`.
 
     grid_sol: np.ndarray
         The grid on which solution is defined
@@ -194,7 +195,10 @@ class SteadyStateLinearPDE(LinearPDE):
     Parameters
     -----------   
     PDE_form : callable function
-        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. 
+        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 and spatial grid as input and returns the observed solution. e.g. `observation_map=lambda u, grid: u**2`.
 
     kwargs: 
         See :class:`~cuqi.pde.LinearPDE` for the remaining keyword arguments. 
@@ -204,8 +208,8 @@ class SteadyStateLinearPDE(LinearPDE):
     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.
     """
 
-    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, *args, **kwargs):
         """Assembles differential operator and rhs according to PDE_form"""
@@ -230,8 +234,8 @@ def observe(self, solution):
             solution_obs = interp1d(self.grid_sol, solution, kind='quadratic')(self.grid_obs)
 
         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):
@@ -251,6 +255,8 @@ 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, the spatial grid, and the time steps as input and returns the observed solution. e.g. `observation_map=lambda u, grid, times: u**2`.
     kwargs: 
         See :class:`~cuqi.pde.LinearPDE` for the remaining keyword arguments 
 
@@ -259,8 +265,9 @@ class TimeDependentLinearPDE(LinearPDE):
     See demos/demo34_TimeDependentLinearPDE.py for 1D heat and 1D wave equations.
     """
 
-    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
@@ -360,16 +367,19 @@ def observe(self, solution):
 
             # 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)
+            solution_obs = _extract_spatial_temporal_obs(solution,
+                                                        self.grid_sol, 
+                                                        self.time_steps, 
+                                                        self.grid_obs, 
+                                                        self._time_obs)
 
         # 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

tests/test_model.py

@@ -671,7 +671,7 @@ def helper_build_steady_state_PDE_test_model():
             poisson_form,
             grid_sol=grid_sol,
             grid_obs=grid_obs,
-            observation_map=lambda u: u**2,
+            observation_map=lambda u, grid_obs: u**2,
         )
 
         # Build the test model

tests/test_pde.py

@@ -170,7 +170,7 @@ def test_observe():
     FOFD_operator = cuqi.operator.FirstOrderFiniteDifference(dim-1, bc_type='zero', dx=dx).get_matrix().todense()
     diff_operator = FOFD_operator.T @ np.diag(kappa) @ FOFD_operator
     poisson_form = lambda x: (diff_operator, source(x[0]))
-    CUQI_pde = cuqi.pde.SteadyStateLinearPDE(poisson_form, grid_sol=grid_sol, grid_obs=grid_obs, observation_map=lambda u:u**2)
+    CUQI_pde = cuqi.pde.SteadyStateLinearPDE(poisson_form, grid_sol=grid_sol, grid_obs=grid_obs, observation_map=lambda u, grid:u**2)
     x_exact = np.array([2]) # [2] is the source term parameter [mag]
     CUQI_pde.assemble(x_exact)
     sol, info = CUQI_pde.solve()
@@ -192,11 +192,12 @@ def test_observe():
 @pytest.mark.parametrize(
     "grid_obs, time_obs, observation_map, expected_obs",
     [(None, 'final', None, 'obs1'),
-     (None, 'final', lambda x: x**2, 'obs2'),
+     (None, 'final', lambda x, grid, times: x**2, 'obs2'),
      ('half_grid', 'FINAL', None, 'obs3'),
      ('half_grid', 'every_5', None, 'obs4'),
-     (None, 'every_5', lambda x: x**2, 'obs5'),
-     (np.array([3, 4.9]), np.array([0.9, 1]), lambda x: x**2, 'obs6')])
+     (None, 'every_5', lambda x, grid, times: x**2, 'obs5'),
+     (np.array([3, 4.9]), np.array([0.9, 1]), lambda x, grid, times: x**2, 'obs6'),
+     ('half_grid', 'every_5', cuqi.pde.spatial_gradient, 'obs7')])
 def test_TimeDependentLinearPDE_heat1D(copy_reference, method, time_steps,
                                        parametrization, expected_sol,
                                        grid_obs, time_obs, observation_map,
@@ -293,7 +294,8 @@ def PDE_form(source_term, t): return (
         expected_observed_sol = sol[idx_x, :][:, idx_t]
 
     if observation_map is not None:
-        expected_observed_sol = observation_map(expected_observed_sol)
+        expected_observed_sol = observation_map(
+            expected_observed_sol, grid_obs, time_obs)
 
     if len(PDE._time_obs) == 1:
         expected_observed_sol = expected_observed_sol.squeeze()