Add notebook showing sensitivity/jacobian calculations via sympy lambidy

[MichaelClerx]

@sanmitraghosh can you explain what you mean a bit more? Maybe with an example?

[sanmitraghosh]

Lets say we have an ODE dX/dt = F(X,p)

Now for calculating sensitivity we need the Jacobians dF/dp and dF/dX. Now we can calculate these quantities fast enough using sympy's lamdify function. This in my opinion is faster than using RMADs such as autgrad for calculating the Jacobians. Here is an example of this to carry out variational inference using PyTorch and Pyro PPL:

[sanmitraghosh]

I will see if I can write a PINTS example with this technique in the context of MALA/HMC

[martinjrobins]

interesting, I know proponents of autograd say that it is faster than symbolic differentiation due to equation expansion, but I've never seen anything to back this up (perhaps if you don't simplify after doing the symbolic differentiation??)

[sanmitraghosh]

@martinjrobins Sympy Lamdify is faster compared to PyTorch's autograd when used for calculating the Jacobians at each time step/solver call. Both running on CPU. But it would be good to compare speeds between sympy and Python autgrad (from HIPS).

[martinjrobins]

and JAX I think. This has an ode solver built in, and my experience of using it with numpyro was good, very fast. No idea how it would compare with sympy lamdify though

[sanmitraghosh]

@martinjrobins Jax should be best because it is the only AD software where you can mix forward mode and reverse mode. So one can use forward for Jacobians and reverse for evaluating the VJPs. But I guess a comparison would interesting.

[sanmitraghosh]

By the way @martinjrobins I have got a grad student working on changing all my current numpy stuff to Jax ! :)

[martinjrobins]

By the way @martinjrobins I have got a grad student working on changing all my current numpy stuff to Jax ! :)

I am both very impressed with JAX, and appalled at the huge amount of effort and convoluted programming that is going into speeding up a language that was designed to favor readability over speed ;)