#899 make inputs more explicit in _integrate

CHANGELOG.md

@@ -24,6 +24,7 @@
 
 ## Breaking changes
 
+-   Changed keyword argument `u` for inputs (when evaluating an object) to `params` 
 -   Removed "set external temperature" and "set external potential" options. Use "external submodels" option instead ([#862](https://github.com/pybamm-team/PyBaMM/pull/862))
 
 # [v0.2.0](https://github.com/pybamm-team/PyBaMM/tree/v0.2.0) - 2020-02-26

pybamm/expression_tree/operations/evaluate.py

@@ -172,7 +172,7 @@ def find_symbols(symbol, constant_symbols, variable_symbols):
         symbol_str = "t"
 
     elif isinstance(symbol, pybamm.InputParameter):
-        symbol_str = "u['{}']".format(symbol.name)
+        symbol_str = "params['{}']".format(symbol.name)
 
     else:
         raise NotImplementedError(

pybamm/solvers/base_solver.py

@@ -227,10 +227,8 @@ def report(string):
                 func_call = Residuals(func, name, model)
             else:
                 func_call = SolverCallable(func, name, model)
-            func_call.set_inputs(inputs)
             if jac is not None:
                 jac_call = SolverCallable(jac, name + "_jac", model)
-                jac_call.set_inputs(inputs)
             else:
                 jac_call = None
             return func, func_call, jac_call
@@ -342,25 +340,6 @@ def report(string):
 
         pybamm.logger.info("Finish solver set-up")
 
-    def set_inputs(self, model, ext_and_inputs):
-        """
-        Set values that are controlled externally, such as external variables and input
-        parameters
-
-        Parameters
-        ----------
-        ext_and_inputs : dict
-            Any external variables or input parameters to pass to the model when solving
-
-        """
-        model.rhs_eval.set_inputs(ext_and_inputs)
-        model.algebraic_eval.set_inputs(ext_and_inputs)
-        model.residuals_eval.set_inputs(ext_and_inputs)
-        for evnt in model.terminate_events_eval:
-            evnt.set_inputs(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, inputs=None):
         """
         Calculate consistent state for the algebraic equations through
@@ -387,16 +366,18 @@ def calculate_consistent_state(self, model, time=0, y0_guess=None, inputs=None):
         if y0_guess is None:
             y0_guess = model.concatenated_initial_conditions.flatten()
 
+        inputs = inputs or {}
+        if model.rhs_eval.form == "casadi":
+            inputs = casadi.vertcat(*[x for x in inputs.values()])
+
         # Split y0_guess into differential and algebraic
-        len_rhs = model.rhs_eval(time, y0_guess).shape[0]
+        len_rhs = model.rhs_eval(time, y0_guess, inputs).shape[0]
         y0_diff, y0_alg_guess = np.split(y0_guess, [len_rhs])
-        inputs = inputs or {}
 
         # Solve using casadi or scipy
         if self.root_method == "casadi":
             # Set up
-            p_stacked = casadi.vertcat(*[x for x in inputs.values()])
-            p = casadi.MX.sym("p", p_stacked.shape[0])
+            p = casadi.MX.sym("p", inputs.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, p)
@@ -408,12 +389,12 @@ def calculate_consistent_state(self, model, time=0, y0_guess=None, inputs=None):
                 {"abstol": self.root_tol},
             )
             try:
-                y0_alg = roots(y0_alg_guess, p_stacked).full().flatten()
+                y0_alg = roots(y0_alg_guess, inputs).full().flatten()
                 success = True
                 message = None
                 # Check final output
                 fun = model.casadi_algebraic(
-                    time, casadi.vertcat(y0_diff, y0_alg), p_stacked
+                    time, casadi.vertcat(y0_diff, y0_alg), inputs
                 )
                 abs_fun = casadi.fabs(fun)
                 max_fun = casadi.mmax(fun)
@@ -429,7 +410,7 @@ def calculate_consistent_state(self, model, time=0, y0_guess=None, inputs=None):
             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)
+                out = algebraic(time, y0, inputs)
                 pybamm.logger.debug(
                     "Evaluating algebraic equations at t={}, L2-norm is {}".format(
                         time * model.timescale, np.linalg.norm(out)
@@ -438,14 +419,14 @@ def root_fun(y0_alg):
                 return out
 
             if jac:
-                if issparse(jac(0, y0_guess)):
+                if issparse(jac(0, y0_guess, inputs)):
 
                     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()
+                        return jac(0, y0, inputs)[:, len_rhs:].toarray()
 
                 else:
 
@@ -454,7 +435,7 @@ 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:]
+                        return jac(0, y0, inputs)[:, len_rhs:]
 
             else:
                 jac_fn = None
@@ -557,8 +538,6 @@ def solve(self, model, t_eval=None, external_variables=None, inputs=None):
         # Non-dimensionalise time
         t_eval_dimensionless = t_eval / model.timescale_eval
         # Solve
-        # Set inputs and external
-        self.set_inputs(model, ext_and_inputs)
 
         # Calculate discontinuities
         discontinuities = [
@@ -749,7 +728,6 @@ def step(
         # Step
         t_eval = np.linspace(t, t + dt_dimensionless, npts)
         # Set inputs and external
-        self.set_inputs(model, ext_and_inputs)
 
         pybamm.logger.info("Calling solver")
         timer.reset()
@@ -827,43 +805,36 @@ def __init__(self, function, name, model):
         self._function = function
         if isinstance(function, casadi.Function):
             self.form = "casadi"
-            self.inputs = casadi.DM()
+            # self.inputs = casadi.DM()
         else:
             self.form = "python"
-            self.inputs = {}
+            # self.inputs = {}
         self.name = name
         self.model = model
-
-    def set_inputs(self, inputs):
-        "Set inputs"
-        if self.form == "python":
-            self.inputs = inputs
-        elif self.form == "casadi":
-            self.inputs = casadi.vertcat(*[x for x in inputs.values()])
         self.timescale = self.model.timescale_eval
 
-    def __call__(self, t, y):
+    def __call__(self, t, y, inputs):
         y = y[:, np.newaxis]
         if self.name in ["RHS", "algebraic", "residuals", "event"]:
             pybamm.logger.debug(
                 "Evaluating {} for {} at t={}".format(
                     self.name, self.model.name, t * self.timescale
                 )
             )
-            return self.function(t, y).flatten()
+            return self.function(t, y, inputs).flatten()
         else:
-            return self.function(t, y)
+            return self.function(t, y, inputs)
 
-    def function(self, t, y):
+    def function(self, t, y, inputs):
         if self.form == "casadi":
-            states_eval = self._function(t, y, self.inputs)
+            states_eval = self._function(t, y, inputs)
             if self.name in ["RHS", "algebraic", "residuals", "event"]:
                 return states_eval.full()
             else:
                 # keep jacobians sparse
                 return states_eval
         else:
-            return self._function(t, y, self.inputs, known_evals={})[0]
+            return self._function(t, y, inputs, known_evals={})[0]
 
 
 class Residuals(SolverCallable):
@@ -873,6 +844,6 @@ def __init__(self, function, name, model):
         super().__init__(function, name, model)
         self.mass_matrix = model.mass_matrix.entries
 
-    def __call__(self, t, y, ydot):
-        states_eval = super().__call__(t, y)
+    def __call__(self, t, y, ydot, inputs):
+        states_eval = super().__call__(t, y, inputs)
         return states_eval - self.mass_matrix @ ydot

pybamm/solvers/casadi_solver.py

@@ -88,6 +88,9 @@ def _integrate(self, model, t_eval, inputs=None):
             Any external variables or input parameters to pass to the model when solving
         """
         inputs = inputs or {}
+        # convert inputs to casadi format
+        inputs = casadi.vertcat(*[x for x in inputs.values()])
+
         if len(model.rhs) == 0:
             # casadi solver won't allow solving algebraic model so we have to raise an
             # error here
@@ -111,7 +114,10 @@ def _integrate(self, model, t_eval, inputs=None):
             t = t_eval[0]
             init_event_signs = np.sign(
                 np.concatenate(
-                    [event(t, model.y0) for event in model.terminate_events_eval]
+                    [
+                        event(t, model.y0, inputs)
+                        for event in model.terminate_events_eval
+                    ]
                 )
             )
             pybamm.logger.info("Start solving {} with {}".format(model.name, self.name))
@@ -152,7 +158,7 @@ def _integrate(self, model, t_eval, inputs=None):
                 new_event_signs = np.sign(
                     np.concatenate(
                         [
-                            event(t, current_step_sol.y[:, -1])
+                            event(t, current_step_sol.y[:, -1], inputs)
                             for event in model.terminate_events_eval
                         ]
                     )
@@ -181,7 +187,6 @@ def get_integrator(self, model, t_eval, inputs):
             y0 = model.y0
             rhs = model.casadi_rhs
             algebraic = model.casadi_algebraic
-            p_stacked = casadi.vertcat(*[x for x in inputs.values()])
 
             options = {
                 "grid": t_eval,
@@ -193,7 +198,7 @@ def get_integrator(self, model, t_eval, inputs):
 
             # set up and solve
             t = casadi.MX.sym("t")
-            p = casadi.MX.sym("p", p_stacked.shape[0])
+            p = casadi.MX.sym("p", inputs.shape[0])
             y_diff = casadi.MX.sym("y_diff", rhs(t_eval[0], y0, p).shape[0])
             problem = {"t": t, "x": y_diff, "p": p}
             if algebraic(t_eval[0], y0, p).is_empty():
@@ -223,12 +228,11 @@ def get_integrator(self, model, t_eval, inputs):
         )
 
     def _run_integrator(self, integrator, model, y0, inputs, t_eval):
-        rhs_size = model.rhs_eval(t_eval[0], y0).shape[0]
+        rhs_size = model.rhs_eval(t_eval[0], y0, inputs).shape[0]
         y0_diff, y0_alg = np.split(y0, [rhs_size])
         try:
             # Try solving
-            p_stacked = casadi.vertcat(*[x for x in inputs.values()])
-            sol = integrator(x0=y0_diff, z0=y0_alg, p=p_stacked, **self.extra_options)
+            sol = integrator(x0=y0_diff, z0=y0_alg, p=inputs, **self.extra_options)
             y_values = np.concatenate([sol["xf"].full(), sol["zf"].full()])
             return pybamm.Solution(t_eval, y_values)
         except RuntimeError as e:

pybamm/solvers/idaklu_solver.py

@@ -1,6 +1,7 @@
 #
 # Solver class using sundials with the KLU sparse linear solver
 #
+import casadi
 import pybamm
 import numpy as np
 import scipy.sparse as sparse
@@ -146,6 +147,8 @@ def _integrate(self, model, t_eval, inputs=None):
         t_eval : numeric type
             The times at which to compute the solution
         """
+        if model.rhs_eval.form == "casadi":
+            inputs = casadi.vertcat(*[x for x in inputs.values()])
 
         if model.jacobian_eval is None:
             raise pybamm.SolverError("KLU requires the Jacobian to be provided")
@@ -161,17 +164,17 @@ def _integrate(self, model, t_eval, inputs=None):
         mass_matrix = model.mass_matrix.entries
 
         if model.jacobian_eval:
-            jac_y0_t0 = model.jacobian_eval(t_eval[0], y0)
+            jac_y0_t0 = model.jacobian_eval(t_eval[0], y0, inputs)
             if sparse.issparse(jac_y0_t0):
 
                 def jacfn(t, y, cj):
-                    j = model.jacobian_eval(t, y) - cj * mass_matrix
+                    j = model.jacobian_eval(t, y, inputs) - cj * mass_matrix
                     return j
 
             else:
 
                 def jacfn(t, y, cj):
-                    jac_eval = model.jacobian_eval(t, y) - cj * mass_matrix
+                    jac_eval = model.jacobian_eval(t, y, inputs) - cj * mass_matrix
                     return sparse.csr_matrix(jac_eval)
 
         class SundialsJacobian:
@@ -207,12 +210,14 @@ def get_jac_col_ptrs(self):
 
         def rootfn(t, y):
             return_root = np.ones((num_of_events,))
-            return_root[:] = [event(t, y) for event in model.terminate_events_eval]
+            return_root[:] = [
+                event(t, y, inputs) for event in model.terminate_events_eval
+            ]
 
             return return_root
 
         # get ids of rhs and algebraic variables
-        rhs_ids = np.ones(model.rhs_eval(0, y0).shape)
+        rhs_ids = np.ones(model.rhs_eval(0, y0, inputs).shape)
         alg_ids = np.zeros(len(y0) - len(rhs_ids))
         ids = np.concatenate((rhs_ids, alg_ids))
 
@@ -221,7 +226,7 @@ def rootfn(t, y):
             t_eval,
             y0,
             ydot0,
-            model.residuals_eval,
+            lambda t, y, ydot: model.residuals_eval(t, y, ydot, inputs),
             jac_class.jac_res,
             jac_class.get_jac_data,
             jac_class.get_jac_row_vals,

pybamm/solvers/scikits_dae_solver.py

@@ -1,6 +1,7 @@
 #
 # Solver class using Scipy's adaptive time stepper
 #
+import casadi
 import pybamm
 
 import numpy as np
@@ -68,17 +69,20 @@ def _integrate(self, model, t_eval, inputs=None):
             Any input parameters to pass to the model when solving
 
         """
+        if model.convert_to_format == "casadi":
+            inputs = casadi.vertcat(*[x for x in inputs.values()])
+
         residuals = model.residuals_eval
         y0 = model.y0
         events = model.terminate_events_eval
         jacobian = model.jacobian_eval
         mass_matrix = model.mass_matrix.entries
 
         def eqsres(t, y, ydot, return_residuals):
-            return_residuals[:] = residuals(t, y, ydot)
+            return_residuals[:] = residuals(t, y, ydot, inputs)
 
         def rootfn(t, y, ydot, return_root):
-            return_root[:] = [event(t, y) for event in events]
+            return_root[:] = [event(t, y, inputs) for event in events]
 
         extra_options = {
             "old_api": False,
@@ -88,17 +92,17 @@ def rootfn(t, y, ydot, return_root):
         }
 
         if jacobian:
-            jac_y0_t0 = jacobian(t_eval[0], y0)
+            jac_y0_t0 = jacobian(t_eval[0], y0, inputs)
             if sparse.issparse(jac_y0_t0):
 
                 def jacfn(t, y, ydot, residuals, cj, J):
-                    jac_eval = jacobian(t, y) - cj * mass_matrix
+                    jac_eval = jacobian(t, y, inputs) - cj * mass_matrix
                     J[:][:] = jac_eval.toarray()
 
             else:
 
                 def jacfn(t, y, ydot, residuals, cj, J):
-                    jac_eval = jacobian(t, y) - cj * mass_matrix
+                    jac_eval = jacobian(t, y, inputs) - cj * mass_matrix
                     J[:][:] = jac_eval
 
             extra_options.update({"jacfn": jacfn})