Training a Neural ODE to Model Gravitational Waveforms

This code is adapted from Astroinformatics/ScientificMachineLearning

The code has been minimally adapted from Keith et. al. 2021 which originally used Flux.jl

Package Imports

julia

using Lux,
    ComponentArrays,
    LineSearches,
    OrdinaryDiffEqLowOrderRK,
    Optimization,
    OptimizationOptimJL,
    Printf,
    Random,
    SciMLSensitivity
using CairoMakie

Define some Utility Functions

Tip

This section can be skipped. It defines functions to simulate the model, however, from a scientific machine learning perspective, isn't super relevant.

We need a very crude 2-body path. Assume the 1-body motion is a Newtonian 2-body position vector and use Newtonian formulas to get , (e.g. Theoretical Mechanics of Particles and Continua 4.3)

julia

function one2two(path, m₁, m₂)
    M = m₁ + m₂
    r₁ = m₂ / M .* path
    r₂ = -m₁ / M .* path
    return r₁, r₂
end

Next we define a function to perform the change of variables:

julia

@views function soln2orbit(soln, model_params=nothing)
    @assert size(soln, 1) ∈ [2, 4] "size(soln,1) must be either 2 or 4"

    if size(soln, 1) == 2
        χ = soln[1, :]
        ϕ = soln[2, :]

        @assert length(model_params) == 3 "model_params must have length 3 when size(soln,2) = 2"
        p, M, e = model_params
    else
        χ = soln[1, :]
        ϕ = soln[2, :]
        p = soln[3, :]
        e = soln[4, :]
    end

    r = p ./ (1 .+ e .* cos.(χ))
    x = r .* cos.(ϕ)
    y = r .* sin.(ϕ)

    orbit = vcat(x', y')
    return orbit
end

This function uses second-order one-sided difference stencils at the endpoints; see https://doi.org/10.1090/S0025-5718-1988-0935077-0

julia

function d_dt(v::AbstractVector, dt)
    a = -3 / 2 * v[1] + 2 * v[2] - 1 / 2 * v[3]
    b = (v[3:end] .- v[1:(end - 2)]) / 2
    c = 3 / 2 * v[end] - 2 * v[end - 1] + 1 / 2 * v[end - 2]
    return [a; b; c] / dt
end

This function uses second-order one-sided difference stencils at the endpoints; see https://doi.org/10.1090/S0025-5718-1988-0935077-0

julia

function d2_dt2(v::AbstractVector, dt)
    a = 2 * v[1] - 5 * v[2] + 4 * v[3] - v[4]
    b = v[1:(end - 2)] .- 2 * v[2:(end - 1)] .+ v[3:end]
    c = 2 * v[end] - 5 * v[end - 1] + 4 * v[end - 2] - v[end - 3]
    return [a; b; c] / (dt^2)
end

Now we define a function to compute the trace-free moment tensor from the orbit

julia

function orbit2tensor(orbit, component, mass=1)
    x = orbit[1, :]
    y = orbit[2, :]

    Ixx = x .^ 2
    Iyy = y .^ 2
    Ixy = x .* y
    trace = Ixx .+ Iyy

    if component[1] == 1 && component[2] == 1
        tmp = Ixx .- trace ./ 3
    elseif component[1] == 2 && component[2] == 2
        tmp = Iyy .- trace ./ 3
    else
        tmp = Ixy
    end

    return mass .* tmp
end

function h_22_quadrupole_components(dt, orbit, component, mass=1)
    mtensor = orbit2tensor(orbit, component, mass)
    mtensor_ddot = d2_dt2(mtensor, dt)
    return 2 * mtensor_ddot
end

function h_22_quadrupole(dt, orbit, mass=1)
    h11 = h_22_quadrupole_components(dt, orbit, (1, 1), mass)
    h22 = h_22_quadrupole_components(dt, orbit, (2, 2), mass)
    h12 = h_22_quadrupole_components(dt, orbit, (1, 2), mass)
    return h11, h12, h22
end

function h_22_strain_one_body(dt::T, orbit) where {T}
    h11, h12, h22 = h_22_quadrupole(dt, orbit)

    h₊ = h11 - h22
    hₓ = T(2) * h12

    scaling_const = √(T(π) / 5)
    return scaling_const * h₊, -scaling_const * hₓ
end

function h_22_quadrupole_two_body(dt, orbit1, mass1, orbit2, mass2)
    h11_1, h12_1, h22_1 = h_22_quadrupole(dt, orbit1, mass1)
    h11_2, h12_2, h22_2 = h_22_quadrupole(dt, orbit2, mass2)
    h11 = h11_1 + h11_2
    h12 = h12_1 + h12_2
    h22 = h22_1 + h22_2
    return h11, h12, h22
end

function h_22_strain_two_body(dt::T, orbit1, mass1, orbit2, mass2) where {T}
    # compute (2,2) mode strain from orbits of BH1 of mass1 and BH2 of mass 2

    @assert abs(mass1 + mass2 - 1.0) < 1.0e-12 "Masses do not sum to unity"

    h11, h12, h22 = h_22_quadrupole_two_body(dt, orbit1, mass1, orbit2, mass2)

    h₊ = h11 - h22
    hₓ = T(2) * h12

    scaling_const = √(T(π) / 5)
    return scaling_const * h₊, -scaling_const * hₓ
end

function compute_waveform(dt::T, soln, mass_ratio, model_params=nothing) where {T}
    @assert mass_ratio ≤ 1 "mass_ratio must be <= 1"
    @assert mass_ratio ≥ 0 "mass_ratio must be non-negative"

    orbit = soln2orbit(soln, model_params)
    if mass_ratio > 0
        m₂ = inv(T(1) + mass_ratio)
        m₁ = mass_ratio * m₂

        orbit₁, orbit₂ = one2two(orbit, m₁, m₂)
        waveform = h_22_strain_two_body(dt, orbit₁, m₁, orbit₂, m₂)
    else
        waveform = h_22_strain_one_body(dt, orbit)
    end
    return waveform
end

Simulating the True Model

RelativisticOrbitModel defines system of odes which describes motion of point like particle in Schwarzschild background, uses

where, , , and are constants

julia

function RelativisticOrbitModel(u, (p, M, e), t)
    χ, ϕ = u

    numer = (p - 2 - 2 * e * cos(χ)) * (1 + e * cos(χ))^2
    denom = sqrt((p - 2)^2 - 4 * e^2)

    χ̇ = numer * sqrt(p - 6 - 2 * e * cos(χ)) / (M * (p^2) * denom)
    ϕ̇ = numer / (M * (p^(3 / 2)) * denom)

    return [χ̇, ϕ̇]
end

mass_ratio = 0.0         # test particle
u0 = Float64[π, 0.0]     # initial conditions
datasize = 250
tspan = (0.0f0, 6.0f4)   # timespace for GW waveform
tsteps = range(tspan[1], tspan[2]; length=datasize)  # time at each timestep
dt_data = tsteps[2] - tsteps[1]
dt = 100.0
const ode_model_params = [100.0, 1.0, 0.5]; # p, M, e

Let's simulate the true model and plot the results using OrdinaryDiffEq.jl

julia

prob = ODEProblem(RelativisticOrbitModel, u0, tspan, ode_model_params)
soln = Array(solve(prob, RK4(); saveat=tsteps, dt, adaptive=false))
waveform = first(compute_waveform(dt_data, soln, mass_ratio, ode_model_params))

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="Time", ylabel="Waveform")

    l = lines!(ax, tsteps, waveform; linewidth=2, alpha=0.75)
    s = scatter!(ax, tsteps, waveform; marker=:circle, markersize=12, alpha=0.5)

    axislegend(ax, [[l, s]], ["Waveform Data"])

    fig
end

Defining a Neural Network Model

Next, we define the neural network model that takes 1 input (time) and has two outputs. We'll make a function ODE_model that takes the initial conditions, neural network parameters and a time as inputs and returns the derivatives.

It is typically never recommended to use globals but in case you do use them, make sure to mark them as const.

We will deviate from the standard Neural Network initialization and use WeightInitializers.jl,

julia

const nn = Chain(
    Base.Fix1(fast_activation, cos),
    Dense(1 => 32, cos; init_weight=truncated_normal(; std=1.0e-4), init_bias=zeros32),
    Dense(32 => 32, cos; init_weight=truncated_normal(; std=1.0e-4), init_bias=zeros32),
    Dense(32 => 2; init_weight=truncated_normal(; std=1.0e-4), init_bias=zeros32),
)
ps, st = Lux.setup(Random.default_rng(), nn)

