Add more info to divergence warnings

[aseyboldt]

This is a patch to store some additional information in the internal warnings about divergences. I used this a couple of times for debugging, and I think it might be of some use for other people as well.
Before, we only stored where the step that diverged started. This patch also adds what the energy, logp, gradient and velocity were where it started, and where it ended up.
I also converted the report.SamplerWarning to a dataclass to make it easier to extend. This requires >=py3.7.

This makes it possible to make interesting plots of where divergences happend:

import pymc3 as pm
import matplotlib.pyplot as plt
import numpy as np

with pm.Model() as model:
    sd = pm.HalfNormal('sd')
    pm.Normal('x', sd=sd)
    
    trace = pm.sample(Emax=1000)

x = 'sd_log__'
y = 'x'
plt.scatter(trace[x], trace[y], marker='.')
for warn in trace.report._warnings:
    if warn.divergence_point_dest is not None:
        source = warn.divergence_point_source
        dest = warn.divergence_point_dest
        source = np.array([source[x], source[y]])
        dest = np.array([dest[x], dest[y]])
        plt.arrow(*source, *(dest - source) / 100)

(Maybe Emax=1000 is a bit high as a default by the way?)

[AlexAndorra]

Loving this @aseyboldt !! Please ping me when it's ready for review 😉

[aseyboldt]

@AlexAndorra @ColCarroll This is ready for review now.
I included an extra dependency for dataclasses so that tests pass on py36. We can get rid of that when we drop py36.
I didn't include it in the release notes yet, because I think we need to figure out if/how we want to use this for things except debugging.
I think it would be nice if arviz could make use of this to show more precise info about divergences at some point.

[AlexAndorra]

Thanks @aseyboldt, LGTM! I don't know much about dataclasses but that looks neat.
I think this should definitely be used by ArviZ down the road -- probably in plot_pair? (cc-ing @OriolAbril about that).
In the meantime, I think it makes sense to mention this in the release notes, so that people can use it for manual debugging, as you did in your example.
Waiting for @ColCarroll's approval to merge 😉

[OriolAbril]

Hey, are we planning on following NEP 29? If so there is no need to keep the changes compatible with 3.6 as next release should drop 3.6 support https://twitter.com/Mbussonn/status/1272667933332758528?s=20

Plot looks great, thanks @aseyboldt, is this covered on the divergences notebook? Maybe it would be useful to have a comparison and guidance on where this info can have an edge wrt "regular" divergences

[AlexAndorra]

Adrian did write a PR yesterday to drop py36, but we don't know yet when we'll merge it (just added that point to today's lab meeting 😉 ). Let's wait for that decision then before merging.

is this covered on the divergences notebook? Maybe it would be useful to have a comparison and guidance on where this info can have an edge wrt "regular" divergences

Are you talking about the "Diagnosing divergences" NB? @rpgoldman is updating it, so maybe Adrian should wait until this PR is merged to add the example? Very good idea though: would be super helpful to show a comparison and explain why / when this plot would be useful 👌

[aseyboldt]

I wouldn't know what to write in the notebook for now to tell you the truth. I just don't know if, why or when this plot would be useful.
The reason I would still like to have this info in the warnings isn't so much that it allows us to make that plot, but because it has turned out to be useful for debugging the bug where the step size was nan, and also a couple of issues in sunode. I think it might turn out to be interesting for diagnosing divergences in addition to that, but I don't know.

[rpgoldman]

@aseyboldt I'd be happy to discuss this with you and see about getting it into the revised notebook.

[AlexAndorra]

Makes sense @aseyboldt. We decided to drop py36 in the lab meeting, so we can do that:

• I'll merge your PR on py36
• You can rebase this branch and remove the dataclasses requirement from this PR
• I'll merge once everything is green
  Sounds good?

[AlexAndorra]

LGTM, thanks Adrian!

[AlexAndorra]

@aseyboldt, since we're finally not dropping py36 (#4003), does this need an extra dependency for dataclasses so that tests pass on py36?
If yes, I think I can revert the merge here, so that you can continue working from this branch (essentially reverting your last commit)? Sorry for all the back and forth...

[junpenglao]

@AlexAndorra I already added the new dependency in the PR.

[AlexAndorra]

Ah true, didn't see that, thanks @junpenglao !