Support sensitivities for Experiments

CHANGELOG.md

@@ -1,5 +1,8 @@
 # [Unreleased](https://github.com/pybamm-team/PyBaMM/)
 
+## Features
+- Added sensitivity calculation support for `pybamm.Simulation` and `pybamm.Experiment` ([#4415](https://github.com/pybamm-team/PyBaMM/pull/4415))
+
 ## Optimizations
 - Removed the `start_step_offset` setting and disabled minimum `dt` warnings for drive cycles with the (`IDAKLUSolver`). ([#4416](https://github.com/pybamm-team/PyBaMM/pull/4416))
 

src/pybamm/simulation.py

@@ -182,6 +182,27 @@ def set_up_and_parameterise_experiment(self):
         reduces simulation time since the model formulation is efficient.
         """
         parameter_values = self._parameter_values.copy()
+
+        # some parameters are used to control the experiment, and should not be
+        # input parameters
+        restrict_list = {"Initial temperature [K]", "Ambient temperature [K]"}
+        for step in self.experiment.steps:
+            if issubclass(step.__class__, pybamm.experiment.step.BaseStepImplicit):
+                restrict_list.update(step.get_parameter_values([]).keys())
+            elif issubclass(step.__class__, pybamm.experiment.step.BaseStepExplicit):
+                restrict_list.update(["Current function [A]"])
+        for key in restrict_list:
+            try:
+                param = parameter_values[key]
+                if isinstance(param, pybamm.InputParameter):
+                    raise pybamm.ModelError(
+                        f"Cannot use '{key}' as an input parameter in this experiment. "
+                        f"This experiment is controlled via the following parameters: {restrict_list}. "
+                        f"None of these parameters are able to be input parameters."
+                    )
+            except KeyError:
+                pass
+
         # Set the initial temperature to be the temperature of the first step
         # We can set this globally for all steps since any subsequent steps will either
         # start at the temperature at the end of the previous step (if non-isothermal

src/pybamm/solvers/base_solver.py

@@ -670,6 +670,33 @@ def calculate_consistent_state(self, model, time=0, inputs=None):
         y0 = root_sol.all_ys[0]
         return y0
 
+    def _solve_process_calculate_sensitivities_arg(
+        inputs, model, calculate_sensitivities
+    ):
+        # get a list-only version of calculate_sensitivities
+        if isinstance(calculate_sensitivities, bool):
+            if calculate_sensitivities:
+                calculate_sensitivities_list = [p for p in inputs.keys()]
+            else:
+                calculate_sensitivities_list = []
+        else:
+            calculate_sensitivities_list = calculate_sensitivities
+
+        calculate_sensitivities_list.sort()
+        if not hasattr(model, "calculate_sensitivities"):
+            model.calculate_sensitivities = []
+
+        # Check that calculate_sensitivites have not been updated
+        sensitivities_have_changed = (
+            calculate_sensitivities_list != model.calculate_sensitivities
+        )
+
+        # save sensitivity parameters so we can identify them later on
+        # (FYI: this is used in the Solution class)
+        model.calculate_sensitivities = calculate_sensitivities_list
+
+        return calculate_sensitivities_list, sensitivities_have_changed
+
     def solve(
         self,
         model,
@@ -722,15 +749,6 @@ def solve(
         """
         pybamm.logger.info(f"Start solving {model.name} with {self.name}")
 
-        # get a list-only version of calculate_sensitivities
-        if isinstance(calculate_sensitivities, bool):
-            if calculate_sensitivities:
-                calculate_sensitivities_list = [p for p in inputs.keys()]
-            else:
-                calculate_sensitivities_list = []
-        else:
-            calculate_sensitivities_list = calculate_sensitivities
-
         # Make sure model isn't empty
         self._check_empty_model(model)
 
@@ -772,6 +790,12 @@ def solve(
             self._set_up_model_inputs(model, inputs) for inputs in inputs_list
         ]
 
+        calculate_sensitivities_list, sensitivities_have_changed = (
+            BaseSolver._solve_process_calculate_sensitivities_arg(
+                model_inputs_list[0], model, calculate_sensitivities
+            )
+        )
+
         # (Re-)calculate consistent initialization
         # Assuming initial conditions do not depend on input parameters
         # when len(inputs_list) > 1, only `model_inputs_list[0]`
@@ -792,13 +816,8 @@ def solve(
                     "for initial conditions."
                 )
 
-        # Check that calculate_sensitivites have not been updated
-        calculate_sensitivities_list.sort()
-        if hasattr(model, "calculate_sensitivities"):
-            model.calculate_sensitivities.sort()
-        else:
-            model.calculate_sensitivities = []
-        if calculate_sensitivities_list != model.calculate_sensitivities:
+        # if any setup configuration has changed, we need to re-set up
+        if sensitivities_have_changed:
             self._model_set_up.pop(model, None)
             # CasadiSolver caches its integrators using model, so delete this too
             if isinstance(self, pybamm.CasadiSolver):
@@ -1066,6 +1085,58 @@ def _check_events_with_initialization(t_eval, model, inputs_dict):
                 f"Events {event_names} are non-positive at initial conditions"
             )
 
+    def _set_sens_initial_conditions_from(
+        self, solution: pybamm.Solution, model: pybamm.BaseModel
+    ) -> tuple:
+        """
+        A restricted version of BaseModel.set_initial_conditions_from that only extracts the
+        sensitivities from a solution object, and only for a model that has been descretised.
+        This is used when setting the initial conditions for a sensitivity model.
+
+        Parameters
+        ----------
+        solution : :class:`pybamm.Solution`
+            The solution to use to initialize the model
+
+        model: :class:`pybamm.BaseModel`
+            The model whose sensitivities to set
+
+        Returns
+        -------
+
+        initial_conditions : tuple of ndarray
+            The initial conditions for the sensitivities, each element of the tuple
+            corresponds to an input parameter
+        """
+
+        ninputs = len(model.calculate_sensitivities)
+        initial_conditions = tuple([] for _ in range(ninputs))
+        solution = solution.last_state
+        for var in model.initial_conditions:
+            final_state = solution[var.name]
+            final_state = final_state.sensitivities
+            final_state_eval = tuple(
+                final_state[key] for key in model.calculate_sensitivities
+            )
+
+            scale, reference = var.scale.value, var.reference.value
+            for i in range(ninputs):
+                scaled_final_state_eval = (final_state_eval[i] - reference) / scale
+                initial_conditions[i].append(scaled_final_state_eval)
+
+        # Also update the concatenated initial conditions if the model is already
+        # discretised
+        # Unpack slices for sorting
+        y_slices = {var: slce for var, slce in model.y_slices.items()}
+        slices = [y_slices[symbol][0] for symbol in model.initial_conditions.keys()]
+
+        # sort equations according to slices
+        concatenated_initial_conditions = [
+            casadi.vertcat(*[eq for _, eq in sorted(zip(slices, init))])
+            for init in initial_conditions
+        ]
+        return concatenated_initial_conditions
+
     def process_t_interp(self, t_interp):
         # set a variable for this
         no_interp = (not self.supports_interp) and (
@@ -1092,6 +1163,7 @@ def step(
         npts=None,
         inputs=None,
         save=True,
+        calculate_sensitivities=False,
         t_interp=None,
     ):
         """
@@ -1117,6 +1189,11 @@ def step(
             Any input parameters to pass to the model when solving
         save : bool, optional
             Save solution with all previous timesteps. Defaults to True.
+        calculate_sensitivities : list of str or bool, optional
+            Whether the solver calculates sensitivities of all input parameters. Defaults to False.
+            If only a subset of sensitivities are required, can also pass a
+            list of input parameter names
+
         t_interp : None, list or ndarray, optional
             The times (in seconds) at which to interpolate the solution. Defaults to None.
             Only valid for solvers that support intra-solve interpolation (`IDAKLUSolver`).
@@ -1188,8 +1265,15 @@ def step(
         # Set up inputs
         model_inputs = self._set_up_model_inputs(model, inputs)
 
+        # process calculate_sensitivities argument
+        calculate_sensitivities_list, sensitivities_have_changed = (
+            BaseSolver._solve_process_calculate_sensitivities_arg(
+                model_inputs, model, calculate_sensitivities
+            )
+        )
+
         first_step_this_model = model not in self._model_set_up
-        if first_step_this_model:
+        if first_step_this_model or sensitivities_have_changed:
             if len(self._model_set_up) > 0:
                 existing_model = next(iter(self._model_set_up))
                 raise RuntimeError(
@@ -1208,18 +1292,45 @@ def step(
         ):
             pybamm.logger.verbose(f"Start stepping {model.name} with {self.name}")
 
+        using_sensitivities = len(model.calculate_sensitivities) > 0
+
         if isinstance(old_solution, pybamm.EmptySolution):
             if not first_step_this_model:
                 # reset y0 to original initial conditions
                 self.set_up(model, model_inputs, ics_only=True)
         elif old_solution.all_models[-1] == model:
-            # initialize with old solution
-            model.y0 = old_solution.all_ys[-1][:, -1]
+            last_state = old_solution.last_state
+            model.y0 = last_state.all_ys[0]
+            if using_sensitivities and isinstance(last_state._all_sensitivities, dict):
+                full_sens = last_state._all_sensitivities["all"][0]
+                model.y0S = tuple(full_sens[:, i] for i in range(full_sens.shape[1]))
+
         else:
             _, concatenated_initial_conditions = model.set_initial_conditions_from(
                 old_solution, return_type="ics"
             )
             model.y0 = concatenated_initial_conditions.evaluate(0, inputs=model_inputs)
+            if using_sensitivities:
+                model.y0S = self._set_sens_initial_conditions_from(old_solution, model)
+
+        # hopefully we'll get rid of explicit sensitivities soon so we can remove this
+        explicit_sensitivities = model.len_rhs_sens > 0 or model.len_alg_sens > 0
+        if (
+            explicit_sensitivities
+            and using_sensitivities
+            and not isinstance(old_solution, pybamm.EmptySolution)
+            and not old_solution.all_models[-1] == model
+        ):
+            y0_list = []
+            if model.len_rhs > 0:
+                y0_list.append(model.y0[: model.len_rhs])
+                for s in model.y0S:
+                    y0_list.append(s[: model.len_rhs])
+            if model.len_alg > 0:
+                y0_list.append(model.y0[model.len_rhs :])
+                for s in model.y0S:
+                    y0_list.append(s[model.len_rhs :])
+            model.y0 = casadi.vertcat(*y0_list)
 
         set_up_time = timer.time()
 

src/pybamm/solvers/c_solvers/idaklu/IDAKLUSolverOpenMP.inl

@@ -445,7 +445,7 @@ Solution IDAKLUSolverOpenMP<ExprSet>::solve(
   }
 
   if (sensitivity) {
-    CheckErrors(IDAGetSens(ida_mem, &t_val, yyS));
+    CheckErrors(IDAGetSensDky(ida_mem, t_val, 0, yyS));
   }
 
   // Store Consistent initialization
@@ -478,7 +478,7 @@ Solution IDAKLUSolverOpenMP<ExprSet>::solve(
     bool hit_adaptive = save_adaptive_steps && retval == IDA_SUCCESS;
 
     if (sensitivity) {
-      CheckErrors(IDAGetSens(ida_mem, &t_val, yyS));
+      CheckErrors(IDAGetSensDky(ida_mem, t_val, 0, yyS));
     }
 
     if (hit_tinterp) {
@@ -499,7 +499,7 @@ Solution IDAKLUSolverOpenMP<ExprSet>::solve(
         // Reset the states and sensitivities at t = t_val
         CheckErrors(IDAGetDky(ida_mem, t_val, 0, yy));
         if (sensitivity) {
-          CheckErrors(IDAGetSens(ida_mem, &t_val, yyS));
+          CheckErrors(IDAGetSensDky(ida_mem, t_val, 0, yyS));
         }
       }
 

src/pybamm/solvers/casadi_algebraic_solver.py

@@ -170,7 +170,7 @@ def _integrate(self, model, t_eval, inputs_dict=None, t_interp=None):
             model,
             inputs_dict,
             termination="final time",
-            sensitivities=explicit_sensitivities,
+            all_sensitivities=explicit_sensitivities,
         )
         sol.integration_time = integration_time
         return sol

src/pybamm/solvers/casadi_solver.py

@@ -193,7 +193,7 @@ def _integrate(self, model, t_eval, inputs_dict=None, t_interp=None):
                     y0,
                     model,
                     inputs_dict,
-                    sensitivities=False,
+                    all_sensitivities=False,
                 )
                 solution.solve_time = 0
                 solution.integration_time = 0
@@ -478,7 +478,7 @@ def integer_bisect():
             np.array([t_event]),
             y_event[:, np.newaxis],
             "event",
-            sensitivities=bool(model.calculate_sensitivities),
+            all_sensitivities=False,
         )
         solution.integration_time = (
             coarse_solution.integration_time + dense_step_sol.integration_time
@@ -696,7 +696,7 @@ def _run_integrator(
                 y_sol,
                 model,
                 inputs_dict,
-                sensitivities=extract_sensitivities_in_solution,
+                all_sensitivities=extract_sensitivities_in_solution,
                 check_solution=False,
             )
             sol.integration_time = integration_time
@@ -736,7 +736,7 @@ def _run_integrator(
                 y_sol,
                 model,
                 inputs_dict,
-                sensitivities=extract_sensitivities_in_solution,
+                all_sensitivities=extract_sensitivities_in_solution,
                 check_solution=False,
             )
             sol.integration_time = integration_time

src/pybamm/solvers/idaklu_solver.py

@@ -818,7 +818,7 @@ def _integrate(self, model, t_eval, inputs_dict=None, t_interp=None):
             np.array([sol.t[-1]]),
             np.transpose(y_event)[:, np.newaxis],
             termination,
-            sensitivities=yS_out,
+            all_sensitivities=yS_out,
         )
         newsol.integration_time = integration_time
         if not self.output_variables: