Learning Orbital Dynamics

Hi everyone, I am currently learning Julia by working on a project based on the paper “Learning orbital dynamics of binary black hole systems from gravitational wave measurements.” The idea is to reconstruct the original orbit of the binary black hole system using measured waveforms, achieved by training a neural network.

I am encountering several issues because the code was written in an older version of Julia, preventing me from defining the FastChain and FastDense neural networks, among other problems.I was wondering if someone could help me with either installing an older version of Julia that is compatible with the code (I am using VSCode on an M2 MacBook), or updating SXS1.jl to make it compatible with the current version.

Also this is the link to the original code Three examples of learning orbital dynamics of binary black hole systems from gravitational wave measurements with Julia
SXS1.jl (8.6 KB)
utils.jl (5.6 KB)
models.jl (2.5 KB)

If you have any question, or required a more detail explanation, please let me know.
Kind regards

An updated version of that code was made into a tutorial and is hosted here:

4 Likes

Thank you for your answer. I was following the tutorial and using as a template. I use the PINNs AbstractNROrbitModel, then define both neural networks as defined in the example. But I get the error type Array has no field layer_1

``````# Defines the PINNs
function AbstractNROrbitModel(u, model_params, t;
NN_chiphi=nothing, NN_chiphi_params=nothing,
NN_pe=nothing, NN_pe_params=nothing)
u[1] = χ
u[2] = ϕ
u[3] = p
u[4] = e

q is the mass ratio
=#
χ, ϕ, p, e = u
q = model_params[1]
M = 1.0

if p <= 0
println("p = ", p)
end

if isnothing(NN_chiphi)
nn_chiphi = [1, 1]
else
nn_chiphi = 1 .+ NN_chiphi(u, NN_chiphi_params, st_chiphi)
end

if isnothing(NN_pe)
nn_pe = [0, 0]
else
nn_pe = NN_pe(u, NN_pe_params, st_pe)
end

numer = (1 + e * cos(χ))^2
denom = M * (abs(p)^(3 / 2))

χ̇ = (numer / denom) * nn_chiphi[1]
ϕ̇ = (numer / denom) * nn_chiphi[2]
ṗ = nn_pe[1]
ė = nn_pe[2]

return [χ̇, ϕ̇, ṗ, ė]
end
``````
``````## Define neural network chi and phi
NN_chiphi = Lux.Chain((x) -> [cos(x[1]),1/abs(x[3]),1/sqrt(abs(x[3])),sqrt(abs(x[3])),x[3],sqrt(abs(x[3]))^3,x[3]^2,x[4],x[4]^2],
Lux.Dense(9, 32, tanh),
Lux.Dense(32, 2))
p_chiphi, st_chiphi = Lux.setup(rng, NN_chiphi)
NN_chiphi_params = ComponentArray{Float64}(p_chiphi)

## Define neural network p and e
NN_pe = Lux.Chain((x) -> [1/sqrt(abs(x[3]))^3,1/abs(x[3]),1/sqrt(abs(x[3])),sqrt(abs(x[3])),x[3],sqrt(abs(x[3]))^3,x[3]^2,x[4],x[4]^2,x[3]*x[4]],
Lux.Dense(10,32, tanh),
Lux.Dense(32,2))
p_pe, st_pe = Lux.setup(rng, NN_pe)
NN_pe_params = ComponentArray{Float64}(p_pe)

# Concatenate parameters
NN_params_vcat = vcat(NN_chiphi_params,NN_pe_params)
NN_params = ComponentArray{Float64}(NN_params_vcat)
l1 = length(NN_chiphi_params)
``````
``````function ODE_model(u, NN_params, t)
NN_params1 = NN_params[1:l1]
NN_params2 = NN_params[l1+1:end]
du = AbstractNROrbitModel(u, model_params, t,
NN_chiphi=NN_chiphi, NN_chiphi_params=NN_params1,
NN_pe=NN_pe, NN_pe_params=NN_params2)
return du
end
``````
``````prob_nn = ODEProblem(ODE_model, u0, tspan, NN_params)
soln_nn = Array(solve(prob_nn, RK4(), u0 = u0, p = NN_params, saveat = tsteps, dt = dt, adaptive=false))
``````

