Plan for DiffEqDiffTools

[ChrisRackauckas]

The idea for this package is to setup differentiation functions which utilize as much of the performance overloads as possible, and make it easily available for all of the native Julia solvers. The idea is pretty much understood by these comments:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/integrators/rosenbrock_integrators.jl#L24

There's no reason to do this same setup in every solver in all of its variations, and instead it needs a library to be a backend for all of the *DiffEq solvers. Essentially, the idea is to make a library here that allows someone choose between autodiff or not, and provide functions that will compute the Jacobian. But also, utilize the performance overloads:

http://docs.juliadiffeq.org/latest/features/performance_overloads.html#Other-Available-Functions-1

So if f(Val{:jac},t,u,J) exists, it will use that for the Jacobian, using the traits to check for existence:

has_jac(f)
has_expjac(f)
has_invjac(f)
has_tgrad(f)
has_hes(f)
has_invhes(f)
has_invW(f)
has_invW_t(f)
has_paramjac(f)
has_paramderiv(f)

Then, the methods will be bootstrapped on top of each other. For example, there will be a method for computing the explicit inverse Rosenbrock-W function (M - γJ)^(-1). This will check has_invW(f), and if not, will go down to computing the Jacobian, using the method which checks if has_jac(f) before finally autodifferentiating via ForwardDiff or using Calculus.jl (/FiniteDiff.jl) for when autodifferentiation won't work (some functions and types).

This will make it very easy to build most stiff solvers, since the set of functions covers the vast majority of cases. So the Rosenbrock solvers can say "give me the Rosenbrock-W inverse", and it will give the computed inverse in the manner that is most computationally efficient given what the user (or ParameterizedFunctions) has specified. Given that all of the has functions are compile-time known traits, this has 0-overhead, and so this design will be killer when done.

Lastly, it should be opened up to allow the user to choose the method by with \ will occur. This shouldn't be hard since you can just let people pass a factorization object, but it'll need some PRs to libraries like IterativeSolvers.jl to allow them to be used in that interface. See this discussion for more details:

https://discourse.julialang.org/t/unified-interface-for-linear-solving/699/32

Together, all of this forms a very self-contained project that would be great for beginners, but would make it a lot easier to implement more solver methods in the most efficient manner possible at a very high level. I think this is a perfect GSoC project.

[ChrisRackauckas]

A few extra details:

Differentiation has to occur on a single variable at a time to use the functions available from other libraries. So what can be done is the one can use call-overloading and use types to internally hold the other variables. This is shown here:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/derivative_wrappers.jl

Also, u is not assumed to be a Vector but can be an arbitrary array, so the top functions handle the resizing automatically using views. Lastly, the function f(t,u,du) writes to du and returns du, so you can effectively think about it as differentiating

(u) -> (f(t,u,du); return du)

The issue here is that du must be the right type, and that can be troublesome to get with autodifferentiation. So everything is setup to find the chunksize beforehand when that internal cache is created. The solution was first provided here:

JuliaDiff/ForwardDiff.jl#136 (comment)

and resides here:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/misc_utils.jl#L1

These all assume that the chunksize is known since the algorithms themselves set the chunksize as type information:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/algorithms.jl#L30

Those are all of the pieces for how to get it working for ForwardDiff. Something similar will need to be done for Caclulus/ReverseDiff.

[sidgupta234]

You are suggesting the user a choice to use F/B autodifferentiation or the symbolic differentiation & if the first choice fails then use the second as the fall back option. But don't the user already have a choice to choose a specific option as he can go on with either of the approach on his own, or am I getting this all wrong?

[ChrisRackauckas]

No, let me take a step back. The approach is this:

