Repository: tinosulzer/PyBaMM

Issue:
refactor solver class
see #16.

This class will take an instance of a model that has been processed by the Parameter (#18) and Spatial Discretisation (#20) classes. That is, the solver can assume that all Parameter, Variable and Operator nodes representing gradients (e.g. Laplace) nodes in the expression tree have been replaced with concrete Vector and Matrix nodes. That is, the expression tree now represents a concrete linear algebra expression

For now we only consider Models representing ODE equations and solvers implementing of methods of lines (MOL). But the implementation should be flexible enough to be extended to DAE equations, and time discretistion with a newton iteration solver rather than MOL.

The solver class needs to:
   1. verify that the model rhs expressions are concrete linear algebra expressions
   2. loop over time:
           2.1. evaluate RHS as a function of current state variable `y` and `t`
           2.2. step forward in time using RHS
 
Solver class should store the current state vector internally, and have an `integrate` function that can move the simulation forward in time, to allow for output/plotting over time, e.g.

```python
model = ReactionDiffusionModel()
param_file = Parameters("file.dat")
param_file.process(model)
disc = FiniteVolumeDiscretisation(mesh_dx=0.1, something_else=2)
y0 = disc.process(model)
solver = RungeKuttaSolver(y0, tolerance=1e-5)
Tf = 1.0
nout = 100
nout_dt = Tf / nout
for i in range(nout):
   solver.integrate(nout_dt)
   plot(solver.y, disc)
```

Discussion:

Patch:
diff --git a/docs/index.rst b/docs/index.rst
index 117f9988..e999b810 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -26,3 +26,4 @@ Contents
     source/models/index
     source/parameters/index
     source/discretisations/index
+    source/solvers/index
diff --git a/docs/source/solvers/base_solver.rst b/docs/source/solvers/base_solver.rst
new file mode 100644
index 00000000..9c0bd4f6
--- /dev/null
+++ b/docs/source/solvers/base_solver.rst
@@ -0,0 +1,5 @@
+Base Solver
+===========
+
+.. autoclass:: pybamm.BaseSolver
+  :members:
diff --git a/docs/source/solvers/index.rst b/docs/source/solvers/index.rst
new file mode 100644
index 00000000..3a1caae4
--- /dev/null
+++ b/docs/source/solvers/index.rst
@@ -0,0 +1,7 @@
+Solvers
+=======
+
+.. toctree::
+
+  base_solver
+  scipy_solver
diff --git a/docs/source/solvers/scipy_solver.rst b/docs/source/solvers/scipy_solver.rst
new file mode 100644
index 00000000..b273ed2b
--- /dev/null
+++ b/docs/source/solvers/scipy_solver.rst
@@ -0,0 +1,5 @@
+Scipy Solver
+============
+
+.. autoclass:: pybamm.ScipySolver
+  :members:
diff --git a/pybamm/__init__.py b/pybamm/__init__.py
index 5bd18991..1b0479ad 100644
--- a/pybamm/__init__.py
+++ b/pybamm/__init__.py
@@ -79,6 +79,18 @@ from .expression_tree.independent_variable import IndependentVariable
 from .expression_tree.independent_variable import t
 from .expression_tree.vector import Vector, StateVector
 
+#
+# Model classes
+#
+from .models.core import BaseModel
+from .models.reaction_diffusion import ReactionDiffusionModel
+from .models.electrolyte_current import ElectrolyteCurrentModel
+
+#
+# Submodel classes
+#
+from .models.submodels import electrolyte, interface
+
 #
 # Parameters class and methods
 #
@@ -107,28 +119,10 @@ from .discretisations.finite_volume_meshes import (
 from .simulation import Simulation
 
 #
-# Solver class and lists of known methods
-#
-from .solver import Solver
-from .solver import KNOWN_INTEGRATORS
-from .solver import KNOWN_SPATIAL_DISCRETISATIONS
-
+# Solver classes
 #
-# Variables class
-#
-from .variables import Variables
-
-#
-# Model classes
-#
-from .models.core import BaseModel
-from .models.reaction_diffusion import ReactionDiffusionModel
-from .models.electrolyte_current import ElectrolyteCurrentModel
-
-#
-# Submodel classes
-#
-from .models.submodels import electrolyte, interface
+from .solvers.base_solver import BaseSolver
+from .solvers.scipy_solver import ScipySolver
 
 #
 # Remove any imported modules, so we don't expose them as part of pybamm
diff --git a/pybamm/discretisations/base_discretisation.py b/pybamm/discretisations/base_discretisation.py
index de878c94..18bbc5e7 100644
--- a/pybamm/discretisations/base_discretisation.py
+++ b/pybamm/discretisations/base_discretisation.py
@@ -157,11 +157,10 @@ class BaseDiscretisation(object):
             )
 
         elif isinstance(symbol, pybamm.BinaryOperator):
-            new_left = self.process_symbol(
-                symbol.children[0], domain, y_slices, boundary_conditions
-            )
+            left, right = symbol.children
+            new_left = self.process_symbol(left, domain, y_slices, boundary_conditions)
             new_right = self.process_symbol(
-                symbol.children[1], domain, y_slices, boundary_conditions
+                right, domain, y_slices, boundary_conditions
             )
             return symbol.__class__(new_left, new_right)
 
@@ -172,6 +171,9 @@ class BaseDiscretisation(object):
             return symbol.__class__(new_child)
 
         elif isinstance(symbol, pybamm.Variable):
+            assert isinstance(y_slices, dict), ValueError(
+                """y_slices should be dict, not {}""".format(type(y_slices))
+            )
             return pybamm.StateVector(y_slices[symbol.id])
 
         else:
diff --git a/pybamm/solver.py b/pybamm/solver.py
deleted file mode 100644
index 6b174fd0..00000000
--- a/pybamm/solver.py
+++ /dev/null
@@ -1,125 +0,0 @@
-#
-# Solver for the simulation
-#
-from __future__ import absolute_import, division
-from __future__ import print_function, unicode_literals
-import pybamm
-
-import numpy as np
-import scipy.integrate as it
-
-KNOWN_INTEGRATORS = ["BDF", "analytical"]
-KNOWN_SPATIAL_DISCRETISATIONS = ["Finite Volumes"]
-
-
-class Solver(object):
-    """Solver for the simulation created in pybamm.sim.
-
-    Parameters
-    ----------
-    integrator : string, optional
-        The method for integration in time:
-            * "BDF" (default): scipy.integrate.solve_ivp with method="BDF".
-            * "analytical": use explicit analytical solution (if possible).
-                Otherwise revert to "BDF".
-    spatial_discretisation : string, optional
-        The spatial discretisation scheme:
-            * "Finite Volumes" (default): Finite Volumes discretisation.
-                Cell edges are on the boundaries between the subdomains
-    """
-
-    def __init__(
-        self, integrator="BDF", spatial_discretisation="Finite Volumes", tol=1e-8
-    ):
-
-        if integrator not in KNOWN_INTEGRATORS:
-            raise NotImplementedError(
-                """Integrator '{}' is not implemented.
-                                      Valid choices: one of '{}'.""".format(
-                    integrator, KNOWN_INTEGRATORS
-                )
-            )
-        self.integrator = integrator
-
-        if spatial_discretisation not in KNOWN_SPATIAL_DISCRETISATIONS:
-            raise NotImplementedError(
-                """Spatial discretisation '{}' is not
-                                      implemented. Valid choices: one of '{}'.
-                                      """.format(
-                    spatial_discretisation, KNOWN_SPATIAL_DISCRETISATIONS
-                )
-            )
-        self.spatial_discretisation = spatial_discretisation
-
-        self.tol = tol
-
-    def __str__(self):
-        return "{}_{}".format(self.integrator, self.spatial_discretisation)
-
-    def operators(self, mesh):
-        """Define the operators in each domain.
-
-        Parameters
-        ----------
-        mesh : :class:`pybamm.mesh.Mesh` instance
-            The mesh on which the operators are defined.
-
-        Returns
-        -------
-        :class: `pybamm.operators.Operators`
-            A class of all the operators.
-
-        """
-        return pybamm.Operators(self.spatial_discretisation, mesh)
-
-    def get_simulation_vars(self, sim):
-        """Run a simulation.
-
-        Parameters
-        ----------
-        sim : :class:`pybamm.simulation.Simulation` instance
-            The simulation to be solved.
-
-        Returns
-        -------
-        vars : pybamm.variables.Variables() instance
-            The variables of the solved model.
-
-        """
-
-        # Initialise variables
-        vars = pybamm.Variables(sim.model)
-
-        # Initialise y for PDE solver
-        yinit = sim.model.initial_conditions()
-
-        # Solve ODEs
-        def derivs(t, y):
-            vars.update(t, y)
-            dydt = sim.model.pdes_rhs(vars)
-            return dydt
-
-        if self.integrator == "analytical":
-            # TODO: implement analytical simulation
-            pass
-        elif self.integrator == "BDF":
-            target_time = sim.mesh.time
-            sol = it.solve_ivp(
-                derivs,
-                (target_time[0], target_time[-1]),
-                yinit,
-                t_eval=target_time,
-                method="BDF",
-                rtol=self.tol,
-                atol=self.tol,
-            )
-            # TODO: implement concentration cut-off event
-
-        # Extract variables from y
-        vars.update(sol.t, sol.y)
-
-        # Post-process (get potentials)
-        vars.dt = np.concatenate([np.array([0]), np.diff(vars.t)])
-        # TODO: write post-processing function
-
-        return vars
diff --git a/pybamm/solvers/__init__.py b/pybamm/solvers/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
new file mode 100644
index 00000000..6cad122d
--- /dev/null
+++ b/pybamm/solvers/base_solver.py
@@ -0,0 +1,81 @@
+#
+# Base solver class
+#
+from __future__ import absolute_import, division
+from __future__ import print_function, unicode_literals
+
+
+class BaseSolver(object):
+    """Solve a discretised model.
+
+    Parameters
+    ----------
+    tolerance : float, optional
+        The tolerance for the solver (default is 1e-8).
+    """
+
+    def __init__(self, tol=1e-8):
+        self._tol = tol
+
+        # Solutions
+        self._t = None
+        self._y = None
+
+    @property
+    def tol(self):
+        return self._tol
+
+    @tol.setter
+    def tol(self, value):
+        self._tol = value
+
+    @property
+    def t(self):
+        return self._t
+
+    @t.setter
+    def t(self, value):
+        self._t = value
+
+    @property
+    def y(self):
+        return self._y
+
+    @y.setter
+    def y(self, value):
+        self._y = value
+
+    def solve(self, model, t_eval):
+        """Calculate the solution of the model at specified times.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel` (or subclass)
+            The model whose solution to calculate. Must have attributes rhs and
+            initial_conditions
+        t_eval : numeric type
+            The times at which to compute the solution
+
+        """
+
+        def dydt(t, y):
+            return model.rhs.evaluate(t, y)
+
+        y0 = model.initial_conditions
+        self.t, self.y = self.integrate(dydt, y0, t_eval)
+
+    def integrate(self, derivs, y0, t_eval):
+        """
+        Solve a model defined by dydt with initial conditions y0.
+
+        Parameters
+        ----------
+        derivs : method
+            A function that takes in t and y and returns the time-derivative dydt
+        y0 : numeric type
+            The initial conditions
+        t_eval : numeric type
+            The times at which to compute the solution
+
+        """
+        raise NotImplementedError
diff --git a/pybamm/solvers/scipy_solver.py b/pybamm/solvers/scipy_solver.py
new file mode 100644
index 00000000..5d174494
--- /dev/null
+++ b/pybamm/solvers/scipy_solver.py
@@ -0,0 +1,65 @@
+#
+# Solver class using Scipy's adaptive time stepper
+#
+from __future__ import absolute_import, division
+from __future__ import print_function, unicode_literals
+import pybamm
+
+import scipy.integrate as it
+
+
+class ScipySolver(pybamm.BaseSolver):
+    """Solve a discretised model, using scipy.integrate.solve_ivp.
+
+    Parameters
+    ----------
+    method : string, optional
+        The method to use in solve_ivp (default is "BDF")
+    tolerance : float, optional
+        The tolerance for the solver (default is 1e-8). Set as the both reltol and
+        abstol in solve_ivp.
+    """
+
+    def __init__(self, method="BDF", tol=1e-8):
+        super().__init__(tol)
+        self.method = method
+
+    @property
+    def method(self):
+        return self._method
+
+    @method.setter
+    def method(self, value):
+        self._method = value
+
+    def integrate(self, derivs, y0, t_eval):
+        """
+        Solve a model defined by dydt with initial conditions y0.
+
+        Parameters
+        ----------
+        derivs : method
+            A function that takes in t (size (1,)), y (size (n,))
+            and returns the time-derivative dydt (size (n,))
+        y0 : :class:`numpy.array`, size (n,)
+            The initial conditions
+        t_eval : :class:`numpy.array`, size (k,)
+            The times at which to compute the solution
+
+        Returns
+        -------
+        object
+            An object containing the times and values of the solution, as well as
+            various diagnostic messages.
+        """
+        sol = it.solve_ivp(
+            derivs,
+            (t_eval[0], t_eval[-1]),
+            y0,
+            t_eval=t_eval,
+            method=self.method,
+            rtol=self.tol,
+            atol=self.tol,
+        )
+        # TODO: implement concentration cut-off event
+        return sol.t, sol.y
diff --git a/pybamm/variables.py b/pybamm/variables.py
deleted file mode 100644
index b704ddf2..00000000
--- a/pybamm/variables.py
+++ /dev/null
@@ -1,116 +0,0 @@
-#
-# Variables of a model or simulation
-#
-from __future__ import absolute_import, division
-from __future__ import print_function, unicode_literals
-
-import numpy as np
-
-
-class Variables(object):
-    """
-    Extracts and stores the model variables.
-
-    Parameters
-    ----------
-    model : A model class instance
-        The simulation model.
-    mesh : :class:`pybamm.mesh.Mesh` instance
-        The simulation mesh.
-    """
-
-    def __init__(self, model):
-        self.model = model
-        self.mesh = model.mesh
-
-    def update(self, t, y):
-        """Update variables with a new t and y.
-        Note that we can't store variables at each timestep since the
-        timestepping is adaptive.
-
-        Parameters
-        ----------
-        t : float or array_like
-            The simulation time.
-        y : array_like
-            An array containing variable values to be extracted.
-        """
-        self.t = t
-        # Unpack y iteratively
-        start = 0
-        for var, mesh in self.model.pde_variables:
-            end = start + mesh.npts
-            self.__dict__[var] = y[start:end]
-            start = end
-            # Split 'x' variables further into n, s and p
-            if mesh == self.mesh.x:
-                (
-                    self.__dict__[var + "n"],
-                    self.__dict__[var + "s"],
-                    self.__dict__[var + "p"],
-                ) = np.split(
-                    self.__dict__[var],
-                    np.cumsum([self.mesh.xn.npts, self.mesh.xs.npts]),
-                )
-
-    @property
-    def neg(self):
-        """Variables in the negative electrode."""
-
-        class NegVars(object):
-            def __init__(self, all_vars):
-                self.t = all_vars.t
-                for attr, value in all_vars.__dict__.items():
-                    if attr[-1] == "n":
-                        self.__dict__[attr[:-1]] = value
-
-        return NegVars(self)
-
-    @property
-    def pos(self):
-        """Variables for the positive electrode."""
-
-        class PosVars(object):
-            def __init__(self, all_vars):
-                self.t = all_vars.t
-                for attr, value in all_vars.__dict__.items():
-                    if attr[-1] == "p":
-                        self.__dict__[attr[:-1]] = value
-
-        return PosVars(self)
-
-    def set_reaction_vars(self, reaction_vars):
-        for name, value in reaction_vars.items():
-            self.__dict__[name] = value
-            (
-                self.__dict__[name + "n"],
-                self.__dict__[name + "s"],
-                self.__dict__[name + "p"],
-            ) = np.split(
-                self.__dict__[name], np.cumsum([self.mesh.xn.npts, self.mesh.xs.npts])
-            )
-
-    def average(self):
-        """Average variable attributes over the relevant (sub)domain."""
-        param = self.model.param.geometric
-        mesh = self.mesh
-        old_keys = list(self.__dict__.keys())
-        for attr in old_keys:
-            if attr in ["c"]:
-                avg = np.dot(self.__dict__[attr].T, mesh.x.d_edges)
-                self.__dict__[attr + "_avg"] = avg
-
-            elif attr[-1] == "n":
-                # Negative
-                avg = np.dot(self.__dict__[attr].T, mesh.xn.d_edges) / param.ln
-                self.__dict__[attr + "_avg"] = avg
-
-            elif attr[-1] == "s":
-                # Separator
-                avg = np.dot(self.__dict__[attr].T, mesh.xs.d_edges) / param.ls
-                self.__dict__[attr + "_avg"] = avg
-
-            elif attr[-1] == "p":
-                # Positive
-                avg = np.dot(self.__dict__[attr].T, mesh.xp.d_edges) / param.lp
-                self.__dict__[attr + "_avg"] = avg


Test patch:
diff --git a/tests/test_discretisations/test_base_discretisations.py b/tests/test_discretisations/test_base_discretisations.py
index 3a3811af..a8044d8f 100644
--- a/tests/test_discretisations/test_base_discretisations.py
+++ b/tests/test_discretisations/test_base_discretisations.py
@@ -96,6 +96,48 @@ class TestDiscretise(unittest.TestCase):
         # non-spatial unary operator
         # TODO: none of these implemented yet
 
+    def test_process_complex_expression(self):
+        var1 = pybamm.Variable("var1")
+        var2 = pybamm.Variable("var2")
+        par1 = pybamm.Parameter("par1")
+        par2 = pybamm.Parameter("par2")
+        scal1 = pybamm.Scalar("scal1")
+        scal2 = pybamm.Scalar("scal2")
+        expression = (scal1 * (par1 + var2)) / ((var1 - par2) + scal2)
+
+        disc = pybamm.BaseDiscretisation(None)
+        y_slices = {var1.id: slice(53), var2.id: slice(53, 59)}
+        exp_disc = disc.process_symbol(expression, None, y_slices)
+        self.assertTrue(isinstance(exp_disc, pybamm.Division))
+        # left side
+        self.assertTrue(isinstance(exp_disc.children[0], pybamm.Multiplication))
+        self.assertTrue(isinstance(exp_disc.children[0].children[0], pybamm.Scalar))
+        self.assertTrue(isinstance(exp_disc.children[0].children[1], pybamm.Addition))
+        self.assertTrue(
+            isinstance(exp_disc.children[0].children[1].children[0], pybamm.Parameter)
+        )
+        self.assertTrue(
+            isinstance(exp_disc.children[0].children[1].children[1], pybamm.StateVector)
+        )
+        self.assertEqual(
+            exp_disc.children[0].children[1].children[1].y_slice, y_slices[var2.id]
+        )
+        # right side
+        self.assertTrue(isinstance(exp_disc.children[1], pybamm.Addition))
+        self.assertTrue(
+            isinstance(exp_disc.children[1].children[0], pybamm.Subtraction)
+        )
+        self.assertTrue(
+            isinstance(exp_disc.children[1].children[0].children[0], pybamm.StateVector)
+        )
+        self.assertEqual(
+            exp_disc.children[1].children[0].children[0].y_slice, y_slices[var1.id]
+        )
+        self.assertTrue(
+            isinstance(exp_disc.children[1].children[0].children[1], pybamm.Parameter)
+        )
+        self.assertTrue(isinstance(exp_disc.children[1].children[1], pybamm.Scalar))
+
     def test_discretise_spatial_operator(self):
         mesh = MeshForTesting()
         disc = DiscretisationForTesting(mesh)
diff --git a/tests/test_solvers/__init__.py b/tests/test_solvers/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/test_solvers/test_base_solver.py b/tests/test_solvers/test_base_solver.py
new file mode 100644
index 00000000..d7e6c3be
--- /dev/null
+++ b/tests/test_solvers/test_base_solver.py
@@ -0,0 +1,27 @@
+#
+# Tests for the Base Solver class
+#
+from __future__ import absolute_import, division
+from __future__ import print_function, unicode_literals
+import pybamm
+
+import unittest
+
+
+class TestBaseSolver(unittest.TestCase):
+    def test_base_solver_init(self):
+        solver = pybamm.BaseSolver(tol=1e-4)
+        self.assertEqual(solver.tol, 1e-4)
+        self.assertEqual(solver.t, None)
+        self.assertEqual(solver.y, None)
+        with self.assertRaises(NotImplementedError):
+            solver.integrate(None, None, None)
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    unittest.main()
diff --git a/tests/test_solvers/test_scipy_solver.py b/tests/test_solvers/test_scipy_solver.py
new file mode 100644
index 00000000..bc77cb4e
--- /dev/null
+++ b/tests/test_solvers/test_scipy_solver.py
@@ -0,0 +1,65 @@
+#
+# Tests for the Scipy Solver class
+#
+from __future__ import absolute_import, division
+from __future__ import print_function, unicode_literals
+import pybamm
+
+import unittest
+import numpy as np
+
+
+class TestScipySolver(unittest.TestCase):
+    def test_integrate(self):
+        # Constant
+        solver = pybamm.ScipySolver(tol=1e-8, method="RK45")
+
+        def constant_growth(t, y):
+            return 0.5 * np.ones_like(y)
+
+        y0 = np.array([0])
+        t_eval = np.linspace(0, 1, 100)
+        t_sol, y_sol = solver.integrate(constant_growth, y0, t_eval)
+        np.testing.assert_array_equal(t_sol, t_eval)
+        np.testing.assert_allclose(0.5 * t_sol, y_sol[0])
+
+        # Exponential decay
+        solver = pybamm.ScipySolver(tol=1e-8, method="BDF")
+
+        def exponential_decay(t, y):
+            return -0.1 * y
+
+        y0 = np.array([1])
+        t_eval = np.linspace(0, 1, 100)
+        t_sol, y_sol = solver.integrate(exponential_decay, y0, t_eval)
+        np.testing.assert_allclose(y_sol[0], np.exp(-0.1 * t_sol))
+
+    def test_model_solver(self):
+        # Create model
+        model = pybamm.BaseModel()
+        var = pybamm.Variable("var", domain=["whole cell"])
+        model.rhs = {var: pybamm.Scalar(0.1) * var}
+        model.initial_conditions = {var: pybamm.Scalar(1)}
+        # No need to set parameters; can use base discretisation (no spatial operators)
+        param = pybamm.ParameterValues(
+            base_parameters={"Ln": 0.1, "Ls": 0.2, "Lp": 0.3}
+        )
+        mesh = pybamm.FiniteVolumeMacroMesh(param)
+        disc = pybamm.BaseDiscretisation(mesh)
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.ScipySolver(tol=1e-8, method="RK45")
+        t_eval = mesh["time"]
+        solver.solve(model, t_eval)
+        np.testing.assert_array_equal(solver.t, t_eval)
+        np.testing.assert_allclose(solver.y[0], np.exp(0.1 * solver.t))
+
+
+if __name__ == "__main__":
+    print("Add -v for more debug output")
+    import sys
+
+    if "-v" in sys.argv:
+        debug = True
+    unittest.main()