Repository: pybamm-team/PyBaMM

Issue:
Recalculate initial conditions during parameter sweeps
During parameter sweeps, initial conditions only get calculated once (at the first solve), and then are fixed for all future solves. They should be recalculated each time if the algebraic equations or initial conditions change when an input changes.

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0883fa719..16c902554 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,7 @@
 
 ## Bug fixes
 
+-   Fixed solver to recompute initial conditions when inputs are changed ([#951](https://github.com/pybamm-team/PyBaMM/pull/951)
 -   Reformatted thermal submodels ([#938](https://github.com/pybamm-team/PyBaMM/pull/938)
 -   Reformatted electrolyte submodels ([#927](https://github.com/pybamm-team/PyBaMM/pull/927))
 -   Reformatted convection submodels ([#635](https://github.com/pybamm-team/PyBaMM/pull/635))
diff --git a/pybamm/simulation.py b/pybamm/simulation.py
index c52eeabf5..2527fe7b3 100644
--- a/pybamm/simulation.py
+++ b/pybamm/simulation.py
@@ -443,6 +443,7 @@ class Simulation:
                     timer.format(timer.time())
                 )
             )
+        return self.solution
 
     def step(
         self, dt, solver=None, npts=2, external_variables=None, inputs=None, save=True
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
index bb16012c8..f091acb20 100644
--- a/pybamm/solvers/base_solver.py
+++ b/pybamm/solvers/base_solver.py
@@ -142,7 +142,9 @@ class BaseSolver(object):
                 )
 
         inputs = inputs or {}
-        y0 = model.concatenated_initial_conditions.evaluate(0, None, inputs=inputs)
+        model.y0 = model.concatenated_initial_conditions.evaluate(
+            0, None, inputs=inputs
+        ).flatten()
 
         # Set model timescale
         model.timescale_eval = model.timescale.evaluate(inputs=inputs)
@@ -169,18 +171,19 @@ class BaseSolver(object):
         if model.convert_to_format != "casadi":
             simp = pybamm.Simplification()
             # Create Jacobian from concatenated rhs and algebraic
-            y = pybamm.StateVector(slice(0, np.size(y0)))
+            y = pybamm.StateVector(slice(0, np.size(model.y0)))
             # set up Jacobian object, for re-use of dict
             jacobian = pybamm.Jacobian()
         else:
             # Convert model attributes to casadi
             t_casadi = casadi.MX.sym("t")
             y_diff = casadi.MX.sym(
-                "y_diff", len(model.concatenated_rhs.evaluate(0, y0, inputs=inputs))
+                "y_diff",
+                len(model.concatenated_rhs.evaluate(0, model.y0, inputs=inputs)),
             )
             y_alg = casadi.MX.sym(
                 "y_alg",
-                len(model.concatenated_algebraic.evaluate(0, y0, inputs=inputs)),
+                len(model.concatenated_algebraic.evaluate(0, model.y0, inputs=inputs)),
             )
             y_casadi = casadi.vertcat(y_diff, y_alg)
             p_casadi = {}
@@ -322,36 +325,69 @@ class BaseSolver(object):
                     "rhs", [t_casadi, y_casadi, p_casadi_stacked], [explicit_rhs]
                 )
             model.casadi_algebraic = algebraic
-        if self.algebraic_solver is True:
-            # we don't calculate consistent initial conditions
-            # for an algebraic solver as this will be the job of the algebraic solver
+        if len(model.rhs) == 0:
+            # No rhs equations: residuals is algebraic only
             model.residuals_eval = Residuals(algebraic, "residuals", model)
             model.jacobian_eval = jac_algebraic
-            model.y0 = y0.flatten()
         elif len(model.algebraic) == 0:
-            # can use DAE solver to solve ODE model
-            # - no initial condition initialization needed
+            # No algebraic equations: residuals is rhs only
             model.residuals_eval = Residuals(rhs, "residuals", model)
             model.jacobian_eval = jac_rhs
-            model.y0 = y0.flatten()
         # Calculate consistent initial conditions for the algebraic equations
         else:
-            if len(model.rhs) > 0:
-                all_states = pybamm.NumpyConcatenation(
-                    model.concatenated_rhs, model.concatenated_algebraic
+            all_states = pybamm.NumpyConcatenation(
+                model.concatenated_rhs, model.concatenated_algebraic
+            )
+            # Process again, uses caching so should be quick
+            residuals_eval, jacobian_eval = process(all_states, "residuals")[1:]
+            model.residuals_eval = residuals_eval
+            model.jacobian_eval = jacobian_eval
+
+        pybamm.logger.info("Finish solver set-up")
+
+    def _set_initial_conditions(self, model, inputs, update_rhs):
+        """
+        Set initial conditions for the model. This is skipped if the solver is an
+        algebraic solver (since this would make the algebraic solver redundant), and if
+        the model doesn't have any algebraic equations (since there are no initial
+        conditions to be calculated in this case).
+
+        Parameters
+        ----------
+        model : :class:`pybamm.BaseModel`
+            The model for which to calculate initial conditions.
+        inputs : dict
+            Any input parameters to pass to the model when solving
+        update_rhs : bool
+            Whether to update the rhs. True for 'solve', False for 'step'.
+
+        """
+        if self.algebraic_solver is True:
+            return None
+        elif len(model.algebraic) == 0:
+            if update_rhs is True:
+                # Recalculate initial conditions for the rhs equations
+                model.y0 = model.concatenated_initial_conditions.evaluate(
+                    0, None, inputs=inputs
+                ).flatten()
+            else:
+                return None
+        else:
+            if update_rhs is True:
+                # Recalculate initial conditions for the rhs equations
+                y0_from_inputs = model.concatenated_initial_conditions.evaluate(
+                    0, None, inputs=inputs
+                ).flatten()
+                # Reuse old solution for algebraic equations
+                y0_from_model = model.y0
+                len_rhs = len(
+                    model.concatenated_rhs.evaluate(0, model.y0, inputs=inputs)
                 )
-                # Process again, uses caching so should be quick
-                residuals_eval, jacobian_eval = process(all_states, "residuals")[1:]
-                model.residuals_eval = residuals_eval
-                model.jacobian_eval = jacobian_eval
+                y0_guess = np.r_[y0_from_inputs[:len_rhs], y0_from_model[len_rhs:]]
             else:
-                model.residuals_eval = Residuals(algebraic, "residuals", model)
-                model.jacobian_eval = jac_algebraic
-            y0_guess = y0.flatten()
+                y0_guess = model.y0
             model.y0 = self.calculate_consistent_state(model, 0, y0_guess, inputs)
 
-        pybamm.logger.info("Finish solver set-up")
-
     def calculate_consistent_state(self, model, time=0, y0_guess=None, inputs=None):
         """
         Calculate consistent state for the algebraic equations through
@@ -480,12 +516,9 @@ class BaseSolver(object):
             )
         else:
             raise pybamm.SolverError(
-                """
-                Could not find consistent initial conditions: solver terminated
-                successfully, but maximum solution error ({}) above tolerance ({})
-                """.format(
-                    max_fun, self.root_tol
-                )
+                "Could not find consistent initial conditions: solver terminated "
+                "successfully, but maximum solution error "
+                "({}) above tolerance ({})".format(max_fun, self.root_tol)
             )
 
     def solve(self, model, t_eval=None, external_variables=None, inputs=None):
@@ -555,6 +588,10 @@ class BaseSolver(object):
             self.models_set_up.add(model)
         else:
             set_up_time = 0
+
+        # (Re-)calculate consistent initial conditions
+        self._set_initial_conditions(model, ext_and_inputs, update_rhs=True)
+
         # Non-dimensionalise time
         t_eval_dimensionless = t_eval / model.timescale_eval
         # Solve
@@ -758,6 +795,9 @@ class BaseSolver(object):
             model.y0 = old_solution.y[:, -1]
             set_up_time = 0
 
+        # (Re-)calculate consistent initial conditions
+        self._set_initial_conditions(model, ext_and_inputs, update_rhs=False)
+
         # Non-dimensionalise dt
         dt_dimensionless = dt / model.timescale_eval
         # Step


Test patch:
diff --git a/tests/integration/test_experiments.py b/tests/integration/test_experiments.py
index 35656d558..a14f3963f 100644
--- a/tests/integration/test_experiments.py
+++ b/tests/integration/test_experiments.py
@@ -30,6 +30,26 @@ class TestExperiments(unittest.TestCase):
             [cap / 2] * 5 + [0] * 4 + [-cap / 2] * 4,
         )
 
+    def test_rest_discharge_rest(self):
+        # An experiment which requires recomputing consistent states
+        experiment = pybamm.Experiment(
+            ["Rest for 5 minutes", "Discharge at 0.1C until 3V", "Rest for 30 minutes"],
+            period="1 minute",
+        )
+        parameter_values = pybamm.ParameterValues(
+            chemistry=pybamm.parameter_sets.Chen2020
+        )
+        model = pybamm.lithium_ion.DFN()
+        sim = pybamm.Simulation(
+            model,
+            parameter_values=parameter_values,
+            experiment=experiment,
+            solver=pybamm.CasadiSolver(),
+        )
+        sol = sim.solve()
+        np.testing.assert_array_almost_equal(sol["Current [A]"].data[:5], 0)
+        np.testing.assert_array_almost_equal(sol["Current [A]"].data[-29:], 0)
+
     def test_gitt(self):
         experiment = pybamm.Experiment(
             ["Discharge at C/20 for 1 hour", "Rest for 1 hour"] * 10, period="6 minutes"
diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
index 6d14c77eb..1018baa68 100644
--- a/tests/unit/test_solvers/test_casadi_solver.py
+++ b/tests/unit/test_solvers/test_casadi_solver.py
@@ -255,6 +255,42 @@ class TestCasadiSolver(unittest.TestCase):
         np.testing.assert_array_equal(solution.t, t_eval[: len(solution.t)])
         np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t), rtol=1e-06)
 
+    def test_model_solver_dae_inputs_in_initial_conditions(self):
+        # Create model
+        model = pybamm.BaseModel()
+        var1 = pybamm.Variable("var1")
+        var2 = pybamm.Variable("var2")
+        model.rhs = {var1: pybamm.InputParameter("rate") * var1}
+        model.algebraic = {var2: var1 - var2}
+        model.initial_conditions = {
+            var1: pybamm.InputParameter("ic 1"),
+            var2: pybamm.InputParameter("ic 2"),
+        }
+
+        # Solve
+        solver = pybamm.CasadiSolver(rtol=1e-8, atol=1e-8)
+        t_eval = np.linspace(0, 5, 100)
+        solution = solver.solve(
+            model, t_eval, inputs={"rate": -1, "ic 1": 0.1, "ic 2": 2}
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[0], 0.1 * np.exp(-solution.t), decimal=5
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[-1], 0.1 * np.exp(-solution.t), decimal=5
+        )
+
+        # Solve again with different initial conditions
+        solution = solver.solve(
+            model, t_eval, inputs={"rate": -0.1, "ic 1": 1, "ic 2": 3}
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[0], 1 * np.exp(-0.1 * solution.t), decimal=5
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[-1], 1 * np.exp(-0.1 * solution.t), decimal=5
+        )
+
     def test_model_solver_with_external(self):
         # Create model
         model = pybamm.BaseModel()
diff --git a/tests/unit/test_solvers/test_scikits_solvers.py b/tests/unit/test_solvers/test_scikits_solvers.py
index 9eae6fe38..c646075a3 100644
--- a/tests/unit/test_solvers/test_scikits_solvers.py
+++ b/tests/unit/test_solvers/test_scikits_solvers.py
@@ -72,6 +72,7 @@ class TestScikitsSolvers(unittest.TestCase):
 
         t_eval = np.linspace(0, 1, 100)
         solver.set_up(model)
+        solver._set_initial_conditions(model, {}, True)
         # check y0
         np.testing.assert_array_equal(model.y0, [0, 0])
         # check dae solutions
@@ -564,6 +565,42 @@ 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))
 
+    def test_model_solver_dae_inputs_in_initial_conditions(self):
+        # Create model
+        model = pybamm.BaseModel()
+        var1 = pybamm.Variable("var1")
+        var2 = pybamm.Variable("var2")
+        model.rhs = {var1: pybamm.InputParameter("rate") * var1}
+        model.algebraic = {var2: var1 - var2}
+        model.initial_conditions = {
+            var1: pybamm.InputParameter("ic 1"),
+            var2: pybamm.InputParameter("ic 2"),
+        }
+
+        # Solve
+        solver = pybamm.ScikitsDaeSolver(rtol=1e-8, atol=1e-8)
+        t_eval = np.linspace(0, 5, 100)
+        solution = solver.solve(
+            model, t_eval, inputs={"rate": -1, "ic 1": 0.1, "ic 2": 2}
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[0], 0.1 * np.exp(-solution.t), decimal=5
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[-1], 0.1 * np.exp(-solution.t), decimal=5
+        )
+
+        # Solve again with different initial conditions
+        solution = solver.solve(
+            model, t_eval, inputs={"rate": -0.1, "ic 1": 1, "ic 2": 3}
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[0], 1 * np.exp(-0.1 * solution.t), decimal=5
+        )
+        np.testing.assert_array_almost_equal(
+            solution.y[-1], 1 * np.exp(-0.1 * solution.t), decimal=5
+        )
+
     def test_model_solver_dae_with_external(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 76a9c2507..e7393927c 100644
--- a/tests/unit/test_solvers/test_scipy_solver.py
+++ b/tests/unit/test_solvers/test_scipy_solver.py
@@ -288,6 +288,29 @@ class TestScipySolver(unittest.TestCase):
         np.testing.assert_array_equal(solution.t, t_eval[: len(solution.t)])
         np.testing.assert_allclose(solution.y[0], np.exp(-0.1 * solution.t))
 
+    def test_model_solver_inputs_in_initial_conditions(self):
+        # Create model
+        model = pybamm.BaseModel()
+        var1 = pybamm.Variable("var1")
+        model.rhs = {var1: pybamm.InputParameter("rate") * var1}
+        model.initial_conditions = {
+            var1: pybamm.InputParameter("ic 1"),
+        }
+
+        # Solve
+        solver = pybamm.ScipySolver(rtol=1e-8, atol=1e-8)
+        t_eval = np.linspace(0, 5, 100)
+        solution = solver.solve(model, t_eval, inputs={"rate": -1, "ic 1": 0.1})
+        np.testing.assert_array_almost_equal(
+            solution.y[0], 0.1 * np.exp(-solution.t), decimal=5
+        )
+
+        # Solve again with different initial conditions
+        solution = solver.solve(model, t_eval, inputs={"rate": -0.1, "ic 1": 1})
+        np.testing.assert_array_almost_equal(
+            solution.y[0], 1 * np.exp(-0.1 * solution.t), decimal=5
+        )
+
 
 if __name__ == "__main__":
     print("Add -v for more debug output")