#1100 starting to get SDAEs working with casadi

Author: valentinsulzer

pybamm/__init__.py

@@ -108,6 +108,7 @@ def version(formatted=False):
     to_python,
     EvaluatorPython,
 )
+
 if system() != "Windows":
     from .expression_tree.operations.evaluate import EvaluatorJax
 

pybamm/solvers/base_solver.py

@@ -287,19 +287,44 @@ def report(string):
                 func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi)
                 # Add sensitivity vectors to the rhs and algebraic equations
                 if self.solve_sensitivity_equations is True:
-                    if name == "rhs":
+                    if name == "rhs" and model.len_rhs > 0:
                         report("Creating sensitivity equations for rhs using CasADi")
                         df_dx = casadi.jacobian(func, y_diff)
                         df_dp = casadi.jacobian(func, p_casadi_stacked)
                         S_x_mat = S_x.reshape(
-                            (model.len_rhs_and_alg, p_casadi_stacked.shape[0])
+                            (model.len_rhs, p_casadi_stacked.shape[0])
                         )
                         if model.len_alg == 0:
                             S_rhs = (df_dx @ S_x_mat + df_dp).reshape((-1, 1))
                         else:
                             df_dz = casadi.jacobian(func, y_alg)
-                            S_rhs = df_dx @ S_x_mat + df_dz @ S_z + df_dp
+                            S_z_mat = S_z.reshape(
+                                (model.len_rhs, p_casadi_stacked.shape[0])
+                            )
+                            S_rhs = (df_dx @ S_x_mat + df_dz @ S_z_mat + df_dp).reshape(
+                                (-1, 1)
+                            )
                         func = casadi.vertcat(func, S_rhs)
+                    if name == "algebraic" and model.len_alg > 0:
+                        report(
+                            "Creating sensitivity equations for algebraic using CasADi"
+                        )
+                        dg_dz = casadi.jacobian(func, y_alg)
+                        dg_dp = casadi.jacobian(func, p_casadi_stacked)
+                        S_z_mat = S_z.reshape(
+                            (model.len_rhs, p_casadi_stacked.shape[0])
+                        )
+                        if model.len_rhs == 0:
+                            S_alg = (dg_dz @ S_z_mat + dg_dp).reshape((-1, 1))
+                        else:
+                            dg_dx = casadi.jacobian(func, y_diff)
+                            S_x_mat = S_x.reshape(
+                                (model.len_rhs, p_casadi_stacked.shape[0])
+                            )
+                            S_alg = (dg_dx @ S_x_mat + dg_dz @ S_z_mat + dg_dp).reshape(
+                                (-1, 1)
+                            )
+                        func = casadi.vertcat(func, S_alg)
                     elif name == "initial_conditions":
                         if model.len_rhs == 0 or model.len_alg == 0:
                             S_0 = casadi.jacobian(func, p_casadi_stacked).reshape(
@@ -309,8 +334,12 @@ def report(string):
                         else:
                             x0 = func[: model.len_rhs]
                             z0 = func[model.len_rhs :]
-                            Sx_0 = casadi.jacobian(x0, p_casadi_stacked)
-                            Sz_0 = casadi.jacobian(z0, p_casadi_stacked)
+                            Sx_0 = casadi.jacobian(x0, p_casadi_stacked).reshape(
+                                (-1, 1)
+                            )
+                            Sz_0 = casadi.jacobian(z0, p_casadi_stacked).reshape(
+                                (-1, 1)
+                            )
                             func = casadi.vertcat(x0, Sx_0, z0, Sz_0)
                 if use_jacobian:
                     report(f"Calculating jacobian for {name} using CasADi")

pybamm/solvers/casadi_algebraic_solver.py

@@ -75,7 +75,11 @@ def _integrate(self, model, t_eval, inputs=None):
             y0_diff = casadi.DM()
             y0_alg = y0
         else:
-            len_rhs = model.concatenated_rhs.size
+            # Check y0 to see if it includes sensitivities
+            if model.len_rhs_and_alg == y0.shape[0]:
+                len_rhs = model.len_rhs
+            else:
+                len_rhs = model.len_rhs * (inputs.shape[0] + 1)
             y0_diff = y0[:len_rhs]
             y0_alg = y0[len_rhs:]
 

pybamm/solvers/solution.py

@@ -69,6 +69,8 @@ def __init__(
             self.sensitivity = {}
         else:
             n_states = model.len_rhs_and_alg
+            n_rhs = model.len_rhs
+            n_alg = model.len_alg
             n_t = len(t)
             n_p = np.vstack(list(inputs.values())).size
             # Get the point where the algebraic equations start
@@ -97,26 +99,19 @@ def __init__(
             #   tn_x1_p0, tn_x1_p1, ..., tn_x1_pn
             #   ...
             #   tn_xn_p0, tn_xn_p1, ..., tn_xn_pn
-            # 1. Extract the relevant parts of y
-            # This makes a (n_states * n_p, n_t) matrix
-            full_sens_matrix = np.vstack(
-                [
-                    y[model.len_rhs : len_rhs_and_sens, :],
-                    y[len_rhs_and_sens + model.len_alg :, :],
-                ]
+            # 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)
+            # 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)
+            # Transpose and reshape into a (n_states * n_t, n_p) matrix
+            full_sens_matrix = full_sens_matrix.transpose(2, 1, 0).reshape(
+                n_t * n_states, n_p
             )
-            # 2. Transpose into a (n_t, n_states * n_p) matrix
-            full_sens_matrix = full_sens_matrix.T
-            # 3. Reshape into a (n_t, n_p, n_states) matrix,
-            # then tranpose n_p and n_states to get (n_t, n_states, n_p) matrix
-            full_sens_matrix = full_sens_matrix.reshape(n_t, n_p, n_states).transpose(
-                0, 2, 1
-            )
-            # 3. Stack time and space to get a (n_t * n_states, n_p) matrix
-            full_sens_matrix = full_sens_matrix.reshape(n_t * n_states, n_p)
 
             # Save the full sensitivity matrix
-
             sensitivity = {"all": full_sens_matrix}
             # also save the sensitivity wrt each parameter (read the columns of the
             # sensitivity matrix)