Sensitivities with respect to initial conditions in DifferentialEquations.jl

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

struct SystemB <: AbstractSystem

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

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

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]

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]

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

1 Like

The best way is to probably follow this except put partials on u0 instead of p.

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

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: . Going to catch a flight though but there’s both forward and backwards mode, so hopefully these notes are parsable.

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

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

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