type Array has no field layer_1

I would appreciate your help with the problem. Please let me know if you need more details or the code file.

Best regards.

Instead of vcating, make a componentarray of componentarrays.

`ComponentArray(NN_params1 = NN_chiphi_params, NN_params2 = NN_pe_params)`

and use that

``````function ODE_model(u, NN_params, t)
NN_params1 = NN_params.NN_params1
NN_params2 = NN_params.NN_params2
du = AbstractNROrbitModel(u, model_params, t,
NN_chiphi=NN_chiphi, NN_chiphi_params=NN_params1,
NN_pe=NN_pe, NN_pe_params=NN_params2)
return du
end
``````

so that you retain the structure instead of turning everything into arrays

2 Likes

Chris, thank you so much for taking the time to answer my question. There has been an improvement, and I followed your suggestion, but now I have the problem MethodError: no method matching +(::Int64, ::Vector{Float64}). For element-wise addition, use broadcasting with dot syntax: scalar .+ array.

This is the current state of the code (with your suggestion)

``````## Define neural network chi and phi
NN_chiphi = Lux.Chain((x) -> [cos(x[1]),1/abs(x[3]),1/sqrt(abs(x[3])),sqrt(abs(x[3])),x[3],sqrt(abs(x[3]))^3,x[3]^2,x[4],x[4]^2],
Lux.Dense(9, 32, tanh),
Lux.Dense(32, 2))
p_chiphi, st_chiphi = Lux.setup(rng, NN_chiphi)
NN_chiphi_params = ComponentArray{Float64}(p_chiphi)

## Define neural network p and e
NN_pe = Lux.Chain((x) -> [1/sqrt(abs(x[3]))^3,1/abs(x[3]),1/sqrt(abs(x[3])),sqrt(abs(x[3])),x[3],sqrt(abs(x[3]))^3,x[3]^2,x[4],x[4]^2,x[3]*x[4]],
Lux.Dense(10,32, tanh),
Lux.Dense(32,2))
p_pe, st_pe = Lux.setup(rng, NN_pe)
NN_pe_params = ComponentArray{Float64}(p_pe)

# Neural Networks parameters
NN_params = ComponentArray(NN_params1 = NN_chiphi_params, NN_params2 = NN_pe_params)
``````
``````# The PINN model
function AbstractNROrbitModel(u, model_params, t; NN_chiphi=nothing, NN_chiphi_params=nothing,NN_pe=nothing, NN_pe_params=nothing)
χ, ϕ, p, e = u
q = model_params[1]
M = 1.0

if p <= 0
println("p = ", p)
end

if isnothing(NN_chiphi)
nn_chiphi = [1, 1]
else
nn_chiphi = 1 .+ NN_chiphi(u, NN_chiphi_params, st_chiphi)
end

if isnothing(NN_pe)
nn_pe = [0, 0]
else
nn_pe = NN_pe(u, NN_pe_params, st_pe)
end

numer = (1 + e * cos(χ))^2
denom = M * (abs(p)^(3 / 2))

χ̇ = (numer / denom) * nn_chiphi[1]
ϕ̇ = (numer / denom) * nn_chiphi[2]
ṗ = nn_pe[1]
ė = nn_pe[2]

return [χ̇, ϕ̇, ṗ, ė]
end
``````
``````function ODE_model(u, NN_params, t)
NN_params1 = NN_params.NN_params1
NN_params2 = NN_params.NN_params2
du = AbstractNROrbitModel(u, model_params, t, NN_chiphi=NN_chiphi, NN_chiphi_params=NN_params1, NN_pe=NN_pe, NN_pe_params=NN_params2)
return du
end
``````
``````prob_nn = ODEProblem(ODE_model, u0, tspan, NN_params)
soln_nn = Array(solve(prob_nn, RK4(), u0 = u0, p = NN_params, saveat = tsteps, dt = dt, adaptive=false))
``````

MethodError: no method matching +(::Int64, ::Vector{Float64})

Kind regards
Jose

what line is it pointing to?

1 Like

The complete error is

``````MethodError: no method matching +(::Int64, ::Vector{Float64})