A user defines a function f(t,u,du) (or f(t,u), let's ignore that complication for now) to solve an ODE/SDE/DAE/DDE. This writes the output to du (inplace updating for better performance). If they know the analytical form for the Jacobian, they can supply a dispatch f(Val{:jac},t,u,J) which writes to J. DiffEqBase.jl is setup so that way f will automatically have a "trait": if that dispatch exists, then at compile-time has_jac(f) is evaluated to true (and other branches compile away).

(The sidenote on symbolic differentiation is that if the function is defined by the macros from ParameterizedFunctions.jl, the function f(Val{:jac},t,u,J) is determined by symbolic differentiation. Symbolic differentiation cannot be applied here in general)

While theoretically every user could use ForwardDiff/ReverseDiff/Calculus.jl to supply the Jacobian, this is an undue burden on the user. This is especially the case for autodifferentiation: give it a try and you'll see what the complications are (supplying the caches for du, overloading a type so that way you don't make a closure every iteration, etc., all the things described above). This is why a Jacobian calculation is provided as a fallback. The routine would then be something like:

caclulate_jacobian!(J,f,... # cache variables)
  if has_jac(f)
    f(Val{:jac},t,u,J)
    return J
  else
     if difftype == :forward
        # Use ForwardDiff
     elseif difftype == :reverse
        # Use ReverseDiff
     else # numerical
        # Use Calculus
     end
  end
end

The end result is that any differential equation solver could just

calculate_jacobian!(J,f,...)

without having to worry about all of the cases (or, this library could create the requested dispatches).

Of course, when you see how you have to handle the autodifferentiation, you'll see that this high level description isn't that simple (since you can't differentiating the full f and thus need some closure/type for it in the first place. There is also an issue is putting together a sleek API so that way the internal caches only need to be defined when necessary).

Then the next part is to bootstrap up from that to get the other forms like (M - γJ)^(-1) (which there also exists a dispatch for f(Val{:invW},t,u,γ,iW) since ParameterizedFunctions.jl will symbolically compute this function), using the calculate_jacobian! function.

The big thing is, if you haven't tried to autodifferentiate the inplace version f(t,u,du) efficiently, give it a try. The whole problem stems from handling that issue.

[ChrisRackauckas]

@sidgupta234 are you still interested in this?

[Ayush-iitkgp]

@ChrisRackauckas this sounds interesting to me :D and I would like to pursue it towards a GSoC project. However, I am pretty new JuliaDiffEq, it will take me some time to get used to it. Any suggestions on how to get started quickly ?

[ChrisRackauckas]

The easiest way to get started would be to get familiar with solving equations. Work through the tutorial:

http://docs.juliadiffeq.org/latest/tutorials/ode_example.html

and you'll be familiar with the interface. From there, the devdocs explain the internals:

http://devdocs.juliadiffeq.org/latest/contributing/diffeq_internals.html

It could probably use more details, so just ask if you need help. I would recommend stepping through the details once. But here's the play-by-play:

Initialization:
https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/solve.jl#L11
Setup the cache variables:
https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/caches.jl#L871
Finish the initialization:
https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/solve.jl#L226
Enter the main loop:
https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/solve.jl#L241
And then, the only part of the main loop which matters for this is the calculation of the steps:
https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/integrators/rosenbrock_integrators.jl#L10

The header/footer etc. take care of all of the saving, callbacks, adaptive stepsize, etc. so you don't have to worry about any of that.

I think the first step to tackling this would be to get the non-autodiff version working. It's commented out:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/integrators/rosenbrock_integrators.jl#L24

It probably needs to do something similar to the ForwardDiff version above it. You can see from the top of that method that the uf and tf are from the cache, so they're made here:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/caches.jl#L891

and where the function they call is from here:

https://github.com/JuliaDiffEq/OrdinaryDiffEq.jl/blob/master/src/derivative_wrappers.jl#L27

Essentially, it builds a type so that way uf(u) acts as f(t,u,du), since the calculus packages require just one argument. With this newer setup (that Caclulus.jl comment is quite old!), you may just need to re-arrange some arguments to get it working.

I think the first thing is to get these working inline in the Rosenbrock methods, then refactor them out to a package. That's a good spot for a PR, and I think you'll start to be familiar with where this is all going.