#943 recalculate initial conditions

Author: valentinsulzer

pybamm/solvers/base_solver.py

@@ -142,7 +142,9 @@ def set_up(self, model, inputs=None):
                 )
 
         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 @@ def set_up(self, model, inputs=None):
         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 @@ def report(string):
                     "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 @@ def jac_fn(y0_alg):
             )
         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 @@ def solve(self, model, t_eval=None, external_variables=None, inputs=None):
             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 @@ def step(
             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

tests/unit/test_solvers/test_casadi_solver.py

@@ -255,6 +255,42 @@ def test_model_solver_with_inputs(self):
         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()

tests/unit/test_solvers/test_scikits_solvers.py

@@ -72,6 +72,7 @@ def test_dae_integrate_bad_ics(self):
 
         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 @@ def test_model_solver_dae_inputs_events(self):
             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()

tests/unit/test_solvers/test_scipy_solver.py

@@ -288,6 +288,29 @@ def test_model_solver_with_inputs_with_casadi(self):
         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")