# Effect of Interpolation on ODE Gradient Computation

**URL:** https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525
**Category:** General Usage
**Created:** [August 5, 2023, 10:57pm UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525 "2023-08-05T22:57:15Z")
**Posts on this page:** 8
**Page:** 1

<div class="post-metadata">

### Author: ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)
#### Post date: [August 5, 2023, 10:57pm UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/1 "2023-08-05T22:57:15Z")

</div>

I am interested in understanding how to optimize the coefficients of an ODE (such as Lotka-Volterra) under some cost function given that this ODE is influenced by a set of controls c1(t), …, cn(t) for which I observe simulated trajectories. As a simple example, here I’ve written a MRE with one simulation of Lotka-Volterra where the second coordinate is subjected to a time-varying control. I want to take the gradient of the square cost of this simulation with respect to the Lotka-Volterra parameters and see how this adjoint-based gradient computation is affected by the control input which is an interpolated function.

```julia
cd(@ __DIR__ )

using Pkg
Pkg.activate(".")
Pkg.instantiate()

using Flux
using Zygote
using ForwardDiff
using DifferentialEquations
using SciMLSensitivity
using Optimization
using Random
using MAT
using Interpolations
using Random
using Profile
using FlameGraphs

# Define Controlled ODE
struct ControlledODE{Of, So, Se, T, Sa}
    odefunc::Of
    solver::So
    sensealg::Se
    tspan::T
    saveat::Sa
end

function ControlledODE(odefunc;
    solver=Tsit5(),
    sensealg=InterpolatingAdjoint(; autojacvec=ZygoteVJP()),
    tspan=(0.0f0, 1.0f0),
    saveat=[])
    return ControlledODE(odefunc, solver, sensealg, tspan, saveat)
end

function (c::ControlledODE)(u0, control; tspan=c.tspan, saveat=c.saveat)
    function ode(du, u, p, t)
        c.odefunc(du, u, t)
        du[length(du)] += control(t)
    end
    prob = ODEProblem(ode, u0, tspan)
    return solve(prob, c.solver; sensealg=c.sensealg, saveat=saveat)
end

Flux.@functor ControlledODE

# Lotka Volterra with control input
tspan = (0.0, 10.0)
Nt = 100
tsteps = LinRange(tspan[1], tspan[2], Nt)

v = sin.(pi.*tsteps)
control = linear_interpolation(tsteps, v)

p = Float64[2.2, 1.0, 2.0, 0.4]
function lotka_volterra(du, u, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = (α - β*y)x
    du[2] = (δ*x - γ)y
end

u0 = Float64[1.0, 1.0]

sensealg = SciMLSensitivity.InterpolatingAdjoint(; autojacvec=ZygoteVJP())
model = ControlledODE(lotka_volterra; solver=Tsit5(), sensealg=sensealg, tspan=tspan, saveat=tsteps)

# Cost for simulation
function cost(model)
    return sum(abs2, model(u0, control))
end

# Compute the gradient
Flux.gradient(m -> cost(m), model)

```

I am currently dealing with a bug in my code which makes the gradient computation fails coming from the fact that Flux cannot see how the parameter vector p has any effect on the model simulation.

Any help would be much appreciated!

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [August 6, 2023, 2:41am UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/2 "2023-08-06T02:41:02Z")

</div>

Don’t do anything fancy. Just make it a function of `p`.

```julia
cd(@ __DIR__ )

using Pkg
Pkg.activate(".")
Pkg.instantiate()

using Flux
using Zygote
using ForwardDiff
using DifferentialEquations
using SciMLSensitivity
using Optimization
using Random
using MAT
using Interpolations
using Random
using Profile
using FlameGraphs

# Define Controlled ODE
struct ControlledODE{Of, So, Se, T, Sa}
    odefunc::Of
    solver::So
    sensealg::Se
    tspan::T
    saveat::Sa
end

function ControlledODE(odefunc;
    solver=Tsit5(),
    sensealg=InterpolatingAdjoint(; autojacvec=ZygoteVJP()),
    tspan=(0.0f0, 1.0f0),
    saveat=[])
    return ControlledODE(odefunc, solver, sensealg, tspan, saveat)
end

function (c::ControlledODE)(u0, control, p; tspan=c.tspan, saveat=c.saveat)
    function ode(du, u, _p, t)
        c.odefunc(du, u, _p, t)
        du[length(du)] += control(t)
    end
    prob = ODEProblem(ode, u0, tspan, p)
    return solve(prob, c.solver; sensealg=c.sensealg, saveat=saveat)
end

Flux.@functor ControlledODE

# Lotka Volterra with control input
tspan = (0.0, 10.0)
Nt = 100
tsteps = LinRange(tspan[1], tspan[2], Nt)

v = sin.(pi.*tsteps)
control = linear_interpolation(tsteps, v)

p = Float64[2.2, 1.0, 2.0, 0.4]
function lotka_volterra(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = (α - β*y)*x
    du[2] = (δ*x - γ)*y
end

u0 = Float64[1.0, 1.0]

sensealg = SciMLSensitivity.InterpolatingAdjoint(; autojacvec=ZygoteVJP())
model = ControlledODE(lotka_volterra; solver=Tsit5(), sensealg=sensealg, 
                      tspan=tspan, saveat=tsteps)

# Cost for simulation
function cost(p)
    return sum(abs2, model(u0, control, p))
end

# Compute the gradient
Flux.gradient(_p -> cost(_p), p)

```

