Sensitivities with respect to initial conditions in DifferentialEquations.jl

Specific Domains Numerics
1

For my research, I perform many differential corrections processes. Often times, for these processes I want the transition matrix, \Phi(t_1, t_0), which is obtained by propagating the linearization of the system alongside the nonlinear differential equations such that the following approximation can be made:

\mathbf{\delta x}(t_1) = \Phi(t_1, t_0) \mathbf{\delta x}(t_0)

where \mathbf{\delta x}(t) is the variation with respect to the reference solution at time t. The linear differential equations are obviously

\dot{\Phi}(t, t_0) = A(t)\Phi(t, t_0)

where A(t) is the jacobian of the nonlinear sytem at time t.

Now to the actual question, I see that there is infrastructure in DifferentialEquations.jl for finding sensitivity of the state to the parameters, but is there any ergonomic way to find this transition matrix given the equations of motion and Jacobian without gluing them together in some function? My goal is being able to do this in general for many different systems which only need to provide their EOM and Jacobian definitions. A contrived small example would be


abstract type AbstractSystem end

struct SystemA <: AbstractSystem
    alpha::Float64
end

struct SystemB <: AbstractSystem
    beta::Float64
end

function eom!(du, u, p::SystemA, t)
    du[1] = u[1] + exp(u[2])
    du[2] = -p.alpha* u[2]
end

function jac!(J, u, p::SystemA, t)
    J[1, 1] = 1
    J[1, 2] = exp(u[2])
    J[2, 1] = 0
    J[2, 2] = -p.alpha
end

function eom!(du, u, p::SystemB, t)
    du[1] = u[2] - u[2]^3 - p.beta* u[1]
    du[2] = u[1] - u[2] - u[1] * u[2]
end

function jac!(J, u, p::SystemB, t)
    J[1,1] = -p.beta
    J[1,2] = 1 - 3 * u[2]^2
    J[2,1] = 1 - u[2]
    J[2,2] = -1 - u[1]
end

function stm(u0, system::T, tspan) where {T <: AbstractSystem}
    # Do something here to get the `sol` as well as the 
    # transition matrices over timespan.
end
1 Like
2

The best way is to probably follow this http://docs.juliadiffeq.org/latest/analysis/sensitivity.html#Examples-using-ForwardDiff.jl-1 except put partials on u0 instead of p.

3

Ok, thanks. Do you forsee this ever being incorporated into DifferentialEquations or is the use case not common enough?

4

it’s in there, but it’s just documented poorly. adjoint_sensitivities_u0 exists too which gives the sensitivites w.r.t. p and u0, but I need to find time to add docs.

Would you mind adding a part to the ForwardDiff example showing how to calculate the forward sensitivities using dual numbers? And add adjoint_sensitivities_u0 which just returns the tuple (du0,dp)? The signature is here and the same as the other adjoints: https://github.com/JuliaDiffEq/DiffEqSensitivity.jl/blob/master/src/adjoint_sensitivity.jl#L325-L336 . Going to catch a flight though but there’s both forward and backwards mode, so hopefully these notes are parsable.

5

Or if you can’t add the docs on this, ping me in like a day so I get a reminder

6

Yeah, I will see if I can get an example up and running.

7

You can do this with jet transport, eg using TaylorIntegration.jl. Cc @PerezHz, @lbenet

1 Like
8

Hello,

I am interested as well in computing the Jacobian of the solution with respect to initial conditions. i.e the matrix dx(t=t)/dx(t =0).

I have found this example from the documentation https://docs.juliadiffeq.org/stable/analysis/sensitivity/
However the sensitivity analysis is only done with respect to the parameter p, how can I modify this code to construct the jacobian with respect to u0?

Another question, I will be using interpolated function for the RHS of the ODE. Since Interpolations.jl and ForwardDiff.jl are not compatible, how can I go around this problem?

function f(du,u,p,t)
  du[1] = dx = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = dy = -p[3]*u[2] + u[1]*u[2]
end

u0 = [1.0;1.0]
p = [1.5,1.0,3.0]
prob = ODEForwardSensitivityProblem(f,u0,(0.0,10.0),p)

sol = solve(prob,DP8())
x,dp = extract_local_sensitivities(sol)

Best,

1 Like
9

It’s shown here: https://docs.juliadiffeq.org/dev/analysis/sensitivity/#concrete_solve-Examples-1

1 Like
10

That example provides the sensitivity of a scalar valued function of the states, \frac{\partial f(t_f, \bar{x}_f)}{\partial \bar{x}_0} and \frac{\partial f(t_f, \bar{x}_f)}{\partial \bar{p}}, however I can not get it to provide the sensitivities of all of the final states with respect to the initial states, i.e., \frac{\partial \bar{x}_f}{\partial \bar{x}_0} – more generally, \frac{\partial \bar{f}(t_f, \bar{x}_f)}{\partial \bar{x}_0}. After reading through all of the Zygote.jl documentation I am unsure as to whether this is possible. However, clearly, to determine the sensitivities of the function f(t_f, \bar{x}_f), these must be known by the library, correct?

11

I don’t understand the question. You noticed that using concrete_solve where you calculate only the end will give df(xf)/dx0?

12

The Zygote.gradient function does will not provide the sensitivities for a vector valued f. For example, using the same problem on the linked page:

using DiffEqSensitivity, OrdinaryDiffEq, Zygote

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 = concrete_solve(prob,Tsit5())

Attempting to get the partials of the final state with respect to the initial state from the Zygote.gradient function yields the following error:

du01,dp1 = Zygote.gradient((u0,p)->last(concrete_solve(prob,Tsit5(),u0,p,saveat=0.1,sensealg=QuadratureAdjoint())),u0,p)
# ERROR: LoadError: Output should be scalar; gradients are not defined for output [1.0337581393337192, 0.9063703701433584]

This makes sense given the function is named gradient and what I desire is a Jacobian.

13

Instead of gradient use pullback to build the whole Jacobian:

out = Zygote.pullback((u0,p)->Array(concrete_solve(prob,Tsit5(),u0,p,saveat=0.1,sensealg=QuadratureAdjoint()))[:,end],u0,p)
[out[2]([1.0,0.0])[2] out[2]([0.0,1.0])[2]]

4×2 Array{Float64,2}:
 2.19418   -6.20335 
 0.199812  -0.703454
 0.575124  -1.70302 
 0.946907  -2.71072
2 Likes
14

Was this problem addressed satisfactorily? I.e., Jacobian w.r.t. initial condition u0, at each time step, in addition to that for parameter p, as provided by ODEForwardSensitivityProblem. I’m looking at the suggested documentation, but missing it. I do see how to obtain it solely for the integration endpoint.

15

What’s the question?

16

Can one obtain the derivative of the ODE solution, at each time step, w.r.t. both the initial condition u0 and the parameter p?

17

I’d just do it the same old way as always:

function f(theta)
  _prob = remake(prob,p=theta[1:n],u0=[n+1:end])
  solve(_prob,alg)
end
ForwardDiff.jacobian(f,x)
18

Thank you for the reply.