Repository: pybamm-team/PyBaMM
Pull request #550 (2019-07-24 13:59:33)

introduce different methods of stepping solution
# Description
Adds a step method for the solvers. At the moment this returns a new solution object each time step is called (this is the same behavior as the step method in scikits-odes), but maybe it would be good to automatically append the next step to the existing solution? 

Fixes # (issue)

## Type of change

Please delete options that are not relevant.

- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)


# Key checklist:

- [ ] No style issues: `$ flake8`
- [ ] All tests pass: `$ python run-tests.py --unit`
- [ ] The documentation builds: `$ cd docs` and then `$ make clean; make html`

You can run all three at once, using `$ python run-tests.py --quick`.

## Further checks: 

- [ ] Code is commented, particularly in hard-to-understand areas
- [ ] Tests added that prove fix is effective or that feature works
- [ ] Any dependent changes have been merged and published in downstream modules



Patch:
diff --git a/examples/notebooks/models/SPM.ipynb b/examples/notebooks/models/SPM.ipynb
index a3f270a8e..b02ca84a0 100644
--- a/examples/notebooks/models/SPM.ipynb
+++ b/examples/notebooks/models/SPM.ipynb
@@ -124,7 +124,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -140,7 +140,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -183,7 +183,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -199,7 +199,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -218,7 +218,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -257,7 +257,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -273,7 +273,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -301,7 +301,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -329,7 +329,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -347,7 +347,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -378,7 +378,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -645,7 +645,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -667,7 +667,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -712,7 +712,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -760,7 +760,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -829,7 +829,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.7.3"
   }
  },
  "nbformat": 4,
diff --git a/examples/scripts/SPMe_step.py b/examples/scripts/SPMe_step.py
new file mode 100644
index 000000000..666976c61
--- /dev/null
+++ b/examples/scripts/SPMe_step.py
@@ -0,0 +1,64 @@
+#
+# Example to compare solving for all times against stepping individually
+#
+import pybamm
+import numpy as np
+import matplotlib.pyplot as plt
+
+pybamm.set_logging_level("INFO")
+
+# load model
+model = pybamm.lithium_ion.SPMe()
+
+# create geometry
+geometry = model.default_geometry
+
+# load parameter values and process model and geometry
+param = model.default_parameter_values
+param.process_model(model)
+param.process_geometry(geometry)
+
+# set mesh
+mesh = pybamm.Mesh(geometry, model.default_submesh_types, model.default_var_pts)
+
+# discretise model
+disc = pybamm.Discretisation(mesh, model.default_spatial_methods)
+disc.process_model(model)
+
+# solve model
+t_eval = np.linspace(0, 0.2, 100)
+solver = model.default_solver
+solution = solver.solve(model, t_eval)
+
+# step model
+dt = 0.05
+time = 0
+end_time = solution.t[-1]
+step_solver = model.default_solver
+step_solution = None
+while time < end_time:
+    current_step_sol = step_solver.step(model, dt=dt, npts=10)
+    if not step_solution:
+        # create solution object on first step
+        step_solution = current_step_sol
+    else:
+        # append solution from the current step to step_solution
+        step_solution.append(current_step_sol)
+    time += dt
+
+# plot
+voltage = pybamm.ProcessedVariable(
+    model.variables["Terminal voltage [V]"], solution.t, solution.y, mesh=mesh
+
+)
+step_voltage = pybamm.ProcessedVariable(
+    model.variables["Terminal voltage [V]"], step_solution.t, step_solution.y, mesh=mesh
+)
+plt.plot(solution.t, voltage(solution.t), "b-", label="SPMe (continuous solve)")
+plt.plot(
+    step_solution.t, step_voltage(step_solution.t), "ro", label="SPMe (steppped solve)"
+)
+plt.xlabel(r"$t$")
+plt.ylabel("Terminal voltage [V]")
+plt.legend()
+plt.show()
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
index 48023269f..a30216f13 100644
--- a/pybamm/solvers/base_solver.py
+++ b/pybamm/solvers/base_solver.py
@@ -1,6 +1,8 @@
 #
 # Base solver class
 #
