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, LuxAMDGPU, LuxCUDA, OrdinaryDiffEq, Optimization,
      OptimizationOptimJL, Printf, Random, SciMLSensitivity
using CairoMakie

CUDA.allowscalar(false)

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)<1e-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, orbit1, mass1, orbit2, mass2)
    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; markershape=:circle,
        markersize=12, markeralpha=0.25, 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(broadcast, cos),
    Dense(1 => 32, cos; init_weight=truncated_normal(; std=1e-4)),
    Dense(32 => 32, cos; init_weight=truncated_normal(; std=1e-4)),
    Dense(32 => 2; init_weight=truncated_normal(; std=1e-4)))
ps, st = Lux.setup(Xoshiro(), nn)

((layer_1 = NamedTuple(), layer_2 = (weight = Float32[0.00011340812; 3.2809297f-5; 1.0014967f-5; -9.990923f-5; -4.3205917f-5; 0.00010386309; 9.164389f-5; -0.00010796392; 0.00015823921; -4.437526f-5; -9.740333f-5; -0.00019613246; -0.00013827338; 9.762513f-5; -8.401355f-5; 8.799325f-5; 8.460665f-5; 9.806187f-5; 8.6452805f-5; -2.4717076f-5; 0.0001535528; -0.0001427423; -0.00011783649; -9.80459f-5; 8.313435f-5; 0.00011971523; -4.46122f-5; 1.77301f-7; 2.5376681f-5; -2.4404124f-5; 9.128665f-5; 0.00014188747;;], 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.00014156256 -9.811367f-5 -1.517552f-6 0.00016699592 -0.00013732523 -6.3737905f-5 3.3648437f-5 0.0001634907 4.093107f-5 -3.6230485f-5 2.6603404f-5 6.914893f-5 1.4580082f-5 -8.376587f-5 1.499435f-5 -9.853641f-5 1.2621165f-5 8.196162f-5 -6.582095f-5 8.569358f-6 -8.837326f-5 8.642007f-5 0.00020047584 -2.9746847f-5 -0.000100567886 6.3697706f-5 -2.4618008f-5 -9.2386406f-5 1.4190906f-5 -2.987285f-5 3.9674404f-5 -2.221928f-5; -7.3961295f-5 5.107348f-5 -1.25083725f-5 5.3003616f-5 0.00010938918 -1.7892273f-5 -3.3245517f-5 -2.01148f-5 -7.30904f-7 7.5736876f-5 -0.00010958888 -0.00011940187 8.635474f-5 -0.00018097984 0.000107903674 -9.916499f-5 -0.0002525818 6.4250606f-5 4.1518964f-5 -1.2521067f-5 -1.4804918f-5 -5.7143166f-5 -7.471726f-5 -0.00012908937 -3.4056216f-6 4.9543356f-5 -9.364649f-5 -5.1094314f-5 -4.6537763f-5 -5.9186354f-5 6.3696025f-5 -2.9774865f-5; 2.0801594f-5 8.362106f-6 0.00014500385 0.00014755265 -9.782925f-6 7.232735f-5 -0.00015232062 -0.00012671648 8.047096f-5 7.517334f-5 4.0161867f-5 -3.5384433f-5 1.1653104f-5 0.00010767147 -1.3131585f-5 2.5852463f-5 4.6637765f-6 -2.4403307f-5 -8.059132f-5 0.00017241148 1.475767f-5 -5.757022f-5 -0.00015111905 -0.00014129387 3.9219325f-5 0.00014003659 -0.00015052085 0.00029057162 0.00014398372 8.677407f-5 -0.00012131474 -1.0865262f-5; 0.00011003546 0.00016107866 2.766693f-5 -2.4925546f-5 -0.00013837212 0.00013506766 7.6011f-5 0.00010802495 5.2456795f-5 -9.955632f-5 9.0362504f-5 7.731734f-5 0.00010679167 -0.00014037582 5.186795f-5 1.2005265f-6 0.00011025447 5.0290117f-5 -5.1358005f-5 -0.00015502913 0.00019844093 6.4749665f-6 -0.00016835763 -4.077826f-6 2.2255883f-5 -6.627248f-5 -7.4372633f-6 -9.2532224f-5 -0.00016317538 -0.00023337106 -6.327392f-5 9.510757f-6; -4.877959f-5 0.000111007335 -0.00012343853 -8.922489f-5 -1.4022108f-5 6.92703f-5 -0.00014855403 -7.743551f-5 -1.1391856f-5 -6.64675f-6 3.4402125f-5 0.0002186161 0.00016305305 -8.375087f-5 9.215641f-5 9.4706666f-5 2.5164243f-5 -0.00019379282 9.986605f-5 9.7672375f-5 8.304275f-5 3.520148f-5 -0.00013444443 6.3760875f-5 -7.7740224f-5 0.0001202884 -0.00015718729 2.4690196f-5 2.6267246f-5 3.7738188f-5 0.00010785975 1.7241578f-5; -0.00021674494 -5.0069117f-5 0.00019511077 -1.3805467f-6 -0.00010837159 6.231287f-5 4.7232938f-5 9.554871f-5 -0.00019214283 -2.3225084f-5 -1.2651437f-5 -3.0432508f-5 3.8927286f-5 4.5957295f-5 -1.1540804f-5 0.00030040392 5.638816f-6 -0.00015406098 0.00014416596 -1.7310642f-5 6.277926f-5 9.007929f-5 6.2393847f-6 8.975524f-5 -1.5287858f-5 -4.1864663f-5 7.428573f-5 -0.00023177551 -9.165949f-5 -8.983296f-5 -0.00015798958 2.3861057f-5; -1.5072078f-5 6.9965005f-5 1.7785236f-5 -2.0398182f-5 -4.162357f-6 8.639045f-5 -1.4841806f-5 2.3637292f-6 0.00011413424 0.000112395384 1.3264921f-5 -1.4758984f-5 -3.45135f-5 0.00012556222 0.00022367231 -4.6766712f-5 4.0619902f-6 -0.00011636409 -4.011377f-5 -6.0061986f-5 -0.0001391922 8.136689f-8 2.5876163f-5 -9.962023f-6 -9.325049f-5 0.00017469504 -4.1053907f-5 -8.129812f-6 0.00015975317 -0.00019158708 -1.7481741f-5 7.9576304f-5; 0.0001245884 4.6486322f-5 5.1412855f-5 -4.674493f-5 3.4289536f-5 -0.00013141884 -9.013261f-5 -5.5951205f-6 4.125992f-5 -4.7334488f-5 0.00021864806 4.893553f-5 -4.846627f-5 -1.5561232f-5 -0.00011433091 -9.1516835f-5 4.5547393f-5 -7.504747f-5 0.00013871712 6.938905f-5 3.4820903f-5 -8.592702f-5 -1.6073507f-5 9.697634f-5 4.9590137f-5 -1.8486593f-5 0.0001345847 -8.8592395f-5 -1.8502293f-5 5.312635f-5 3.663749f-5 -3.887912f-5; 4.3029413f-5 -5.8115053f-5 -9.79173f-5 -6.771011f-5 3.297839f-6 2.289137f-5 -0.00014672716 6.326552f-5 7.89383f-6 -3.6980447f-5 -3.1390686f-5 0.00012616182 0.0001697994 -0.00014502708 -0.00022832723 7.680645f-5 0.00017640064 -0.00013793133 -4.8998256f-5 9.581965f-5 -0.00010461029 0.00017992317 7.827547f-5 7.7142555f-5 7.782721f-5 -2.0406335f-6 0.000101972335 0.00018598197 -0.00011490433 8.835141f-6 3.1883206f-5 -0.00012890752; 6.808792f-5 -0.0001839431 0.00015006302 -5.5762335f-5 -0.00013817698 -5.9380316f-5 0.00016808703 -3.998663f-5 -4.8634196f-5 3.1601572f-5 3.23283f-5 -1.4613938f-5 2.5012218f-5 -0.00018407426 -3.6703437f-5 3.737339f-5 7.8831305f-5 0.00011071217 -5.4118635f-5 8.520669f-5 6.0180628f-5 0.00024431307 5.273774f-5 2.4208919f-5 9.477329f-5 2.6201908f-6 2.911732f-5 -0.00022157382 6.011246f-6 8.044288f-5 0.00015570794 -9.233129f-5; -0.00012052653 1.3465535f-6 0.0001783217 -8.839138f-5 -3.1912245f-5 -6.282885f-5 -6.349209f-5 -8.533851f-5 -0.00016356591 -3.9394647f-5 1.2402778f-5 7.976498f-5 -0.00012994005 3.4589724f-5 1.6261467f-5 0.00019368868 7.675898f-5 0.00016140526 5.9786616f-6 -9.10683f-5 -0.000105497595 -5.1691066f-5 -0.00012642182 -6.489419f-5 0.0001751558 9.6908196f-5 1.4118654f-6 -5.8078174f-5 4.3729433f-6 -1.1362517f-5 6.206006f-5 1.331932f-5; -3.3587065f-5 -0.000113246635 -0.00012955924 -2.292071f-5 3.307507f-5 0.00011393347 -6.0897135f-5 -0.00019180034 1.4472878f-5 -0.000117850104 8.72241f-5 4.7593698f-5 9.964195f-5 -2.4952073f-5 -0.00018551225 -6.616688f-5 5.006724f-5 -2.8952132f-5 -0.00010744806 -5.2000494f-5 0.00011208395 -0.0002498601 5.9776947f-5 -0.00012753709 -0.00011752931 -8.1820845f-6 -0.00017713547 -6.695523f-5 0.00015603926 0.0001292478 4.1312014f-5 -0.0002221211; -7.128793f-5 0.00016443687 -3.0863495f-5 -7.669469f-5 -2.6215224f-5 2.867582f-6 -0.00016141444 -5.6848035f-5 0.00017235398 -2.3746074f-5 6.130454f-5 -0.00017731883 -6.232084f-5 1.8692841f-5 0.00015196079 6.157048f-5 -3.4617482f-5 1.3108995f-5 -6.496418f-5 3.506001f-5 -7.1566596f-5 6.247823f-5 -0.00010291936 5.7846795f-5 3.210926f-5 -5.6527493f-5 7.247856f-5 9.33015f-5 6.622545f-5 -6.516572f-5 4.4051878f-5 -3.3282136f-5; -0.00011829952 0.00016207603 -2.6408405f-5 -1.9043713f-5 -2.7642885f-5 -0.00013119672 -7.4777396f-5 8.691534f-5 8.218466f-5 -8.744666f-5 -3.098265f-5 6.316316f-5 -2.3710376f-5 1.7484153f-5 5.4146276f-5 -0.000101633916 0.00015427092 0.00028224068 -4.0874795f-7 8.089582f-5 0.00013113428 2.7629396f-5 9.6891854f-5 -6.6886045f-5 8.954415f-5 -4.3712847f-5 2.7109858f-5 -2.3283614f-5 -0.0001643385 -5.124151f-5 3.839928f-5 -7.731988f-5; -4.2265336f-5 -0.00015316042 8.888119f-5 -0.000104785584 8.268285f-6 -6.8925634f-5 8.682539f-5 9.083212f-5 -2.5445628f-5 7.052769f-5 -6.0680675f-5 0.00013888312 3.7213493f-5 -4.4190903f-5 0.000109429544 -4.1258423f-5 -1.5919257f-5 -0.000118554584 -8.714114f-5 -1.3423698f-5 -7.128542f-5 -0.00022807624 -4.6775214f-5 2.4986768f-5 8.095359f-5 -1.901928f-5 2.452534f-5 0.00018246549 -6.445913f-5 -1.131712f-5 -2.6922156f-5 0.00010507092; 0.00010393506 -0.00016572044 0.0001826992 0.00013126047 0.00013866162 0.00020592136 0.00012411547 -4.3536842f-5 -0.00015657167 9.439269f-5 -7.966156f-5 -0.00010378239 -8.540378f-5 -6.30376f-5 3.7696911f-6 -0.00026096636 3.6424062f-5 -9.3506795f-5 -0.00013874442 -3.702843f-5 -6.161214f-5 3.4262068f-6 0.00026881188 -6.840278f-5 9.9414225f-5 4.540923f-5 0.00011233505 4.8393056f-5 0.000112785354 -1.1175712f-5 0.00017866767 -4.222517f-5; -0.00012568003 -2.7524296f-5 -6.6743894f-5 -6.423323f-5 -1.977272f-5 -9.576463f-5 7.666163f-5 4.2683565f-5 2.7817929f-5 1.0548998f-5 -0.00011866199 -0.00013129206 0.00013543369 -3.0368214f-5 -3.1198837f-5 -0.0001777863 -3.7585996f-5 -9.196156f-5 0.000110632456 -6.518952f-5 -3.3768007f-5 -8.026831f-5 -8.575026f-5 2.3256734f-5 -2.972639f-5 -5.7257435f-5 -0.00022663553 -0.000116349984 -9.252055f-5 8.771717f-5 6.1138235f-5 0.00010377394; -0.00015122101 0.00022250555 -0.00013489451 8.7480366f-5 -8.549556f-5 0.00012857793 -7.220271f-5 6.827511f-5 -6.4384585f-5 -0.00016480038 0.00011284755 8.00845f-5 9.7739205f-5 -1.4331893f-5 1.1986961f-5 -0.000116840754 -3.8461792f-5 -1.3898441f-5 0.00019769951 -5.0751463f-5 -9.483562f-6 0.00012867231 0.00013268231 -0.00021376186 4.682368f-6 1.9103143f-5 -8.316697f-5 -8.809467f-5 -6.1333856f-5 -9.858935f-6 5.000858f-6 -0.00012287246; 2.3604498f-5 4.6526624f-5 -2.9173987f-5 -0.00024712266 -0.00017876402 -0.00016922623 7.47458f-5 -6.33877f-5 -3.6697078f-5 -7.583508f-5 0.00015662325 2.5254267f-5 -3.059348f-5 0.00020411164 7.626975f-5 9.448594f-6 1.2125071f-5 6.024337f-5 -5.851504f-6 -3.5932815f-5 -6.660372f-5 7.399119f-5 -0.00013992749 -0.00022908897 -0.000113892536 -1.986681f-5 -8.934934f-5 0.00012021076 -1.5911586f-5 -8.780692f-5 -1.151107f-5 -7.827197f-7; -0.0002055156 -5.3613163f-5 0.00013410521 -6.9706966f-6 7.0229376f-6 1.4096539f-6 -0.00016749899 7.566419f-5 -0.00011994494 5.185088f-5 1.4498425f-5 -2.3627972f-5 -3.8528935f-5 3.4861212f-5 6.829045f-5 -0.000119699784 -0.00011499884 -7.4220266f-6 0.00018222106 3.7699153f-6 -5.089868f-6 -5.9110487f-5 0.00010094063 -0.00011955277 -6.401306f-5 6.397591f-5 -2.888949f-5 9.627232f-5 5.5537603f-6 -9.948773f-5 4.51412f-5 9.71577f-7; 5.6067725f-5 -6.931405f-5 0.000117395284 0.00022472789 -9.100562f-5 -0.00010641615 -6.1037346f-5 -0.00010733551 -3.1581912f-5 -6.691269f-5 4.8335336f-5 -7.307645f-6 5.599756f-5 3.312475f-5 -3.895588f-5 -2.7871527f-5 -3.266601f-5 -7.9796875f-5 -5.3999636f-5 2.7975346f-5 1.2234575f-5 -1.682895f-5 -0.00016058375 9.995087f-6 1.2557403f-5 -0.00018376406 -4.9031816f-5 -7.507692f-5 -0.00013441994 -0.0001444401 0.00017147485 3.7063794f-5; -7.033317f-5 -2.9558596f-5 -0.00015901076 2.4425042f-6 2.3967461f-5 7.187309f-5 0.00025936795 -0.00012734935 -1.3547716f-6 2.4444207f-5 2.5813597f-5 0.00012840524 0.0003317028 -0.00013460938 -9.3497845f-5 0.00012017578 9.427198f-5 -0.0002199859 3.9750255f-5 0.00020362379 -4.9279024f-5 -9.9389836f-5 -8.3618426f-5 9.2390226f-5 -0.000107847154 0.00010294984 -0.00012627833 -8.869408f-5 -0.00014664406 4.2412885f-5 -0.00016417194 3.2001804f-5; -6.3029074f-5 -0.0001527952 1.0535898f-5 -5.391492f-5 -0.0002278912 2.325454f-5 3.177855f-5 -8.3847466f-5 -4.5754314f-5 0.000106730775 -0.00023069885 -2.0556507f-5 0.00025303726 -7.335827f-5 -4.3911936f-5 -6.4138905f-5 -1.5417554f-5 -3.7610098f-5 8.352618f-5 0.00017013338 -0.00011105443 -3.261541f-5 0.00017519973 5.3404947f-6 -5.69709f-5 -6.496197f-5 -3.7815687f-5 0.0001502521 0.00013033353 -3.1277777f-5 5.206057f-5 2.9687285f-6; -0.00013140556 6.9060803f-7 -8.8772336f-5 1.8929286f-5 3.397261f-5 4.7382306f-5 -5.7502097f-5 3.26821f-5 -4.9978058f-5 1.4620945f-5 7.814715f-6 4.1244384f-5 2.1999316f-5 0.00010458455 -4.2806812f-5 0.00016115289 -0.00013172816 -2.8587483f-5 1.9958095f-5 -6.708505f-5 -4.7145564f-5 -0.000141155 -8.094968f-5 -0.00014231665 5.280408f-5 -3.1761105f-5 9.366611f-6 -0.00013961366 5.2050287f-5 3.7327314f-5 3.2658998f-5 -0.0001371651; -4.125543f-5 -0.00022209117 0.00011682745 -8.491548f-6 -1.095696f-5 -7.349185f-5 9.675288f-5 4.8146678f-5 0.00024390296 2.9749595f-5 0.000106366155 -6.5939246f-5 0.00020534662 -4.4923087f-5 -2.3404999f-5 -6.351213f-5 1.6676477f-5 0.00013496123 0.00014743135 6.582297f-5 3.9581784f-5 -1.4164465f-5 -5.242885f-5 8.270567f-5 -0.00013805758 -4.6702688f-5 -5.7950623f-5 -1.6153153f-5 0.00012497827 -8.811363f-5 4.9019705f-5 9.599656f-5; 5.9385216f-5 -1.7540049f-6 9.621014f-5 -5.1853032f-5 0.00015560212 2.0128033f-5 5.1846146f-5 1.0333128f-5 1.7179013f-5 -0.000188536 -4.6423895f-5 5.2459025f-5 0.0001509285 -5.895016f-5 -8.778701f-5 -1.1115667f-5 4.53824f-5 -0.000113025315 -0.00011277421 -5.2061547f-5 0.00014770738 -9.4210554f-5 -5.05245f-5 -0.00010159458 -0.00012724199 -0.000101610676 0.00012396892 -9.794783f-5 -3.2778567f-5 -1.1444896f-5 7.075737f-5 -4.143076f-5; -1.9212586f-5 9.686188f-5 -8.232209f-5 -3.7457467f-5 4.8496735f-5 0.00012407255 -3.4657776f-5 3.9476115f-5 0.00016354468 9.990037f-5 1.401645f-5 -0.000115941664 -5.5375993f-5 -0.00021627567 -1.9354005f-5 5.7542962f-5 -1.9661864f-5 -0.00028484195 2.198932f-5 -5.884804f-6 -0.00018816453 -5.4484724f-5 -6.63086f-5 3.916278f-5 -0.00015688162 -2.999979f-5 -1.4491449f-5 0.00017595665 -6.969547f-5 -0.00012779546 8.1196915f-5 8.858673f-5; 0.00016343874 -3.8318827f-5 -9.389531f-5 -5.0160186f-5 2.3524695f-5 1.5061712f-5 1.806115f-5 7.480334f-5 8.23867f-5 0.0001427907 -6.971522f-7 -0.00025831407 4.0868515f-5 -7.1229015f-6 4.9699207f-5 -0.00020396945 -0.00020271681 8.4812345f-5 -9.938254f-6 8.841552f-5 -1.1265471f-5 -0.00016805819 -8.867307f-5 -4.6431633f-5 2.944485f-5 -8.879436f-5 -0.00016809237 -2.0559086f-5 0.00016534043 -6.289881f-5 -9.4184754f-5 -0.00013409012; -5.7644473f-5 -3.541472f-5 -4.6580404f-5 -7.146827f-5 4.285381f-5 -1.1417459f-6 -0.00010710779 -6.32022f-5 0.00012676835 5.16021f-5 0.00016374915 0.00017052848 6.891614f-5 -0.00014135256 0.00012690522 8.2927414f-5 -0.00011360938 0.00021434131 6.099151f-5 9.301906f-6 0.00010037649 -7.0248548f-6 -0.00014967707 0.000101035046 -5.6734293f-6 -0.00012880597 4.3167292f-5 4.4839267f-6 5.901098f-5 0.00018500042 0.00016087679 -9.2648166f-5; -0.00017676917 7.6590615f-5 8.920116f-5 -0.0001488689 0.00010899476 -0.00017075292 -4.0270985f-5 0.00014151317 2.8666898f-5 -4.8055837f-5 -3.801341f-5 3.1675696f-5 -9.472108f-5 1.3138095f-5 -0.0001239027 5.5682784f-5 7.649597f-5 6.4410124f-5 -1.8557983f-5 0.00011456157 -0.00014423669 0.00022011256 -2.6755824f-5 0.00016565724 8.5450236f-5 1.5266f-5 -2.317029f-5 7.68452f-6 -3.2184525f-5 -3.691155f-5 -0.000159565 0.00015980682; 0.00017030508 7.528614f-5 -2.6625972f-5 -8.7292356f-5 -1.2368286f-5 7.2868264f-5 -7.815586f-5 -0.00016535269 -9.28558f-5 0.00010033605 -6.573564f-5 8.424681f-5 2.3016755f-5 2.4647039f-5 -8.655332f-5 -0.00020018968 -0.00012179129 -3.831024f-5 -0.00022175073 2.123916f-5 -5.8231737f-5 -0.00014432642 5.5496002f-6 4.4627515f-5 5.7364647f-5 -3.3952617f-6 -4.217787f-5 6.2551786f-5 2.2906937f-5 3.956708f-5 0.00014557954 8.7029406f-8; 8.146902f-6 2.4373032f-5 -8.608345f-5 -9.678789f-5 -6.3482395f-5 8.274202f-5 1.576273f-5 2.3338283f-5 4.1236966f-5 -4.9489718f-5 -0.00010362939 -1.8899595f-5 -3.361941f-5 1.1177696f-5 -5.0133585f-5 -1.4017355f-5 4.7044003f-5 -0.00010790271 -4.333274f-5 -5.6098954f-5 6.785915f-5 7.4130505f-5 -4.8344824f-5 0.0002574895 0.00015002121 5.365681f-5 8.5888234f-5 -0.00018845937 -0.0001366969 -4.1488904f-5 -4.8940314f-5 2.5306545f-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.140263f-5 -2.7525864f-5 -8.824687f-6 -0.00014753721 6.928443f-5 -5.20971f-5 2.1341175f-5 5.1338044f-5 -5.1185245f-5 -6.0365743f-5 -2.1162485f-5 -5.4177315f-5 -1.8738348f-5 -7.366614f-5 9.997543f-5 7.272992f-5 0.00010180698 2.8969784f-5 -6.6538814f-5 0.00012844187 0.000117294636 4.7614296f-5 -1.6270359f-6 0.000110104265 9.9959594f-5 4.9830465f-5 6.245717f-5 0.00015353833 -3.2936343f-5 -1.1067261f-5 -1.163227f-5 -0.00015293476; 7.992012f-5 0.000109838955 -6.561947f-5 -1.1899505f-5 6.17777f-6 2.0398946f-5 5.6770627f-5 1.651491f-5 0.00012576806 5.9921418f-5 -3.7098922f-5 9.494895f-5 -5.61198f-5 -7.666328f-6 3.772619f-5 2.2888104f-5 7.2048904f-5 2.2989754f-5 1.02762915f-5 -6.580987f-5 -0.000109561755 8.1869715f-5 5.8212125f-5 4.088206f-5 2.0078927f-5 -0.00013463644 9.703037f-5 6.324683f-5 0.00010554408 -1.1328338f-5 5.601645f-5 6.554657f-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{Float64}(ps)

const nn_model = StatefulLuxLayer(nn, st)

Lux.StatefulLuxLayer{true, Lux.Chain{@NamedTuple{layer_1::Lux.WrappedFunction{Base.Fix1{typeof(broadcast), typeof(cos)}}, layer_2::Lux.Dense{true, typeof(cos), PartialFunctions.PartialFunction{nothing, nothing, typeof(WeightInitializers.truncated_normal), Tuple{}, @NamedTuple{std::Float64}}, typeof(WeightInitializers.zeros32)}, layer_3::Lux.Dense{true, typeof(cos), PartialFunctions.PartialFunction{nothing, nothing, typeof(WeightInitializers.truncated_normal), Tuple{}, @NamedTuple{std::Float64}}, typeof(WeightInitializers.zeros32)}, layer_4::Lux.Dense{true, typeof(identity), PartialFunctions.PartialFunction{nothing, nothing, typeof(WeightInitializers.truncated_normal), Tuple{}, @NamedTuple{std::Float64}}, typeof(WeightInitializers.zeros32)}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}, layer_4::@NamedTuple{}}}(Chain(), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple()), nothing)

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; markershape=:circle, markersize=12,
        markeralpha=0.25, alpha=0.5, strokewidth=2)

    l2 = lines!(ax, tsteps, waveform_nn; linewidth=2, alpha=0.75)
    s2 = scatter!(ax, tsteps, waveform_nn; markershape=:circle,
        markersize=12, markeralpha=0.25, 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

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))
    loss = sum(abs2, waveform .- pred_waveform)
    return loss, pred_waveform
