Speed up Adjoint Sensitivities

[mjohnson541]

The current implementation of adjoint sensitivities is functional, but very poorly optimized. Currently it just evaluates the full derivative function for the domain and then extracts the one derivative it needs from that in the function g. We then calculate gradients dgdu and dgdp using forward differentiation.

In this scheme in each evaluation of g we evaluate many reactions that are irrelevant to the calculation of the derivative of the species of interest. Ideally we should determine the reactions involving the species of interest and only evaluate those reactions. Additionally it should be much faster to calculate dgdu and dgdp using reverse mode differentiation with ReverseDiff, Tracker or Zygote than forward mode differentiation.

It should also be possible to do even better than the above for specific domains by using analytic dgdu and dgdp.

[jiweiqi]

Add a reference on comparing different Adjoint sensitivity approaches in Julia eco-system https://arxiv.org/abs/1812.01892. I am actually not sure whether reducing the number of parameters affect the cost a lot. I guess the cost should be dominantly by the dimension of outputs involved in the loss function g.

[jiweiqi]

Where is the current implementation of adjoint sensitivities you are referring to? @mjohnson541

[mjohnson541]

Sorry, I still need to add it to the wiki, but the ConstantTP H2 example notebook on master has examples for both forward and adjoint sensitivities.

[mjohnson541]

The adjoint code is in Simulation.jl

[jiweiqi]

I am not sure why I can not ping Christopher Rackauckas here. If you can, it will be great to ping him here. He might have some suggestions.

[mjohnson541]

I'm not sure, I can't either, but I think for the moment the path forward is relatively clear...we have some very clear inefficiencies in the implementation that should net us >100x speed up...we can ping him once we've fixed these...

[ChrisRackauckas]

In this scheme in each evaluation of g we evaluate many reactions that are irrelevant to the calculation of the derivative of the species of interest. Ideally we should determine the reactions involving the species of interest and only evaluate those reactions. Additionally it should be much faster to calculate dgdu and dgdp using reverse mode differentiation with ReverseDiff, Tracker or Zygote than forward mode differentiation.

Oof that's not good. Is the issue in the vjps, i.e. #84?

[mjohnson541]

So I believe these two bits are separate from #84. The first bit is just a programming problem, just need to make a smaller phase object specifically for calculating g. The second, I've already had some success with dgdp switching to ReverseDiff was ~100x faster (without even fixing the issue with g), I'm working the infrastructure to run ReverseDiff for dgdu properly, but it seems straightforward and I think it should be cheaper than dgdp. Put together I anticipate a lot of speed up once I have it working.

[mjohnson541]

@ChrisRackauckas @hwpang has gotten ReverseDiff for the smaller phase object mostly working non-preallocating for dgdu and dgdp, we're checking some aspects. But we're a bit at a lost of how to use caching to preallocate for ReverseDiff. As you noted earlier caching might also help avoid allocations in the main function evaluations. Are there good templates for how to do this generally and for ReverseDiff?