#1100 explicit foward sensitivity working

Author: valentinsulzer

examples/notebooks/DFN-sensitivity.ipynb

@@ -42,7 +42,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -64,7 +64,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -85,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -106,27 +106,27 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
-    "solver = pybamm.CasadiSolver(mode=\"fast\", sensitivity=True)\n",
+    "solver = pybamm.CasadiSolver(mode=\"fast\", sensitivity=\"casadi\")\n",
     "sim = pybamm.Simulation(model, parameter_values=param, solver=solver)\n",
     "solution = sim.solve(t_eval=np.linspace(0,3600), inputs={\"Dsn\": 1, \"Dsp\": 1})"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "1.7811190900000042"
+       "0.015949444000000312"
       ]
      },
-     "execution_count": 20,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -142,6 +142,44 @@
     "Since we have not specified the parameter values when solving, the resulting solution contains _symbolic_ variables, such as the voltage"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CPU times: user 3.73 s, sys: 194 ms, total: 3.92 s\n",
+      "Wall time: 3.91 s\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'all': DM(sparse: 1500-by-100, 73500 nnz\n",
+       "  (30, 0) -> -8.80308e-05\n",
+       "  (31, 0) -> -9.35556e-05\n",
+       "  (32, 0) -> -0.000107887\n",
+       "  ...\n",
+       "  (1497, 98) -> 0.0170617\n",
+       "  (1498, 98) -> 0.021474\n",
+       "  (1499, 98) -> 0.0260439),\n",
+       " 'Dsn': DM([00, 00, 00, ..., 0.0170617, 0.021474, 0.0260439]),\n",
+       " 'Dsp': DM([00, 00, 00, ..., 00, 00, 00])}"
+      ]
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%%time\n",
+    "solution[\"X-averaged negative particle concentration\"].sensitivity"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 17,

pybamm/solvers/base_solver.py

@@ -204,7 +204,10 @@ def set_up(self, model, inputs=None):
             model.convert_to_format = "casadi"
 
         # Only allow solving sensitivity equations with the casadi format for now
-        if self.sensitivity is True and model.convert_to_format != "casadi":
+        if (
+            self.sensitivity is "explicit forward"
+            and model.convert_to_format != "casadi"
+        ):
             raise NotImplementedError(
                 "model should be converted to casadi format in order to solve "
                 "sensitivity equations"
@@ -231,7 +234,7 @@ def set_up(self, model, inputs=None):
                     p_casadi[name] = casadi.MX.sym(name, value.shape[0])
             p_casadi_stacked = casadi.vertcat(*[p for p in p_casadi.values()])
             # sensitivity vectors
-            if self.sensitivity is True:
+            if self.sensitivity == "explicit forward":
                 S_x = casadi.MX.sym("S_x", model.len_rhs * p_casadi_stacked.shape[0])
                 S_z = casadi.MX.sym("S_z", model.len_alg * p_casadi_stacked.shape[0])
                 y_and_S = casadi.vertcat(y_diff, S_x, y_alg, S_z)
@@ -286,7 +289,7 @@ def report(string):
                 report(f"Converting {name} to CasADi")
                 func = func.to_casadi(t_casadi, y_casadi, inputs=p_casadi)
                 # Add sensitivity vectors to the rhs and algebraic equations
-                if self.sensitivity is True:
+                if self.sensitivity == "explicit forward":
                     if name == "rhs" and model.len_rhs > 0:
                         report("Creating sensitivity equations for rhs using CasADi")
                         df_dx = casadi.jacobian(func, y_diff)
@@ -408,7 +411,7 @@ def report(string):
         )[0]
         init_eval = InitialConditions(initial_conditions, model)
 
-        if self.sensitivity is True:
+        if self.sensitivity == "explicit forward":
             init_eval.y_dummy = np.zeros(
                 (
                     model.len_rhs_and_alg * (np.vstack(list(inputs.values())).size + 1),
@@ -456,7 +459,7 @@ def report(string):
         ):
             # can use DAE solver to solve model with algebraic equations only
             if len(model.rhs) > 0:
-                if self.sensitivity is True:
+                if self.sensitivity == "explicit forward":
                     # Copy mass matrix blocks diagonally
                     single_mass_matrix_inv = model.mass_matrix_inv.entries.toarray()
                     n_inputs = p_casadi_stacked.shape[0]
@@ -946,7 +949,7 @@ def _set_up_ext_and_inputs(self, model, external_variables, inputs):
             name = input_param.name
             if name not in inputs:
                 # Don't allow symbolic inputs if using `sensitivity`
-                if self.sensitivity is True:
+                if self.sensitivity == "explicit forward":
                     raise pybamm.SolverError(
                         "Cannot have symbolic inputs if explicitly solving forward"
                         "sensitivity equations"