Forward sensitivity and adjoints for DDEs

[devmotion]

For simple cases we can just differentiate through the DDE solver but if the DDE system contains parameter-dependent C1-discontinuities the forward sensitivities have jump discontinuities which, e.g., ForwardDiff can't deal with (see SciML/DelayDiffEq.jl#183 and https://epubs.siam.org/doi/abs/10.1137/100814949).

Hence it would be great if we could add forward sensitivity equations and adjoint methods for DDEs. I would like to help with getting them in here, I've just never worked on DiffEqSensitivity and hence might need some guidance and/or time to get started. I would assume that (hopefully) one can exploit the existing implementations for ODEs.

[devmotion]

http://www.cs.toronto.edu/pub/reports/na/adjoint_paperNSERC.pdf

[ChrisRackauckas]

http://www.naturalspublishing.com/files/published/4av943i134gz84.pdf is a good resource.

[andrschl]

I would have a first implementation of an interpolating adjoint method for DDEs with constant delays. The code can be found here:
https://github.com/andrschl/dde_adjoint_method

However, I am sure the structure and efficiency could be improved a lot. Also, I am new to julia and DiffEqSensitivity and I am not sure how this would fit best into the existing sensitivity framework.

As described on the github repo, there is also quite some difference between my sensitivity method and the AD method as soon as there are any delays. So I am not sure, whether this is still a bug in my code or because of the issues mentioned above.

[andrschl]

@ChrisRackauckas, @devmotion Apparently, performance and memory usage of my approach become terrible for increasing number of sample times. I assume this is because of the adjoint state history which I constructed as a dictionary of dense solutions of previous steps. And as I saw the solution itself contains the initial history. So I am essentially storing all solutions. Any suggestions how to do this on a lower level and more efficiently?

The second point which I am not sure about are the discontinuities. Currently I account for discontinuities in the adjoint state through passing them to the solver by constant_lags. I assume it would be better to do this directly, but I am struggling a bit to find the related code in the library..

Would it be an option to use DiscreteCallbacks for C⁰ and C¹ discontinuities and to use a single solve(...) command? This would require, that the method of steps can handle jumps in the history function.

[devmotion]

Yes, as mentioned in the initial discussion in DelayDiffEq, IMO we should just use callbacks. The solver handles discontinuities fine, a common example in the tests is actually a discontinuity at the initial time point which is then propagated to discontinuities in higher order derivatives (or derivatives of the same order in the case of neutral DDEs) which have to be considered by the solver as well.

[andrschl]

Ok, great. This should simplify things a lot and then we can also add discontinuities coming from non-smoothness at the initial point.

[andrschl]

Adding the callbacks improved things a lot :D The code looks much simpler and runtime and memory usage are better than for ReverseDiffAdjoint() in my example.

I pushed the code onto my repo. Currently I am ignoring the discontinuities due to the initial condition, but it shouldn't be hard to add them.

Also I feel like my adjoint function has already a similar form as adjoint_sensitivities(sol,alg,dg,ts;sensealg=InterpolatingAdjoint()), but I would need some guidance to properly include it into the DiffEqSensitivity framework :)

[ChrisRackauckas]

So what needs to be done to get this into the DiffEqSensitivity framework? The dispatches currently live here:

https://github.com/SciML/DiffEqSensitivity.jl/blob/master/src/local_sensitivity/interpolating_adjoint.jl

[andrschl]

Thanks for the reference to the code and sorry for my late reply. I am currently adding some things such as a continuous loss with delays. I will give it a try to include it as soon as I am done and find some time. Would you prefer a similar struct and function as for ODEInterpolatingAdjointSensitivityFunction just for DDEs (I did not look at SDDEs)?

[ChrisRackauckas]

Would you prefer a similar struct and function as for ODEInterpolatingAdjointSensitivityFunction just for DDEs (I did not look at SDDEs)?

I think that might be what's needed. Don't worry about SDDEs: I think from the SDE adjoint work it's clear that "SDDEs should just work", but it'll probably take like 5 years before someone has a clear proof that it actually does work, so you might as well just make sure the current SDE handling code works and let us address it in the future.

[andrschl]

Ok, yeah makes sense. However, I am still not sure whether the DDE adjoint is actually correct. In my tests, there is an average mismatch compared to the ReverseDiffAdjoint of 1-2%. In some gradient components it is sometimes much larger. And my test case does not include any parameter dependent discontinuities. Without delays the mismatch is much smaller (approx. 0.1%).
@devmotion @ChrisRackauckas do you know of any existing implementations to compare it to? I only found this C++ library for forward sensitivities.

Also for the discontinuities I am currently using the following two callbacks:

    # callback for adjoint jumps
    a_jumps = ξ[2:end-1]
    a_jump_idx = 1
    affect1!(integrator) = begin
        a_jump_idx += 1
        integrator.u[1:data_dim] += dl(sol(T - integrator.t), data[:, N-a_jump_idx])
    end
    cb1 = PresetTimeCallback(a_jumps,affect1!, save_positions=(true,true))

    # callback for adjoint C¹,C² discontinuities
    higher_order_disc = get_higher_order_disc(ξ, lags, order=order)
    affect2!(integrator) = nothing
    cb2 = PresetTimeCallback(higher_order_disc,affect2!, save_positions=(false,true))

    cb = CallbackSet(cb1,cb2)

Where cb1 is for the adjoint jumps in the presence of a discrete loss and cb2 is supposed to handle the other discontinuities. Is it enough to do affect2!(integrator) = nothing for the higher order discontinuities?

[devmotion]

I'm sorry, I can't focus on this right now and I would have to look up the adjoint equations first.

However, in my opinion, the major question when implementing the forward sensitivities equations is what would be a good interface and API for all the derivatives that are required here. As soon as one has decided on an API for all the delayed Jacobians, derivatives, Jacobians, and parameter derivatives of the delayed arguments that appear in the extended system as shown in eq 5 in the SIAM paper linked above (they are not part of the standard DiffEqFunctions API), the implementation of the extended DDE system should hopefully be quite straightforward.

[andrschl]

Thanks for your answer and sorry for my late one. Well, for the gradients I just defined the functions,

# we need the functions fi: xi->f(x0,...,xi,...,xk)
    ndelays=length(lags)
    fs = []
    for i in 0:ndelays
        g = function (x, xt, p)
            y = vcat(xt[1:data_dim*i],x,xt[(i+1)*data_dim+1:end])
            return f(y, p)
        end
        push!(fs, g)
    end

and then calculated the VJP in the adjoint DDE function. The same could also be used for forward sensitivities.

[devmotion]

My point was that it is not possible to specify the Jacobians manually 🤷 Similar to the ODE case, there should be an interface that allows users to specify all these required derivatives and Jacobians, and would use AD in case they are not provided. We don't want to hardcode AD, and in particular not of a particular AD system (your code always uses Zygote it seems).