Repository: pybamm-team/PyBaMM

Issue:
Casadi events
Hack events with the casadi solver by doing step-and-check (if possible). May slow down the solver too much.

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 037defc39..d59641db3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,7 @@
 *.tmp
 *.png
 /local/
+*.DS_Store
 
 # don't ignore important .txt files
 !requirements*
diff --git a/examples/scripts/compare-dae-solver.py b/examples/scripts/compare-dae-solver.py
index e4dff68a3..a092a7865 100644
--- a/examples/scripts/compare-dae-solver.py
+++ b/examples/scripts/compare-dae-solver.py
@@ -32,6 +32,6 @@ scikits_sol = pybamm.ScikitsDaeSolver(atol=1e-8, rtol=1e-8).solve(model, t_eval)
 
 # plot
 models = [model, model, model]
-solutions = [scikits_sol, klu_sol, casadi_sol]
+solutions = [casadi_sol, klu_sol, casadi_sol]
 plot = pybamm.QuickPlot(models, mesh, solutions)
 plot.dynamic_plot()
diff --git a/examples/scripts/compare_lithium_ion.py b/examples/scripts/compare_lithium_ion.py
index 1aca4c781..d92fb46f8 100644
--- a/examples/scripts/compare_lithium_ion.py
+++ b/examples/scripts/compare_lithium_ion.py
@@ -45,7 +45,7 @@ for model in models:
 
 # solve model
 solutions = [None] * len(models)
-t_eval = np.linspace(0, 0.17, 100)
+t_eval = np.linspace(0, 0.3, 100)
 for i, model in enumerate(models):
     solutions[i] = model.default_solver.solve(model, t_eval)
 
diff --git a/pybamm/quick_plot.py b/pybamm/quick_plot.py
index b2470f0ac..472dc1237 100644
--- a/pybamm/quick_plot.py
+++ b/pybamm/quick_plot.py
@@ -8,7 +8,7 @@ from collections import defaultdict
 
 def ax_min(data):
     "Calculate appropriate minimum axis value for plotting"
-    data_min = np.min(data)
+    data_min = np.nanmin(data)
     if data_min <= 0:
         return 1.04 * data_min
     else:
@@ -17,7 +17,7 @@ def ax_min(data):
 
 def ax_max(data):
     "Calculate appropriate maximum axis value for plotting"
-    data_max = np.max(data)
+    data_max = np.nanmax(data)
     if data_max <= 0:
         return 0.96 * data_max
     else:
@@ -256,14 +256,26 @@ class QuickPlot(object):
             # Get min and max y values
             y_min = np.min(
                 [
-                    ax_min(var(self.ts[i], **{spatial_var_name: spatial_var_value}))
+                    ax_min(
+                        var(
+                            self.ts[i],
+                            **{spatial_var_name: spatial_var_value},
+                            warn=False
+                        )
+                    )
                     for i, variable_list in enumerate(variable_lists)
                     for var in variable_list
                 ]
             )
             y_max = np.max(
                 [
-                    ax_max(var(self.ts[i], **{spatial_var_name: spatial_var_value}))
+                    ax_max(
+                        var(
+                            self.ts[i],
+                            **{spatial_var_name: spatial_var_value},
+                            warn=False
+                        )
+                    )
                     for i, variable_list in enumerate(variable_lists)
                     for var in variable_list
                 ]
diff --git a/pybamm/solvers/algebraic_solver.py b/pybamm/solvers/algebraic_solver.py
index 3f085daf0..ff69455ef 100644
--- a/pybamm/solvers/algebraic_solver.py
+++ b/pybamm/solvers/algebraic_solver.py
@@ -24,6 +24,7 @@ class AlgebraicSolver(object):
     def __init__(self, method="lm", tol=1e-6):
         self.method = method
         self.tol = tol
+        self.name = "Algebraic solver ({})".format(method)
 
     @property
     def method(self):
@@ -51,7 +52,7 @@ class AlgebraicSolver(object):
             equations.
 
         """
-        pybamm.logger.info("Start solving {}".format(model.name))
+        pybamm.logger.info("Start solving {} with {}".format(model.name, self.name))
 
         # Set up
         timer = pybamm.Timer()
@@ -87,7 +88,6 @@ class AlgebraicSolver(object):
 
         # Assign times
         solution.solve_time = timer.time() - solve_start_time
-        solution.total_time = timer.time() - start_time
         solution.set_up_time = set_up_time
 
         pybamm.logger.info("Finish solving {}".format(model.name))
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
index 3b95b10f7..37f2fd819 100644
--- a/pybamm/solvers/base_solver.py
+++ b/pybamm/solvers/base_solver.py
@@ -20,6 +20,7 @@ class BaseSolver(object):
         self._method = method
         self._rtol = rtol
         self._atol = atol
+        self.name = "Base solver"
 
     @property
     def method(self):
@@ -64,7 +65,7 @@ class BaseSolver(object):
             If an empty model is passed (`model.rhs = {}` and `model.algebraic={}`)
 
         """
-        pybamm.logger.info("Start solving {}".format(model.name))
+        pybamm.logger.info("Start solving {} with {}".format(model.name, self.name))
 
         # Make sure model isn't empty
         if len(model.rhs) == 0 and len(model.algebraic) == 0:
@@ -84,7 +85,6 @@ class BaseSolver(object):
 
         # 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))
@@ -97,7 +97,7 @@ class BaseSolver(object):
         )
         return solution
 
-    def step(self, model, dt, npts=2):
+    def step(self, model, dt, npts=2, log=True):
         """
         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
@@ -129,36 +129,40 @@ class BaseSolver(object):
 
         # Run set up on first step
         if not hasattr(self, "y0"):
-            start_time = timer.time()
+            pybamm.logger.info(
+                "Start stepping {} with {}".format(model.name, self.name)
+            )
+
             if model.convert_to_format == "casadi" or isinstance(
                 self, pybamm.CasadiSolver
             ):
                 self.set_up_casadi(model)
             else:
+                pybamm.logger.debug(
+                    "Start stepping {} with {}".format(model.name, self.name)
+                )
                 self.set_up(model)
             self.t = 0.0
-            set_up_time = timer.time() - start_time
+            set_up_time = timer.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))
+        pybamm.logger.debug("Finish stepping {} ({})".format(model.name, termination))
         if set_up_time:
-            pybamm.logger.info(
+            pybamm.logger.debug(
                 "Set-up time: {}, Step time: {}, Total time: {}".format(
                     timer.format(solution.set_up_time),
                     timer.format(solution.solve_time),
@@ -166,7 +170,7 @@ class BaseSolver(object):
                 )
             )
         else:
-            pybamm.logger.info(
+            pybamm.logger.debug(
                 "Step time: {}".format(timer.format(solution.solve_time))
             )
         return solution
diff --git a/pybamm/solvers/casadi_solver.py b/pybamm/solvers/casadi_solver.py
index 373bbc056..b31b5f32e 100644
--- a/pybamm/solvers/casadi_solver.py
+++ b/pybamm/solvers/casadi_solver.py
@@ -9,28 +9,171 @@ import numpy as np
 class CasadiSolver(pybamm.DaeSolver):
     """Solve a discretised model, using CasADi.
 
+    **Extends**: :class:`pybamm.DaeSolver`
+
     Parameters
     ----------
+    method : str, optional
+        The method to use for solving the system ('cvodes', for ODEs, or 'idas', for
+        DAEs). Default is 'idas'.
+    mode : str
+            How to solve the model (default is "safe"):
+
+            - "fast": perform direct integration, without accounting for events. \
+            Recommended when simulating a drive cycle or other simulation where \
+            no events should be triggered.
+            - "safe": perform step-and-check integration, checking whether events have \
+            been triggered. Recommended for simulations of a full charge or discharge.
     rtol : float, optional
         The relative tolerance for the solver (default is 1e-6).
     atol : float, optional
         The absolute tolerance for the solver (default is 1e-6).