((layer_1 = NamedTuple(), layer_2 = (weight = Float32[3.700288f-5; -4.113668f-6; 9.356703f-5; 0.00011446823; 9.597043f-6; 8.4145875f-5; -0.00014828012; -0.000107239; -0.00016230089; 7.52283f-5; -5.9782087f-5; 5.654403f-5; 0.00011483511; 6.110723f-5; -4.5487108f-5; 9.459609f-6; -3.5151254f-5; 2.2677343f-5; 0.0001192291; 0.000116700474; -2.066938f-5; 7.2016934f-5; 2.0718682f-5; -1.3356756f-6; 5.4201687f-6; 0.00013726388; 0.00016450048; 0.0001380765; -5.428903f-5; -6.3552674f-5; -2.0597556f-5; 0.00010115165;;], bias = Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), layer_3 = (weight = Float32[9.881067f-5 -0.0001471758 -4.9051946f-6 1.7691636f-5 6.0974162f-5 -7.258757f-5 -8.5116786f-5 3.5989175f-5 8.63789f-6 -0.00017725273 1.1907103f-5 -8.148734f-5 -8.012505f-5 5.6806275f-5 7.930484f-5 -5.5281045f-5 5.590133f-5 -0.00025607384 -1.05714735f-5 2.4802703f-5 -7.4567746f-5 4.812379f-5 -0.00011039869 3.5762873f-6 -5.9758422f-5 -7.299722f-5 0.00019947015 8.5604f-5 0.00011056237 -0.00010472044 -3.9154133f-5 2.9866499f-6; 0.00025126818 -6.977072f-5 0.00013264081 8.040638f-5 2.7539598f-5 -1.2414686f-5 1.9604428f-5 0.00027974186 1.9923225f-5 2.1915841f-5 9.455386f-5 -0.00014495786 -2.7384729f-5 -6.024254f-5 8.355607f-5 8.777235f-5 6.579751f-5 0.00014912104 -4.6067074f-5 9.91923f-5 2.5276433f-5 -9.2973874f-5 8.835674f-5 6.6557186f-5 -6.7259665f-5 -0.00013745173 -6.9524685f-5 1.7084056f-5 -9.3701645f-5 3.0243664f-5 5.762544f-5 8.530714f-5; -1.4354671f-5 -7.880687f-7 -4.8876085f-5 4.940143f-5 0.00015061958 -7.5654745f-5 2.413094f-5 -0.00020419613 -3.8293114f-5 2.1692109f-5 5.5993292f-5 -2.698027f-5 8.962141f-5 -3.7372363f-5 -5.4021413f-5 -0.00012231032 -0.00011220433 0.00017631084 -7.708734f-6 8.4454194f-5 -0.00011791752 -0.00011385891 -1.3252726f-5 0.00012518138 6.825469f-5 -2.4365383f-5 0.0001277659 -0.00014759133 7.717262f-6 -8.294962f-5 -8.467847f-5 -9.610131f-6; 0.000111765425 9.915711f-5 2.1732845f-5 0.00012877594 2.1372769f-5 -0.00011464579 -0.00016062215 -5.7670808f-5 -0.00010609946 -1.5125922f-5 1.8111541f-5 -0.00015987885 2.9438295f-5 2.7662241f-5 -4.589893f-5 -7.076036f-5 0.00011715815 0.00011495805 -0.0001330699 0.00012140304 -4.9583934f-5 0.00014066325 0.00010260247 0.00018230316 -5.0776493f-5 -4.4903067f-5 -5.356226f-5 -7.2060815f-5 -6.479097f-5 0.00012130829 1.0922779f-5 3.9889346f-5; 6.62586f-5 -2.07597f-5 0.00018795088 0.0001574439 0.00016945135 3.0097106f-5 2.9796296f-5 4.691297f-6 -4.3790347f-5 7.26634f-5 -6.8950816f-5 -2.8408062f-5 -1.1915455f-5 -4.8984002f-5 7.372661f-5 -9.81665f-5 -7.694304f-5 4.5771303f-5 0.00014321112 1.6895054f-5 -4.1667336f-5 8.449894f-5 -0.00015311824 -3.0381218f-5 5.0475515f-5 -9.35214f-5 0.00017292907 6.9561414f-5 6.0109774f-5 -9.452414f-5 0.00010379013 -6.6358654f-5; 9.0135414f-5 7.495609f-5 8.84303f-5 -0.00021417 -3.726954f-5 9.860314f-5 -3.201995f-5 -0.00017851387 0.00017613832 6.3744854f-5 2.8967806f-5 -0.000108117245 -7.202048f-5 -8.0333615f-5 0.00012163102 0.00011211016 -0.00013718926 -0.0001084557 -3.2579046f-5 -4.634402f-5 -8.217444f-5 0.00011119046 -0.00011301282 -0.00020864647 -2.0861902f-5 7.343962f-5 -3.104788f-5 -8.9056535f-5 2.843116f-5 3.7288006f-5 2.0847152f-5 -9.512038f-5; -5.7210123f-5 2.472469f-6 -5.636596f-5 -5.323594f-5 4.655511f-5 -5.047377f-5 -5.950634f-5 2.2490205f-5 -0.00014122264 -0.00020025976 1.2172527f-5 -0.00010213863 0.00011293164 2.836364f-5 2.289591f-5 -2.3735487f-5 -0.00017503618 1.7444332f-5 -7.487197f-5 0.00012158209 -1.1293873f-5 -0.00013279906 7.765833f-5 1.8023076f-5 -5.4343924f-5 0.000112784364 0.00016841374 2.371867f-5 6.058776f-7 6.364905f-5 -2.8960609f-5 -0.00010185781; 0.00014587678 6.852852f-5 7.1711554f-5 -1.406588f-5 -6.3825028f-6 -0.00017217438 6.3564745f-5 -4.994066f-6 0.000105497544 -9.385411f-5 -1.5280979f-5 -5.5771445f-5 -2.5311745f-6 -3.910222f-6 -5.0391955f-5 4.6327576f-5 -0.00016147511 7.294608f-5 -8.6171305f-5 0.00027172716 -5.243094f-5 -0.00014818528 -8.541602f-5 -0.000104099 -2.248921f-5 -0.00012916992 -3.0934876f-5 3.335118f-5 -0.00021195115 -5.8525216f-5 1.1884171f-5 7.801469f-5; -7.216985f-5 -4.1661977f-5 -3.337237f-5 9.7707496f-5 -6.7290064f-5 -2.775493f-5 -9.3701834f-5 -7.05171f-5 0.00021073323 -9.474861f-5 7.307714f-5 -6.242675f-5 6.77171f-6 -3.636654f-5 -0.00011189406 -3.993446f-6 -6.9304675f-5 6.56607f-5 -0.00011586617 0.00019301217 -0.00014648869 4.5490022f-5 -4.91699f-5 -7.0112364f-5 6.334912f-5 0.00010215128 -2.7660024f-5 -0.00012177362 -8.631069f-5 3.071492f-6 0.000111581605 -5.7248028f-5; -7.197442f-5 -2.561531f-5 0.00012431185 0.00014837792 -2.9673627f-5 -7.686588f-6 0.00013004306 -6.378278f-5 -0.000121447294 3.973344f-5 0.000102279184 -6.8030786f-5 -2.0415026f-5 -0.0001515569 -0.00010703435 -2.1422022f-5 -0.00011228455 0.00010242451 -0.00017377679 0.0001050216 -1.8106271f-5 5.8331916f-5 -1.0320825f-5 -6.969278f-5 -0.00026298148 -0.000104946885 1.7360995f-5 4.9729417f-5 -0.00016601068 -0.0001563121 4.2810218f-5 -8.135927f-5; -0.00012529454 2.6343138f-5 -2.8169608f-5 2.5189056f-5 -9.33244f-5 -0.00019270276 1.5784164f-5 -1.5488467f-5 -2.1613357f-6 -3.408934f-5 3.488155f-5 -8.007969f-5 -0.00010065668 -4.9124406f-5 -3.6347552f-5 0.00010916393 -0.00014880323 0.00014311084 -2.8604338f-5 -3.9859082f-5 -3.0650347f-6 0.00016772 7.504759f-5 8.0687925f-5 7.446156f-5 3.119854f-5 -5.5374574f-5 6.732037f-5 2.7608627f-5 2.5482068f-5 3.341979f-5 -5.5871107f-5; -3.0249985f-5 5.4107597f-5 4.924289f-5 2.1291151f-5 -0.00013100094 -4.5762226f-5 0.00028993713 1.2913129f-6 8.468675f-5 1.8190402f-5 -1.3899782f-5 4.4726785f-5 -8.966285f-6 -7.829256f-5 6.605625f-5 6.8889385f-5 -2.2961818f-5 5.7881025f-5 -1.454203f-5 9.552061f-5 -3.2428467f-5 3.9281f-5 1.2706576f-5 0.00010839479 2.6276231f-5 -2.020767f-5 0.00017467483 0.00011590086 1.567353f-5 2.6606253f-5 -0.00015691046 -4.5739f-5; -0.00032256875 6.211332f-5 -0.00013475149 -9.4021925f-6 3.7114947f-5 0.0001516498 -6.858148f-6 0.00011554444 9.9445264f-5 -8.7385466f-5 1.9801044f-5 -0.00012938994 2.2372786f-5 0.0002563744 0.000114281094 0.00018587602 1.8123175f-5 -0.00013735854 0.00014017563 -1.0507487f-5 3.5795456f-5 1.950906f-5 4.573966f-5 -0.00015903937 1.7563625f-5 1.3170727f-6 -8.471229f-5 0.00021979366 1.29568525f-5 9.924875f-5 0.00012946244 -7.912245f-5; 1.03538845f-7 6.303624f-5 -6.796391f-5 -4.964159f-5 0.00011433407 0.000174112 -9.0492926f-5 5.5459295f-5 8.268506f-5 0.00016724937 -5.258283f-5 -0.00013463746 -4.4898694f-5 9.419925f-5 0.00013007916 -0.00012242836 -3.6759415f-5 -1.0061261f-5 1.52420125f-5 -3.9802748f-5 7.379727f-5 -4.3966334f-5 7.2054504f-6 0.00013432321 0.00011799234 1.21626f-6 -2.08583f-5 -7.956517f-5 0.00015924664 0.00012652315 0.000107833046 9.213831f-5; 6.1845734f-5 4.9809863f-5 0.00012626912 4.4230044f-5 2.6763444f-5 -0.0001462021 7.486341f-5 0.00014361388 -4.0635307f-5 6.514654f-5 1.7315024f-5 3.9892857f-5 -1.4074708f-5 5.7213696f-5 0.00012721252 -8.5694745f-7 -0.00019068192 -4.9061644f-5 7.143295f-5 -1.766488f-5 -8.677715f-5 0.00010102872 4.7510148f-5 0.00018247339 5.8780533f-5 0.00012812302 -1.2909259f-5 3.6358848f-5 0.00016419846 -4.177821f-5 -2.381302f-5 0.0001167209; -5.089125f-5 -0.00010678429 8.97534f-5 6.752199f-5 -0.00014772789 0.00018824953 0.000105443694 -4.4353106f-5 -6.169265f-5 -3.6763086f-5 -0.00013076671 -2.9779549f-5 -9.307892f-5 9.9317236f-5 0.00014428269 2.393499f-6 9.45182f-6 -3.5089677f-5 -0.00014437147 3.0455218f-5 -9.116882f-5 -7.347149f-5 0.00011113723 -4.8613158f-5 7.725848f-5 -9.229881f-5 4.2161813f-5 -6.7039275f-7 0.00010392604 2.0530977f-5 -9.210392f-6 4.5394234f-5; -0.00015743745 7.149998f-5 -7.537442f-5 1.591947f-5 0.0001109207 0.00016128292 1.0165206f-6 -0.00025247602 5.1640076f-5 -5.6203944f-6 5.021376f-5 -1.8233877f-5 0.0002110788 -0.00010517637 -0.00018276392 7.0529626f-5 -2.718574f-5 -0.00011870454 -3.033389f-5 -0.00013936046 0.00015240168 -3.8954608f-5 1.414975f-5 4.797772f-5 8.232333f-5 2.5100482f-5 5.5314016f-5 0.00013990531 9.581742f-5 0.00022401076 -7.42183f-5 0.00020260504; -1.9456587f-5 7.224755f-5 4.8723676f-5 5.2831063f-5 7.381585f-5 3.7187972f-5 -9.709522f-5 -2.3027694f-6 -9.347537f-5 -8.622485f-5 0.000104521525 5.3364896f-5 -2.7092346f-5 1.17112595f-5 -9.240865f-5 0.000105808474 -5.513361f-5 -0.00018747193 -0.00011270921 -1.932986f-5 -2.2953995f-5 -9.807545f-6 -0.00011079728 3.0645275f-5 -3.855851f-5 -0.00025350173 -7.75084f-5 0.00012790813 -2.7779543f-5 -0.00014507335 -3.012833f-5 0.00019734379; -0.00011845354 -7.508007f-5 -2.7992435f-5 -7.800473f-5 -0.00014423122 0.00012155474 -1.1386977f-6 -5.9983104f-5 -8.831918f-5 0.00018718802 -0.00012409297 4.2396423f-5 -2.6211854f-5 1.922595f-5 -0.00016296198 2.9059574f-5 -2.1883185f-5 3.6240664f-5 1.8259168f-5 7.225779f-5 0.00018740896 8.377947f-5 0.00029425125 0.00014808557 7.0638846f-5 0.00012296859 -4.995293f-6 0.00016424178 -3.1881274f-5 6.159489f-5 2.485919f-5 3.704355f-5; 3.920016f-5 -0.00016207865 -7.201213f-5 0.0001746343 -3.4656474f-5 -1.1319661f-5 0.00010219991 -0.00013184088 -0.0002314381 0.00011406893 2.6409707f-5 8.774785f-5 -3.6012654f-5 -0.00014698032 5.0937924f-5 -5.353267f-5 8.891027f-5 0.0001600207 -1.649121f-6 2.5746376f-5 4.3463686f-5 3.458344f-5 4.256769f-5 0.00017646998 8.841286f-5 5.4472377f-5 -0.0001482924 0.00013040897 0.00031332782 0.00016698691 9.099141f-7 -0.00010058158; -5.4744753f-5 -0.0001607139 -0.00010237587 5.159719f-5 -0.00019383014 0.00012196537 9.485992f-6 -6.3303356f-5 -0.00012120556 0.00012596394 -0.00024071184 9.980426f-5 0.00012130102 0.00017801748 -5.508867f-5 2.9157462f-5 -2.5373336f-5 9.272182f-5 -2.841186f-5 -1.5992677f-5 -3.4458957f-5 -1.0007353f-5 -4.0462328f-5 0.00020930717 -3.0762174f-5 -7.341282f-5 -0.00014254566 -0.00014459061 2.164509f-5 0.00020332195 -0.00025072673 -2.1918815f-6; 0.00014951786 3.745084f-5 -1.844009f-5 3.7200447f-5 -7.663113f-5 0.00011831125 5.8594444f-5 -0.00017070145 -3.7994185f-5 -0.00010313767 -0.000141965 -6.115176f-5 -0.00016608629 -3.3147462f-5 1.9924779f-5 -0.00010888381 -3.078276f-5 3.4677054f-5 8.764434f-5 6.801612f-5 -0.00014626973 -0.000118249656 -9.275998f-5 0.00015544957 -2.9180175f-5 -4.300968f-5 0.00013421112 0.00019334022 -0.0001602313 -2.1692782f-5 -7.482682f-5 8.6852604f-5; -6.001022f-5 3.3260578f-5 -0.0001847593 -4.074499f-5 1.6895752f-5 -1.1473168f-5 -5.842434f-5 -8.7285785f-5 2.4336328f-5 1.9660263f-5 8.383561f-6 -5.4761567f-5 6.021864f-5 0.00013729405 -0.00012296604 -8.11427f-5 -4.078438f-5 -1.5620805f-6 -5.7717447f-5 1.7131313f-5 -0.000217294 0.00012708755 -0.00013274164 -3.5154637f-5 -6.466461f-5 -0.0003199532 -0.00016136355 1.5410642f-5 9.542301f-5 -6.0087837f-5 1.1501878f-6 -4.0365805f-5; -5.0003346f-5 2.4938447f-5 -5.876273f-5 5.2162806f-5 -1.2775139f-5 6.559545f-5 4.0542054f-5 0.00016125978 9.564049f-5 -5.1272567f-5 0.00016641212 -5.5442728f-5 3.9010203f-5 1.7205606f-5 -9.0796115f-5 9.477055f-5 -1.3109616f-6 -0.00013678164 0.00015702577 0.00011861976 3.1128977f-6 0.00014935486 5.890203f-5 6.3639372f-6 1.4839328f-5 6.984175f-5 0.00010606637 -0.00010060597 -4.8020424f-6 7.03042f-5 -0.00012530865 7.9304f-6; -0.00012761286 3.829487f-5 -9.2319235f-5 -0.000121199424 -2.954021f-5 -0.00019395692 6.683778f-5 1.5977186f-6 4.3477517f-5 -6.595504f-5 2.619378f-5 0.00016182628 -2.537815f-5 -0.00019138245 -4.9661696f-5 -6.7219225f-5 -4.2723466f-5 -0.00010628128 -0.00015566261 -1.2007005f-5 -0.00019421277 -0.00014503713 0.00013727852 5.068265f-6 -1.0222597f-5 3.0774725f-5 -1.3026869f-5 -6.387822f-5 0.00016850476 2.280247f-5 -1.5313373f-5 9.710236f-5; 5.907802f-5 4.4920525f-5 6.868866f-5 8.828209f-5 -9.162904f-5 9.134826f-5 -1.1148319f-5 1.632717f-7 7.119294f-5 0.000119572025 7.293182f-5 0.00021981925 -3.420978f-5 0.00010437276 -3.060414f-5 -6.859265f-5 -7.983431f-5 0.0001239 0.00012049327 0.0001078821 -7.297228f-5 0.00010022674 7.897796f-5 5.1960626f-5 -0.00013449987 0.000121643796 -5.5495657f-5 -0.0001615421 2.8792254f-5 -9.236468f-5 2.7272241f-5 7.998117f-5; 4.769264f-5 -2.9264329f-5 -5.692907f-5 -4.38092f-5 -1.2298325f-5 0.00023629899 0.00013189278 2.1412056f-5 0.00011569873 3.9089064f-6 0.00014313794 5.0481623f-5 2.086636f-5 0.0001275838 -0.00015307289 -0.00014465813 9.90181f-5 6.482552f-5 -1.1288205f-5 0.00016257889 -0.000102422346 9.689416f-6 -2.6328802f-5 3.245279f-6 -0.00016282545 5.3307125f-5 0.00014412723 -8.2670675f-5 -9.5295065f-5 6.565833f-5 0.00017639475 6.804423f-5; -4.5855893f-5 1.8586189f-8 -0.00011045143 3.80849f-5 1.997963f-5 -0.0001826499 -7.9930614f-5 -5.161819f-5 3.818113f-5 -9.1628564f-5 -1.5615798f-5 -2.993341f-5 -2.3042287f-5 5.6679284f-5 -0.00017848647 -3.2369924f-5 -0.0001409952 -0.00012494951 -2.2995213f-5 0.00012248979 1.0509755f-5 -0.00015084492 0.0001396451 -4.921549f-5 -4.292621f-6 0.00012433098 -1.0909147f-5 0.00016000797 8.1791084f-5 -0.00011309016 1.0091114f-5 -0.00013290919; 2.170866f-5 0.00014476122 -0.00013535355 0.00029736874 2.4500705f-5 0.000115941606 0.00010981927 4.6362464f-5 8.9741625f-5 1.6801036f-5 -6.788401f-5 9.4404595f-6 6.4294663f-6 5.7917383f-5 3.2005228f-5 7.024924f-5 6.28105f-5 3.285441f-5 -0.00011498297 0.00018928373 2.0208328f-5 5.5081702f-5 7.343542f-5 6.195593f-5 -7.110745f-5 3.6790403f-5 -0.000105282175 -0.0001639142 5.8785026f-5 -2.25028f-5 -6.0274073f-5 3.552366f-5; -0.00028695766 -7.831894f-5 -3.1225223f-5 -6.820796f-5 -4.808259f-5 6.3649815f-5 0.00021295634 -8.1131286f-5 5.090619f-5 -1.9017545f-5 -3.8152273f-5 0.00021346298 -5.109224f-5 0.0002477126 -9.6699994f-5 -0.00018816926 -0.00010956748 -3.884591f-5 -9.566459f-5 -3.680908f-5 -2.168836f-5 -0.00017768618 8.2512044f-5 -3.4403103f-5 0.00016120107 -0.00010045437 -7.6490986f-5 3.4885163f-6 9.679702f-5 -3.3126693f-5 0.00011120526 -3.6731076f-6; 0.00024580362 -0.00012994492 -6.4185726f-5 0.00012444264 -5.932479f-5 0.00016992452 0.00019625902 1.1661377f-5 0.0001622589 4.5949095f-5 -0.00019550184 2.4711882f-5 0.00021814129 -3.399797f-5 0.00010822056 -0.0001766398 -0.00020148257 -2.0116202f-5 -9.468069f-5 2.1678727f-5 0.0001673984 2.5920506f-5 5.469778f-5 -1.6477288f-5 -9.058147f-5 -1.942265f-5 -6.9692265f-5 -7.6010176f-5 0.00022230143 -0.0001509336 9.820281f-5 3.0050835f-5; -9.464777f-5 4.8808175f-5 -5.7763936f-7 5.16259f-6 -0.00013001442 -2.2826145f-5 3.9975523f-5 0.00023709405 6.103565f-5 0.0001481177 4.373834f-5 8.541154f-5 -7.105519f-5 -1.2451289f-5 3.7429877f-6 1.1094695f-5 5.9660662f-5 4.152157f-5 0.0001282051 -0.0001918384 0.000107954686 1.3231187f-5 -0.00014040928 0.0001065315 -6.700831f-5 5.714649f-5 -2.237536f-5 -1.4552344f-5 -0.00016780535 -6.958537f-5 -0.00015830786 8.042422f-6], bias = Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), layer_4 = (weight = Float32[-0.00014011745 -0.00011215393 -7.086878f-5 -1.1604832f-5 0.00011111224 -0.00016877231 6.687057f-5 3.93343f-5 -0.00012315334 -0.000121348814 0.00012097552 1.0058994f-5 8.908395f-5 -5.6189558f-5 -0.00010959909 -1.5718922f-5 4.2901407f-5 -0.00012098371 0.00010315528 0.0001606372 2.992205f-5 3.350543f-5 -9.727116f-5 4.0764036f-5 7.3267656f-5 -9.829378f-5 -6.151407f-5 0.00020967897 -7.221835f-5 -7.940053f-5 -2.4017345f-5 6.3448947f-6; 8.185206f-5 2.0311636f-5 -4.1457835f-5 2.5407666f-5 0.0002924119 -9.2269256f-5 -1.0931032f-5 -1.6485224f-5 0.00012732424 -4.4107463f-5 -4.415277f-5 -6.831065f-5 3.0010917f-5 4.0175124f-5 0.00029966483 -1.063542f-5 9.370819f-5 -4.9487094f-6 -9.842093f-5 3.8527098f-5 -0.00015987779 0.00017303074 -3.0254036f-5 -1.3368402f-5 4.4196177f-6 -1.6532751f-5 -2.4901752f-5 -8.5863045f-5 2.162175f-5 6.2652066f-6 -9.574145f-5 0.00013992997], bias = Float32[0.0, 0.0])), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple()))

