Add more statistics and move jitter to transition

[sethaxen]

See #123 for discussion

Adds the following statistics:

• hamiltonian_energy_error
• max_hamiltonian_energy_error
• is_adapt
• nom_step_size

To-do:

• test that stats are returned, in particular that is_adapt starts as true and switches to false after adaptation is over
• refactor code around jitter so that the step size used is accessible for stats. I'd appreciate input on this.

I'll submit a PR to Turing so that the new stats fields show up in internals instead of parameters. Looks like Turing doesn't use AdvancedHMC's sample function, so I'll need to add an is_adapt over there.

[sethaxen]

I'm not certain why Travis isn't detecting this PR

[xukai92]

Just see if reopening it will trigger Travis.

[sethaxen]

Thanks! Any advice on how to handle jitter?

[sethaxen]

What if the integrator steps returned a Transition instead of a PhasePoint? That would enable other experimental integrators to also return stats we don't know about.

[xukai92]

I prefer not because Transition is designed to store a transition from a MCMC kernel.

How about the following

1. define stats(::AbstractLeapfrog) = ... （similar to L10 and L11 for jitter and temper)
   • This can then be used in transition to grab info from integrator
2. add another filed ϵ0 to JitteredLeapfrog and make it mutable
3. change jitter to jitter! to "refresh" ϵ using ϵ0

How do you think?

[sethaxen]

That's a convenient way to do it. I implemented it.

I misunderstood what PyMC3's step_size_bar is. It's an internal parameter during the adaptation phase. step_size in PyMC3 and Stan indicates the nominal (unjittered) step size, and they don't output the jittered one. Still, I think that would be useful to have in stats. I think I'll keep the nominal step size as step_size and for JitteredLeapfrog only output step_size_jittered.

Adaptation uses the nominal step size, so I'll also add a nominal_step_size method. It also seems that a common pattern is NesterovDualAveraging(0.8, integrator.ϵ). I'm also planning to add something like

NesterovDualAveraging(δ::AbstractFloat, i::AbstractIntegrator) = NesterovDualAveraging(δ, nominal_step_size(i))

[xukai92]

Adaptation uses the nominal step size, so I'll also add a nominal_step_size method.

Can you explain a bit more? I don't quite catch this point.

It also seems that a common pattern is NesterovDualAveraging(0.8, integrator.ϵ). I'm also planning to add something like

Sounds good to me.

[sethaxen]

Adaptation uses the nominal step size, so I'll also add a nominal_step_size method.

Can you explain a bit more? I don't quite catch this point.

So when one sets the step size, it's really setting a "nominal" step size, since if one jitters for example, the step size can change from step to step. It's the nominal step size that should get adapted (this is how Stan does it). You'd also want to use the nominal step size for HMCDA. If there's a more descriptive term than "nominal", I'm happy to use it; that's just the term Stan uses.

[sethaxen]

One more point tangential to this PR is that the way Stan and (I think) PyMC3 implement HMC, the step size is only jittered at the beginning of the transition, not at every step in the transition. As currently implemented, JitteredLeapfrog is going to jitter every step.

[xukai92]

One more point tangential to this PR is that the way Stan and (I think) PyMC3 implement HMC, the step size is only jittered at the beginning of the transition, not at every step in the transition. As currently implemented, JitteredLeapfrog is going to jitter every step.

You are right. It's the case for NUTS but HMC and HMCDA are fine as there is only a single call for the integrator. But let's also do jitter one time for a transition by moving it to the beginning of transition. It's fine now as we only have leapfrog integrator and we can worry about generalise it later.

Also I think what we could improve is that we don't need make JitteredLeapfrog mutable anymore because for jitter we can just return a new leapfrog and use that within the transition function. So it's acutally has a nice side effect :)

[sethaxen]

But let's also do jitter one time for a transition by moving it to the beginning of transition.

Also I think what we could improve is that we don't need make JitteredLeapfrog mutable anymore because for jitter we can just return a new leapfrog and use that within the transition function. So it's acutally has a nice side effect :)

So just to be clear, you're recommending 1) making JitteredLeapfrog immutable again and renaming jitter! back to jitter 2) removing jitter from step, and 3) adding jitter to the beginning of transition, returning a new JitteredLeapfrog to be used within that transition. Did I get that right?

[xukai92]

Yes exactly!

[sethaxen]

Okay, unless there's anything else, I think this is ready to merge.

[xukai92]

PR looks good to me now!

[xukai92]

I will merge it in 2 hours if no objection.