Learning Orbital Dynamics

New to Julia
1

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

2

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

4 Likes
3

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.

4

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
5

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})
For element-wise addition, use broadcasting with dot syntax: scalar .+ array

If you need more information or my file, please let me know. I truly appreciate your help with the problem.

Kind regards
Jose

6

what line is it pointing to?

1 Like
7

The complete error is

MethodError: no method matching +(::Int64, ::Vector{Float64})
For element-wise addition, use broadcasting with dot syntax: scalar .+ array

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:
  [1] _broadcast_getindex_evalf
    @ ./broadcast.jl:709 [inlined]
  [2] _broadcast_getindex
    @ ./broadcast.jl:682 [inlined]
  [3] (::Base.Broadcast.var"#31#32"{Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(+), Tuple{Int64, Tuple{Vector{Float64}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}}})(k::Int64)
    @ Base.Broadcast ./broadcast.jl:1118
  [4] ntuple
    @ ./ntuple.jl:49 [inlined]
  [5] copy
    @ ./broadcast.jl:1118 [inlined]
  [6] materialize(bc::Base.Broadcast.Broadcasted{Base.Broadcast.Style{Tuple}, Nothing, typeof(+), Tuple{Int64, Tuple{Vector{Float64}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}})
...
    @ ~/.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 OrdinaryDiffEq, ModelingToolkit, DataDrivenDiffEq, SciMLSensitivity, DataDrivenSparse
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)

Thank for your interest in the topic. If you need more information, please let me know.

Best regards
Jose

8

What line is it pointing to?

1 Like
9

The line is

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

What line in the solve related to prob_nn?

11

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?

12

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
13

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
14

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

15

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
adtype = Optimization.AutoZygote()
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.

16

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
17

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