Similar to most deep learning frameworks, Lux defaults to using Float32. However, in this case we need Float64

julia

const params = ComponentArray(f64(ps))

const nn_model = StatefulLuxLayer(nn, nothing, st)

StatefulLuxLayer{Val{true}()}(
    Chain(
        layer_1 = WrappedFunction(Base.Fix1{typeof(LuxLib.API.fast_activation), typeof(cos)}(LuxLib.API.fast_activation, cos)),
        layer_2 = Dense(1 => 32, cos),            # 64 parameters
        layer_3 = Dense(32 => 32, cos),           # 1_056 parameters
        layer_4 = Dense(32 => 2),                 # 66 parameters
    ),
)         # Total: 1_186 parameters,
          #        plus 0 states.

Now we define a system of odes which describes motion of point like particle with Newtonian physics, uses

where, , , and are constants

julia

function ODE_model(u, nn_params, t)
    χ, ϕ = u
    p, M, e = ode_model_params

    # In this example we know that `st` is am empty NamedTuple hence we can safely ignore
    # it, however, in general, we should use `st` to store the state of the neural network.
    y = 1 .+ nn_model([first(u)], nn_params)

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

    χ̇ = (numer / denom) * y[1]
    ϕ̇ = (numer / denom) * y[2]

    return [χ̇, ϕ̇]