+import pybamm
+import numpy as np
 
 
 class BaseSolver(object):
@@ -33,7 +35,9 @@ class BaseSolver(object):
         self._tol = value
 
     def solve(self, model, t_eval):
-        """Calculate the solution of the model at specified times.
+        """
+        Execute the solver setup and calculate the solution of the model at
+        specified times.
 
         Parameters
         ----------
@@ -44,7 +48,125 @@ class BaseSolver(object):
             The times at which to compute the solution
 
         """
+        pybamm.logger.info("Start solving {}".format(model.name))
+
+        # Set up
+        timer = pybamm.Timer()
+        start_time = timer.time()
+        self.set_up(model)
+        set_up_time = timer.time() - start_time
+
+        # Solve
+        solution, solve_time, termination = self.compute_solution(model, t_eval)
+
+        # Assign times
+        solution.solve_time = solve_time
+        solution.total_time = timer.time() - start_time
+        solution.set_up_time = set_up_time
+
+        pybamm.logger.info("Finish solving {} ({})".format(model.name, termination))
+        pybamm.logger.info(
+            "Set-up time: {}, Solve time: {}, Total time: {}".format(
+                timer.format(solution.set_up_time),
+                timer.format(solution.solve_time),
+                timer.format(solution.total_time),
+            )
+        )
+        return solution
+
+    def step(self, model, dt, npts=2):
+        """
+        Step the solution of the model forward by a given time increment. The
+        first time this method is called it executes the necessary setup by
+        calling `self.set_up(model)`.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            The model whose solution to calculate. Must have attributes rhs and
+            initial_conditions
+        dt : numeric type
+            The timestep over which to step the solution
+        npts : int, optional
+            The number of points at which the solution will be returned during
+            the step dt. Defualt is 2 (returns the solution at t0 and t0 + dt).
+
+        """
+        # Set timer
+        timer = pybamm.Timer()
+
+        # Run set up on first step
+        if not hasattr(self, 'y0'):
+            start_time = timer.time()
+            self.set_up(model)
+            self.t = 0.0
+            set_up_time = timer.time() - start_time
+        else:
+            set_up_time = None
+
+        # Step
+        pybamm.logger.info("Start stepping {}".format(model.name))
+        t_eval = np.linspace(self.t, self.t + dt, npts)
+        solution, solve_time, termination = self.compute_solution(model, t_eval)
+
+        # Assign times
+        solution.solve_time = solve_time
+        if set_up_time:
+            solution.total_time = timer.time() - start_time
+            solution.set_up_time = set_up_time
+
+        # Set self.t and self.y0 to their values at the final step
+        self.t = solution.t[-1]
+        self.y0 = solution.y[:, -1]
+
+        pybamm.logger.info("Finish stepping {} ({})".format(model.name, termination))
+        if set_up_time:
+            pybamm.logger.info(
+                "Set-up time: {}, Step time: {}, Total time: {}".format(
+                    timer.format(solution.set_up_time),
+                    timer.format(solution.solve_time),
+                    timer.format(solution.total_time),
+                )
+            )
+        else:
+            pybamm.logger.info(
+                "Step time: {}".format(
+                    timer.format(solution.solve_time),
+                )
+            )
+        return solution
 
+    def compute_solution(self, model, t_eval):
+        """Calculate the solution of the model at specified times. Note: this
+        does *not* execute the solver setup.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            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
+
+        """
+        raise NotImplementedError
+
+    def set_up(self, model):
+        """Unpack model, perform checks, simplify and calculate jacobian.
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            The model whose solution to calculate. Must have attributes rhs and
+            initial_conditions
+
+        Raises
+        ------
+        :class:`pybamm.SolverError`
+            If the model contains any algebraic equations (in which case a DAE solver
+            should be used instead)
+
+        """
         raise NotImplementedError
 
     def get_termination_reason(self, solution, events):
diff --git a/pybamm/solvers/dae_solver.py b/pybamm/solvers/dae_solver.py
index df76b88ea..035a31f05 100644
--- a/pybamm/solvers/dae_solver.py
+++ b/pybamm/solvers/dae_solver.py
@@ -55,7 +55,7 @@ class DaeSolver(pybamm.BaseSolver):
     def max_steps(self, max_steps):
         self._max_steps = max_steps
 
-    def solve(self, model, t_eval):
+    def compute_solution(self, model, t_eval):
         """Calculate the solution of the model at specified times.
 
         Parameters
