# Gradient of loss function from ODE solution

**URL:** <https://discourse.julialang.org/t/gradient-of-loss-function-from-ode-solution/113910>\
**Category:** General Usage\
**Tags:** question, autodiff\
**Created:** [May 6, 2024, 2:29pm UTC](https://discourse.julialang.org/t/gradient-of-loss-function-from-ode-solution/113910 "2024-05-06T14:29:36Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Enrico\_Camporeale](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/enrico_camporeale/32/206102_2.png) [@Enrico\_Camporeale](https://discourse.julialang.org/u/Enrico_Camporeale)\
**Post date:** [May 6, 2024, 2:29pm UTC](https://discourse.julialang.org/t/gradient-of-loss-function-from-ode-solution/113910/1 "2024-05-06T14:29:36Z")

</div>

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!

```julia
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)

```

---

<div class="post-metadata">

**Author:** ![Enrico\_Camporeale](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/enrico_camporeale/32/206102_2.png) [@Enrico\_Camporeale](https://discourse.julialang.org/u/Enrico_Camporeale)\
**Post date:** [May 7, 2024, 8:53am UTC](https://discourse.julialang.org/t/gradient-of-loss-function-from-ode-solution/113910/2 "2024-05-07T08:53:56Z")

</div>

@ChrisRackauckas Could you please give a look at this whenever you got the chance? It should be really straightforward and I have tried to follow your example here (just a slightly more complicated ODE)

> [@Correcting ODE initial conditions using automatic differentiation](https://discourse.julialang.org/t/correcting-ode-initial-conditions-using-automatic-differentiation/66018):
>
> Hi all, I’m trying to find a way to correct initial conditions of a system of ODEs by using automatic differentiation through the ODE solver. I have implemented a similar algorithm in Python using [TorchDIffEq](https://github.com/rtqichen/torchdiffeq) but I’m new to Julia and having trouble debugging the errors. I’ve gone through the docs and examples of both Zygote and ReverseDiff (which I think are best suited to my purpose) but I’m having difficulty adapting them to my problem. What I want to do is first solve the system of ODEs with…

Something is obviously escaping me 🙂  
Many thanks!