added 1) Uniform sampling test and 2) basic SBC test
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:
- Doesn't work for parameters that are arrays yet.
- Builds the likelihood and samplers etc.. at every iteration so can't be batched. As the result the python loop is very slow (around 1h10min for 1K rank statistics
- the MCMC chain length is variable as the test needs independent samples. I used emcee
integrated_timefunction to get the IAT and used that to determine whether the chain should be run for longer. For some generated datasets the sampler needs to be run for longer. It might be helpful to generate less data. That way the IAT will (probably) always be below some desired threshold. The test would then still check the IAT, but could perhaps discard that sample
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.