end

loss (generic function with 1 method)

Warmup the loss function

julia

loss(params)

(0.18949756604333226, [-0.02429727185962314, -0.0235102925974439, -0.022723313335264506, -0.02139355378875865, -0.019493164612733593, -0.01698171994136396, -0.013806140491030875, -0.009897723597621932, -0.00517442329348454, 0.0004594738063954026, 0.007104042213695339, 0.014842507661817292, 0.023692989213615828, 0.03347845021113924, 0.043544413402668895, 0.05209151462826409, 0.054764767855073324, 0.04236775346304852, 0.0013043998898729582, -0.06728901633070034, -0.11057401712371051, -0.07546148541391867, -0.005887630729846096, 0.03932667315389012, 0.05442310173379479, 0.052905432975379196, 0.04471137669164596, 0.03466090109428131, 0.02476568828310879, 0.015772452150263037, 0.00789036212236598, 0.001116189862411408, -0.004631390714853905, -0.009453363489426823, -0.013447591303105941, -0.0166977970767595, -0.01927508907835867, -0.02123494527881264, -0.022619401221224, -0.023457944297036203, -0.023768482592627253, -0.02355768784340814, -0.02282227140057929, -0.02154660792920708, -0.019704873527110516, -0.017258659030370615, -0.014156201564156347, -0.010331825004053198, -0.005705177255188662, -0.00018286587053454147, 0.006333776274715846, 0.013932125056292688, 0.022638199964946086, 0.03230884755775537, 0.04236857819526888, 0.05121137745343169, 0.054906371974836596, 0.044980946295508174, 0.008001970000435003, -0.058791397157871045, -0.10925131501527802, -0.08313775639354039, -0.013545220610642297, 0.03583095525785254, 0.05385892196671321, 0.0536417315126012, 0.045865641028061205, 0.035854307953790764, 0.02585657347665434, 0.016721687967197512, 0.00869291444195693, 0.0017889536518412967, -0.004077439031083817, -0.008997826970796647, -0.013080970725947594, -0.01640542251473956, -0.01905140855877257, -0.021070128992769, -0.02251076972866397, -0.0234005670487927, -0.02376149888491097, -0.023600841588783193, -0.022915693709195932, -0.021693979699879955, -0.019911176810095067, -0.017527472999980736, -0.014498998766166813, -0.010754945760908416, -0.006225164725013257, -0.0008102574717086608, 0.005579398561956084, 0.013040610108677844, 0.02160134253415504, 0.0311536778672496, 0.041186966993006925, 0.050274695988428264, 0.054871204925638906, 0.04719671232463329, 0.014180706287670965, -0.05012615176395675, -0.10663673878050633, -0.09016247848514769, -0.021605565676766338, 0.03185178545575579, 0.053045969131764785, 0.054288499068007465, 0.0470019533124743, 0.03705654117903216, 0.02696772992843871, 0.017687100260025474, 0.009516538604083228, 0.002472936472377069, -0.0035069925151516522, -0.008534909597366256, -0.012702569241957047, -0.016108182752455336, -0.018818982062905835, -0.02090135416075699, -0.022395720778795366, -0.02333906488176042, -0.023749690444849886, -0.02363864739629126, -0.023005226648567374, -0.021835819072261668, -0.0201106592725722, -0.01779018488265458, -0.014832659273801091, -0.011169323602887471, -0.006732717839547147, -0.0014247484121192413, 0.004841936613332766, 0.01216627423785485, 0.02058431781770221, 0.030011792552975078, 0.04000444238008188, 0.04929053029068499, 0.05467649838402125, 0.049049188597438176, 0.01982511505820106, -0.04142873511696463, -0.10280933207056323, -0.09636834868640363, -0.029999485983500072, 0.027363885920557503, 0.051959834822802135, 0.0548309809938564, 0.04811515721973632, 0.038268251051774796, 0.028092256053686957, 0.018677056830096864, 0.010354145007761863, 0.0031762187454195327, -0.0029296138688008895, -0.008055415510480538, -0.012318184015753327, -0.01580073601538584, -0.01858360173813001, -0.020723585514771443, -0.02227690387493277, -0.023272167412017528, -0.023733314144086273, -0.023672182711379188, -0.023089040984422885, -0.021973299564956476, -0.020302041683609044, -0.018048365377556566, -0.01515740412559374, -0.01157328138626246, -0.007229872314061868, -0.0020272358031986743, 0.004121378381024548, 0.011311067939462526, 0.019584122987321154, 0.028885276490355185, 0.03882326760898941, 0.04826586460438811, 0.05434473291331943, 0.050563034798303104, 0.024945435733898326, -0.03284195720900771, -0.09786029114031475, -0.10160747975806254, -0.038640281532031674, 0.022354760604880387, 0.05056566181378314, 0.055257155743325494, 0.049198911450358006, 0.03948317049434824, 0.029240583824140007, 0.019681339431580742, 0.011211427281488906, 0.0038914382928126514, -0.0023306980828522866, -0.007570457940146855, -0.011923491073441733, -0.015487689776334114, -0.018336151100332213, -0.02054352208348352, -0.022152190384822707, -0.02320063892552883, -0.023711849268721528, -0.023701008326641702, -0.023168090848201067, -0.02210353932988246, -0.020492315323040643, -0.018296134819202014, -0.01547526338934666, -0.01196898931419845, -0.007714999103791401, -0.0026154743476823547, 0.003415329411901827, 0.010471889463006229, 0.01860428045695935, 0.027774730631214887, 0.037645047772113575, 0.04720898896345348, 0.053889423994895456, 0.051770839824570855, 0.029545059303703138, -0.024466897545252372, -0.09192320884851156, -0.10575238585783545, -0.04740849124735235, 0.01680303145489248, 0.04884084017970412, 0.05554725057721764, 0.05024624121809597, 0.040703733232586504, 0.03040172703074383, 0.020707465741728097, 0.012086872928219432, 0.004625103882669917, -0.0017226635012785253, -0.007071976228598408, -0.011519078141911235, -0.015165110665122083, -0.018084970171400046, -0.02035588379825906, -0.02202214689616498, -0.023124011887267926, -0.023685988466837775, -0.023724981324067432, -0.023242279935107187, -0.022229891186186297, -0.020674119930980426, -0.01853869176677316, -0.015785369614264997, -0.01235568385311043, -0.008189444328552559, -0.004023204803994489])

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

