IDAKLU solver sensitivities using IDAS

[martinjrobins]

Summary

Extend the IDAKLU solver to return sensitivities of y wrt the parameters p

Motivation

It is useful to calculate sensitivities in, for example, model fitting to data. The IDAKLU solver already using IDA in sundials, and IDAS is another part of the sundials package which can calculate sensitivities, so in theory this should be easily achievable.

Additional context

Work on incorporating sensitivities for the casadi solver is already underway in #1100

[martinjrobins]

@tinosulzer: were you planning that the API for sensitivities you have implemented in #1100 is just for the casadi solver, or should I also follow this for the idaklu solver? I'm worried that this API seems rather casadi specific, what are your thoughts on this?

[valentinsulzer]

Ideally the API should just be something like

sol = sim.solve(t_eval, inputs={"p": 2, "q": 3}
sol["variable"].data # returns the values
sol["variable"].sensitivity["q"] # returns the sensitivities wrt q

which then works with any solver. This involves getting rid of the ProcessedSymbolicVariable class which is casadi-specific

[martinjrobins]

Should we put a flag or something going into sim.solve which indicates if sensitivities are required?

[valentinsulzer]

IIRC in the 1100 branch I made it a flag in the Solver class (so you can decide between explicitly defining the SDAEs or letting the solver do its own fancy forward sensitivity / adjoint stuff)

[martinjrobins]

is there any reason that the inputs arg to solve has keys of strings, rather than InputParameters? I feel that something like this would be more natural:

a = pybamm.InputParameter("a")
# ...
sol = solver.solve(model, t_eval, inputs={a: 1.0})

It also helps when I'm calculating the sensititivites of rhs and algebraic wrt the input parameters using Jacobian.jac, as I don't have to re-create the InputParameter. @tinosulzer : what do you think?

[martinjrobins]

actually, maybe if the user is using a pre-built model they won't have access to the InputParameter instance, maybe this is why we need strings?

Anyhoo, I've drafted out some changes to the base solver that I'm thinking of for the sensitivities. I'm thinking that if sensitivities are requested, then a dict of {p: drhs / dp} for each input parameter p is calculated, as well as {p: dalg / dp} of course. This is calculated using python, jax or casadi depending on the convert_to_format. Each solver can then use these to calculate the sensitivities. I've split up the derivatives wrt each parameter separately, rather than stacking them, cause this is how idas wants them. If individual solvers need them stacked, then of course they can do so.

[valentinsulzer]

Sounds good

This is calculated using python, jax or casadi depending on the convert_to_format

Won't this depend on which solver is used rather than the convert_to_format?

[martinjrobins]

Won't this depend on which solver is used rather than the convert_to_format?

But some of the solvers support multiple convert_to_format, right? Like the scipy solver. Generally the base solver seems to rely on convert_to_format to decide what format to convert the rhs and alg functions into.

[martinjrobins]

also need to add flag to solve function to say which sensitivities will be required:

# calculate sensitivities for "a" input
sol = solver.solve(model, t_eval, inputs={"a": a_value}, calculate_sensitivites=["a"])
# calculate sensitivities for all input parameters
sol = solver.solve(model, t_eval, inputs={"a": a_value}, calculate_sensitivites=True)