end

Let us now simulate the neural network model and plot the results. We'll use the untrained neural network parameters to simulate the model.

julia

prob_nn = ODEProblem(ODE_model, u0, tspan, params)
soln_nn = Array(solve(prob_nn, RK4(); u0, p=params, saveat=tsteps, dt, adaptive=false))
waveform_nn = first(compute_waveform(dt_data, soln_nn, mass_ratio, ode_model_params))

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="Time", ylabel="Waveform")

    l1 = lines!(ax, tsteps, waveform; linewidth=2, alpha=0.75)
    s1 = scatter!(
        ax, tsteps, waveform; marker=:circle, markersize=12, alpha=0.5, strokewidth=2
    )

    l2 = lines!(ax, tsteps, waveform_nn; linewidth=2, alpha=0.75)
    s2 = scatter!(
        ax, tsteps, waveform_nn; marker=:circle, markersize=12, alpha=0.5, strokewidth=2
    )

    axislegend(
        ax,
        [[l1, s1], [l2, s2]],
        ["Waveform Data", "Waveform Neural Net (Untrained)"];
        position=:lb,
    )

    fig
end

Setting Up for Training the Neural Network

Next, we define the objective (loss) function to be minimized when training the neural differential equations.

julia

const mseloss = MSELoss()

function loss(θ)
    pred = Array(solve(prob_nn, RK4(); u0, p=θ, saveat=tsteps, dt, adaptive=false))
    pred_waveform = first(compute_waveform(dt_data, pred, mass_ratio, ode_model_params))
    return mseloss(pred_waveform, waveform)
end

Warmup the loss function

julia

loss(params)

0.0007146129702641741

Now let us define a callback function to store the loss over time

julia

const losses = Float64[]

function callback(θ, l)
    push!(losses, l)
    @printf "Training \t Iteration: %5d \t Loss: %.10f\n" θ.iter l
    return false
end

Training the Neural Network

Training uses the BFGS optimizers. This seems to give good results because the Newtonian model seems to give a very good initial guess

julia

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, params)
res = Optimization.solve(
    optprob,
    BFGS(; initial_stepnorm=0.01, linesearch=LineSearches.BackTracking());
    callback,
    maxiters=1000,
)

