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

Precompiling OrdinaryDiffEqLowOrderRK...
   4050.3 ms  ✓ OrdinaryDiffEqLowOrderRK
  1 dependency successfully precompiled in 4 seconds. 98 already precompiled.

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 $r = r_{1} - r_{2}$ and use Newtonian formulas to get $r_{1}$ , $r_{2}$ (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

one2two (generic function with 1 method)

Next we define a function to perform the change of variables: $(χ (t), ϕ (t)) \mapsto (x (t), y (t))$

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

soln2orbit (generic function with 2 methods)

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

d_dt (generic function with 1 method)

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

d2_dt2 (generic function with 1 method)

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 BH 1 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

compute_waveform (generic function with 2 methods)

Simulating the True Model

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

u [1] = χ

u [2] = ϕ

where, $p$ , $M$ , and $e$ 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

Defiing 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 incase 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[7.860439f-6; -5.2525433f-5; 0.00013166785; 7.766766f-5; -0.00014561972; -0.00017803736; -2.7755737f-5; 0.00010047881; 6.7428086f-6; -0.0001594555; -9.549311f-5; 4.0078412f-5; -9.89078f-5; 0.00018193704; 8.555898f-5; -4.970338f-5; -3.8765254f-5; -1.3489131f-5; -0.00013164493; 4.297612f-5; 1.1798019f-5; -0.00011137081; 3.6243822f-5; -9.186421f-5; -6.508293f-5; 4.663855f-5; -1.4155018f-5; 0.00016139026; -0.0001249956; -8.825055f-5; 6.864034f-6; -0.000113623086;;], 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.00016437678 -3.0016356f-5 0.00021092324 -1.9295398f-5 6.400007f-5 3.5518326f-5 -7.671253f-5 7.6223337f-6 -1.7718496f-5 -0.00013926686 -1.9196865f-5 -0.00014257817 4.0983534f-5 7.9018006f-5 -0.00010902154 -8.941748f-5 8.6309075f-5 0.00011615165 0.00013435379 -7.870236f-5 -8.417054f-5 -6.1048006f-5 -0.000110034176 -1.0304273f-5 0.00010379674 0.000118325755 0.00018107906 -5.199068f-5 -0.00013425709 -1.3520461f-5 -7.018773f-5 -8.1540806f-5; -0.0001138165 0.00013198743 6.216527f-5 0.00010475709 -0.00012259047 -3.5811176f-5 -5.9808615f-5 4.6188074f-5 0.0001017403 8.599425f-5 2.4665938f-5 0.00016514427 0.000119936565 0.00017141782 0.00018569987 8.3548126f-5 -0.000119457975 -0.000102200815 -2.5700943f-5 -3.5522182f-6 -0.00014269249 -1.3331167f-5 9.086315f-6 -3.1403084f-5 -0.00010148137 -8.164521f-5 -7.184492f-5 9.113478f-5 -5.43189f-5 -0.0001400215 -3.2325417f-5 -6.1269043f-6; 9.969454f-6 5.1681636f-5 0.0001323837 -6.091742f-5 -2.84636f-5 -5.0008377f-5 -1.2507815f-5 -6.465501f-5 1.4994268f-6 -0.00013679356 2.4531822f-5 0.00012220755 0.00018068936 4.203786f-5 -3.4607397f-6 -0.00023079172 5.7677495f-5 0.00012361121 9.097084f-5 0.0001160965 -0.00015204407 0.00011802255 -4.591044f-5 3.2231765f-5 6.239905f-5 -3.0868632f-5 -0.00014959059 9.025776f-5 -4.874964f-5 3.5693734f-5 4.917927f-5 0.00010072045; -2.0711632f-6 -3.1827927f-5 0.00013867537 0.00013392804 8.84111f-5 6.3304474f-6 -4.025227f-5 -0.00011997796 -0.0001582737 8.382454f-5 3.6359375f-5 -9.078856f-5 -4.752343f-5 0.00019082347 0.000115616705 7.918959f-5 -6.270233f-5 -0.00022068835 0.00016179885 2.8037526f-5 9.4876465f-5 5.1381292f-5 4.4210457f-5 7.896579f-5 -6.684568f-5 5.9662198f-5 -7.461792f-5 1.0225014f-5 -4.5663954f-5 -2.357594f-5 -7.383484f-5 -3.5562636f-5; -0.000116719355 0.00016689862 -0.000100075136 8.6156906f-5 0.00010011918 2.8148857f-5 5.4008928f-5 7.3757816f-5 6.9377784f-5 5.9282804f-5 1.8332275f-5 -0.00015216526 0.0002312166 4.3590888f-5 -0.00019778108 -3.14341f-5 -3.1108688f-5 0.00014344706 -0.00011801217 2.2652677f-5 -1.20078785f-5 0.0001323587 -2.9821753f-5 -0.0002715281 -5.2604533f-5 -7.438716f-5 -0.0002244406 0.000118518656 1.0915055f-5 7.141488f-5 7.474548f-5 -0.00011599659; 0.00018576678 9.744388f-5 -2.086525f-5 -0.00010667362 2.5712141f-5 -1.9456742f-5 -1.2227168f-5 7.856932f-6 0.00014285228 1.5197984f-5 0.00010530933 -9.205047f-5 5.8769456f-5 -0.00027841583 1.1447948f-5 -5.0622653f-5 3.1692068f-6 -5.100443f-5 -5.075083f-5 0.00010642333 -3.5123762f-6 -9.171245f-6 -5.301154f-5 -9.309107f-6 -0.0001233679 -0.00015029332 -0.00016648695 -3.7564741f-6 0.00010693836 -4.192232f-5 -6.602926f-5 -1.7919037f-5; -0.00010453267 0.00015391129 9.915401f-5 8.348953f-5 -7.9054924f-5 5.6505087f-5 0.00012845981 1.7698028f-5 1.3949116f-5 -9.806561f-5 -6.420396f-5 0.0002217041 4.7849775f-5 -5.236996f-5 4.9019938f-5 -6.746362f-5 0.00014889111 5.7057696f-5 -0.00010691402 0.0001795911 -4.0313258f-5 8.079401f-6 7.2666726f-6 -5.686987f-6 6.0907998f-5 -5.126703f-5 8.3842715f-5 0.00016367433 -5.5189397f-5 -0.0001455218 -0.00010544503 -0.00015418504; 0.00022932114 9.020424f-5 3.0043677f-5 -5.2220155f-5 0.00010644842 -5.0049846f-5 3.6261197f-5 -3.728025f-5 2.4976409f-5 -0.00023991911 -0.00021004882 -2.6648035f-5 -0.00015839229 4.5194247f-5 -0.00010322864 -9.639724f-5 0.00012596064 1.3371975f-5 -0.00015947327 -6.847464f-5 9.8559f-5 -3.992497f-5 -3.391051f-5 3.467761f-5 -5.4849126f-5 -4.343793f-5 -0.00010100763 -0.00013983452 0.00018997362 -0.00026679982 -0.00013637258 5.1325955f-5; -7.251158f-5 -2.21497f-5 -5.078739f-5 8.177907f-5 -0.00014436334 6.0707163f-5 -4.0421226f-5 0.0001102568 0.00015287085 6.8625064f-5 -7.154448f-5 7.953224f-5 0.0001660816 0.00014572355 0.000100684876 -5.7809317f-5 -0.00022931145 -0.00020820695 -6.144196f-5 -0.00012487781 -5.33979f-5 9.256115f-5 1.6038002f-5 0.0001195173 4.6324552f-5 -0.00027295144 7.2264105f-5 8.911913f-5 -6.401373f-7 0.00014419894 0.00012912118 1.9567065f-5; 6.820378f-5 2.12501f-5 0.0001519751 0.0001277628 -4.5427285f-5 9.45922f-5 2.8407487f-5 3.679055f-5 -9.1179245f-5 -8.759315f-5 -4.233391f-5 -1.0148495f-5 -2.8013366f-5 0.00019726023 4.0231207f-5 -2.0327785f-5 0.000199447 -0.000111626985 -0.00013777778 -0.000117939824 -6.2143245f-6 0.00020680211 -1.3994887f-5 8.546733f-5 0.000104675535 3.492703f-5 -8.7907225f-5 0.00010485748 7.1407274f-5 0.00011894936 0.00013858061 -2.8788549f-5; -1.5867355f-6 4.1940072f-7 0.00021171928 4.52169f-6 6.012732f-5 9.992477f-5 -3.561707f-5 -6.138835f-5 -0.00017997247 6.2459614f-5 -1.34978145f-5 -1.5011756f-5 -0.00012783983 0.00011435254 0.00019095208 2.1907656f-5 -4.116716f-5 -1.17440095f-5 7.0702874f-5 -9.540582f-5 7.0204645f-5 0.00013954662 0.00010261999 -0.0001211858 -7.0892784f-5 -0.0001676153 -1.7646533f-5 -3.4530327f-5 -2.2219998f-5 -2.18182f-5 -4.4124412f-5 0.000111558154; 0.00016322506 0.00014593737 3.3885233f-5 0.0001275012 8.391951f-5 0.00018750902 8.676476f-6 0.00020314452 2.0860696f-5 1.0924117f-5 -0.0001489237 -3.377215f-5 0.0001377311 -6.447554f-5 1.8870135f-6 -1.525362f-5 -0.00018388667 -0.00020160212 7.516864f-5 -8.7843495f-5 -9.126322f-5 -1.8796603f-5 -1.542366f-5 2.639844f-5 -5.7380897f-5 -4.7694493f-6 0.00019579113 0.00010669442 -2.5035256f-6 -5.628761f-5 -0.00012336455 -1.8165401f-5; 4.386997f-5 3.080651f-5 0.00018174354 -2.854481f-5 -5.467102f-5 3.1228672f-5 5.778987f-5 -3.1693147f-5 -5.9755175f-6 7.561135f-5 4.518769f-5 6.331077f-6 -5.9239874f-6 -0.00021513706 1.7105978f-5 -2.608389f-5 6.7201822f-6 9.437318f-5 9.7142896f-5 3.587049f-5 2.990647f-5 9.455785f-5 0.00013769187 7.590747f-5 0.00012126878 -4.2956963f-5 3.1407562f-5 0.00016296902 -2.3624254f-5 5.0848314f-5 2.0582893f-5 -9.444387f-5; -9.685379f-5 4.364414f-5 -8.146248f-5 0.00012298896 -0.0001297601 5.0669637f-6 -0.00020389126 -8.497968f-5 -3.2198812f-5 5.7340447f-5 5.6329452f-5 0.00012570899 -6.903591f-5 9.564116f-5 0.0001340695 -4.0525847f-6 -0.00014826849 -9.674515f-5 -0.00012002463 -1.9548655f-5 -6.245278f-5 -0.00010349815 7.256741f-6 5.4170167f-5 -9.72962f-5 -0.00011858146 -4.4183424f-5 3.905559f-5 7.1765775f-5 2.2052234f-5 -9.916257f-5 -0.00018513115; 2.0909301f-5 -4.8683356f-5 -8.092678f-5 8.124017f-5 3.441194f-5 -0.00010688532 -0.00017534495 -9.5536416f-5 5.324465f-5 1.4815658f-5 -2.5158256f-6 2.1041446f-5 -0.00013107157 0.00015172818 9.03989f-5 1.0895201f-5 -3.224868f-5 -0.00013421122 5.294083f-6 3.5997127f-5 6.929181f-6 2.8395696f-5 2.9743112f-5 -0.0001417928 6.4311724f-5 -0.00013529892 0.00019494956 1.6890172f-5 -0.00019581386 -0.0001004215 0.00013518709 0.0002651039; 0.00011439457 -0.00017873818 -8.155978f-5 -2.3657272f-5 5.2572363f-5 1.871703f-5 -8.4813604f-5 -0.00020245861 -6.970776f-5 0.00015375929 -0.00022528024 1.1745393f-5 2.5853018f-5 -1.4499828f-6 0.00012424964 2.6088617f-6 -3.924224f-5 -1.40888005f-5 -0.00012934036 -4.0908646f-5 -0.000109726425 7.257933f-5 0.00026414773 8.74527f-6 3.3685497f-5 9.373089f-5 -2.3758246f-5 3.0599972f-6 -9.0555855f-5 0.00014242473 -6.925916f-6 9.393081f-5; 0.00018411597 -0.00010941034 0.00015787876 9.781739f-5 0.00012580647 -3.7076934f-5 5.9199956f-5 5.7521836f-5 6.913701f-5 -1.6429325f-5 -8.043186f-5 -6.6614945f-7 6.404517f-5 0.0001124876 -8.0944104f-5 1.560465f-5 -0.00016145661 -0.00020957802 0.00012761129 0.00012393577 8.609023f-5 0.00027474988 -9.616394f-5 -9.7552016f-5 9.3259354f-5 8.2833896f-5 0.00016256991 5.094344f-5 -0.00016081546 0.00013049027 2.501641f-5 9.764774f-6; 4.4500546f-5 -0.00011666155 -5.058511f-5 -0.0001115021 -0.00020778585 4.6901692f-5 2.2114928f-5 -2.4246567f-5 -7.744055f-6 1.7532006f-5 -9.8236895f-5 -4.661263f-5 7.069244f-5 -2.470423f-5 8.236617f-5 -1.9313466f-5 7.0327774f-6 6.528813f-5 -6.1844024f-5 5.0168937f-6 0.00016198569 -0.00014750911 -3.796759f-5 7.6624834f-5 0.00013080529 0.000107561624 8.559652f-5 0.00015994655 3.5554945f-5 9.201267f-5 -4.8966453f-5 -0.00013019545; 0.00010356599 9.386258f-5 -9.555804f-5 -0.00011302738 -2.982038f-5 -2.6537766f-5 0.0001326902 1.3503379f-5 0.0002321677 -0.00014570077 -6.160482f-5 9.383279f-5 -0.00012899985 0.00026124864 7.559704f-5 -3.245845f-5 2.2868871f-5 7.393422f-5 6.939985f-5 -8.073464f-5 -6.493277f-5 -7.971229f-5 -3.690797f-5 7.159614f-5 0.0002240174 -4.9012317f-5 -6.309731f-5 6.5474334f-5 -3.9521747f-5 2.5880536f-5 9.168814f-6 7.0612776f-5; 3.6504764f-5 -0.00014354469 -0.000107286105 0.00016915759 6.8394984f-5 -3.302745f-5 -0.00015622316 0.00012367406 4.196102f-5 -0.00024222277 2.7071996f-5 -0.00016572379 -1.2416313f-6 9.093682f-5 0.00013993763 -2.9302382f-5 -4.238206f-5 6.964645f-6 0.00017332078 -0.0001184762 2.8117662f-5 -0.00012941372 -2.5065112f-5 1.2464733f-5 -0.00010473354 0.0001495244 0.00016974339 -1.9494899f-5 -3.4896857f-5 -5.3749554f-5 -3.168685f-5 0.00017349978; -5.132354f-6 -4.424082f-5 -4.964353f-5 -1.5468098f-5 -2.581242f-5 3.4385128f-5 -0.00016474552 3.214969f-5 -2.6949356f-6 -7.3306846f-6 -0.00023628463 -2.713697f-5 4.682057f-5 -5.0788974f-5 -6.5953805f-5 -2.250798f-5 -0.00015827961 -8.265855f-6 0.00013737188 -6.6480216f-6 3.0908624f-5 2.7139635f-5 0.00023415721 -8.188687f-5 -6.4403386f-5 -8.165859f-5 -9.382276f-5 -2.3535324f-5 6.069784f-5 -7.942674f-5 0.00018662795 2.8620603f-5; -0.000110345616 4.7111524f-5 4.532873f-5 -5.5340144f-5 -0.00020332819 -1.6361653f-5 -8.1452934f-5 0.00020983846 -1.7723463f-5 6.42776f-6 -0.00011324887 -4.4252903f-7 -8.6307926f-5 -6.0750586f-5 3.6340498f-5 -0.0001346239 1.4279817f-5 -0.00015716022 -0.00011565898 0.00010061201 -4.9437135f-6 3.4787143f-5 -4.7063295f-5 -0.00018093284 1.5014159f-5 -4.8186812f-5 0.00010728006 -4.107313f-6 8.443386f-6 0.00018245523 0.00015710696 0.00010806732; -4.593428f-5 3.7794541f-6 0.0001901828 -0.00015562437 0.00010617965 2.7162823f-5 -0.0001961611 -6.838603f-5 0.00010460149 1.8424496f-5 -1.7457225f-5 0.00010073535 -1.3157878f-6 4.764597f-5 -7.2148105f-5 -6.362214f-5 -8.791657f-5 2.551204f-5 -7.157915f-5 7.802277f-5 -6.758359f-5 0.00011234907 -2.9502333f-5 4.234131f-5 1.3768293f-5 0.00015046699 -2.074506f-5 1.5611389f-6 -8.091321f-5 -2.2058419f-5 -9.851729f-5 4.396451f-5; -7.404357f-5 -0.000116137926 -3.919746f-5 6.9903613f-6 0.00015672618 -0.00012052637 -5.075656f-5 -4.31624f-5 -9.2455695f-5 3.0230767f-6 7.4591815f-5 -0.00017187528 8.2586484f-5 1.3629483f-5 4.858484f-5 -1.0701159f-5 -9.76107f-5 8.8346635f-5 5.066664f-6 0.00011652417 -0.00016424122 6.7113804f-5 6.106476f-5 9.281236f-5 3.989169f-5 -0.00013496888 9.74449f-5 -0.000106641215 0.0002279588 0.00017059782 -7.062001f-5 0.00013929975; -4.043872f-5 -4.6237346f-5 -1.505656f-5 -8.941849f-5 -7.745332f-5 -1.8620416f-5 -6.3493855f-5 8.04013f-6 4.0744246f-5 5.7554855f-5 1.9408004f-5 4.646881f-5 0.00020582223 -0.00015895322 0.00014299557 0.00029195202 -6.145626f-5 -5.6737663f-5 -0.00015056907 -8.007406f-5 3.0197156f-5 -0.000121206926 -0.00016406784 -0.0001453998 -7.0228336f-5 -0.00016744592 -2.524225f-5 -3.979123f-5 -5.5559358f-5 3.8262624f-5 1.5647951f-5 0.00022073666; 5.5627712f-5 6.8776084f-5 -7.352222f-5 1.0002345f-6 9.808353f-5 -0.000106885396 -7.6797754f-5 3.6088128f-5 -3.8107824f-5 -0.00020009115 -3.370073f-5 7.5047414f-5 -0.00011180497 3.9220264f-5 3.158692f-5 1.054442f-6 -0.00023945759 -0.00012433676 -2.6348956f-5 -5.522512f-5 -0.00022014593 -0.000115676325 -3.7359912f-6 -7.85893f-5 -7.1170274f-5 -3.0461179f-5 8.644194f-5 -2.413197f-5 -5.5292185f-5 0.00013858915 5.4612905f-5 -1.0997543f-5; 2.031275f-5 6.189822f-5 -0.00016143479 -8.91556f-5 1.1044021f-5 1.8451106f-5 7.726851f-6 0.00010280489 0.00027500797 -0.00015673127 8.405811f-6 -0.00011637053 7.788022f-5 8.792073f-5 -2.3665384f-6 -0.00023994221 9.9212775f-5 5.095716f-5 -0.00016145442 7.010363f-5 -9.213432f-5 5.1421102f-5 0.000120278026 -0.00019157163 3.3669326f-6 2.5090276f-5 -0.00010461801 8.769612f-6 8.415159f-5 0.00012264286 -6.644246f-5 -2.309465f-6; -8.704746f-5 8.5311745f-5 4.154299f-5 0.0001244581 0.0002487938 -5.0771534f-5 2.9506034f-5 7.959508f-5 -6.330686f-5 8.9787296f-5 0.000145846 -1.31013385f-5 -4.477698f-5 -0.00015706575 -2.0324507f-5 2.814357f-5 7.020316f-5 1.8658266f-5 3.0784297f-5 0.00023089771 -0.00014576453 -6.0703285f-5 0.00017205071 -4.7142464f-7 -3.542533f-5 -8.493326f-5 7.959889f-5 0.00012569474 -0.00010559724 -0.00021669785 0.00016136313 0.0001259058; -1.6658827f-5 -3.0758212f-5 3.1607575f-5 5.1340565f-5 -0.00016384062 -5.1519546f-5 -2.1732601f-5 -6.084497f-5 3.4906643f-5 5.818177f-5 -4.396652f-5 0.00018607815 -8.900506f-5 0.000112422706 -7.666776f-5 3.1738666f-5 1.0709513f-5 -9.913219f-5 6.0799746f-5 -1.1407148f-5 -0.00013092186 -1.05260415f-5 1.4557312f-5 -5.050489f-5 -0.00010791435 -7.1867944f-5 -0.0002690436 -3.3812503f-5 -2.2607754f-5 -1.2320319f-5 -0.0001482694 2.9437142f-5; 6.869186f-5 -1.277427f-8 -6.626843f-5 -2.1587384f-5 7.173915f-5 -3.0154533f-5 3.0754265f-5 -5.243435f-5 -0.00017200179 4.5476485f-5 -9.203399f-5 -0.00013684381 7.4456216f-6 2.2765864f-6 5.940914f-6 -0.000103125516 -0.00010338021 7.6328126f-7 -4.4208286f-6 3.710244f-5 -1.5903246f-5 -0.00020041747 5.2565396f-5 1.0064219f-5 0.00013255388 0.00018796537 -1.46149805f-5 4.7379745f-5 -5.66258f-5 7.8980884f-5 0.00012663766 4.566003f-5; 6.0409773f-5 0.00010700812 -0.0001300802 0.00013264797 -9.973728f-5 0.00018789402 6.6095f-5 0.00029767348 -0.00012588724 -3.6550168f-6 2.0192834f-5 9.936858f-6 4.4869426f-5 3.8916845f-5 -7.401763f-5 -0.00018950703 -0.00011452327 -3.1347794f-5 7.583299f-5 -1.9396692f-5 -0.00015249474 -3.6658774f-5 1.4171241f-5 0.00012923034 0.00012168421 -9.645912f-5 -8.4772786f-5 -0.00020042581 3.073528f-5 6.52481f-5 2.0169764f-5 -8.7430744f-5; -0.00015594509 -6.431986f-5 8.600898f-5 -6.847904f-5 1.2529786f-5 -8.806717f-6 3.4579807f-5 -2.468049f-5 -8.22499f-5 -9.3672985f-5 -9.418263f-6 -5.1334155f-5 0.0001442151 -4.4076656f-5 8.811325f-5 7.1457915f-5 -5.6878343f-5 0.00011270518 5.951645f-5 0.0001358126 -7.364864f-5 -0.00011210586 8.181958f-5 6.882281f-5 -5.7000256f-5 7.8650046f-5 1.134115f-5 4.324312f-5 -0.00015913602 -1.3885723f-5 3.2554562f-5 -0.00013899442], 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[4.720289f-5 -6.1033745f-5 -0.000117455354 6.339205f-5 9.4987015f-5 0.00010637434 -0.00014580597 -7.928112f-5 0.00012856306 5.526573f-5 5.67344f-5 -3.319828f-5 1.4563848f-5 -1.8394529f-5 -6.5796544f-6 -9.174684f-5 3.709855f-5 5.6026445f-5 -5.2116233f-5 -3.3011038f-5 -0.0002144286 -6.186036f-5 0.00013419699 -3.4802364f-5 6.544735f-5 7.683924f-5 4.327698f-5 0.00013636211 -3.126764f-5 -4.0090868f-6 -0.0001071095 -8.382162f-5; 2.4680015f-5 -9.9803234f-5 3.5336463f-5 0.00012122493 2.7823971f-5 0.00014654118 -3.1051786f-5 7.1056456f-6 9.648971f-5 8.8834306f-5 -8.428908f-5 2.0494661f-5 -0.00013468543 -0.00011850743 0.00015842692 -0.0001187949 0.00015368695 -3.531974f-8 2.317611f-5 -5.6460435f-6 7.135657f-5 -3.1484447f-5 0.000107748965 -8.069987f-5 -4.6652836f-5 -0.00012347964 -0.00013556932 -2.4259138f-5 -9.283368f-5 -0.00015578445 7.999232f-5 -6.526465f-5], bias = Float32[0.0, 0.0])), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple()))

