#1100 reformat Solution syntax

Author: valentinsulzer

pybamm/solvers/algebraic_solver.py

@@ -59,9 +59,9 @@ def _integrate(self, model, t_eval, inputs=None):
         inputs : dict, optional
             Any input parameters to pass to the model when solving
         """
-        inputs = inputs or {}
+        inputs_dict = inputs or {}
         if model.convert_to_format == "casadi":
-            inputs = casadi.vertcat(*[x for x in inputs.values()])
+            inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
 
         y0 = model.y0
         if isinstance(y0, casadi.DM):
@@ -210,4 +210,7 @@ def jac_norm(y):
         y_diff = np.r_[[y0_diff] * len(t_eval)].T
         y_sol = np.r_[y_diff, y_alg]
         # Return solution object (no events, so pass None to t_event, y_event)
-        return pybamm.Solution(t_eval, y_sol, termination="success")
+        return pybamm.Solution(
+            t_eval, y_sol, termination="success", model=model, inputs=inputs_dict
+        )
+

pybamm/solvers/base_solver.py

@@ -34,11 +34,11 @@ class BaseSolver(object):
     sensitivity : str, optional
         Whether (and how) to calculate sensitivities when solving. Options are:
 
-        - "explicit forward": explicitly formulate the sensitivity equations.
-        The formulation is as per "Park, S., Kato, D., Gima, Z.,
-        Klein, R., & Moura, S. (2018). Optimal experimental design for parameterization
-        of an electrochemical lithium-ion battery model. Journal of The Electrochemical
-        Society, 165(7), A1309.". See #1100 for details
+        - "explicit forward": explicitly formulate the sensitivity equations. \
+        The formulation is as per "Park, S., Kato, D., Gima, Z., \
+        Klein, R., & Moura, S. (2018). Optimal experimental design for parameterization\
+        of an electrochemical lithium-ion battery model. Journal of The Electrochemical\
+        Society, 165(7), A1309.". See #1100 for details \
         - see specific solvers for other options
     """
 