@@ -67,78 +67,24 @@ class DaeSolver(pybamm.BaseSolver):
             The times at which to compute the solution
 
         """
-        pybamm.logger.info("Start solving {}".format(model.name))
-
-        # Set up
         timer = pybamm.Timer()
-        start_time = timer.time()
-        concatenated_rhs, concatenated_algebraic, y0, model_events, jac = self.set_up(
-            model
-        )
-        set_up_time = timer.time() - start_time
-
-        # get mass matrix entries
-        mass_matrix = model.mass_matrix.entries
-
-        def residuals(t, y, ydot):
-            pybamm.logger.debug(
-                "Evaluating residuals for {} at t={}".format(model.name, t)
-            )
-            y = y[:, np.newaxis]
-            rhs_eval, known_evals = concatenated_rhs.evaluate(t, y, known_evals={})
-            # reuse known_evals
-            alg_eval = concatenated_algebraic.evaluate(t, y, known_evals=known_evals)[0]
-            # turn into 1D arrays
-            rhs_eval = rhs_eval[:, 0]
-            alg_eval = alg_eval[:, 0]
-            return np.concatenate((rhs_eval, alg_eval)) - mass_matrix @ ydot
-
-        # Create event-dependent function to evaluate events
-        def event_fun(event):
-            def eval_event(t, y):
-                return event.evaluate(t, y)
-
-            return eval_event
-
-        events = [event_fun(event) for event in model_events.values()]
-
-        # Create function to evaluate jacobian
-        if jac is not None:
 
-            def jacobian(t, y):
-                return jac.evaluate(t, y, known_evals={})[0]
-
-        else:
-            jacobian = None
-
-        # Solve
         solve_start_time = timer.time()
         pybamm.logger.info("Calling DAE solver")
         solution = self.integrate(
-            residuals,
-            y0,
+            self.residuals,
+            self.y0,
             t_eval,
-            events=events,
-            mass_matrix=mass_matrix,
-            jacobian=jacobian,
+            events=self.event_funs,
+            mass_matrix=model.mass_matrix.entries,
+            jacobian=self.jacobian,
         )
-        # Assign times
-        solution.solve_time = timer.time() - solve_start_time
-        solution.total_time = timer.time() - start_time
-        solution.set_up_time = set_up_time
+        solve_time = timer.time() - solve_start_time
 
         # Identify the event that caused termination
-        termination = self.get_termination_reason(solution, model_events)
-
-        pybamm.logger.info("Finish solving {} ({})".format(model.name, termination))
-        pybamm.logger.info(
-            "Set-up time: {}, Solve time: {}, Total time: {}".format(
-                timer.format(solution.set_up_time),
-                timer.format(solution.solve_time),
-                timer.format(solution.total_time),
-            )
-        )
-        return solution
+        termination = self.get_termination_reason(solution, self.events)
+
+        return solution, solve_time, termination
 
     def set_up(self, model):
         """Unpack model, perform checks, simplify and calculate jacobian.
@@ -149,26 +95,13 @@ class DaeSolver(pybamm.BaseSolver):
             The model whose solution to calculate. Must have attributes rhs and
             initial_conditions
 