Closest candidates are:
+(::Any, ::Any, !Matched::Any, !Matched::Any...)
@ Base operators.jl:587
+(!Matched::ChainRulesCore.NoTangent, ::Any)
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/I1EbV/src/tangent_arithmetic.jl:59
+(::Any, !Matched::MutableArithmetics.Zero)
@ MutableArithmetics ~/.julia/packages/MutableArithmetics/SXYDN/src/rewrite.jl:65
...
Stacktrace:
[4] ntuple
@ ./ntuple.jl:49 [inlined]
[5] copy
...
@ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1066 [inlined]
[19] #solve#51
@ ~/.julia/packages/DiffEqBase/PBhFc/src/solve.jl:1003 [inlined]
[20] top-level scope
@ ~/Desktop/julia_iGW/SXS_BBH_0217/SXS0217.ipynb:2
``````

the libraries that I am using are

``````# SciML Tools
using Optimization, OptimizationOptimisers, OptimizationOptimJL

# Standard Libraries
using LinearAlgebra, Statistics

# External Libraries
using ComponentArrays, Lux, Zygote, Plots, StableRNGs, DataFrames, CSV, LineSearches
gr()

# Set a random seed for reproducible behaviour
rng = StableRNG(1111)
``````

Best regards
Jose

What line is it pointing to?

1 Like

The line is

``````soln_nn = Array(solve(prob_nn, RK4(), u0 = u0, p = NN_params, saveat = tsteps, dt = dt, adaptive=false))
``````

What line in the solve related to `prob_nn`?

I am sorry, I don’t follow your question…however, I did manage to solve the problem by using

``````# Define NN chi-phi
NN_chiphi = Chain((x) -> [cos(x[1]),1/abs(x[3]),1/sqrt(abs(x[3])),sqrt(abs(x[3])),x[3],sqrt(abs(x[3]))^3,x[3]^2,x[4],x[4]^2],
Dense(9, 32, tanh),
Dense(32, 2))
NN_chiphi_params, st_chiphi  = Lux.setup(Xoshiro(), NN_chiphi)

params_chiphi = ComponentArray{Float64}(NN_chiphi_params)
nn_model_chiphi = StatefulLuxLayer(NN_chiphi, st_chiphi)

# Define NN p-e
NN_pe = Chain((x) -> [1/sqrt(abs(x[3]))^3,1/abs(x[3]),1/sqrt(abs(x[3])),sqrt(abs(x[3])),x[3],sqrt(abs(x[3]))^3,x[3]^2,x[4],x[4]^2,x[3]*x[4]],
Dense(10, 32, tanh),
Dense(32, 2))
NN_pe_params, st_pe = Lux.setup(Xoshiro(), NN_pe)

params_pe = ComponentArray{Float64}(NN_pe_params)
nn_model_pe = StatefulLuxLayer(NN_pe, st_pe)

# Parameters of the NN
NN_params = ComponentVector(chiphi = params_chiphi, pe = params_pe)
NN_params_final = ComponentArray{Float64}(NN_params)

# Define PINNs
function AbstractNROrbitModel(u, model_params, t; NN_chiphi=nothing, chiphi=nothing, NN_pe=nothing, pe=nothing)
χ, ϕ, p, e = u
q = model_params[1]
M=1.0

if p <= 0
println("p = ", p)
end

if isnothing(NN_chiphi)
nn_chiphi = [1,1]
else
nn_chiphi = 1 .+ NN_chiphi(u, chiphi)
end

if isnothing(NN_pe)
nn_pe = [0,0]
else
nn_pe = NN_pe(u, pe)
end

numer = (1+e*cos(χ))^2
denom = M*(abs(p)^(3/2))

χ̇ = (numer / denom) * nn_chiphi[1]
ϕ̇ = (numer / denom) * nn_chiphi[2]
ṗ = nn_pe[1]
ė = nn_pe[2]

return [χ̇, ϕ̇, ṗ, ė]
end

function ODE_model(u, NN_params_final, t)
NN_params1 = NN_params_final.chiphi
NN_params2 = NN_params_final.pe
du = AbstractNROrbitModel(u, model_params, t, NN_chiphi = nn_model_chiphi, chiphi = NN_params1, NN_pe = nn_model_pe, pe = NN_params2)
return du
end

