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.00011973267; -6.019302f-5; 9.096157f-5; 4.2661748f-5; 9.349409f-7; 0.00016044428; -9.652911f-6; -0.00013854839; 9.681807f-6; -2.8657123f-5; 0.00017334803; 2.4411342f-5; 9.011604f-5; 5.220362f-5; 0.00024906726; -2.4143723f-5; 0.00012566449; -4.8479058f-5; 4.074634f-5; 1.0247664f-5; -0.0001859722; 3.7356087f-5; -7.656956f-5; -1.2043876f-5; -8.47999f-5; -2.4505294f-5; -7.9587626f-5; -0.00022258626; -0.0001703838; 2.2777765f-5; 4.271127f-5; 6.0953906f-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[-1.0062973f-5 -0.0001435936 7.171802f-5 3.4863042f-6 0.00017328089 6.824237f-5 -0.00012931744 7.675739f-6 0.00022924854 -0.00011960628 0.00011614925 -0.00013504615 4.9548744f-5 1.507356f-5 -0.00010520048 9.820607f-5 0.00011530446 -5.298501f-5 3.924124f-5 -8.876832f-5 6.5198125f-5 4.548969f-5 5.3590033f-5 -0.000111658905 8.5903535f-5 -3.4463697f-5 0.00013809335 1.6326525f-5 6.29639f-5 6.3808046f-5 5.5576777f-5 -0.0001310525; -7.0357426f-5 9.636652f-5 -2.7267588f-5 -5.7131452f-5 0.00015933493 9.020348f-5 -1.280962f-5 -1.5835829f-5 1.2373991f-5 -0.0003102026 -1.5455973f-5 -5.0626495f-5 -7.3990675f-5 0.00022303218 -5.265653f-5 5.5802695f-5 0.000119704986 2.3875706f-5 2.6456584f-5 -1.1514318f-5 3.4455326f-5 8.091209f-5 -1.6922031f-5 -9.741731f-5 -0.00013727894 -9.3055656f-5 0.00014561473 -0.00012389639 5.852755f-5 -0.000109424866 -4.5462173f-5 0.00021897774; -5.4261145f-5 0.000121771845 0.00020005245 2.6634425f-5 -2.619291f-5 -0.00020373413 -1.2680036f-6 3.037937f-5 -5.047399f-5 -1.8179611f-6 9.946187f-5 5.8790112f-5 -0.00014307268 3.3849017f-5 -0.00016031059 2.021655f-6 -3.0627314f-5 4.2110256f-5 -2.9767696f-5 0.00021924138 2.025092f-5 0.00011184174 -0.00019683011 2.243284f-5 6.390769f-5 1.2823031f-6 0.00012127408 8.413545f-5 5.0939067f-5 -3.1260803f-5 7.489051f-5 -0.000117571406; 8.286847f-5 5.6977842f-5 -0.00011781059 -9.9578145f-5 -1.48217305f-5 7.6619406f-5 7.078823f-5 -1.2595202f-5 -1.8840798f-5 0.00013955703 7.530427f-5 -3.0709045f-5 0.0002204093 0.00010372281 0.00010360437 -4.6901005f-6 0.000103564205 9.927554f-5 0.00015315072 0.00016233182 -6.378829f-5 0.00027685237 -0.00022597678 0.00013754105 9.243972f-5 -0.00012381934 -0.00010605406 -0.0001433673 -2.6276277f-5 -0.00013443631 -0.00011772916 5.814016f-5; -0.00016425819 3.8754493f-5 -0.00016896137 6.288182f-5 -0.00023414425 -9.0244575f-5 -5.2501702f-5 1.4775569f-5 4.904329f-5 -1.5052479f-5 -2.7051257f-5 -3.614626f-5 0.00012460808 -0.00018959433 5.600226f-5 -5.7855483f-5 3.128997f-5 9.41794f-5 8.772447f-6 -0.00014428626 -3.1018924f-5 1.3908119f-6 -4.760862f-7 -0.00012216585 1.242881f-5 4.9815384f-7 -0.00010849579 -0.00010241591 0.00014788071 -8.60535f-5 -8.829287f-5 -5.122713f-5; 0.00011356168 2.7072892f-5 -9.253048f-5 -9.7144235f-5 -5.3558386f-5 -7.102739f-5 -0.00018623754 0.00013162836 -1.5184637f-5 -9.944188f-5 0.00014779605 -2.3400253f-5 -6.563826f-5 -7.489051f-5 -0.00012160472 -7.8382036f-5 9.002087f-5 -1.4130239f-5 3.6039422f-5 -0.00011180349 -6.0440257f-6 -0.0001540761 -5.739455f-6 5.1871833f-5 6.2550454f-5 -2.9444225f-5 4.1011826f-6 7.9901416f-5 5.603711f-5 -1.4256181f-5 3.4426583f-5 4.9856906f-5; -3.286154f-5 7.116674f-5 -4.1574076f-6 -0.00015561515 4.6678346f-5 0.0001754534 3.814063f-7 -0.00022762116 5.648911f-5 4.255361f-5 -8.870969f-5 1.8850651f-5 3.7530056f-6 -4.5058187f-5 -7.086632f-5 -0.0001447943 -9.620589f-5 0.000102021324 0.00017278506 -1.789651f-5 -0.00019770837 -0.00012196603 -0.00020703816 0.00014147113 0.000104440216 0.00011035788 1.3612266f-6 8.4339816f-5 0.00018007657 8.51515f-6 -2.6068446f-5 -8.29233f-5; -0.00017980425 0.00024446344 7.396711f-5 -0.00021361737 3.724509f-5 1.4738594f-5 0.000136625 2.9866022f-5 3.47192f-5 9.010391f-5 -0.00019570945 4.4907054f-5 9.772386f-6 -7.9038015f-5 -0.00033689776 -4.9428443f-5 4.496193f-5 -5.6542092f-5 -6.9646762f-6 -0.00019470662 5.1095558f-5 -0.00014013774 -9.7984426f-5 -0.00011226721 0.00031234138 -9.5950694f-5 -0.00014004567 4.6424328f-5 -4.7736758f-5 7.010866f-5 1.6203121f-5 4.6829777f-5; -8.708018f-5 -0.00016045521 -7.2275757f-6 6.393291f-5 0.00011965771 0.00011464074 -4.5697686f-7 8.0861464f-5 0.0001022533 5.646368f-5 0.000116467454 0.00017300091 2.5471847f-5 0.00020660755 -4.3006254f-5 0.00012930758 -1.0878323f-5 0.00011079276 0.00020138327 -0.00015660196 9.4770485f-6 3.581487f-5 -8.650456f-7 -0.00014728917 -4.715184f-5 -0.00010080215 2.8409338f-5 -6.403153f-5 0.00015295464 -0.00017042554 1.3564003f-6 -0.00011035926; -9.573557f-5 -6.596612f-5 4.345059f-5 0.00010860645 0.00019122523 5.4775028f-5 8.066031f-5 -6.302869f-5 2.8398154f-5 5.9679318f-5 4.9511655f-5 0.00012019725 7.560769f-5 3.4006407f-5 -7.405292f-6 -3.3514407f-5 -0.00013196601 8.2634935f-5 -1.6556698f-5 0.00011844044 -2.1374497f-5 -5.9172984f-5 0.0002714602 9.1502596f-5 -2.4370893f-5 -6.784604f-5 2.7011944f-5 -0.00012285201 -0.00018385517 9.90055f-5 1.537271f-5 -9.662465f-5; -0.00013863033 0.00012131699 4.4665827f-5 4.008011f-5 -0.00029301443 0.00010328879 -0.0001005913 6.607562f-5 8.064893f-5 -1.5013011f-5 -8.330572f-5 -4.4715736f-5 6.1474464f-5 -0.000120070705 -4.839504f-5 -1.6226278f-6 -7.70251f-5 -5.7063782f-5 6.19982f-5 6.2265244f-5 -5.091939f-5 -0.000121028796 -8.224875f-5 -6.8355905f-5 -5.8808615f-5 6.492918f-5 8.943161f-5 6.348348f-5 -0.00019190085 -0.00010643892 3.577843f-6 -4.6126414f-5; 0.00012505539 -3.8034304f-5 -0.00018813847 2.5143841f-5 -1.5800473f-5 5.7095695f-6 0.00012570035 8.26293f-5 -2.4939525f-5 5.9741706f-5 1.7747216f-5 -1.7496443f-5 8.440886f-5 9.904175f-5 -2.7824995f-5 6.900068f-5 8.423556f-5 -0.00012994438 0.00016383445 -8.6308886f-5 0.0001719667 8.853987f-5 -0.00011011876 -1.8455012f-5 8.94854f-5 -0.00011606892 7.799474f-6 -5.584804f-5 -8.175717f-5 -1.3835867f-5 -9.989124f-5 3.248445f-5; 0.0001143422 0.00012731152 0.00018125362 -6.4191668f-6 0.00016003738 -2.7765425f-5 -0.0001555723 -4.8895436f-5 0.0001035261 -0.00012203125 -2.6826072f-5 1.8942063f-5 4.4934506f-5 -0.000106232124 2.2164802f-6 -0.00018457786 -2.8980878f-5 -2.7131715f-5 3.6822305f-5 -0.00012809213 0.00017954524 6.7248395f-5 6.213277f-5 7.3543066f-5 -8.91384f-5 8.050668f-5 9.043197f-5 0.00011347448 -2.4456815f-5 -0.00014946179 0.00019021088 -0.00018729462; 9.272404f-5 1.8320108f-5 7.896093f-5 0.00017118001 -0.0001250922 2.291453f-5 6.254659f-5 0.000116823874 -3.7232323f-5 -7.725545f-5 -1.972598f-5 -1.9868286f-5 2.5172509f-5 -5.481218f-5 0.00016046854 4.3901264f-6 -6.932835f-5 7.142715f-5 2.5496958f-5 -3.063386f-5 0.00016012261 -7.974759f-5 3.226735f-5 0.000104737555 -4.9273913f-5 -3.3438133f-5 0.00020110856 -0.00013596313 0.00015645802 -0.00012764387 0.00010554395 7.904172f-6; 7.832622f-5 0.00011860157 5.266906f-5 -7.044353f-5 3.933173f-5 1.6427433f-5 -0.00011026127 -4.351902f-5 0.00022462536 7.477401f-5 -9.0549554f-5 0.00016367993 -3.6606558f-5 5.467878f-5 -2.690087f-5 4.4372475f-5 -1.575863f-5 3.407016f-5 6.962714f-6 -0.00011319172 -1.3054113f-5 -4.122088f-5 3.6056736f-5 9.1263464f-5 7.2483024f-5 -0.00016864711 7.967744f-5 2.6180001f-5 -0.00015273914 -1.9276631f-5 6.645485f-5 -9.559493f-5; 4.3957913f-5 0.00014572433 -1.4306217f-5 -7.624112f-5 -2.9588346f-5 0.00014455158 4.0845338f-5 6.834219f-5 -9.741042f-5 -0.00014014456 -1.6901933f-5 1.7843891f-5 -0.00010838938 4.6790734f-5 0.00019295757 -6.0569873f-5 -0.00013637965 -3.0474872f-5 0.0001037594 0.00019073563 6.419302f-7 -0.00013129206 -0.000105330306 -1.7867922f-5 0.000115244446 -1.6840724f-5 -0.00016233283 0.000109826695 -5.968982f-7 -0.00013197932 -0.00013100801 8.5669435f-6; 9.623581f-6 -5.6304347f-5 -4.3870004f-5 0.0002447548 0.00017727955 -4.4379878f-5 -8.2319966f-5 0.000109945286 -6.691387f-6 4.3810637f-6 4.286829f-5 -0.00012954461 -0.000112196445 -6.5902816f-5 6.441595f-5 0.00012012445 1.5023467f-5 -0.00011764802 0.00012402602 -0.00021996978 6.3675914f-5 -9.014398f-5 0.00010708526 -0.00020586811 0.00012778126 7.886614f-5 -0.00013825842 -6.0239498f-5 2.3244516f-5 -3.6498983f-5 8.609585f-5 2.1549386f-5; 7.344093f-5 9.49406f-5 -7.0200425f-5 -3.4793175f-5 -1.9189849f-5 -3.8277554f-5 0.00011976373 -2.1197337f-5 -8.067376f-5 -3.2448952f-5 -7.225808f-6 -2.4172796f-5 -3.2152413f-5 3.2865668f-5 -2.211482f-5 -4.6703088f-5 -0.00023568759 -3.9639992f-5 -3.4805715f-5 0.00017792893 7.392999f-5 -0.000108427135 1.9616336f-5 1.4932766f-5 0.00014690605 -6.202372f-5 -4.172182f-5 0.00020422452 3.4323224f-5 1.6616243f-5 -0.00014446728 6.3973195f-5; -2.8892682f-5 -9.694493f-5 2.6234446f-5 -3.5557365f-5 1.9830559f-5 -0.00016544068 -7.2985036f-5 -1.26402165f-5 0.00016205726 -0.00019165275 3.345725f-5 2.944733f-5 -7.4804484f-5 -6.582739f-6 2.4713332f-5 0.0001649914 7.26928f-6 -0.00022154477 9.388105f-5 5.7786277f-5 0.00014243826 -3.4910674f-5 9.038743f-5 -5.1882926f-6 -9.168539f-5 0.00022661644 2.6507185f-5 0.00025639276 -4.074251f-5 -6.5140506f-5 8.830005f-5 0.00012997884; 3.4252414f-5 1.7978846f-5 0.000169445 0.00018952793 -8.011082f-5 6.781402f-5 0.000120731675 1.5522899f-5 2.1491327f-5 -0.00013182944 -1.638258f-5 -9.910352f-5 1.763488f-5 -5.268952f-5 0.0001341999 0.00012931575 2.0700243f-5 -7.521594f-5 -8.710873f-5 -0.00025061367 -4.814495f-5 -5.369676f-5 8.906928f-5 -0.00012929876 3.5379262f-6 0.00019203326 -0.0001413457 0.00011003606 2.02384f-5 2.5250322f-6 -4.823328f-6 0.00017709933; 0.00010651335 3.9165167f-5 7.6804965f-5 0.00018133184 8.555229f-5 6.652526f-5 -0.00012029219 -7.9235f-5 -0.00020872283 -1.1949973f-5 -0.00013770128 -5.876042f-5 -4.9586248f-5 -9.1145936f-5 -0.00010791469 2.2212389f-5 2.3586554f-5 0.00013769098 0.000111135545 -9.670577f-5 -9.732586f-5 8.423925f-6 0.00017610678 9.734474f-5 1.8881644f-6 1.411364f-5 -5.9431986f-5 -1.2221734f-5 -8.04621f-5 -6.966804f-5 0.00023346557 -6.003869f-5; -3.876047f-5 0.00014561284 -4.148563f-5 8.296341f-5 -6.987597f-5 0.00013355879 -0.00011680688 4.320796f-6 3.716612f-5 -4.7806792f-5 -9.4781724f-5 5.3489697f-5 -1.4912613f-5 8.296775f-5 2.3379346f-6 -1.0566036f-5 -0.000108239285 -7.448106f-5 -3.0216563f-5 7.662209f-5 0.00014964533 -0.00011516985 -3.5049095f-5 -3.0671647f-5 -3.2881521f-6 -3.6293895f-5 7.330256f-7 -3.5557747f-5 4.4640456f-5 8.841141f-5 6.586005f-5 -0.00017228896; -1.24846265f-5 0.0001208953 0.00014398887 4.5776145f-5 -0.000115045936 -5.9687086f-6 6.0125035f-6 6.8744295f-5 8.819092f-6 6.5838617f-6 -8.074785f-6 8.18883f-6 8.1181395f-5 0.00016108702 2.6273632f-5 -3.1020445f-5 9.508686f-6 4.0793813f-5 -0.0001408116 3.1298189f-6 2.3511593f-5 -8.186044f-5 -0.00010960179 -0.00016496978 4.405605f-6 -6.639267f-5 -0.000101537495 7.49069f-6 0.00010797909 3.2122553f-5 0.00019854389 6.8400455f-5; -0.00017579607 4.488868f-5 -7.397126f-8 2.2734788f-5 4.2988877f-6 -0.00012200552 4.1339466f-5 2.328481f-5 4.1465668f-5 8.132256f-6 8.115436f-5 1.3009662f-5 -1.9977986f-5 -0.0001287793 -9.6333824f-5 2.874979f-5 -5.3080348f-5 -5.785651f-5 -6.20475f-5 5.9302507f-5 2.9213074f-5 9.509549f-5 -3.0428158f-5 0.00017109046 0.00020654775 5.3617234f-5 0.00010196593 0.0002069366 2.0894107f-5 1.6600687f-5 -7.48392f-5 0.00012764451; -1.6030723f-5 -1.9729712f-5 2.562178f-5 0.00022121618 0.00018367641 4.139218f-5 7.6101554f-5 0.00019930796 7.185698f-5 7.8943165f-5 -7.644431f-5 4.3318218f-5 -1.3253722f-5 -2.1225085f-5 -0.00016582398 -0.00018792605 -1.3799033f-5 0.0002778923 -0.000105594445 0.000120767974 3.4850607f-5 -4.4116554f-5 3.988118f-5 -4.367971f-5 0.00016983895 -4.00821f-5 -8.422913f-5 1.4080487f-5 9.899427f-6 4.9401973f-5 -3.4165874f-5 -7.69171f-5; 4.1773757f-5 0.00014913613 0.000126567 0.00017125023 -6.546571f-5 0.00033155672 0.00011466273 7.809752f-5 6.979409f-5 3.1921016f-5 6.510741f-5 -8.501913f-5 6.257237f-5 0.00011156401 -9.156247f-5 9.0500354f-5 9.2775146f-5 -6.766979f-6 -1.7866389f-5 -2.7544174f-5 -4.399843f-5 7.5294745f-5 0.00011452333 -8.9779605f-5 -0.0001744599 -0.00013224669 9.717873f-5 0.00013802775 2.7223847f-5 6.938874f-5 -0.00015714747 -3.729033f-5; -0.00010186369 1.8786994f-5 1.1669726f-5 3.7432466f-5 0.00014849601 -8.1198996f-5 1.2223435f-5 3.6235957f-5 -6.6038585f-5 -0.00020149031 -8.3387975f-5 -1.1616641f-5 3.861739f-5 5.2266863f-5 0.00017672828 -3.6429665f-5 0.0001213096 -0.0001294468 -9.729489f-5 0.00011058528 -4.197558f-6 0.00019134328 8.317239f-5 6.726744f-5 2.998288f-5 -0.0001831322 7.862825f-5 0.00014545645 -3.136986f-5 0.00011239314 0.0001090515 7.624094f-5; 0.00010192903 -6.524324f-5 7.143097f-5 -8.02076f-5 -1.0315073f-5 -7.236077f-5 -0.00015450438 -7.5828546f-5 -0.00017704595 5.233977f-5 -3.8811017f-5 -7.597749f-5 0.00018876392 -0.00013852713 -1.1569201f-5 -2.537716f-5 -0.00022445363 -4.2580206f-5 -0.00011457844 -0.00017382255 -1.9114601f-5 3.0382655f-5 9.44476f-5 -4.1854648f-5 0.00017476828 9.583724f-5 8.147809f-5 9.057342f-5 -0.00016546999 8.475024f-5 4.4872453f-5 8.55354f-5; 4.9008333f-5 -0.00015482772 -2.9896846f-5 -7.672894f-5 -7.277536f-5 -1.7218908f-5 -2.0766205f-5 0.00015757453 0.00016821099 0.00012837019 -0.00016522144 2.8738874f-5 -3.516303f-5 -0.000107145315 8.35585f-5 -1.5251731f-5 1.805957f-5 0.00013188834 -0.000100679004 0.00015572732 -8.925097f-5 0.00017002202 0.00010146622 -0.00024448716 0.00016385132 1.8185434f-5 -3.342352f-5 7.46513f-5 -0.00013326635 6.0554674f-5 7.850088f-5 -5.4996348f-5; -0.000119494725 -0.00025112068 0.00011172785 9.278077f-6 1.7621252f-5 -1.15393295f-5 3.6610567f-5 -2.8657137f-6 -4.93761f-5 3.655143f-5 -0.00012682205 -8.361162f-6 0.00018711382 -2.9714225f-5 3.2761087f-5 0.00015201842 -0.00010104727 -0.00012077274 -9.988359f-6 -0.00014403982 -7.473175f-5 -2.1082627f-5 -0.00021641343 2.9726623f-5 0.00014701443 -0.00015106272 2.60452f-5 -0.00020045234 7.8442135f-5 -6.923768f-5 5.397001f-5 4.615466f-6; -1.7326232f-5 -4.5282093f-5 -1.1445137f-5 6.4935935f-5 9.0735455f-5 -1.1834731f-5 -2.7228747f-5 -1.252255f-6 -8.46312f-5 -4.548237f-6 -2.6095266f-5 3.601256f-5 7.6172396f-6 -4.7420755f-5 -0.00021223628 1.2169157f-5 0.0001208721 -4.9286413f-5 -4.5233795f-5 -0.00019525501 2.342413f-5 -0.0001615138 -4.1510673f-5 4.5838387f-5 -0.00022700398 -6.700035f-5 -7.5483396f-5 -4.520957f-6 6.766729f-5 4.7984697f-5 0.00010202208 -7.772921f-5; 0.0001495528 0.00020418811 3.9540133f-5 -6.3646266f-6 6.240287f-5 0.00015583959 -9.5992895f-5 5.6563535f-5 7.875602f-5 -8.022817f-6 -6.209155f-5 3.4326014f-5 6.877887f-5 -8.3817555f-5 -9.155829f-5 -2.6342583f-5 -1.791582f-5 -7.3701696f-5 3.1804808f-5 -0.00022330634 -8.917969f-6 -0.00024460425 0.00014764017 7.990559f-5 1.008664f-5 0.00013233775 -0.00021250913 -6.5767315f-5 -0.00010873815 -2.6074382f-5 -8.926506f-5 -8.4303254f-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.00011602808 -0.000120417775 7.054768f-5 -0.0001364321 1.0047987f-5 -8.540841f-5 -2.2970164f-5 -3.7272024f-5 7.4761874f-6 1.3028075f-5 1.06746675f-5 -4.7827394f-5 3.6491794f-5 -8.778504f-5 -3.844625f-5 9.919472f-5 -9.7661235f-5 2.755268f-5 1.7365865f-5 -4.661768f-5 -0.00011277104 -2.358428f-5 0.00014879108 -0.00028839387 -8.063303f-5 -1.4004848f-5 5.652013f-5 4.9286264f-5 0.000105819985 4.3599415f-5 -0.00013917121 -0.000105863226; -9.1772155f-5 -0.00017995205 4.675251f-7 -5.2920554f-5 0.00014703989 0.00018397685 -6.5224354f-5 3.4467295f-5 5.9178594f-5 -0.00023862447 -6.239535f-5 -7.785182f-5 1.7593879f-5 -1.3778444f-5 6.928703f-5 -8.063013f-5 4.069519f-5 1.616499f-5 0.00013895672 8.854975f-5 0.00012469526 -4.8969636f-5 6.860495f-5 -1.4344437f-5 9.547716f-5 5.0288338f-5 -2.761656f-5 9.424831f-5 -0.0001829136 -8.509033f-5 -5.9762977f-5 2.7484823f-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.17215171721421207, [-0.024249801535431388, -0.02346585418572156, -0.022681906836011435, -0.021357317180048352, -0.019464416988440935, -0.016963046554685536, -0.013800484708468138, -0.00990849440296376, -0.00520559267126083, 0.0004032969641172105, 0.007017656891419924, 0.014720493495480241, 0.02353090231748062, 0.033276169354516875, 0.043314705833896336, 0.05188106697558526, 0.054700333055434175, 0.0427238357065267, 0.002423105357175596, -0.06575883215773011, -0.11027317642844461, -0.07678253350835433, -0.0072353830951188206, 0.03869619541269424, 0.054302679477350145, 0.053010890392771556, 0.044889237694299845, 0.034846679863281424, 0.024936638661368058, 0.015923083057834027, 0.008021668513541608, 0.0012311811425605268, -0.004529475034590188, -0.009361716233634384, -0.013364035775378412, -0.016620815901878253, -0.019203788197126054, -0.021169013129862177, -0.022559076807699988, -0.023404004269245684, -0.023722253289358303, -0.023521083092865745, -0.022797843963709284, -0.021537669920586207, -0.019715602687456514, -0.017294276511174298, -0.014223184402618666, -0.010438139233305582, -0.005860520949963716, -0.00039881752053247047, 0.006043962037684947, 0.013554593203308008, 0.022161982187542492, 0.03173508798113943, 0.041733947747754976, 0.050644402047554236, 0.05475043416740677, 0.045981405841294294, 0.011152289366379492, -0.05420462270255706, -0.10786003980798775, -0.08685899277206673, -0.01774186110866609, 0.0338250226807328, 0.053499201452751846, 0.05400660035342886, 0.04644819440068384, 0.03644347334895938, 0.026379742162639767, 0.01716279566038105, 0.009056918301880643, 0.0020871637830344454, -0.003833000649596514, -0.008796465842887504, -0.012913937856553289, -0.01626587145143468, -0.018934307254695595, -0.02097197454005451, -0.0224295200316549, -0.023335406713336643, -0.023712900244009478, -0.02357041802509943, -0.022906394593631228, -0.021710058883234875, -0.019958459751178312, -0.01761360308066557, -0.014633651841338422, -0.010950330269378201, -0.006496278517170611, -0.00117520713687785, 0.00509967877361141, 0.012423985228668584, 0.02082949459404453, 0.03022588244507418, 0.04015436535793723, 0.049323788875630895, 0.05450738977182879, 0.04854762140980133, 0.018910188751749105, -0.042502139586782596, -0.10315104867281276, -0.09558611000472392, -0.028939522297792744, 0.02803995234799488, 0.052223130646726956, 0.054837231555989026, 0.04797944252126462, 0.03806268827600256, 0.027862360747976808, 0.018437229466984544, 0.010129000748480839, 0.002967092500939512, -0.0031100197038603306, -0.0082161156668996, -0.012446121061225964, -0.015901568755078654, -0.018652617240000398, -0.020768350144362, -0.02229141599601115, -0.023260958394085672, -0.023697093505961463, -0.02361282939105417, -0.023009218865298837, -0.021874751751256145, -0.020191745724326714, -0.017923246391015904, -0.015030312218619087, -0.011447676185503672, -0.0071116709655967255, -0.0019285321933857303, 0.00418515298414508, 0.011326237661996849, 0.01953487303756195, 0.02874787822057803, 0.03858428047507269, 0.047938884837366906, 0.05401544627420644, 0.05050337220632601, 0.0256911577239574, -0.030974995521538738, -0.09638900641676833, -0.102573566871294, -0.040619015798842875, 0.021284462092519427, 0.05040555779361737, 0.05546518394540094, 0.04946935076603893, 0.03970253564114501, 0.029377636383265602, 0.01975531095702029, 0.01123161885375543, 0.003879775858755852, -0.0023695397955581397, -0.00761096830444328, -0.011966053312152423, -0.015522255130645694, -0.018364387021327692, -0.020552878493211888, -0.022147379737285706, -0.02317926869799286, -0.02367518135491305, -0.023649303371970062, -0.02310461742752934, -0.02203298018800196, -0.020414371431591624, -0.018224930107454487, -0.015413676223328199, -0.011928852014086397, -0.007709193266232499, -0.0026602412566282306, 0.0032996827586000317, 0.010262587199066909, 0.018274577343534007, 0.027303342429091142, 0.03702853083260528, 0.0465057213069404, 0.05331917019944852, 0.051918238954459836, 0.03153306060791138, -0.019910122722075096, -0.08788362950254597, -0.10750094646131855, -0.05250722342833319, 0.013532854861997918, 0.0479619737149416, 0.055852146956217485, 0.050901076175073945, 0.04135375329743446, 0.03093541376814808, 0.02110756232480277, 0.01237125722871237, 0.004818557596569256, -0.0015962956792808102, -0.006991680144950654, -0.011469023119692376, -0.015132311698807245, -0.018060132121416765, -0.020332234338905318, -0.021995136036621243, -0.023091139245199024, -0.023646409526520064, -0.023679578179550936, -0.023193543086071973, -0.022182036958480172, -0.02063327847752358, -0.018513107066246644, -0.015785984296997794, -0.012396365770967076, -0.008287656846264849, -0.003368634740061504, 0.0024402092108903213, 0.00922922779343203, 0.017051490948011293, 0.02589298414209541, 0.03549058692942278, 0.04503937218323815, 0.05245167905094164, 0.05286401786344559, 0.0364818425938837, -0.009505948227704146, -0.07802438838691865, -0.11013670289983736, -0.06425595915170446, 0.004772010596095392, 0.04482024983889405, 0.055946769355377436, 0.05225496995452235, 0.04301460045508618, 0.03252435619257789, 0.022501895668115328, 0.013547176658980737, 0.0057907331212782714, -0.0008020753538423008, -0.006349674377069929, -0.01095520031828327, -0.014727439711552081, -0.017748123886680903, -0.02010087532816729, -0.021835155048116412, -0.0229959341098181, -0.023611652701661365, -0.02370339152416624, -0.023275956637742234, -0.022325259983575547, -0.020841440939984492, -0.018793020521969617, -0.016146637502150406, -0.012849766519993163, -0.008848805202157648, -0.004847843884322281])

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.00011973267100973714; -6.019301872584702e-5; 9.096157009471292e-5; 4.266174801152644e-5; 9.349408855980382e-7; 0.00016044428048176378; -9.652911103319343e-6; -0.0001385483919874571; 9.681806659489127e-6; -2.8657123039017627e-5; 0.00017334803123936343; 2.4411341655644033e-5; 9.011603833644094e-5; 5.220362072567235e-5; 0.0002490672632118383; -2.4143722839640925e-5; 0.00012566449004216483; -4.847905802303613e-5; 4.074634125572134e-5; 1.0247664249599429e-5; -0.0001859721960498307; 3.735608697743877e-5; -7.656955858682703e-5; -1.2043876267849151e-5; -8.479990356138734e-5; -2.4505294277351263e-5; -7.95876258051809e-5; -0.00022258625540396467; -0.00017038380610754774; 2.2777765479984303e-5; 4.27112718170935e-5; 6.095390563132088e-5;;], bias = [-1.164825726562541e-16; 7.32997919745033e-18; 9.983766760539752e-17; 7.770460759065174e-17; 1.7788999500159933e-18; 1.4588208534933488e-16; -1.2507218408227132e-17; -1.251678622168931e-16; 1.2733169472004439e-17; 7.135391304264406e-18; 2.484754549107743e-16; 6.723620038833896e-18; -2.2857880932878122e-17; 1.1538087757854018e-16; 2.463886399288879e-16; -1.3610381677146602e-17; 9.87350838026452e-17; -4.603549365892236e-17; 5.6417112338565525e-18; 2.1427353925903686e-17; -1.7516937468067382e-16; 9.686404129065536e-17; -1.1068338528151916e-16; -4.679508920209203e-18; 2.5985377908932743e-17; 2.0807714940635198e-17; -1.2025065573329291e-16; -2.2061612577673057e-16; -1.8567883300781978e-17; 2.482936980343263e-18; 3.4047360819198803e-18; 2.173035420130597e-18;;]), layer_3 = (weight = [-1.0061051374231235e-5 -0.0001435916774186004 7.171994459566595e-5 3.4882259809725193e-6 0.0001732828104890276 6.824429240293579e-5 -0.00012931552098380893 7.677660529568793e-6 0.00022925046598756195 -0.00011960435856381964 0.00011615116911923022 -0.00013504423165942132 4.9550665631964246e-5 1.5075481338853033e-5 -0.00010519855911629895 9.820798825457177e-5 0.00011530637950852913 -5.2983089493510066e-5 3.924316208794128e-5 -8.876639819719384e-5 6.520004628907118e-5 4.54916126058287e-5 5.359195440759037e-5 -0.00011165698285695655 8.590545715049844e-5 -3.446177558872437e-5 0.0001380952753846818 1.632844671472133e-5 6.296582526044495e-5 6.380996731094503e-5 5.557869828989134e-5 -0.0001310505749403198; -7.035733811956914e-5 9.636661037635938e-5 -2.7267500046794955e-5 -5.7131364302980244e-5 0.0001593350176682785 9.020356868692558e-5 -1.2809531983142202e-5 -1.5835740998076307e-5 1.2374078904182701e-5 -0.0003102025069545866 -1.5455885078912935e-5 -5.062640717536254e-5 -7.399058755421144e-5 0.00022303227218632536 -5.265644300546716e-5 5.5802782879406165e-5 0.00011970507388302038 2.3875793556877868e-5 2.6456672129067134e-5 -1.1514230538371901e-5 3.4455414277695585e-5 8.091217808957168e-5 -1.692194324444572e-5 -9.741722111524122e-5 -0.0001372788532188075 -9.305556828788522e-5 0.0001456148171547888 -0.000123896300794363 5.852763628187465e-5 -0.00010942477762167383 -4.5462085255925074e-5 0.00021897783211731302; -5.4259894708870824e-5 0.00012177309585758837 0.0002000537050168512 2.663567542704249e-5 -2.6191660394305017e-5 -0.00020373288239562497 -1.2667532285363746e-6 3.0380619906002633e-5 -5.047274058078849e-5 -1.8167106912177227e-6 9.94631180635115e-5 5.8791362649300245e-5 -0.00014307143388059756 3.3850267614707524e-5 -0.00016030934162183824 2.0229054786807387e-6 -3.062606370286228e-5 4.211150630517336e-5 -2.976644569052881e-5 0.00021924263168848616 2.025217040289681e-5 0.00011184299259703118 -0.00019682885718275532 2.2434090029220385e-5 6.390894017904274e-5 1.2835534750604354e-6 0.00012127533304528576 8.413669897497975e-5 5.094031698325898e-5 -3.125955222849255e-5 7.489175809105353e-5 -0.00011757015559333028; 8.287123114749126e-5 5.6980601816385334e-5 -0.00011780782892421003 -9.957538508861785e-5 -1.4818970835409752e-5 7.662216581117524e-5 7.079098782450589e-5 -1.2592442312985069e-5 -1.8838037926986118e-5 0.00013955978852645128 7.530702919645007e-5 -3.0706285310959205e-5 0.00022041206307656967 0.00010372556630935908 0.0001036071282712616 -4.687340829338635e-6 0.00010356696498527327 9.927829728875734e-5 0.00015315347468894462 0.00016233458359674242 -6.378553166108842e-5 0.00027685512964562096 -0.00022597401925414693 0.0001375438098464596 9.244247690256669e-5 -0.00012381657665646984 -0.0001060513002314226 -0.00014336454273257543 -2.6273517273907963e-5 -0.00013443355214367962 -0.00011772640368254384 5.8142917855357026e-5; -0.00016426227840953705 3.8750406823709946e-5 -0.00016896545741135337 6.287773693085625e-5 -0.00023414833150699074 -9.024866186905255e-5 -5.250578862277114e-5 1.4771482605049232e-5 4.903920388070612e-5 -1.5056565749805115e-5 -2.7055343813336912e-5 -3.6150345400787905e-5 0.00012460399273889032 -0.00018959841428684904 5.59981734875745e-5 -5.7859569202523714e-5 3.128588185403491e-5 9.417531226401265e-5 8.768360097489355e-6 -0.00014429034228331664 -3.1023010811836416e-5 1.3867254517356232e-6 -4.801726652450645e-7 -0.00012216993228893042 1.2424723430490129e-5 4.940673727242762e-7 -0.00010849987767530952 -0.00010241999290539078 0.00014787662338108665 -8.605758659203735e-5 -8.829695899865612e-5 -5.123121547174407e-5; 0.00011356015108418405 2.7071366793328487e-5 -9.253200317374982e-5 -9.714576065648294e-5 -5.355991145820394e-5 -7.102891452717263e-5 -0.0001862390612868223 0.00013162683405340648 -1.518616266452973e-5 -9.944340620763597e-5 0.00014779452118916057 -2.3401778372270137e-5 -6.563978734533861e-5 -7.489203329073166e-5 -0.0001216062423108395 -7.838356162296618e-5 9.001934319627299e-5 -1.4131764538385097e-5 3.603789679714524e-5 -0.00011180501808753271 -6.045551282049059e-6 -0.00015407762605804915 -5.7409806058285545e-6 5.1870307806031467e-5 6.254892884413106e-5 -2.9445750834336582e-5 4.0996570251584335e-6 7.989989041478042e-5 5.6035586175374077e-5 -1.4257706817214027e-5 3.442505710068212e-5 4.985538050191626e-5; -3.286226457697843e-5 7.116601226021708e-5 -4.15813339481313e-6 -0.00015561588044642346 4.6677620032421246e-5 0.00017545267654219408 3.806805023551872e-7 -0.0002276218856581526 5.648838303619153e-5 4.2552883298919666e-5 -8.8710416109675e-5 1.8849925238607527e-5 3.752279765248412e-6 -4.505891230786999e-5 -7.086704683154194e-5 -0.00014479503262686448 -9.620661696084314e-5 0.00010202059818999055 0.00017278433550250842 -1.789723502719685e-5 -0.00019770909968195506 -0.0001219667592681094 -0.00020703888550764658 0.00014147040381663563 0.00010443949029882698 0.00011035715759367784 1.360500830340939e-6 8.433908986496139e-5 0.00018007584911407747 8.514423809681148e-6 -2.6069172251884462e-5 -8.292402891806152e-5; -0.00017980657479471705 0.00024446112025782406 7.396478418980086e-5 -0.0002136196960576644 3.7242764176536686e-5 1.4736269133794072e-5 0.0001366226827145707 2.9863697148027907e-5 3.47168754730061e-5 9.010158488547151e-5 -0.00019571177448356433 4.4904729532047025e-5 9.770061535735622e-6 -7.90403389486311e-5 -0.0003369000876203486 -4.9430767592073575e-5 4.495960480438479e-5 -5.654441678177726e-5 -6.96700066089494e-6 -0.00019470894379753953 5.109323343517696e-5 -0.00014014006204650042 -9.7986750677022e-5 -0.0001122695355092022 0.00031233905270949217 -9.595301866045517e-5 -0.00014004799207884675 4.642200316326408e-5 -4.7739081980902574e-5 7.010633659103584e-5 1.6200796619726213e-5 4.682745226330068e-5; -8.707791502237676e-5 -0.00016045294259091404 -7.225309299964217e-6 6.393517461257629e-5 0.00011965997319786588 0.00011464300397201168 -4.54710484138598e-7 8.086373046389994e-5 0.00010225556971856165 5.646594573628806e-5 0.00011646972044254061 0.00017300317623304866 2.547411341803254e-5 0.00020660982037080955 -4.300398731141527e-5 0.00012930984295091323 -1.0876056957927053e-5 0.00011079502647627679 0.00020138553728467438 -0.0001565996973018465 9.479314858608428e-6 3.5817138187136965e-5 -8.627792131726496e-7 -0.00014728690811314856 -4.714957330141975e-5 -0.00010079988359355636 2.8411604862788284e-5 -6.40292601834721e-5 0.00015295690892380145 -0.00017042327462111866 1.3586666549750098e-6 -0.00011035699216298325; -9.573348773018575e-5 -6.596403084302646e-5 4.3452674684080424e-5 0.0001086085367870683 0.00019122731925846927 5.4777113357978367e-5 8.066239849824373e-5 -6.302660306291173e-5 2.8400239055481803e-5 5.968140346623821e-5 4.951374018955726e-5 0.00012019933371731891 7.560977708009373e-5 3.4008492614810506e-5 -7.40320631568984e-6 -3.3512321285234095e-5 -0.00013196392221915754 8.263702063515544e-5 -1.6554612490727632e-5 0.0001184425246475442 -2.1372411471843687e-5 -5.917089850016113e-5 0.0002714622724607726 9.150468128893434e-5 -2.436880725043461e-5 -6.784395639158907e-5 2.7014029067791572e-5 -0.00012284992674441712 -0.00018385308333281611 9.90075832970615e-5 1.5374796087707215e-5 -9.662256609488611e-5; -0.0001386334265873324 0.00012131389618733525 4.466273409779281e-5 4.0077018155701514e-5 -0.000293017518316645 0.00010328569575795903 -0.00010059439276230318 6.60725267778643e-5 8.064584059755521e-5 -1.5016103617397162e-5 -8.330881294676601e-5 -4.471882905929094e-5 6.147137125371105e-5 -0.00012007379748976322 -4.8398131667590425e-5 -1.6257205572562293e-6 -7.702819178230204e-5 -5.706687491621391e-5 6.199510923468868e-5 6.226215143104288e-5 -5.092248150611338e-5 -0.00012103188831244262 -8.225184186013564e-5 -6.83589974243825e-5 -5.881170775975113e-5 6.492609051636674e-5 8.942851806640524e-5 6.348038684230833e-5 -0.00019190394280593512 -0.00010644201441320344 3.5747502522860244e-6 -4.612950627344092e-5; 0.0001250565109943609 -3.803318398477016e-5 -0.00018813734644263657 2.5144961880325354e-5 -1.5799352725113988e-5 5.7106899179346915e-6 0.00012570146642920253 8.263042397392838e-5 -2.4938404778911087e-5 5.9742826094773044e-5 1.7748336147079163e-5 -1.7495322056308855e-5 8.44099776872094e-5 9.904286704786649e-5 -2.7823874945736187e-5 6.900179866297774e-5 8.423667892875199e-5 -0.00012994326246607512 0.00016383557519239006 -8.6307765771014e-5 0.00017196782570625947 8.854098993415589e-5 -0.00011011763857150038 -1.8453891273199884e-5 8.948652245398976e-5 -0.00011606780029305586 7.800594146283569e-6 -5.584691828754987e-5 -8.175605037887088e-5 -1.3834746853351565e-5 -9.989011599154583e-5 3.248557183487831e-5; 0.00011434342265331033 0.00012731274434102046 0.00018125484599934054 -6.417941888661596e-6 0.00016003860470575705 -2.7764200452097264e-5 -0.00015557107438596475 -4.889421073401694e-5 0.00010352732603407882 -0.00012203002300609445 -2.6824847039266858e-5 1.8943287402204494e-5 4.4935731027201774e-5 -0.00010623089950329312 2.217705075053717e-6 -0.0001845766357591615 -2.8979652809521424e-5 -2.7130490009744172e-5 3.68235294145288e-5 -0.00012809090297466653 0.00017954646570343495 6.724961964848815e-5 6.213399571336061e-5 7.354429068712e-5 -8.913717213815875e-5 8.05079059928416e-5 9.043319256945682e-5 0.00011347570650014396 -2.4455589694285724e-5 -0.00014946056511050618 0.0001902121027951816 -0.0001872933910208075; 9.272676705521143e-5 1.8322835406517704e-5 7.896366021967646e-5 0.00017118274134490438 -0.00012508947967644816 2.2917257031868047e-5 6.254931812081441e-5 0.00011682660131853179 -3.7229595742887346e-5 -7.725272136367971e-5 -1.972325245945651e-5 -1.986555927642713e-5 2.5175235881113238e-5 -5.4809451664102274e-5 0.0001604712657201476 4.3928535736040505e-6 -6.932562375070117e-5 7.142987904743973e-5 2.5499685383050703e-5 -3.063113164599101e-5 0.00016012533759136348 -7.974486053781748e-5 3.227007853937703e-5 0.00010474028257094274 -4.927118554106664e-5 -3.3435405763798376e-5 0.00020111128704528649 -0.00013596040063635267 0.00015646074333132401 -0.00012764114345982755 0.00010554667695332115 7.906899175661163e-6; 7.832724304684747e-5 0.00011860259337631721 5.2670081005678974e-5 -7.044250878457183e-5 3.933275229578696e-5 1.6428454274902315e-5 -0.00011026025231180859 -4.351799861829582e-5 0.00022462637798676126 7.477503322811628e-5 -9.054853353369269e-5 0.0001636809488460592 -3.6605536932576863e-5 5.4679802362124905e-5 -2.689984967070993e-5 4.437349571312617e-5 -1.5757608251401417e-5 3.40711799268775e-5 6.963734905775995e-6 -0.00011319070244858188 -1.3053092059964766e-5 -4.121985830201956e-5 3.6057756462149665e-5 9.12644850374371e-5 7.24840451781786e-5 -0.00016864608687533223 7.967846113538016e-5 2.618102235067694e-5 -0.00015273812232258918 -1.927561015540892e-5 6.645587063908374e-5 -9.559390809804468e-5; 4.395724791089245e-5 0.00014572366963243983 -1.430688170961995e-5 -7.624178651617007e-5 -2.9589010225158635e-5 0.0001445509162322753 4.084467327907547e-5 6.83415248232083e-5 -9.741108333123053e-5 -0.000140145227121088 -1.690259777836923e-5 1.7843226672134975e-5 -0.00010839004498570191 4.6790069248396125e-5 0.00019295690181304044 -6.057053796843751e-5 -0.000136380311957456 -3.047553654773248e-5 0.00010375873280559563 0.00019073496987680766 6.412655721426998e-7 -0.00013129272956297842 -0.00010533097047022557 -1.7868586653043936e-5 0.00011524378096788207 -1.684138878494626e-5 -0.0001623334926961828 0.0001098260301687941 -5.975628454852906e-7 -0.00013197998741538395 -0.0001310086761777206 8.566278850645178e-6; 9.623619802120242e-6 -5.630430780572467e-5 -4.3869965453516974e-5 0.00024475484212755575 0.00017727959027408977 -4.4379839097267536e-5 -8.231992717471553e-5 0.00010994532465131353 -6.6913481771491425e-6 4.381102711455937e-6 4.286832961381729e-5 -0.00012954457372531874 -0.00011219640612241469 -6.590277655620708e-5 6.441599077672971e-5 0.00012012449121696296 1.5023506306285124e-5 -0.00011764798284846008 0.0001240260561405437 -0.00021996973839693348 6.367595312779148e-5 -9.014393902022972e-5 0.00010708529850828807 -0.00020586807397764752 0.0001277812942805455 7.886618084169345e-5 -0.00013825837697938672 -6.023945854965306e-5 2.3244555335373996e-5 -3.649894381512534e-5 8.609589246928112e-5 2.154942455829471e-5; 7.344085085115635e-5 9.494051979764576e-5 -7.020050773952866e-5 -3.47932569051327e-5 -1.9189931240717726e-5 -3.827763660900955e-5 0.00011976364973143876 -2.1197418878173714e-5 -8.067384127151951e-5 -3.2449034311313e-5 -7.22589041778535e-6 -2.4172878088748843e-5 -3.2152495382789145e-5 3.2865585548833235e-5 -2.211490258140586e-5 -4.670317006038299e-5 -0.0002356876744082825 -3.964007422418125e-5 -3.4805797018080755e-5 0.0001779288481562698 7.392990434219423e-5 -0.00010842721761527265 1.9616254089313873e-5 1.4932683793275775e-5 0.00014690596932688015 -6.202380112462474e-5 -4.1721903958560575e-5 0.00020422443546031936 3.432314175789525e-5 1.6616160305246882e-5 -0.00014446736398475972 6.397311286910648e-5; -2.889110084245857e-5 -9.694335068663901e-5 2.623602693754153e-5 -3.555578342959509e-5 1.9832140137146786e-5 -0.00016543909925125145 -7.298345491610604e-5 -1.263863515808677e-5 0.00016205883974587405 -0.0001916511731137912 3.345883186235264e-5 2.9448911559007385e-5 -7.480290270498576e-5 -6.5811576091851305e-6 2.4714913419702998e-5 0.00016499298606911512 7.270861428393781e-6 -0.00022154319350695494 9.388262941181445e-5 5.77878587905397e-5 0.0001424398457715654 -3.490909267886745e-5 9.038901287974353e-5 -5.186711237478518e-6 -9.168380830990655e-5 0.00022661801875481217 2.6508766208710216e-5 0.0002563943368699689 -4.074092913973853e-5 -6.513892487012006e-5 8.830162795172779e-5 0.00012998042319347757; 3.4253720196827494e-5 1.798015150689254e-5 0.00016944630956550123 0.00018952923314919035 -8.010951237246969e-5 6.78153261569533e-5 0.00012073298025916682 1.5524204764910925e-5 2.149263277806693e-5 -0.00013182813231560087 -1.6381275007925054e-5 -9.910221374320468e-5 1.7636186067648297e-5 -5.268821437220784e-5 0.00013420120708449297 0.0001293170604649312 2.0701548829489878e-5 -7.521463371540164e-5 -8.710742654238543e-5 -0.00025061236042256115 -4.8143643970444945e-5 -5.369545428856913e-5 8.907058841123349e-5 -0.00012929745239398095 3.539231921423221e-6 0.00019203456363448267 -0.00014134440093493936 0.00011003736622204161 2.023970560093449e-5 2.5263379038769733e-6 -4.822022082404726e-6 0.0001771006315275394; 0.00010651349173860798 3.9165310846895345e-5 7.680510923966967e-5 0.00018133198198631633 8.555243283046649e-5 6.65254034917068e-5 -0.00012029204783027752 -7.923485219942257e-5 -0.00020872269062090432 -1.1949828438789084e-5 -0.00013770113702628667 -5.8760274731300263e-5 -4.9586103449089304e-5 -9.114579126469563e-5 -0.00010791454525895475 2.2212533332661305e-5 2.3586697972516312e-5 0.00013769112291049437 0.00011113568933217088 -9.670562696408242e-5 -9.732571317579523e-5 8.424068875340751e-6 0.00017610692764867306 9.734488292599335e-5 1.8883087537894795e-6 1.4113784056119665e-5 -5.9431841981110864e-5 -1.222159000315467e-5 -8.046195700297032e-5 -6.96678992593451e-5 0.0002334657171392091 -6.003854406879061e-5; -3.876086405363789e-5 0.000145612443948611 -4.148602304632778e-5 8.296301463849755e-5 -6.987636218547931e-5 0.00013355839407302655 -0.00011680727105124031 4.3204024863198634e-6 3.7165725434877764e-5 -4.780718584663334e-5 -9.478211789393522e-5 5.348930360048908e-5 -1.4913006206934854e-5 8.296735838519333e-5 2.3375410095351004e-6 -1.0566429163756406e-5 -0.0001082396781654274 -7.448145400269038e-5 -3.021695622885222e-5 7.662169743560146e-5 0.00014964493972946297 -0.0001151702460716674 -3.5049489024114195e-5 -3.067204009277854e-5 -3.2885456945411235e-6 -3.6294288956579837e-5 7.326320500207757e-7 -3.555814030697972e-5 4.464006203341562e-5 8.841101886935429e-5 6.585965390548024e-5 -0.00017228935096647653; -1.2483372642958471e-5 0.00012089655740529559 0.00014399012673057686 4.577739912311878e-5 -0.00011504468181991542 -5.967454739947039e-6 6.013757342964474e-6 6.874554912902037e-5 8.820345670040062e-6 6.585115557038457e-6 -8.073531148330143e-6 8.190084031605252e-6 8.11826490691684e-5 0.00016108827787794465 2.6274885954248864e-5 -3.101919116924134e-5 9.509940014327816e-6 4.0795067093505825e-5 -0.0001408103393030505 3.1310727310599164e-6 2.3512846951414083e-5 -8.18591846607777e-5 -0.00010960053384659647 -0.0001649685267464404 4.406858920637707e-6 -6.639141554368278e-5 -0.00010153624077683168 7.491944075073239e-6 0.00010798034293220528 3.212380636360881e-5 0.00019854514050640496 6.840170920005323e-5; -0.00017579407514072337 4.489067422460228e-5 -7.197556366128503e-8 2.2736783911255665e-5 4.300883382399852e-6 -0.00012200352440355486 4.1341462100405944e-5 2.3286806289109545e-5 4.146766358522396e-5 8.13425166475701e-6 8.115635345223869e-5 1.301165801911046e-5 -1.997599039109328e-5 -0.0001287773026991497 -9.633182863707405e-5 2.8751785329092656e-5 -5.3078352458537016e-5 -5.785451295318533e-5 -6.204550091874262e-5 5.9304502691390684e-5 2.9215069197227447e-5 9.509748841870187e-5 -3.0426162731007075e-5 0.0001710924573568783 0.00020654974131380363 5.361922948273493e-5 0.00010196792781121315 0.000206938597592508 2.0896103059054496e-5 1.660268234323378e-5 -7.48372019291711e-5 0.0001276465066336227; -1.60281667731119e-5 -1.9727156303778913e-5 2.5624336795805126e-5 0.0002212187339510293 0.0001836789690615149 4.1394736636921116e-5 7.610411035879564e-5 0.00019931051998448856 7.185953496540945e-5 7.894572105309916e-5 -7.644175433505005e-5 4.332077356688684e-5 -1.3251165672921795e-5 -2.122252930253917e-5 -0.00016582142210519454 -0.00018792349668634117 -1.3796476868238477e-5 0.00027789486399437154 -0.00010559188898601651 0.00012077053021225195 3.485316325260262e-5 -4.4113997991150856e-5 3.98837348812216e-5 -4.367715313397303e-5 0.00016984150921510918 -4.0079544977704e-5 -8.422657059687666e-5 1.4083042752643364e-5 9.901982831267487e-6 4.940452914619832e-5 -3.41633183224602e-5 -7.691454606079752e-5; 4.177785845965673e-5 0.00014914023373420454 0.0001265711008818302 0.00017125432819738403 -6.546161133883859e-5 0.0003315608165939517 0.00011466683386969963 7.810161801895931e-5 6.979819332627229e-5 3.1925116818449625e-5 6.511150892033518e-5 -8.501502710715354e-5 6.257647070013419e-5 0.00011156811266117955 -9.155836854924752e-5 9.050445547472268e-5 9.27792472429886e-5 -6.762877856099075e-6 -1.786228753190584e-5 -2.754007331532779e-5 -4.3994327679736175e-5 7.529884638642367e-5 0.0001145274337977617 -8.97755042751728e-5 -0.00017445579836932387 -0.00013224258404080248 9.718283141408786e-5 0.00013803185462158617 2.7227947757592813e-5 6.939284245162795e-5 -0.00015714337027021338 -3.728622945346778e-5; -0.00010186136050951435 1.8789323529485844e-5 1.1672055460358934e-5 3.7434796193989666e-5 0.0001484983383295165 -8.119666590692177e-5 1.222576492975152e-5 3.6238286792196844e-5 -6.603625506730367e-5 -0.00020148798494236216 -8.338564503110973e-5 -1.1614311280284342e-5 3.86197186332639e-5 5.2269193200768884e-5 0.00017673061092743052 -3.642733490895133e-5 0.00012131192643092138 -0.00012944446489317128 -9.729256116887287e-5 0.00011058760874104137 -4.195228303898123e-6 0.00019134560870835842 8.317471810979716e-5 6.72697730794587e-5 2.9985210584265498e-5 -0.00018312987076686354 7.863057698951943e-5 0.00014545877793232335 -3.1367528482684385e-5 0.00011239547320539483 0.00010905382960001258 7.624326981268932e-5; 0.00010192715870302892 -6.524511247321899e-5 7.142910197230408e-5 -8.020947263490455e-5 -1.0316943025699973e-5 -7.236264338629764e-5 -0.00015450624761784601 -7.583041571692478e-5 -0.00017704781914277154 5.233789868738397e-5 -3.881288632313899e-5 -7.597936184525612e-5 0.00018876205302162434 -0.00013852900139668946 -1.1571070735788319e-5 -2.5379029302584267e-5 -0.00022445549968310753 -4.258207611188304e-5 -0.00011458031060088479 -0.00017382442440053598 -1.91164707028623e-5 3.0380784828722702e-5 9.444572804581578e-5 -4.1856517618598036e-5 0.00017476641385669176 9.583537046650304e-5 8.147621718321227e-5 9.057154982672157e-5 -0.00016547185789007933 8.474836811741688e-5 4.487058338794814e-5 8.553353125548059e-5; 4.900922637414417e-5 -0.000154826828050608 -2.9895952805709688e-5 -7.672804507212121e-5 -7.277446432755613e-5 -1.7218015030562925e-5 -2.0765311758912264e-5 0.00015757541974987927 0.0001682118802515854 0.00012837107921466852 -0.00016522054805439193 2.8739767246117575e-5 -3.516213631275628e-5 -0.00010714442190266579 8.355939486742737e-5 -1.5250838017271136e-5 1.8060464086990835e-5 0.0001318892375517694 -0.00010067811102391554 0.0001557282141827111 -8.925007726550458e-5 0.00017002290975750095 0.00010146711271470889 -0.0002444862709429051 0.000163852213760656 1.818632723978417e-5 -3.342262583220846e-5 7.465219665997806e-5 -0.00013326545898355744 6.0555567203375304e-5 7.850177486641925e-5 -5.499545453250801e-5; -0.00011949762882137913 -0.00025112358742333346 0.00011172494816421646 9.275173960504754e-6 1.7618349028651443e-5 -1.1542232987953863e-5 3.660766344111685e-5 -2.8686171215320368e-6 -4.9379003376799704e-5 3.654852809560827e-5 -0.00012682495469630388 -8.364065598271848e-6 0.00018711091945008474 -2.9717128546307235e-5 3.275818346017579e-5 0.00015201551252104665 -0.00010105017316233472 -0.0001207756435003642 -9.991262578096234e-6 -0.00014404272258687495 -7.473465701036349e-5 -2.1085530880355043e-5 -0.00021641633828110176 2.972371987127045e-5 0.00014701153087955912 -0.00015106562239523654 2.60422963212703e-5 -0.00020045524468580904 7.843923196491998e-5 -6.924058141592579e-5 5.3967105140345805e-5 4.612562515123608e-6; -1.732914151749008e-5 -4.528500230468802e-5 -1.1448045914096435e-5 6.493302626234313e-5 9.073254605221222e-5 -1.1837639793015286e-5 -2.723165616577822e-5 -1.2551640546472034e-6 -8.463411231593561e-5 -4.551146137483005e-6 -2.6098174746802816e-5 3.6009650468376435e-5 7.6143304814989134e-6 -4.742366454399268e-5 -0.00021223918756532482 1.2166248236216557e-5 0.00012086919141633674 -4.928932193967629e-5 -4.523670449804695e-5 -0.00019525791740113269 2.342122152382682e-5 -0.00016151671170325227 -4.1513582434950056e-5 4.5835478342386437e-5 -0.0002270068931446655 -6.700325920780472e-5 -7.54863054661333e-5 -4.523866298604078e-6 6.766438437094556e-5 4.798178761314311e-5 0.00010201917157440278 -7.773211709524491e-5; 0.0001495518057810956 0.00020418711702514773 3.953914111772107e-5 -6.365618527335629e-6 6.240187704971275e-5 0.00015583859695763455 -9.599388666974163e-5 5.6562542714540895e-5 7.875502834504688e-5 -8.023809267608356e-6 -6.209254442072207e-5 3.4325022544369584e-5 6.87778787425581e-5 -8.38185465807251e-5 -9.155928509132557e-5 -2.6343574599006666e-5 -1.791681151079214e-5 -7.370268812227703e-5 3.1803815919564016e-5 -0.00022330733426205592 -8.918961237923293e-6 -0.00024460524090390393 0.0001476391748030524 7.990460054425742e-5 1.0085648146932791e-5 0.00013233675909869608 -0.00021251011875282548 -6.576830727569789e-5 -0.00010873914524270921 -2.6075373706400867e-5 -8.926604877050693e-5 -8.430424585530198e-5], bias = [1.921786307209035e-9; 8.789190050134396e-11; 1.2504148812771212e-9; 2.759638658565189e-9; -4.0864654563959214e-9; -1.5256145600285591e-9; -7.257926362306835e-10; -2.32442960482532e-9; 2.26637932249765e-9; 2.085537853776543e-9; -3.0927651263003428e-9; 1.1204678255460235e-9; 1.224877475321225e-9; 2.7271956201550424e-9; 1.020909304245852e-9; -6.646559494382461e-10; 3.899342734955543e-11; -8.234145849406738e-11; 1.581324318981226e-9; 1.3057169819360422e-9; 1.4432009549384312e-10; -3.935652982223644e-10; 1.2538523267575744e-9; 1.9956939078962535e-9; 2.55591752995915e-9; 4.101057123727002e-9; 2.329851411471181e-9; -1.8697572701351324e-9; 8.933663432767605e-10; -2.903453407028585e-9; -2.9091078170073337e-9; -9.91884728601028e-10;;]), layer_4 = (weight = [-0.0005572566820208557 -0.0007937026262524729 -0.0006027371350159636 -0.0008097167570915016 -0.0006632364353912129 -0.0007586931955229081 -0.0006962550015859167 -0.0007105567364154836 -0.0006658085334180905 -0.0006602566700924333 -0.0006626099449352349 -0.0007211122134551628 -0.0006367930196168518 -0.0007610696972092854 -0.0007117310752981213 -0.0005740901170262534 -0.0007709460867192456 -0.0006457321711362864 -0.0006559189228036226 -0.00071990248582105 -0.0007860558927556543 -0.00069686912726478 -0.0005244937356517977 -0.0009616786115207227 -0.0007539177071414011 -0.0006872892690566108 -0.0006167645875060232 -0.0006239984997886235 -0.0005674648472694767 -0.0006296852286443204 -0.0008124558416854573 -0.0007791480513692754; 0.00014216770893870395 5.398784074717275e-5 0.00023440740675234172 0.00018101927179873937 0.00038097963229907544 0.0004179167254411029 0.00016871553604570189 0.0002684071421866321 0.00029311844363497764 -4.684615903553573e-6 0.0001715444594282229 0.00015608806149225427 0.0002515337605158771 0.00022016138424090162 0.00030322691643611805 0.00015330975900653944 0.0002746350862441426 0.0002501048844999265 0.0003728965915869229 0.00032248962985238777 0.0003586351545196296 0.00018497025779880998 0.0003025448352555981 0.00021959541944993648 0.0003294169957206736 0.00028422808379670095 0.00020632328939831497 0.00032818817391118334 5.102628667107437e-5 0.00014884949250286642 0.00017417684072532018 0.00026142470856713477], bias = [-0.0006732848512731586; 0.00023393989498430928;;]))

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: 48 default, 0 interactive, 24 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-1/julialang/lux-dot-jl/docs
  JULIA_AMDGPU_LOGGING_ENABLED = true
  JULIA_DEBUG = Literate
  JULIA_CPU_THREADS = 2
  JULIA_NUM_THREADS = 48
  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