+    root_method : str, optional
+        The method to use for finding consistend initial conditions. Default is 'lm'.
+    root_tol : float, optional
+        The tolerance for root-finding. Default is 1e-6.
+    max_step_decrease_counts : float, optional
+        The maximum number of times step size can be decreased before an error is
+        raised. Default is 5.
+    extra_options : keyword arguments, optional
+        Any extra keyword-arguments; these are passed directly to the CasADi integrator.
+        Please consult `CasADi documentation <https://tinyurl.com/y5rk76os>`_ for
+        details.
 
-    **Extends**: :class:`pybamm.DaeSolver`
     """
 
     def __init__(
         self,
         method="idas",
+        mode="safe",
         rtol=1e-6,
         atol=1e-6,
         root_method="lm",
         root_tol=1e-6,
-        max_steps=1000,
+        max_step_decrease_count=5,
         **extra_options,
     ):
-        super().__init__(method, rtol, atol, root_method, root_tol, max_steps)
+        super().__init__(method, rtol, atol, root_method, root_tol)
+        if mode in ["safe", "fast"]:
+            self.mode = mode
+        else:
+            raise ValueError(
+                """
+                invalid mode '{}'. Must be either 'safe', for solving with events,
+                or 'fast', for solving quickly without events""".format(
+                    mode
+                )
+            )
+        self.max_step_decrease_count = max_step_decrease_count
         self.extra_options = extra_options
+        self.name = "CasADi solver ({}) with '{}' mode".format(method, mode)
+
+    def solve(self, model, t_eval):
+        """
+        Execute the solver setup and calculate the solution of the model at
+        specified times.
+
+        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
+
+        Raises
+        ------
+        :class:`pybamm.ValueError`
+            If an invalid mode is passed.
+        :class:`pybamm.ModelError`
+            If an empty model is passed (`model.rhs = {}` and `model.algebraic={}`)
+
+        """
+        if self.mode == "fast":
+            # Solve model normally by calling the solve method from parent class
+            return super().solve(model, t_eval)
+        elif model.events == {}:
+            pybamm.logger.info("No events found, running fast mode")
+            # Solve model normally by calling the solve method from parent class
+            return super().solve(model, t_eval)
+        elif self.mode == "safe":
+            # Step-and-check
+            timer = pybamm.Timer()
+            self.set_up_casadi(model)
+            set_up_time = timer.time()
+            init_event_signs = np.sign(
+                np.concatenate([event(0, self.y0) for event in self.event_funs])
+            )
+            solution = None
+            pybamm.logger.info(
+                "Start solving {} with {} in 'safe' mode".format(model.name, self.name)
+            )
+            self.t = 0.0
+            for dt in np.diff(t_eval):
+                # Step
+                solved = False
+                count = 0
+                while not solved:
+                    # Try to solve with the current step, if it fails then halve the
+                    # step size and try again. This will make solution.t slightly
+                    # different to t_eval, but shouldn't matter too much as it should
+                    # only happen near events.
+                    try:
+                        current_step_sol = self.step(model, dt)
+                        solved = True
+                    except pybamm.SolverError:
+                        dt /= 2
+                    count += 1
+                    if count >= self.max_step_decrease_count:
+                        if solution is None:
+                            t = 0
+                        else:
+                            t = solution.t[-1]
+                        raise pybamm.SolverError(
+                            """
+                            Maximum number of decreased steps occurred at t={}. Try
+                            solving the model up to this time only
+                            """.format(
+                                t
+                            )
+                        )
+                # Check most recent y
+                new_event_signs = np.sign(
+                    np.concatenate(
+                        [
+                            event(0, current_step_sol.y[:, -1])
+                            for event in self.event_funs
+                        ]
+                    )
+                )
+                # Exit loop if the sign of an event changes
+                if (new_event_signs != init_event_signs).any():
+                    solution.termination = "event"
+                    solution.t_event = solution.t[-1]
+                    solution.y_event = solution.y[:, -1]
+                    break
+                else:
+                    if not solution:
+                        # create solution object on first step
+                        solution = current_step_sol
+                        solution.set_up_time = set_up_time
+                    else:
+                        # append solution from the current step to solution
+                        solution.append(current_step_sol)
+
+            # Calculate more exact termination reason
+            solution.termination = self.get_termination_reason(solution, self.events)
+            pybamm.logger.info(
+                "Finish solving {} ({})".format(model.name, solution.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 compute_solution(self, model, t_eval):
         """Calculate the solution of the model at specified times. In this class, we
