Added plotting utilities and refactored hyperparams #143
Conversation
|
@ThibeauWouters The checks seem to have failed due to some missing hyperparameters, would you mind have a look of that? |
|
@kazewong I have pushed some updates for the failed checks, and have also added some additional postprocessing analysis code to check the Gelman-Rubin R statistic. |
|
My tests show something is still off, I need to double check things |
|
It should be fixed now |
There was a problem hiding this comment.
I would not include r_hat related functionality in this PR. I am more inclined to have that as a notebook in example, since statistical metric like r_hat or correlation length needs to be used with context of user understanding the meaning of what they are measuring or monitoring.
src/flowMC/utils/postprocessing.py
Outdated
| import jax.numpy as jnp | ||
| from flowMC.sampler.Sampler import Sampler | ||
|
|
||
| def plot_summary(sampler: Sampler, which: str = "training", **plotkwargs) -> None: |
There was a problem hiding this comment.
Instead of taking the entire Sampler as input, these post-processing function should only interact with serialized output from the sampler to avoid run time complication.
This might need a bit more work in terms of tidying up what kind of output we are serialize from the sampler. I would suggest holding these modifications off or open a separate issue. I will get to it soonish
src/flowMC/sampler/Sampler.py
Outdated
| @@ -113,11 +128,16 @@ def __init__( | |||
| production["log_prob"] = jnp.empty((self.n_chains, 0)) | |||
| production["local_accs"] = jnp.empty((self.n_chains, 0)) | |||
| production["global_accs"] = jnp.empty((self.n_chains, 0)) | |||
|
|
|||
| if self.track_gelman_rubin: | |||
There was a problem hiding this comment.
Since R_hat can be compute in post-processing, I don't think we should have it in the main sampler.
src/flowMC/sampler/Sampler.py
Outdated
| @@ -264,6 +284,15 @@ def sampling_loop( | |||
| global_acceptance[:, 1::self.output_thinning], | |||
| axis=1, | |||
| ) | |||
|
|
|||
| if self.track_gelman_rubin: | |||
There was a problem hiding this comment.
Same as above. Please remove this from the main sampler
src/flowMC/utils/postprocessing.py
Outdated
| plt.ylim(0-eps, 1+eps) | ||
| plt.savefig(f"{outdir}{name}_{which}.png", bbox_inches='tight') | ||
|
|
||
| def gelman_rubin(chains: Float[Array, "n_chains n_steps n_dim"], discard_fraction: float = 0.1) -> Array: |
There was a problem hiding this comment.
Arviz actually support r_hat and many more different statistical metric. I would suggest to use that instead of writing our own function.
This PR adds a few basic functionalities to simplify the treatment of the hyperparameters (now moved to a separate file) and by adding a few functionalities that allow to easily plot a few key quantities such as acceptance rates by a simple function call from any flowMC or jim script.