search

rtqichen / torchdiffeq

Differentiable ODE solvers with full GPU support and O(1)-memory backpropagation.
MIT License
5.45k stars 909 forks source link

`odeint_adjoint` with constant derivative / external controls #191

Closed jaivardhankapoor closed 2 years ago

jaivardhankapoor commented 2 years ago

Hi, I am trying to integrate a differential equation with the following formulation: dx/dt = v(t | z), where v is independent of x and is outputted from a neural net with input z. In this, v can be thought of as an external control.

The issue is that odeint_adjoint only supports input as func of the form func(t, x). How do I include v(t|z) into the odeint_adjoint arguments without augmenting the differential equation itself (d[x;v]/dt=[v(t|z);0])?

Apologies if this is a common question that I might have somehow missed.

EDIT:

If v doesn't depend on x, couldn't you compute the exact solution?

In any case, it's not clear to me if v is a function or a variable, but you can include external inputs as long as they're in the scope of func. This should work automatically with odeint, and with odeint_adjoint you'd have to specify the extra parameters through the adjoint_params option. Maybe #145 (comment) can help.

the ODE also depends on x, along with an external z: dx/dt = v(t, x| z). So It cannot be solved exactly. Apologies for the confusion.

rtqichen commented 2 years ago

If v doesn't depend on x, couldn't you compute the exact solution?

In any case, it's not clear to me if v is a function or a variable, but you can include external inputs as long as they're in the scope of func. This should work automatically with odeint, and with odeint_adjoint you'd have to specify the extra parameters through the adjoint_params option. Maybe https://github.com/rtqichen/torchdiffeq/issues/145#issuecomment-762480592 can help.

xxh0523 commented 2 years ago

Hi. I am a power system researcher. Your question is very similar to the differential-algebraic equations (DAE) of power systems. I have developed a neural DAE module based on torchdiffeq. Maybe you can look to my work at https://github.com/xxh0523/Py_PSNODE and https://arxiv.org/abs/2110.12981. However, the adjoint approach is analyzed in the paper but not implemented since in power system simulations, the integration step is usually fixed.

jaivardhankapoor commented 2 years ago

Hi. I am a power system researcher. Your question is very similar to the differential-algebraic equations (DAE) of power systems. I have developed a neural DAE module based on torchdiffeq. Maybe you can look to my work at https://github.com/xxh0523/Py_PSNODE and https://arxiv.org/abs/2110.12981. However, the adjoint approach is analyzed in the paper but not implemented since in power system simulations, the integration step is usually fixed.

Thank you for directing me toward this @xxh0523. I am confident that dx/dt = v(t, x| z) is a form of DAE . Indeed, the most straightforward way to interface extra controls with odeint_adjoint seems to be augmenting the ODE equation using d[x(t);v(x,t|z)] / dt = [v(t,x|z);0] as mentioned in the original question. In case this might lead to greater memory consumption and wierd adjoint issues (these might be actually happening to me right now) I will look into this solution as mentioned by @rtqichen. Closing this issue as current interface seems to be able to handle this for now. Thanks for your help!

juliusaka commented 8 months ago

@jaivardhankapoor I'm facing a similar question. I'm curious, how did you solve that in the end?