Neural Networks to describe binary-black-holes trajectories

New to Julia
1

Hello everyone! I am new to Julia and currently working on a project that involves using neural networks to describe geodesics around a black hole. I am encountering an error ‘type Array has no field layer_1,’ and I am unsure how to resolve it. I would appreciate your help to solve the problem. Please find part of my code below.

The neural network architecture

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)

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(9, 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)

NN_params = vcat(params_chiphi, params_pe)
NN_paramsf = ComponentArray{Float64}(NN_pe_params)
l1 = length(params_chiphi)

The system of differential equations

function AbstractNROrbitModel(u, model_params, t; NN_chiphi=nothing, params_chiphi=nothing, NN_pe=nothing, params_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, params_chiphi)
    end

    if isnothing(NN_pe)
        nn_pe = [0,0]
    else
        nn_pe = NN_pe(u, params_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

AbstractNROrbitModel (generic function with 1 method)

function ODE_model(u, NN_paramsf, t)
    NN_params1 = NN_paramsf[1:l1]
    NN_params2 = NN_paramsf[l1+1:end]
    du = AbstractNROrbitModel(u, model_params, t, NN_chiphi=nn_model_chiphi, params_chiphi=NN_params1,
        NN_pe=nn_model_pe, params_pe=NN_params2)
    return du
end

ODE_model (generic function with 1 method)

prob_nn = ODEProblem(ODE_model, u0, tspan, NN_paramsf)
soln_nn = Array(solve(prob_nn, RK4(), u0 = u0, p = NN_paramsf, saveat = tsteps, dt = dt, adaptive=false))

type Array has no field layer_1

If you required more information or the files, let me know.
Best regards

1 Like
2

Wild stab in the dark here since I’m not familiar with these packages. But the output of the solve function may not be convertible to an Array type. Try removing Array(...) from the last line of code and just do:

soln_nn = solve(prob_nn, RK4(), u0 = u0, p = NN_paramsf, saveat = tsteps, dt = dt, adaptive=false)

and see if that works.

3

Thank you for replying to my question. Unfortunately, it does not work. I don’t know if this makes any difference, but the original code defines the neural network using FastChain and FastDense, which does not currently work. See the original code below:

## Define neural network models (seems to help to build up the NNs with a bunch of ansatz functions)
NN_chiphi = FastChain((x, NN_chiphi_params) -> [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],
                FastDense(9, 32, tanh),
                FastDense(32, 2))
NN_chiphi_params = initial_params(NN_chiphi) .* 0
NN_pe = FastChain((x, NN_pe_params) -> [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]],
                FastDense(10, 32, tanh),
                FastDense(32, 2))
NN_pe_params = initial_params(NN_pe) .* 0
NN_params = vcat(NN_chiphi_params,NN_pe_params)
l1 = length(NN_chiphi_params)

Best regards.

4

Seems you are mixing component and regular vectors:

julia> vcat(params_chiphi, params_pe) |> typeof
Vector{Float64} (alias for Array{Float64, 1})

Thus, concatenating two component vectors is no longer a component vector. Accordingly, when you fetch part of the vector as NN_params1 = NN_paramsf[1:l1] in ODE_model this is also just a regular vector and does not have a component layer_1 etc.

Instead, you might want to construct a combined component vector, e.g.,

NN_params = ComponentVector((chiphi = params_chiphi, pe = params_pe))

Then, there is also no need to manually slice the parts in ODE_model as you can just fetch the named components:

 NN_params1 = NN_paramsf.chiphi
 NN_params2 = NN_paramsf.pe
2 Likes
5

Thank you for your reply. You are correct! I am mixing component and regular vectors. Your proposal works!!. This is the correct code for the neural network

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)

NN_pe = 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_pe_params, st_pe = Lux.setup(Xoshiro(), NN_pe)

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

NN_params = ComponentVector(chiphi = params_chiphi, pe = params_pe)
NN_paramsf = ComponentArray{Float64}(NN_params)

and the ODE_model

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

Thank you so much for your help!.
Best regards