Gradient of ODESystem containing a loop

Specific Domains Numerics
1

Hi! I am trying to take the gradient of an ODESystem containing a loop using either Zygote or ForwardDiff. Both seem to fail when I do so. Here is a minimal working example (taken from a previous JuliaForum post)


using Zygote, SciMLSensitivity, ForwardDiff, DifferentialEquations

# succeeds

function fiip(du,u,p,t)
    du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
    du[2] = dy = -p[3]*u[2] + p[4]*u[1]*u[2]
  end
  
  p = [1.5,1.0,3.0,1.0]; u0 = [1.0;1.0]
  prob = ODEProblem(fiip,u0,(0.0,10.0),p)
  sol = solve(prob,Tsit5())
  
  plot(sol)
  
  function sum_of_solution(x)
      _prob = remake(prob,p=x)
      sum(solve(_prob,Tsit5(),saveat=0.1))
  end
  
  ForwardDiff.gradient(sum_of_solution,p)

# fails with both Zygote and ForwardDiff

function fiip(du,u,p,t,N)
    sv = zeros(Int64(2))
    for i in 1:N
        sv[i] = u[i]
    end
    du[1] = dx = p[1]*sv[1] - p[2]*u[1]*u[2]
    du[2] = dy = -p[3]*sv[2] + p[4]*u[1]*u[2]
  end

  p = [1.5,1.0,3.0,1.0]; u0 = [1.0;1.0]

  fiip2(du,u0,p,t) = fiip(du,u0,p,t, 2)
  prob = ODEProblem(fiip2,u0,(0.0,10.0),p)
  
  sol = solve(prob,Tsit5())
  
  plot(sol)
  
  function sum_of_solution(x)
      _prob = remake(prob,p=x)
      sum(solve(_prob,Tsit5(),saveat=0.1))
  end
  
  ForwardDiff.gradient(sum_of_solution,p)

The loop really is necessary in my minimum working example, as using a state variable that is a sum of lots of the us… Any idea how to get this working?

2

With ForwardDiff, this is likely part of the problem, you are creating an array of Float64 values. Create this as

zeros(eltype(u), 2)

and then make sure that u0 has the same element type as p, i.e.,

_prob = remake(prob,p=x, u0 = eltype(p).(u0))

or something like that.

1 Like
3

Thanks so much! This works perfectly