#1100 fixed some solver tests

valentinsulzer · valentinsulzer · commit e70e05750de2 · 2020-12-28T17:21:53.000+01:00
diff --git a/pybamm/solvers/casadi_algebraic_solver.py b/pybamm/solvers/casadi_algebraic_solver.py
@@ -73,7 +73,7 @@ def _integrate(self, model, t_eval, inputs=None):
         y0 = model.y0
 
         # If y0 already satisfies the tolerance for all t then keep it
-        if all(
+        if self.sensitivity != "casadi" and all(
             np.all(abs(model.casadi_algebraic(t, y0, inputs).full()) < self.tol)
             for t in t_eval
         ):
diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py
@@ -54,48 +54,21 @@ def __init__(self, base_variable, base_variable_casadi, solution, warn=True):
         self.base_variable = base_variable
         self.base_variable_casadi = base_variable_casadi
         self.t_sol = solution.t
-        self.u_sol = solution.y
         self.mesh = base_variable.mesh
-        self.inputs = solution.inputs
         self.domain = base_variable.domain
         self.auxiliary_domains = base_variable.auxiliary_domains
         self.warn = warn
 
+        self.inputs = solution.inputs
+        self.symbolic_inputs = solution._symbolic_inputs
+
+        self.u_sol = solution.y
+        self.y_sym = solution._y_sym
+
         # Sensitivity starts off uninitialized, only set when called
         self._sensitivity = None
         self.solution_sensitivity = solution.sensitivity
 
-        # Special case: symbolic solution, with casadi
-        if isinstance(solution.y, casadi.Function):
-            # Evaluate solution at specific inputs value
-            inputs_stacked = casadi.vertcat(*solution.inputs.values())
-            self.u_sol = solution.y(inputs_stacked).full()
-            # Convert variable to casadi
-            t_MX = casadi.MX.sym("t")
-            y_MX = casadi.MX.sym("y", self.u_sol.shape[0])
-            # Make all inputs symbolic first for converting to casadi
-            symbolic_inputs_dict = {
-                name: casadi.MX.sym(name, value.shape[0])
-                for name, value in solution.inputs.items()
-            }
-
-            # The symbolic_inputs will be used for sensitivity
-            symbolic_inputs = casadi.vertcat(*symbolic_inputs_dict.values())
-            var_casadi = base_variable.to_casadi(
-                t_MX, y_MX, inputs=symbolic_inputs_dict
-            )
-            self.base_variable_sym = casadi.Function(
-                "variable", [t_MX, y_MX, symbolic_inputs], [var_casadi]
-            )
-            # Store symbolic inputs for sensitivity
-            self.symbolic_inputs = symbolic_inputs
-            self.y_sym = solution.y(symbolic_inputs)
-        else:
-            self.u_sol = solution.y
-            self.base_variable_sym = None
-            self.symbolic_inputs = None
-            self.y_sym = None
-
         # Set timescale
         self.timescale = solution.timescale_eval
         self.t_pts = self.t_sol * self.timescale
@@ -565,17 +538,16 @@ def sensitivity(self):
             return {}
         # Otherwise initialise and return sensitivity
         if self._sensitivity is None:
-            # Check that we can compute sensitivities
-            if self.base_variable_sym is None and self.solution_sensitivity == {}:
+            if self.solution_sensitivity != {}:
+                self.initialise_sensitivity_explicit_forward()
+            elif self.y_sym is not None:
+                self.initialise_sensitivity_casadi()
+            else:
                 raise ValueError(
                     "Cannot compute sensitivities. The 'sensitivity' argument of the "
                     "solver should be changed from 'None' to allow sensitivity "
                     "calculations. Check solver documentation for details."
                 )
-            if self.base_variable_sym is None:
-                self.initialise_sensitivity_explicit_forward()
-            else:
-                self.initialise_sensitivity_casadi()
         return self._sensitivity
 
     def initialise_sensitivity_explicit_forward(self):
@@ -639,7 +611,7 @@ def initialise_0D_symbolic():
             for idx in range(len(self.t_sol)):
                 t = self.t_sol[idx]
                 u = self.y_sym[:, idx]
-                next_entries = self.base_variable_sym(t, u, self.symbolic_inputs)
+                next_entries = self.base_variable_casadi(t, u, self.symbolic_inputs)
                 if idx == 0:
                     entries = next_entries
                 else:
@@ -653,7 +625,7 @@ def initialise_1D_symbolic():
             for idx in range(len(self.t_sol)):
                 t = self.t_sol[idx]
                 u = self.y_sym[:, idx]
-                next_entries = self.base_variable_sym(t, u, self.symbolic_inputs)
+                next_entries = self.base_variable_casadi(t, u, self.symbolic_inputs)
                 if idx == 0:
                     entries = next_entries
                 else:
@@ -662,7 +634,7 @@ def initialise_1D_symbolic():
             return entries
 
         inputs_stacked = casadi.vertcat(*self.inputs.values())
-        self.base_eval = self.base_variable_sym(
+        self.base_eval = self.base_variable_casadi(
             self.t_sol[0], self.u_sol[:, 0], inputs_stacked
         )
         if (
diff --git a/pybamm/solvers/solution.py b/pybamm/solvers/solution.py
@@ -53,12 +53,7 @@ def __init__(
         inputs=None,
     ):
         self.t = t
-        if isinstance(y, casadi.DM):
-            y = y.full()
-
-        # if inputs are None, initialize empty, to be populated later
-        inputs = inputs or pybamm.FuzzyDict()
-        self.set_inputs(inputs)
+        self.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
@@ -107,8 +102,12 @@ def __init__(
             #   tn_xn_p0, tn_xn_p1, ..., tn_xn_pn
             # 1, Extract rhs and alg sensitivities and reshape into 3D matrices
             # with shape (n_p, n_states, n_t)
-            ode_sens = y[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t)
-            alg_sens = y[len_rhs_and_sens + n_alg :, :].reshape(n_p, n_alg, n_t)
+            if isinstance(y, casadi.DM):
+                y_full = y.full()
+            else:
+                y_full = y
+            ode_sens = y_full[n_rhs:len_rhs_and_sens, :].reshape(n_p, n_rhs, n_t)
+            alg_sens = y_full[len_rhs_and_sens + n_alg :, :].reshape(n_p, n_alg, n_t)
             # 2. Concatenate into a single 3D matrix with shape (n_p, n_states, n_t)
             # i.e. along first axis
             full_sens_matrix = np.concatenate([ode_sens, alg_sens], axis=1)
@@ -163,8 +162,16 @@ def y(self):
 
     @y.setter
     def y(self, y):
-        self._y = y
-        self._y_MX = casadi.MX.sym("y", y.shape[0])
+        if isinstance(y, casadi.Function):
+            self._y_fn = None
+            inputs_stacked = casadi.vertcat(*self.inputs.values())
+            self._y = y(inputs_stacked)
+            self._y_sym = y(self._symbolic_inputs)
+        else:
+            self._y = y
+            self._y_fn = None
+            self._y_sym = None
+        self._y_MX = casadi.MX.sym("y", self._y.shape[0])
 
     @property
     def model(self):
@@ -196,8 +203,12 @@ def inputs(self):
         "Values of the inputs"
         return self._inputs
 
-    def set_inputs(self, inputs):
+    @inputs.setter
+    def inputs(self, inputs):
         "Updates the input values"
+        # if inputs are None, initialize empty, to be populated later
+        inputs = inputs or pybamm.FuzzyDict()
+
         # self._inputs = {}
         # for name, inp in inputs.items():
         #     # Convert number to vector of the right shape
@@ -233,13 +244,13 @@ def set_inputs(self, inputs):
                         inp = inp[:, np.newaxis]
                     inp = np.tile(inp, len(self.t))
                 self._inputs[name] = inp
-        self._all_inputs_as_MX_dict = {}
-        for key, value in self._inputs.items():
-            self._all_inputs_as_MX_dict[key] = casadi.MX.sym("input", value.shape[0])
+        self._symbolic_inputs_dict = {
+            name: casadi.MX.sym(name, value.shape[0])
+            for name, value in self.inputs.items()
+        }
 
-        self._all_inputs_as_MX = casadi.vertcat(
-            *[p for p in self._all_inputs_as_MX_dict.values()]
-        )
+        # The symbolic_inputs will be used for sensitivity
+        self._symbolic_inputs = casadi.vertcat(*self._symbolic_inputs_dict.values())
 
     @property
     def t_event(self):
@@ -298,12 +309,12 @@ def update(self, variables):
                     # Convert variable to casadi
                     # Make all inputs symbolic first for converting to casadi
                     var_sym = var_pybamm.to_casadi(
-                        self._t_MX, self._y_MX, inputs=self._all_inputs_as_MX_dict
+                        self._t_MX, self._y_MX, inputs=self._symbolic_inputs_dict
                     )
 
                     var_casadi = casadi.Function(
                         "variable",
-                        [self._t_MX, self._y_MX, self._all_inputs_as_MX],
+                        [self._t_MX, self._y_MX, self._symbolic_inputs],
                         [var_sym],
                     )
                     self.model._variables_casadi[key] = var_casadi
@@ -359,8 +370,8 @@ def clear_casadi_attributes(self):
         "Remove casadi objects for pickling, will be computed again automatically"
         self._t_MX = None
         self._y_MX = None
-        self._all_inputs_as_MX = None
-        self._all_inputs_as_MX_dict = None
+        self._symbolic_inputs = None
+        self._symbolic_inputs_dict = None
 
     def save(self, filename):
         """Save the whole solution using pickle"""
diff --git a/tests/unit/test_solvers/test_algebraic_solver.py b/tests/unit/test_solvers/test_algebraic_solver.py
@@ -38,7 +38,7 @@ def test_wrong_solver(self):
 
     def test_simple_root_find(self):
         # Simple system: a single algebraic equation
-        class Model:
+        class Model(pybamm.BaseModel):
             y0 = np.array([2])
             rhs = {}
             timescale_eval = 1
@@ -61,7 +61,7 @@ def algebraic_eval(self, t, y, inputs):
         self.assertNotEqual(solution.y, -2)
 
     def test_root_find_fail(self):
-        class Model:
+        class Model(pybamm.BaseModel):
             y0 = np.array([2])
             rhs = {}
             timescale_eval = 1
@@ -92,7 +92,7 @@ def test_with_jacobian(self):
         A = np.array([[4, 3], [1, -1]])
         b = np.array([0, 7])
 
-        class Model:
+        class Model(pybamm.BaseModel):
             y0 = np.zeros(2)
             rhs = {}
             timescale_eval = 1
diff --git a/tests/unit/test_solvers/test_casadi_algebraic_solver.py b/tests/unit/test_solvers/test_casadi_algebraic_solver.py
diff --git a/tests/unit/test_solvers/test_processed_variable.py b/tests/unit/test_solvers/test_processed_variable.py