@@ -49,7 +192,7 @@ class CasadiSolver(pybamm.DaeSolver):
         timer = pybamm.Timer()
 
         solve_start_time = timer.time()
-        pybamm.logger.info("Calling DAE solver")
+        pybamm.logger.debug("Calling DAE solver")
         solution = self.integrate_casadi(
             self.casadi_problem, self.y0, t_eval, mass_matrix=model.mass_matrix.entries
         )
diff --git a/pybamm/solvers/dae_solver.py b/pybamm/solvers/dae_solver.py
index 39fe18f36..74001b7b9 100644
--- a/pybamm/solvers/dae_solver.py
+++ b/pybamm/solvers/dae_solver.py
@@ -39,6 +39,7 @@ class DaeSolver(pybamm.BaseSolver):
         self.root_method = root_method
         self.root_tol = root_tol
         self.max_steps = max_steps
+        self.name = "Base DAE solver"
 
     @property
     def root_method(self):
diff --git a/pybamm/solvers/idaklu_solver.py b/pybamm/solvers/idaklu_solver.py
index 131f0fdd8..d7e985d83 100644
--- a/pybamm/solvers/idaklu_solver.py
+++ b/pybamm/solvers/idaklu_solver.py
@@ -43,6 +43,7 @@ class IDAKLUSolver(pybamm.DaeSolver):
             raise ImportError("KLU is not installed")
 
         super().__init__("ida", rtol, atol, root_method, root_tol, max_steps)
+        self.name = "IDA KLU solver"
 
     def integrate(self, residuals, y0, t_eval, events, mass_matrix, jacobian):
         """
diff --git a/pybamm/solvers/ode_solver.py b/pybamm/solvers/ode_solver.py
index 9b650241f..a1655d745 100644
--- a/pybamm/solvers/ode_solver.py
+++ b/pybamm/solvers/ode_solver.py
@@ -19,6 +19,7 @@ class OdeSolver(pybamm.BaseSolver):
 
     def __init__(self, method=None, rtol=1e-6, atol=1e-6):
         super().__init__(method, rtol, atol)
+        self.name = "Base ODE solver"
 
     def compute_solution(self, model, t_eval):
         """Calculate the solution of the model at specified times.
diff --git a/pybamm/solvers/scikits_dae_solver.py b/pybamm/solvers/scikits_dae_solver.py
index b45425e81..1e572f4af 100644
--- a/pybamm/solvers/scikits_dae_solver.py
+++ b/pybamm/solvers/scikits_dae_solver.py
@@ -48,6 +48,7 @@ class ScikitsDaeSolver(pybamm.DaeSolver):
             raise ImportError("scikits.odes is not installed")
 
         super().__init__(method, rtol, atol, root_method, root_tol, max_steps)
