Neural Networks to describe binary-black-holes trajectories

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

Best regards

1 Like

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.

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.

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

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