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[-0.00012740369; 0.0002829192; 0.00011747041; -7.6054515f-5; 0.00011791948; 1.6768246f-5; 2.1529293f-5; 1.2365288f-5; -3.3090422f-5; -1.4934125f-5; -0.00018906317; -0.00015978911; -5.2236803f-5; 6.619379f-5; 2.8871762f-5; -6.985896f-5; -0.00012750016; -4.9145605f-5; -2.265392f-5; -4.8042086f-5; -6.933823f-5; 3.0399304f-5; -1.9543932f-5; -7.762098f-5; 1.4231987f-5; -0.00013111791; -1.1848226f-6; 9.695231f-5; 2.6603f-5; 8.262481f-6; -4.2200292f-5; 0.00013653966;;], 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[-0.00011883832 0.00016211007 7.530082f-5 3.6053512f-5 -4.573183f-5 8.8933644f-5 4.4013555f-5 -2.281334f-5 0.00014123098 -2.4297544f-5 0.00019308172 -0.00014420117 0.00014794337 0.00010906441 -6.69199f-5 2.0631649f-5 0.00013781223 0.00026044034 -9.503834f-5 -1.4156406f-5 1.2446943f-6 -0.00023032943 7.963334f-5 2.1737129f-5 5.183522f-5 -0.00012601748 -7.024304f-5 -4.0446568f-5 -9.036837f-5 -0.0001453705 8.865468f-5 0.00016911232; 4.6706733f-5 -0.00011201001 0.00012273871 5.084209f-5 2.6279f-5 3.822703f-5 -5.8278827f-5 -7.798874f-6 4.9336133f-5 9.281781f-5 -9.9753706f-5 -5.4496315f-5 -2.319941f-6 -9.3692805f-5 1.8234272f-5 -5.6362842f-5 -2.6276417f-5 -3.7362028f-5 -0.00021638282 -0.00019387268 2.631816f-5 0.00012948803 0.00013422848 5.932876f-5 -7.190098f-5 0.00019395938 -0.00014219151 8.827311f-6 0.00015165919 -9.639158f-5 -6.9325206f-6 0.00024369956; 1.6531829f-5 3.324146f-5 1.0770538f-5 0.00032846673 -0.00010278821 -0.00012828407 -2.4500674f-5 0.00010181535 -9.09991f-5 -5.7240923f-5 2.3821507f-5 7.2236144f-5 -0.00017421869 -1.6259803f-5 -0.0001503551 -0.00016787657 -0.00010444166 -5.3394717f-5 6.448141f-5 -0.00011835448 -1.6765971f-5 9.703217f-5 0.0001549397 -0.000110294925 -0.0001500862 -0.00011111564 -6.55644f-5 -0.00011254008 8.9033536f-5 -1.9381969f-5 -3.242799f-5 6.893542f-5; -0.00014275202 9.9729354f-5 0.0002748568 1.7739214f-5 0.00020850572 -0.00015730661 -0.00018554807 2.874012f-5 -1.7650054f-5 -5.6411718f-5 8.5230306f-5 -2.7026928f-5 -0.00017837374 0.00012095432 -0.00012397548 4.710292f-5 9.573434f-5 -0.00021673366 5.7221667f-5 6.547384f-5 -0.0002702296 0.00016617394 -0.00021189428 -4.2864453f-6 4.1075673f-6 -7.8536876f-5 1.2022175f-5 4.190836f-5 4.4570666f-6 -7.9446625f-5 -0.00011560316 -0.00013503202; -3.078594f-5 -0.00016294724 0.00011326176 -6.572744f-5 -7.108519f-5 2.314764f-5 4.507882f-5 -0.00016388444 -3.089808f-5 6.324521f-5 -1.3164919f-5 0.00017959482 -2.3044584f-5 0.0001509715 6.8137466f-5 -6.668825f-5 5.5475633f-5 0.00018898686 -0.00018197733 3.0050021f-5 0.00017996471 -7.87114f-6 -3.451523f-5 8.54307f-5 -0.00023151627 6.216509f-5 5.9033755f-5 -6.2324376f-5 0.0001102082 0.00021602848 7.7606775f-5 1.4618241f-5; 9.206881f-5 5.9961396f-5 7.1755067f-6 -0.00014432182 -0.00017698832 0.0001585406 2.5399855f-5 -9.850883f-5 7.0495924f-5 6.412567f-5 -0.00013826661 -9.3917624f-5 3.7673253f-5 5.3099416f-6 4.5290446f-5 7.514679f-5 -0.00012829428 0.00010447201 -4.065499f-5 0.00018594146 1.6332884f-5 0.00010614698 6.490749f-5 -0.00010899853 0.00013092482 9.737774f-5 -1.8231993f-5 -0.000134424 -8.559609f-5 2.620763f-5 -0.00017186062 3.7317433f-5; 7.6358454f-5 -3.9395563f-5 -0.0001612821 5.4471766f-6 2.6774207f-5 -9.189757f-5 -9.633577f-5 6.0510967f-5 6.118395f-5 0.00021276085 4.027092f-5 0.00010298319 -0.00025871178 2.7570077f-5 4.5452813f-5 -6.0397437f-5 0.0002522944 0.00021370963 0.00018339598 6.321723f-5 -0.00012879871 -8.762518f-6 -5.8970087f-5 5.202664f-5 3.624049f-5 -3.0087149f-5 3.0803287f-5 6.697916f-6 -9.7831515f-5 2.9483746f-5 -1.6436474f-5 -3.472948f-5; -0.00014166838 -2.4888806f-5 3.9980073f-6 -0.00010593238 0.00015422335 -0.000291447 5.448907f-5 3.129295f-5 1.7705948f-5 7.985548f-5 -0.00017656408 -4.0946226f-5 2.439601f-5 4.541847f-5 0.00013723804 -1.798299f-5 -5.566121f-5 -1.2338974f-5 -0.00011110796 -6.6668595f-6 -8.722685f-5 -8.287055f-5 -8.5696265f-6 3.2533455f-5 -4.42852f-5 0.00010856262 2.9701905f-5 0.00018288789 -6.944528f-5 4.9842238f-5 2.8572602f-5 -7.068699f-5; 7.656304f-5 -1.2038126f-5 -8.333188f-5 -1.921461f-5 -5.0217466f-5 0.00011668638 1.7529896f-7 2.8149838f-5 2.0734207f-5 -7.558573f-5 -4.244396f-5 7.3614065f-5 5.1631312f-5 4.467191f-6 -3.1148586f-5 0.00015302071 -0.00014994133 0.00015237811 -2.56135f-5 0.00010538345 -1.517395f-5 8.023177f-5 4.210785f-5 -5.0196526f-5 3.5190656f-5 0.00016596122 3.744718f-5 0.00016471746 -0.0002583961 8.714164f-5 3.700902f-5 -9.44848f-8; -0.00013893115 5.529618f-5 7.912546f-5 -3.5857825f-5 -7.9682395f-5 -9.393053f-6 -3.4598037f-5 -9.9752135f-5 -6.774488f-5 -3.6205358f-5 -5.918434f-6 8.4171385f-5 -3.074355f-5 -2.3808225f-5 -0.0001094997 0.00016207276 0.00016943546 4.6223715f-5 4.4158947f-5 1.5926276f-5 7.730441f-5 -0.000117423166 2.9555804f-5 -6.9854555f-5 0.00016211405 0.00012636333 0.00013799321 7.932004f-5 -0.000181433 5.677608f-5 6.190729f-5 -5.34332f-5; 1.091316f-5 6.898843f-5 -5.3287233f-5 4.443628f-5 -0.00014640561 -4.310104f-5 0.00016032069 2.7929049f-5 -1.1572988f-5 0.00011948562 -8.971406f-5 -0.00017052988 -0.00012518509 9.313835f-6 2.1012194f-5 8.142811f-5 9.524531f-5 0.00016773543 1.4860113f-5 4.031107f-5 -0.0001116714 2.5877323f-5 -5.9532846f-5 -7.3677f-5 -0.00013412829 -9.703796f-5 -2.476725f-5 -0.00011859284 -0.00013061453 4.1761545f-5 -2.7915237f-5 -7.363147f-5; -0.000103807 6.201734f-5 -5.544839f-5 2.3652956f-5 0.00018599398 -0.00012593916 -0.00023664696 3.4347377f-5 0.000119874596 -7.129401f-5 4.4952536f-5 -0.0001223195 -2.2027161f-5 -3.1292904f-5 -0.00011942254 2.9111163f-5 6.970712f-5 -8.301896f-6 5.3726428f-5 9.902478f-5 -5.992248f-5 9.827444f-5 7.367511f-5 0.00011540036 -1.1428022f-5 -0.00016597832 -5.183176f-6 -3.0457222f-6 -8.83031f-6 7.312629f-5 0.00015557157 3.6311736f-5; -3.6410653f-5 -3.4849163f-5 -1.9353332f-5 6.0067207f-5 4.855823f-5 0.00013734757 0.00011851531 7.9817895f-5 0.00012610013 0.00013653122 1.0763614f-6 -0.00017003289 2.3878196f-5 -1.2759429f-5 6.145687f-5 0.00015128728 9.730206f-5 6.1051585f-5 5.273011f-5 -0.00010460802 -0.00013594287 -0.00012197394 -2.166086f-5 -9.7288634f-5 4.1505056f-5 7.83626f-5 3.1741776f-5 -0.00010521191 3.1580654f-5 0.00015126221 -7.423819f-5 0.000114573355; -0.00010096334 2.5505697f-5 0.00013573456 -0.00011650252 -3.594842f-5 -4.6425645f-5 -6.4124186f-5 0.0001650304 4.284348f-5 -0.0001036861 -7.88737f-6 5.696396f-5 -0.0002721436 6.916334f-5 -1.5119892f-5 -1.9959392f-5 3.7423426f-5 -2.9957613f-5 1.1500234f-6 -2.7685153f-5 -0.000120601835 -5.48023f-6 5.938496f-5 6.0568582f-5 -7.63975f-5 -3.43356f-5 8.1877755f-5 -4.6699482f-5 -1.2054909f-5 0.00010623455 0.00012914876 0.00011553858; 6.1279425f-5 -4.431227f-5 -5.116625f-5 5.0682695f-5 0.00013331852 4.674396f-5 -0.000111019624 -0.00015239141 -0.0001074623 -0.00014403673 0.000118654374 0.0001224792 -5.9730697f-5 -4.2208572f-5 3.96425f-5 -0.00017723927 -6.7507f-5 -5.8360045f-5 2.0410068f-6 -0.00030277498 -0.00011315536 6.395578f-5 -3.9493447f-5 6.754751f-5 9.786596f-5 0.00017523357 -6.135887f-5 6.42346f-5 -4.100864f-5 -2.1818307f-5 1.799318f-5 -0.00015556677; -5.831849f-5 -0.00013116228 -4.1019448f-5 2.3371338f-5 1.7581664f-5 2.6676395f-5 6.519024f-5 -4.637574f-5 -2.000425f-5 2.7907858f-5 -0.00011388356 -6.6019497f-6 1.4657718f-5 -7.904789f-5 -0.000121811296 -5.0444803f-5 0.00014095817 0.000104856954 -2.9173769f-5 0.00010031487 -2.424529f-5 0.00010942139 -8.705766f-6 5.3678636f-5 -0.0002882384 0.00011575385 1.2718286f-6 -4.908749f-5 3.6347406f-5 5.5572415f-5 -8.051652f-5 0.00021660334; -4.4692453f-5 7.049371f-5 -0.00017258218 -5.024602f-5 -4.1327574f-5 -6.3993946f-5 -0.00011652188 -0.00011184075 -1.2802917f-5 -1.4745073f-5 7.661019f-5 -0.000105205516 6.302245f-5 -2.253077f-5 -1.3354157f-5 -9.5626914f-5 -4.7998372f-5 -0.00018742848 -7.0078545f-6 1.0734203f-5 1.9317843f-6 0.00012874554 0.00020353618 -7.747934f-5 -0.0001484861 1.4414527f-5 -1.3310939f-5 5.1022416f-5 -0.00017990098 6.542906f-5 8.741957f-5 -2.359333f-5; 9.881722f-5 -0.00018412365 5.3179978f-5 -7.314854f-5 -0.00018796134 5.2748885f-5 6.4184633f-6 -0.00022235027 -1.07798105f-5 -0.00016381011 2.4730862f-5 1.8719355f-5 -0.00013298455 5.27565f-6 -8.447046f-5 -0.00013032263 8.301421f-5 0.00015812389 -9.081025f-5 1.576368f-5 -0.000110590714 -0.00013774354 7.136221f-5 0.00017792535 -0.00012306841 -6.031452f-6 -0.000111355694 3.351387f-5 0.00013012504 2.1386604f-5 -6.792882f-5 -0.00019728256; -0.00013415128 -5.4049073f-5 -3.738039f-5 -1.664742f-5 0.00018081663 5.0219496f-5 -0.00010525832 -5.308501f-5 -3.7414025f-5 -0.00014410626 -0.00016036273 -5.6471188f-5 1.0689242f-5 -5.6242523f-5 -1.6086766f-5 -2.3015946f-5 4.1077234f-5 4.290033f-5 9.8332115f-5 6.0066348f-5 -0.00017003951 -1.3480143f-5 -1.6541278f-5 -4.5134067f-5 2.2371405f-5 9.04125f-5 1.0516738f-7 5.9817943f-5 0.00014558063 -5.8828216f-5 5.8181216f-5 5.7291298f-5; -5.0592294f-6 0.00012211058 -0.00025906868 1.246755f-5 4.78124f-5 -0.00012147193 3.825797f-5 -1.8440833f-5 -8.760192f-6 4.196563f-5 0.00012015885 0.00010762258 4.2500556f-5 -6.582965f-5 -2.4215813f-5 0.00013932877 0.00012746888 -3.894018f-5 -8.0681f-6 9.210975f-6 -0.00014211301 3.486637f-5 5.4501445f-5 -0.00012538771 8.213844f-5 -5.9459973f-5 -1.21185085f-5 0.00015861518 0.00020224642 6.898103f-5 -1.0166967f-5 5.0113344f-5; -2.8129567f-5 -0.00010227668 5.5015364f-5 3.3228116f-5 -6.1948044f-6 -8.354939f-6 -4.7726057f-6 -9.406278f-5 -0.00015331886 7.337104f-5 1.4396411f-5 7.3479176f-5 6.408236f-5 9.9273384f-5 -0.00011662606 -4.1228344f-5 -4.8512935f-5 6.371833f-5 9.100913f-5 -6.8328634f-5 6.1945306f-5 -9.731179f-6 1.1041556f-5 5.0564286f-5 -4.037148f-5 -3.5543264f-6 -6.969499f-6 9.41184f-5 -4.01786f-5 -2.7077498f-5 -0.00012816639 -7.101698f-5; -0.00013893766 -3.6935384f-5 0.00012717156 3.889238f-5 0.00019494042 -0.00023578813 -1.8655486f-5 0.00010796451 -0.00011551302 -0.0001285786 -4.9468807f-5 -0.000113068105 -3.0530107f-5 4.8436432f-5 -2.8621464f-5 -2.1881331f-5 3.9398f-5 -2.7309956f-5 -8.555159f-5 0.00011638427 0.00013287517 3.376408f-5 -1.6491314f-5 6.9282185f-5 -0.0001685218 -7.202947f-5 5.0425333f-5 0.00016450853 -6.322706f-5 -0.00015934715 -5.325392f-5 0.00014156898; 6.9513924f-5 0.00017001157 -1.9957135f-5 -0.00019666016 -3.0069468f-5 -2.8054343f-5 2.3510336f-5 2.0577886f-6 -6.147511f-5 -6.8404246f-5 -3.5530884f-5 -0.00025652008 3.0617088f-5 3.162067f-5 -1.235201f-5 -1.8362088f-7 -6.18344f-5 2.8786582f-5 0.00025447056 6.453014f-5 0.00010893362 7.084508f-5 4.144305f-5 7.6587225f-5 -6.185558f-5 0.000114900686 -3.6040474f-5 -3.0604613f-5 1.8648767f-5 1.0163714f-6 0.00020527198 -8.5257016f-5; -0.00020263731 5.779788f-5 9.77862f-5 2.144957f-5 4.01292f-5 -7.66713f-5 -9.3635994f-5 2.0310998f-5 2.2400192f-5 1.4265192f-5 -1.9146615f-5 -5.1478888f-5 6.421394f-5 4.011086f-5 -8.190776f-5 8.9714194f-5 -1.132885f-5 -5.858288f-5 -0.00016909184 -7.128159f-5 0.00010621516 4.6267083f-5 -6.844735f-5 -7.802731f-5 5.7676036f-5 -1.6105354f-5 -6.1402054f-5 -3.845399f-5 -0.00013553616 -0.00016456167 0.00011255683 -3.2607055f-5; 3.0035302f-5 0.00011284146 3.08702f-5 -0.00013158226 -6.900297f-5 -0.00018175328 -7.4122574f-5 0.00013386089 9.37536f-6 3.3115342f-5 3.6192025f-5 -0.00020746594 -4.2870886f-7 9.2499096f-5 0.00017354553 0.000121299505 0.0001772704 -0.00012464817 -1.0922418f-5 1.0695058f-5 -5.4165183f-5 -9.1385795f-5 0.00013952663 -2.5701962f-5 1.887395f-5 2.6840658f-5 4.5049816f-5 0.00012716079 -6.9047215f-5 9.291953f-5 -0.00011380827 0.00019439925; 0.000102187994 1.22124175f-5 9.578381f-5 -2.4888235f-5 0.00011589589 -5.3192394f-5 8.777504f-5 -0.00012867774 -1.3358861f-5 -9.501309f-5 -8.934559f-6 3.3132103f-5 -3.5258294f-5 -1.0429714f-5 8.426464f-5 0.0001677649 6.131737f-5 -8.0191945f-5 -8.1952865f-5 0.00010651573 0.00012543048 -7.1446026f-5 -0.0001274963 4.382758f-5 0.00010078965 7.164009f-5 1.3454015f-5 6.114415f-5 -6.925468f-5 -0.00012379461 -0.0001052033 0.00021505856; 9.098349f-5 -2.613176f-5 6.134684f-5 0.00019247418 -8.3756575f-5 -8.8636734f-5 4.8438145f-5 0.0001622886 -3.7102094f-5 -0.00020098386 1.8178856f-5 1.2072439f-5 -7.4134274f-5 -8.538014f-5 -0.00015448223 0.0001433431 9.8347424f-5 3.4221663f-5 0.00018170968 -1.737234f-5 -8.25139f-5 2.9273187f-5 -5.2558422f-5 -6.1004914f-5 -0.000119143035 5.090022f-5 -3.6541965f-7 0.0001849076 -0.00014541746 -0.00021117581 5.0580045f-5 -6.045963f-5; -0.00011323827 -3.7466794f-5 6.2076346f-5 0.000117488 9.58791f-5 8.5028114f-5 7.3585805f-5 4.013032f-5 -6.7635694f-5 5.5743596f-5 -1.7162274f-5 0.000110323635 6.477205f-5 0.000104142004 -1.7331828f-5 -7.518243f-5 5.9634796f-5 -2.5763058f-5 0.00010183134 0.00029511002 0.000112319074 2.1055117f-5 -3.092933f-5 2.3800909f-5 5.351867f-5 -2.951273f-5 9.360001f-5 -1.2538956f-5 5.981233f-7 -7.668104f-5 1.29889f-5 -7.67461f-5; 0.00010979223 -9.928308f-5 1.1121979f-5 1.26235145f-5 -8.877297f-5 -2.7202404f-5 0.00012182741 2.991497f-5 -0.00014495649 -6.004588f-5 -0.00010957989 4.272193f-5 -0.00014663153 0.000120303135 0.00021010895 -6.990061f-5 -0.00018528069 -6.415023f-5 -0.00019884191 -0.00014836082 0.00019364378 -5.671753f-6 0.00020367738 -1.9169927f-5 5.6286488f-5 -1.9594829f-5 -0.00022525775 6.2441984f-5 0.00012309616 -6.4388114f-5 0.00016684852 -9.344109f-5; -1.435831f-5 -1.3322077f-5 -0.00010599113 0.0001792288 -0.00015659232 -2.979809f-5 -6.999217f-5 -7.479394f-5 9.350338f-5 -3.870658f-5 -5.0384184f-5 -5.2107345f-5 -8.007198f-5 -0.00014156659 2.2518785f-5 1.2255664f-5 7.11975f-5 0.0001746381 1.5197651f-5 -6.845224f-6 7.13367f-5 0.00022415811 9.27296f-6 -9.463379f-5 -2.0652195f-5 2.2166772f-5 9.871981f-6 -0.00011002039 -7.9424644f-5 -2.6838708f-5 -2.7669135f-5 9.956699f-5; -4.0160892f-5 7.4887954f-5 -6.041852f-5 3.864095f-6 -1.99871f-5 -8.428105f-5 -3.8486483f-5 -3.0455762f-5 4.69744f-5 1.4922765f-5 9.120226f-5 7.112418f-5 9.060635f-5 -0.00018214685 -0.00012527332 3.45784f-6 -5.795882f-5 -0.0001350609 -6.0024624f-5 4.0583065f-5 -0.000101955426 -7.285324f-6 3.2168413f-5 3.042974f-5 1.09882985f-5 6.0050814f-5 -6.0161194f-5 -2.4558261f-5 -4.5166646f-5 3.9216244f-5 -4.9553273f-5 0.00012331163; -7.046862f-5 4.923621f-5 -0.00012298256 0.00021496882 -2.8726486f-6 5.197222f-5 4.4965513f-5 0.00026292793 -6.309988f-5 5.201896f-5 -0.00010623832 -1.4434183f-5 -7.88551f-5 5.1061692f-5 5.4896176f-5 -7.834609f-5 0.00016829337 -8.304238f-5 -0.00010564956 -0.00010120729 8.7216904f-5 -5.069088f-5 0.0001180843 -9.4969415f-5 8.8898865f-5 -5.7046953f-5 0.000103187616 2.7159414f-5 9.4795374f-5 -1.2122212f-5 4.1789943f-5 -1.8311583f-5], 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[-8.160605f-5 -0.00014014129 1.6142134f-5 4.837073f-5 0.00014416904 -0.00017337615 3.7603797f-5 -5.0842273f-6 -0.00016575365 3.376203f-5 9.086136f-5 -0.00018491983 -5.232008f-5 -6.9222006f-5 -0.000103856655 -0.00011687642 0.0001272306 -0.0001553088 -9.144664f-5 -1.9513627f-5 0.00015141648 -7.650931f-5 -0.00018771677 -4.542587f-5 -0.0002980241 -0.00017604699 0.00018420962 -7.76931f-5 -8.667553f-6 2.9419683f-5 -4.594343f-5 -4.4018987f-5; -0.00014548194 5.078884f-5 -7.379462f-5 -0.00014307145 -3.9339946f-5 -5.2547508f-5 0.00012576576 -3.0024143f-5 3.7822367f-5 -0.00013821798 0.0001273618 -4.027058f-5 -0.00011041099 7.814346f-5 0.00012398047 1.8913224f-5 -0.00018123214 0.00010965697 -6.030613f-5 -3.8959293f-5 9.851384f-5 -2.4436618f-5 7.820212f-5 7.545386f-5 -4.335593f-5 -1.9903797f-5 -5.5371387f-5 -5.816531f-5 4.0047958f-6 -0.00014871397 -7.721321f-5 -0.00025879807], 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.0006397797478215037

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 = [-0.0001274036912945828; 0.0002829192089846989; 0.00011747040844062081; -7.605451537531695e-5; 0.0001179194805444842; 1.6768246496191283e-5; 2.1529293007894457e-5; 1.2365288057470481e-5; -3.309042222097028e-5; -1.4934124919799524e-5; -0.00018906316836363104; -0.00015978910960253882; -5.223680273042042e-5; 6.619378837054671e-5; 2.8871761969613746e-5; -6.985895743118057e-5; -0.00012750015594049593; -4.914560486207761e-5; -2.2653919586435945e-5; -4.8042085836576994e-5; -6.933823169666997e-5; 3.0399303796033538e-5; -1.954393155756621e-5; -7.762097811793457e-5; 1.4231986824579016e-5; -0.00013111790758553454; -1.1848226222340873e-6; 9.695230983197427e-5; 2.660300015122377e-5; 8.262481060222132e-6; -4.2200292227780564e-5; 0.00013653966016116158;;], bias = [-2.3218522983997685e-17, 1.4944583844049526e-16, -6.273126540658058e-17, 3.099710797194866e-17, 3.555677496291091e-17, 2.5301060159807463e-17, 2.5529588998557122e-17, 1.2508243090969457e-17, -2.767624341686358e-17, -2.2651948043763224e-17, -1.7002897538033226e-17, -3.3444051247373727e-17, -6.65211613263293e-17, 4.365593071241921e-17, 3.101681590240296e-17, -1.1735344692918633e-16, -1.8555452119905322e-16, -6.545285063649581e-17, -1.1918440054662317e-17, -7.59345773990745e-17, -9.132936398659969e-17, -3.686065128200532e-17, 2.3360282271375335e-18, 1.7433587527933702e-17, 2.2201708911824488e-17, -4.885373234371807e-17, -1.6856972048143176e-18, 3.5125271662750796e-18, 2.550973436081198e-18, 7.22217799136172e-18, -1.0768987830578323e-17, 2.3639582429527842e-16]), layer_3 = (weight = [-0.0001188363375958447 0.00016211204730189714 7.530280073393102e-5 3.605549228365089e-5 -4.572985107092229e-5 8.893562427244063e-5 4.401553543272331e-5 -2.2811359475721443e-5 0.0001412329593322878 -2.4295563879544735e-5 0.00019308370338971194 -0.00014419919078950637 0.00014794535037350592 0.00010906638727743758 -6.691791985661844e-5 2.063362900213419e-5 0.00013781420512810154 0.0002604423216426444 -9.503636009770547e-5 -1.4154426015862109e-5 1.2466742974199402e-6 -0.00023032745138055772 7.963532062700856e-5 2.173910889820202e-5 5.183720078190708e-5 -0.00012601549586992966 -7.024106070178628e-5 -4.044458818401889e-5 -9.036638856644856e-5 -0.00014536852448958048 8.865665678155425e-5 0.00016911429794259016; 4.670711148509215e-5 -0.0001120096347424988 0.0001227390895816726 5.084246747873314e-5 2.6279378451080533e-5 3.8227408529283805e-5 -5.827844854348681e-5 -7.79849581004692e-6 4.933651159962289e-5 9.281819091274062e-5 -9.975332779714693e-5 -5.4495936476371316e-5 -2.3195623928671565e-6 -9.369242600069774e-5 1.8234650431341743e-5 -5.636246334898684e-5 -2.6276038463321338e-5 -3.7361649285511997e-5 -0.00021638243908665486 -0.00019387229969862284 2.631853765496121e-5 0.00012948841338374304 0.00013422885984045113 5.9329137163959045e-5 -7.190059825216761e-5 0.0001939597570323609 -0.00014219112910933198 8.827689432737058e-6 0.00015165956590653052 -9.639120237107729e-5 -6.932142079958094e-6 0.00024369993496083081; 1.65293102007704e-5 3.3238941582495955e-5 1.076801960594938e-5 0.0003284642141487475 -0.00010279072851984765 -0.00012828658621870186 -2.4503192648945366e-5 0.0001018128302938284 -9.100161808887627e-5 -5.724344121986274e-5 2.3818988452040902e-5 7.223362514639297e-5 -0.00017422120548430832 -1.626232118282794e-5 -0.0001503576215647177 -0.0001678790879260731 -0.00010444417533575368 -5.339723541978771e-5 6.447889496926039e-5 -0.00011835699955113265 -1.6768489545162894e-5 9.70296521380529e-5 0.00015493717979993477 -0.00011029744312977142 -0.00015008871672321319 -0.00011111815659672688 -6.556692208973913e-5 -0.00011254259992269246 8.903101730979815e-5 -1.9384487345770316e-5 -3.243050854066463e-5 6.89329015266282e-5; -0.000142754310453917 9.972706151604671e-5 0.00027485451550162104 1.7736921782456775e-5 0.00020850342820337696 -0.00015730890323469786 -0.00018555035809116795 2.8737827725090024e-5 -1.7652346522699396e-5 -5.641401026809029e-5 8.522801339567132e-5 -2.702922052705534e-5 -0.00017837603105293767 0.00012095202668275623 -0.0001239777705079536 4.710062761729559e-5 9.573205146877924e-5 -0.00021673595643663984 5.721937463091335e-5 6.54715474781931e-5 -0.0002702318818221431 0.00016617164486957392 -0.0002118965715082957 -4.288737508645831e-6 4.105275130095868e-6 -7.853916782480698e-5 1.2019882652348198e-5 4.190606681801693e-5 4.454774389851495e-6 -7.944891718510924e-5 -0.00011560544977259149 -0.00013503431569845338; -3.078395036634089e-5 -0.00016294524846372654 0.00011326375226720361 -6.572545260193108e-5 -7.108319857858935e-5 2.3149631145625195e-5 4.508081167236034e-5 -0.00016388245001204227 -3.089609106305964e-5 6.324719708876761e-5 -1.3162928895949355e-5 0.0001795968097424723 -2.304259339931505e-5 0.00015097348352659302 6.813945640109519e-5 -6.668625735679866e-5 5.5477623958112355e-5 0.00018898884866151892 -0.00018197534098058884 3.0052011992064666e-5 0.00017996670487567561 -7.869149500638845e-6 -3.451323880350149e-5 8.543269052400888e-5 -0.00023151428047336558 6.216708118093104e-5 5.903574549242015e-5 -6.232238536237324e-5 0.00011021019292760688 0.00021603046900358144 7.760876598998882e-5 1.4620231580308617e-5; 9.206904067613072e-5 5.996162579465991e-5 7.175736556371216e-6 -0.0001443215907951634 -0.00017698809389768262 0.0001585408288352379 2.540008497961806e-5 -9.850860371299076e-5 7.049615376817068e-5 6.412590008184325e-5 -0.00013826637879757467 -9.39173944229092e-5 3.7673483149399906e-5 5.3101714744009355e-6 4.529067589973106e-5 7.514701863473256e-5 -0.00012829405315498 0.00010447224192975085 -4.0654761704364704e-5 0.00018594169230854132 1.6333114018914984e-5 0.00010614721102828165 6.490771627940247e-5 -0.00010899830087365616 0.00013092504812592653 9.737797225315005e-5 -1.823176282269868e-5 -0.0001344237709585397 -8.559585987169289e-5 2.6207859974955327e-5 -0.00017186039187291636 3.731766335621347e-5; 7.635993565691134e-5 -3.9394081658459226e-5 -0.00016128062529696533 5.448658294887378e-6 2.677568877607196e-5 -9.189608940595482e-5 -9.633428530741214e-5 6.0512448863855616e-5 6.118543128742144e-5 0.0002127623276019311 4.027240097547355e-5 0.00010298467273150545 -0.00025871029615608815 2.7571558485767653e-5 4.545429433256639e-5 -6.0395955050801315e-5 0.00025229589405262337 0.00021371111247481953 0.0001833974648898308 6.321871219461965e-5 -0.00012879722894376086 -8.761035905072784e-6 -5.89686049797731e-5 5.2028122036978576e-5 3.6241971566605414e-5 -3.0085667139246032e-5 3.0804768494826876e-5 6.699397682500494e-6 -9.78300330179904e-5 2.948522810675959e-5 -1.6434992038731272e-5 -3.4727999126516894e-5; -0.00014166961327205033 -2.489003881702957e-5 3.9967748351504626e-6 -0.00010593360920679801 0.0001542221188218268 -0.0002914482421103054 5.4487839320103824e-5 3.129171901022742e-5 1.7704715814802305e-5 7.985425045703267e-5 -0.00017656530971647933 -4.0947458815157294e-5 2.4394776942104274e-5 4.541723765655318e-5 0.0001372368032251367 -1.798422189963597e-5 -5.566244280557701e-5 -1.2340206099138085e-5 -0.00011110919430862439 -6.668091914575313e-6 -8.722808260713787e-5 -8.287178488458747e-5 -8.570858967828412e-6 3.253222248972358e-5 -4.4286432144104425e-5 0.00010856138842869037 2.9700672625808045e-5 0.00018288665606164476 -6.944651531298895e-5 4.984100533405973e-5 2.857136942552195e-5 -7.068822296346061e-5; 7.65650400541896e-5 -1.2036125717104425e-5 -8.332987652400477e-5 -1.9212610076716113e-5 -5.021546547075709e-5 0.00011668838107719087 1.7729968997462937e-7 2.8151838542479176e-5 2.073620800795445e-5 -7.558373197686855e-5 -4.244195968503631e-5 7.36160653293244e-5 5.1633312815035e-5 4.469191772591618e-6 -3.114658482379317e-5 0.0001530227089649315 -0.00014993933307713033 0.00015238011094036367 -2.5611500113220113e-5 0.00010538545013398728 -1.5171948852434975e-5 8.023377433412666e-5 4.210985191159875e-5 -5.0194525264738234e-5 3.519265721545559e-5 0.00016596321786047696 3.7449179054172897e-5 0.00016471946566591365 -0.0002583940878860376 8.7143642140953e-5 3.701102088664559e-5 -9.24840715903171e-8; -0.00013893015145120265 5.529717876079215e-5 7.912645654025517e-5 -3.585682586096869e-5 -7.96813955905938e-5 -9.392053443000148e-6 -3.4597037455816776e-5 -9.975113513907877e-5 -6.774387745215117e-5 -3.620435833842954e-5 -5.917434436273679e-6 8.417238407739082e-5 -3.074255161597935e-5 -2.3807225233462084e-5 -0.00010949869917501246 0.0001620737557738958 0.0001694364573547518 4.622471445211609e-5 4.415994683834255e-5 1.5927275881701454e-5 7.730540804070742e-5 -0.00011742216608515495 2.955680341140367e-5 -6.985355591430048e-5 0.00016211505410931545 0.00012636433078179983 0.000137994206880401 7.93210374747327e-5 -0.00018143200301826476 5.677707943571276e-5 6.190828847204754e-5 -5.3432199792295154e-5; 1.091126393730613e-5 6.898653698303395e-5 -5.328912850673136e-5 4.4434385538619464e-5 -0.0001464075070791301 -4.310293700401411e-5 0.00016031879529035924 2.7927153107875694e-5 -1.1574883889997525e-5 0.0001194837273695066 -8.971595928158928e-5 -0.000170531774008116 -0.00012518698752247314 9.31193933210387e-6 2.101029856275131e-5 8.14262126121259e-5 9.524341619253486e-5 0.0001677335345200768 1.4858217193969018e-5 4.0309175867870597e-5 -0.00011167329323798805 2.587542764812543e-5 -5.953474140115747e-5 -7.367889741290673e-5 -0.00013413018716820256 -9.703985818050416e-5 -2.4769146354820007e-5 -0.00011859473709337512 -0.0001306164235094999 4.175964893219308e-5 -2.7917133072341043e-5 -7.363336447014926e-5; -0.0001038065955625628 6.201773898202404e-5 -5.544798737980997e-5 2.3653357074293316e-5 0.00018599438148077935 -0.00012593875622849372 -0.000236646554968389 3.43477778543493e-5 0.00011987499705276052 -7.129360788974369e-5 4.495293718638897e-5 -0.00012231909828267786 -2.2026760107907045e-5 -3.129250295467189e-5 -0.00011942213937879938 2.9111564016714822e-5 6.970752396148024e-5 -8.30149461842329e-6 5.3726828936418313e-5 9.902518026058145e-5 -5.992207921815556e-5 9.827483985566136e-5 7.3675511102123e-5 0.00011540076333323955 -1.1427620721293326e-5 -0.00016597791442189054 -5.182774985656529e-6 -3.0453210100652716e-6 -8.829909221163733e-6 7.31266928952428e-5 0.0001555719728813745 3.6312137253045274e-5; -3.64084115111956e-5 -3.4846921885641526e-5 -1.935109121896963e-5 6.006944767668333e-5 4.8560472527777854e-5 0.0001373498091148575 0.00011851754988547296 7.982013648379497e-5 0.000126102371900378 0.0001365334612224614 1.0786025719162873e-6 -0.0001700306454999196 2.388043703340174e-5 -1.275718823799573e-5 6.145911192525031e-5 0.0001512895252695635 9.730429901045395e-5 6.105382653422392e-5 5.273235019067179e-5 -0.00010460578032998622 -0.00013594063039400124 -0.00012197170126946797 -2.1658619339541734e-5 -9.728639252480647e-5 4.150729745979321e-5 7.836484309753914e-5 3.1744017355413256e-5 -0.00010520967026005345 3.158289491202545e-5 0.0001512644523196286 -7.423594780029946e-5 0.00011457559612116248; -0.0001009635037546335 2.5505534395759622e-5 0.0001357343986826659 -0.00011650267911975424 -3.5948580588615625e-5 -4.6425806757930606e-5 -6.412434809974574e-5 0.00016503024166083719 4.284331738815467e-5 -0.00010368626168131077 -7.887532191693208e-6 5.6963797433832214e-5 -0.000272143765286479 6.9163177286804e-5 -1.5120054157911241e-5 -1.9959554592083828e-5 3.742326374847348e-5 -2.9957775367826044e-5 1.1498612054981231e-6 -2.768531554395502e-5 -0.0001206019972953309 -5.4803921186487625e-6 5.938479957037426e-5 6.0568419613240666e-5 -7.639765895803639e-5 -3.433576272002511e-5 8.187759307545428e-5 -4.669964469869752e-5 -1.2055071564851707e-5 0.00010623438505788152 0.0001291485947520176 0.00011553842126957058; 6.127718625891316e-5 -4.4314508332606216e-5 -5.1168489509046804e-5 5.0680456123757e-5 0.00013331628252587403 4.674172197070011e-5 -0.00011102186289548824 -0.00015239364910570876 -0.00010746453807854927 -0.00014403897256756167 0.00011865213565132466 0.00012247696832589857 -5.97329355431978e-5 -4.221081070767373e-5 3.9640260009999235e-5 -0.00017724151001942894 -6.750923529664689e-5 -5.836228337188055e-5 2.038768397194761e-6 -0.0003027772156119142 -0.00011315759656184912 6.395353812791494e-5 -3.949568526053756e-5 6.754527094841467e-5 9.786372066542536e-5 0.00017523132811135975 -6.136110932033494e-5 6.423236009965435e-5 -4.1010877967839376e-5 -2.1820545217399275e-5 1.7990940986291134e-5 -0.0001555690079757617; -5.831861401786533e-5 -0.00013116240209312346 -4.1019573680772254e-5 2.3371212384132518e-5 1.758153827866093e-5 2.6676268836650617e-5 6.519011164671239e-5 -4.637586446590294e-5 -2.00043761910129e-5 2.7907731938808776e-5 -0.00011388368350564857 -6.602075453359193e-6 1.4657591938499548e-5 -7.904801371152999e-5 -0.00012181142140291201 -5.04449292097196e-5 0.00014095804887945327 0.00010485682812175219 -2.9173894498329973e-5 0.00010031474609747964 -2.424541546899844e-5 0.00010942126188781542 -8.705891767407677e-6 5.367850987711442e-5 -0.00028823852538009676 0.0001157537242502955 1.2717029190533967e-6 -4.908761750668753e-5 3.634728070353588e-5 5.557228884897622e-5 -8.05166439280802e-5 0.00021660321160341188; -4.469515664094636e-5 7.049100856245009e-5 -0.00017258488190417898 -5.024872410285783e-5 -4.1330277264926254e-5 -6.399664965284107e-5 -0.0001165245816373549 -0.00011184345606307562 -1.2805620310250436e-5 -1.4747776214795254e-5 7.660748412311354e-5 -0.00010520821927637484 6.30197496939188e-5 -2.2533473777004186e-5 -1.3356860501521728e-5 -9.562961752817607e-5 -4.800107529444194e-5 -0.0001874311823776257 -7.010557903418361e-6 1.0731499576646738e-5 1.9290808431821717e-6 0.00012874283453856601 0.00020353347969648386 -7.748204046220061e-5 -0.00014848880203499515 1.4411823547571387e-5 -1.331364222278891e-5 5.1019712889686316e-5 -0.00017990368035680624 6.542635998629896e-5 8.741686503344298e-5 -2.3596032799094155e-5; 9.881372911691154e-5 -0.00018412714039365367 5.3176489152928314e-5 -7.315203024705739e-5 -0.00018796482968541644 5.274539594022982e-5 6.414974584851693e-6 -0.00022235375568204147 -1.0783299196252516e-5 -0.00016381359805752103 2.472737300920662e-5 1.8715865870002776e-5 -0.00013298804319862497 5.272161295521953e-6 -8.447395143963838e-5 -0.0001303261195535583 8.30107202527435e-5 0.00015812040160011244 -9.081374069143563e-5 1.576019180714349e-5 -0.00011059420271709538 -0.000137747028795336 7.135871796702474e-5 0.00017792186203945592 -0.00012307190250073142 -6.034940590367697e-6 -0.00011135918234874102 3.351037971510176e-5 0.00013012155627432754 2.1383115516010957e-5 -6.793230918635013e-5 -0.0001972860444330261; -0.00013415221014471672 -5.4049999513103154e-5 -3.738131540191383e-5 -1.664834685104236e-5 0.0001808157003816029 5.0218569462257723e-5 -0.0001052592442156045 -5.308593512924829e-5 -3.7414952153968095e-5 -0.00014410718990436267 -0.00016036365827490863 -5.6472114871178186e-5 1.0688315391262167e-5 -5.624344971326491e-5 -1.6087692480114764e-5 -2.3016872497436758e-5 4.1076306892168306e-5 4.289940357378313e-5 9.833118843201118e-5 6.0065421215794965e-5 -0.00017004043451922314 -1.3481069750897099e-5 -1.6542205181368303e-5 -4.5134994203152874e-5 2.237047834026188e-5 9.041157050345173e-5 1.042406546198648e-7 5.98170163848701e-5 0.00014557970739954656 -5.8829143150406665e-5 5.818028971955721e-5 5.7290370981750394e-5; -5.057351421268999e-6 0.0001221124555800287 -0.0002590668002025569 1.2469427753396631e-5 4.781427633730968e-5 -0.00012147005253921079 3.825984895827421e-5 -1.8438954787983504e-5 -8.758314100691176e-6 4.19675096114371e-5 0.00012016073088175471 0.00010762445963933615 4.2502434377251515e-5 -6.582777058922755e-5 -2.4213935097054012e-5 0.00013933064547615956 0.0001274707618054574 -3.8938301434814405e-5 -8.066222283945586e-6 9.21285282545316e-6 -0.0001421111366593568 3.4868248627913686e-5 5.450332246853216e-5 -0.0001253858346839732 8.214031623226779e-5 -5.945809534153832e-5 -1.2116630562944657e-5 0.00015861705542790326 0.00020224829568776462 6.898291104023541e-5 -1.0165088673079406e-5 5.01152215352373e-5; -2.813029794647728e-5 -0.00010227741193824684 5.5014632666865734e-5 3.322738513365233e-5 -6.195535374859645e-6 -8.355670329328238e-6 -4.773336681827701e-6 -9.40635106280382e-5 -0.00015331959337796353 7.337030939656022e-5 1.439567978654854e-5 7.347844467862829e-5 6.408162545543867e-5 9.927265301943568e-5 -0.00011662679361048429 -4.1229075291046775e-5 -4.8513665828936705e-5 6.371760202004337e-5 9.100840208251678e-5 -6.832936531733029e-5 6.19445748407469e-5 -9.731909526598561e-6 1.1040824861120845e-5 5.056355468245074e-5 -4.037221122867717e-5 -3.5550573874887513e-6 -6.970229774159986e-6 9.41176661536184e-5 -4.017932923030384e-5 -2.70782290362689e-5 -0.00012816711722276426 -7.10177079644877e-5; -0.0001389386409660647 -3.693636907843047e-5 0.0001271705738967431 3.88913947157682e-5 0.0001949394397437719 -0.00023578911649527865 -1.865647144731036e-5 0.00010796352281469007 -0.00011551400829338351 -0.00012857958362436761 -4.9469791783463776e-5 -0.00011306909008415431 -3.053109257015228e-5 4.843544657925369e-5 -2.8622448816885073e-5 -2.188231636827104e-5 3.939701556229644e-5 -2.731094108602336e-5 -8.555257526458141e-5 0.00011638328279369668 0.0001328741825185833 3.3763096062476945e-5 -1.649229970061125e-5 6.928119974910562e-5 -0.0001685227864887463 -7.203045542117499e-5 5.042434778298334e-5 0.00016450754329563506 -6.322804555528692e-5 -0.0001593481386313113 -5.325490500544057e-5 0.00014156799144113196; 6.951497583696508e-5 0.0001700126197077099 -1.9956083417753943e-5 -0.0001966591041879703 -3.006841667473811e-5 -2.8053291100771337e-5 2.3511387769137278e-5 2.0588401978939e-6 -6.147405998734337e-5 -6.840319452993121e-5 -3.552983231447501e-5 -0.0002565190332122191 3.0618139351023265e-5 3.162172309868342e-5 -1.2350958752616207e-5 -1.8256929211915125e-7 -6.183334677433707e-5 2.878763392556367e-5 0.000254471615551044 6.45311922538132e-5 0.00010893467355100728 7.084613411035816e-5 4.1444102168786145e-5 7.658827620372422e-5 -6.185452708694632e-5 0.00011490173732209806 -3.6039422195865845e-5 -3.060356153827032e-5 1.8649818446025864e-5 1.0174229485683248e-6 0.00020527303080354286 -8.52559640067649e-5; -0.00020263942307266887 5.779576721200052e-5 9.778408828949842e-5 2.1447458571649196e-5 4.012709003816651e-5 -7.667341007093629e-5 -9.363810560374239e-5 2.0308885801357973e-5 2.2398080623833017e-5 1.42630797979728e-5 -1.914872692778283e-5 -5.1480999822354054e-5 6.421183032741281e-5 4.01087473490145e-5 -8.190987311011802e-5 8.971208270465615e-5 -1.133096149189406e-5 -5.858499017834242e-5 -0.0001690939551865374 -7.12837008124828e-5 0.00010621304507892233 4.626497146616074e-5 -6.844946066324322e-5 -7.802942301553002e-5 5.767392402573755e-5 -1.6107465775547912e-5 -6.14041654573449e-5 -3.845610073242359e-5 -0.00013553827185590453 -0.0001645637839048796 0.00011255471880373922 -3.260916665951097e-5; 3.003673428688174e-5 0.00011284289085671023 3.0871632233216426e-5 -0.0001315808271021069 -6.900153802339832e-5 -0.00018175184358474062 -7.412114190733261e-5 0.0001338623174968715 9.376791794104515e-6 3.3116774474231484e-5 3.619345680711765e-5 -0.000207464510268583 -4.2727676022729213e-7 9.250052845241113e-5 0.0001735469619314723 0.00012130093692471791 0.00017727183022438528 -0.0001246467376323819 -1.0920986270454889e-5 1.0696490434249295e-5 -5.416375042036584e-5 -9.138436270509136e-5 0.00013952806203529038 -2.5700529832784106e-5 1.887538256492277e-5 2.6842090491158188e-5 4.505124762440852e-5 0.00012716222282401588 -6.904578312164384e-5 9.292096238719118e-5 -0.00011380684135578836 0.00019440068591288151; 0.00010218951387916686 1.2213936994963191e-5 9.578533240248566e-5 -2.4886715667046992e-5 0.00011589741074833199 -5.319087472623347e-5 8.777656216528124e-5 -0.00012867622104119637 -1.3357341461479076e-5 -9.501157296933366e-5 -8.933039910598327e-6 3.313362207965449e-5 -3.5256774256660375e-5 -1.0428194630615089e-5 8.426615999899642e-5 0.0001677664174593949 6.131888835345388e-5 -8.019042548137027e-5 -8.195134546719638e-5 0.00010651724619213257 0.00012543200022894534 -7.144450615039262e-5 -0.00012749478014763583 4.3829097873573615e-5 0.00010079116941317785 7.164160956283044e-5 1.3455534596623596e-5 6.114566963053647e-5 -6.925316233998321e-5 -0.00012379309305567417 -0.00010520177716268436 0.0002150600837439193; 9.098299959606051e-5 -2.6132252982443845e-5 6.134634348728216e-5 0.00019247368798988265 -8.375706785426952e-5 -8.863722724908725e-5 4.843765235437319e-5 0.00016228811181081694 -3.7102587170188615e-5 -0.0002009843574471189 1.8178363448724768e-5 1.2071945989927651e-5 -7.413476670462361e-5 -8.538063231225754e-5 -0.00015448272280463703 0.0001433426095771262 9.834693081375704e-5 3.422116968595694e-5 0.0001817091851860868 -1.7372833112126313e-5 -8.251439250142935e-5 2.9272694506112173e-5 -5.2558914843935e-5 -6.100540669121451e-5 -0.00011914352763978528 5.0899727271651344e-5 -3.659126109098036e-7 0.0001849071140766741 -0.0001454179564590914 -0.00021117630059433866 5.0579552394815536e-5 -6.046012459973273e-5; -0.0001132349432552561 -3.746346990718539e-5 6.207967018808019e-5 0.00011749132610957801 9.588242321141767e-5 8.503143837543539e-5 7.358912918817566e-5 4.013364306984818e-5 -6.763236959131189e-5 5.5746920425289925e-5 -1.7158949758352877e-5 0.00011032695985730163 6.477537610435923e-5 0.00010414532809618745 -1.7328503217587304e-5 -7.51791019681113e-5 5.963812074723167e-5 -2.575973374377371e-5 0.00010183466585686739 0.0002951133421632286 0.0001123223985089718 2.1058441434714788e-5 -3.0926003800860175e-5 2.3804233224202715e-5 5.3521994432512374e-5 -2.9509405776216954e-5 9.360333135219657e-5 -1.2535631667289092e-5 6.014476824103276e-7 -7.667771640598415e-5 1.2992224783654627e-5 -7.674277801899056e-5; 0.00010979120559532538 -9.928410583631434e-5 1.1120955768642314e-5 1.2622491503499837e-5 -8.877399233860629e-5 -2.720342747161358e-5 0.00012182638891281373 2.991394756427303e-5 -0.00014495751094764262 -6.0046903691182006e-5 -0.00010958091008468459 4.272090807681631e-5 -0.0001466325528057489 0.00012030211217243168 0.00021010792550405316 -6.990163530681321e-5 -0.0001852817120565253 -6.415124955346928e-5 -0.00019884293289615027 -0.00014836184420382724 0.00019364275356529241 -5.672776133360426e-6 0.00020367635732291078 -1.9170950345504092e-5 5.6285464942993496e-5 -1.959585171825341e-5 -0.0002252587772276459 6.244096146438385e-5 0.00012309514129774468 -6.438913698766581e-5 0.00016684749714804674 -9.34421157766277e-5; -1.4358953412496398e-5 -1.3322720624017987e-5 -0.0001059917733434615 0.0001792281516466245 -0.0001565929589921636 -2.9798732830147265e-5 -6.999281617003118e-5 -7.479458439751159e-5 9.350273900078067e-5 -3.8707224339066096e-5 -5.038482708162442e-5 -5.210798848541481e-5 -8.007262225850063e-5 -0.00014156723340396312 2.2518141813873423e-5 1.2255020701799104e-5 7.119685757085458e-5 0.00017463745894953094 1.5197007778863462e-5 -6.845867337448803e-6 7.133605391596911e-5 0.0002241574664005955 9.272316456429207e-6 -9.463443278922206e-5 -2.0652837738214087e-5 2.2166129165513264e-5 9.871337861352532e-6 -0.00011002103132317212 -7.942528758552478e-5 -2.6839351666910624e-5 -2.766977853728986e-5 9.95663474534621e-5; -4.0162191427038527e-5 7.488665462827664e-5 -6.0419818029000067e-5 3.862795919304197e-6 -1.998839840870963e-5 -8.428234693474399e-5 -3.8487782577250845e-5 -3.045706077591785e-5 4.697310171559225e-5 1.4921466144185644e-5 9.120095958588966e-5 7.112288178997e-5 9.060505138131632e-5 -0.00018214814596922226 -0.00012527461919580824 3.456540779552801e-6 -5.7960118711884595e-5 -0.00013506219372266145 -6.002592314962551e-5 4.058176594207818e-5 -0.00010195672482132051 -7.286623232177259e-6 3.216711368531662e-5 3.0428439939897328e-5 1.0986999340273563e-5 6.00495145854493e-5 -6.016249287404396e-5 -2.4559560616578574e-5 -4.516794476421796e-5 3.9214944562443595e-5 -4.955457231608419e-5 0.00012331033533792533; -7.046708939830794e-5 4.923773652171303e-5 -0.00012298102971632866 0.00021497035034466098 -2.8711207678217744e-6 5.197374760522351e-5 4.496704044053165e-5 0.0002629294611606732 -6.309834877262089e-5 5.202048835693705e-5 -0.00010623679569924317 -1.4432655160795168e-5 -7.885356914717981e-5 5.106321971746731e-5 5.489770394982586e-5 -7.834456498036546e-5 0.00016829489307477116 -8.304085365031177e-5 -0.00010564803248504344 -0.00010120576204732318 8.721843171864365e-5 -5.068935221876206e-5 0.0001180858306097782 -9.496788713382999e-5 8.890039301244388e-5 -5.704542506370762e-5 0.00010318914412986344 2.716094225528038e-5 9.479690182255404e-5 -1.2120684158582301e-5 4.179147056769349e-5 -1.8310054810146353e-5], bias = [1.979980030356884e-9, 3.785114049758386e-10, -2.5186046980311424e-9, -2.2921624259912312e-9, 1.9905413324006684e-9, 2.299035155572659e-10, 1.4817042119291613e-9, -1.232451864521743e-9, 2.0007252974429652e-9, 9.995626240397066e-10, -1.8957755038418508e-9, 4.012136960745953e-10, 2.24117193285951e-9, -1.6221677134588285e-10, -2.2384401222938604e-9, -1.2571801936180382e-10, -2.703411160764293e-9, -3.4886871332112935e-9, -9.26725961996573e-10, 1.8779306413243247e-9, -7.309472124584666e-10, -9.85246141545734e-10, 1.0515917138426022e-9, -2.1117686920800426e-9, 1.432098409230967e-9, 1.519535480741874e-9, -4.92959038824215e-10, 3.3244035056991004e-9, -1.0230181945273302e-9, -6.432307940662419e-10, -1.2991868080193617e-9, 1.5277974293306535e-9]), layer_4 = (weight = [-0.0007311498608190174 -0.0007896852101794138 -0.0006334016116668214 -0.000601173050119404 -0.0005053747781325372 -0.0008229200702320116 -0.0006119400654223587 -0.000654628109674041 -0.0008152974503057086 -0.0006157818688392291 -0.000558682465426874 -0.0008344637519756282 -0.0007018638599351963 -0.0007187659310758401 -0.0007534004272321894 -0.0007664203469270862 -0.0005223131336777111 -0.0008048523378003989 -0.00074099054439484 -0.0006690574511233565 -0.0004981274354857056 -0.0007260532038179166 -0.0008372606583159075 -0.0006949696629397304 -0.0009475679568020365 -0.0008255908443002022 -0.0004653342953853413 -0.0007272366964570933 -0.0006582114485809905 -0.0006201242313122631 -0.000695487307548721 -0.0006935628471232684; 0.00011228636689403971 0.00030855719342611854 0.00018397366633587692 0.00011469685005575407 0.00021842836661456285 0.00020522084671247335 0.0003835340868287802 0.0002277441952441551 0.0002955906727297834 0.00011955036803063669 0.00038513011105919333 0.0002174977723923004 0.000147357306915197 0.0003359118123489575 0.00038174876393228354 0.0002766815793809704 7.653613511583295e-5 0.000367425172886657 0.00019746221443046392 0.00021880902149437454 0.0003562821889114495 0.00023333172554683242 0.0003359704648802438 0.0003332221658847545 0.00021441240006290767 0.000237864530229569 0.00020239696559389106 0.00019960291464459322 0.00026177313901098906 0.0001090543836542154 0.00018055512420855682 -1.0297415807600106e-6], bias = [-0.0006495439263193634, 0.00025776835530455207]))

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 9V74 80-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-18.1.7 (ORCJIT, znver4)
  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