---

<div class="post-metadata">

### Author: ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)
#### Post date: [August 6, 2023, 3:08am UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/3 "2023-08-06T03:08:20Z")

</div>

Thanks @ChrisRackauckas! @Oscar_Smith not sure what the best way is to plot a flamegraph, I have it either as an svg or html file.

---

<div class="post-metadata">

### Author: ![Oscar\_Smith](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oscar_smith/32/25343_2.png) [@Oscar\_Smith](https://discourse.julialang.org/u/Oscar_Smith)
#### Post date: [August 6, 2023, 3:12am UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/4 "2023-08-06T03:12:50Z")

</div>

The image appears to be missing.

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [August 6, 2023, 3:51am UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/5 "2023-08-06T03:51:09Z")

</div>

Use [GitHub - tkluck/StatProfilerHTML.jl: Show Julia profiling data in an explorable HTML page](https://github.com/tkluck/StatProfilerHTML.jl) and then share the whole folder.

---

<div class="post-metadata">

### Author: ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)
#### Post date: [August 6, 2023, 8:13pm UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/6 "2023-08-06T20:13:37Z")

</div>

Perfect, here is the whole folder from StatProfiler in a google drive folder: [https://drive.google.com/drive/folders/1BRddREe0AdeUPkxvexiZTXV5BIcmK74e?usp=drive\_link](https://drive.google.com/drive/folders/1BRddREe0AdeUPkxvexiZTXV5BIcmK74e?usp=drive_link)

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [August 6, 2023, 8:26pm UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/7 "2023-08-06T20:26:49Z")

</div>

It looks like you profiled compilation. Run it then profile it.

---

<div class="post-metadata">

### Author: ![George\_Stepaniants](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/george_stepaniants/32/47230_2.png) [@George\_Stepaniants](https://discourse.julialang.org/u/George_Stepaniants)
#### Post date: [August 6, 2023, 8:36pm UTC](https://discourse.julialang.org/t/effect-of-interpolation-on-ode-gradient-computation/102525/8 "2023-08-06T20:36:52Z")

</div>

Darn it, I made sure to run the gradient computation once before. And then I ran @profilehtml “gradient computation”.

But somehow it seems to have again profiled the compilation. Hopefully now the updated google drive has the right thing.

Let me know if it’s still incorrect. I’m running this exactly:

```julia
cd(@ __DIR__ )

using Pkg
Pkg.activate(".")
Pkg.instantiate()

using Flux
using Zygote
using ForwardDiff
using DifferentialEquations
using SciMLSensitivity
using Optimization
using Random
using MAT
using Interpolations
using Random
using StatProfilerHTML
using FlameGraphs
using ImageShow

# Define Controlled ODE
struct ControlledODE{Of, So, Se, T, Sa}
    odefunc::Of
    solver::So
    sensealg::Se
    tspan::T
    saveat::Sa
end

function ControlledODE(odefunc;
    solver=Tsit5(),
    sensealg=InterpolatingAdjoint(; autojacvec=ZygoteVJP()),
    tspan=(0.0f0, 1.0f0),
    saveat=[])
    return ControlledODE(odefunc, solver, sensealg, tspan, saveat)
end

function (c::ControlledODE)(u0, control, p; tspan=c.tspan, saveat=c.saveat)
    function ode(du, u, _p, t)
        c.odefunc(du, u, _p, t)
        du[length(du)] += control(t)
    end
    prob = ODEProblem(ode, u0, tspan, p)
    return solve(prob, c.solver; sensealg=c.sensealg, saveat=saveat)
end

Flux.@functor ControlledODE

# Lotka Volterra with control input
tspan = (0.0, 10.0)
Nt = 100
tsteps = LinRange(tspan[1], tspan[2], Nt)

v = sin.(pi.*tsteps)
control = linear_interpolation(tsteps, v)

p = Float64[2.2, 1.0, 2.0, 0.4]
function lotka_volterra(du, u, p, t)
    x, y = u
    α, β, δ, γ = p
    du[1] = (α - β*y)*x
    du[2] = (δ*x - γ)*y
end

u0 = Float64[1.0, 1.0]

sensealg = SciMLSensitivity.InterpolatingAdjoint(; autojacvec=ZygoteVJP())
model = ControlledODE(lotka_volterra; solver=Tsit5(), sensealg=sensealg, 
                      tspan=tspan, saveat=tsteps)

# Cost for simulation
function cost(p)
    return sum(abs2, model(u0, control, p))
end

# Compute the gradient
Flux.gradient(_p -> cost(_p), p)

@profilehtml Flux.gradient(_p -> cost(_p), p)

```