-        Returns
-        -------
-        concatenated_rhs : :class:`pybamm.Concatenation`
-            Right-hand side of differential equations
-        concatenated_algebraic : :class:`pybamm.Concatenation`
-            Algebraic equations, which should evaluate to zero
-        y0 : :class:`numpy.array`
-            Vector of initial conditions
-        events : dict
-            Dicitonary of events at which the model should terminate
-        jac : :class:`pybamm.SparseStack`
-            Jacobian matrix for the differential and algebraic equations
-
         Raises
         ------
         :class:`pybamm.SolverError`
             If the model contains any algebraic equations (in which case a DAE solver
             should be used instead)
         """
-        # create simplified rhs algebraic and event expressions
+        # create simplified rhs, algebraic and event expressions
         concatenated_rhs = model.concatenated_rhs
         concatenated_algebraic = model.concatenated_algebraic
         events = model.events
@@ -236,7 +169,48 @@ class DaeSolver(pybamm.BaseSolver):
             # can use DAE solver to solve ODE model
             y0 = model.concatenated_initial_conditions[:, 0]
 
-        return concatenated_rhs, concatenated_algebraic, y0, events, jac
+        # Create functions to evaluate residuals
+        def residuals(t, y, ydot):
+            pybamm.logger.debug(
+                "Evaluating residuals for {} at t={}".format(model.name, t)
+            )
+            y = y[:, np.newaxis]
+            rhs_eval, known_evals = concatenated_rhs.evaluate(t, y, known_evals={})
+            # reuse known_evals
+            alg_eval = concatenated_algebraic.evaluate(t, y, known_evals=known_evals)[0]
+            # turn into 1D arrays
+            rhs_eval = rhs_eval[:, 0]
+            alg_eval = alg_eval[:, 0]
+            return (
+                np.concatenate((rhs_eval, alg_eval)) - model.mass_matrix.entries @ ydot
+            )
+
+        # Create event-dependent function to evaluate events
+        def event_fun(event):
+            def eval_event(t, y):
+                return event.evaluate(t, y)
+
+            return eval_event
+
+        event_funs = [event_fun(event) for event in events.values()]
+
+        # Create function to evaluate jacobian
+        if jac is not None:
+
+            def jacobian(t, y):
+                return jac.evaluate(t, y, known_evals={})[0]
+
+        else:
+            jacobian = None
+
+        # Add the solver attributes
+        self.y0 = y0
+        self.rhs = rhs
+        self.algebraic = algebraic
+        self.residuals = residuals
+        self.events = events
+        self.event_funs = event_funs
+        self.jacobian = jacobian
 
     def calculate_consistent_initial_conditions(
         self, rhs, algebraic, y0_guess, jac=None
diff --git a/pybamm/solvers/ode_solver.py b/pybamm/solvers/ode_solver.py
index 4e325873a..7d16dcb99 100644
--- a/pybamm/solvers/ode_solver.py
+++ b/pybamm/solvers/ode_solver.py
@@ -17,7 +17,7 @@ class OdeSolver(pybamm.BaseSolver):
     def __init__(self, method=None, tol=1e-8):
         super().__init__(method, tol)
 
-    def solve(self, model, t_eval):
+    def compute_solution(self, model, t_eval):
         """Calculate the solution of the model at specified times.
 
         Parameters