retcode: Success
u: ComponentVector{Float64}(layer_1 = Float64[], layer_2 = (weight = [3.7002879253077236e-5; -4.113668182979994e-6; 9.356703230871707e-5; 0.00011446823191341839; 9.597042662787482e-6; 8.414587500715396e-5; -0.0001482801162635547; -0.00010723900049898563; -0.00016230088658620431; 7.522830128419142e-5; -5.978208719169223e-5; 5.654403139486686e-5; 0.00011483510752436315; 6.110723188609757e-5; -4.5487107854489615e-5; 9.459608918400417e-6; -3.51512535416539e-5; 2.2677342712970073e-5; 0.00011922910198309118; 0.0001167004738817521; -2.06693803192545e-5; 7.201693369994235e-5; 2.0718682208061568e-5; -1.3356756198816952e-6; 5.420168690757882e-6; 0.00013726387987813784; 0.0001645004813325952; 0.00013807650248057313; -5.428903023122473e-5; -6.355267396418224e-5; -2.059755570370211e-5; 0.00010115165059675127;;], bias = [5.1029355324289433e-17, -1.3403819659999655e-18, 7.137991163097836e-17, 2.358782506783353e-16, 2.6203340116467114e-18, 2.4929496950311076e-16, -2.4495257837866566e-16, -2.9145300442035613e-16, -1.1875329438434076e-16, 1.9437652936949358e-16, -1.5532077612244388e-17, 1.3223200447531018e-17, 5.5031891313922073e-17, 3.140882085598784e-17, -8.745650792917747e-17, 5.163752015351205e-18, -5.0181646254158673e-17, 2.830400851554461e-17, 3.0064339921303785e-16, 4.6343364402108684e-17, -5.568403888364979e-17, 1.3824350868292455e-16, 4.128553152758892e-17, -3.1667193581841404e-18, 4.243084866834942e-18, 3.434127195292828e-16, 6.423215165129219e-17, -3.427519505432735e-17, -8.421770238973515e-17, -1.2484790023320713e-16, -1.5707912178788563e-17, 1.3816095735258424e-16]), layer_3 = (weight = [9.88087059555701e-5 -0.0001471777640075306 -4.9071580873222125e-6 1.768967297144376e-5 6.097219841071647e-5 -7.25895359904343e-5 -8.51187495539639e-5 3.5987211161824333e-5 8.635926063713242e-6 -0.00017725469558749682 1.190513933002017e-5 -8.148930544516561e-5 -8.012701334914045e-5 5.680431161058612e-5 7.930287495005312e-5 -5.528300869296062e-5 5.589936802029926e-5 -0.0002560757992581344 -1.057343695387497e-5 2.4800739196196085e-5 -7.456970968301142e-5 4.812182771606939e-5 -0.00011040065044321786 3.574323840119224e-6 -5.976038562617026e-5 -7.299918695602058e-5 0.00019946818167560832 8.560203525454327e-5 0.00011056040341250093 -0.00010472240580961544 -3.9156096023681725e-5 2.9846863741615634e-6; 0.0002512717694396369 -6.976713463633794e-5 0.00013264439642691426 8.040996929659378e-5 2.7543185052065956e-5 -1.2411099095767255e-5 1.9608014764075646e-5 0.0002797454437593096 1.9926812665925438e-5 2.191942824604259e-5 9.455745042098732e-5 -0.00014495426855252964 -2.7381141780355486e-5 -6.023895163641284e-5 8.355965858821047e-5 8.777593778243739e-5 6.580109787203688e-5 0.00014912462959794725 -4.606348717905841e-5 9.91958879519328e-5 2.5280020264046655e-5 -9.297028715925433e-5 8.8360328146093e-5 6.65607733305776e-5 -6.725607797240381e-5 -0.00013744814329519738 -6.952109812964575e-5 1.7087643055409344e-5 -9.36980575436528e-5 3.024725104362696e-5 5.7629026042133964e-5 8.531072692741844e-5; -1.4356013379205266e-5 -7.894107752187318e-7 -4.8877427288007764e-5 4.9400089084030295e-5 0.00015061824193349403 -7.56560872459352e-5 2.41295987153224e-5 -0.00020419746910096292 -3.8294455988700184e-5 2.1690766844638368e-5 5.5991950249387654e-5 -2.6981611553258644e-5 8.962006937605259e-5 -3.737370538050789e-5 -5.402275532009632e-5 -0.0001223116593236633 -0.00011220567433969708 0.0001763095027500236 -7.710076387354908e-6 8.445285162548485e-5 -0.00011791886539771846 -0.00011386024892903353 -1.3254067779548191e-5 0.0001251800389227859 6.825334672708354e-5 -2.4366725155173925e-5 0.00012776456093076214 -0.00014759267254606636 7.715919638082352e-6 -8.295096234991686e-5 -8.467980991458213e-5 -9.611473286840732e-6; 0.00011176622130759078 9.915790353287063e-5 2.1733641842983815e-5 0.0001287767335454905 2.1373565614596524e-5 -0.00011464499046243856 -0.00016062135189688346 -5.767001137403082e-5 -0.00010609866520073887 -1.512512524615832e-5 1.8112337711755735e-5 -0.00015987805461892884 2.9439091959173732e-5 2.76630379815116e-5 -4.5898134048655626e-5 -7.075956462332928e-5 0.00011715895008845277 0.00011495884602507617 -0.00013306910484945714 0.00012140383834801318 -4.9583137369807386e-5 0.00014066404757639314 0.00010260326591404197 0.00018230396010065507 -5.0775695926602404e-5 -4.490227009202785e-5 -5.356146364622461e-5 -7.206001835559e-5 -6.479017254579836e-5 0.00012130908355200613 1.0923576083116196e-5 3.9890142870971705e-5; 6.626113323070976e-5 -2.0757169789447897e-5 0.00018795340880094363 0.00015744642447022627 0.00016945387911326297 3.0099636621139297e-5 2.97988267054839e-5 4.6938275379679775e-6 -4.3787816633305174e-5 7.266593316378589e-5 -6.894828542792482e-5 -2.840553092321128e-5 -1.1912924076778414e-5 -4.898147157584574e-5 7.372913974611901e-5 -9.816396918105852e-5 -7.694051013751939e-5 4.577383375070065e-5 0.00014321364729473187 1.6897584334191447e-5 -4.166480497880825e-5 8.450147148030433e-5 -0.0001531157045457976 -3.037868695466264e-5 5.047804593849043e-5 -9.351886687294213e-5 0.00017293159767792254 6.956394505028964e-5 6.0112305124474626e-5 -9.452161024748632e-5 0.00010379266173620724 -6.635612366447082e-5; 9.013332965185971e-5 7.495400859842087e-5 8.84282162609384e-5 -0.0002141720797204509 -3.727162402139026e-5 9.860105274346217e-5 -3.202203515465739e-5 -0.00017851595839384155 0.00017613623276920915 6.374276927307643e-5 2.8965721869452595e-5 -0.00010811932952275273 -7.202256165110537e-5 -8.033569974723938e-5 0.0001216289347234011 0.00011210807497480899 -0.00013719134221684096 -0.0001084577852430935 -3.258113065167291e-5 -4.634610377034174e-5 -8.217652429958765e-5 0.00011118837210450683 -0.00011301490667175461 -0.00020864855104956768 -2.08639867174624e-5 7.343753896057027e-5 -3.104996357942026e-5 -8.905861977618751e-5 2.842907633966046e-5 3.7285921282445134e-5 2.084506741301663e-5 -9.512246105951486e-5; -5.7211943262062735e-5 2.4706493019951952e-6 -5.636777938370325e-5 -5.323776064350678e-5 4.65532895174878e-5 -5.047558885443785e-5 -5.9508159087532014e-5 2.248838530768604e-5 -0.00014122446088029637 -0.00020026158096179278 1.2170707246425685e-5 -0.00010214044674675968 0.00011292981720153768 2.836182023720791e-5 2.2894090885257993e-5 -2.3737306823295127e-5 -0.00017503800414879047 1.7442512340222212e-5 -7.487378831164322e-5 0.00012158026869265916 -1.1295693060946506e-5 -0.00013280088102355538 7.765651310883244e-5 1.8021256560769292e-5 -5.434574346900229e-5 0.00011278254454347233 0.0001684119189446553 2.3716850715311582e-5 6.04057798411541e-7 6.364723152337588e-5 -2.896242846354595e-5 -0.00010185963116263614; 0.00014587470276207943 6.85264420874529e-5 7.17094770928084e-5 -1.406795774047827e-5 -6.384580151213236e-6 -0.0001721764530109082 6.356266738774983e-5 -4.996143265742304e-6 0.00010549546630981466 -9.385618988264714e-5 -1.52830563044738e-5 -5.5773522139792946e-5 -2.5332518495972015e-6 -3.912299339139398e-6 -5.039403196381076e-5 4.6325498156214224e-5 -0.00016147718644308586 7.294400337311663e-5 -8.617338254654571e-5 0.0002717250789694691 -5.2433017221780165e-5 -0.00014818735882247259 -8.541809449045702e-5 -0.00010410108017342908 -2.2491287486197366e-5 -0.0001291719928437606 -3.0936953512270514e-5 3.334910423772011e-5 -0.00021195322562082283 -5.852729383579927e-5 1.188209342696922e-5 7.801261007014605e-5; -7.217165489221564e-5 -4.1663779256694936e-5 -3.337417160748643e-5 9.770569363370652e-5 -6.729186657123915e-5 -2.7756732151669024e-5 -9.370363636645784e-5 -7.051889929227509e-5 0.00021073142329459038 -9.475041383650472e-5 7.307534019689563e-5 -6.242855101829861e-5 6.769907713696641e-6 -3.636834094306887e-5 -0.00011189586598432223 -3.995248562049299e-6 -6.930647737119927e-5 6.565889839309577e-5 -0.00011586797131696534 0.00019301036286395188 -0.0001464904905786428 4.548821948462749e-5 -4.9171701675575906e-5 -7.011416687403527e-5 6.334731938311178e-5 0.00010214947747061479 -2.7661826379527313e-5 -0.00012177542316732465 -8.631249306729543e-5 3.069689683046815e-6 0.00011157980247860148 -5.72498299787385e-5; -7.197787389331746e-5 -2.5618764346562617e-5 0.00012430839225245283 0.0001483744660724526 -2.9677080681993835e-5 -7.690041719151982e-6 0.00013003960585749897 -6.378623701337776e-5 -0.00012145074769092978 3.972998703075801e-5 0.00010227573030688795 -6.803423938280948e-5 -2.0418479187612902e-5 -0.00015156035559657014 -0.00010703780325302094 -2.142547535940933e-5 -0.00011228800330017957 0.00010242105300831264 -0.0001737802424834457 0.00010501814787104914 -1.8109725068187178e-5 5.832846246604052e-5 -1.0324278526512275e-5 -6.969623544891768e-5 -0.0002629849380248732 -0.00010495033828945061 1.7357541502619278e-5 4.9725963173591105e-5 -0.00016601413270863778 -0.0001563155576953194 4.280676395110429e-5 -8.13627241662164e-5; -0.00012529512843081079 2.6342545776638885e-5 -2.8170200142414753e-5 2.5188464302699458e-5 -9.332499267963064e-5 -0.00019270335416273382 1.5783572404544063e-5 -1.548905912692279e-5 -2.1619274442477847e-6 -3.40899301712514e-5 3.488095661544936e-5 -8.0080283272185e-5 -0.00010065727348170298 -4.9124998089030856e-5 -3.634814367011455e-5 0.00010916333611624944 -0.0001488038202418136 0.0001431102476877969 -2.860492951491416e-5 -3.985967360982717e-5 -3.065626391038579e-6 0.0001677194135360306 7.504699659569823e-5 8.068733349012837e-5 7.446096914157899e-5 3.1197946906686355e-5 -5.537516573291063e-5 6.731977681432907e-5 2.7608034894580332e-5 2.5481476210767924e-5 3.341919854086528e-5 -5.587169890837854e-5; -3.0247032147223184e-5 5.411054957986839e-5 4.924584431901705e-5 2.1294103542824247e-5 -0.00013099798666333759 -4.575927373625047e-5 0.00028994008466578825 1.2942654833778344e-6 8.468970049615342e-5 1.8193354161129856e-5 -1.389682943371433e-5 4.472973746343243e-5 -8.96333272212242e-6 -7.828961092999528e-5 6.605920357143342e-5 6.889233777511889e-5 -2.2958865476163036e-5 5.788397761390526e-5 -1.4539077302833701e-5 9.552356500389655e-5 -3.2425513883559134e-5 3.928395239964598e-5 1.2709528571953898e-5 0.00010839773899389463 2.6279184065438157e-5 -2.020471724936634e-5 0.00017467778384196828 0.00011590381331952371 1.5676482188075157e-5 2.6609205131909084e-5 -0.00015690751165637752 -4.5736048879549006e-5; -0.0003225662680921261 6.211580830895216e-5 -0.0001347490056065087 -9.399706410125432e-6 3.711743329804259e-5 0.00015165228751756494 -6.855661918617738e-6 0.00011554692673997841 9.944775002111314e-5 -8.738297992350177e-5 1.980352999853027e-5 -0.00012938745430401909 2.2375271786937816e-5 0.0002563768771364219 0.00011428358039156244 0.00018587850855001118 1.8125661430892174e-5 -0.00013735605397924018 0.0001401781169119072 -1.0505000769275259e-5 3.579794202268897e-5 1.9511545819471794e-5 4.5742146089103356e-5 -0.00015903687971840735 1.7566111186531974e-5 1.3195588093198606e-6 -8.470980037202551e-5 0.00021979614166810395 1.2959338623109405e-5 9.925123368192948e-5 0.00012946492416121837 -7.911996589938189e-5; 1.0701340032031527e-7 6.303971560002465e-5 -6.79604378875334e-5 -4.963811414754398e-5 0.00011433754508645426 0.000174115481344334 -9.048945110963173e-5 5.546276980922358e-5 8.26885337245464e-5 0.0001672528417783835 -5.257935452944838e-5 -0.00013463398835127438 -4.489521933109887e-5 9.420272209710184e-5 0.00013008263732801123 -0.00012242488781892544 -3.6755940242148696e-5 -1.0057786026738003e-5 1.5245487078483765e-5 -3.979927322334115e-5 7.380074777203729e-5 -4.396285993251438e-5 7.208924956938696e-6 0.0001343266888524515 0.000117995816539332 1.2197345439707938e-6 -2.0854826003518805e-5 -7.956169764484399e-5 0.0001592501178619528 0.00012652662945939868 0.0001078365206139267 9.214178527500424e-5; 6.184983268168083e-5 4.981396178619446e-5 0.00012627322049365577 4.423414266920269e-5 2.6767542485240495e-5 -0.00014619800439550113 7.48675103836629e-5 0.0001436179829483928 -4.063120813091575e-5 6.515063633696791e-5 1.7319122170651906e-5 3.989695524933763e-5 -1.4070609532163196e-5 5.721779435544218e-5 0.00012721662115791981 -8.528490779917581e-7 -0.00019067782503997627 -4.9057545575570305e-5 7.143704916369573e-5 -1.7660780750256798e-5 -8.677305394529909e-5 0.00010103282166665169 4.7514246225100136e-5 0.0001824774901338152 5.878463180989467e-5 0.000128127119941821 -1.2905160283962397e-5 3.6362946238552555e-5 0.00016420255470575814 -4.1774111914943015e-5 -2.3808921278478442e-5 0.00011672499586847616; -5.0891475054536105e-5 -0.0001067845156077372 8.97531712533836e-5 6.752176652559341e-5 -0.00014772811171236277 0.00018824930104789544 0.00010544346825815134 -4.435333228423939e-5 -6.169287423939083e-5 -3.676331159761898e-5 -0.00013076693774304157 -2.977977471994976e-5 -9.30791434958644e-5 9.931701008360094e-5 0.0001442824645190425 2.3932729371747116e-6 9.45159392619217e-6 -3.5089903191996896e-5 -0.00014437169791302372 3.0454991631692155e-5 -9.116904455107467e-5 -7.347171829418073e-5 0.00011113700695465489 -4.86133836394718e-5 7.725825655403714e-5 -9.229903714834406e-5 4.216158724636205e-5 -6.706188306755059e-7 0.00010392581263919548 2.0530750973980987e-5 -9.21061798829647e-6 4.5394007814016753e-5; -0.00015743531624349367 7.15021080101796e-5 -7.537228596658556e-5 1.592160075931863e-5 0.00011092283163424132 0.000161285053007016 1.0186517262548583e-6 -0.0002524738891518798 5.164220707926507e-5 -5.618263305510793e-6 5.021589019423594e-5 -1.8231746100640954e-5 0.00021108093566697602 -0.00010517423728043913 -0.00018276178964907883 7.053175719892157e-5 -2.718360959289657e-5 -0.00011870241072057294 -3.0331758200484555e-5 -0.0001393583305183005 0.00015240381541633508 -3.895247649080919e-5 1.4151881412260636e-5 4.797985017620516e-5 8.232545875915815e-5 2.510261326067274e-5 5.5316147306864396e-5 0.00013990744552545357 9.58195516138273e-5 0.0002240128880363817 -7.421616540551658e-5 0.00020260716625460204; -1.9458797750781913e-5 7.224534202621065e-5 4.872146529167332e-5 5.282885250424719e-5 7.381363831015988e-5 3.7185761912476934e-5 -9.709742734185963e-5 -2.3049798433187027e-6 -9.347758022113932e-5 -8.62270593547899e-5 0.00010451931449140425 5.336268587641394e-5 -2.7094556925946456e-5 1.1709049077103984e-5 -9.241085935129889e-5 0.00010580626404632567 -5.51358192656392e-5 -0.0001874741414117369 -0.000112711421379132 -1.933207057801751e-5 -2.2956205056768283e-5 -9.809755637909971e-6 -0.00011079948889303243 3.064306439270943e-5 -3.856072215924292e-5 -0.00025350393697364096 -7.751060765213512e-5 0.00012790592300960973 -2.7781752932719384e-5 -0.0001450755647870429 -3.0130541189943816e-5 0.00019734157977880452; -0.00011845059321829834 -7.507711830883873e-5 -2.7989485316076476e-5 -7.800177678336352e-5 -0.00014422827058341016 0.00012155768803450886 -1.135747821086112e-6 -5.998015381964665e-5 -8.831623250490335e-5 0.00018719097301028479 -0.00012409002061802884 4.239937283229127e-5 -2.6208903873790678e-5 1.9228900273004477e-5 -0.0001629590302040786 2.906252433560195e-5 -2.1880234781277872e-5 3.624361437643244e-5 1.826211831894637e-5 7.226073961601583e-5 0.00018741191473920204 8.378242222489642e-5 0.00029425420092549475 0.00014808852160014559 7.064179539112718e-5 0.00012297153756617063 -4.992342892829318e-6 0.00016424472727476094 -3.187832458978451e-5 6.159783828435431e-5 2.4862139470785185e-5 3.704649810925437e-5; 3.920325537128732e-5 -0.00016207555461059762 -7.200903644170987e-5 0.00017463739925678012 -3.465337874567582e-5 -1.1316565971905128e-5 0.00010220300748858774 -0.0001318377807116201 -0.000231435002999774 0.00011407202248610488 2.6412802187587594e-5 8.77509475025366e-5 -3.600955903731171e-5 -0.00014697722951732454 5.09410193693853e-5 -5.352957384511343e-5 8.89133617669381e-5 0.00016002379845977274 -1.6460259054035204e-6 2.5749471321786555e-5 4.346678099628058e-5 3.4586536573804706e-5 4.257078683386076e-5 0.00017647307970710592 8.84159554765609e-5 5.447547221354694e-5 -0.0001482893029539134 0.00013041206250580495 0.0003133309174949323 0.00016699000400640165 9.130092188601032e-7 -0.00010057848419147135; -5.474678854731801e-5 -0.000160715934354058 -0.00010237790283205395 5.159515549970668e-5 -0.00019383217894862446 0.00012196333537652572 9.483956463413691e-6 -6.330539198748996e-5 -0.00012120759364838903 0.00012596190336701784 -0.00024071388051303847 9.980222367524332e-5 0.00012129898223869385 0.00017801544291508586 -5.5090704873845974e-5 2.9155425553498383e-5 -2.5375372016396963e-5 9.271978470572257e-5 -2.8413895589584058e-5 -1.599471269506585e-5 -3.44609930754548e-5 -1.0009388998848727e-5 -4.04643638750447e-5 0.0002093051311101853 -3.076420997892551e-5 -7.34148548791905e-5 -0.00014254769356844156 -0.00014459264511165873 2.1643053886537962e-5 0.00020331991382484545 -0.00025072877050767673 -2.1939175359545803e-6; 0.00014951613510133977 3.7449118357517004e-5 -1.8441813986960397e-5 3.7198723552189794e-5 -7.663285115627364e-5 0.00011830952379027981 5.859272012827485e-5 -0.0001707031687291813 -3.7995908238312274e-5 -0.00010313939029942679 -0.00014196673047622263 -6.11534841627403e-5 -0.0001660880142625886 -3.3149185506725294e-5 1.992305540993663e-5 -0.00010888553416948558 -3.078448473020186e-5 3.467533050211379e-5 8.764261391121029e-5 6.801439402266435e-5 -0.00014627144893524942 -0.00011825137964110328 -9.276170650623512e-5 0.00015544784499438684 -2.9181898677007106e-5 -4.3011404082734625e-5 0.0001342093974781766 0.000193338500215815 -0.00016023302206663053 -2.169450537253208e-5 -7.482854463519302e-5 8.685088058342067e-5; -6.001501138897588e-5 3.3255787581501554e-5 -0.00018476408409994787 -4.074978174982612e-5 1.68909620121756e-5 -1.1477958526737275e-5 -5.84291291154413e-5 -8.729057559618995e-5 2.433153816742058e-5 1.9655473021629638e-5 8.378770996128616e-6 -5.476635748280836e-5 6.021385082634157e-5 0.00013728925581446306 -0.00012297083127422125 -8.114749185851018e-5 -4.078917014638889e-5 -1.5668706403535005e-6 -5.772223709142615e-5 1.7126522958945146e-5 -0.00021729879050625607 0.00012708276075201725 -0.00013274642957484466 -3.515942704624683e-5 -6.46693978190254e-5 -0.0003199579831409739 -0.00016136834425491933 1.5405851751658783e-5 9.541822249658673e-5 -6.0092627666821045e-5 1.145397629179884e-6 -4.0370595218767396e-5; -5.000014915057141e-5 2.494164311284655e-5 -5.8759532641120965e-5 5.2166002768353026e-5 -1.2771942585082898e-5 6.559864652299981e-5 4.0545250720596525e-5 0.00016126298089002466 9.564368758072687e-5 -5.126937083477278e-5 0.00016641531930727975 -5.543953137165738e-5 3.901339970834585e-5 1.7208802654213075e-5 -9.079291829183453e-5 9.477374498457164e-5 -1.3077650524677708e-6 -0.00013677844743286785 0.0001570289701870915 0.00011862295888786191 3.1160942566325596e-6 0.0001493580590417042 5.890522754126643e-5 6.367133772449952e-6 1.484252481387917e-5 6.984494631832849e-5 0.0001060695673171571 -0.00010060277180092052 -4.798845913309688e-6 7.030739890834362e-5 -0.00012530545553234924 7.933596515069824e-6; -0.00012761642281400407 3.829130875676353e-5 -9.232279697168137e-5 -0.00012120298631922292 -2.9543771942533063e-5 -0.00019396048153718642 6.683422108912977e-5 1.5941562613415253e-6 4.3473955175482493e-5 -6.595860466329252e-5 2.6190217996215272e-5 0.00016182271358122893 -2.5381713184434333e-5 -0.0001913860149026485 -4.966525810157233e-5 -6.722278774684806e-5 -4.27270286119674e-5 -0.00010628484304031819 -0.00015566617073924102 -1.2010567642537724e-5 -0.00019421633331073918 -0.00014504069693354818 0.0001372749568152567 5.064702891852398e-6 -1.0226159046558242e-5 3.0771162729112046e-5 -1.3030431345344805e-5 -6.38817844257151e-5 0.0001685011986599658 2.2798906885170734e-5 -1.5316935586431123e-5 9.709879961595588e-5; 5.9081573085505286e-5 4.492408024344725e-5 6.869221815104656e-5 8.828564441027609e-5 -9.162548249622684e-5 9.135181298443368e-5 -1.1144763816642302e-5 1.6682653931439985e-7 7.119649362246403e-5 0.00011957557998794545 7.293537473270394e-5 0.00021982280722550814 -3.4206226159953365e-5 0.0001043763155346893 -3.060058534646511e-5 -6.858909492183354e-5 -7.983075502596804e-5 0.0001239035524075113 0.00012049682536125546 0.0001078856554111509 -7.296872664045458e-5 0.00010023029298718556 7.898151361115362e-5 5.196418041560933e-5 -0.00013449631970178206 0.00012164735071090812 -5.5492102005902636e-5 -0.00016153855122430445 2.8795808789080754e-5 -9.2361125466771e-5 2.727579575986221e-5 7.998472264698595e-5; 4.7695753657516443e-5 -2.9261215111564066e-5 -5.692595755865946e-5 -4.3806085393984134e-5 -1.2295210815882072e-5 0.00023630210466032864 0.0001318958963947409 2.141516937805846e-5 0.00011570184118859416 3.91202023999928e-6 0.00014314104895854743 5.048473735039533e-5 2.0869474375982414e-5 0.00012758691422457233 -0.00015306977753245678 -0.00014465501613632984 9.902121359298614e-5 6.482863658855418e-5 -1.1285091439308593e-5 0.00016258200025000903 -0.00010241923178729843 9.69253027685662e-6 -2.6325688175365898e-5 3.2483928789411217e-6 -0.00016282234015126915 5.3310239074525906e-5 0.0001441303463634091 -8.266756072593175e-5 -9.529195109100576e-5 6.566144269707006e-5 0.00017639786505422312 6.804734561407641e-5; -4.585903194875359e-5 1.5447644056257053e-8 -0.00011045456891788264 3.808176291732458e-5 1.9976491906884742e-5 -0.0001826530450363856 -7.993375299173108e-5 -5.162132844470262e-5 3.817799109476421e-5 -9.163170294978131e-5 -1.5618937019283503e-5 -2.9936548222579992e-5 -2.3045425101547744e-5 5.667614562644025e-5 -0.0001784896111223178 -3.2373062700759945e-5 -0.00014099833321434363 -0.0001249526493358704 -2.2998351474772347e-5 0.00012248664746385463 1.0506616612286325e-5 -0.0001508480589711339 0.000139641958534224 -4.921862897954922e-5 -4.295759359781311e-6 0.00012432784309004415 -1.0912285566680452e-5 0.00016000483186309235 8.178794554059648e-5 -0.0001130932968100776 1.0087975291565541e-5 -0.00013291232868704894; 2.171214819599724e-5 0.0001447647040599771 -0.00013535005982204392 0.00029737222521860544 2.450419324804556e-5 0.00011594509433551828 0.00010982276162891965 4.636595204418163e-5 8.974511311449869e-5 1.6804523880172692e-5 -6.788051959529918e-5 9.443947788433795e-6 6.432954619849625e-6 5.7920871227444633e-5 3.2008715975729316e-5 7.025272807868193e-5 6.281398546479508e-5 3.2857896626861225e-5 -0.00011497948125408698 0.00018928722002372475 2.0211816632123215e-5 5.508519043859443e-5 7.343890629772054e-5 6.195941696705302e-5 -7.110396526923401e-5 3.679389138185158e-5 -0.00010527868707371202 -0.00016391071093906003 5.8788514621021666e-5 -2.2499311484815415e-5 -6.027058461268477e-5 3.552714716123197e-5; -0.00028695942745282453 -7.83207098578032e-5 -3.122699416446222e-5 -6.820973356690664e-5 -4.8084362046506574e-5 6.364804432057957e-5 0.00021295456537308275 -8.11330566505765e-5 5.0904419200097946e-5 -1.9019316081832097e-5 -3.8154044145196885e-5 0.0002134612048536628 -5.1094010981902754e-5 0.000247710825365791 -9.670176512765806e-5 -0.0001881710297169093 -0.00010956925251259622 -3.884767937445277e-5 -9.56663599793647e-5 -3.681085143790956e-5 -2.1690130947103855e-5 -0.0001776879464017127 8.251027321937294e-5 -3.440487415384726e-5 0.00016119929643143674 -0.00010045613742869423 -7.649275647446121e-5 3.486745359338766e-6 9.679524811479207e-5 -3.3128463823525166e-5 0.0001112034922841078 -3.674878539021467e-6; 0.00024580563479512557 -0.0001299429044319346 -6.418370900965224e-5 0.00012444465540292462 -5.9322772694847714e-5 0.00016992653548195203 0.00019626103460731128 1.1663394279907336e-5 0.00016226092123368616 4.595111170610686e-5 -0.00019549982095659306 2.4713899283688195e-5 0.0002181433025089909 -3.399595368127871e-5 0.00010822258048401721 -0.00017663778850864916 -0.0002014805562520437 -2.011418507199593e-5 -9.46786734583909e-5 2.1680743548215727e-5 0.00016740041030980328 2.5922523154894134e-5 5.469979590098055e-5 -1.647527130435046e-5 -9.057944987025405e-5 -1.9420633516266245e-5 -6.969024831952711e-5 -7.600815876677298e-5 0.00022230344769022872 -0.0001509315828751828 9.820482376208212e-5 3.0052851461819145e-5; -9.464760383811938e-5 4.8808337886704514e-5 -5.774765201285999e-7 5.162752621664123e-6 -0.000130014258467007 -2.2825982306266194e-5 3.997568550298657e-5 0.00023709421206044755 6.103581385000322e-5 0.00014811786517279982 4.3738501552844606e-5 8.541170522116888e-5 -7.105502751178327e-5 -1.245112613209924e-5 3.743150536003095e-6 1.1094857690160536e-5 5.9660825207949006e-5 4.152173190855271e-5 0.00012820525720494576 -0.00019183823038579185 0.00010795484835293923 1.3231349875105651e-5 -0.0001404091135201711 0.00010653166376764736 -6.700814716273139e-5 5.7146654434155066e-5 -2.2375196714301707e-5 -1.455218076548555e-5 -0.0001678051894133647 -6.958520403818399e-5 -0.000158307692764673 8.042585285557743e-6], bias = [-1.96348660388163e-9, 3.587256954007262e-9, -1.3420858019069292e-9, 7.966448558288452e-10, 2.5306297005438753e-9, -2.0845596863624143e-9, -1.819775479484989e-9, -2.0773902423320085e-9, -1.8023909693566758e-9, -3.453613588899937e-9, -5.917294074444723e-10, 2.952627633077269e-9, 2.4861077672550896e-9, 3.4745552686458593e-9, 4.098373674563021e-9, -2.260795870747409e-10, 2.1310985975449724e-9, -2.2104263982948423e-9, 2.9498910565961772e-9, 3.0951263432409247e-9, -2.0359868095274e-9, -1.723416017197214e-9, -4.790184257559881e-9, 3.196526826279751e-9, -3.5622899100045564e-9, 3.5548462732108944e-9, 3.113875191522221e-9, -3.138544531947625e-9, 3.488280978092165e-9, -1.7709538269865677e-9, 2.016922828631347e-9, 1.628348868420308e-10]), layer_4 = (weight = [-0.0008233182089716139 -0.0007953544686150181 -0.0007540695901527111 -0.0006948056703372506 -0.0005720884632729438 -0.0008519730579876156 -0.0006163302082697752 -0.0006438664533573394 -0.0008063541107629101 -0.0008045493752316895 -0.0005622253242866364 -0.0006731416561211679 -0.0005941167606944272 -0.0007393901194632896 -0.000792799515400627 -0.0006989197743611716 -0.0006402993395963954 -0.0008041844394689014 -0.0005800453758618655 -0.0005225634439755112 -0.0006532787082311566 -0.0006496953543906337 -0.0007804714524711213 -0.0006424365800815346 -0.000609932905738381 -0.0007814943236314041 -0.0007447146875010547 -0.00047352167126873875 -0.0007554189051747461 -0.0007626013089346017 -0.0007072181025069589 -0.0006768559577955712; 0.0002987898444671396 0.0002372493543989048 0.00017547996932648123 0.00024234547939569047 0.0005093496599314523 0.00012466852818971834 0.00020600676210580187 0.00020045256224647233 0.00034426203051451417 0.00017283026345556675 0.00017278504586373597 0.00014862710598230849 0.0002469486904355213 0.00025711285024051324 0.0005166025090617313 0.00020630239791395139 0.00031064597785612395 0.00021198907099172283 0.00011851682380926622 0.0002554648484669155 5.705999938450123e-5 0.00038996853733940217 0.00018668360641142903 0.0002035693419058394 0.00022135734392974448 0.00020040497004095785 0.00019203599265519138 0.00013107470621024893 0.00023855947557646003 0.00022320300089260435 0.00012119633588819062 0.00035686778350698994], bias = [-0.0006832008531285604, 0.00021693781833579523]))

