using pybamm variables to implement cost functions

[martinjrobins]

Feature description

PyBaMM uses model variables to do post-processing calculations on solved models, e.g. calculating the voltage from the internal state vector returned by the solver. This is quite similar in concept to the cost functions implemented in PyBop: "take the output of a simulation and calculate some function from this", so I would suggest that we can simply use PyBaMM variables to implement each cost function.

Motivation

This was inspired by this problem suggested by @NicolaCourtier:

Another example is: Imagine we want to optimise the energy density so we are scaling the simulation result by the cell mass. See this example. The mass of each component is not included in the standard parameter set because they are not used in typical models. When optimising the volume fractions, we then require parameters which are not used by PyBaMM but are used in the cost function to compute the cell mass.
(We currently solve the problem by making assumptions and neglecting the parameters not in the PyBaMM parameter set)

The advantage of using PyBaMM's variables are:

• more efficient calculation. Using the pybamm solvers output_variables argument the calculation of the cost function can be moved into the solver itself
• use of pybamm's flexible expression trees to specify cost functions, allowing us to use additional parameters in the cost function that are not part of the base model
• simplifies the pybop code as there is no need to calculate any of the cost functions, you just use the solution object like so: sol["my_cost_function"].data

Possible implementation

This would require some additions on the pybamm side. The pybamm.ExplicitTimeIntegral allows us to implement some of the design cost functions, but you would need something like pybamm.DiscreteSumOverTime to implement the fitting cost functions. You would also need to support the calculation of sensitivities for these symbol classes.

Additional context

No response

[BradyPlanden]

This is an interesting proposal - it seems to offer a number of improvements. A couple of questions come to mind:

1. For the fitting costs and likelihoods, do you foresee any problems in translating the time-series reference into the solver? I.e. the solver would need to store the reference object and integrate over a distance function defined by the forward model prediction and the reference. I suspect PyBOP would need to interact more with the solver metrics to parse potential errors and feed them back to the user. Catching and parsing the error from the solver function isn't new to PyBOP, but I suspect we would need to do more.
2. Does the implementation change for meta-cost classes, i.e. WeightedCost or LogPosterior?
3. Same as 2, but for the transformation class. I think the underlying question is how difficult is it to implement this functionality in the solvers. Whether not this would be doable for most of our developers, or if barrier of entry increases majorly.
4. Any thoughts on how difficult it would be to robustly capture the hessian information for the variables? I think it would be easier than the current implementation (since we wouldn't have to convert first and second order sensitivities into corresponding cost derivatives).
5. "additional parameters in the cost function that are not part of the base model" - this sounds nice, but I'm not sure how it differs from our "additional variables" approach in the problem class. I suspect it's just nicely integrated into the PyBaMM model without PyBOP having to define those variables, is that correct?

Overall this sounds like a sensible direction. I'm keen to make sure we can still use the Jax-based solvers, but I don't think this would interfere with that. Do you think a small toy example could be spun up for us to discuss around?

[martinjrobins]

1. For the fitting costs and likelihoods, do you foresee any problems in translating the time-series reference into the solver? I.e. the solver would need to store the reference object and integrate over a distance function defined by the forward model prediction and the reference. I suspect PyBOP would need to interact more with the solver metrics to parse potential errors and feed them back to the user. Catching and parsing the error from the solver function isn't new to PyBOP, but I suspect we would need to do more.

by "time-series reference" do you mean the data? My first thought would be that it would be held by the pybamm.DiscreteSumOverTime node in the expression tree and would be converted into a casadi matrix to be embedded within the final casadi function (in much the same way as we treat discretisation matrices in the expressions)

yea, this would mean interacting much more with the pybamm model class, so there would be increased error handling due to that

2. Does the implementation change for meta-cost classes, i.e. WeightedCost or LogPosterior?

I don't think so, you can add any number of variables to a pybamm model. Then the WeightedCost would then extract all of them and sum and weight them appropriately.

3. Same as 2, but for the transformation class. I think the underlying question is how difficult is it to implement this functionality in the solvers. Whether not this would be doable for most of our developers, or if barrier of entry increases majorly.

The transformation classes transform the input parameters from the optimizer before you simulate the model right? If so then I don't think this would affect the transformation classes at all (this would be downstream of the transformation)

4. Any thoughts on how difficult it would be to robustly capture the hessian information for the variables? I think it would be easier than the current implementation (since we wouldn't have to convert first and second order sensitivities into corresponding cost derivatives).

not difficult, just slow. pybamm only has forward sensitivity analysis, so we'd still need to add backwards adjoint features, and then hessian calculation. Sundials has all this however, so its just a matter of exposing the functionality. Its slow going, but its on the roadmap

5. "additional parameters in the cost function that are not part of the base model" - this sounds nice, but I'm not sure how it differs from our "additional variables" approach in the problem class. I suspect it's just nicely integrated into the PyBaMM model without PyBOP having to define those variables, is that correct?

additional variables are output variables right? I'm talking about input paramers here. So in the example @NicolaCourtier mentioned above this would be the mass of each component (I believe).

Overall this sounds like a sensible direction. I'm keen to make sure we can still use the Jax-based solvers, but I don't think this would interfere with that. Do you think a small toy example could be spun up for us to discuss around?

yea, I think this is a good idea. I'll put something together

[BradyPlanden]

Thanks for the detailed response, that helps clarify.

by "time-series reference" do you mean the data?

Yep, the reference provided by the user. At the moment we construct this within the pybop.Dataset class as Numpy objects. We may need to rethink what the construction order needs to be for this design, as the cost function will need to be defined before / within the model construction.

I don't think so, you can add any number of variables to a pybamm model. Then the WeightedCost would then extract all of them and sum and weight them appropriately.

This sounds reasonable. In this situation the meta-cost classes would still be child classes of pybop.BaseCost, but the composed costs would just be output variables from the pybamm solution. I think this could clean up meta-costs quite a bit, as it would remove the complex interaction of differentiating from normal classes.

The transformation classes transform the input parameters from the optimizer before you simulate the model right? If so then I don't think this would affect the transformation classes at all (this would be downstream of the transformation)

Yup, duh. Don't mind me over here.

additional variables are output variables right? I'm talking about input paramers here. So in the example @NicolaCourtier mentioned above this would be the mass of each component (I believe).

Right, ok, so is the idea here is that these parameters are exposed at the right level (within the expression tree), so they don't need to be funnelled up to the cost classes in their current definition. That's an understandable improvement.

yea, I think this is a good idea. I'll put something together

Cool, I'm looking forward to discussing in more depth!

One area that I think we will need to work on is ensuring the testing suite on the PyBaMM side covers this functionality (assuming they are happy with these additions) as I suspect small changes in the solver / expression tree codebase could break functionality on our side without it effecting the majority of users of PyBaMM. Just to reiterate, overall this sounds promising though.

[martinjrobins]

FYI I put together a toy example just in pybamm. The general idea is that cost functions just add new variables and optionally parameters to the original model

import pybamm

model = pybamm.lithium_ion.SPMe()
parameters = ["Positive electrode thickness [m]", "Positive particle radius [m]", "Cell mass [kg]"]
experiment = pybamm.Experiment(
    ["Discharge at 1C until 2.5 V (5 seconds period)"],

)

# add cell mass to the model
parameter_set = model.default_parameter_values
parameter_set.update({"Cell mass [kg]": 0.1}, check_already_exists=False)

# add gravimetrical energy density cost function to the model
cell_mass = pybamm.Parameter("Cell mass [kg]")
voltage = model.variables["Voltage [V]"]
current = model.variables["Current [A]"]
model.variables["Gravimetrical energy density [Wh kg-1]"] = pybamm.ExplicitTimeIntegral(current * voltage / 3600 / cell_mass, pybamm.Scalar(0.0))

# define cost function
def cost_function(args):
    new_parameters = parameter_set.copy()
    for i, parameter in enumerate(parameters):
        new_parameters[parameter] = args[i]
    sim = pybamm.Simulation(model, parameter_values=new_parameters, experiment=experiment)
    sol = sim.solve()
    return sol["Gravimetrical energy density [Wh kg-1]"](sol.t[-1])

print('eval', cost_function([8.27568156e-05, 2.07975193e-06, 1.0]))

[martinjrobins]

I note that there is also a proposal to use jax functions to calculate cost functions (#481).