Issue 804 initialize casadi

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+-   Changed rootfinding algorithm to CasADi, scipy.optimize.root still accessible as an option ([#844](https://github.com/pybamm-team/PyBaMM/pull/844))
 -   Added NCA parameter set ([#824](https://github.com/pybamm-team/PyBaMM/pull/824))
 -   Added functionality to `Solution` that automatically gets `t_eval` from the data when simulating drive cycles and performs checks to ensure the output has the required resolution to accurately capture the input current ([#819](https://github.com/pybamm-team/PyBaMM/pull/819))
 -   Added options to export a solution to matlab or csv ([#811](https://github.com/pybamm-team/PyBaMM/pull/811))

pybamm/solvers/base_solver.py

@@ -23,7 +23,10 @@ class BaseSolver(object):
     atol : float, optional
         The absolute tolerance for the solver (default is 1e-6).
     root_method : str, optional
-        The method to use to find initial conditions (default is "lm")
+        The method to use to find initial conditions (default is "casadi"). If "casadi",
+        the solver uses casadi's Newton rootfinding algorithm to find initial
+        conditions. Otherwise, the solver uses 'scipy.optimize.root' with method
+        specified by 'root_method' (e.g. "lm", "hybr", ...)
     root_tol : float, optional
         The tolerance for the initial-condition solver (default is 1e-6).
     max_steps: int, optional
@@ -36,7 +39,7 @@ def __init__(
         method=None,
         rtol=1e-6,
         atol=1e-6,
-        root_method="lm",
+        root_method="casadi",
         root_tol=1e-6,
         max_steps=1000,
     ):
@@ -125,10 +128,11 @@ def set_up(self, model, inputs=None):
                 "Cannot use ODE solver '{}' to solve DAE model".format(self.name)
             )
 
+        if self.ode_solver is True:
+            self.root_method = None
         if (
-            isinstance(self, pybamm.CasadiSolver)
-            and model.convert_to_format != "casadi"
-        ):
+            isinstance(self, pybamm.CasadiSolver) or self.root_method == "casadi"
+        ) and model.convert_to_format != "casadi":
             pybamm.logger.warning(
                 f"Converting {model.name} to CasADi for solving with CasADi solver"
             )
@@ -268,6 +272,16 @@ def report(string):
         model.terminate_events_eval = terminate_events_eval
         model.discontinuity_events_eval = discontinuity_events_eval
 
+        # Save CasADi functions for the CasADi solver
+        # Note: when we pass to casadi the ode part of the problem must be in explicit
+        # form so we pre-multiply by the inverse of the mass matrix
+        if self.root_method == "casadi" or isinstance(self, pybamm.CasadiSolver):
+            mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
+            explicit_rhs = mass_matrix_inv @ rhs(t_casadi, y_casadi, u_casadi_stacked)
+            model.casadi_rhs = casadi.Function(
+                "rhs", [t_casadi, y_casadi, u_casadi_stacked], [explicit_rhs]
+            )
+            model.casadi_algebraic = algebraic
         # Calculate consistent initial conditions for the algebraic equations
         if len(model.algebraic) > 0:
             all_states = pybamm.NumpyConcatenation(
@@ -278,26 +292,13 @@ def report(string):
             model.residuals_eval = residuals_eval
             model.jacobian_eval = jacobian_eval
             y0_guess = y0.flatten()
-            model.y0 = self.calculate_consistent_state(model, 0, y0_guess)
+            model.y0 = self.calculate_consistent_state(model, 0, y0_guess, inputs)
         else:
             # can use DAE solver to solve ODE model
             model.residuals_eval = Residuals(rhs, "residuals", model)
             model.jacobian_eval = jac_rhs
             model.y0 = y0.flatten()
 
-        # Save CasADi functions for the CasADi solver
-        # Note: when we pass to casadi the ode part of the problem must be in explicit
-        # form so we pre-multiply by the inverse of the mass matrix
-        if model.convert_to_format == "casadi" and isinstance(
-            self, pybamm.CasadiSolver
-        ):
-            mass_matrix_inv = casadi.MX(model.mass_matrix_inv.entries)
-            explicit_rhs = mass_matrix_inv @ rhs(t_casadi, y_casadi, u_casadi_stacked)
-            model.casadi_rhs = casadi.Function(
-                "rhs", [t_casadi, y_casadi, u_casadi_stacked], [explicit_rhs]
-            )
-            model.casadi_algebraic = algebraic
-
         pybamm.logger.info("Finish solver set-up")
 
     def set_inputs(self, model, ext_and_inputs):
@@ -319,7 +320,7 @@ def set_inputs(self, model, ext_and_inputs):
         if model.jacobian_eval:
             model.jacobian_eval.set_inputs(ext_and_inputs)
 
-    def calculate_consistent_state(self, model, time=0, y0_guess=None):
+    def calculate_consistent_state(self, model, time=0, y0_guess=None, inputs=None):
         """
         Calculate consistent state for the algebraic equations through
         root-finding
@@ -328,6 +329,12 @@ def calculate_consistent_state(self, model, time=0, y0_guess=None):
         ----------
         model : :class:`pybamm.BaseModel`
             The model for which to calculate initial conditions.
+        time : float
+            The time at which to calculate the states
+        y0_guess : :class:`np.array`
+            Guess for the rootfinding
+        inputs : dict, optional
+            Any input parameters to pass to the model when solving
 
         Returns
         -------
@@ -336,73 +343,109 @@ def calculate_consistent_state(self, model, time=0, y0_guess=None):
             of the algebraic equations)
         """
         pybamm.logger.info("Start calculating consistent states")
-        rhs = model.rhs_eval
-        algebraic = model.algebraic_eval
-        jac = model.jac_algebraic_eval
         if y0_guess is None:
             y0_guess = model.concatenated_initial_conditions.flatten()
 
         # Split y0_guess into differential and algebraic
-        len_rhs = rhs(time, y0_guess).shape[0]
+        len_rhs = model.rhs_eval(time, y0_guess).shape[0]
         y0_diff, y0_alg_guess = np.split(y0_guess, [len_rhs])
+        inputs = inputs or {}
 
-        def root_fun(y0_alg):
-            "Evaluates algebraic using y0_diff (fixed) and y0_alg (changed by algo)"
-            y0 = np.concatenate([y0_diff, y0_alg])
-            out = algebraic(time, y0)
-            pybamm.logger.debug(
-                "Evaluating algebraic equations at t={}, L2-norm is {}".format(
-                    time * model.timescale, np.linalg.norm(out)
-                )
+        # Solve using casadi or scipy
+        if self.root_method == "casadi":
+            # Set up
+            u_stacked = casadi.vertcat(*[x for x in inputs.values()])
+            u = casadi.MX.sym("u", u_stacked.shape[0])
+            y_alg = casadi.MX.sym("y_alg", y0_alg_guess.shape[0])
+            y = casadi.vertcat(y0_diff, y_alg)
+            alg_root = model.casadi_algebraic(time, y, u)
+            # Solve
+            # set error_on_fail to False and just check the final output is small
+            # enough
+            roots = casadi.rootfinder(
+                "roots",
+                "newton",
+                dict(x=y_alg, p=u, g=alg_root),
+                {"abstol": self.root_tol},
             )
-            return out
-
-        if jac:
-            if issparse(jac(0, y0_guess)):
+            try:
+                y0_alg = roots(y0_alg_guess, u_stacked).full().flatten()
+                success = True
+                message = None
+                # Check final output
+                fun = model.casadi_algebraic(
+                    time, casadi.vertcat(y0_diff, y0_alg), u_stacked
+                )
+            except RuntimeError as err:
+                success = False
+                message = err.args[0]
+                fun = None
+        else:
+            algebraic = model.algebraic_eval
+            jac = model.jac_algebraic_eval
+
+            def root_fun(y0_alg):
+                "Evaluates algebraic using y0_diff (fixed) and y0_alg (changed by algo)"
+                y0 = np.concatenate([y0_diff, y0_alg])
+                out = algebraic(time, y0)
+                pybamm.logger.debug(
+                    "Evaluating algebraic equations at t={}, L2-norm is {}".format(
+                        time * model.timescale, np.linalg.norm(out)
+                    )
+                )
+                return out
 
-                def jac_fn(y0_alg):
-                    """
-                    Evaluates jacobian using y0_diff (fixed) and y0_alg (varying)
-                    """
-                    y0 = np.concatenate([y0_diff, y0_alg])
-                    return jac(0, y0)[:, len_rhs:].toarray()
+            if jac:
+                if issparse(jac(0, y0_guess)):
 
-            else:
+                    def jac_fn(y0_alg):
+                        """
+                        Evaluates jacobian using y0_diff (fixed) and y0_alg (varying)
+                        """
+                        y0 = np.concatenate([y0_diff, y0_alg])
+                        return jac(0, y0)[:, len_rhs:].toarray()
 
-                def jac_fn(y0_alg):
-                    """
-                    Evaluates jacobian using y0_diff (fixed) and y0_alg (varying)
-                    """
-                    y0 = np.concatenate([y0_diff, y0_alg])
-                    return jac(0, y0)[:, len_rhs:]
+                else:
 
-        else:
-            jac_fn = None
-        # Find the values of y0_alg that are roots of the algebraic equations
-        sol = optimize.root(
-            root_fun,
-            y0_alg_guess,
-            jac=jac_fn,
-            method=self.root_method,
-            tol=self.root_tol,
-        )
-        # Return full set of consistent initial conditions (y0_diff unchanged)
-        y0_consistent = np.concatenate([y0_diff, sol.x])
+                    def jac_fn(y0_alg):
+                        """
+                        Evaluates jacobian using y0_diff (fixed) and y0_alg (varying)
+                        """
+                        y0 = np.concatenate([y0_diff, y0_alg])
+                        return jac(0, y0)[:, len_rhs:]
 
-        if sol.success and np.all(sol.fun < self.root_tol * len(sol.x)):
+            else:
+                jac_fn = None
+            # Find the values of y0_alg that are roots of the algebraic equations
+            sol = optimize.root(
+                root_fun,
+                y0_alg_guess,
+                jac=jac_fn,
+                method=self.root_method,
+                tol=self.root_tol,
+            )
+            # Set outputs
+            y0_alg = sol.x
+            success = sol.success
+            fun = sol.fun
+            message = sol.message
+
+        if success and np.all(fun < self.root_tol * len(y0_alg)):
+            # Return full set of consistent initial conditions (y0_diff unchanged)
+            y0_consistent = np.concatenate([y0_diff, y0_alg])
             pybamm.logger.info("Finish calculating consistent initial conditions")
             return y0_consistent
-        elif not sol.success:
+        elif not success:
             raise pybamm.SolverError(
-                "Could not find consistent initial conditions: {}".format(sol.message)
+                "Could not find consistent initial conditions: {}".format(message)
             )
         else:
             raise pybamm.SolverError(
                 """
                 Could not find consistent initial conditions: solver terminated
                 successfully, but maximum solution error ({}) above tolerance ({})
                 """.format(
-                    np.max(sol.fun), self.root_tol * len(sol.x)
+                    np.max(fun), self.root_tol * len(y0_alg)
                 )
             )
 
@@ -514,7 +557,7 @@ def solve(self, model, t_eval, external_variables=None, inputs=None):
                 )
             else:
                 t_eval_dimensionless = np.insert(
-                    t_eval_dimensionless, dindex, [dtime - eps, dtime + eps],
+                    t_eval_dimensionless, dindex, [dtime - eps, dtime + eps]
                 )
         end_indices.append(len(t_eval_dimensionless))
 
@@ -548,7 +591,10 @@ def solve(self, model, t_eval, external_variables=None, inputs=None):
                 last_state = solution.y[:, -1]
                 if len(model.algebraic) > 0:
                     model.y0 = self.calculate_consistent_state(
-                        model, t_eval_dimensionless[end_index], last_state
+                        model,
+                        t_eval_dimensionless[end_index],
+                        last_state,
+                        ext_and_inputs,
                     )
                 else:
                     model.y0 = last_state

pybamm/solvers/casadi_solver.py

@@ -47,7 +47,7 @@ def __init__(
         mode="safe",
         rtol=1e-6,
         atol=1e-6,
-        root_method="lm",
+        root_method="casadi",
         root_tol=1e-6,
         max_step_decrease_count=5,
         **extra_options,

pybamm/solvers/idaklu_solver.py

@@ -36,7 +36,7 @@ class IDAKLUSolver(pybamm.BaseSolver):
     """
 
     def __init__(
-        self, rtol=1e-6, atol=1e-6, root_method="lm", root_tol=1e-6, max_steps=1000
+        self, rtol=1e-6, atol=1e-6, root_method="casadi", root_tol=1e-6, max_steps=1000
     ):
 
         if idaklu_spec is None:

pybamm/solvers/scikits_dae_solver.py

@@ -40,7 +40,7 @@ def __init__(
         method="ida",
         rtol=1e-6,
         atol=1e-6,
-        root_method="lm",
+        root_method="casadi",
         root_tol=1e-6,
         max_steps=1000,
     ):

setup.py

@@ -492,7 +492,7 @@ def load_version():
     # List of dependencies
     install_requires=[
         "numpy>=1.16",
-        "scipy>=1.0",
+        "scipy>=1.3",
         "pandas>=0.24",
         "anytree>=2.4.3",
         "autograd>=1.2",