Visualizing the Results

Let us now plot the loss over time

julia

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="Iteration", ylabel="Loss")

    lines!(ax, losses; linewidth=4, alpha=0.75)
    scatter!(ax, 1:length(losses), losses; marker=:circle, markersize=12, strokewidth=2)

    fig
end

Finally let us visualize the results

julia

prob_nn = ODEProblem(ODE_model, u0, tspan, res.u)
soln_nn = Array(solve(prob_nn, RK4(); u0, p=res.u, saveat=tsteps, dt, adaptive=false))
waveform_nn_trained = first(
    compute_waveform(dt_data, soln_nn, mass_ratio, ode_model_params)
)

begin
    fig = Figure()
    ax = CairoMakie.Axis(fig[1, 1]; xlabel="Time", ylabel="Waveform")

    l1 = lines!(ax, tsteps, waveform; linewidth=2, alpha=0.75)
    s1 = scatter!(
        ax, tsteps, waveform; marker=:circle, alpha=0.5, strokewidth=2, markersize=12
    )

    l2 = lines!(ax, tsteps, waveform_nn; linewidth=2, alpha=0.75)
    s2 = scatter!(
        ax, tsteps, waveform_nn; marker=:circle, alpha=0.5, strokewidth=2, markersize=12
    )

    l3 = lines!(ax, tsteps, waveform_nn_trained; linewidth=2, alpha=0.75)
    s3 = scatter!(
        ax,
        tsteps,
        waveform_nn_trained;
        marker=:circle,
        alpha=0.5,
        strokewidth=2,
        markersize=12,
    )

    axislegend(
        ax,
        [[l1, s1], [l2, s2], [l3, s3]],
        ["Waveform Data", "Waveform Neural Net (Untrained)", "Waveform Neural Net"];
        position=:lb,
    )

    fig
end

Appendix

julia

using InteractiveUtils
InteractiveUtils.versioninfo()

if @isdefined(MLDataDevices)
    if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice)
        println()
        CUDA.versioninfo()
    end

    if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice)
        println()
        AMDGPU.versioninfo()
    end
end

Julia Version 1.12.7
Commit 6d172b025e4 (2026-08-15 08:05 UTC)
Build Info:
  Official https://julialang.org release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-18.1.7 (ORCJIT, znver3)
  GC: Built with stock GC
Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores)
Environment:
  JULIA_DEBUG = Literate
  LD_LIBRARY_PATH = 
  JULIA_NUM_THREADS = 4
  JULIA_CPU_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0

This page was generated using Literate.jl.

Trusted by

Training a Neural ODE to Model Gravitational Waveforms ​

Package Imports ​

Define some Utility Functions ​

Simulating the True Model ​

Defining a Neural Network Model ​

Setting Up for Training the Neural Network ​

Training the Neural Network ​

Visualizing the Results ​

Appendix ​

Training a Neural ODE to Model Gravitational Waveforms

Package Imports

Define some Utility Functions

Simulating the True Model

Defining a Neural Network Model

Setting Up for Training the Neural Network

Training the Neural Network

Visualizing the Results

Appendix