Add support for Hamiltonian Neural Networks

[avik-pal]

Fixes #53
Fixes #357
Fixes #80 when this PR is used together with #355

TODOs:

• Implement Neural Hamiltonian DE Layer
• Example for using the layer
• Working on GPU? Will be fixed with Keno's AD Add support for Hamiltonian Neural Networks #370 (comment)
  • HNN layer
  • Neural Hamiltonian DE Layer
• Documentation
• Tests

[avik-pal]

Zygote seems to complain about mutating arrays when backpropagating through the gradient of the NN. Here's a small repro:

using Flux, Zygote

model = Chain(
    Dense(2, 10),
    Dense(10, 1)
)

forward(m, x) = gradient(x -> m(x)[1], x)[1]

x = rand(2, 1)

gradient(() -> sum(forward(model, x) - x), Flux.params(model))

Is this to be done in some other way?

[ChrisRackauckas]

Nesting Zygote is an issue. I think you need to use ReverseDiff over Zygote, or Zygote over ReverseDiff.

[avik-pal]

Yes, ReverseDiff over Zygote does solve this issue.

[avik-pal]

For training HNNs the standard way (torchdyn, original paper) seems to be to supervise on the gradients of the network, rather than on the solution from the ODE Solver. I was wondering if we should have a HamiltonianNN layer separate from the NeuralHamiltonianDE which simply returns the gradient and the latter just acts as a wrapper over the former to compute the solutions of the ODE.

[ChrisRackauckas]

Yeah, that would make sense. It would then match well with our collocation tools:

https://diffeqflux.sciml.ai/dev/examples/collocation/
https://diffeqflux.sciml.ai/dev/Collocation/

We should also do a non-smoothed collocation method via DataInterpolations, and demonstrate that on this HamiltonianNN

[ChrisRackauckas]

Also, just a note, some of the models from the HNN paper are impossible to fit directly from the solve on the time-series. Specifically the double pendulum, since it's chaotic so the adjoints diverge exponentially fast.

[avik-pal]

ReverseDiff seems to bug out on GPU with CuArray only supports bits types.

hnn = HamiltonianNN(
    Chain(Dense(2, 64, relu), Dense(64, 1)) |> gpu
)

p = hnn.p;

hnn(rand(2, 1) |> gpu)

ReverseDiff.gradient(p -> sum(hnn(rand(2, 1) |> gpu, p)), p)

[ChrisRackauckas]

ReverseDiff seems to bug out on GPU with CuArray only supports bits types.

Yes it does. I think we'll need to just skip that for now, open an issue, and grab what we can get. Keno's AD should fix this.

[Vaibhavdixit02]

Looks great! 🎉

[avik-pal]

We should also do a non-smoothed collocation method via DataInterpolations, and demonstrate that on this HamiltonianNN

@ChrisRackauckas could you elaborate a bit on this?

[ChrisRackauckas]

Essentially one way to perform this fitting is to fit spliens like a cublic spline from DataInterpolations, and then take the derivative of the spline to give you estimated derivatives, and then perform the fit directly against those derivative approximations. It would be nice to have this version under the same interface as the smoothed one, which is essentially the same idea except using regression splines under some smoothing kernel to get derivatives of very noisy and dense data.