Repair ode API after refactor broke it

[michaelosthege]

Problem

Fix #3634, which broke the ODE API. It became apparent because NUTS went nuts.

Cause

When I refactored, I changed the order of inputs to the DifferentialEquation such that y0 comes before theta. But because I neglected to account for this change also in the forward sensitivities, the resulting gradient was broken.

Solution & Changes

In this PR, I fixed the problem by:

• re-arranging the forward sensitivities such that y0 comes before theta
• making sure that y0 before theta is done everywhere

A few more improvements:

• explicitly set dtype=floatX wherever possible
• use float64 in the forward integration (augment_system), because odeint can't be set to float32
• in utils.py line 51, only the theta part of the parameters is passed to the user-defined ODE function. In the original implementation it was stack(theta, y0) so, the function actually had access to y0 which was unintended (I suppose, @Dpananos ?)
• introduce and use DtypeError where applicable
• re-run ODE example notebooks

[Dpananos]

broke the ODE API, which became apparent because NUTS went nuts.

Oh, my sides.

in utils.py line 51, only the theta part of the parameters is passed to the user-defined ODE function. In the original implementation it was stack(theta, y0) so, the function actually had access to y0 which was unintended (I suppose, @Dpananos ?)

Intended actually, but not important to the implementation. Better that you have fixed it. Thanks for all your hard work.

So does this fix the sampling issue? Can you verify that the API now samples the example notebook in a reasonable time? I can't recall how fast they were at the end of GSOC, but I think the first example sampled within 10 mins or so.

[codecov]

Codecov Report

Merging #3684 into master will decrease coverage by 0.01%.
The diff coverage is 95%.

@@            Coverage Diff             @@
##           master    #3684      +/-   ##
==========================================
- Coverage    89.9%   89.88%   -0.02%     
==========================================
  Files         134      134              
  Lines       20166    20176      +10     
==========================================
+ Hits        18130    18135       +5     
- Misses       2036     2041       +5

Impacted Files	Coverage Δ
pymc3/exceptions.py	84.21% <ø> (-15.79%)	⬇️
pymc3/ode/utils.py	100% <100%> (ø)	⬆️
pymc3/tests/test_ode.py	100% <100%> (ø)	⬆️
pymc3/ode/ode.py	94.05% <94.05%> (+0.3%)	⬆️
pymc3/step_methods/hmc/base_hmc.py	95.32% <0%> (-0.94%)	⬇️
pymc3/distributions/distribution.py	95.02% <0%> (-0.3%)	⬇️

[michaelosthege]

broke the ODE API, which became apparent because NUTS went nuts.

Oh, my sides.

in utils.py line 51, only the theta part of the parameters is passed to the user-defined ODE function. In the original implementation it was stack(theta, y0) so, the function actually had access to y0 which was unintended (I suppose, @Dpananos ?)

Intended actually, but not important to the implementation. Better that you have fixed it. Thanks for all your hard work.
So does this fix the sampling issue? Can you verify that the API now samples the example notebook in a reasonable time? I can't recall how fast they were at the end of GSOC, but I think the first example sampled within 10 mins or so.

Yes, the performance is now restored.

[ColCarroll]

Nothing that requires immediate action, but one optional change.

This looks like it was tricky to find, and seems to generally improve how debuggable the codebase is. Thanks for keeping on it!

[guilhermeprokisch]

I'm still having problems with the sample speed in the notebooks example. Is this normal?

[michaelosthege]

I'm still having problems with the sample speed in the notebooks example. Is this normal?

I can confirm that the example from the current master seems to be waay slower than it was in the run shown on docs.pymc.io.

Unfortunately I won't be able to investigate before the weekend. But my next step would be to compare the merge commit of this PR against current master & look for changes in the ode submodule.

Link to the compare tool: 5e9f349...master