#1477 got some more casadi tests working

pybamm-team · Jul 16, 2021 · 72560c5 · 72560c5
1 parent f5699c4
commit 72560c5
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/pybamm/solvers/base_solver.py b/pybamm/solvers/base_solver.py
@@ -11,6 +11,7 @@
 from scipy.sparse import block_diag
 import multiprocessing as mp
 import warnings
+import numbers
 
 
 class BaseSolver(object):
@@ -231,8 +232,15 @@ def set_up(self, model, inputs=None, t_eval=None,
         # (FYI: this is used in the Solution class)
         model.calculate_sensitivities = calculate_sensitivites
         if calculate_sensitivites_explicit:
-            model.len_rhs_sens = model.len_rhs * len(calculate_sensitivites)
-            model.len_alg_sens = model.len_alg * len(calculate_sensitivites)
+            num_parameters = 0
+            for name in calculate_sensitivites:
+                # if not a number, assume its a vector
+                if isinstance(inputs[name], numbers.Number):
+                    num_parameters += 1
+                else:
+                    num_parameters += len(inputs[name])
+            model.len_rhs_sens = model.len_rhs * num_parameters
+            model.len_alg_sens = model.len_alg * num_parameters
 
         if model.convert_to_format != "casadi":
             # Create Jacobian from concatenated rhs and algebraic
@@ -615,7 +623,7 @@ def jacp(*args, **kwargs):
         # if we have changed the equations to include the explicit sensitivity
         # equations, then we also need to update the mass matrix
         if calculate_sensitivites_explicit:
-            n_inputs = len(calculate_sensitivites)
+            n_inputs = model.len_rhs_sens // model.len_rhs
             model.mass_matrix_inv = pybamm.Matrix(
                 block_diag(
                     [model.mass_matrix_inv.entries] * (n_inputs + 1), format="csr"

diff --git a/pybamm/solvers/processed_variable.py b/pybamm/solvers/processed_variable.py
@@ -49,8 +49,6 @@ def __init__(self, base_variables, base_variables_casadi, solution, warn=True):
 
         self.symbolic_inputs = solution.has_symbolic_inputs
 
-        self.u_sol = solution.y
-
         # Sensitivity starts off uninitialized, only set when called
         self._sensitivities = None
         self.solution_sensitivities = solution.sensitivities
@@ -518,7 +516,7 @@ def initialise_sensitivity_explicit_forward(self):
 
         # Set up symbolic variables
         t_casadi = casadi.MX.sym("t")
-        y_casadi = casadi.MX.sym("y", self.u_sol.shape[0])
+        y_casadi = casadi.MX.sym("y", self.all_ys[0].shape[0])
         p_casadi = {
             name: casadi.MX.sym(name, value.shape[0])
             for name, value in self.all_inputs[0].items()
@@ -539,10 +537,9 @@ def initialise_sensitivity_explicit_forward(self):
         )
         for idx in range(len(self.all_ts[0])):
             t = self.all_ts[0][idx]
-            u = self.u_sol[:, idx]
-            inp = inputs_stacked[:, idx]
-            next_dvar_dy_eval = dvar_dy_func(t, u, inp)
-            next_dvar_dp_eval = dvar_dp_func(t, u, inp)
+            u = self.all_ys[0][:, idx]
+            next_dvar_dy_eval = dvar_dy_func(t, u, inputs_stacked)
+            next_dvar_dp_eval = dvar_dp_func(t, u, inputs_stacked)
             if idx == 0:
                 dvar_dy_eval = next_dvar_dy_eval
                 dvar_dp_eval = next_dvar_dp_eval

diff --git a/tests/unit/test_solvers/test_casadi_solver.py b/tests/unit/test_solvers/test_casadi_solver.py
@@ -823,7 +823,6 @@ def test_solve_sensitivity_scalar_var_vector_input(self):
     print("Add -v for more debug output")
     import sys
 
-    pybamm.set_logging_level('DEBUG')
     if "-v" in sys.argv:
         debug = True
     pybamm.settings.debug_mode = True