Add Interface / example notebook with numpyro

[ben18785]

Numpyro is a lightweight probabilistic programming language allowing automatic differentiation via Jax and JIT compilation to GPU/TPU/CPU.

We would like to potentially build an interface to numpyro but, first, it'd be good to see if we could interface with their models and run our inference algorithms on them. The following then seem like a reasonable order of things to try:

• Code up basic (i.e. non-ODE) model in Pyro; figure out how to extract the log-probability and sensitivities of these models (@martinjrobins any ideas here? I couldn't find this when I looked). Then run a Pints sampler on it.
• Code up a ODE model in Pyro and do the same as above.
• Compare numPyro's NUTS with Pints' NUTS.

[martinjrobins]

to extract the log-likelihood you make use of numpyro's effects handlers, see an example here http://num.pyro.ai/en/stable/handlers.html

The effect handlers are really cool, they allow users to implement custom side effects for the primitive operations (i.e. numpyro's sample or param). So you can set all the free parameters in a function, extract a trace, or list of all the sample operations (which define the priors and measurement model) and then sum up the log prob of each to arrive at a total log probability.

[ben18785]

@martinjrobins thanks. So I'm guessing getting the gradients should also be possible? (I can see how to get the log_prob now.)

[martinjrobins]

I think getting the gradient is what jax does (https://github.com/google/jax). Once you have a function which takes in a parameter vector and calculates the log probability, jax can automatically calculate the gradient of that function (since the ode integration is done in Jax, it knows how to calculate its gradient automatically)

[DavAug]

I'll give it a go.

[DavAug]

Where should this notebook go? Is there a directory for interfacing examples? Where can I find the stan interface for example?

[ben18785]

The Stan interface is currently sitting in a PR -- #1077. I reckon start with a folder called "interfaces" within the examples folder and put an example notebook there illustrating how it works.

[MichaelClerx]

Can we discuss again whether all these interfaces should live in the main repo or just have a repo of their own? There's a million cool packages out there if we add them all to PINTS, even as optional extras, our testing will start taking forever

[ben18785]

I’d vote for having them in the same repo. Keeps things simple and stops users having to look in lots of places to find out how things with. Actually, is there a way to have a vote on this (then the issue is addressed and we can move on it either way)? Does Github allow a poll?

…

On 6 May 2020, at 20:42, Michael Clerx ***@***.***> wrote: Can we discuss again whether all these interfaces should live in the main repo or just have a repo of their own? There's a million cool packages out there if we add them all to PINTS, even as optional extras, our testing will start taking forever — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub, or unsubscribe.

[ben18785]

Here's a Heath Robinson version of the poll. Please edit the below to add your name next to the answer that best represents your response:

• Interfaces with other packages should be within the main Pints repo. @ben18785 @martinjrobins
• Interfaces in a separate repo. @MichaelClerx

[MichaelClerx]

And whenever any single one of these interfaces breaks, due to some amazing new improvements in numpyro, theano, pymc3, stan etc., then none of our PRs will build and nobody can do any work until it's fixed.
I think these interfaces are a great idea, and they'll be super useful, but would much rather see them live in their own repo where they can be maintained by whoever finds them most interesting, rather than whoever is the first to try to work on PINTS after they break.
Coding up an interface is one thing, maintaining it for the next N years is something else.

[martinjrobins]

I think its important they be in the same repo, so that anyone browsing the examples to see what you can do with pints will find them. I agree that we'll have to think carefully about testing, my suggestion would be to only test the interface against a specific version of numpyro, stan etc. That is, the github action which tests the interface notebook would pip install a particular version of the other package. Whoever wants to maintain that particular interface would be responsible for incrementing this version in step with the third-party package. That way, if we don't maintain it, all that happens is that we end up supporting an older version.

[chonlei]

For browsing example purposes, I think we could simply add a link in the examples to a separate repo which contains all these interface notebooks? In that way we don't have to maintain and/or add extra dependences/packages in the main pints?

[DavAug]

I don't feel like I have I vote on this. I can see pros and cons for both sides, and @martinjrobins and @chonlei suggestions seem to be feasible leaning towards either direction. What exactly would be the disadvantage of supporting only specific versions of other packages? Does pip installing pints then mess with the versions of numpyro if you had installed it in the environment perviously, and would that be a problem?

[MichaelClerx]

Does pip installing pints then mess with the versions of numpyro if you had installed it in the environment perviously, and would that be a problem?

Nah pip would just complain that it can't do that, I think. Alright let's just all do it in this repo for now. Have a look at the Stan PR #1077 first, there'll be some overlap I'm sure.

[martinjrobins]

Note I don't think numpyro or the others should be a dependency of pints! I think this issue should just be a notebook showing how you can use the pints samplers using a numpyro model. The github actions test runner will need to install numpyro so that this notebook can be tested, but that is all.

[MichaelClerx]

Note I don't think numpyro or the others should be a dependency of pints! I think this issue should just be a notebook showing how you can use the pints samplers using a numpyro model. The github actions test runner will need to install numpyro so that this notebook can be tested, but that is all.

Ah OK. So we'd need some place to store a list of deps for the notebooks? Or just stick it directly in the CI files?

[martinjrobins]

stick it directly in the CI files I think? We could have a separate CI file for the interface notebooks if you think that would be a cleaner separation

[MichaelClerx]

Sticking it in https://github.com/pints-team/pints/blob/master/.github/workflows/notebook-test.yml is fine by me. Probably add a line to the top of the notebook saying which additional dependencies are needed?

[DavAug]

Let me retreat from this issue for the moment. I got stuck with a number of other things for the next days.