Track model logp (in HMC and NUTS)

[eigenfoo]

Close #2971.

WIP. Following up from #3121. I'm opening up this PR so I can get Travis builds/tests/feedback from maintainers 😄

@ColCarroll I've done a first pass at tracking the logp. The sampler stats stack actually goes deeper than you outlined before, since the model logp is actually kept in integration.py for HMC. What I've done is modify the State so it keeps track of the model logp, and percolate that up as a sampler stat to HMC and NUTS.

@junpenglao currently this PR falls a bit short of what you were expecting; I've only made the model logp a sampling stat for HMC and NUTS. If I'm not mistaken, no other sampler computes the logp (e.g. Metropolis doesn't), so there would be overhead to require these samplers to track the logp. However, if this is a feature the dev team wants, I'd be happy to do that too - probably in a separate PR, though.

I'm also unsure about the compoundstep changes... states is a list, and cannot be indexed by a string. I must have misunderstood @junpenglao's comments here: I'll take another look sometime in the next few days. For now, tests will fail for that.

[junpenglao]

If I'm not mistaken, no other sampler computes the logp (e.g. Metropolis doesn't), so there would be overhead to require these samplers to track the logp. However, if this is a feature the dev team wants, I'd be happy to do that too - probably in a separate PR, though.

Oh you are right - in Metropolis we only computed the delta_logp. Hmmm, it does complicate things...

[eigenfoo]

This seems to pass tests now. The test it fails is due to SMC and appears unrelated. Ready for review!

What I've done for CompoundStep is to remove logp items from all but the last state in states. In other words, only the last state can be the model logp. Sometimes, the last state may not have a logp at all (eg if the last sampler is Metropolis), but if the last state has a logp, then it is the model logp. Does that sound reasonable?

[junpenglao]

Need a new test to check if the logp is saved correctly at least.

Also, add a line to release-note.

I think this is a good first step. Maybe in another PR we can also add it to other samplers. I have an idea to do this also for Metropolis where we are computing the delta_logp instead of the logp. We can compute the model logp at sample 0, then from sample 1 onward we do state.logp = logp_tm1 - delta_logp which should recover the logp at time t while also cheap to compute.
Careful tests are of course needed to make sure it is correct especially for CompoundSteps.

[twiecki]

This is more maintenance.

[eigenfoo]

@junpenglao I'm actually not sure where the tests for sampler stats are. I would've thought test_stats.py, but that's actually to test the pm.stats module... Any pointers?

[junpenglao]

pymc3/tests/test_step.py

[eigenfoo]

Sorry for the lull, I was taking a break from my laptop!

I've added a test to assert the existence of various sampler stats for NUTS, and also assert that their shapes are correct. Let's see if this passes tests.

[eigenfoo]

Passes tests! @junpenglao ready for review

[junpenglao]

Maybe we should rename it to lp__, see arviz-devs/arviz#176 (comment)

[eigenfoo]

Hm, my concern is how intuitive a name lp__ is... I understand Stan has it as lp__ but with Python there is the convention of having self-descriptive names, even if this means verbosity. I'd be fairly confused at both what lp stands for, and why there are two underscores after it (AFAIK nothing else in PyMC3 has two underscores after it...)

I don't really know much about Arviz though. Is "compatibility" with Stan/Arviz a high priority for the PyMC devs?

[twiecki]

I agree with @eigenfoo, lp__ is not a great name. Compatibility to Arviz is important, but not to Stan. But since we also have control over arviz I would suggest standardizing on a reasonable name choice there (like logp) and here.

[junpenglao]

In that case, either logp or model_logp would be a good name choice.
FYI @eigenfoo, in PyMC3 the transformed parameters all have __ which is generated internally. These RVs are meant to be "hidden" and we use the double underscore to identify these parameters and hide them when generating summary and figure.

[ColCarroll]

Yeah, agree on using a descriptive name - I think arviz is going to keep trying to decide on a "standard" name for sampler statistics, which may change over time (I'm still in favor of a descriptive name there as well, but it should be a friendly spot for stan users too). So no worries on what's going on there :)

[eigenfoo]

@junpenglao I think I like model_logp better - it's more explicit, and harder to confuse it with all the other log-probabilities that you could talk about 🙂

I'll need some help with the test though - see my comment above.

[springcoil]

Awesome work @eigenfoo

[aseyboldt]

Maybe better sample_logp? It isn't really a property of the model, and it is different for each sample. Or maybe posterior_logp? Otherwise I think this looks good.

[fonnesbeck]

It gets a little semantic, I suppose. Its the log-probability of the model at a particular point, so I don't think model_logp is necessarily bad. The sample doesn't have a logp without a model!

[aseyboldt]

How about value_that_is_proportional_to_the_posterior_logp_given_model_and_dataset_at_the_point_of_that_sample. We have tab completion after all. :-)
Edit: I forgot the parametrization. It depends on that as well.

[aseyboldt]

But I guess you have a point. They all probably work. (except that last one)

[ColCarroll]

This looks good to me (and something I will use right away) - thanks, @eigenfoo! Is there something specific you would still like review on?

[eigenfoo]

For the record, the problem that stalled me up was this:

The way CompoundStep works is essentially a glorified for-loop over a list of step methods. This has various implications for sampler stats, since each step method in this list may or may not support stats, and may support different stats, etc. etc. According to these docs, we should be returning every single stat generated by every single sampler. If there happen to be name collisions, we stick the numpy arrays together, along a new axis.

The "bad thing" about this PR is best explained by example:

with pm.Model():
    x = pm.Normal('x')
    y = pm.Normal('y')
    step1 = pm.NUTS([mu1])
    step2 = pm.NUTS([mu2])
    trace = pm.sample(step=[step1, step2])

Here, trace.get_sampler_stats('model_logp').shape is (1000, 2). This may be counter-intuitive, since there should only be one model logp, so it may make sense to insist that model_logp is always a 1-D numpy array.

Eventually, @ColCarroll and I decided that it's better to keep this quirk and document it, rather than having an inconsistent API. This means that if you wanted the overall logp for with a CompoundStep, you'd need trace.get_sampler_stats('model_logp')[:, -1], and if you (for whatever reason) wanted one of the intermediate logp's, you'd need to index one of the columns of trace.get_sampler_stats('model_logp').

I think we're going to go ahead and merge this once tests pass.

[twiecki]

Hm yeah that's unfortunate but I agree with your assessment.

[aseyboldt]

Thank @eigenfoo!

[ColCarroll]

Thanks for this!