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, 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(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[4.5613913f-5; 0.00019234157; 2.7182374f-5; -0.00020461941; 9.7079566f-5; 1.5926082f-5; 6.3381216f-5; -5.7683428f-5; 1.641379f-5; -2.2179664f-5; 9.268546f-5; -5.223045f-5; 0.0001204457; -0.00014333909; 0.00010610596; -3.505027f-5; -2.3483948f-5; -6.5862885f-5; -7.428432f-5; -4.335849f-5; 6.996508f-5; 4.3917116f-6; 5.3925933f-6; 4.8963153f-5; -4.0818966f-5; 6.404497f-5; 2.49128f-5; 7.115352f-6; 5.3040326f-6; -1.5418673f-5; 9.545954f-5; 9.762855f-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_3 = (weight = Float32[-0.000103871644 -0.0001955706 0.00018313213 -3.348735f-5 2.3271165f-5 8.464796f-5 6.153664f-5 -2.7462604f-5 4.6800833f-5 0.00018489499 8.390951f-5 -4.8735405f-5 -2.582679f-5 -0.0002018403 -0.000112296984 0.00015591274 9.1649f-5 -0.0001734883 8.097614f-6 -0.000113494454 -7.404662f-5 -6.148624f-5 -0.00017750943 -0.00025185096 -4.8596918f-5 -3.175463f-5 6.4319036f-5 3.1680283f-5 -6.941056f-5 -7.608777f-5 -1.0845364f-5 0.000101776684; 6.5321045f-5 0.00011945266 -0.00015950488 0.00028341066 -3.198685f-5 -0.00016115831 -0.0001617231 5.421374f-5 6.7700494f-5 -2.8499186f-5 -3.8963408f-5 -5.8950267f-5 7.745632f-5 9.9532655f-5 -8.094304f-5 -0.0001279928 -0.00016509272 -0.00025744553 5.4018547f-6 2.9885645f-5 -2.0733689f-5 -0.00014916636 -0.00013712114 -5.852838f-5 1.8889572f-5 -3.4464887f-5 4.645891f-5 -1.2841978f-5 0.00012506856 3.8151276f-5 -4.8214788f-5 4.2537253f-5; -0.00010645085 9.8868f-6 -3.131619f-5 -0.0002503466 -0.00022683004 2.4377357f-5 -3.247971f-5 -0.00014600155 -8.1020065f-5 -0.00016472612 -4.1419236f-5 4.8352464f-5 -2.3080012f-5 -6.659357f-5 0.00014902632 -5.510511f-5 8.5940184f-5 -0.00020290818 -1.3958499f-5 -0.00020731182 -1.2720678f-6 0.00021396055 1.3502881f-5 0.00016699231 -9.935121f-5 4.4258322f-5 -0.00012692524 -6.117882f-5 -4.0852858f-5 4.571077f-5 3.2980763f-5 -7.183391f-5; 0.00014823547 -8.25422f-5 -5.0683975f-5 -0.00022103773 -4.3117067f-5 9.581825f-5 -0.00023411808 -0.00013024657 6.348267f-5 2.986057f-5 -1.9693516f-5 -2.5281617f-5 -5.3500073f-5 5.8060457f-5 7.567137f-5 -5.396586f-5 4.7463243f-5 -0.00012089248 -5.6908334f-6 9.079194f-5 7.095525f-5 -6.9431394f-6 0.00014009644 -7.3530064f-5 9.5834184f-5 5.6188387f-6 -4.6421213f-5 0.00012088091 7.572651f-5 7.3147145f-5 -4.5287274f-5 -0.00011222701; 0.00015719175 -3.9704453f-5 -8.19519f-5 -0.00014138929 -9.621486f-5 -5.2346622f-5 9.671993f-7 9.443894f-5 -7.783321f-5 -3.379931f-5 -8.567318f-6 3.515014f-5 8.190372f-5 -9.961566f-5 -7.827495f-5 2.8160423f-5 -2.3362818f-5 -4.413887f-5 3.616325f-5 -7.3819305f-5 -0.00012590917 3.6275735f-5 9.05954f-5 0.0001123973 0.00017515731 6.6351306f-5 5.8705722f-5 8.368102f-5 9.784018f-5 7.6341355f-5 2.4668725f-5 -8.310276f-5; -4.9484856f-6 -5.674502f-5 0.00016316204 -3.4085584f-5 3.2051674f-5 -0.00012594806 -0.00014219274 3.2745764f-5 -7.1911396f-5 7.2211496f-6 -2.8571001f-5 2.570706f-5 -9.618583f-6 0.00017646738 0.00015900779 0.00016743942 -1.8788027f-5 -2.582109f-5 5.136481f-5 -6.0678554f-5 5.6842226f-5 8.121977f-5 -3.4063527f-5 -0.00013555364 -0.00012540103 7.4801724f-6 0.00012265045 -4.2004358f-5 0.00016969099 1.1933902f-5 0.00011712605 5.790966f-6; -0.000118732045 8.413054f-5 -2.173228f-5 -8.037065f-5 -5.79429f-5 3.8050064f-6 0.00024216047 -4.6636666f-5 -6.231574f-5 -1.0642348f-5 -0.000102071346 0.00010325349 4.1388404f-5 0.000106330255 -5.9297752f-5 6.990904f-5 0.00018916787 7.047474f-5 1.7748642f-5 2.9404706f-5 -0.00011297875 5.061542f-5 4.406678f-5 -5.7820316f-6 -5.1122202f-5 -0.00020620992 -7.7225326f-5 3.4396653f-5 -0.00021106994 -9.702334f-5 -5.7567995f-6 -5.1196486f-5; -0.000104118626 3.8273167f-5 -0.00014140256 9.840149f-5 -5.1962015f-5 -9.202226f-5 -0.00021619054 0.00011011333 5.318181f-5 9.672245f-6 0.0001652877 -0.00017707245 -7.4622796f-5 -0.00012672161 -4.218254f-6 4.946875f-7 4.3026725f-5 -6.489172f-5 3.9384457f-5 -0.0001731442 -0.00013290226 1.6823615f-5 0.00018841554 1.0568395f-5 0.00010316793 -5.307746f-5 -2.7167538f-5 9.810383f-6 -8.895507f-5 7.119211f-5 -1.8949928f-5 0.0001159132; -0.00013456795 6.257102f-5 -1.8263256f-5 1.2943337f-5 -4.7650963f-5 -4.3779286f-7 0.00016948406 -0.00017115746 -7.740652f-5 7.779009f-5 -4.458766f-5 0.00029745715 -9.039387f-6 -2.120973f-5 -5.6139896f-5 -9.500459f-5 0.0001661828 -4.2006348f-5 -3.2613723f-5 -4.670542f-5 -3.9607457f-6 3.7047328f-5 1.5051981f-5 0.00015239745 -3.5815046f-5 4.0842256f-6 1.4253509f-5 -0.00010424212 6.0472674f-5 -9.869268f-5 -0.000113540555 -0.0001069724; 6.421568f-5 3.5515907f-6 9.1632384f-5 0.00015352182 -4.2550975f-5 -4.6556226f-5 -6.0556253f-5 1.709983f-6 3.0830917f-5 0.00011723425 0.000108781904 0.00013236875 -9.780927f-5 9.008057f-5 1.453051f-5 -0.00015229145 0.00011206262 2.3733488f-5 -4.4362783f-5 3.815401f-5 5.08667f-5 -1.8604422f-5 -9.251175f-5 0.00026782195 -0.00030297594 0.00014862379 0.0002071232 -3.909452f-5 -0.00014309469 7.676019f-5 2.5200404f-6 0.0001646384; 2.9700266f-5 2.8444434f-5 0.00019739816 0.0002317098 2.9815817f-5 -9.9889396f-5 -8.687724f-6 -9.4445444f-5 4.7030385f-6 9.706607f-5 9.309902f-5 -0.000105647196 1.6553382f-5 6.362703f-5 9.076525f-5 -7.556437f-5 -3.2020547f-5 -5.8463494f-5 0.0001652476 1.501295f-5 -0.00013082243 -4.982664f-5 0.0001346891 0.00017365858 -0.0001754459 9.060477f-5 0.00011840552 -4.314156f-5 5.8626236f-5 6.387139f-5 4.7680137f-6 -8.233125f-5; -0.00014077549 -0.00012616758 -0.00012311428 -7.562119f-5 6.8019246f-5 -3.0012747f-5 0.00015323244 -8.77207f-5 4.868926f-5 6.656286f-5 -5.5218567f-5 0.00014223045 9.1740745f-5 3.4047964f-5 6.5659486f-5 -4.380196f-5 0.00012576475 3.9408613f-5 -7.0845417f-6 -1.3209486f-5 6.748422f-5 -8.1334394f-5 4.871015f-6 -9.2910836f-5 2.8922546f-5 -3.2402735f-5 4.6354013f-5 0.00034995133 -3.563906f-5 -0.000117792915 -9.733355f-6 -1.7405887f-7; -4.2183015f-5 -0.00014544866 -2.2169665f-5 -4.9087546f-5 -0.00020466198 0.00013750528 4.4607663f-5 -0.00011775543 -8.9586705f-5 -8.772235f-5 3.2311327f-5 3.2844397f-5 0.00015565187 -5.134928f-6 8.359333f-5 -3.5979512f-5 4.231141f-5 0.0001033584 7.5263546f-5 -2.6378191f-5 -3.4901892f-5 -0.00012262928 -4.2274663f-5 -5.807761f-5 -6.534819f-5 5.065244f-6 9.570354f-5 -2.3585188f-5 9.967649f-5 -0.000115291215 3.2885728f-5 6.30274f-5; -1.784474f-5 0.00016716533 0.00013522562 2.7284397f-5 0.00018954468 0.00012215787 4.0621955f-5 -7.0255715f-5 -0.00020354782 -3.973168f-5 3.2283388f-5 -0.00025144868 -8.867015f-6 3.2458214f-5 -1.6341302f-5 0.00010525662 6.6974004f-5 -8.4292646f-5 -6.842958f-5 0.00012808724 2.7158732f-5 -2.0473515f-5 -0.00016923048 8.047684f-6 -1.2472717f-5 -9.998895f-5 3.2780317f-6 -2.011538f-5 -0.00021364958 -5.7479378f-5 -0.00013140052 0.00010124617; 3.598549f-5 0.00018573985 -4.5854056f-5 -0.000103030456 -0.00012523838 6.392053f-5 0.0002904314 5.6947425f-5 3.568999f-5 -6.474586f-5 0.00017313732 0.00014978308 -0.00016178607 4.7431506f-5 2.0081565f-5 -0.00030161836 3.5265228f-5 0.00021190854 2.165826f-5 6.955934f-5 0.00017211995 -0.00017726164 6.906744f-5 -3.633269f-5 -1.364417f-5 0.00014624215 -5.6651115f-5 5.1554245f-5 0.00010683004 3.4827222f-5 -0.00012373741 -7.4654f-6; 6.6592176f-5 0.00017464855 -2.0248022f-5 0.00012050241 -6.172815f-5 0.00010145147 8.757429f-5 -8.845553f-5 9.751475f-5 -5.9425336f-5 0.00024525844 -4.9402344f-5 0.00013712006 0.00016522835 4.9981798f-5 5.924653f-5 0.00010154913 2.8310385f-5 6.242242f-5 0.00019429241 9.834293f-5 -5.1121795f-5 0.000102888815 -5.5732336f-5 0.00025003296 -2.7793081f-5 -5.877569f-5 -3.9588846f-5 4.977986f-5 -2.2239496f-5 -0.00010202718 -9.674124f-5; 1.9970608f-5 0.00015335804 -7.2259245f-5 -1.5337622f-5 0.0001311377 6.003373f-6 6.6739114f-5 -5.0642342f-5 -9.337545f-6 -3.666317f-5 2.5593095f-5 -5.9685513f-5 -6.83737f-5 -4.701907f-5 -4.956655f-5 6.8791174f-5 -4.2020292f-6 -7.973303f-5 -5.5327982f-5 -0.00019852685 -0.00012468142 -4.506905f-5 -7.2899595f-5 0.00013579603 -6.292252f-5 -4.724099f-5 2.5051915f-5 -2.2560087f-5 8.246151f-5 -1.7605547f-5 4.1634277f-5 7.043f-5; -2.7860484f-5 4.0157007f-5 2.5572192f-5 4.2198394f-6 0.00013586643 3.0988183f-5 -0.00015247102 -5.737023f-5 4.50524f-5 -0.00011112012 -3.887098f-5 4.1266867f-5 1.7571909f-5 9.6882184f-5 9.5647745f-5 -5.412185f-5 0.00011897714 0.00012867979 3.7931422f-5 -3.816678f-5 3.91085f-5 -4.7974703f-5 -9.194247f-5 5.5232296f-5 0.00010835611 0.00018883761 -0.0001661525 3.5989015f-5 -6.961442f-5 -5.7026755f-5 -1.4721738f-5 -7.297958f-5; -5.5844528f-5 -8.401536f-7 -6.082532f-5 -1.566273f-5 2.2536638f-5 5.4783395f-5 -0.000103604034 0.000106000385 -0.00014363602 2.1625465f-6 0.00014277475 -0.00011138319 5.2935266f-5 0.00015257286 -3.40413f-5 -0.00010913862 -9.203976f-5 -2.6056807f-5 -9.182499f-6 -0.00026183247 4.7883113f-6 0.00010597877 -6.0906055f-5 3.1774336f-5 2.7010085f-5 0.00010255491 2.8069583f-6 -2.5494011f-5 -3.6253412f-5 7.794546f-5 -7.1460825f-5 2.0900628f-5; -0.00012873723 1.7922592f-5 3.8715796f-5 -7.603868f-5 0.00014745434 -1.8547784f-5 -6.422266f-5 4.2134274f-5 -5.5788518f-5 4.942262f-5 3.1514795f-5 -2.5933035f-5 -9.211242f-5 2.0263304f-5 3.5156307f-5 -0.00012823306 -4.1771964f-5 -3.8861184f-5 5.4423625f-5 -7.327029f-5 1.3729199f-5 -2.639187f-6 2.0826052f-5 -4.763766f-5 0.00014795834 9.7051954f-5 -0.00012869094 -2.3304232f-5 0.00011829985 -5.937176f-5 5.9864036f-5 -0.00013736583; 2.180481f-5 -8.123618f-5 -0.00014011333 1.4233662f-5 2.433651f-5 0.000102360056 -0.00010620481 9.358219f-5 -7.3564996f-5 -9.410647f-5 6.489953f-5 -4.4232598f-5 4.6864956f-5 2.0784997f-5 -2.0957914f-6 5.4281416f-5 -6.5949585f-5 -8.6320295f-5 -0.000114325485 4.1727348f-5 -0.00013972555 1.2164693f-6 0.00011663185 8.037948f-5 -3.7258364f-5 -0.00010105982 -5.63612f-5 6.0645187f-5 -2.815638f-5 -4.1242984f-5 -0.00011109857 0.00018393296; 5.3116906f-5 -1.3695551f-5 -9.531769f-6 -0.00021912767 -2.24407f-5 0.00016122968 -0.00012455962 -8.092892f-5 -5.004883f-5 0.0002666798 -8.926439f-5 9.6578005f-5 0.00013982608 -8.886394f-5 -0.00015066164 -4.6925394f-5 1.5905927f-5 -0.00017775345 2.4965235f-5 8.275607f-5 9.140425f-5 7.377923f-5 0.00012107462 -0.00017933655 -3.5049223f-5 -0.00011901651 -1.0169942f-5 -2.2998729f-5 0.0001417768 2.6261136f-5 1.5357653f-5 0.00022100903; 9.4026334f-5 -3.4737764f-6 -2.0320434f-5 -0.00020660806 2.9263533f-6 -3.5373218f-5 5.9627942f-5 -1.4201123f-5 -6.455315f-5 3.4917106f-5 -8.920421f-5 -7.3197094f-5 3.6076985f-5 0.00010649656 6.750297f-6 0.00014223455 -7.908192f-5 -5.0413682f-6 -0.00010380356 -1.8049608f-5 -0.00018501068 0.00010492831 1.848246f-5 -0.00023566502 2.1200774f-5 -6.463226f-5 -8.8342415f-5 -1.13131355f-5 7.631597f-5 6.2858046f-5 5.972136f-6 1.476415f-5; -0.00014214734 -7.699401f-5 3.3692886f-5 0.00017862913 -0.00017032777 4.437086f-6 0.00012322691 -0.00014017796 6.930749f-5 -0.00010637234 6.565595f-6 0.00013693226 2.9269724f-5 0.00020584834 5.576879f-6 0.00011752201 3.105668f-5 0.00017149371 -7.497793f-5 4.561877f-5 0.00014838777 -1.0161487f-5 0.00013273327 -0.00018838371 7.884449f-5 0.00011210641 -4.4961307f-6 -0.00022302338 -1.041289f-5 2.9903872f-6 -1.5910617f-5 9.470503f-5; 3.1147316f-5 2.4947294f-5 -0.00013268163 -2.2986149f-5 5.4021435f-5 4.7259005f-7 0.00014617301 -3.4689336f-5 3.488622f-5 1.2864741f-5 -0.00013664899 6.045228f-5 -0.00011268217 -5.363126f-5 -5.7416295f-5 -4.672077f-5 -5.683104f-5 0.00010045006 0.00021820313 -8.110475f-6 -1.9840036f-5 -0.00017428449 -2.1205327f-5 1.622919f-5 5.264644f-5 -7.806471f-5 -8.211358f-5 -9.6781114f-5 -0.00012555937 -9.733848f-5 0.00011540612 0.00019545856; 6.7294523f-6 -5.927878f-6 3.4103192f-5 2.0403872f-5 -5.5492703f-5 -0.000113607755 0.00025140075 5.3416636f-5 -6.31842f-5 -9.133063f-5 7.2696675f-5 0.000112977614 0.00017463426 0.00026482547 -6.633254f-5 -7.67012f-6 -5.4247972f-5 0.00013494371 2.9024415f-5 -1.8850804f-5 2.8594688f-5 9.7001306f-5 6.367495f-5 -6.6978886f-5 0.00020421755 1.0977296f-5 -5.5108398f-5 -0.00012717482 -2.6717778f-5 3.3656423f-5 9.958974f-6 -2.541052f-5; 0.00013180247 -0.000117841904 -0.0001037626 1.4175618f-5 9.9137585f-5 0.0001278101 0.00014203222 8.6084794f-5 4.2703203f-5 3.862967f-5 -5.5256834f-5 0.00018781707 -0.00018881807 -0.00017852102 6.624268f-5 -7.949009f-5 8.735013f-5 -1.9856843f-5 -1.1712102f-5 -0.000119610464 0.00010579567 -3.8690963f-5 5.9955033f-5 -0.00021388044 -0.00019713247 -4.1131076f-5 -0.00016749231 3.810381f-5 -3.3941953f-5 -0.00015804195 3.36636f-5 5.70613f-5; 0.00013124249 1.5553143f-5 -0.00013162228 -0.00014369623 0.00016641208 9.3571776f-5 -0.00013470046 9.4162664f-5 -4.4893422f-5 -5.4486474f-5 -6.2578714f-5 0.00016160615 4.8792164f-5 -4.602861f-5 -0.000102899976 5.0947776f-5 9.656752f-7 -6.9735834f-5 -7.460267f-5 -1.2348849f-5 -7.3945506f-5 1.0299032f-5 -3.0231668f-6 -7.335786f-5 -2.1581973f-5 0.000101336984 4.718243f-5 9.6694705f-5 -8.909754f-5 -1.662915f-5 7.927897f-5 0.00011558249; -0.000108068794 1.7804667f-5 0.00015914912 -9.40703f-5 0.00010490888 8.515163f-5 -5.4336324f-5 0.00013162906 -1.2885078f-5 -8.06865f-5 8.86226f-6 -4.2136802f-5 8.72406f-5 1.0896734f-5 -0.00013646045 -3.4768f-5 -2.322952f-5 -6.72176f-5 4.6964874f-6 7.466218f-5 -1.6269203f-5 -9.223468f-5 5.1435312f-5 -3.284774f-5 -8.796306f-5 6.616878f-5 6.7991554f-5 -6.417984f-6 -0.00014266343 0.00018078339 -6.0926166f-5 0.00022324671; -7.21427f-5 5.6235753f-5 -4.6437555f-5 -2.0618174f-5 -7.3526164f-5 9.0357265f-5 1.45437525f-5 -0.00014352406 -2.0464759f-5 -5.9956088f-5 -3.3575194f-5 -7.0983406f-5 -0.000110765744 -1.2466829f-5 -0.00011097892 5.852028f-5 2.5351857f-5 -0.000111092835 9.4148025f-5 -2.096683f-5 9.8670826f-5 6.681318f-5 2.8959062f-6 -2.5940853f-5 0.00020203266 6.565604f-5 2.1751293f-5 0.0001094773 -0.00011084944 0.00010418308 -0.00019687781 -2.6789226f-5; -8.6187414f-5 -7.478814f-5 -4.2830114f-5 0.00011412184 -0.000102177575 -8.8856f-5 7.834966f-5 2.1123937f-5 0.00011583531 4.1557825f-5 4.7376498f-6 -9.3757575f-5 4.4060358f-5 -7.087315f-5 3.933878f-8 0.000106495696 0.00013665667 -9.5831136f-5 0.00014605466 -1.922161f-5 -4.470068f-5 2.3831266f-5 0.00011071551 -5.5295416f-8 4.8978312f-5 -4.7679085f-5 2.9585004f-5 2.3562083f-5 9.6548116f-5 0.00012634521 -0.00012654194 0.00015285025; -6.271246f-5 4.2192503f-5 -2.3929791f-5 -3.5316873f-5 5.5183296f-5 3.2361924f-5 0.00010549702 0.00017395306 -7.074031f-5 -5.176886f-5 2.4076153f-5 5.0187635f-5 5.5780296f-5 -1.0270028f-5 3.158136f-5 -0.00016007594 2.3709885f-5 -0.000110878114 -9.1927155f-5 8.0392754f-5 5.5490997f-5 7.1797476f-5 4.5129007f-5 1.3230255f-5 -0.00010379153 0.0002044453 -5.265454f-5 1.7560209f-5 5.7919417f-5 5.0075047f-5 0.00013856626 8.389799f-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[-0.00010211615 4.5702032f-5 0.00016922949 -5.2617543f-5 1.6188984f-5 4.0480095f-6 7.027658f-5 3.3580774f-5 -6.6078637f-6 -6.354011f-5 -0.0003016353 0.00017922268 4.848966f-6 0.00012365358 1.04195005f-5 3.857842f-5 -0.00012581384 -1.1232468f-5 -1.6825843f-5 0.00015860924 -4.750097f-5 -3.8824677f-5 -8.6670785f-5 -5.516875f-6 -5.198854f-5 2.0511432f-5 -7.153182f-5 -0.00011863276 -9.214741f-5 0.00012905186 -1.6364242f-5 -7.5833614f-6; 0.00010482918 0.00013501798 -0.00011757469 5.1049294f-5 6.456481f-5 -0.00014273108 0.00012486709 -1.8624412f-5 9.9893405f-5 0.00015555946 -3.0443587f-5 -1.2448053f-5 -1.38608175f-5 -0.00013083847 1.1939891f-5 1.5137468f-5 -0.00015853927 -8.817892f-5 -5.0861363f-5 -3.3352708f-6 -1.7160803f-5 6.4941785f-5 -3.4432273f-5 -0.00016817883 -7.184156f-5 -8.3322906f-5 -3.6261423f-5 1.4379693f-6 4.8730795f-5 -0.00011180734 1.49375555f-5 0.0001340473], 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, @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; 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

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.17596729573929765, [-0.024249511737021664, -0.02346557233375969, -0.02268163293049792, -0.021357031502375664, -0.019464059580601663, -0.016962495203849615, -0.013799527296815524, -0.00990679139505417, -0.00520263093368503, 0.0004082625335704292, 0.007025663440072411, 0.014732898602534912, 0.023549292804195095, 0.03330185627272482, 0.04334708017329918, 0.05191292703117135, 0.054706617512887724, 0.04264402570772163, 0.0021823998585432636, -0.06605161777364818, -0.11023285217597345, -0.07648279930248536, -0.007143836123777531, 0.038549688410399095, 0.05410690446725175, 0.0528747554382779, 0.044837856300484744, 0.03486989173693428, 0.0250152882321383, 0.0160386051430862, 0.008158699275717607, 0.0013777088961848278, -0.0043826211387860085, -0.00922147267663768, -0.013235631502155048, -0.01650817260446587, -0.01910982798987551, -0.02109587471158973, -0.02250826886259425, -0.02337650685011654, -0.023718578112946632, -0.023541300352401557, -0.022841559271287924, -0.0216040088206955, -0.01980309963682163, -0.01740076011668611, -0.014345585333126482, -0.010572226124343803, -0.006000537522803671, -0.0005369999427192718, 0.005917980062900897, 0.013454452188797188, 0.022105019073778254, 0.03174169621106318, 0.04182329075164026, 0.05081986084141631, 0.05496026322933571, 0.04604134292845008, 0.010728008781769586, -0.05508894455012823, -0.10809157578836075, -0.0860452504844271, -0.017173681350306783, 0.03370804472436024, 0.05311516265317079, 0.05368112306184197, 0.04628565370142437, 0.036441132108472406, 0.02650206382156838, 0.017370786234859676, 0.00931709862300528, 0.002372990669217658, -0.0035418242243770376, -0.00851527822245443, -0.012654248055944173, -0.016036245827040066, -0.01874112355265182, -0.02081982553526652, -0.022321695511926434, -0.023273987289905113, -0.023699073502974878, -0.02360437806000454, -0.02298752962928748, -0.021836719618247378, -0.020127942631233364, -0.017821910164340454, -0.014875087191249902, -0.011217137044454373, -0.0067779033535089435, -0.0014574986679524798, 0.004835606252273685, 0.012203014266536107, 0.020683516842671806, 0.03019343163719816, 0.04027481319330133, 0.04961656263328953, 0.05491247646109636, 0.04879851351346202, 0.018442821778558768, -0.04390846287341647, -0.10389366849573446, -0.09441874559909318, -0.027770756891389088, 0.028071816787458414, 0.05167971413398861, 0.05430550195359678, 0.0476781832867683, 0.03801025543034434, 0.02801019757058053, 0.018725786367479583, 0.010505451578871675, 0.0033892533275977346, -0.002674614955516442, -0.007791958139298379, -0.012051553027275475, -0.015550284289552291, -0.018354681794185996, -0.02053116088501017, -0.022120099179203774, -0.02315898850533838, -0.023666305373015374, -0.023653818949032884, -0.023121105801802506, -0.022055381993782267, -0.020437255742432804, -0.018228004567088108, -0.015386251791242323, -0.011843824037987977, -0.007533117656215219, -0.0023553017823767716, 0.0037796895787796125, 0.010977227664842116, 0.019287615715407876, 0.02865791835100307, 0.03871118562881354, 0.048321685170250935, 0.05460004498324608, 0.05097683854547746, 0.02530415174717753, -0.03278193415236342, -0.09782131673178374, -0.10128653438439512, -0.03876476354139301, 0.021605223616113636, 0.04974954114506092, 0.05471643367441676, 0.049002221310868226, 0.03957408948836692, 0.02953133688380591, 0.020111089763991154, 0.01171647584208507, 0.0044345092683757605, -0.0017905508158968207, -0.007042199230417546, -0.011433324517070326, -0.015044783937794702, -0.017956403442235022, -0.0202246088344575, -0.021906230269877362, -0.02303006765290954, -0.023620764104290266, -0.02369046189952832, -0.023240639469637153, -0.02226111829792256, -0.020729902691240722, -0.018620579923401343, -0.015879401478521307, -0.0124506894011241, -0.008268265704195395, -0.003231322049547835, 0.0027502823147817217, 0.009779369477051998, 0.017915047027417107, 0.027138990946473696, 0.03713857508653474, 0.046951056082231886, 0.0540619679653111, 0.05262887610394695, 0.031328303161251826, -0.021963827002839865, -0.09010528327054197, -0.10638203258582012, -0.04994084045309375, 0.014297715531384364, 0.047261076411863404, 0.05488372238540476, 0.05024206268326892, 0.041122266106784794, 0.031073516019412155, 0.021515965031225756, 0.012955487442968894, 0.005501306604754775, -0.0008749750136038499, -0.0062770959202858275, -0.010795136503722803, -0.014524418042805631, -0.017536773388471837, -0.01990709317909398, -0.021677824068008788, -0.022888183079409274, -0.023561447993879513, -0.02371423765832153, -0.02334698943713218, -0.022451255054907002, -0.0210125485931189, -0.018994092060087543, -0.016356565020376284, -0.013039967840465903, -0.0089818096499875, -0.004083338195153797, 0.001745038020527868, 0.008606535159259581, 0.016570003368195797, 0.025638853697055944, 0.035561911837727146, 0.04552004398156009, 0.05332751765689544, 0.05381258510457067, 0.036536585298594736, -0.011638581758143522, -0.08104730190499097, -0.10949942750894243, -0.061026762315831, 0.006143444167745352, 0.044164994677517916, 0.054766482163907206, 0.05138014632412671, 0.042651961131422556, 0.032624030270304166, 0.022946684042904085, 0.014220607889827161, 0.006596048675867091, 5.969495688381685e-5, -0.005488523828641274, -0.010137491656702345, -0.013985034032142215, -0.017104397393043796, -0.01957307098650171, -0.021435435204472556, -0.02273261772147718, -0.023489508216812, -0.023724722444562734, -0.023440106118067516, -0.02262897901489553, -0.02127829373701074, -0.019353541442006548, -0.01681697879391493, -0.013610967453410176, -0.009675144806447978, -0.005739322159485783])

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 = [4.561391324384266e-5; 0.00019234156934529715; 2.7182373742074915e-5; -0.0002046194131249506; 9.707956633066732e-5; 1.5926081687220352e-5; 6.33812160230167e-5; -5.768342816734539e-5; 1.641378912607507e-5; -2.217966357420638e-5; 9.268546273217645e-5; -5.223045081942289e-5; 0.0001204457003039917; -0.00014333908620721932; 0.00010610595927569126; -3.5050270525993545e-5; -2.3483948098132994e-5; -6.586288509409532e-5; -7.428431854346062e-5; -4.335849007464685e-5; 6.996507727296759e-5; 4.39171162725249e-6; 5.392593266144086e-6; 4.896315294898198e-5; -4.08189662266264e-5; 6.404497253232309e-5; 2.491280065439525e-5; 7.115352218536593e-6; 5.304032583803837e-6; -1.5418672774091344e-5; 9.54595379880778e-5; 9.762855188457295e-5;;], bias = [2.974158761221016e-17; 1.6645605984361132e-16; 3.73347041555877e-17; -4.353344808965943e-16; -2.1793564409504928e-17; -9.843219390702251e-18; 1.1360305347124183e-16; -2.2086896977869573e-17; 1.7791473931397966e-17; -2.962495624019279e-19; 1.7814250856275522e-16; -6.385865288882224e-17; 1.435690599779515e-16; -3.9113269249410313e-16; 5.276415197448219e-17; 2.4941979783392954e-17; -1.49439805669238e-17; -1.728575745174089e-16; -5.1513229484087866e-17; -1.1927733169024466e-16; 1.2949428395296997e-16; -3.286701803577857e-18; 8.323729396827602e-18; 1.9843451016435512e-17; -2.9470659613855625e-17; 1.2011879442220345e-16; 2.794086905440294e-17; -2.948061015572763e-18; 5.106943346962861e-18; -1.3809286218008013e-17; -1.703035830590902e-17; -2.2790082713960694e-17;;]), layer_3 = (weight = [-0.00010387433245818529 -0.0001955732844773691 0.00018312943906017063 -3.349003997319052e-5 2.3268475738188896e-5 8.464527093997598e-5 6.153394827811751e-5 -2.7465292410164237e-5 4.679814413498527e-5 0.00018489230172668043 8.390681944980461e-5 -4.8738093360265414e-5 -2.582947887752436e-5 -0.00020184299170544065 -0.00011229967309667316 0.00015591005252168038 9.164631377166536e-5 -0.0001734909859489949 8.094925428630299e-6 -0.00011349714292486614 -7.404930908273122e-5 -6.14889253167773e-5 -0.00017751211668405332 -0.0002518536460822764 -4.859960642101679e-5 -3.175731796113063e-5 6.43163472293544e-5 3.167759462917576e-5 -6.94132507899182e-5 -7.609045530011506e-5 -1.0848052957260615e-5 0.00010177399566124452; 6.531850444058165e-5 0.00011945012296690924 -0.000159507421684469 0.00028340811617580966 -3.198938871595189e-5 -0.00016116085394846216 -0.0001617256428823024 5.421119853177422e-5 6.769795358318895e-5 -2.8501725736208404e-5 -3.896594795026997e-5 -5.895280706558112e-5 7.745378310698395e-5 9.95301153501484e-5 -8.094558221149218e-5 -0.00012799533656893158 -0.00016509525526875546 -0.0002574480684949231 5.3993145605645386e-6 2.9883104745165522e-5 -2.0736228960813243e-5 -0.00014916889709515457 -0.00013712368023211096 -5.853092157725939e-5 1.8887031857373094e-5 -3.446742708423728e-5 4.645637012923225e-5 -1.2844517967681913e-5 0.0001250660199196827 3.8148736295039405e-5 -4.821732805670377e-5 4.253471264870735e-5; -0.00010645570506216617 9.881942279313157e-6 -3.132104899027255e-5 -0.0002503514675107685 -0.00022683490311367806 2.4372499357062e-5 -3.248456920022433e-5 -0.00014600640639304577 -8.102492352822738e-5 -0.0001647309814497447 -4.142409448734985e-5 4.834760551698948e-5 -2.3084870516797933e-5 -6.659843125718208e-5 0.000149021465805228 -5.5109967032827245e-5 8.593532624313708e-5 -0.00020291303877360424 -1.3963357158843951e-5 -0.00020731668115235558 -1.2769259537337787e-6 0.0002139556913253294 1.3498023182412328e-5 0.00016698744952105371 -9.935606572884792e-5 4.4253464185257484e-5 -0.00012693009699337073 -6.118367815265408e-5 -4.085771575678122e-5 4.570591267206757e-5 3.297590450626379e-5 -7.183876510583219e-5; 0.00014823496327106056 -8.254271073430051e-5 -5.068448284921113e-5 -0.00022103823561384466 -4.311757424195028e-5 9.581774225760366e-5 -0.0002341185884147492 -0.00013024707572232763 6.348216425929333e-5 2.9860062349579552e-5 -1.9694023867959416e-5 -2.5282124843688985e-5 -5.350058039813882e-5 5.805994966020329e-5 7.567086300649047e-5 -5.3966369014707013e-5 4.746273567387023e-5 -0.00012089298819821146 -5.691341121309975e-6 9.079143070219226e-5 7.09547400760856e-5 -6.943647113549578e-6 0.00014009593156468235 -7.353057139018401e-5 9.583367660477896e-5 5.618330996621558e-6 -4.6421721199768317e-5 0.00012088040399201324 7.572600021328887e-5 7.314663685927282e-5 -4.528778139549045e-5 -0.00011222751913057371; 0.00015719250649656727 -3.9703699538021574e-5 -8.195114360912114e-5 -0.00014138853594684465 -9.621410873275806e-5 -5.234586870553301e-5 9.679529896216973e-7 9.443969310471384e-5 -7.783245683437543e-5 -3.379855780419298e-5 -8.566564527219854e-6 3.515089401137625e-5 8.19044768761403e-5 -9.961490587354535e-5 -7.827419477305046e-5 2.8161176207713968e-5 -2.3362064264592353e-5 -4.413811557949019e-5 3.6164001987553264e-5 -7.381855112523211e-5 -0.00012590841225371206 3.6276488292269895e-5 9.059615394310701e-5 0.00011239805161697407 0.00017515806820686541 6.635205926816487e-5 5.870647573122416e-5 8.368177504256042e-5 9.784093407891366e-5 7.634210914350716e-5 2.4669478709464075e-5 -8.310200365283154e-5; -4.946652295653001e-6 -5.6743185466673976e-5 0.00016316387316907415 -3.4083750704359175e-5 3.2053507113559925e-5 -0.00012594622987092787 -0.00014219091119015194 3.274759709016859e-5 -7.190956259149435e-5 7.222982987846735e-6 -2.85691678233339e-5 2.57088939018688e-5 -9.616749791933626e-6 0.0001764692131313064 0.00015900962151350418 0.0001674412485793848 -1.878619352067683e-5 -2.5819256022838815e-5 5.136664165762562e-5 -6.0676720947990536e-5 5.684405894713822e-5 8.122160196068402e-5 -3.406169363882688e-5 -0.00013555180359402184 -0.00012539919427366388 7.4820057146674635e-6 0.00012265228704779076 -4.20025246218156e-5 0.000169692822766925 1.1935735330885913e-5 0.0001171278852619799 5.792799386071969e-6; -0.00011873317175438491 8.412941047330875e-5 -2.1733406307859424e-5 -8.03717766272303e-5 -5.794402654525419e-5 3.8038795567700623e-6 0.00024215934631986752 -4.6637792317297364e-5 -6.231686615744537e-5 -1.0643474476339001e-5 -0.00010207247300665452 0.00010325236603723997 4.138727768198247e-5 0.00010632912840568343 -5.92988789746118e-5 6.990791103210404e-5 0.00018916674257828006 7.047361673660791e-5 1.7747514951525074e-5 2.9403579040562148e-5 -0.0001129798756179935 5.061429424434028e-5 4.406565226720705e-5 -5.783158436926015e-6 -5.1123329236971885e-5 -0.0002062110497233064 -7.722645291021015e-5 3.4395526248161345e-5 -0.00021107106926744742 -9.702446454564042e-5 -5.757926325412257e-6 -5.1197613126230984e-5; -0.00010412046309586073 3.8271329810142745e-5 -0.0001414043980397051 9.839965435885953e-5 -5.196385233222458e-5 -9.20240962507905e-5 -0.00021619238019713474 0.0001101114941199558 5.3179974326374606e-5 9.67040808671135e-6 0.0001652858592357766 -0.0001770742854781204 -7.462463314669784e-5 -0.00012672345097899652 -4.220091216952576e-6 4.928504715086778e-7 4.302488754244971e-5 -6.48935564307242e-5 3.938261955834742e-5 -0.00017314603961796898 -0.00013290410048586536 1.6821777659677427e-5 0.0001884136983214175 1.0566557772726305e-5 0.00010316609053614135 -5.307929665224895e-5 -2.71693751195214e-5 9.808545779995166e-6 -8.895690904257401e-5 7.119027226206347e-5 -1.8951764889894666e-5 0.00011591136366103832; -0.0001345686756495503 6.257029634881432e-5 -1.826398293915004e-5 1.2942610182028891e-5 -4.765169033916725e-5 -4.385201053519982e-7 0.0001694833339469395 -0.0001711581859227384 -7.740724650929959e-5 7.778936595868981e-5 -4.458838849231678e-5 0.00029745642713251315 -9.040114137946461e-6 -2.1210456553574445e-5 -5.614062284721075e-5 -9.500531415242023e-5 0.00016618207190631615 -4.20070751816611e-5 -3.261445054804126e-5 -4.6706146905396644e-5 -3.961472985881177e-6 3.7046600836130765e-5 1.5051253639196023e-5 0.00015239672246834967 -3.581577367285509e-5 4.083498323691792e-6 1.425278186516336e-5 -0.00010424284811151495 6.047194655266309e-5 -9.869340975555868e-5 -0.00011354128183583795 -0.00010697312937836218; 6.421892409817457e-5 3.554833728382036e-6 9.163562733120569e-5 0.00015352506104060072 -4.254773214670541e-5 -4.655298310425728e-5 -6.055300967163472e-5 1.7132261011337034e-6 3.083416028785853e-5 0.00011723749573248748 0.000108785146739319 0.00013237198961106207 -9.780602907158684e-5 9.008381109123152e-5 1.453375331875263e-5 -0.00015228821051168015 0.00011206586147564767 2.3736731026660586e-5 -4.4359540180081905e-5 3.8157251555464915e-5 5.086994294215766e-5 -1.8601178663439234e-5 -9.250850619609108e-5 0.0002678251892295583 -0.0003029726960728951 0.00014862703189392508 0.00020712644415580967 -3.909127756814115e-5 -0.00014309144374309677 7.676343172462062e-5 2.523283473418871e-6 0.00016464164743338065; 2.970300945081671e-5 2.8447177347587605e-5 0.00019740089944308672 0.00023171254021788057 2.9818560752670302e-5 -9.98866523594959e-5 -8.684981120394277e-6 -9.44427008370329e-5 4.705781767475123e-6 9.706881271187282e-5 9.310176416975708e-5 -0.00010564445231365175 1.6556125293483093e-5 6.36297702577349e-5 9.076799348904772e-5 -7.556162720801307e-5 -3.2017803940912094e-5 -5.8460751214523834e-5 0.0001652503490469007 1.501569319872359e-5 -0.00013081968766426996 -4.982389848817963e-5 0.00013469183638436827 0.00017366132694506914 -0.00017544315026396924 9.06075149679158e-5 0.00011840826507169808 -4.313881775386527e-5 5.8628979123673546e-5 6.387413329422565e-5 4.770756978456469e-6 -8.232850789581246e-5; -0.00014077445498813338 -0.00012616654398997946 -0.0001231132465848361 -7.562015357807136e-5 6.802028096600731e-5 -3.0011711854313994e-5 0.0001532334734717711 -8.771966363681913e-5 4.869029530411256e-5 6.65638961861249e-5 -5.5217531692764296e-5 0.00014223148341144953 9.174177998530802e-5 3.404899857442749e-5 6.566052056470189e-5 -4.3800924825774394e-5 0.00012576578760374 3.9409647658173926e-5 -7.0835068557089446e-6 -1.3208450721578317e-5 6.748525797476507e-5 -8.133335918799363e-5 4.87204974589897e-6 -9.290980065387144e-5 2.8923581200359352e-5 -3.240170022152822e-5 4.6355047604009824e-5 0.0003499523607813163 -3.563802434044279e-5 -0.00011779188039625392 -9.73232013464833e-6 -1.730240030924281e-7; -4.218411251041715e-5 -0.0001454497598502721 -2.21707616334381e-5 -4.908864340231871e-5 -0.00020466307452135136 0.00013750418455618692 4.460656581376141e-5 -0.00011775652657216978 -8.958780218801202e-5 -8.772344718874164e-5 3.231023018760232e-5 3.2843299584224565e-5 0.00015565077209193096 -5.136025004146557e-6 8.359223446622214e-5 -3.598060922069181e-5 4.231031360849075e-5 0.00010335730056557278 7.526244897907115e-5 -2.6379288294824404e-5 -3.490298896632664e-5 -0.00012263037771173503 -4.227576047252613e-5 -5.807870749105365e-5 -6.534928817256044e-5 5.064146796381569e-6 9.570244018267183e-5 -2.35862846353625e-5 9.967538952640509e-5 -0.00011529231252335685 3.288463066145422e-5 6.302630371227521e-5; -1.7845926089719664e-5 0.00016716414461910687 0.00013522443692265515 2.72832119262047e-5 0.000189543491393056 0.00012215668608034445 4.0620769828749834e-5 -7.025690069356398e-5 -0.00020354900993341162 -3.9732865156182506e-5 3.2282202260762365e-5 -0.00025144986943276086 -8.868200206118323e-6 3.245702897032053e-5 -1.634248769193981e-5 0.0001052554368979343 6.697281856388284e-5 -8.429383091795615e-5 -6.843076629964015e-5 0.00012808605311504117 2.7157547043246576e-5 -2.0474700471894468e-5 -0.00016923166480781104 8.046498583937395e-6 -1.2473901955743733e-5 -9.999013786371232e-5 3.276846433365368e-6 -2.011656510519852e-5 -0.00021365076031004603 -5.748056286763916e-5 -0.00013140170562480754 0.00010124498726866048; 3.598842222063213e-5 0.0001857427784642464 -4.585112338018133e-5 -0.00010302752280326298 -0.00012523544801580153 6.3923464674098e-5 0.0002904343352355637 5.695035788166561e-5 3.569292375404903e-5 -6.474292540870112e-5 0.00017314025235181394 0.00014978601456905067 -0.00016178313608462838 4.743443850827842e-5 2.0084497643335963e-5 -0.0003016154290008688 3.526816062449035e-5 0.0002119114730675345 2.166119220136679e-5 6.956227361742772e-5 0.00017212288430245665 -0.00017725870502827148 6.907037522696499e-5 -3.6329757952537983e-5 -1.36412375496569e-5 0.00014624508248461558 -5.6648182456320916e-5 5.1557177991434716e-5 0.00010683297359437291 3.483015525207442e-5 -0.0001237344761676561 -7.462467058858809e-6; 6.659782930861954e-5 0.00017465420361022343 -2.0242369183077992e-5 0.00012050806227251801 -6.172249440218948e-5 0.00010145712413711605 8.757994483809357e-5 -8.844987341084021e-5 9.752040178634354e-5 -5.941968292111443e-5 0.0002452640980749996 -4.9396691147678984e-5 0.00013712571645505095 0.00016523400673525528 4.9987451083734304e-5 5.925218257240904e-5 0.00010155478204021305 2.8316038614742766e-5 6.242807258737366e-5 0.00019429806680770425 9.834858038581184e-5 -5.111614181307901e-5 0.00010289446773592553 -5.572668332254778e-5 0.0002500386105652699 -2.7787428089628182e-5 -5.877003813159588e-5 -3.958319238384085e-5 4.978551415613094e-5 -2.2233842437824753e-5 -0.00010202152797368924 -9.673558842273144e-5; 1.996941934900379e-5 0.00015335684727158847 -7.22604338314667e-5 -1.5338811159216064e-5 0.00013113650927531838 6.002184158354726e-6 6.673792520283818e-5 -5.064353101791745e-5 -9.338733639957416e-6 -3.666435734506737e-5 2.5591906500419535e-5 -5.968670232126443e-5 -6.837489056654705e-5 -4.702026055913612e-5 -4.956774066689048e-5 6.878998535663806e-5 -4.203218146209062e-6 -7.973421745722075e-5 -5.5329171323895776e-5 -0.00019852803527631793 -0.00012468260977205694 -4.507023843492439e-5 -7.290078358513364e-5 0.00013579483927438993 -6.292370674221773e-5 -4.72421772663751e-5 2.5050726230041695e-5 -2.2561275933036236e-5 8.246032373259716e-5 -1.7606736242520204e-5 4.163308838011989e-5 7.042880749774475e-5; -2.7859533549321287e-5 4.015795722634157e-5 2.5573141936643995e-5 4.220789758703302e-6 0.00013586738070271457 3.098913377298804e-5 -0.0001524700738463207 -5.7369280586403645e-5 4.505334883843378e-5 -0.000111119169634452 -3.887003102308567e-5 4.126781724890449e-5 1.757285910398554e-5 9.68831344999666e-5 9.56486955311471e-5 -5.4120900205720334e-5 0.00011897809319462447 0.0001286807427442044 3.793237275940964e-5 -3.8165831103373654e-5 3.9109451716380786e-5 -4.7973752845695e-5 -9.19415206971222e-5 5.523324663782257e-5 0.00010835705772306556 0.00018883856402508625 -0.0001661515517133865 3.59899649248349e-5 -6.96134694243425e-5 -5.702580445530993e-5 -1.4720787550592577e-5 -7.29786289247357e-5; -5.584572460270774e-5 -8.413500720402394e-7 -6.08265177672856e-5 -1.5663926717239824e-5 2.2535441928766666e-5 5.478219867449827e-5 -0.00010360523043550764 0.00010599918863385677 -0.00014363721453469603 2.1613499827671414e-6 0.00014277355188619311 -0.00011138438600233086 5.2934069060713174e-5 0.00015257166199967695 -3.404249542859312e-5 -0.0001091398185619744 -9.20409543707629e-5 -2.60580032429538e-5 -9.183695112598722e-6 -0.00026183366168175187 4.787114846974709e-6 0.00010597757176378376 -6.090725179297285e-5 3.1773139596904676e-5 2.7008888025872176e-5 0.00010255371711332143 2.8057617671456547e-6 -2.5495207921498387e-5 -3.625460846001649e-5 7.794426022678256e-5 -7.146202148055669e-5 2.0899431945281146e-5; -0.00012873821016163138 1.7921610368128405e-5 3.871481475308921e-5 -7.603966424700974e-5 0.00014745335737481445 -1.8548765194086397e-5 -6.422364004882495e-5 4.213329247087706e-5 -5.57894991395958e-5 4.9421637402897736e-5 3.15138140542496e-5 -2.593401678658109e-5 -9.211340421805267e-5 2.026232231239406e-5 3.515532533876384e-5 -0.00012823404450662963 -4.177294523450168e-5 -3.886216564913713e-5 5.4422643177728865e-5 -7.327127151418987e-5 1.372821769856899e-5 -2.640168315142387e-6 2.0825070340120123e-5 -4.763864036411541e-5 0.00014795736295874647 9.705097271600299e-5 -0.00012869192051929038 -2.330521330057747e-5 0.00011829887170116782 -5.937274091755603e-5 5.986305494681651e-5 -0.0001373668119520348; 2.1803298745770886e-5 -8.123769430779439e-5 -0.00014011484518519111 1.4232150983726171e-5 2.4334999120449492e-5 0.00010235854505922175 -0.0001062063215660146 9.358067699836249e-5 -7.356650667602665e-5 -9.410798293643159e-5 6.489801534812842e-5 -4.423410928679976e-5 4.686344481784155e-5 2.0783485974660827e-5 -2.097302517673648e-6 5.42799048115271e-5 -6.595109653531099e-5 -8.632180607050927e-5 -0.00011432699603133112 4.172583657677828e-5 -0.00013972706574818356 1.2149581560587126e-6 0.00011663033591944244 8.037797169422367e-5 -3.7259874720940635e-5 -0.00010106133186595796 -5.6362712262081766e-5 6.0643675619619775e-5 -2.815789185219924e-5 -4.1244494700687966e-5 -0.00011110007972560854 0.0001839314522793329; 5.3117254075995296e-5 -1.3695203513557696e-5 -9.531421263009977e-6 -0.00021912732473348663 -2.2440351593024357e-5 0.00016123002432513132 -0.00012455927345209717 -8.092857161310457e-5 -5.004848345189276e-5 0.0002666801433334414 -8.926403962315075e-5 9.657835334300115e-5 0.00013982642712297762 -8.886359274394031e-5 -0.00015066129117977854 -4.692504581536098e-5 1.590627515914925e-5 -0.0001777531010774613 2.4965582717494296e-5 8.275642039015019e-5 9.14045945065378e-5 7.377957671641453e-5 0.00012107496740290045 -0.0001793362039351551 -3.504887491354656e-5 -0.00011901616517926531 -1.0169593686360588e-5 -2.299838116224309e-5 0.0001417771550150876 2.6261483519273663e-5 1.5358000830099696e-5 0.0002210093793947526; 9.402439566209363e-5 -3.4757151069174082e-6 -2.0322373234119073e-5 -0.00020661000205553955 2.9244145708511335e-6 -3.5375156289699586e-5 5.9626003644629786e-5 -1.4203061868934917e-5 -6.455508598712005e-5 3.4915167202672965e-5 -8.920614679929005e-5 -7.319903277213261e-5 3.607504579585286e-5 0.00010649462303726628 6.7483582143628705e-6 0.00014223261343893924 -7.908385639431599e-5 -5.043306987240761e-6 -0.0001038055012713033 -1.8051546862731697e-5 -0.0001850126169983148 0.00010492637040904975 1.8480521544080905e-5 -0.0002356669607933477 2.11988356796455e-5 -6.463419747425664e-5 -8.834435327564075e-5 -1.1315074220800081e-5 7.631403088915415e-5 6.28561068197832e-5 5.970197276008699e-6 1.4762211617461124e-5; -0.0001421451440647474 -7.69918103570597e-5 3.3695084556635845e-5 0.0001786313234960147 -0.0001703255681887551 4.439284356446481e-6 0.00012322911176753588 -0.0001401757606173173 6.93096892691001e-5 -0.00010637014086667766 6.567793451641105e-6 0.0001369344547755782 2.9271922602589095e-5 0.0002058505354115378 5.579077668615693e-6 0.00011752420802523942 3.105887778375907e-5 0.00017149590827913357 -7.49757298136384e-5 4.562096843884841e-5 0.0001483899698259059 -1.0159288404868506e-5 0.0001327354705326044 -0.00018838151184453155 7.884668709167808e-5 0.00011210861090959403 -4.493932220687644e-6 -0.0002230211818924389 -1.0410691839827288e-5 2.9925856951349627e-6 -1.5908418146036585e-5 9.470722695269591e-5; 3.114618837656405e-5 2.494616663256966e-5 -0.0001326827548101088 -2.2987276423964463e-5 5.4020307544248546e-5 4.714625273965272e-7 0.00014617188597270311 -3.46904632525053e-5 3.488509199172144e-5 1.2863613193015445e-5 -0.00013665011621839175 6.045115176760568e-5 -0.00011268330101216735 -5.363238571504431e-5 -5.741742253945849e-5 -4.672189581384962e-5 -5.6832166336850375e-5 0.00010044893159004896 0.0002182020037079704 -8.111602909653992e-6 -1.9841163265541564e-5 -0.00017428561815819823 -2.1206454875088598e-5 1.6228063287539053e-5 5.2645311626232524e-5 -7.806583718690063e-5 -8.211471114833943e-5 -9.678224134319586e-5 -0.00012556049452456763 -9.733960880027268e-5 0.00011540498988409736 0.00019545743296561058; 6.732657958522741e-6 -5.9246724857126734e-6 3.410639757133678e-5 2.040707723716768e-5 -5.54894971071681e-5 -0.00011360454963211263 0.00025140395583660586 5.3419841343596034e-5 -6.318099193664562e-5 -9.132742278667484e-5 7.269988119444581e-5 0.00011298081945936574 0.0001746374661808842 0.00026482867971136634 -6.63293351761023e-5 -7.666914178812151e-6 -5.4244766296291223e-5 0.00013494691366871102 2.9027620904553665e-5 -1.8847598120160355e-5 2.8597893752913625e-5 9.700451183674589e-5 6.367815213819696e-5 -6.697568031883384e-5 0.00020422075314057559 1.0980502075864994e-5 -5.510519193991417e-5 -0.0001271716130656962 -2.671457267639495e-5 3.365962830797479e-5 9.9621799665568e-6 -2.5407314104790347e-5; 0.00013180092331051401 -0.00011784345430756503 -0.00010376414890604584 1.4174068138841197e-5 9.913603547058953e-5 0.00012780854715994976 0.00014203067233349803 8.608324399713071e-5 4.270165275695222e-5 3.8628120575160374e-5 -5.525838448054248e-5 0.0001878155232875755 -0.00018881962047859817 -0.00017852256884768087 6.624113278260498e-5 -7.949164161662105e-5 8.734857668198143e-5 -1.9858393232858072e-5 -1.1713651718741881e-5 -0.00011961201404890476 0.00010579411876417419 -3.8692513288404464e-5 5.9953483043065637e-5 -0.00021388199117470941 -0.00019713401731539962 -4.113262572761943e-5 -0.00016749386147100974 3.810226165253146e-5 -3.394350302958454e-5 -0.00015804349847574304 3.3662050310819696e-5 5.7059751026397136e-5; 0.00013124270662127078 1.555336347417201e-5 -0.00013162205687796674 -0.00014369601377358933 0.00016641229921413552 9.357199632257249e-5 -0.00013470023805813863 9.416288411637771e-5 -4.489320236563656e-5 -5.448625416567525e-5 -6.25784941886369e-5 0.00016160636913890028 4.879238439651128e-5 -4.6028389996645316e-5 -0.00010289975581064876 5.09479959790783e-5 9.658953021076207e-7 -6.973561358710755e-5 -7.460245070353163e-5 -1.2348628851321852e-5 -7.394528621181854e-5 1.029925169030465e-5 -3.0229466657314495e-6 -7.335764349510881e-5 -2.158175267060737e-5 0.00010133720387864855 4.718265153395035e-5 9.66949246420714e-5 -8.909732242420289e-5 -1.662893010608826e-5 7.927919249676127e-5 0.0001155827067040089; -0.00010806810931478175 1.7805351687112276e-5 0.00015914980141736116 -9.406961797447389e-5 0.00010490956739579613 8.515231567492269e-5 -5.433563890926133e-5 0.0001316297432064214 -1.2884392875430382e-5 -8.068581408530919e-5 8.862945120770733e-6 -4.2136117175114634e-5 8.72412867616972e-5 1.0897419285015559e-5 -0.00013645976904024125 -3.476731470379508e-5 -2.3228835546769938e-5 -6.721691787185677e-5 4.697172441073002e-6 7.466286589681627e-5 -1.6268518068837786e-5 -9.22339957318126e-5 5.143599739006286e-5 -3.2847054884196484e-5 -8.796237560397484e-5 6.616946595079673e-5 6.799223885230272e-5 -6.4172989217421704e-6 -0.00014266274118505853 0.00018078407557971335 -6.0925480996386015e-5 0.00022324739367525376; -7.214371921858943e-5 5.623473211759729e-5 -4.643857587493262e-5 -2.0619194539750903e-5 -7.352718435126309e-5 9.035624484400012e-5 1.45427319222606e-5 -0.00014352507619318705 -2.0465779154480164e-5 -5.9957108671229915e-5 -3.3576214402131504e-5 -7.098442634000616e-5 -0.00011076676505748898 -1.2467849184612766e-5 -0.000110979943339628 5.851925911715965e-5 2.535083681155219e-5 -0.00011109385573203795 9.414700420907618e-5 -2.0967851152678152e-5 9.866980494567092e-5 6.681216277516185e-5 2.894885572598743e-6 -2.5941874038695316e-5 0.00020203164408317563 6.56550163040763e-5 2.1750272797502036e-5 0.00010947628274894059 -0.0001108504603979297 0.00010418205588225953 -0.00019687882936690894 -2.6790247058527942e-5; -8.618538905391107e-5 -7.478611715596253e-5 -4.28280885982282e-5 0.0001141238631960709 -0.00010217555009983602 -8.885397747614867e-5 7.835168306958166e-5 2.112596174886264e-5 0.0001158373366623234 4.1559850243018027e-5 4.739674833055371e-6 -9.37555500142244e-5 4.406238312259788e-5 -7.087112809057276e-5 4.136385189868777e-8 0.00010649772101790133 0.00013665869718427302 -9.58291106227764e-5 0.0001460566878366298 -1.9219585200674798e-5 -4.4698653627324945e-5 2.383329100742038e-5 0.00011071753632752723 -5.327034368642096e-8 4.8980337479248826e-5 -4.7677059600732794e-5 2.9587029158235464e-5 2.3564107860536627e-5 9.655014090715475e-5 0.00012634723915728291 -0.00012653991635413944 0.0001528522721772069; -6.271002451115945e-5 4.2194939734901314e-5 -2.392735507613596e-5 -3.5314436107157874e-5 5.5185732773548614e-5 3.236436066059818e-5 0.00010549945625085505 0.00017395549263879835 -7.073787230923352e-5 -5.1766424010689995e-5 2.4078589260040206e-5 5.019007118959204e-5 5.578273237174592e-5 -1.026759139508975e-5 3.158379592773979e-5 -0.00016007350598387074 2.3712321191727436e-5 -0.0001108756779359644 -9.19247187340632e-5 8.039519059080921e-5 5.549343302105604e-5 7.179991268627787e-5 4.513144346843914e-5 1.3232691227914943e-5 -0.0001037890916706032 0.00020444773587938514 -5.26521045031918e-5 1.7562645436431937e-5 5.792185299639099e-5 5.007748302147216e-5 0.00013856869815922988 8.390042682718166e-5], bias = [-2.6887985284062243e-9; -2.5400901395784496e-9; -4.858119562128704e-9; -5.077167961859526e-10; 7.536911932090038e-10; 1.833343381842169e-9; -1.1268154339642759e-9; -1.837054989784954e-9; -7.27242057241593e-10; 3.2430482040132063e-9; 2.7432825986238184e-9; 1.0348655623921325e-9; -1.0970439830075968e-9; -1.185293584717905e-9; 2.9328447448272375e-9; 5.653154833158208e-9; -1.1889149720291435e-9; 9.503474815360532e-10; -1.1964969001209423e-9; -9.81355501454887e-10; -1.5111301328652096e-9; 3.478745286115998e-10; -1.9387470439948917e-9; 2.1984932729557978e-9; -1.1275179230870286e-9; 3.2057061951168225e-9; -1.5500231512323779e-9; 2.2008942467907355e-10; 6.85046534371975e-10; -1.0205911591665232e-9; 2.025072560991214e-9; 2.436419767091229e-9;;]), layer_4 = (weight = [-0.0007893282852901119 -0.0006415101279760752 -0.0005179823116540244 -0.0007398298466119793 -0.0006710233129970365 -0.0006831642241075428 -0.0006169356980489487 -0.0006536314584349429 -0.000693820161339734 -0.0007507521639045023 -0.0009888474161253485 -0.0005079896037074415 -0.0006823633162492375 -0.0005635587039401767 -0.0006767926104754778 -0.0006486331588551722 -0.0008130261127902464 -0.000698444757106122 -0.0007040381198509438 -0.0005286030486368119 -0.000734713225541152 -0.0007260369843723611 -0.0007738830056120352 -0.0006927290748807667 -0.0007392008200763788 -0.0006667006438875932 -0.0007587440696546584 -0.0008058450661140341 -0.0007793597116865891 -0.0005581604315318933 -0.0007035764558895867 -0.0006947955310358888; 0.0003460058881653231 0.00037619469688130375 0.0001236019030534122 0.00029222605983410504 0.0003057415742731584 9.844566347025145e-5 0.0003660438468343081 0.00022255232890851234 0.0003410701686315797 0.00039673614377831224 0.00021073311331912614 0.00022872870716288704 0.0002273159410757316 0.0001103382903995264 0.0002531165897115727 0.0002563139801402738 8.263748182901884e-5 0.00015299784453278988 0.0001903153937071314 0.0002378414901313768 0.00022401594689828752 0.00030611855252645625 0.00020674446447233293 7.299790268489578e-5 0.0001693351965889523 0.00015785378043152684 0.00020491532560186783 0.00024261473721069108 0.0002899075595743852 0.0001293694188743239 0.00025611429031496553 0.00037522401463935213], bias = [-0.0006872123100655278; 0.00024117676828951215;;]))

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(LuxCUDA) && CUDA.functional()
    println()
    CUDA.versioninfo()
end

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

Julia Version 1.10.3
Commit 0b4590a5507 (2024-04-30 10:59 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: 8 default, 0 interactive, 4 GC (on 2 virtual cores)
Environment:
  JULIA_CPU_THREADS = 2
  JULIA_AMDGPU_LOGGING_ENABLED = true
  JULIA_DEPOT_PATH = /root/.cache/julia-buildkite-plugin/depots/01872db4-8c79-43af-ab7d-12abac4f24f6
  LD_LIBRARY_PATH = /usr/local/nvidia/lib:/usr/local/nvidia/lib64
  JULIA_PKG_SERVER = 
  JULIA_NUM_THREADS = 8
  JULIA_CUDA_HARD_MEMORY_LIMIT = 100%
  JULIA_PKG_PRECOMPILE_AUTO = 0
  JULIA_DEBUG = Literate

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

CUDA libraries: 
- CUBLAS: 12.4.5
- CURAND: 10.3.5
- CUFFT: 11.2.1
- CUSOLVER: 11.6.1
- CUSPARSE: 12.3.1
- CUPTI: 22.0.0
- NVML: 12.0.0+550.54.15

Julia packages: 
- CUDA: 5.3.4
- CUDA_Driver_jll: 0.8.1+0
- CUDA_Runtime_jll: 0.12.1+0

Toolchain:
- Julia: 1.10.3
- LLVM: 15.0.7

Environment:
- JULIA_CUDA_HARD_MEMORY_LIMIT: 100%

1 device:
  0: NVIDIA A100-PCIE-40GB MIG 1g.5gb (sm_80, 4.735 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