From 9572e8e9ff646bad4de53391d7745c8745dd694e Mon Sep 17 00:00:00 2001 From: Valentin Sulzer Date: Fri, 18 Oct 2019 21:08:59 -0400 Subject: [PATCH] #664 add casadi set up to solver --- examples/scripts/compare_lithium_ion.py | 6 +- pybamm/discretisations/discretisation.py | 2 +- pybamm/models/base_model.py | 18 +++-- pybamm/solvers/algebraic_solver.py | 4 +- pybamm/solvers/base_solver.py | 26 ++++-- pybamm/solvers/dae_solver.py | 4 +- pybamm/solvers/ode_solver.py | 81 ++++++++++++++++++- .../test_models/standard_model_tests.py | 3 +- 8 files changed, 120 insertions(+), 24 deletions(-) diff --git a/examples/scripts/compare_lithium_ion.py b/examples/scripts/compare_lithium_ion.py index 41ad311860..7fcfae8788 100644 --- a/examples/scripts/compare_lithium_ion.py +++ b/examples/scripts/compare_lithium_ion.py @@ -18,9 +18,9 @@ # load models options = {"thermal": "isothermal"} models = [ - # pybamm.lithium_ion.SPM(options), - # pybamm.lithium_ion.SPMe(options), - pybamm.lithium_ion.DFN(options) + pybamm.lithium_ion.SPM(options), + pybamm.lithium_ion.SPMe(options), + pybamm.lithium_ion.DFN(options), ] diff --git a/pybamm/discretisations/discretisation.py b/pybamm/discretisations/discretisation.py index a00e8f6aec..b96f04f7ba 100644 --- a/pybamm/discretisations/discretisation.py +++ b/pybamm/discretisations/discretisation.py @@ -127,7 +127,7 @@ def process_model(self, model, inplace=True): model_disc.options = model.options model_disc.use_jacobian = model.use_jacobian model_disc.use_simplify = model.use_simplify - model_disc.use_to_python = model.use_to_python + model_disc.convert_to_format = model.convert_to_format model_disc.bcs = self.bcs diff --git a/pybamm/models/base_model.py b/pybamm/models/base_model.py index 0470d86773..c2e4ab60bc 100644 --- a/pybamm/models/base_model.py +++ b/pybamm/models/base_model.py @@ -59,11 +59,17 @@ class BaseModel(object): Whether to simplify the expression tress representing the rhs and algebraic equations, Jacobain (if using) and events, before solving the model (default is True) - use_to_python : bool - Whether to convert the expression tress representing the rhs and - algebraic equations, Jacobain (if using) and events into pure python code - that will calculate the result of calling `evaluate(t, y)` on the given - expression tree (default is True) + convert_to_format : str + Whether to convert the expression trees representing the rhs and + algebraic equations, Jacobain (if using) and events into a different format: + + - None: keep PyBaMM expression tree structure. + - "python": convert into pure python code that will calculate the result of + calling `evaluate(t, y)` on the given expression treeself. + - "casadi": convert into CasADi expression tree, which then uses CasADi's + algorithm to calculate the Jacobian. + + Default is "python". """ def __init__(self, name="Unnamed model"): @@ -86,7 +92,7 @@ def __init__(self, name="Unnamed model"): # Default behaviour is to use the jacobian and simplify self.use_jacobian = True self.use_simplify = True - self.use_to_python = True + self.convert_to_format = "python" def _set_dictionary(self, dict, name): """ diff --git a/pybamm/solvers/algebraic_solver.py b/pybamm/solvers/algebraic_solver.py index baea9ebb5d..f3be30d360 100644 --- a/pybamm/solvers/algebraic_solver.py +++ b/pybamm/solvers/algebraic_solver.py @@ -199,14 +199,14 @@ def set_up(self, model): pybamm.logger.info("Simplifying jacobian") jac = jac.simplify() - if model.use_to_python: + if model.convert_to_format == "python": pybamm.logger.info("Converting jacobian to python") jac = pybamm.EvaluatorPython(jac) else: jac = None - if model.use_to_python: + if model.convert_to_format == "python": pybamm.logger.info("Converting algebraic to python") concatenated_algebraic = pybamm.EvaluatorPython(concatenated_algebraic) diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py index ad1ad2b4a8..1e1a0d0f4a 100644 --- a/pybamm/solvers/base_solver.py +++ b/pybamm/solvers/base_solver.py @@ -73,7 +73,10 @@ def solve(self, model, t_eval): # Set up timer = pybamm.Timer() start_time = timer.time() - self.set_up(model) + if model.convert_to_format == "casadi": + self.set_up_casadi(model) + else: + self.set_up(model) set_up_time = timer.time() - start_time # Solve @@ -127,7 +130,10 @@ def step(self, model, dt, npts=2): # Run set up on first step if not hasattr(self, "y0"): start_time = timer.time() - self.set_up(model) + if model.convert_to_format == "casadi": + self.set_up_casadi(model) + else: + self.set_up(model) self.t = 0.0 set_up_time = timer.time() - start_time else: @@ -187,11 +193,17 @@ def set_up(self, model): 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 set_up_casadi(self, model): + """Convert model to casadi format and use their inbuilt functionalities. + + Parameters + ---------- + model : :class:`pybamm.BaseModel` + The model whose solution to calculate. Must have attributes rhs and + initial_conditions """ raise NotImplementedError diff --git a/pybamm/solvers/dae_solver.py b/pybamm/solvers/dae_solver.py index 003403254f..656413c9d0 100644 --- a/pybamm/solvers/dae_solver.py +++ b/pybamm/solvers/dae_solver.py @@ -140,7 +140,7 @@ def set_up(self, model): jac_algebraic = simp.simplify(jac_algebraic) jac = simp.simplify(jac) - if model.use_to_python: + if model.convert_to_format == "python": pybamm.logger.info("Converting jacobian to python") jac_algebraic = pybamm.EvaluatorPython(jac_algebraic) jac = pybamm.EvaluatorPython(jac) @@ -152,7 +152,7 @@ def jac_alg_fn(t, y): jac = None jac_alg_fn = None - if model.use_to_python: + if model.convert_to_format == "python": pybamm.logger.info("Converting RHS to python") concatenated_rhs = pybamm.EvaluatorPython(concatenated_rhs) pybamm.logger.info("Converting algebraic to python") diff --git a/pybamm/solvers/ode_solver.py b/pybamm/solvers/ode_solver.py index db279b2577..cbe2514c99 100644 --- a/pybamm/solvers/ode_solver.py +++ b/pybamm/solvers/ode_solver.py @@ -99,13 +99,13 @@ def set_up(self, model): pybamm.logger.info("Simplifying jacobian") jac_rhs = simp.simplify(jac_rhs) - if model.use_to_python: + if model.convert_to_format == "python": pybamm.logger.info("Converting jacobian to python") jac_rhs = pybamm.EvaluatorPython(jac_rhs) else: jac_rhs = None - if model.use_to_python: + if model.convert_to_format == "python": pybamm.logger.info("Converting RHS to python") concatenated_rhs = pybamm.EvaluatorPython(concatenated_rhs) pybamm.logger.info("Converting events to python") @@ -145,6 +145,83 @@ def jacobian(t, y): self.event_funs = event_funs self.jacobian = jacobian + def set_up_casadi(self, model): + """Convert model to casadi format and use their inbuilt functionalities. + + 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) + + """ + # 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 + + y0 = model.concatenated_initial_conditions[:, 0] + + t_casadi = casadi.SX.sym("t") + y_casadi = casadi.SX.sym("y", len(y0)) + pybamm.logger.info("Converting RHS to CasADi") + concatenated_rhs = concatenated_rhs.to_casadi(t_casadi, y_casadi) + pybamm.logger.info("Converting events to CasADi") + events = { + name: event.to_casadi(t_casadi, y_casadi) for name, event in events.items() + } + + # Create function to evaluate rhs + concatenated_rhs_fn = casadi.Function( + "rhs", [t_casadi, y_casadi], [concatenated_rhs] + ) + + def dydt(t, y): + pybamm.logger.debug("Evaluating RHS for {} at t={}".format(model.name, t)) + dy = concatenated_rhs_fn(t, y).full() + return dy[:, 0] + + # Create event-dependent function to evaluate events + def event_fun(event): + event_fn = casadi.Function("event", [t_casadi, y_casadi], [event]) + + def eval_event(t, y): + return event_fun(t, y) + + return eval_event + + event_funs = [event_fun(event) for event in events.values()] + + # Create function to evaluate jacobian + if model.use_jacobian: + + casadi_jac = casadi.jacobian() + casadi_jac_fn = casadi.Function("jac", [t_casadi, y_casadi], [casadi_jac]) + + def jacobian(t, y): + return casadi_jac_fn(t, y) + + 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/tests/integration/test_models/standard_model_tests.py b/tests/integration/test_models/standard_model_tests.py index 2b5afa1244..ba827e9cd3 100644 --- a/tests/integration/test_models/standard_model_tests.py +++ b/tests/integration/test_models/standard_model_tests.py @@ -172,6 +172,7 @@ def evaluate_model(self, simplify=False, use_known_evals=False, to_python=False) def set_up_model(self, simplify=False, to_python=False): self.model.use_simplify = simplify - self.model.use_to_python = to_python + if to_python is True: + self.model.convert_to_format = "python" self.model.default_solver.set_up(self.model) return None