Similar to most DL frameworks, Lux defaults to using Float32, however, in this case we need Float64

julia

const params = ComponentArray(f64(ps))

const nn_model = StatefulLuxLayer{true}(nn, nothing, st)

StatefulLuxLayer{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

u [1] = χ

u [2] = ϕ

where, $p$ , $M$ , and $e$ 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

ODE_model (generic function with 1 method)

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

loss (generic function with 1 method)

Warmup the loss function

julia

loss(params)

0.0007120370537351161

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

callback (generic function with 1 method)

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 = [7.860438927306752e-6; -5.2525432693069914e-5; 0.0001316678535656434; 7.766766066189403e-5; -0.00014561972056968258; -0.00017803735681794764; -2.7755737391948064e-5; 0.00010047880641644231; 6.742808636772987e-6; -0.0001594555069458531; -9.549310925642917e-5; 4.0078411984731166e-5; -9.890780347622999e-5; 0.00018193703726856453; 8.555898239135396e-5; -4.970337977284799e-5; -3.876525443043724e-5; -1.3489130651576472e-5; -0.0001316449342992339; 4.2976120312192544e-5; 1.1798018931580595e-5; -0.00011137081310125814; 3.624382225096552e-5; -9.186421084440545e-5; -6.50829315417017e-5; 4.663854997486031e-5; -1.4155018106957071e-5; 0.00016139025683481036; -0.00012499559670698556; -8.825054828772068e-5; 6.8640338213228945e-6; -0.00011362308578100378;;], bias = [3.1690976299135775e-18, -4.722969324672089e-18, 2.7588708788696755e-16, 8.7796646292919e-17, -1.4630496965090196e-16, -1.6769321880840118e-16, -5.820014156550568e-17, 6.977471489405705e-17, -2.890529698946684e-19, -1.2224450295050418e-16, -7.643036073865617e-17, -9.501276662725863e-18, -1.0611435143202708e-16, 9.369849162608388e-17, 2.5226215178435307e-17, 1.3438622433120839e-17, -4.123813299284313e-17, -2.6249650815877978e-18, -1.913405972317005e-16, 4.5660264392232926e-17, 5.99414400413261e-18, -3.074086997943475e-16, 1.7365310931321055e-17, -1.1415485472436397e-16, -1.5490152008516075e-16, 2.986866709460698e-17, -2.160552188993909e-17, 3.833659073928962e-16, 1.337248393465031e-16, -4.397377880384665e-17, 7.926768205002992e-18, 3.816684556720878e-18]), layer_3 = (weight = [-0.00016437806084832564 -3.001764094188149e-5 0.00021092195924757385 -1.9296682925416568e-5 6.399878725097083e-5 3.551704025631777e-5 -7.671381650107335e-5 7.621048322107302e-6 -1.7719781002601727e-5 -0.00013926814904515887 -1.9198150088414602e-5 -0.00014257945190729893 4.098224848894425e-5 7.90167202885735e-5 -0.00010902282270217395 -8.94187662688671e-5 8.630779006670836e-5 0.0001161503630519248 0.0001343525025240552 -7.870364538939978e-5 -8.417182585072153e-5 -6.104929093817117e-5 -0.00011003546137908805 -1.0305558365872467e-5 0.00010379545232899278 0.00011832447011874692 0.00018107777099975937 -5.199196425928032e-5 -0.00013425837574140834 -1.352174627209143e-5 -7.01890141990633e-5 -8.154209104213711e-5; -0.00011381601295711315 0.00013198791778136162 6.216576214369268e-5 0.00010475758092666106 -0.0001225899811046961 -3.5810686016207306e-5 -5.9808125728022707e-5 4.61885632162634e-5 0.00010174078700088122 8.599474160887054e-5 2.4666427931650673e-5 0.00016514476221941542 0.00011993705477520825 0.0001714183111530878 0.0001857003611135487 8.354861559068904e-5 -0.0001194574849009314 -0.00010220032497586527 -2.5700453691981974e-5 -3.5517286264167534e-6 -0.00014269199680115872 -1.3330677278147003e-5 9.086804608659799e-6 -3.1402594389375056e-5 -0.00010148087828697545 -8.16447188057738e-5 -7.184443318096766e-5 9.113527114690955e-5 -5.4318412070144876e-5 -0.00014002100731290844 -3.2324927518356876e-5 -6.126414676557068e-6; 9.971092314449045e-6 5.168327419720576e-5 0.00013238534387833572 -6.0915782395239954e-5 -2.8461962809749666e-5 -5.0006739055474116e-5 -1.2506177025566816e-5 -6.465337270005083e-5 1.501064790649838e-6 -0.00013679192403164993 2.4533460365418516e-5 0.00012220918955914726 0.00018069099400359748 4.203949805864354e-5 -3.459101716216215e-6 -0.00023079007865693416 5.7679132570559695e-5 0.00012361285275016106 9.097247761142576e-5 0.00011609813458686727 -0.00015204243305437138 0.00011802418970670021 -4.590880331182789e-5 3.223340258169601e-5 6.240068700332101e-5 -3.086699421931972e-5 -0.0001495889510430383 9.025939738544304e-5 -4.874800202619058e-5 3.569537233612217e-5 4.918090866464321e-5 0.00010072208619145118; -2.0699493046023594e-6 -3.182671328369754e-5 0.00013867658587576022 0.00013392925636315366 8.841231344438132e-5 6.331661252869799e-6 -4.0251057115999584e-5 -0.00011997674421702097 -0.00015827248837860415 8.382575349121165e-5 3.636058924774163e-5 -9.078734849874109e-5 -4.752221623614002e-5 0.0001908246837689771 0.00011561791931915623 7.919080114396076e-5 -6.270111892201181e-5 -0.00022068713708842295 0.00016180005938683768 2.8038740011363268e-5 9.487767844864978e-5 5.138250587213179e-5 4.4211671293562466e-5 7.896700723966134e-5 -6.684446391711295e-5 5.966341147603265e-5 -7.46167073242006e-5 1.0226228069469978e-5 -4.566273976645236e-5 -2.3574725973340262e-5 -7.383362511001644e-5 -3.5561422309295024e-5; -0.00011671944717607035 0.00016689853022610273 -0.00010007522847233795 8.615681385192531e-5 0.0001001190854752392 2.8148765197624996e-5 5.4008835506760133e-5 7.37577235143584e-5 6.937769161806867e-5 5.928271151999513e-5 1.8332183246752113e-5 -0.00015216535153529177 0.00023121651124013114 4.359079561465944e-5 -0.000197781167442077 -3.143419358609743e-5 -3.1108780019781914e-5 0.0001434469692342864 -0.0001180122611528308 2.2652585032414297e-5 -1.200797065182629e-5 0.00013235861355708443 -2.9821845016773224e-5 -0.0002715281783436473 -5.260462545054527e-5 -7.438724911269299e-5 -0.00022444069814598694 0.00011851856346977445 1.0914962915489514e-5 7.141478512705753e-5 7.474539110473782e-5 -0.00011599668265050746; 0.00018576495818151595 9.744205777124159e-5 -2.086707213137709e-5 -0.00010667544524594361 2.5710317962796483e-5 -1.9458564971573887e-5 -1.2228990677704538e-5 7.855108882636743e-6 0.00014285045795431792 1.5196160851624561e-5 0.0001053075061953723 -9.20522911165789e-5 5.876763275764955e-5 -0.00027841765081218847 1.1446125021306062e-5 -5.062447639475007e-5 3.1673837874223278e-6 -5.100625316672525e-5 -5.075265330205151e-5 0.00010642150623780278 -3.5141992576435353e-6 -9.173068470857155e-6 -5.3013362454385145e-5 -9.310929677754246e-6 -0.00012336972981401976 -0.00015029513856106813 -0.00016648877176167004 -3.758297176394942e-6 0.00010693653489753093 -4.192414457620167e-5 -6.603108017416245e-5 -1.792086045142471e-5; -0.00010453094441133783 0.00015391301272513888 9.915573460358847e-5 8.349125997240086e-5 -7.905319657100862e-5 5.650681410525768e-5 0.00012846153836774366 1.7699755660324737e-5 1.3950843055979353e-5 -9.806388410896795e-5 -6.420223220873438e-5 0.00022170583102463424 4.785150227446258e-5 -5.2368232757265255e-5 4.902166543372262e-5 -6.746189032371526e-5 0.0001488928348019967 5.705942308607271e-5 -0.0001069122925788937 0.0001795928211843841 -4.0311530795109046e-5 8.081127959776209e-6 7.2683998562993805e-6 -5.68525987789858e-6 6.090972525028043e-5 -5.126530308432299e-5 8.384444223095423e-5 0.0001636760527229046 -5.518766997073027e-5 -0.00014552007345871738 -0.00010544330584041074 -0.0001541833088073696; 0.0002293175000499449 9.020059383823668e-5 3.004003315933052e-5 -5.2223798995832364e-5 0.00010644477311640511 -5.005348997885512e-5 3.625755358138017e-5 -3.728389332768434e-5 2.4972765238488035e-5 -0.00023992275804860482 -0.00021005246366489718 -2.6651678347261686e-5 -0.00015839592946808743 4.519060326627472e-5 -0.00010323228183783848 -9.640088523388743e-5 0.00012595699241268427 1.3368331344680036e-5 -0.00015947691849081775 -6.847828466794432e-5 9.855535768787195e-5 -3.99286129709671e-5 -3.3914153611986386e-5 3.4673966872513915e-5 -5.4852769828864104e-5 -4.344157263636527e-5 -0.00010101127394824462 -0.00013983816301563646 0.00018996997983545798 -0.00026680346331209136 -0.00013637621872709328 5.1322311044218384e-5; -7.251044679850486e-5 -2.214856375775544e-5 -5.078625271398424e-5 8.178020708310985e-5 -0.00014436220170882525 6.0708298960706135e-5 -4.042009039966561e-5 0.00011025793411118394 0.0001528719882206232 6.862619976327607e-5 -7.154334106434613e-5 7.953337681993639e-5 0.00016608274073387928 0.00014572468408896152 0.00010068601169205792 -5.7808181181999396e-5 -0.00022931031652678 -0.00020820581939023136 -6.14408267121654e-5 -0.00012487667790100188 -5.339676444254032e-5 9.256228860049838e-5 1.6039137890229888e-5 0.00011951843826846697 4.632568799007385e-5 -0.00027295030248768854 7.226524086018874e-5 8.912026588386682e-5 -6.390016689651213e-7 0.00014420007993049493 0.00012912231669343866 1.956820106880121e-5; 6.820772991093462e-5 2.1254050935816552e-5 0.00015197904589627165 0.0001277667544320766 -4.542333395761709e-5 9.459615115814153e-5 2.8411437725732396e-5 3.679449959350978e-5 -9.117529394527309e-5 -8.758919999196077e-5 -4.2329960577931454e-5 -1.0144543792852978e-5 -2.800941534472828e-5 0.00019726418408138737 4.0235158068087194e-5 -2.032383404978274e-5 0.0001994509440385404 -0.00011162303404373478 -0.00013777382979128283 -0.00011793587314853088 -6.210373484211795e-6 0.00020680606407155213 -1.3990936039339685e-5 8.547127812677502e-5 0.00010467948608608965 3.4930981329375304e-5 -8.790327397830215e-5 0.00010486142868216187 7.141122490907534e-5 0.00011895331421443785 0.00013858456090136087 -2.8784597687936806e-5; -1.5860273009870596e-6 4.201089126071218e-7 0.00021171998620556242 4.522397981048049e-6 6.012802865879095e-5 9.99254795402399e-5 -3.561636306986259e-5 -6.138764486823493e-5 -0.00017997176231293035 6.246032232010124e-5 -1.3497106314760225e-5 -1.5011047557350911e-5 -0.000127839118553682 0.00011435324510383563 0.00019095278791209487 2.190836372901134e-5 -4.1166452624025464e-5 -1.1743301305983084e-5 7.070358211930417e-5 -9.540511140076516e-5 7.020535364571925e-5 0.000139547332460635 0.00010262069796735362 -0.00012118509524444342 -7.08920755045315e-5 -0.00016761459604313673 -1.7645824604556307e-5 -3.452961876465634e-5 -2.2219289653011745e-5 -2.1817491264671302e-5 -4.4123703750670265e-5 0.00011155886265040982; 0.00016322664453061945 0.00014593895468740305 3.38868136389157e-5 0.0001275027840386179 8.392109305010668e-5 0.00018751060417731248 8.67805613259002e-6 0.0002031460986693403 2.086227652787417e-5 1.0925697665026366e-5 -0.000148922123458514 -3.377056998554046e-5 0.00013773267858075127 -6.447396151610769e-5 1.8885938282181468e-6 -1.5252040044038412e-5 -0.00018388509340704775 -0.00020160053679839997 7.517021879199652e-5 -8.784191447896325e-5 -9.126163638550143e-5 -1.8795023039029765e-5 -1.5422080082975986e-5 2.640002060095067e-5 -5.737931712767606e-5 -4.767868973181286e-6 0.0001957927103142318 0.00010669599725975899 -2.501945232991478e-6 -5.6286028098591877e-5 -0.00012336296492842307 -1.8163820983073793e-5; 4.387362079610815e-5 3.081015913080804e-5 0.0001817471895428147 -2.8541160447322837e-5 -5.4667372553077835e-5 3.123232110548518e-5 5.7793520880254545e-5 -3.168949822980465e-5 -5.97186843337489e-6 7.561500043986861e-5 4.519134037577175e-5 6.334726293019989e-6 -5.920338282573348e-6 -0.00021513340814825556 1.7109627349733646e-5 -2.608024058831441e-5 6.723831318553467e-6 9.437682839957751e-5 9.714654535667561e-5 3.58741383187894e-5 2.9910119535957843e-5 9.456149947976864e-5 0.00013769551534819968 7.591111736287072e-5 0.0001212724275482746 -4.2953313624863676e-5 3.141121143740353e-5 0.0001629726723130462 -2.362060493584466e-5 5.085196264003189e-5 2.0586541921107583e-5 -9.444022342182586e-5; -9.685758408346632e-5 4.364034981512374e-5 -8.146627072896034e-5 0.0001229851696388734 -0.0001297638881970072 5.063172977297584e-6 -0.00020389505510871625 -8.498346863966554e-5 -3.220260254062204e-5 5.7336656083836206e-5 5.632566177332536e-5 0.00012570519908143912 -6.903969909949384e-5 9.563736972448336e-5 0.00013406571093759972 -4.0563754174032545e-6 -0.00014827228000109327 -9.67489394843712e-5 -0.00012002842407859088 -1.9552446190628368e-5 -6.245657272126004e-5 -0.00010350193764502016 7.2529501830457144e-6 5.4166375832163306e-5 -9.729999141023401e-5 -0.00011858525243775148 -4.4187214314718115e-5 3.9051799887576886e-5 7.176198420161946e-5 2.2048443553965248e-5 -9.916636184976858e-5 -0.00018513494427501925; 2.090883834711751e-5 -4.86838190354821e-5 -8.092724339262777e-5 8.123970736591713e-5 3.441147681400975e-5 -0.00010688578651211203 -0.00017534540994716325 -9.553687879274902e-5 5.324418715276299e-5 1.4815194741412272e-5 -2.5162886359376654e-6 2.1040982470317132e-5 -0.00013107203732402376 0.0001517277149886922 9.039843542809253e-5 1.0894737793882831e-5 -3.224914427306006e-5 -0.0001342116857941215 5.293620129984716e-6 3.59966642335573e-5 6.928717890281797e-6 2.8395233026409616e-5 2.974264942708453e-5 -0.00014179326273192986 6.431126065394449e-5 -0.0001352993832497807 0.0001949491005561362 1.6889708500410728e-5 -0.00019581432308227978 -0.00010042196401086974 0.00013518662570817869 0.00026510344731463784; 0.00011439415779256173 -0.00017873858938167853 -8.156019371875417e-5 -2.3657684648637738e-5 5.2571950293501925e-5 1.871661756246687e-5 -8.481401651573446e-5 -0.00020245902611116832 -6.970817535824231e-5 0.00015375887837087057 -0.0002252806492446639 1.1744980513471436e-5 2.5852605761033913e-5 -1.4503951728611104e-6 0.00012424922410224576 2.608449351614134e-6 -3.924265307199279e-5 -1.4089212852409973e-5 -0.00012934077436278432 -4.0909058368411555e-5 -0.00010972683717903768 7.257891651767383e-5 0.0002641473207791608 8.744857625575049e-6 3.3685085002220845e-5 9.373047454806915e-5 -2.3758658569418754e-5 3.0595848716693464e-6 -9.055626774846455e-5 0.00014242431883991713 -6.926328167024485e-6 9.393039603543552e-5; 0.0001841201183237734 -0.0001094061926737797 0.00015788291330359793 9.782154211686514e-5 0.0001258106231075092 -3.7072783921362604e-5 5.9204106041289164e-5 5.7525986453098875e-5 6.91411578413698e-5 -1.6425175086220383e-5 -8.042771241483175e-5 -6.61999418316991e-7 6.40493190116716e-5 0.0001124917498640867 -8.093995438278304e-5 1.5608799917097923e-5 -0.0001614524635401907 -0.00020957387186674022 0.00012761543894562787 0.00012393991619726798 8.609437918901831e-5 0.00027475403021932237 -9.615978796820858e-5 -9.75478660580217e-5 9.326350391754714e-5 8.283804618666058e-5 0.00016257405787153546 5.094758944728082e-5 -0.00016081130615524798 0.0001304944189582514 2.5020560319246074e-5 9.768924432558082e-6; 4.450085401528978e-5 -0.00011666123970301385 -5.0584801614812454e-5 -0.00011150179267517818 -0.00020778554386581975 4.6902000063500914e-5 2.211523536586659e-5 -2.424625887599876e-5 -7.743747632254398e-6 1.7532313391503186e-5 -9.82365868951343e-5 -4.6612321570343514e-5 7.069274800939851e-5 -2.470392206690141e-5 8.2366476304535e-5 -1.931315779459223e-5 7.0330851811134024e-6 6.528843590401575e-5 -6.184371665242768e-5 5.0172015433174444e-6 0.00016198599990800488 -0.00014750880433369088 -3.796728339168724e-5 7.662514184246219e-5 0.0001308055965962676 0.00010756193154136161 8.559682686225042e-5 0.00015994685457896646 3.555525330354768e-5 9.201297728920743e-5 -4.8966144790235e-5 -0.00013019513932179325; 0.00010356828599895037 9.386487974977538e-5 -9.555574185955084e-5 -0.00011302508326055312 -2.9818081452814296e-5 -2.6535467691564783e-5 0.00013269249641296085 1.3505676838274635e-5 0.00023217000009451678 -0.0001456984766949469 -6.160252203274579e-5 9.383508470334614e-5 -0.00012899754917755833 0.00026125094038857146 7.559933758299872e-5 -3.2456152689815777e-5 2.2871169425957352e-5 7.393651933368904e-5 6.940214982447468e-5 -8.073233882471435e-5 -6.493047228589041e-5 -7.970999402035335e-5 -3.690567361929861e-5 7.15984412860683e-5 0.00022401970520574907 -4.9010018551894524e-5 -6.30950101143061e-5 6.547663244842151e-5 -3.951944861486171e-5 2.588283380161544e-5 9.171112401123766e-6 7.061507378663521e-5; 3.650465740143488e-5 -0.00014354479637326503 -0.00010728621120428441 0.0001691574828154053 6.839487767216057e-5 -3.302755496328422e-5 -0.00015622326893130153 0.00012367395836996112 4.196091256822626e-5 -0.00024222287603983404 2.7071889490668772e-5 -0.0001657238942417672 -1.2417374083659221e-6 9.093671113156062e-5 0.00013993752807514928 -2.930248840051995e-5 -4.23821663563495e-5 6.964538698281893e-6 0.00017332066934691966 -0.00011847630659680641 2.811755555818023e-5 -0.0001294138243967088 -2.5065218068578445e-5 1.2464626962225503e-5 -0.0001047336452021618 0.00014952430072376674 0.00016974328471483824 -1.9495005034643464e-5 -3.489696311511811e-5 -5.3749660509440975e-5 -3.16869561348932e-5 0.00017349967245614364; -5.134221679304932e-6 -4.4242687577910665e-5 -4.964539897261032e-5 -1.5469965943313e-5 -2.5814287651847827e-5 3.438326035761383e-5 -0.00016474738862232237 3.214782060607047e-5 -2.6968031633945474e-6 -7.332552227997665e-6 -0.00023628649563394832 -2.7138837536744497e-5 4.681870137940785e-5 -5.0790842065865095e-5 -6.595567304973015e-5 -2.250984783674882e-5 -0.00015828147792124072 -8.267722874711865e-6 0.00013737001660414194 -6.649889142840209e-6 3.090675687783429e-5 2.7137767177949126e-5 0.0002341553432088635 -8.188873965625111e-5 -6.440525379364097e-5 -8.166045648610994e-5 -9.382462798018906e-5 -2.3537191223984183e-5 6.069597162789717e-5 -7.942860741965902e-5 0.0001866260814167841 2.8618735599623768e-5; -0.0001103470353325598 4.711010438881455e-5 4.53273092364272e-5 -5.5341563382858625e-5 -0.00020332961179405889 -1.6363072462575026e-5 -8.145435316578936e-5 0.00020983703831120378 -1.772488252641652e-5 6.426340916054237e-6 -0.00011325028793976707 -4.439482393152609e-7 -8.63093450232133e-5 -6.075200548251977e-5 3.6339078689020515e-5 -0.0001346253197323289 1.4278397940181219e-5 -0.00015716163801585804 -0.00011566039796775824 0.0001006105881617644 -4.945132731116753e-6 3.4785723584034426e-5 -4.706471412533389e-5 -0.0001809342579133539 1.5012739423803873e-5 -4.818823112032762e-5 0.00010727863951728201 -4.108732299523927e-6 8.441966712097673e-6 0.00018245381175915584 0.00015710553958455547 0.00010806589813114257; -4.593441173670279e-5 3.7793233154343817e-6 0.00019018267647010885 -0.00015562449681925028 0.00010617951937133138 2.7162692440875975e-5 -0.00019616122865004856 -6.838615792690198e-5 0.0001046013568888565 1.842436557166921e-5 -1.7457355915716704e-5 0.00010073521860289928 -1.3159186160506811e-6 4.764584098706697e-5 -7.214823546706135e-5 -6.362226833476142e-5 -8.791669808447752e-5 2.5511908625453473e-5 -7.15792774095046e-5 7.802264024616761e-5 -6.758372166546458e-5 0.00011234894207559183 -2.9502463639396427e-5 4.234118039716605e-5 1.3768161971206715e-5 0.00015046686183308797 -2.0745190176562644e-5 1.5610080985844744e-6 -8.091334149826649e-5 -2.2058549638636323e-5 -9.851741908237734e-5 4.3964377419293234e-5; -7.40427903563168e-5 -0.00011613714464124472 -3.919667956302493e-5 6.991142572812671e-6 0.00015672696013120283 -0.00012052558754454351 -5.07557787839243e-5 -4.316161807742336e-5 -9.245491396518537e-5 3.0238580168242475e-6 7.45925963658433e-5 -0.00017187449425341714 8.258726579719799e-5 1.3630264266102851e-5 4.8585620834498505e-5 -1.070037792465123e-5 -9.760991538291944e-5 8.834741588566336e-5 5.067445317985583e-6 0.00011652495145449258 -0.00016424044310929368 6.711458540871316e-5 6.106554342318183e-5 9.281314401073058e-5 3.989247038611925e-5 -0.000134968101474216 9.744568074006471e-5 -0.00010664043379598136 0.00022795957470429556 0.00017059860242647743 -7.061922674415746e-5 0.00013930053232820265; -4.044082957334265e-5 -4.6239454925667574e-5 -1.5058668716651286e-5 -8.942060080952913e-5 -7.745542675618877e-5 -1.8622524615033393e-5 -6.349596365252841e-5 8.038021084948716e-6 4.0742137250346976e-5 5.7552746190796185e-5 1.9405895406825072e-5 4.646670063012006e-5 0.00020582012225278006 -0.00015895532701758647 0.00014299346516709894 0.0002919499107949414 -6.14583681024639e-5 -5.6739771346472196e-5 -0.00015057118285111284 -8.007616848116997e-5 3.019504733602143e-5 -0.00012120903407120613 -0.00016406994687058322 -0.00014540190600454618 -7.023044450106193e-5 -0.00016744802847169573 -2.5244359087808026e-5 -3.9793338467259885e-5 -5.5561466390632696e-5 3.826051547294631e-5 1.56458424018658e-5 0.0002207345547933181; 5.562366295291923e-5 6.877203504357721e-5 -7.352626908300165e-5 9.961856414851366e-7 9.807948317416799e-5 -0.00010688944512705955 -7.680180333752213e-5 3.608407878004278e-5 -3.811187325825106e-5 -0.0002000951970364765 -3.370477843760378e-5 7.504336481018075e-5 -0.0001118090183842998 3.9216214823932526e-5 3.158286944766428e-5 1.0503931173218698e-6 -0.00023946163474081018 -0.0001243408076368469 -2.6353005345069815e-5 -5.522916745009717e-5 -0.00022014997721614426 -0.00011568037353186304 -3.740040077691801e-6 -7.859334782893724e-5 -7.117432306019161e-5 -3.0465227431495027e-5 8.643789278378611e-5 -2.413601924102731e-5 -5.5296233589387884e-5 0.00013858510301019026 5.4608856040638585e-5 -1.1001591523157642e-5; 2.0312441915057575e-5 6.18979157244491e-5 -0.0001614350930375328 -8.915590472121852e-5 1.1043714106683836e-5 1.845079902899112e-5 7.726544094230562e-6 0.00010280457906822585 0.000275007665617641 -0.00015673157934167355 8.405503716977726e-6 -0.00011637083837364973 7.787991314592678e-5 8.792042178732485e-5 -2.3668457523570387e-6 -0.00023994251562350593 9.92124678979331e-5 5.095685089520924e-5 -0.00016145472357117588 7.010332599214107e-5 -9.21346308367625e-5 5.142079505652123e-5 0.00012027771851911034 -0.0001915719348524987 3.3666252185238254e-6 2.508996847075708e-5 -0.00010461831873350964 8.769304326172683e-6 8.415127929232049e-5 0.00012264255026287287 -6.644276611575116e-5 -2.309772231336127e-6; -8.704423278478563e-5 8.531497398715528e-5 4.1546218128375996e-5 0.00012446133580149802 0.0002487970322565734 -5.0768305370865204e-5 2.95092631152206e-5 7.9598305021387e-5 -6.330363437674596e-5 8.979052465498004e-5 0.0001458492314871154 -1.3098109846136582e-5 -4.4773753031880795e-5 -0.00015706251913593465 -2.032127856884567e-5 2.8146798215206692e-5 7.020638979364379e-5 1.8661494252489454e-5 3.078752517673589e-5 0.00023090094138959578 -0.00014576129804273736 -6.070005663580276e-5 0.00017205394208057065 -4.681960060427926e-7 -3.542210295355973e-5 -8.49300286770041e-5 7.960211762318715e-5 0.00012569797210949896 -0.00010559401023577317 -0.00021669462337533475 0.00016136636070067408 0.0001259090330879849; -1.6662365804252582e-5 -3.076175102981278e-5 3.1604035986664635e-5 5.1337026026862645e-5 -0.00016384416353321847 -5.152308490141233e-5 -2.1736140250170154e-5 -6.0848508790532144e-5 3.4903104326046616e-5 5.817823062170568e-5 -4.397005882130676e-5 0.000186074609748147 -8.90085966130033e-5 0.00011241916677318738 -7.667130098345636e-5 3.1735126914982294e-5 1.070597413851729e-5 -9.913572552980692e-5 6.0796207482770946e-5 -1.1410686918707744e-5 -0.00013092540298024035 -1.0529580275049113e-5 1.45537733733078e-5 -5.050842714628955e-5 -0.00010791788640565542 -7.187148271262255e-5 -0.00026904714736522935 -3.381604160313152e-5 -2.2611292431987064e-5 -1.2323857798058405e-5 -0.00014827294485078978 2.9433603278433405e-5; 6.869139095382899e-5 -1.324076598870732e-8 -6.626889891609154e-5 -2.1587850749853613e-5 7.173868569393166e-5 -3.0154999802246147e-5 3.075379852984246e-5 -5.2434815114025904e-5 -0.00017200225640240922 4.5476018460109414e-5 -9.20344545358206e-5 -0.00013684427623789598 7.445155080716845e-6 2.2761199323646083e-6 5.940447483083813e-6 -0.00010312598253054374 -0.00010338067742682837 7.628147676242766e-7 -4.421295141723941e-6 3.710197332278986e-5 -1.5903712590586907e-5 -0.00020041794133990686 5.2564929393922574e-5 1.0063752671925168e-5 0.00013255340953204978 0.00018796490538247097 -1.4615446996829499e-5 4.737927800474284e-5 -5.662626510289128e-5 7.898041720349595e-5 0.00012663719742871956 4.565956176688293e-5; 6.040968241798923e-5 0.00010700802908907883 -0.00013008028668188204 0.00013264787798910653 -9.973736791941042e-5 0.00018789393311529747 6.609491197089898e-5 0.00029767339242319825 -0.00012588732693181139 -3.655107549243098e-6 2.019274347625e-5 9.936767354551735e-6 4.486933562369346e-5 3.891675464623005e-5 -7.401772359222851e-5 -0.0001895071216050977 -0.00011452335798018376 -3.1347884715765265e-5 7.583290084661267e-5 -1.939678236446106e-5 -0.00015249483453877967 -3.665886447485268e-5 1.4171150040815142e-5 0.0001292302442824028 0.00012168411581008939 -9.64592143201711e-5 -8.477287679048954e-5 -0.0002004259038141432 3.0735187691639164e-5 6.524801233248109e-5 2.0169673233644556e-5 -8.743083503865367e-5; -0.00015594568523967804 -6.432045322300979e-5 8.600838353002613e-5 -6.847963797783528e-5 1.2529190623792945e-5 -8.807312244712547e-6 3.45792120947208e-5 -2.4681085253650954e-5 -8.225049370823592e-5 -9.367357981547387e-5 -9.418858301178444e-6 -5.133474971673586e-5 0.00014421450194047788 -4.407725120701351e-5 8.811265595566455e-5 7.145731958869727e-5 -5.687893846927348e-5 0.000112704585060308 5.9515854243255995e-5 0.0001358120078088889 -7.36492349665442e-5 -0.00011210645445546061 8.18189817232252e-5 6.882221512340005e-5 -5.7000850777077406e-5 7.864945089520762e-5 1.1340554717642497e-5 4.324252384092998e-5 -0.0001591366146590361 -1.388631810416017e-5 3.255396748270674e-5 -0.0001389950177531495], bias = [-1.2853470297038883e-9, 4.896051452503787e-10, 1.6379466293720113e-9, 1.2138759146976405e-9, -9.218426927451368e-11, -1.8230331087992903e-9, 1.727273527097074e-9, -3.643656384310119e-9, 1.135604007608576e-9, 3.9509734757092134e-9, 7.081922327030148e-10, 1.5803600726850138e-9, 3.649090973270635e-9, -3.7907175855678225e-9, -4.630365361585973e-10, -4.123261859208235e-10, 4.150028400466572e-9, 3.0780554282139766e-10, 2.2980427059247367e-9, -1.0615544228007823e-10, -1.8675891400887143e-9, -1.4192107834140116e-9, -1.3080312335908598e-10, 7.813016500553129e-10, -2.1085295864516417e-9, -4.0488919442204864e-9, -3.073417879929848e-10, 3.228635512057411e-9, -3.5387966208828077e-9, -4.664963470446139e-10, -9.072315647865157e-11, -5.950119176721757e-10]), layer_4 = (weight = [-0.0006446651666109494 -0.0007529018339533018 -0.0008093233860063905 -0.0006284760118885408 -0.0005968810790197654 -0.000585493681187751 -0.000837673994838595 -0.0007711489063037479 -0.0005633050097127248 -0.0006366020175565896 -0.0006351336841893969 -0.0007250663187550846 -0.0006773039539336318 -0.0007102623027784417 -0.0006984477441356452 -0.0007836149331594133 -0.0006547691546035754 -0.0006358416473538963 -0.0007439842062632546 -0.0007248791323072195 -0.0009062966064743499 -0.000753728407564923 -0.0005576711033216811 -0.0007266704450221633 -0.0006264206496902814 -0.0006150284993179414 -0.0006485911114611045 -0.0005555057596760502 -0.000723135451689339 -0.0006958771766227974 -0.0007989775974893779 -0.0007756897085017334; 0.0002635064491133733 0.00013902321060088398 0.00027416288755319483 0.00036005136845857904 0.00026665041748154093 0.0003853675979829293 0.00020777463582142915 0.00024593198624691553 0.0003353161483059413 0.0003276606322724664 0.00015453736205649262 0.0002593210875760926 0.00010414091948133386 0.00012031890720450069 0.0003972533638995509 0.00012003154294842141 0.0003925132655046808 0.00023879112576484445 0.00026200251414252603 0.0002331804026158072 0.00031018298463556074 0.0002073419829842547 0.0003465754111874278 0.00015812656915209847 0.00019217357651452952 0.00011534668831958337 0.00010325712161342742 0.00021456723190503042 0.0001459926682109346 8.304199293465108e-5 0.00031881876799978747 0.00017356179484362758], bias = [-0.0006918680946778557, 0.00023882644622278632]))

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.11.5
Commit 760b2e5b739 (2025-04-14 06:53 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 128 × AMD EPYC 7502 32-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver2)
Threads: 16 default, 0 interactive, 8 GC (on 16 virtual cores)
Environment:
  JULIA_CPU_THREADS = 16
  JULIA_DEPOT_PATH = /cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6
  JULIA_PKG_SERVER = 
  JULIA_NUM_THREADS = 16
  JULIA_CUDA_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0
  JULIA_DEBUG = Literate

