Correcting ODE initial conditions using automatic differentiation

General Usage
, ,
2

Try

function twobody(du,u,p,t)
    r = norm(u[1:3])
    mu = p[1]
    du[1] = dx = u[4]
    du[2] = dy = u[5]
    du[3] = dz = u[6]
    du[4] = dvx = -mu*u[1]/(r^3)
    du[5] = dvy = -mu*u[2]/(r^3)
    du[6] = dvz = -mu*u[3]/(r^3)
end

# known quantities
r1 = [1., 0., 0.]
r2 = [1., 1/8, 1/8]
Δt = 0.125
p = [1.]

v1 = [0.05, 1., 1.]  #initial guess for v1
u0 = [r1; v1] # initial state

prob = ODEProblem(twobody, u0, (0,Δt), p)
sol = solve(prob, reltol=1e-8, abstol=1e-8)

function loss_r(v)
    u = [r1; v]
    _prob = remake(prob, u0=u)
    sol = solve(_prob,Tsit5(),saveat=Δt,sensealg=QuadratureAdjoint())
    norm(r2-sol[end][1:3])
end

dLdv = Zygote.gradient(loss_r, v1)
5 Likes