rlouf / mcx

Express & compile probabilistic programs for performant inference on CPU & GPU. Powered by JAX.
https://rlouf.github.io/mcx
Apache License 2.0
325 stars 17 forks source link

added 1) Uniform sampling test and 2) basic SBC test #64

Open jeremiecoullon opened 3 years ago

jeremiecoullon commented 3 years ago

Here are some changes to get the ball rolling with MCMC tests. There's still a lot left to do though. I include a notebook to look at the empirical distributions of the test outputs. Though obviously that will need to be automated at some point.

Though it might be fun at some point to have a similar notebook where a user (they would probably be an mcmc-nerd..) could interact with the outputs of the tests themselves. That way they could convince themselves that the sampler works correctly.

Uniform sampling

This is similar to the hmc_test.py except that it gets a lot more samples. I set a simple model with 1 parameter where the prior is uniform and the data doesn't depend on the prior. That way the posterior is also uniform. I use HMC to get 500K samples and in the notebook I plot the empirical CDF along with the line f(y)=x. This is a very simple test to do and has saved me from a few bugs in the past! Perhaps this could be removed and the hmc_test.py could simply run for longer (to make sure that the variance is correct).

SBC

This is a first implementation of SBC. It will work for any model with scalar parameters. Problems/things to think about:

rlouf commented 3 years ago

Now that we're moving samplers to blackjax, I think tests based on SBC would be better there. It will also be a lot simpler to implement than in MCX.