Merge pull request #2920 from pybamm-team/i2858-init-cond-sundials

I2858-init-cond-sundials

CHANGELOG.md

@@ -39,6 +39,7 @@
 - Parameter sets can now contain the key "chemistry", and will ignore its value (this previously would give errors in some cases) ([#2901](https://github.com/pybamm-team/PyBaMM/pull/2901))
 - Fixed a bug in the discretisation of initial conditions of a scaled variable ([#2856](https://github.com/pybamm-team/PyBaMM/pull/2856))
 - Fixed keyerror on "all" when getting sensitivities from IDAKLU solver([#2883](https://github.com/pybamm-team/PyBaMM/pull/2883))
+- Initial conditions for sensitivity equations calculated correctly ([#2920](https://github.com/pybamm-team/PyBaMM/pull/2920))
 
 ## Breaking changes
 

pybamm/solvers/base_solver.py

@@ -130,13 +130,14 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
         )
 
         # Process initial conditions
-        initial_conditions = process(
+        initial_conditions, _, jacp_ic, _ = process(
             model.concatenated_initial_conditions,
             "initial_conditions",
             vars_for_processing,
             use_jacobian=False,
-        )[0]
+        )
         model.initial_conditions_eval = initial_conditions
+        model.jacp_initial_conditions_eval = jacp_ic
 
         # evaluate initial condition
         y0_total_size = (
@@ -146,9 +147,25 @@ def set_up(self, model, inputs=None, t_eval=None, ics_only=False):
         if model.convert_to_format == "casadi":
             # stack inputs
             inputs_casadi = casadi.vertcat(*[x for x in inputs.values()])
-            model.y0 = initial_conditions(0, y_zero, inputs_casadi)
+            model.y0 = initial_conditions(0.0, y_zero, inputs_casadi)
+            if jacp_ic is None:
+                model.y0S = None
+            else:
+                model.y0S = jacp_ic(0.0, y_zero, inputs_casadi)
         else:
-            model.y0 = initial_conditions(0, y_zero, inputs)
+            model.y0 = initial_conditions(0.0, y_zero, inputs)
+            if jacp_ic is None:
+                model.y0S = None
+            else:
+                # we are calculating the derivative wrt the inputs
+                # so need to make sure we convert int -> float
+                # This is to satisfy JAX jacfwd function which requires
+                # float inputs
+                inputs_float = {
+                    key: float(value) if isinstance(value, int) else value
+                    for key, value in inputs.items()
+                }
+                model.y0S = jacp_ic(0.0, y_zero, inputs_float)
 
         if ics_only:
             pybamm.logger.info("Finish solver set-up")

pybamm/solvers/c_solvers/idaklu/casadi_solver.cpp

@@ -1,6 +1,7 @@
 #include "casadi_solver.hpp"
 #include "casadi_sundials_functions.hpp"
 #include "common.hpp"
+#include <idas/idas.h>
 #include <memory>
 
 
@@ -291,7 +292,27 @@ Solution CasadiSolver::solve(np_array t_np, np_array y0_np, np_array yp0_np,
                              np_array_dense inputs)
 {
   DEBUG("CasadiSolver::solve");
+
   int number_of_timesteps = t_np.request().size;
+  auto t = t_np.unchecked<1>();
+  realtype t0 = RCONST(t(0));
+  auto y0 = y0_np.unchecked<1>();
+  auto yp0 = yp0_np.unchecked<1>();
+
+
+  if (y0.size() != number_of_states + number_of_parameters * number_of_states) {
+    throw std::domain_error(
+      "y0 has wrong size. Expected " + 
+      std::to_string(number_of_states + number_of_parameters * number_of_states) + 
+      " but got " + std::to_string(y0.size()));
+  }
+
+  if (yp0.size() != number_of_states + number_of_parameters * number_of_states) {
+    throw std::domain_error(
+      "yp0 has wrong size. Expected " + 
+      std::to_string(number_of_states + number_of_parameters * number_of_states) + 
+      " but got " + std::to_string(yp0.size()));
+  }
 
   // set inputs
   auto p_inputs = inputs.unchecked<2>();
@@ -300,26 +321,38 @@ Solution CasadiSolver::solve(np_array t_np, np_array y0_np, np_array yp0_np,
     functions->inputs[i] = p_inputs(i, 0);
   }
 
+  // set initial conditions
   realtype *yval = N_VGetArrayPointer(yy);
   realtype *ypval = N_VGetArrayPointer(yp);
   std::vector<realtype *> ySval(number_of_parameters);
-  for (int is = 0 ; is < number_of_parameters; is++) {
-    ySval[is] = N_VGetArrayPointer(yyS[is]);
-    N_VConst(RCONST(0.0), yyS[is]);
-    N_VConst(RCONST(0.0), ypS[is]);
+  std::vector<realtype *> ypSval(number_of_parameters);
+  for (int p = 0 ; p < number_of_parameters; p++) {
+    ySval[p] = N_VGetArrayPointer(yyS[p]);
+    ypSval[p] = N_VGetArrayPointer(ypS[p]);
+    for (int i = 0; i < number_of_states; i++) {
+      ySval[p][i] = y0[i + (p + 1) * number_of_states];
+      ypSval[p][i] = yp0[i + (p + 1) * number_of_states];
+    }
   }
 
-  auto t = t_np.unchecked<1>();
-  auto y0 = y0_np.unchecked<1>();
-  auto yp0 = yp0_np.unchecked<1>();
   for (int i = 0; i < number_of_states; i++)
   {
     yval[i] = y0[i];
     ypval[i] = yp0[i];
   }
 
-  realtype t0 = RCONST(t(0));
   IDAReInit(ida_mem, t0, yy, yp);
+  if (number_of_parameters > 0) {
+    IDASensReInit(ida_mem, IDA_SIMULTANEOUS, yyS, ypS);
+  }
+
+  // calculate consistent initial conditions
+  DEBUG("IDACalcIC");
+  IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1));
+  if (number_of_parameters > 0)
+  {
+    IDAGetSens(ida_mem, &t0, yyS);
+  }
 
   int t_i = 1;
   realtype tret;
@@ -368,9 +401,7 @@ Solution CasadiSolver::solve(np_array t_np, np_array y0_np, np_array yp0_np,
     }
   }
 
-  // calculate consistent initial conditions
-  DEBUG("IDACalcIC");
-  IDACalcIC(ida_mem, IDA_YA_YDP_INIT, t(1));
+
 
   int retval;
   while (true)

pybamm/solvers/c_solvers/idaklu/casadi_sundials_functions.cpp

@@ -24,9 +24,9 @@ int residual_casadi(realtype tres, N_Vector yy, N_Vector yp, N_Vector rr,
   const int ns = p_python_functions->number_of_states;
   casadi::casadi_axpy(ns, -1., tmp, NV_DATA_OMP(rr));
 
-  DEBUG_VECTOR(yy);
-  DEBUG_VECTOR(yp);
-  DEBUG_VECTOR(rr);
+  //DEBUG_VECTOR(yy);
+  //DEBUG_VECTOR(yp);
+  //DEBUG_VECTOR(rr);
 
   // now rr has rhs_alg(t, y) - mass_matrix * yp
   return 0;
@@ -263,6 +263,7 @@ int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp,
                          N_Vector tmp2, N_Vector tmp3)
 {
 
+  DEBUG("sensitivities_casadi");
   CasadiFunctions *p_python_functions =
       static_cast<CasadiFunctions *>(user_data);
 
@@ -278,7 +279,7 @@ int sensitivities_casadi(int Ns, realtype t, N_Vector yy, N_Vector yp,
   }
   // resvalsS now has (∂F/∂p i )
   p_python_functions->sens();
-
+  
   for (int i = 0; i < np; i++)
   {
     // put (∂F/∂y)s i (t) in tmp

pybamm/solvers/idaklu_solver.py

@@ -177,6 +177,17 @@ def inputs_to_dict(inputs):
             y0 = y0.full()
         y0 = y0.flatten()
 
+        y0S = model.y0S
+        # only casadi solver needs sensitivity ics
+        if model.convert_to_format != "casadi":
+            y0S = None
+        if y0S is not None:
+            if isinstance(y0S, casadi.DM):
+                y0S = (y0S,)
+
+            y0S = (x.full() for x in y0S)
+            y0S = [x.flatten() for x in y0S]
+
         if ics_only:
             return base_set_up_return
 
@@ -484,8 +495,26 @@ def _integrate(self, model, t_eval, inputs_dict=None):
             y0 = y0.full()
         y0 = y0.flatten()
 
+        y0S = model.y0S
+        # only casadi solver needs sensitivity ics
+        if model.convert_to_format != "casadi":
+            y0S = None
+        if y0S is not None:
+            if isinstance(y0S, casadi.DM):
+                y0S = (y0S,)
+
+            y0S = (x.full() for x in y0S)
+            y0S = [x.flatten() for x in y0S]
+
         # solver works with ydot0 set to zero
         ydot0 = np.zeros_like(y0)
+        if y0S is not None:
+            ydot0S = [np.zeros_like(y0S_i) for y0S_i in y0S]
+            y0full = np.concatenate([y0, *y0S])
+            ydot0full = np.concatenate([ydot0, *ydot0S])
+        else:
+            y0full = y0
+            ydot0full = ydot0
 
         try:
             atol = model.atol
@@ -499,8 +528,8 @@ def _integrate(self, model, t_eval, inputs_dict=None):
         if model.convert_to_format == "casadi":
             sol = self._setup["solver"].solve(
                 t_eval,
-                y0,
-                ydot0,
+                y0full,
+                ydot0full,
                 inputs,
             )
         else:

tests/integration/test_models/standard_model_tests.py

@@ -122,7 +122,7 @@ def test_sensitivities(self, param_name, param_value, output_name="Voltage [V]")
         output_sens = self.solution[output_name].sensitivities[param_name]
 
         # check via finite differencing
-        h = 1e-4 * param_value
+        h = 1e-2 * param_value
         inputs_plus = {param_name: (param_value + 0.5 * h)}
         inputs_neg = {param_name: (param_value - 0.5 * h)}
         sol_plus = self.solver.solve(self.model, t_eval, inputs=inputs_plus)

tests/integration/test_models/test_full_battery_models/test_lithium_ion/base_lithium_ion_tests.py

@@ -20,7 +20,13 @@ def test_basic_processing(self):
     def test_sensitivities(self):
         model = self.model()
         param = pybamm.ParameterValues("Ecker2015")
-        modeltest = tests.StandardModelTest(model, parameter_values=param)
+        if pybamm.have_idaklu():
+            solver = pybamm.IDAKLUSolver()
+        else:
+            solver = pybamm.CasadiSolver()
+        modeltest = tests.StandardModelTest(
+            model, parameter_values=param, solver=solver
+        )
         modeltest.test_sensitivities("Current function [A]", 0.15652)
 
     def test_basic_processing_1plus1D(self):

tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_dfn.py

@@ -13,10 +13,6 @@ class TestDFN(BaseIntegrationTestLithiumIon, TestCase):
     def setUp(self):
         self.model = pybamm.lithium_ion.DFN
 
-    def test_sensitivities(self):
-        # skip sensitivities test for now as it is failing with casadi 3.6
-        pass
-
     def test_particle_distribution_in_x(self):
         model = pybamm.lithium_ion.DFN()
         param = model.default_parameter_values

tests/unit/test_solvers/test_idaklu_solver.py

@@ -181,6 +181,33 @@ def test_input_params(self):
             true_solution = b_value * sol.t
             np.testing.assert_array_almost_equal(sol.y[1:3], true_solution)
 
+    def test_sensitivites_initial_condition(self):
+        model = pybamm.BaseModel()
+        model.convert_to_format = "casadi"
+        u = pybamm.Variable("u")
+        v = pybamm.Variable("v")
+        a = pybamm.InputParameter("a")
+        model.rhs = {u: -u}
+        model.algebraic = {v: a * u - v}
+        model.initial_conditions = {u: 1, v: 1}
+        model.variables = {"2v": 2 * v}
+
+        disc = pybamm.Discretisation()
+        disc.process_model(model)
+
+        solver = pybamm.IDAKLUSolver()
+
+        t_eval = np.linspace(0, 3, 100)
+        a_value = 0.1
+
+        sol = solver.solve(
+            model, t_eval, inputs={"a": a_value}, calculate_sensitivities=True
+        )
+
+        np.testing.assert_array_almost_equal(
+            sol["2v"].sensitivities["a"].full().flatten(), np.exp(-sol.t) * 2, decimal=4
+        )
+
     def test_ida_roberts_klu_sensitivities(self):
         # this test implements a python version of the ida Roberts
         # example provided in sundials