Deep Bayesian Model Discovery without using NeuralODE object

[gurtajbir]

Hi everyone. I am trying to implement the Deep Bayesian Model Discovery on the Lotka-Volterra model discussed in the Automated Discovery of Missing Physics example. The problem I am facing is that I am not able to figure out a way to pass the parameters of the neural network embedded in the ODE of the Lotka-Volterra model to the Hamiltonian as done here. The main issue here is that the hamiltonian is fed a vector of parameters and they are updated naturally as the optimization is carried out. I am having trouble achieving the same with the missing physics example.
Any pointers as to how this can be achieved will be very helpful. Thanks.

[ChrisRackauckas]

What have you tried so far? You'd do the exact same thing except change the optimization to the Bayesian fitting routine.

[gurtajbir]

Hi Chris. I was trying to extract the parameters of the Flux model as a vector using a function like the following for a model U,

`p_model = []
for i in 1:length(U.layers)
weights = Float64.(Flux.params(U.layers[i])[1])
for row in weights
p_model = [p_model; row]
end
biases = Float64.(Flux.params(U.layers[i])[2])
p_model = [p_model; biases]
end'

This was now in the correct format to be fed into metric = DiagEuclideanMetric(length(p)) and integrator = Leapfrog(find_good_stepsize(h, p)). But the trouble was with converting the update to the parameter vector back to something that could be used by the neural network in the UDE dynamics.

[gurtajbir]

I had posted the same on Julia discourse. I was advised to use θ, re = Flux.destructure(chain) which now enables me to make predictions in the UDE dynamics using the new set of parameters to re . This implementation on my part is inefficient but it able to get the job done as in it is able to take the new parameters and make a prediction based on that
function ude_dynamics!(du, u, p, t, p_true)
U = re(p) # Reconstruct with passed parameters p
û = U(u) # Network prediction
du[1] = p_true[1] * u[1] + û[1]
du[2] = -p_true[4] * u[2] + û[2]
end

[gurtajbir]

What have you tried so far? You'd do the exact same thing except change the optimization to the Bayesian fitting routine.

What change are you referring to here? Do you mean a change to the loss l(θ) = -sum(abs2, ode_data .- predict_neuralode(θ)) - sum(θ .* θ) or somewhere here
integrator = Leapfrog(find_good_stepsize(h, p)) prop = AdvancedHMC.NUTS{MultinomialTS, GeneralisedNoUTurn}(integrator) adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.45, integrator)) samples, stats = sample(h, prop, p, 500, adaptor, 500; progress = true)

I do not have a strong background in Bayesian Inference so most of this stuff is new to me. I did go on try exactly the same as in the tutorial and got the following result

[ChrisRackauckas]

I highly recommend just changing to a Lux network instead of a Flux one. Is the tutorial using Flux? If so we should just update it

[Vaibhavdixit02]

@gurtajbir it looks like you are on the right track, how many samples is that plot from? Are you dropping the warmup samples?

[gurtajbir]

@gurtajbir it looks like you are on the right track, how many samples is that plot from? Are you dropping the warmup samples?

Hi @Vaibhavdixit02 . This plot was using the below
samples, stats = sample(h, prop, p, 500, adaptor, 2000; progress = true). These particular values were used according to the code provided here.

[gurtajbir]

I highly recommend just changing to a Lux network instead of a Flux one. Is the tutorial using Flux? If so we should just update it

Hi @ChrisRackauckas. Just changing to Lux seemed to considerably improve how the plot looked. Also, with same network size and activation function, the code with Lux faster than that with Flux (almost 3 times as fast).

[ChrisRackauckas]

Yeah that's expected. Could you make a PR to update that code?

[gurtajbir]

Would you like me to update the existing deep bayesian discovery example in the PR or create a code file that resembles what I am trying out on Lotka-Volterra ?

[ChrisRackauckas]

Updating the deep bayesian discovery example in a PR that changes it to Lux (and ultimately improves its stability) would be perfect.

[gurtajbir]

Sounds good. I'll get started on it.

[Doublonmousse]

I've made a gist redoing https://docs.sciml.ai/Overview/stable/showcase/bayesian_neural_ode/#Step-5:-Plot-diagnostics with Lux and no NeuralODE object here : https://gist.github.com/Doublonmousse/64c5226edc8419dc6bf8e0594f2cb89f.

If this good enough ? I'll do a PR to update the example in this repo if it's okay

[ChrisRackauckas]

Yeah that looks good, though I don't think you need the Metal import?

[Doublonmousse]

Yeah, I'm not using it (it's also imported twice ..).

Maybe vector_to_parameters is superfluous or can be improved as well?

I'll open a PR for it then soon so that we can discuss further things there.

[Doublonmousse]

Wait, the github version is already up to date with Lux

SciMLDocs/docs/src/showcase/bayesian_neural_ode.md

Lines 21 to 22 in df454f4

	# ML Tools
	using Lux, Zygote

With the change dating back from #217

but the live web version isn't up to date yet.

[ChrisRackauckas]

There is currently an issue with the doc building due to a bug in CUDA #224