julia

const losses = Float64[]

function callback(θ, l, pred_waveform)
    push!(losses, l)
    @printf "Training %10s Iteration: %5d %10s Loss: %.10f\n" "" length(losses) "" 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 = [0.00011340811761321875; 3.2809297408683195e-5; 1.001496730167261e-5; -9.990922990248364e-5; -4.3205916881549116e-5; 0.00010386308713350357; 9.164388757199785e-5; -0.00010796391870820414; 0.00015823921421528563; -4.437526149557137e-5; -9.740333189248127e-5; -0.0001961324596775095; -0.00013827337534179173; 9.762513218450057e-5; -8.401354716618553e-5; 8.799324859855969e-5; 8.460664685103047e-5; 9.80618715401212e-5; 8.645280468043319e-5; -2.4717075575583542e-5; 0.00015355279901988085; -0.00014274229761177725; -0.00011783649097184172; -9.8045900813142e-5; 8.313434955192135e-5; 0.00011971523053932135; -4.4612199417227915e-5; 1.7730100410047748e-7; 2.5376681151084425e-5; -2.4404123905682574e-5; 9.128665260506961e-5; 0.00014188747445579124;;], bias = [7.513204057573499e-17; -7.787086929018599e-18; 1.402926148569226e-17; -7.177620183787e-17; -1.3350383231859155e-17; 1.6085378522491906e-18; -8.086706801412573e-18; -6.555426704705998e-17; 1.5381813026955948e-16; -1.8940718113853102e-17; -8.848488555086551e-17; -2.5993408063458054e-16; 1.3144452049959313e-16; -2.742029473178671e-17; -6.970997732742797e-17; 9.574731945671133e-17; 1.4942439515562105e-16; 1.5600504270098344e-16; 1.6258218472420664e-17; -2.7237882395331327e-17; 8.649267533719286e-17; -3.764878608619026e-16; -1.345674332871704e-16; -1.7320825794467118e-16; 6.665564613005018e-17; 1.291388812065768e-16; -9.27710993001687e-17; 1.6177075365311863e-19; 1.8345767410297008e-17; -1.5579686255320013e-17; -2.2777243387074113e-17; 8.470114035731309e-17;;]), layer_3 = (weight = [0.00014156343757024603 -9.81127908049247e-5 -1.5166737726473322e-6 0.0001669967947931839 -0.00013732435247545182 -6.373702701564249e-5 3.3649315515893966e-5 0.0001634915731087201 4.0931948821183716e-5 -3.622960714307666e-5 2.660428224446919e-5 6.914981138508617e-5 1.4580960556873341e-5 -8.376499529146854e-5 1.4995228117576701e-5 -9.853553121826827e-5 1.262204357447369e-5 8.196249625728198e-5 -6.582006819696373e-5 8.570236492624831e-6 -8.837237903533234e-5 8.642094842473407e-5 0.00020047672085313665 -2.97459685712707e-5 -0.0001005670076879785 6.3698583905013e-5 -2.461712967812602e-5 -9.238552804487755e-5 1.4191784136011116e-5 -2.9871971786241185e-5 3.9675281801854735e-5 -2.221840106683417e-5; -7.396458379530677e-5 5.107019030532149e-5 -1.2511661158502664e-5 5.3000327237391075e-5 0.00010938588827542393 -1.7895561474283148e-5 -3.324880519954361e-5 -2.011808822003651e-5 -7.341926575808287e-7 7.573358750363255e-5 -0.00010959216881615798 -0.00011940515826292741 8.635145065768957e-5 -0.00018098313005480667 0.00010790038511317209 -9.916827541615756e-5 -0.00025258510134022695 6.424731698047627e-5 4.1515675614793153e-5 -1.2524355885549467e-5 -1.48082062210411e-5 -5.71464551526727e-5 -7.472054851546175e-5 -0.00012909265926878643 -3.4089102214775824e-6 4.954006733300922e-5 -9.364978170794059e-5 -5.1097602342058164e-5 -4.654105207606916e-5 -5.91896422761842e-5 6.369273621516423e-5 -2.9778153408682206e-5; 2.080392184726604e-5 8.364434483539173e-6 0.00014500618269669093 0.00014755497975276247 -9.780596906765628e-6 7.23296761629298e-5 -0.00015231828731483268 -0.00012671414881477796 8.047329172259393e-5 7.517566698372422e-5 4.016419550172481e-5 -3.5382105148172634e-5 1.1655431966944003e-5 0.00010767379715108738 -1.312925635428024e-5 2.5854791578618812e-5 4.6661047912647175e-6 -2.440097889628774e-5 -8.058899404712374e-5 0.00017241381281057533 1.4759998497905472e-5 -5.756789325968748e-5 -0.00015111672112251527 -0.00014129154444670814 3.9221653400470194e-5 0.00014003891553730774 -0.00015051852099129873 0.0002905739529291374 0.0001439860497834363 8.67763955961184e-5 -0.00012131241239838717 -1.0862933800666548e-5; 0.00011003540320985874 0.00016107860328362227 2.7666877300598605e-5 -2.4925599703910645e-5 -0.0001383721780291489 0.0001350676077856954 7.60109444820593e-5 0.00010802489604999232 5.2456741556180096e-5 -9.955637024517266e-5 9.03624507339011e-5 7.731728446402553e-5 0.00010679161395843074 -0.00014037587489268823 5.186789830646392e-5 1.2004731354501794e-6 0.00011025441681000318 5.02900632420826e-5 -5.135805868461814e-5 -0.00015502918786636653 0.0001984408784630982 6.474913149086592e-6 -0.00016835767992573865 -4.077879297019334e-6 2.225582947745367e-5 -6.627253124292083e-5 -7.437316711493553e-6 -9.253227717976185e-5 -0.00016317543732269766 -0.00023337111152044597 -6.327397086692127e-5 9.51070369070332e-6; -4.8778078363677966e-5 0.00011100884845161404 -0.00012343701413685817 -8.922338051109481e-5 -1.402059510058004e-5 6.927181644511462e-5 -0.00014855251387546598 -7.743399359372436e-5 -1.1390342789526896e-5 -6.645236991450499e-6 3.4403637672375185e-5 0.00021861761311679718 0.00016305455906928163 -8.374935745580939e-5 9.215792639094638e-5 9.470817863854779e-5 2.5165756286305805e-5 -0.00019379130421060025 9.986756018276539e-5 9.767388806592288e-5 8.304426561022541e-5 3.5202992565758127e-5 -0.00013444291866238325 6.376238858007937e-5 -7.773871069860403e-5 0.00012028991446510772 -0.00015718577227874175 2.4691709457851195e-5 2.6268758718809916e-5 3.7739700618270674e-5 0.00010786126191175715 1.7243091550484012e-5; -0.00021674562541111256 -5.0069800098847145e-5 0.00019511008630697447 -1.3812300804604095e-6 -0.00010837227401598483 6.231218921211623e-5 4.7232254101708734e-5 9.554802872198726e-5 -0.00019214351724176286 -2.32257674807462e-5 -1.2652120321127148e-5 -3.043319103116418e-5 3.8926602965968414e-5 4.5956611384917504e-5 -1.1541487228078867e-5 0.0003004032338057179 5.638132816692738e-6 -0.00015406166440616316 0.00014416527173207141 -1.731132525410655e-5 6.277857809518819e-5 9.007860407191615e-5 6.238701272105341e-6 8.975455474765087e-5 -1.528854173475424e-5 -4.186534661957844e-5 7.42850466766848e-5 -0.00023177619678708893 -9.166017062771287e-5 -8.983364333207872e-5 -0.00015799025951228086 2.3860373940717423e-5; -1.5070638536403707e-5 6.996644396278593e-5 1.7786675645058085e-5 -2.039674228288594e-5 -4.160917623066188e-6 8.639188917297006e-5 -1.4840366301343856e-5 2.3651686580848603e-6 0.00011413568308053878 0.0001123968237982086 1.3266360557478735e-5 -1.4757544985330242e-5 -3.4512061855628554e-5 0.00012556365863130556 0.00022367374794268288 -4.676527288100474e-5 4.063429688556385e-6 -0.0001163626523600972 -4.011233005147206e-5 -6.0060546555886234e-5 -0.00013919075837183338 8.280634116201309e-8 2.587760235766339e-5 -9.960583437335389e-6 -9.324905232350095e-5 0.00017469648327388706 -4.105246744873458e-5 -8.128372709305402e-6 0.00015975460582123678 -0.00019158564219540374 -1.7480302002498e-5 7.957774386451214e-5; 0.00012458964485070258 4.6487562225442984e-5 5.141409493611293e-5 -4.6743691915407e-5 3.42907762513592e-5 -0.00013141760135674247 -9.013137313210637e-5 -5.593880620080865e-6 4.126115819765292e-5 -4.733324820896586e-5 0.00021864929586975555 4.893676961166286e-5 -4.8465030691816545e-5 -1.5559992594870712e-5 -0.00011432967292982331 -9.151559496437405e-5 4.554863263713946e-5 -7.504623207402699e-5 0.00013871836128040018 6.939028682048511e-5 3.4822143073893166e-5 -8.592578231962174e-5 -1.6072267304635452e-5 9.697758215573333e-5 4.9591376604287096e-5 -1.848535319929479e-5 0.00013458593749768118 -8.859115476858699e-5 -1.8501052896831343e-5 5.312759023021053e-5 3.663873141143243e-5 -3.887787969340339e-5; 4.303014925686997e-5 -5.811431633185891e-5 -9.79165620693778e-5 -6.77093762962012e-5 3.298575123033708e-6 2.289210665627786e-5 -0.00014672642878134272 6.326625721418681e-5 7.894566308879932e-6 -3.697971128270348e-5 -3.138994956965124e-5 0.0001261625524901853 0.00016980013738508544 -0.00014502634307712815 -0.00022832649715122275 7.680718540802141e-5 0.00017640138089779323 -0.00013793059638113282 -4.899751961341577e-5 9.5820383151357e-5 -0.0001046095554076032 0.00017992390480945993 7.82762085262812e-5 7.714329099405058e-5 7.782794405363065e-5 -2.0398972966495897e-6 0.00010197307115310486 0.00018598271113011718 -0.00011490359479205969 8.835876954316101e-6 3.188394244067833e-5 -0.00012890677912504232; 6.808961107100173e-5 -0.0001839414116117778 0.0001500647114001545 -5.576064076802722e-5 -0.00013817528897259903 -5.937862160561679e-5 0.00016808872815392037 -3.9984934934807296e-5 -4.8632501508741786e-5 3.1603266600830844e-5 3.2329992885310185e-5 -1.461224373697728e-5 2.5013912002858073e-5 -0.0001840725680237357 -3.670174273491706e-5 3.737508368732501e-5 7.883299904580354e-5 0.00011071386607301411 -5.411694009734209e-5 8.52083822822526e-5 6.018232225260856e-5 0.000244314759678043 5.2739435980867034e-5 2.4210613540448153e-5 9.477498597636026e-5 2.6218852811425456e-6 2.9119014564747344e-5 -0.0002215721259714705 6.012940596810234e-6 8.044457272564545e-5 0.00015570963214787714 -9.232959646096108e-5; -0.00012052728766016628 1.345794774662154e-6 0.00017832094249874074 -8.839214046625454e-5 -3.1913003337339036e-5 -6.282960589183511e-5 -6.349284580815657e-5 -8.533927234262223e-5 -0.0001635666724238512 -3.939540533488847e-5 1.2402019595668644e-5 7.976421852922727e-5 -0.00012994080520607344 3.4588965493952394e-5 1.6260708210523107e-5 0.0001936879250509655 7.675822212453826e-5 0.00016140450111683407 5.9779029046175644e-6 -9.106905985996391e-5 -0.00010549835337465332 -5.169182463455728e-5 -0.0001264225741093973 -6.489495194423989e-5 0.00017515504237357413 9.690743695804918e-5 1.4111066123782307e-6 -5.8078933076702974e-5 4.372184551667584e-6 -1.1363276140999835e-5 6.205930344853473e-5 1.3318561062393504e-5; -3.359187637935045e-5 -0.00011325144685354637 -0.00012956404923709819 -2.2925522605796415e-5 3.3070258106883014e-5 0.00011392865901333302 -6.090194681557243e-5 -0.0001918051518058854 1.446806647628459e-5 -0.00011785491613220838 8.721928799976417e-5 4.7588885829221925e-5 9.963713575607555e-5 -2.4956884479175176e-5 -0.00018551706550890205 -6.617169372286164e-5 5.006242765564864e-5 -2.895694404097898e-5 -0.00010745287243291294 -5.200530602262982e-5 0.00011207913969165798 -0.0002498649216048734 5.9772134884125966e-5 -0.00012754190054503427 -0.00011753411916099134 -8.186896300601134e-6 -0.000177140284611531 -6.696004373037046e-5 0.00015603444382765637 0.00012924298546028624 4.130720237488501e-5 -0.00022212591527740852; -7.12879512357908e-5 0.00016443685007827246 -3.086352017278893e-5 -7.669471533889053e-5 -2.621524918511618e-5 2.8675571647177333e-6 -0.00016141446695117382 -5.6848060288830726e-5 0.0001723539505255027 -2.3746098399060926e-5 6.130451647179097e-5 -0.0001773188517327801 -6.232086462121553e-5 1.869281613464365e-5 0.0001519607655572322 6.157045272258937e-5 -3.4617506848081115e-5 1.310896977512995e-5 -6.496420547053323e-5 3.5059984562321405e-5 -7.156662041241395e-5 6.247820847049261e-5 -0.0001029193838840733 5.784677041076022e-5 3.21092345038046e-5 -5.652751797613786e-5 7.247853199798505e-5 9.330147867869475e-5 6.622542850520353e-5 -6.516574222048855e-5 4.405185308853061e-5 -3.328216126108635e-5; -0.00011829816435527815 0.00016207738439954578 -2.640705065748855e-5 -1.904235910222566e-5 -2.7641531463067956e-5 -0.00013119536675111127 -7.477604163176012e-5 8.691669131332304e-5 8.218601646768502e-5 -8.74453028274542e-5 -3.098129605239868e-5 6.31645168567075e-5 -2.3709021891761634e-5 1.7485507439066125e-5 5.414762995025569e-5 -0.00010163256161211592 0.00015427227549191056 0.00028224203628889603 -4.073939436985382e-7 8.089717516379208e-5 0.00013113563249754422 2.7630749852334466e-5 9.689320790750103e-5 -6.688469079852722e-5 8.954550207528721e-5 -4.371149339311647e-5 2.7111211917891125e-5 -2.328225969284038e-5 -0.0001643371499568991 -5.12401559291709e-5 3.840063504577986e-5 -7.731852315663503e-5; -4.226616725942328e-5 -0.00015316125278628953 8.888036030542488e-5 -0.00010478641558007625 8.267452937375857e-6 -6.892646557686185e-5 8.682456030941239e-5 9.083129192434941e-5 -2.5446459733991653e-5 7.05268590870543e-5 -6.068150684155822e-5 0.00013888228353311783 3.7212661151662046e-5 -4.4191734890102586e-5 0.00010942871234577717 -4.1259254761155415e-5 -1.5920088447512935e-5 -0.00011855541594128422 -8.714197098312321e-5 -1.3424529573298754e-5 -7.128625508199369e-5 -0.00022807707351899713 -4.6776045895295154e-5 2.4985936430983092e-5 8.09527606513594e-5 -1.9020112052440478e-5 2.4524507932015364e-5 0.0001824646584616267 -6.445995892433053e-5 -1.1317951578589392e-5 -2.6922988002536386e-5 0.00010507008631679217; 0.00010393709879115209 -0.00016571840326088657 0.0001827012317136757 0.0001312625105839984 0.0001386636583248969 0.0002059233963138249 0.00012411750565495617 -4.353480410084066e-5 -0.00015656962870874323 9.439473131186093e-5 -7.965951892567597e-5 -0.0001037803516026942 -8.540174105458389e-5 -6.30355637638996e-5 3.7717290326133536e-6 -0.0002609643264093591 3.6426100157931285e-5 -9.350475677076532e-5 -0.00013874238295267927 -3.702639089334902e-5 -6.161010180388364e-5 3.428244716059569e-6 0.0002688139217554397 -6.840073849327667e-5 9.941626256311682e-5 4.541126752717308e-5 0.00011233708999332734 4.839509408901338e-5 0.00011278739173411242 -1.1173673848876791e-5 0.00017866971091114268 -4.22231308419493e-5; -0.00012568483486277395 -2.752909983132334e-5 -6.674869779703655e-5 -6.423803403479837e-5 -1.9777523781835938e-5 -9.57694368177998e-5 7.665682517951254e-5 4.267876159544504e-5 2.7813125264669276e-5 1.054419465913742e-5 -0.00011866679539410329 -0.00013129686828204548 0.00013542888212512869 -3.03730170041781e-5 -3.120364032541709e-5 -0.0001777911105516387 -3.759079969926099e-5 -9.196636653244408e-5 0.00011062765282373632 -6.519432042000254e-5 -3.377281006279023e-5 -8.027311684301238e-5 -8.575506341614948e-5 2.3251931118709897e-5 -2.9731193869098598e-5 -5.726223877602375e-5 -0.00022664033291834563 -0.00011635478710263501 -9.25253565280828e-5 8.771236628786288e-5 6.113343136645743e-5 0.00010376913872577096; -0.00015122185301332875 0.00022250471214466263 -0.00013489534627478134 8.747952800878098e-5 -8.549639620650072e-5 0.0001285770906524806 -7.220354906347848e-5 6.827427408826356e-5 -6.43854234676697e-5 -0.0001648012200912412 0.0001128467103972169 8.00836626131235e-5 9.773836631030797e-5 -1.4332731336243498e-5 1.1986122661264833e-5 -0.00011684159268223966 -3.846263076592538e-5 -1.3899278894304324e-5 0.00019769867343530695 -5.075230161491328e-5 -9.484400618007393e-6 0.00012867147437465432 0.00013268147289455752 -0.00021376269593200876 4.6815296063960455e-6 1.910230479371626e-5 -8.316780600761503e-5 -8.809550650948221e-5 -6.133469447558291e-5 -9.859773637792222e-6 5.0000196442947294e-6 -0.00012287329606077582; 2.360132065268529e-5 4.652344658869336e-5 -2.9177164389254398e-5 -0.00024712583900646305 -0.0001787671985871521 -0.00016922940768518402 7.474262610093528e-5 -6.339087623155503e-5 -3.670025536074621e-5 -7.583825709981733e-5 0.0001566200758029422 2.525108947204553e-5 -3.0596658245027004e-5 0.0002041084612651444 7.626657542322463e-5 9.445416697864462e-6 1.2121894077145867e-5 6.024019102232495e-5 -5.854681159242991e-6 -3.593599241095624e-5 -6.66069004287152e-5 7.398801470890775e-5 -0.0001399306688758784 -0.00022909214522906618 -0.00011389571364953497 -1.986998678285182e-5 -8.935251890899502e-5 0.00012020758230253962 -1.5914763490869912e-5 -8.78100959303233e-5 -1.1514247045404563e-5 -7.858970538968355e-7; -0.00020551709087612304 -5.361466094200328e-5 0.00013410371400368528 -6.972194682926488e-6 7.021439500931643e-6 1.4081558674704415e-6 -0.00016750048882840612 7.566269130166367e-5 -0.00011994643784872798 5.184938061120203e-5 1.449692706501452e-5 -2.3629470092583397e-5 -3.853043301659532e-5 3.485971399186755e-5 6.828895209310104e-5 -0.000119701281732884 -0.00011500033645006365 -7.423524607471215e-6 0.00018221956349862283 3.76841728666422e-6 -5.0913660061164744e-6 -5.911198525153868e-5 0.00010093913522269429 -0.00011955427100928851 -6.401455823379118e-5 6.397441109130032e-5 -2.8890987891827427e-5 9.627082260976665e-5 5.552262223317858e-6 -9.948922445346455e-5 4.513970350217593e-5 9.700789860377595e-7; 5.606457073773832e-5 -6.931720105749238e-5 0.00011739212953434545 0.00022472473522831035 -9.100877338605537e-5 -0.00010641930254459017 -6.104050095708916e-5 -0.00010733866344488706 -3.158506712459091e-5 -6.691584400650639e-5 4.833218138847734e-5 -7.310799508157558e-6 5.599440504049223e-5 3.3121595545171355e-5 -3.895903552583429e-5 -2.787468162531907e-5 -3.266916298124941e-5 -7.980002971416092e-5 -5.400279093700431e-5 2.797219115782334e-5 1.223141992561542e-5 -1.6832104454423037e-5 -0.00016058690068128771 9.991932013153696e-6 1.255424870797174e-5 -0.00018376721397324677 -4.903497080485858e-5 -7.508007235410415e-5 -0.0001344230955215609 -0.0001444432533095098 0.00017147169585988043 3.706063892643342e-5; -7.033361583384278e-5 -2.9559043862575152e-5 -0.00015901120462589668 2.4420563178493688e-6 2.3967013082399655e-5 7.187264015297785e-5 0.00025936750495671635 -0.00012734980230819974 -1.3552194630741628e-6 2.4443759291132616e-5 2.5813149083805134e-5 0.000128404793547055 0.0003317023397968666 -0.0001346098255752084 -9.349829309289019e-5 0.0001201753289634861 9.427153014908871e-5 -0.00021998634877749357 3.9749807517326934e-5 0.0002036233377722581 -4.927947196196201e-5 -9.939028353332027e-5 -8.36188741530715e-5 9.238977833516327e-5 -0.00010784760200553251 0.00010294939017254359 -0.00012627878134739332 -8.869453125302459e-5 -0.00014664450685162794 4.241243693027877e-5 -0.0001641723906076013 3.200135649163787e-5; -6.303013307136981e-5 -0.00015279625605421108 1.0534838827657502e-5 -5.391597752554006e-5 -0.00022789225860128122 2.3253481947685498e-5 3.177749002839721e-5 -8.384852500261346e-5 -4.575537260483101e-5 0.00010672971601503032 -0.0002306999051254397 -2.055756541470155e-5 0.0002530361996721355 -7.335932988302461e-5 -4.391299463554048e-5 -6.41399639900934e-5 -1.5418612940485992e-5 -3.761115677865829e-5 8.352512285251422e-5 0.00017013232335353286 -0.00011105549148145262 -3.261647017301301e-5 0.0001751986745036999 5.339435836284861e-6 -5.697195975903189e-5 -6.49630275913652e-5 -3.781674623700726e-5 0.0001502510456591658 0.00013033247140606884 -3.127883629735544e-5 5.205951291623806e-5 2.967669650172885e-6; -0.00013140824427080761 6.879190560602488e-7 -8.877502528902734e-5 1.8926596966244122e-5 3.39699212769189e-5 4.737961666125205e-5 -5.750478572862095e-5 3.267941014075129e-5 -4.9980747063609975e-5 1.4618255989815027e-5 7.81202601375745e-6 4.124169521550574e-5 2.1996627254906386e-5 0.00010458186028573887 -4.280950139344658e-5 0.0001611501970170259 -0.0001317308456795102 -2.8590172006611258e-5 1.9955406458951072e-5 -6.708773849949787e-5 -4.714825314419769e-5 -0.00014115769277975667 -8.095236677161747e-5 -0.000142319343068616 5.2801391064906496e-5 -3.17637937417935e-5 9.363922289237784e-6 -0.0001396163539187856 5.2047598218092865e-5 3.732462522621653e-5 3.265630897532339e-5 -0.00013716779045478233; -4.1252708179688107e-5 -0.00022208844922789092 0.0001168301675344947 -8.488827298437556e-6 -1.0954239151766738e-5 -7.348912942763055e-5 9.675560402609983e-5 4.81493986437126e-5 0.00024390568403684587 2.9752316258230415e-5 0.00010636887571234735 -6.593652535306478e-5 0.00020534933690139663 -4.492036561788516e-5 -2.340227758469301e-5 -6.350941141209863e-5 1.667919757818553e-5 0.00013496394938482435 0.00014743406762062289 6.582569010715321e-5 3.958450513816618e-5 -1.4161744112083752e-5 -5.242613071533757e-5 8.270839287359735e-5 -0.00013805486407437594 -4.669996662489083e-5 -5.79479022424555e-5 -1.6150432494036837e-5 0.00012498098629220707 -8.811090840099988e-5 4.902242624592614e-5 9.599928278401043e-5; 5.9383588612713695e-5 -1.755632703190326e-6 9.620851249450232e-5 -5.185466017906627e-5 0.0001556004883047889 2.012640486895542e-5 5.1844517819283844e-5 1.0331499814805724e-5 1.717738466936875e-5 -0.00018853762396362167 -4.6425522506303215e-5 5.2457397194993824e-5 0.00015092687879473374 -5.895178930700853e-5 -8.778863941302747e-5 -1.1117294726552934e-5 4.5380771715592257e-5 -0.00011302694276424215 -0.00011277583491506257 -5.206317457237841e-5 0.00014770575415127613 -9.421218221291434e-5 -5.052612852637526e-5 -0.00010159620962554679 -0.00012724361824553696 -0.00010161230404378691 0.0001239672875988062 -9.79494560135774e-5 -3.278019523464701e-5 -1.144652362314614e-5 7.075574301665491e-5 -4.143238931760667e-5; -1.921490506303115e-5 9.685955815855382e-5 -8.232440954400333e-5 -3.7459785329171656e-5 4.8494416165864364e-5 0.00012407023537578757 -3.466009506466285e-5 3.947379669315008e-5 0.00016354236364040417 9.989804898994078e-5 1.4014131426467928e-5 -0.00011594398305841759 -5.537831161146601e-5 -0.00021627799064246342 -1.9356324213203558e-5 5.7540643371413444e-5 -1.9664182712785406e-5 -0.00028484427234017615 2.198700208880918e-5 -5.887122864282661e-6 -0.00018816684818685276 -5.448704318351568e-5 -6.631092177796186e-5 3.916046121647255e-5 -0.00015688394118425965 -3.0002108854581353e-5 -1.4493768088344592e-5 0.00017595433242292963 -6.969778545992512e-5 -0.00012779777675261792 8.119459603820507e-5 8.8584407837144e-5; 0.0001634354332115923 -3.8322138419298555e-5 -9.389862024660148e-5 -5.016349750189852e-5 2.352138399759568e-5 1.505840028289748e-5 1.8057838321281815e-5 7.480003042178597e-5 8.23833899692252e-5 0.00014278738600000078 -7.00463470593901e-7 -0.0002583173852992481 4.086520421584365e-5 -7.126212761477542e-6 4.96958957289706e-5 -0.00020397276311382698 -0.00020272011969905458 8.480903416671008e-5 -9.941565434248948e-6 8.841221206852362e-5 -1.1268781957582078e-5 -0.00016806150305750416 -8.867638170454046e-5 -4.643494393644703e-5 2.944153785096812e-5 -8.879767191796842e-5 -0.00016809568550639513 -2.0562397179002253e-5 0.000165337121149368 -6.290212403952018e-5 -9.418806511644259e-5 -0.000134093433001215; -5.764134506850435e-5 -3.5411591983080344e-5 -4.657727653780524e-5 -7.146514066650173e-5 4.285693908654403e-5 -1.1386182618180196e-6 -0.00010710466553379258 -6.31990707401819e-5 0.0001267714822975151 5.160522674172425e-5 0.0001637522790192947 0.00017053160887053518 6.89192664654217e-5 -0.00014134943298219296 0.00012690834306025602 8.293054166327819e-5 -0.00011360625025505486 0.00021434444248055506 6.099463904005158e-5 9.305033796903715e-6 0.00010037961952231741 -7.021727145436364e-6 -0.00014967394340007598 0.00010103817372194889 -5.670301692100751e-6 -0.00012880284443187603 4.3170420082372196e-5 4.487054281112227e-6 5.9014108825562515e-5 0.00018500354677284517 0.0001608799201281142 -9.26450385518159e-5; -0.00017676844854247153 7.659133862364978e-5 8.920188284140042e-5 -0.00014886817789044494 0.00010899548524651599 -0.00017075219202194645 -4.0270261339510293e-5 0.00014151389350775328 2.8667621703449052e-5 -4.805511398114343e-5 -3.801268812491813e-5 3.167641935180873e-5 -9.472035580654177e-5 1.3138818215911507e-5 -0.0001239019808020835 5.5683506983432634e-5 7.649669296700307e-5 6.441084715664127e-5 -1.8557259496783133e-5 0.00011456229131329756 -0.00014423596857920872 0.00022011328419796618 -2.675510094956853e-5 0.00016565796559337275 8.545095934426217e-5 1.5266722951881627e-5 -2.3169567244957086e-5 7.68524333000847e-6 -3.21838020648719e-5 -3.691082801776055e-5 -0.00015956427214075105 0.00015980754269878813; 0.00017030300779866025 7.528406542391636e-5 -2.6628044615932215e-5 -8.729442804919775e-5 -1.2370358199545344e-5 7.286619194915354e-5 -7.815793531349254e-5 -0.00016535475759169198 -9.285787262355123e-5 0.00010033397968058651 -6.573771559501319e-5 8.424473464390085e-5 2.3014682774245e-5 2.4644966579081795e-5 -8.655539448241187e-5 -0.00020019175161010023 -0.00012179335924897751 -3.831231395465566e-5 -0.00022175280053502336 2.123708811254327e-5 -5.8233809504725604e-5 -0.00014432849155140566 5.547527798493813e-6 4.462544286057109e-5 5.736257418889129e-5 -3.3973341761749283e-6 -4.217994381192893e-5 6.254971351140226e-5 2.2904864926972383e-5 3.9565007063942874e-5 0.00014557747028434318 8.495695838894703e-8; 8.146005968961156e-6 2.437213632474753e-5 -8.608434356319649e-5 -9.678878333688357e-5 -6.34832910847745e-5 8.274112704639703e-5 1.5761833567743514e-5 2.3337387831638905e-5 4.123607075854067e-5 -4.949061323459135e-5 -0.00010363028624596664 -1.890049017418853e-5 -3.3620305400465896e-5 1.1176800656180084e-5 -5.013448091376301e-5 -1.401825076499894e-5 4.7043107117864265e-5 -0.00010790360894788642 -4.333363605582802e-5 -5.609984938673589e-5 6.785825760999516e-5 7.412960920447214e-5 -4.83457194761233e-5 0.0002574886086975351 0.00015002031367003255 5.365591576519906e-5 8.588733843025875e-5 -0.0001884602613401499 -0.00013669778920047827 -4.1489799257031715e-5 -4.894120931312132e-5 2.5305648962610405e-5], bias = [8.78277597981429e-10; -3.2886660548089774e-9; 2.3282831741638547e-9; -5.339074046958915e-11; 1.5130886586339904e-9; -6.834025976963312e-10; 1.4394492805895284e-9; 1.2398357606550287e-9; 7.361930954175261e-10; 1.694482598032661e-9; -7.587429163944486e-10; -4.811843547738572e-9; -2.483291661334235e-11; 1.354007346458545e-9; -8.316085449633215e-10; 2.0378849365286676e-9; -4.803375075565768e-9; -8.383376480161608e-10; -3.177330215442913e-9; -1.4980514615009741e-9; -3.154643855759494e-9; -4.4786506110345795e-10; -1.0588448681405896e-9; -2.6889769233885316e-9; 2.720916377398536e-9; -1.627832955083342e-9; -2.318796237193203e-9; -3.311282181487449e-9; 3.1276218078120953e-9; 7.234153893101507e-10; -2.0724475286069135e-9; -8.955950748849676e-10;;]), layer_4 = (weight = [-0.0006275460440796759 -0.0007364743304650615 -0.0007177732659908027 -0.000856485899052216 -0.0006396642091025211 -0.0007610457783847459 -0.0006876074708222356 -0.0006576106131910669 -0.0007601339209891256 -0.0007693143708210133 -0.0007301111603465727 -0.0007631255252619243 -0.0007276870356567966 -0.0007826147882128564 -0.0006089732472065934 -0.0006362186864793315 -0.0006071412556209678 -0.0006799788895739869 -0.000775487294345312 -0.0005805067758268412 -0.0005916538616064965 -0.0006613343873279421 -0.0007105757004987259 -0.0005988442816654813 -0.0006089889493713836 -0.0006591181703954606 -0.0006464914112885562 -0.0005554101427461218 -0.0007418848297930867 -0.0007200159375660575 -0.0007205808702151817 -0.000861883434801961; 0.0002881161825082743 0.00031803495459163824 0.0001425765648505006 0.00019629655950804698 0.00021437382135917587 0.000228595007341271 0.0002649666788485552 0.0002247109659436449 0.0003339641171754653 0.0002681174652970925 0.00017109713914613642 0.0003031448726378125 0.00015207626500963098 0.00020052972522055382 0.0002459222499665832 0.00023108414418714516 0.0002802448371272439 0.00023118581448045524 0.00021847229538660688 0.00014238617944146545 9.863425357712819e-5 0.0002900657786168766 0.0002664081825282318 0.0002490780831273305 0.0002282749496627817 7.355960508659139e-5 0.00030522640686647163 0.0002714428321037283 0.00031374008659704243 0.0001968677236724259 0.00026421249073420744 0.00027374262955368387], bias = [-0.0007089486872246964; 0.0002081960644234756;;]))

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; markershape=:circle,
        markersize=12, markeralpha=0.25, 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; markershape=:circle,
        markeralpha=0.25, alpha=0.5, strokewidth=2, markersize=12)

    l2 = lines!(ax, tsteps, waveform_nn; linewidth=2, alpha=0.75)
    s2 = scatter!(ax, tsteps, waveform_nn; markershape=:circle,
        markeralpha=0.25, 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; markershape=:circle,
        markeralpha=0.25, 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(LuxCUDA) && CUDA.functional(); println(); CUDA.versioninfo(); end
if @isdefined(LuxAMDGPU) && LuxAMDGPU.functional(); println(); AMDGPU.versioninfo(); end

Julia Version 1.10.2
Commit bd47eca2c8a (2024-03-01 10:14 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 48 × AMD EPYC 7402 24-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, znver2)
Threads: 6 default, 0 interactive, 3 GC (on 2 virtual cores)
Environment:
  LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64
  JULIA_DEPOT_PATH = /root/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6
  JULIA_PROJECT = /var/lib/buildkite-agent/builds/gpuci-3/julialang/lux-dot-jl/docs
  JULIA_AMDGPU_LOGGING_ENABLED = true
  JULIA_DEBUG = Literate
  JULIA_CPU_THREADS = 2
  JULIA_NUM_THREADS = 6
  JULIA_LOAD_PATH = @:@v#.#:@stdlib
  JULIA_CUDA_HARD_MEMORY_LIMIT = 25%

CUDA runtime 12.3, artifact installation
CUDA driver 12.4
NVIDIA driver 550.54.15

CUDA libraries: 
- CUBLAS: 12.3.4
- CURAND: 10.3.4
- CUFFT: 11.0.12
- CUSOLVER: 11.5.4
- CUSPARSE: 12.2.0
- CUPTI: 21.0.0
- NVML: 12.0.0+550.54.15

Julia packages: 
- CUDA: 5.2.0
- CUDA_Driver_jll: 0.7.0+1
- CUDA_Runtime_jll: 0.11.1+0

Toolchain:
- Julia: 1.10.2
- LLVM: 15.0.7

Environment:
- JULIA_CUDA_HARD_MEMORY_LIMIT: 25%

1 device:
  0: NVIDIA A100-PCIE-40GB MIG 1g.5gb (sm_80, 4.600 GiB / 4.750 GiB available)
┌ Warning: LuxAMDGPU is loaded but the AMDGPU is not functional.
└ @ LuxAMDGPU ~/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6/packages/LuxAMDGPU/sGa0S/src/LuxAMDGPU.jl:19

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 ​

Defiing 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

Defiing a Neural Network Model

Setting Up for Training the Neural Network

Training the Neural Network

Visualizing the Results

Appendix