Sprint16 time dep refactor

cuqi/array/__init__.py

@@ -0,0 +1 @@
+from ._array import MultiVector, Samples, TimeSeries, CUQIarray2

cuqi/array/_array.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Optional
+import cuqi
+
+__all__ = ["MultiVector", "Samples", "TimeSeries", "CUQIarray2"]
+
+class MultiVector(ABC):
+    pass
+
+class Samples(MultiVector):
+    pass
+
+class TimeSeries(MultiVector):
+    """
+    Abstract base class for time series.
+
+    times: Time points of the time series.
+    values: Values of the time series.
+    """
+    pass
+
+class TimeDependantPDESolution(TimeSeries):
+    """
+    Abstract base class for time series of PDEs.
+
+    rejected_times: Time points of the time series that were rejected.
+    local_errors: Estimated local errors of the time steps.
+    u_list: List of solutions at the time points.
+    t_list: List of time steps.
+    """
+    pass
+
+class CUQIarray2(ABC):
+    pass

cuqi/forward_model_gallery/__init__.py

@@ -0,0 +1 @@
+from ._pde_gallary import Heat1D, Poisson1D, Poisson2D, Wave2D

cuqi/forward_model_gallery/_pde_gallery.py

@@ -0,0 +1,11 @@
+class Heat1D():
+    pass
+
+class Poisson1D():
+    pass
+
+class Poisson2D():
+    pass
+
+class Wave2D():
+    pass

cuqi/pde/__init__.py

@@ -4,3 +4,9 @@
     SteadyStateLinearPDE,
     TimeDependentLinearPDE
 )
+from ._time_integrator import (
+    TimeIntegrator,
+    ForwardEuler,
+    BackwardEuler,
+    ThetaMethod,
+    Trapezoidal)

cuqi/pde/_observer.py

@@ -0,0 +1,17 @@
+class Observer():
+    pass
+
+class SteadyObserver(Observer):
+    """
+    Attributes:
+    observation_operator: Observation operator (lambda function or x coordinates).
+    observed_solution: Observed solution (CUQIarray)."""
+    pass
+
+class TimeDependantObserver(Observer):
+    """
+    Attributes:
+    observation_operator: Observation operator (lambda function or x and t coordinates).
+    observed_solution: Observed solution (TimeSeries or CUQIArray)."""
+
+    pass

cuqi/pde/_pde.py

@@ -3,6 +3,7 @@
 from inspect import getsource
 from scipy.interpolate import interp1d
 import numpy as np
+from ._time_integrator import ForwardEuler, BackwardEuler
 
 class PDE(ABC):
     """
@@ -92,7 +93,7 @@ def grids_equal(self):
         return self._grids_equal
 
 
-class LinearPDE(PDE):
+class LinearPDE(PDE, ABC):
     """
     Parametrized Linear PDE base class
 
