Merge pull request #610 from CUQI-DTU/mi_fwd_model

Multiple Inputs: update the Model module

Author: amal-ghamdi

cuqi/experimental/geometry/_productgeometry.py

@@ -172,10 +172,10 @@ def fun2vec(self, *funvals):
         return tuple(funvecs)
 
 
-    def __repr__(self) -> str:
+    def __repr__(self, pad="") -> str:
         """Representation of the product geometry."""
         string = "{}(".format(self.__class__.__name__) + "\n"
         for g in self.geometries:
-            string += "\t{}\n".format(g.__repr__())
-        string += ")"
+            string += pad + "    {}\n".format(g.__repr__())
+        string += pad + ")"
         return string

cuqi/model/_model.py

@@ -4,6 +4,7 @@
 from scipy.interpolate import interp1d
 import numpy as np
 
+
 class PDE(ABC):
     """
     Parametrized PDE abstract base class
@@ -30,7 +31,7 @@ def __init__(self, PDE_form, grid_sol=None, grid_obs=None, observation_map=None)
         self.observation_map = observation_map
 
     @abstractmethod
-    def assemble(self,parameter):
+    def assemble(self, *args, **kwargs):
         pass
 
     @abstractmethod
@@ -155,9 +156,9 @@ class SteadyStateLinearPDE(LinearPDE):
     def __init__(self, PDE_form, **kwargs):
         super().__init__(PDE_form, **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)
+        self.diff_op, self.rhs = self.PDE_form(*args, **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."""
@@ -178,7 +179,7 @@ def observe(self, solution):
             solution_obs = self.observation_map(solution_obs)
 
         return solution_obs
-        
+
 class TimeDependentLinearPDE(LinearPDE):
     """Time Dependent Linear PDE with fixed time stepping using Euler method (backward or forward).
     
@@ -234,13 +235,16 @@ def method(self, value):
                 "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
+        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, t)
+        self.diff_op, self.rhs, self.initial_condition = self.PDE_form(
+            *self._parameter_args, **self._parameter_kwargs, t=t
+        )
 
     def solve(self):
         """Solve PDE by time-stepping"""
@@ -279,15 +283,15 @@ def observe(self, solution):
         # 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 "+
                                  "grid_obs and time_obs are not supported. "+
                                  "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(
@@ -297,7 +301,7 @@ def observe(self, solution):
         # Apply observation map
         if self.observation_map is not None:
             solution_obs = self.observation_map(solution_obs)
-        
+
         # squeeze if only one time observation
         if len(self._time_obs) == 1:
             solution_obs = solution_obs.squeeze()

cuqi/pde/_pde.py

@@ -863,10 +863,9 @@ def PDE_form(IC, t): return (Dxx, np.zeros(N), IC)
         # Bayesian model
         x = cuqi.distribution.Gaussian(np.zeros(model.domain_dim), 1)
         y = cuqi.distribution.Gaussian(model(x), sigma2)
-        
-        # Initialize Deconvolution as BayesianProblem problem
-        super().__init__(y, x, y=data)
 
+        # Initialize Heat1D as BayesianProblem problem
+        super().__init__(y, x, y=data)
         # Store exact values
         self.exactSolution = x_exact
         self.exactData = y_exact

cuqi/testproblem/_testproblem.py

@@ -677,7 +677,7 @@ def beta_likelihood():
         lambda x: x**2,
         range_geometry=1,
         domain_geometry=1,
-        gradient=lambda direction, wrt: 2*wrt*direction)
+        gradient=lambda direction, x: 2*x*direction)
 
     # set a gaussian prior
     x = cuqi.distribution.Gaussian(0, 1)