prob_nn = ODEProblem(ODE_model, u0, tspan, NN_params_final)
soln_nn = Array(solve(prob_nn, RK4(), u0 = u0, p = NN_params_final, saveat = tsteps, dt = dt, adaptive=false))
waveform_nn_real, waveform_nn_imag = compute_waveform(dt_data, soln_nn, mass_ratio)

plot(tsteps, waveform_nn_real, markershape=:circle, markeralpha = 0.25, linewidth = 1, alpha = 0.5, label="waveform NN (Re)")
``````

but the solution of the PINNs without training is really, really bad, and with training there isn’t an improvement. I am not sure how to continue… is it possible to install Julia version 1.5 on vscode on a MacBook?

Sure, just `juliaup add 1.5` and then run `julia-1.5` in a terminal. You can connect a terminal REPL to VSCode by opening the palette (shift-command-p) and typing “Connect external REPL”.

1 Like

You missed the part of the tutorial where it shows how to choose a better neural network initialization

``````NN_params = NN_params .* 0 + Float64(1e-4) * randn(StableRNG(2031), eltype(NN_params), size(NN_params))
``````

Starting with parameters near zero so that it’s a perturbation is really necessary to help the training process.

1 Like

Thank you for your answer. It actually works and was able to run Julia 1.5.4.

Thank you so much for all your help and time!! Setting up the initial parameters of the network actually improved the initial solution without training to a reasonable guess. As for the training part, the original code uses

``````using DiffEqFlux
## Train with BFGS (gives best results because the Newtonian model seems to give a very good initial guess)
optimization_increments = [collect(40:10:num_optimization_increments-5)..., num_optimization_increments-1,  num_optimization_increments]
for i in optimization_increments
println("optimization increment :: ", i, " of ", num_optimization_increments)
tsteps_increment = tsteps[tsteps .<= tspan[1]+i*(tspan[2]-tspan[1])/num_optimization_increments]
tmp_loss(p) = loss(p,saveat=tsteps_increment)
global NN_params = NN_params + Float64(1e-6)*randn(eltype(NN_params), size(NN_params))
if i < optimization_increments[end-1]
local res = DiffEqFlux.sciml_train(tmp_loss, NN_params, BFGS(initial_stepnorm=0.001, linesearch = LineSearches.BackTracking()), cb=callback, maxiters = 50)
elseif i == optimization_increments[end-1]
local res = DiffEqFlux.sciml_train(tmp_loss, NN_params, BFGS(initial_stepnorm=0.001, linesearch = LineSearches.BackTracking()), cb=callback, maxiters = 100, allow_f_increases=true)
else
local res = DiffEqFlux.sciml_train(tmp_loss, NN_params, BFGS(initial_stepnorm=0.001, linesearch = LineSearches.BackTracking()), cb=callback, maxiters = 250, allow_f_increases=true)
end
global NN_params = res.minimizer
local plt = plot(losses, yaxis=:log, linewidth = 2, xlabel = "Iteration", ylabel = "Objective value", legend = false)
display(plot(plt))
end
``````

however I get the error UndefVarError: `sciml_train` not defined. Following the tutorial I use

``````# Training the PINNs
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentVector{Float64}(NN_params_final))
res2 = Optimization.solve(optprob, BFGS(initial_stepnorm=0.01, linesearch=LineSearches.BackTracking()), callback=callback, maxiters=100)
``````

and the loss functions goes from 0.022163269295360187 to 0.018281040920477987 which does not improve that much.

`sciml_train` was removed a few years ago when Optimization.jl was created as a complete optimization package. In 2020 `sciml_train` was created in DiffEqFlux as a SciML-focused training system that mixed BFGS methods with neural network methods (ADAM), and by 2021 that system became a full-fledged optimization library Optimization.jl. So `sciml_train` was deprecated and the deprecation message I think was removed in late 2023.

I recommend looking at the tutorial that I linked in this section:

which shows how to use the Optimization.jl package for doing the optimization process. You define an OpimizationProblem with the loss function (uses the same definition as you have already), and then call `solve` on the optimization problem using the same algorithms as you had before, so the update should be rather smooth.

2 Likes

Chris, thank you so much for all your help and assistance. I was finally able to reproduce the code and obtain the same results! Thank you again.

1 Like