+        self.name = "Scikits DAE solver ({})".format(method)
 
     def integrate(
         self, residuals, y0, t_eval, events=None, mass_matrix=None, jacobian=None
diff --git a/pybamm/solvers/scikits_ode_solver.py b/pybamm/solvers/scikits_ode_solver.py
index 542b15d19..6fbe70127 100644
--- a/pybamm/solvers/scikits_ode_solver.py
+++ b/pybamm/solvers/scikits_ode_solver.py
@@ -40,6 +40,7 @@ class ScikitsOdeSolver(pybamm.OdeSolver):
 
         super().__init__(method, rtol, atol)
         self.linsolver = linsolver
+        self.name = "Scikits ODE solver ({})".format(method)
 
     def integrate(
         self, derivs, y0, t_eval, events=None, mass_matrix=None, jacobian=None
diff --git a/pybamm/solvers/scipy_solver.py b/pybamm/solvers/scipy_solver.py
index 23e11483c..fd2af8d76 100644
--- a/pybamm/solvers/scipy_solver.py
+++ b/pybamm/solvers/scipy_solver.py
@@ -22,6 +22,7 @@ class ScipySolver(pybamm.OdeSolver):
 
     def __init__(self, method="BDF", rtol=1e-6, atol=1e-6):
         super().__init__(method, rtol, atol)
+        self.name = "Scipy solver ({})".format(method)
 
     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 d185e5e45..a25d53d1a 100644
--- a/pybamm/solvers/solution.py
+++ b/pybamm/solvers/solution.py
@@ -93,3 +93,8 @@ class Solution(object):
         """
         self.t = np.concatenate((self.t, solution.t[1:]))
         self.y = np.concatenate((self.y, solution.y[:, 1:]), axis=1)
+        self.solve_time += solution.solve_time
+
+    @property
+    def total_time(self):
+        return self.set_up_time + self.solve_time


Test patch:
diff --git a/tests/unit/test_solvers/test_algebraic_solver.py b/tests/unit/test_solvers/test_algebraic_solver.py
index 629eb2952..772426bb6 100644
--- a/tests/unit/test_solvers/test_algebraic_solver.py
+++ b/tests/unit/test_solvers/test_algebraic_solver.py
@@ -107,11 +107,6 @@ class TestAlgebraicSolver(unittest.TestCase):
             model.variables["var2"].evaluate(t=None, y=solution.y), sol[100:]
         )
 
-        # Test time
-        self.assertGreater(
-            solution.total_time, solution.solve_time + solution.set_up_time
-        )
-
         # Test without jacobian
         model.use_jacobian = False
         solution_no_jac = solver.solve(model)
diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
index 6d6b45eec..89a3f3a17 100644
--- a/tests/unit/test_solvers/test_casadi_solver.py
+++ b/tests/unit/test_solvers/test_casadi_solver.py
@@ -5,7 +5,7 @@ import casadi
 import pybamm
 import unittest
 import numpy as np
-from tests import get_mesh_for_testing
+from tests import get_mesh_for_testing, get_discretisation_for_testing
 import warnings
 
 
@@ -45,21 +45,48 @@ class TestCasadiSolver(unittest.TestCase):
 
         y0 = np.array([1])
         t_eval = np.linspace(0, 3, 100)
-        solver = pybamm.CasadiSolver(
-            rtol=1e-8,
-            atol=1e-8,
-            method="idas",
-            disable_internal_warnings=True,
-            regularity_check=False,
-        )
+        solver = pybamm.CasadiSolver()
         problem = {"x": y, "ode": sqrt_decay}
         # Expect solver to fail when y goes negative
         with self.assertRaises(pybamm.SolverError):
             solver.integrate_casadi(problem, y0, t_eval)
 
+        # Set up as a model and solve
+        # Create model
+        model = pybamm.BaseModel()
+        domain = ["negative electrode", "separator", "positive electrode"]
+        var = pybamm.Variable("var", domain=domain)
+        model.rhs = {var: -pybamm.Function(np.sqrt, var)}
+        model.initial_conditions = {var: 1}
+        # add events so that safe mode is used (won't be triggered)
+        model.events = {"10": var - 10}
+        # 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)
+        # Solve with failure at t=2
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="idas")
+        t_eval = np.linspace(0, 20, 100)
+        with self.assertRaises(pybamm.SolverError):
+            solver.solve(model, t_eval)
+        # Solve with failure at t=0
+        model.initial_conditions = {var: 0}
+        disc.process_model(model)
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="idas")
+        t_eval = np.linspace(0, 20, 100)
+        with self.assertRaises(pybamm.SolverError):
+            solver.solve(model, t_eval)
+
         # Turn warnings back on
         warnings.simplefilter("default")
 
+    def test_bad_mode(self):
+        with self.assertRaisesRegex(ValueError, "invalid mode"):
+            pybamm.CasadiSolver(mode="bad mode")
+
     def test_model_solver(self):
         # Create model
         model = pybamm.BaseModel()
@@ -81,9 +108,42 @@ class TestCasadiSolver(unittest.TestCase):
         np.testing.assert_array_equal(solution.t, t_eval)
         np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
 
-        # Test time
-        self.assertGreater(
-            solution.total_time, solution.solve_time + solution.set_up_time
+        # Safe mode (enforce events that won't be triggered)
+        model.events = {"an event": var + 1}
+        disc.process_model(model)
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8, method="idas")
+        t_eval = np.linspace(0, 1, 100)
+        solution = solver.solve(model, t_eval)
+        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_solver_events(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.events = {
+            "var1 = 1.5": pybamm.min(var1 - 1.5),
+            "var2 = 2.5": pybamm.min(var2 - 2.5),
+        }
+        disc = get_discretisation_for_testing()
+        disc.process_model(model)
+
+        # Solve
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
+        t_eval = np.linspace(0, 5, 100)
+        solution = solver.solve(model, t_eval)
+        np.testing.assert_array_less(solution.y[0], 1.5)
+        np.testing.assert_array_less(solution.y[-1], 2.5)
+        np.testing.assert_array_almost_equal(
+            solution.y[0], np.exp(0.1 * solution.t), decimal=5
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[-1], 2 * np.exp(0.1 * solution.t), decimal=5
         )
 
     def test_model_step(self):
diff --git a/tests/unit/test_solvers/test_scikits_solvers.py b/tests/unit/test_solvers/test_scikits_solvers.py
index c0977fd70..b26492120 100644
--- a/tests/unit/test_solvers/test_scikits_solvers.py
+++ b/tests/unit/test_solvers/test_scikits_solvers.py
@@ -411,11 +411,6 @@ 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))
 
-        # Test time
-        self.assertGreater(
-            solution.total_time, solution.solve_time + solution.set_up_time
-        )
-
     def test_model_solver_ode_events(self):
         # Create model
         model = pybamm.BaseModel()
@@ -508,11 +503,6 @@ class TestScikitsSolvers(unittest.TestCase):
         np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
         np.testing.assert_allclose(solution.y[-1], 2 * np.exp(0.1 * solution.t))
 
-        # Test time
-        self.assertGreater(
-            solution.total_time, solution.solve_time + solution.set_up_time
-        )
-
     def test_model_solver_dae_bad_ics(self):
         # Create model
         model = pybamm.BaseModel()
diff --git a/tests/unit/test_solvers/test_scipy_solver.py b/tests/unit/test_solvers/test_scipy_solver.py
index f5d7e783c..c083dbabc 100644
--- a/tests/unit/test_solvers/test_scipy_solver.py
+++ b/tests/unit/test_solvers/test_scipy_solver.py
@@ -154,7 +154,7 @@ class TestScipySolver(unittest.TestCase):
         np.testing.assert_allclose(solution.y[0], np.exp(0.1 * solution.t))
 
         # Test time
-        self.assertGreater(
+        self.assertEqual(
             solution.total_time, solution.solve_time + solution.set_up_time
         )