@@ -891,10 +891,6 @@ def step(
         solution.set_up_time = set_up_time
         solution.solve_time = timer.time()
 
-        # Add model and inputs to solution
-        solution.model = model
-        solution.inputs = ext_and_inputs
-
         # Identify the event that caused termination
         termination = self.get_termination_reason(solution, model.events)
 

pybamm/solvers/casadi_algebraic_solver.py

@@ -25,7 +25,7 @@ class CasadiAlgebraicSolver(pybamm.BaseSolver):
         Whether (and how) to calculate sensitivities when solving. Options are:
 
         - None: no sensitivities
-        - "explicit forward": explicitly formulate the sensitivity equations.
+        - "explicit forward": explicitly formulate the sensitivity equations. \
         See :class:`pybamm.BaseSolver`
         - "casadi": use casadi to differentiate through the rootfinding operator
 
@@ -66,7 +66,7 @@ def _integrate(self, model, t_eval, inputs=None):
         # Record whether there are any symbolic inputs
         inputs_dict = inputs or {}
         # Create casadi objects for the root-finder
-        inputs = casadi.vertcat(*[v for v in inputs.values()])
+        inputs = casadi.vertcat(*[v for v in inputs_dict.values()])
 
         # Create symbolic inputs
         symbolic_inputs = casadi.MX.sym("inputs", inputs.shape[0])

pybamm/solvers/casadi_solver.py

@@ -58,11 +58,11 @@ class CasadiSolver(pybamm.BaseSolver):
         Any options to pass to the CasADi integrator when calling the integrator.
         Please consult `CasADi documentation <https://tinyurl.com/y5rk76os>`_ for
         details.
-    sensitivity : bool, optional
+    sensitivity : str, optional
         Whether (and how) to calculate sensitivities when solving. Options are:
 
         - None: no sensitivities
-        - "explicit forward": explicitly formulate the sensitivity equations.
+        - "explicit forward": explicitly formulate the sensitivity equations. \
         See :class:`pybamm.BaseSolver`
         - "casadi": use casadi to differentiate through the integrator
     """

pybamm/solvers/dummy_solver.py

@@ -33,4 +33,6 @@ def _integrate(self, model, t_eval, inputs=None):
 
         """
         y_sol = np.zeros((1, t_eval.size))
-        return pybamm.Solution(t_eval, y_sol, termination="final time")
+        return pybamm.Solution(
+            t_eval, y_sol, termination="final time", model=model, inputs=inputs
+        )

pybamm/solvers/idaklu_solver.py

@@ -154,6 +154,7 @@ def _integrate(self, model, t_eval, inputs=None):
         t_eval : numeric type
             The times at which to compute the solution
         """
+        inputs_dict = inputs
         if model.rhs_eval.form == "casadi":
             # stack inputs
             inputs = casadi.vertcat(*[x for x in inputs.values()])
@@ -272,6 +273,8 @@ def rootfn(t, y):
                 t[-1],
                 np.transpose(y_out[-1])[:, np.newaxis],
                 termination,
+                model=model,
+                inputs=inputs_dict,
             )
         else:
             raise pybamm.SolverError(sol.message)

pybamm/solvers/jax_solver.py

@@ -45,16 +45,17 @@ class JaxSolver(pybamm.BaseSolver):
         for details.
     """
 
-    def __init__(self, method='RK45', root_method=None,
-                 rtol=1e-6, atol=1e-6, extra_options=None):
+    def __init__(
+        self, method="RK45", root_method=None, rtol=1e-6, atol=1e-6, extra_options=None
+    ):
         # note: bdf solver itself calculates consistent initial conditions so can set
         # root_method to none, allow user to override this behavior
         super().__init__(method, rtol, atol, root_method=root_method)
-        method_options = ['RK45', 'BDF']
+        method_options = ["RK45", "BDF"]
         if method not in method_options:
-            raise ValueError('method must be one of {}'.format(method_options))
+            raise ValueError("method must be one of {}".format(method_options))
         self.ode_solver = False
-        if method == 'RK45':
+        if method == "RK45":
             self.ode_solver = True
         self.extra_options = extra_options or {}
         self.name = "JAX solver ({})".format(method)
@@ -80,8 +81,9 @@ def get_solve(self, model, t_eval):
         """
         if model not in self._cached_solves:
             if model not in self.models_set_up:
-                raise RuntimeError("Model is not set up for solving, run"
-                                   "`solver.solve(model)` first")
+                raise RuntimeError(
+                    "Model is not set up for solving, run" "`solver.solve(model)` first"
+                )
 
             self._cached_solves[model] = self.create_solve(model, t_eval)
 
@@ -106,32 +108,35 @@ def create_solve(self, model, t_eval):
 
         """
         if model.convert_to_format != "jax":
-            raise RuntimeError("Model must be converted to JAX to use this solver"
-                               " (i.e. `model.convert_to_format = 'jax')")
+            raise RuntimeError(
+                "Model must be converted to JAX to use this solver"
+                " (i.e. `model.convert_to_format = 'jax')"
+            )
 
         if model.terminate_events_eval:
-            raise RuntimeError("Terminate events not supported for this solver."
-                               " Model has the following events:"
-                               " {}.\nYou can remove events using `model.events = []`."
-                               " It might be useful to first solve the model using a"
-                               " different solver to obtain the time of the event, then"
-                               " re-solve using no events and a fixed"
-                               " end-time".format(model.events))
+            raise RuntimeError(
+                "Terminate events not supported for this solver."
+                " Model has the following events:"
+                " {}.\nYou can remove events using `model.events = []`."
+                " It might be useful to first solve the model using a"
+                " different solver to obtain the time of the event, then"
+                " re-solve using no events and a fixed"
+                " end-time".format(model.events)
+            )
 
         # Initial conditions, make sure they are an 0D array
         y0 = jnp.array(model.y0).reshape(-1)
         mass = None
-        if self.method == 'BDF':
+        if self.method == "BDF":
             mass = model.mass_matrix.entries.toarray()
 
         def rhs_ode(y, t, inputs):
-            return model.rhs_eval(t, y, inputs),
+            return (model.rhs_eval(t, y, inputs),)
 
         def rhs_dae(y, t, inputs):
-            return jnp.concatenate([
-                model.rhs_eval(t, y, inputs),
-                model.algebraic_eval(t, y, inputs),
-            ])
+            return jnp.concatenate(
+                [model.rhs_eval(t, y, inputs), model.algebraic_eval(t, y, inputs)]
+            )
 
         def solve_model_rk45(inputs):
             y = odeint(
@@ -158,7 +163,7 @@ def solve_model_bdf(inputs):
             )
             return jnp.transpose(y)
 
-        if self.method == 'RK45':
+        if self.method == "RK45":
             return jax.jit(solve_model_rk45)
         else:
             return jax.jit(solve_model_bdf)
@@ -194,5 +199,7 @@ def _integrate(self, model, t_eval, inputs=None):
         termination = "final time"
         t_event = None
         y_event = onp.array(None)
-        return pybamm.Solution(t_eval, y,
-                               t_event, y_event, termination)
+        return pybamm.Solution(
+            t_eval, y, t_event, y_event, termination, model=model, inputs=inputs
+        )
+

pybamm/solvers/scikits_dae_solver.py

@@ -81,6 +81,7 @@ def _integrate(self, model, t_eval, inputs=None):
             Any input parameters to pass to the model when solving
 
         """
+        inputs_dict = inputs
         if model.convert_to_format == "casadi":
             inputs = casadi.vertcat(*[x for x in inputs.values()])
 
@@ -150,6 +151,8 @@ def jacfn(t, y, ydot, residuals, cj, J):
                 t_root,
                 np.transpose(sol.roots.y),
                 termination,
+                model=model,
+                inputs=inputs_dict,
             )
         else:
             raise pybamm.SolverError(sol.message)

pybamm/solvers/scikits_ode_solver.py

@@ -80,6 +80,7 @@ def _integrate(self, model, t_eval, inputs=None):
             Any input parameters to pass to the model when solving
 
         """
+        inputs_dict = inputs
         if model.rhs_eval.form == "casadi":
             inputs = casadi.vertcat(*[x for x in inputs.values()])
 
@@ -167,6 +168,8 @@ def jac_times_setupfn(t, y, fy, userdata):
                 t_root,
                 np.transpose(sol.roots.y),
                 termination,
+                model=model,
+                inputs=inputs_dict,
             )
         else:
             raise pybamm.SolverError(sol.message)