@@ -29,68 +29,24 @@ class OdeSolver(pybamm.BaseSolver):
             The times at which to compute the solution
 
         """
-        pybamm.logger.info("Start solving {}".format(model.name))
-
-        # Set up
         timer = pybamm.Timer()
-        start_time = timer.time()
-        concatenated_rhs, y0, model_events, jac_rhs = self.set_up(model)
-        set_up_time = timer.time() - start_time
-
-        # Create function to evaluate rhs
-        def dydt(t, y):
-            pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t))
-            y = y[:, np.newaxis]
-            dy = concatenated_rhs.evaluate(t, y, known_evals={})[0]
-            return dy[:, 0]
-
-        # Create event-dependent function to evaluate events
-        def event_fun(event):
-            def eval_event(t, y):
-                return event.evaluate(t, y)
-
-            return eval_event
 
-        events = [event_fun(event) for event in model_events.values()]
-
-        # Create function to evaluate jacobian
-        if jac_rhs is not None:
-
-            def jacobian(t, y):
-                return jac_rhs.evaluate(t, y, known_evals={})[0]
-
-        else:
-            jacobian = None
-
-        # Solve
         solve_start_time = timer.time()
         pybamm.logger.info("Calling ODE solver")
         solution = self.integrate(
-            dydt,
-            y0,
+            self.dydt,
+            self.y0,
             t_eval,
-            events=events,
+            events=self.event_funs,
             mass_matrix=model.mass_matrix.entries,
-            jacobian=jacobian,
+            jacobian=self.jacobian,
         )
-
-        # Assign times
-        solution.solve_time = timer.time() - solve_start_time
-        solution.total_time = timer.time() - start_time
-        solution.set_up_time = set_up_time
+        solve_time = timer.time() - solve_start_time
 
         # Identify the event that caused termination
-        termination = self.get_termination_reason(solution, model_events)
-
-        pybamm.logger.info("Finish solving {} ({})".format(model.name, termination))
-        pybamm.logger.info(
-            "Set-up time: {}, Solve time: {}, Total time: {}".format(
-                timer.format(solution.set_up_time),
-                timer.format(solution.solve_time),
-                timer.format(solution.total_time),
-            )
-        )
-        return solution
+        termination = self.get_termination_reason(solution, self.events)
+
+        return solution, solve_time, termination
 
     def set_up(self, model):
         """Unpack model, perform checks, simplify and calculate jacobian.
@@ -101,17 +57,6 @@ class OdeSolver(pybamm.BaseSolver):
             The model whose solution to calculate. Must have attributes rhs and
             initial_conditions
 
-        Returns
-        -------
-        concatenated_rhs : :class:`pybamm.Concatenation`
-            Right-hand side of differential equations
-        y0 : :class:`numpy.array`
-            Vector of initial conditions
-        events : dict
-            Dictionary of events at which the model should terminate
-        jac_rhs : :class:`pybamm.SparseStack`
-            Jacobian matrix for the differential equations
-
         Raises
         ------
         :class:`pybamm.SolverError`
