Hello dear community, I have a question regarding the following code I have written. I am trying to calculate the adjoint of the simple (linear) systems forward simulation with respect to the input u.
using LinearAlgebra
using DifferentialEquations
using SciMLSensitivity
using Interpolations
using Zygote
using ForwardDiff
using Plots
T = 4.0
N = 40
dt = T / N
np = N + 1
A = [0 1 0 0; 0 0 1 0; 0 0 0 1;-2 0 1 0]
B = [0,0,0,1]
uk = rand(np)
function f(x,u)
return A*x + B*u
end
function S_x(u)
function forward(x, p, t)
return f(x,p(t))
end
u0 = [1.0, 0.0, -1.0, 0.0]
tspan = (0.0, T)
prob = ODEProblem(forward, u0, tspan, linear_interpolation(0.0:dt:T,u))
sol = solve(prob,Midpoint(),saveat = T/N)
return sol.u
end
x_test = S_x(uk)
Zygote.jacobian(S_x, uk)
ForwardDiff.jacobian(S_x, uk)
Now I can evaluate the solution easily, however calculating the derivative w.r.t. u does not work with neither Zygote.jl nor Enzyme.jl. Zygote yields
ERROR: MethodError: no method matching zero(::Type{Vector{Float64}})
ForwardDiff.jl works but yields weird results, a huge matrix filled with entries of type
Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(S_x), Float64}, Float64, 11}}
Maybe there is a better way for these kinds of problems, since my code is very naive right now.
Hi,
It is a bit hard to understand what you are trying to do there. The type to the result from ForwardDiff is correct. This weird type is what calculates the derivatives. Guess zygote is not compatible with your linear interpolation function.
I dont think you need to compute the adjoint by hand unless you want to deeply understand it. the solution should be differentiable (ScimlSensitivity documentation has more information on that). There are some lectures about adjoints in the SciML ecosystem around there, I can’t search it now, but you should take a look.
I will explain what I am trying to achieve. The function S_x computes the solution of the differential equation \dot x = Ax + Bu, ~~ x(0) = x_0
which is a linear differential equation for now but could also be nonlinear.
S_x takes as arguent u, which is the control function (a vector). I want to compute the derivative with respect to u. Now this can be easily writen down and calculated, even analytically utilizing the adjoint system, however I want to use automatic differentiation.
The result still seems odd to me, I need the matrix with the floating point values not the Dual type. This example is actually quite similar to the one in the SciMLSensitivitiy documentation, just that I use interpolation of the input paramters. I think Interpolation is not the problem, I can easily differentiate thorugh interpolations and it gives matrices with numeric floating point values.
I’m not sure if this is going to work. p is expected to be a SciMLStructure, generally an Array. A linear_interpolation from Interpolations.jl is not an Array.
Instead you’d need to do prob = ODEProblem(forward, u0, tspan, u) and
function forward(x, p, t)
return f(x,linear_interpolation(0.0:dt:T,p)(t))
end
should work. Check that first. Caching that computation can come second.
Hello Chris, thank you for your answer.
The code changes did not result in a solution, the error messages still persist for Zygote and ForwardDiff calculates weird results. Here is the part of the changed code
uk = rand(np)#zeros(np)
function f(x,u)
return A*x + B*u
end
function S_x(u)
function forward(x, p, t)
return f(x,linear_interpolation(0.0:dt:T,p)(t))
end
u0 = [1.0, 0.0, -1.0, 0.0]
tspan = (0.0, T)
prob = ODEProblem(forward, u0, tspan, u)
sol = solve(prob,Midpoint(),saveat = T/N)
return sol.u
end
S_x(uk) #works fine
ForwardDiff.jacobian(S_x,uk) #weird result
Zygote.jacobian(S_x,uk) #error
The output for ForwardDiff is still the arrangement of type
You were just returning a vector of vectors. AD libraries do not know how to represent Jacobians of such outputs. Instead, just return a normal array and it’s fine:
using LinearAlgebra
using OrdinaryDiffEq
using SciMLSensitivity
using Interpolations
using Zygote
using ForwardDiff
using Plots
T = 4.0
N = 40
dt = T / N
np = N + 1
A = [0 1 0 0; 0 0 1 0; 0 0 0 1;-2 0 1 0]
B = [0,0,0,1]
uk = rand(np)#zeros(np)
function f(x,u)
return A*x + B*u
end
function S_x(u)
function forward(x, p, t)
return f(x,linear_interpolation(0.0:dt:T,p)(t))
end
u0 = [1.0, 0.0, -1.0, 0.0]
tspan = (0.0, T)
prob = ODEProblem(forward, u0, tspan, u)
sol = solve(prob,Midpoint(),saveat = T/N)
return Array(sol)
end
S_x(uk)
ForwardDiff.jacobian(S_x,uk)
Zygote.jacobian(S_x,uk)