pybamm/solvers/scipy_solver.py

@@ -23,8 +23,11 @@ class ScipySolver(pybamm.BaseSolver):
         Any options to pass to the solver.
         Please consult `SciPy documentation <https://tinyurl.com/yafgqg9y>`_ for
         details.
-    sensitivity : bool, optional
-        Whether to explicitly formulate and solve the forward sensitivity equations.
+    sensitivity : str, optional
+        Whether (and how) to calculate sensitivities when solving. Options are:
+
+        - None: no sensitivities
+        - "explicit forward": explicitly formulate the sensitivity equations. \
         See :class:`pybamm.BaseSolver`
     """
 

pybamm/solvers/solution.py

@@ -57,9 +57,9 @@ def __init__(
         if isinstance(y, casadi.DM):
             y = y.full()
 
-        # if model or inputs are None, initialize empty, to be populated later
-        self.inputs = inputs or pybamm.FuzzyDict()
-        self._model = model or pybamm.BaseModel()
+        # if inputs are None, initialize empty, to be populated later
+        inputs = inputs or pybamm.FuzzyDict()
+        self.set_inputs(inputs)
 
         # If the model has been provided, split up y into solution and sensitivity
         # Don't do this if the sensitivity equations have not been computed (i.e. if
@@ -70,6 +70,7 @@ def __init__(
             model is None
             or isinstance(y, casadi.Function)
             or model.len_rhs_and_alg == y.shape[0]
+            or model.len_rhs_and_alg == 0  # for the dummy solver
         ):
             self._y = y
             self.sensitivity = {}
@@ -129,6 +130,8 @@ def __init__(
                 start = end
             self.sensitivity = sensitivity
 
+        model = model or pybamm.BaseModel()
+        self.set_model(model)
         self._t_event = t_event
         self._y_event = y_event
         self._termination = termination
@@ -163,39 +166,28 @@ def model(self):
         "Model used for solution"
         return self._model
 
-    @model.setter
-    def model(self, value):
+    def set_model(self, value):
         "Updates the model"
-        assert isinstance(value, pybamm.BaseModel)
         self._model = value
 
     @property
     def inputs(self):
         "Values of the inputs"
         return self._inputs
 
-    @inputs.setter
-    def inputs(self, inputs):
+    def set_inputs(self, inputs):
         "Updates the input values"
-        # If there are symbolic inputs, just store them as given
-        if any(isinstance(v, casadi.MX) for v in inputs.values()):
-            self.has_symbolic_inputs = True
-            self._inputs = inputs
-        # Otherwise, make them the same size as the time vector
-        else:
-            self.has_symbolic_inputs = False
-            self._inputs = {}
-            for name, inp in inputs.items():
-                # Convert number to vector of the right shape
-                if isinstance(inp, numbers.Number):
-                    inp = inp * np.ones((1, len(self.t)))
-                # Tile a vector
-                else:
-                    if inp.ndim == 1:
-                        inp = np.tile(inp, (len(self.t), 1)).T
-                    else:
-                        inp = np.tile(inp, len(self.t))
-                self._inputs[name] = inp
+        self._inputs = {}
+        for name, inp in inputs.items():
+            # Convert number to vector of the right shape
+            if isinstance(inp, numbers.Number):
+                inp = inp * np.ones((1, len(self.t)))
+            # Otherwise, tile a vector
+            elif inp.ndim == 1:
+                inp = np.tile(inp, (len(self.t), 1)).T
+            elif inp.shape[1] != len(self.t):
+                inp = np.tile(inp, len(self.t))
+            self._inputs[name] = inp
 
     @property
     def t_event(self):
@@ -434,6 +426,6 @@ def append(self, solution, start_index=1, create_sub_solutions=False):
                     solution.termination,
                     copy_this=solution,
                     model=self.model,
-                    inputs=copy.copy(self.inputs),
+                    inputs=copy.copy(solution.inputs),
                 )
             )

tests/unit/test_solvers/test_algebraic_solver.py

@@ -44,6 +44,7 @@ class Model:
             timescale_eval = 1
             jac_algebraic_eval = None
             convert_to_format = "python"
+            len_rhs_and_alg = 1
 
             def algebraic_eval(self, t, y, inputs):
                 return y + 2
@@ -66,6 +67,7 @@ class Model:
             timescale_eval = 1
             jac_algebraic_eval = None
             convert_to_format = "casadi"
+            len_rhs_and_alg = 1
 
             def algebraic_eval(self, t, y, inputs):
                 # algebraic equation has no real root
@@ -95,6 +97,7 @@ class Model:
             rhs = {}
             timescale_eval = 1
             convert_to_format = "python"
+            len_rhs_and_alg = 2
 
             def algebraic_eval(self, t, y, inputs):
                 return A @ y - b

tests/unit/test_solvers/test_base_solver.py

@@ -119,6 +119,7 @@ def __init__(self):
                 )
                 self.convert_to_format = "casadi"
                 self.bounds = (np.array([-np.inf]), np.array([np.inf]))
+                self.len_rhs_and_alg = 1
 
             def rhs_eval(self, t, y, inputs):
                 return np.array([])
@@ -154,6 +155,8 @@ def __init__(self):
                 )
                 self.convert_to_format = "casadi"
                 self.bounds = (-np.inf * np.ones(4), np.inf * np.ones(4))
+                self.len_rhs = 1
+                self.len_rhs_and_alg = 4
 
             def rhs_eval(self, t, y, inputs):
                 return y[0:1]

tests/unit/test_solvers/test_scikits_solvers.py

@@ -97,6 +97,7 @@ class Model:
             terminate_events_eval = []
             timescale_eval = 1
             convert_to_format = "python"
+            len_rhs_and_alg = 2
 
             def residuals_eval(self, t, y, ydot, inputs):
                 return np.array(