Using machine learning packages such as pytorch in pybamm

[satishrapol]

Hello,
I am interested in using a machine learning model which uses a pybamm variable as input as shown below.

x = pybamm.Variable("x")
dx/dt = 4 * x + pytorch_model(x)

where consider pytorch_model as a simple feedforward model.

But I am unable to simulate as shown above as the input to the pytorch_model is a symbolic term and not a numerical value which the pytorch_model usually expects.

Please let me know how do I achieve this

@martinjrobins @brosaplanella @valentinsulzer

[brosaplanella]

I have never used pytorch, but PyBaMM is not designed to be compatible with it so I don't think what you suggest is possible.

[martinjrobins]

If you can convert your pytorch model to JAX, then you might be able to use the JAX solver in pybamm to solve it and calculate gradients for training. We're also trying to get the IDAKLU solver to use JAX models as well, so this might be a possibility in the future. If you go down this route I'd be keen to hear about it!

[martinjrobins]

there is no user-facing API for adding arbitrary JAX terms to the rhs equations however, so this would have to be a new addition to pybamm (it shouldn't be too difficult to add)

[satishrapol]

Okay, what about the below approach? Would this be a reasonable approximation?
I am using inbuilt pybamm functions here, I understand that the symbolic differentiation of x is not accounted in the approach below but could be a reasonable solution ,if the resolution of time array below is high enough?

time = np.linspace(0, 100, 100000)
predictions = nn_model(time)
model = pybamm.BaseModel()
x = pybamm.Variable("x")
dx/dt = 4 * x + pybamm.Interpolant(time, predictions, x)

[martinjrobins]

so your nn_model does not depend on the current state of the model? only on the value of time? I ask because your pytorch_model above does depend on x. But if it just depents on time, then yea that looks reasonable. But then I'm not sure what you are trying to achieve with this model?

[satishrapol]

I am also pondering if it is mathematically correct or not, but I see it like this, since nn_model are abstract functions, the idea here is to add a neural network term to improve the overall ode.

With a loss function of my choice, I would like to tune the parameters of the nn_model; here, the time can be simply called x_values. Here, nn_model should be a function of x through interpolated x_values. During training nn_model accessing x through x_values will also exist, so I am guessing the tuning will also inherit this mapping.

x_values = np.linspace(0, 100, 100000)
predictions = nn_model(x_values)
model = pybamm.BaseModel()
x = pybamm.Variable("x")
dx/dt = 4 * x + pybamm.Interpolant(x_values, predictions, x)

Let me know if it sounds fine :-)

[martinjrobins]

I think this would remove any advantage of using neural networks in the first place. NN's are useful because they are a function approximator that is good at capturing the non-linearities of your data. A linear (or cubic) interpolator is also a function approximator, but they are not good at capturing these non-linearities (otherwise we might have papers entitled "interpolants are all you need" ;) ). You are taking the output of your nn and feeding it into your interpolant, and therefore taking a worse approximation of your approximation of the data. You might as well just directly train your interpolant model on the data like this:

x_values = np.linspace(0, 100, 100000)
model = pybamm.BaseModel()
x = pybamm.Variable("x")
dx/dt = 4 * x + pybamm.Interpolant(x_values, parameters_to_train, x)

I think the best way to approach this would be to modify pybamm so that it can accept a pytorch nn and use this directly in the model. I think the best way to do this would be to create a new pybamm.Symbol derived class that holds a pytorch model. Then modify the find_symbols function here to convert that node in the expression tree to valid JAX code, perhaps using something like this