This page was generated using Literate.jl.

Layout Switch

Page Layout Max Width

Content Layout Max Width

Spotlight

Trusted by

Training a Neural ODE to Model Gravitational Waveforms

Package Imports

Define some Utility Functions

Simulating the True Model

Defiing a Neural Network Model

Setting Up for Training the Neural Network

Training the Neural Network

Visualizing the Results

Appendix

Layout Switch

Expand all

Expand sidebar with adjustable values

Expand all with adjustable values

Original width

Page Layout Max Width

Adjust the maximum width of the page layout

Content Layout Max Width

Adjust the maximum width of the content layout

Spotlight

ONOn

OFFOff

Layout Switch

Page Layout Max Width

Content Layout Max Width

Spotlight

Trusted by

Training a Neural ODE to Model Gravitational Waveforms ​

Package Imports ​

Define some Utility Functions ​

Simulating the True Model ​

Defiing a Neural Network Model ​

Setting Up for Training the Neural Network ​

Training the Neural Network ​

Visualizing the Results ​

Appendix ​

Layout Switch

Expand all

Expand sidebar with adjustable values

Expand all with adjustable values

Original width

Page Layout Max Width

Adjust the maximum width of the page layout

Content Layout Max Width

Adjust the maximum width of the content layout

Spotlight

ONOn

OFFOff

Training a Neural ODE to Model Gravitational Waveforms

Package Imports

Define some Utility Functions

Simulating the True Model

Defiing a Neural Network Model

Setting Up for Training the Neural Network

Training the Neural Network

Visualizing the Results

Appendix