@@ -122,9 +123,9 @@ def __init__(self, PDE_form, linalg_solve=None, linalg_solve_kwargs=None, **kwar
         self._linalg_solve = linalg_solve
         self._linalg_solve_kwargs = linalg_solve_kwargs
 
-    def _solve_linear_system(self, A, b, linalg_solve, kwargs):
+    def _solve_linear_system(self, A, b):
         """Helper function that solves the linear system `A*x=b` using the provided solve method `linalg_solve` and its keyword arguments `kwargs`. It then returns the output in the format: `solution`, `info`"""
-        returned_values = linalg_solve(A, b, **kwargs)
+        returned_values = self._linalg_solve(A, b, **self._linalg_solve_kwargs)
         if isinstance(returned_values, tuple):
             solution = returned_values[0]
             info = returned_values[1:]
@@ -162,7 +163,7 @@ def solve(self):
         if not hasattr(self, "diff_op") or not hasattr(self, "rhs"):
             raise Exception("PDE is not assembled.")
 
-        return self._solve_linear_system(self.diff_op, self.rhs, self._linalg_solve, self._linalg_solve_kwargs)
+        return self._solve_linear_system(self.diff_op, self.rhs)
 
 
     def observe(self, solution):
@@ -199,53 +200,32 @@ class TimeDependentLinearPDE(LinearPDE):
     See demos/demo34_TimeDependentLinearPDE.py for 1D heat and 1D wave equations.
     """
 
-    def __init__(self, PDE_form, time_steps, method='forward_euler', **kwargs):
+    def __init__(self, PDE_form, time_steps, method=ForwardEuler(), **kwargs):
         super().__init__(PDE_form, **kwargs)
 
         self.time_steps = time_steps
         self.method = method
 
-    @property
-    def method(self):
-        return self._method
-
-    @method.setter
-    def method(self, value):
-        if value.lower() != 'forward_euler' and value.lower() != 'backward_euler':
-            raise ValueError(
-                "method can be set to either `forward_euler` or `backward_euler`")
-        self._method = value
-
     def assemble(self, parameter):
         """Assemble PDE"""
         self._parameter = parameter
 
-    def assemble_step(self, t):
+    def assemble_time_dependant_step(self, t):
         """Assemble time step at time t"""
         self.diff_op, self.rhs, self.initial_condition = self.PDE_form(self._parameter, t)
 
     def solve(self):
         """Solve PDE by time-stepping"""
+        #TODO: use time-stepping class
 
-        if self.method == 'forward_euler':
-            for idx, t in enumerate(self.time_steps[:-1]):
-                dt = self.time_steps[idx+1] - t
-                self.assemble_step(t)
-                if idx == 0:
-                    u = self.initial_condition
-                u = (dt*self.diff_op + np.eye(len(u)))@u + dt*self.rhs  # from u at time t, gives u at t+dt
-            info = None
-
-        if self.method == 'backward_euler':
-            for idx, t in enumerate(self.time_steps[1:]):
-                dt = t - self.time_steps[idx]
-                self.assemble_step(t)
-                if idx == 0:
-                    u = self.initial_condition
-                A = np.eye(len(u)) - dt*self.diff_op
-                # from u at time t-dt, gives u at t
-                u, info = self._solve_linear_system(
-                    A, u + dt*self.rhs, self._linalg_solve, self._linalg_solve_kwargs)
+        for idx, t in enumerate(self.time_steps[:-1]):
+            dt = self.time_steps[idx+1] - t
+            self.assemble_time_dependant_step(t)
+            # can do self.observe_time_dependant_step(t) here
+            if idx == 0:
+                u = self.initial_condition
+            u, _ = self.method.propagate(u, dt, self.rhs, self.diff_op, self._solve_linear_system)
+        info = None
 
         return u, info
 
@@ -259,4 +239,4 @@ def observe(self, solution):
         if self.observation_map is not None:
             solution_obs = self.observation_map(solution_obs)
 
-        return solution_obs
+        return solution_obs

cuqi/pde/_time_integrator.py

@@ -0,0 +1,37 @@
+# do not import numpy
+from __future__ import annotations
+from abc import ABC, abstractmethod
+import numpy as np
+from typing import Optional
+import cuqi
+
+class TimeIntegrator(ABC):
+    @abstractmethod
+    def propagate(self, u, dt, rhs, diff_op, solve_linear_system, I=None, apply_bc=None):
+        """Propagate the solution using the explicit Euler method."""
+        pass
+
+class ThetaMethod(TimeIntegrator):
+    def __init__(self, theta):
+	    # solver details
+        self.theta = theta # 0: backward Euler, 1: forward Euler, 0.5: trapezoidal  
+
+    def propagate(self, u, dt, rhs, diff_op, solver, I=None, apply_bc=None):
+        if I is None:
+            I = np.eye(len(u))
+        rhs_op = I + self.theta*dt*diff_op
+        lhs_op = I - (1-self.theta)*dt*diff_op
+        u_next, info = solver(rhs_op, lhs_op@u + dt*rhs)
+        return u_next, info
+
+class ForwardEuler(ThetaMethod):
+    def __init__(self):
+        super().__init__(1.0)
+
+class BackwardEuler(ThetaMethod):
+    def __init__(self):
+        super().__init__(0.0)
+
+class Trapezoidal(ThetaMethod):
+    def __init__(self):
+        super().__init__(0.5)

cuqi/pde/_time_step_selector.py

@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+
+class Static():
+    pass 
+
+class Dynamic(ABC):
+    pass
+
+class Standard(Dynamic):
+    pass