@@ -119,13 +64,16 @@ class OdeSolver(pybamm.BaseSolver):
             should be used instead)
 
         """
+        # Check for algebraic equations
         if len(model.algebraic) > 0:
             raise pybamm.SolverError(
                 """Cannot use ODE solver to solve model with DAEs"""
             )
 
+        # create simplified rhs and event expressions
         concatenated_rhs = model.concatenated_rhs
         events = model.events
+
         if model.use_simplify:
             # set up simplification object, for re-use of dict
             simp = pybamm.Simplification()
@@ -137,10 +85,10 @@ class OdeSolver(pybamm.BaseSolver):
             events = {name: simp.simplify(event) for name, event in events.items()}
 
         y0 = model.concatenated_initial_conditions[:, 0]
+
         if model.use_jacobian:
             # Create Jacobian from simplified rhs
             y = pybamm.StateVector(slice(0, np.size(y0)))
-
             pybamm.logger.info("Calculating jacobian")
             jac_rhs = concatenated_rhs.jac(y)
             model.jacobian = jac_rhs
@@ -152,7 +100,6 @@ class OdeSolver(pybamm.BaseSolver):
             if model.use_to_python:
                 pybamm.logger.info("Converting jacobian to python")
                 jac_rhs = pybamm.EvaluatorPython(jac_rhs)
-
         else:
             jac_rhs = None
 
@@ -164,7 +111,37 @@ class OdeSolver(pybamm.BaseSolver):
                 name: pybamm.EvaluatorPython(event) for name, event in events.items()
             }
 
-        return concatenated_rhs, y0, events, jac_rhs
+        # Create function to evaluate rhs
+        def dydt(t, y):
+            pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t))
+            y = y[:, np.newaxis]
+            dy = concatenated_rhs.evaluate(t, y, known_evals={})[0]
+            return dy[:, 0]
+
+        # Create event-dependent function to evaluate events
+        def event_fun(event):
+            def eval_event(t, y):
+                return event.evaluate(t, y)
+
+            return eval_event
+
+        event_funs = [event_fun(event) for event in events.values()]
+
+        # Create function to evaluate jacobian
+        if jac_rhs is not None:
+
+            def jacobian(t, y):
+                return jac_rhs.evaluate(t, y, known_evals={})[0]
+
+        else:
+            jacobian = None
+
+        # Add the solver attributes
+        self.y0 = y0
+        self.dydt = dydt
+        self.events = events
+        self.event_funs = event_funs
+        self.jacobian = jacobian
 
     def integrate(
         self, derivs, y0, t_eval, events=None, mass_matrix=None, jacobian=None
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
index 4b5f7b515..d185e5e45 100644
--- a/pybamm/solvers/solution.py
+++ b/pybamm/solvers/solution.py
@@ -1,6 +1,7 @@
 #
 # Solution class
 #
+import numpy as np
 
 
 class Solution(object):
@@ -81,3 +82,14 @@ class Solution(object):
     def termination(self, value):
         "Updates the reason for termination"
         self._termination = value
+
+    def append(self, solution):
+        """
+        Appends solution.t and solution.y onto self.t and self.y.
+        Note: this process removes the initial time and state of solution to avoid
+        duplicate times and states being stored (self.t[-1] is equal to solution.t[0],
+        and self.y[:, -1] is equal to solution.y[:, 0]).
+
+        """
+        self.t = np.concatenate((self.t, solution.t[1:]))
+        self.y = np.concatenate((self.y, solution.y[:, 1:]), axis=1)


Test patch:
diff --git a/tests/unit/test_solvers/test_base_solver.py b/tests/unit/test_solvers/test_base_solver.py
index 822c10980..ca64a7ad2 100644
--- a/tests/unit/test_solvers/test_base_solver.py
+++ b/tests/unit/test_solvers/test_base_solver.py
@@ -14,8 +14,12 @@ class TestBaseSolver(unittest.TestCase):
         solver.tol = 1e-5
         self.assertEqual(solver.tol, 1e-5)
 
+        model = pybamm.BaseModel()
         with self.assertRaises(NotImplementedError):
-            solver.solve(None, None)
+            solver.solve(model, None)
+            solver.step(model, None)
+            solver.compute_solution(None, None)
+            solver.set_up(None)
 
 
 if __name__ == "__main__":
diff --git a/tests/unit/test_solvers/test_scikits_solvers.py b/tests/unit/test_solvers/test_scikits_solvers.py
index 912af5da1..597ecf12e 100644
--- a/tests/unit/test_solvers/test_scikits_solvers.py
+++ b/tests/unit/test_solvers/test_scikits_solvers.py
@@ -614,6 +614,72 @@ class TestScikitsSolvers(unittest.TestCase):
         np.testing.assert_array_equal(solution.t, t_eval)
         np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
 
+    def test_model_step_ode(self):
+        # Create model
+        model = pybamm.BaseModel()
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        var = pybamm.Variable("var", domain=whole_cell)
+        model.rhs = {var: 0.1 * var}
+        model.initial_conditions = {var: 1}
+        disc = get_discretisation_for_testing()
+        disc.process_model(model)
+
+        solver = pybamm.ScikitsOdeSolver(tol=1e-9)
+
+        # Step once
+        dt = 0.1
+        step_sol = solver.step(model, dt)
+        np.testing.assert_array_equal(step_sol.t, [0, dt])
+        np.testing.assert_allclose(step_sol.y[0], np.exp(0.1 * step_sol.t))
+
+        # Step again (return 5 points)
+        step_sol_2 = solver.step(model, dt, npts=5)
+        np.testing.assert_array_equal(step_sol_2.t, np.linspace(dt, 2 * dt, 5))
+        np.testing.assert_allclose(step_sol_2.y[0], np.exp(0.1 * step_sol_2.t))
+
+        # Check steps give same solution as solve
+        t_eval = np.concatenate((step_sol.t, step_sol_2.t[1:]))
+        solution = solver.solve(model, t_eval)
+        concatenated_steps = np.concatenate((step_sol.y[0], step_sol_2.y[0, 1:]))
+        np.testing.assert_allclose(solution.y[0], concatenated_steps)
+
+    def test_model_step_dae(self):
+        # Create model
+        model = pybamm.BaseModel()
+        whole_cell = ["negative electrode", "separator", "positive electrode"]
+        var1 = pybamm.Variable("var1", domain=whole_cell)
+        var2 = pybamm.Variable("var2", domain=whole_cell)
+        model.rhs = {var1: 0.1 * var1}
+        model.algebraic = {var2: 2 * var1 - var2}
+        model.initial_conditions = {var1: 1, var2: 2}
+        model.use_jacobian = False
+        disc = get_discretisation_for_testing()
+        disc.process_model(model)
+
+        solver = pybamm.ScikitsDaeSolver(tol=1e-8)
+
+        # Step once
+        dt = 0.1
+        step_sol = solver.step(model, dt)
+        np.testing.assert_array_equal(step_sol.t, [0, dt])
+        np.testing.assert_allclose(step_sol.y[0], np.exp(0.1 * step_sol.t))
+        np.testing.assert_allclose(step_sol.y[-1], 2 * np.exp(0.1 * step_sol.t))
+
+        # Step again (return 5 points)
+        step_sol_2 = solver.step(model, dt, npts=5)
+        np.testing.assert_array_equal(step_sol_2.t, np.linspace(dt, 2 * dt, 5))
+        np.testing.assert_allclose(step_sol_2.y[0], np.exp(0.1 * step_sol_2.t))
+        np.testing.assert_allclose(step_sol_2.y[-1], 2 * np.exp(0.1 * step_sol_2.t))
+
+        # append solutions
+        step_sol.append(step_sol_2)
+
+        # Check steps give same solution as solve
+        t_eval = step_sol.t
+        solution = solver.solve(model, t_eval)
+        np.testing.assert_allclose(solution.y[0], step_sol.y[0, :])
+        np.testing.assert_allclose(solution.y[-1], step_sol.y[-1, :])
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")
diff --git a/tests/unit/test_solvers/test_scipy_solver.py b/tests/unit/test_solvers/test_scipy_solver.py
index 6749d61d3..cadc2d12b 100644
--- a/tests/unit/test_solvers/test_scipy_solver.py
+++ b/tests/unit/test_solvers/test_scipy_solver.py
@@ -234,6 +234,42 @@ class TestScipySolver(unittest.TestCase):
             np.ones((N, T.size)) * (T[np.newaxis, :] - np.exp(T[np.newaxis, :])),
         )
 
+    def test_model_step(self):
+        # Create model
+        model = pybamm.BaseModel()
+        domain = ["negative electrode", "separator", "positive electrode"]
+        var = pybamm.Variable("var", domain=domain)
+        model.rhs = {var: 0.1 * var}
+        model.initial_conditions = {var: 1}
+        # No need to set parameters; can use base discretisation (no spatial operators)
+
+        # create discretisation
+        mesh = get_mesh_for_testing()
+        spatial_methods = {"macroscale": pybamm.FiniteVolume}
+        disc = pybamm.Discretisation(mesh, spatial_methods)
+        disc.process_model(model)
+
+        solver = pybamm.ScipySolver(tol=1e-8, method="RK45")
+
+        # Step once
+        dt = 0.1
+        step_sol = solver.step(model, dt)
+        np.testing.assert_array_equal(step_sol.t, [0, dt])
+        np.testing.assert_allclose(step_sol.y[0], np.exp(0.1 * step_sol.t))
+
+        # Step again (return 5 points)
+        step_sol_2 = solver.step(model, dt, npts=5)
+        np.testing.assert_array_equal(step_sol_2.t, np.linspace(dt, 2 * dt, 5))
+        np.testing.assert_allclose(step_sol_2.y[0], np.exp(0.1 * step_sol_2.t))
+
+        # append solutions
+        step_sol.append(step_sol_2)
+
+        # Check steps give same solution as solve
+        t_eval = step_sol.t
+        solution = solver.solve(model, t_eval)
+        np.testing.assert_allclose(solution.y[0], step_sol.y[0])
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")