I am interested in the autodiff feature of Julia. Specifically I have an ODE problem that solves for the motion of solar system planets and I have defined a loss function. I want to calculate the gradient of the loss function wrt the initial conditions of the ODE.
Using ReverseDiff.gradient it gives me a zero vector, which is obviously wrong.
I am new to Julia, and any help is appreciated!
using CSV
using DataFrames
using ModelingToolkit
using Optimization
using OptimizationPolyalgorithms
using OrdinaryDiffEq
using Plots
using SciMLSensitivity
using Statistics
using Zygote
using ForwardDiff
using ReverseDiff
gr()
#planets = ["Sun", "Jupiter", "Saturn", "Uranus", "Neptune", "Pluto"]
AU = 1.496e11;
x=[0 1.563021412664830E+07 -9.030189258080004E+07 -1.018974476358996E+08 -2.443763125844157E+08 -2.35165468275322006E+08 -1.011712827283427E+09 2.934840841770302E+09 4.055112581124043E+09 9.514009594170194E+08 0.5*AU/1e3].*1e3;
y=[0 4.327888220902108E+07 5.802615456116644E+07 1.065689158175689E+08 4.473211564076996E+07 7.421837640432589E+08 -1.077496255617324E+09 6.048399137411513E+08 -1.914578873112663E+09 -4.776029500570151E+09 0].*1e3;
z=[0 2.102123103174893E+06 6.006513603716755E+06 -3.381951053601424E+03 6.935657388967808E+06 2.179850895804323E+06 5.901251900068215E+07 -3.576451387567792E+07 -5.400973716179796E+07 2.358627841705075E+08 0].*1e3;
pos=[x;y;z]
ux=[0 -5.557001175482630E+01 -1.907374632532257E+01 -2.201749257051057E+01 -3.456935754608896E+00 -1.262559929908801E+01 6.507898648442419E+00 -1.433852081777671E+00 2.275119229131818E+00 5.431808363374300E+00 rand(1)*3].*1e3;
uy=[0 1.840863017229157E+01 -2.963461693326599E+01 -2.071074857788741E+01 -2.176307370133160E+01 -3.332552395475581E+00 -6.640809674126991E+00 6.347897341634990E+00 4.942356914027413E+00 -2.387056445508962E-02 rand(1)*3].*1e3;
uz=[0 6.602621285552567E+00 6.946391255404438E-01 1.575245213712245E-03 -3.711433859326417E-01 2.962741332356101E-01 -1.434198106014633E-01 4.228261484335974E-02 -1.548950389954096E-01 -1.551877289694926E+00 rand(1)*3].*1e3;
vel=[ux;uy;uz]
tspan = (0.0, 1.0*365.25*24*3600)
G=6.67e-11;
m=1.989e30;# the mass of the sun
Gm=[G.*m./1e9; 22032.09; 324858.63; 398600.440;
42828.3; 126686511; 37931207.8;
5793966; 6835107; 872.4; G/1e9].*1e9;
function many_body!(du, u, p, t)
interactions = vec([1 2 3 4 5 6 7 8 9 10 11]);
for i in 1:11
ind = vec([1 2 3 4 5 6 7 8 9 10 11]); # bodies we are considering interactions between
i2 = findall(ind .== i) # a body can't interact with itself
deleteat!(ind, i2)
# square of distance between this planet and the other objects
R2 = (u[(i-1)*6+1] .- u[(ind .- 1).*6 .+ 1]).^2 .+
(u[(i-1)*6+2] .- u[(ind .- 1).*6 .+ 2]).^2 .+
(u[(i-1)*6+3] .- u[(ind .- 1).*6 .+ 3]).^2
# rate of change of position, because we are solving a second order ODE
# - always the same
du[(i-1)*6+1] = u[(i-1)*6+4] # dx/dt
du[(i-1)*6+2] = u[(i-1)*6+5] # dy/dt
du[(i-1)*6+3] = u[(i-1)*6+6] # dz/dt
# inverse square law between body i and the rest of them
du[(i-1)*6+4] = - sum(Gm[ind] ./ R2 .* (u[(i-1)*6+1] .- u[(ind .- 1).*6 .+ 1]) ./ sqrt.(R2)) # dvx/dt
du[(i-1)*6+5] = - sum(Gm[ind] ./ R2 .* (u[(i-1)*6+2] .- u[(ind .- 1).*6 .+ 2]) ./ sqrt.(R2)) # dvy/dt
du[(i-1)*6+6] = - sum(Gm[ind] ./ R2 .* (u[(i-1)*6+3] .- u[(ind .- 1).*6 .+ 3]) ./ sqrt.(R2)) # dvz/dt
end
end
u0=reshape([x;y;z;ux;uy;uz],66)
prob = ODEProblem(many_body!, u0, tspan)
#sol = solve(prob, Tsit5(),reltol=1e-8);
function loss(u)
_prob = remake(prob, u0=u)
sol = solve(_prob,Tsit5(),reltol=1e-8)
Earth_pos=reduce(hcat,[sol.u[i][19:21] for i in 1:length(sol)]);
Sun_pos=reduce(hcat,[sol.u[i][1:3] for i in 1:length(sol)]);
probe_pos=reduce(hcat,[sol.u[i][61:63] for i in 1:length(sol)]);
mid_point = 0.5*(Earth_pos+Sun_pos);
mean(sum(((mid_point+probe_pos)/AU).^2,dims=1));
end
dLdu = ReverseDiff.gradient(loss,u0)
df=DataFrame(sol.u,:auto)
CSV.write("sol.csv",df)
