From 4b9d9804b8143d7325c72950d94d844430d3f840 Mon Sep 17 00:00:00 2001 From: Alberto Cabezas Gonzalez Date: Fri, 12 May 2023 15:29:59 +0100 Subject: [PATCH] Guidelines with MH example and revised skeletons --- README.md | 4 +++ docs/developer/approximate_inf_algorithm.py | 12 +++---- docs/developer/guidelines.md | 36 +++++++++++---------- docs/developer/sampling_algorithm.py | 24 ++++++++++---- docs/index.md | 2 +- 5 files changed, 47 insertions(+), 31 deletions(-) diff --git a/README.md b/README.md index 8bfe02ebf..8a9977878 100644 --- a/README.md +++ b/README.md @@ -126,6 +126,10 @@ information related to the transition are returned separately. They can thus be easily composed and exchanged. We specialize these kernels by closure instead of passing parameters. +### New algorithms + +We hope to make implementing and testing new algorithms easy with BlackJAX. Many basic methods are already implemented in the library, and you can use them to test new algorithms. Follow the [guidelines](https://blackjax-devs.github.io/blackjax/developer/guidelines.html) to implement your own method and test new ideas on existing methods without writing everything from scratch! + ## Contributions ### What contributions? diff --git a/docs/developer/approximate_inf_algorithm.py b/docs/developer/approximate_inf_algorithm.py index 67b9bb2e0..52f5a8c24 100644 --- a/docs/developer/approximate_inf_algorithm.py +++ b/docs/developer/approximate_inf_algorithm.py @@ -32,7 +32,7 @@ class ApproxInfState(NamedTuple): - """State of the approximate inference algorithm. + """State of your approximate inference algorithm. Give an overview of the variables needed at each step and for sampling. """ @@ -41,9 +41,9 @@ class ApproxInfState(NamedTuple): class ApproxInfInfo(NamedTuple): - """Additional information on the algorithm transition. + """Additional information on your algorithm transition. - Given an overview of the collected values at each step of the approximation. + Give an overview of the collected values at each step of the approximation. """ ... @@ -63,7 +63,7 @@ def step( *args, **kwargs, ) -> Tuple[ApproxInfState, ApproxInfInfo]: - """Approximate the target density using the some approximation. + """Approximate the target density using your approximation. Parameters ---------- @@ -81,14 +81,14 @@ def step( def sample(rng_key: PRNGKey, state: ApproxInfState, num_samples: int = 1): - """Sample from the approximation.""" + """Sample from your approximation.""" # the sample should be a PyTree of the same structure as the `position` in the init function samples = ... return samples class approx_inf_algorithm: - """Implements the (basic) user interface for the approximate inference method. + """Implements the (basic) user interface for your approximate inference method. Describe in detail the inner mechanism of the method and its use. diff --git a/docs/developer/guidelines.md b/docs/developer/guidelines.md index 16b4d9a4e..6d1d98695 100644 --- a/docs/developer/guidelines.md +++ b/docs/developer/guidelines.md @@ -1,41 +1,43 @@ # Developer Guidelines -## Style -In its broadest sense, an algorithm that belongs in the blackjax library should approximate integrals on a probability space. An introduction to probability theory is outside the scope of this document, but the Monte Carlo method is ever-present and important to understand. In simple terms, we want to approximate an integral with a sum. To do this, generate samples with probabilities defined by a density (continuous variable) or measure (discrete variable) function. The idea is to sample more from areas with higher probability but also from areas with low probability, just at a lower rate. You can also approximate the target density directly, using an approximation that is easier to handle, then do inference, i.e. solve integrals, with the approximation directly and use importance sampling to correct its bias. +In the broadest sense, an algorithm that belongs in the BlackJAX library should provide the tools to approximate integrals on a probability space. An introduction to probability theory is outside the scope of this document, but the Monte Carlo method is ever-present and important to understand. In simple terms, we want to approximate an integral with a sum. To do this, generate samples with [relative likelihood](https://en.wikipedia.org/wiki/Relative_likelihood) given by a target probability density function (known up to a normalization constant). The idea is to sample more from areas with higher likelihood but also from areas with low likelihood, just at a lower rate. You can also approximate the target density directly, using a density that is tractable and easy to sample from, then do inference with the approximation instead of the target, potentially using [importance sampling](https://en.wikipedia.org/wiki/Importance_sampling) to correct the approximation error. -In the following section, we’ll explain blackjax’s design of different algorithms for Monte Carlo integration. Keep in mind some basic principles: +In the following section, we’ll explain BlackJAX’s design of different algorithms for Monte Carlo integration. Keep in mind some basic principles: -Leverage JAX's unique strengths: functional programming and composable function-transformation approach. -Write small and general functions, compose them to create complex methods, reuse the same building blocks for similar algorithms. -Consider compatibility with the broader JAX ecosystem (Flax, Optax, GPJax). -Write code that is easy to read and understand. -Write code that is well documented, describe in detail the inner mechanism of the algorithm and its use. +- Leverage JAX's unique strengths: functional programming and composable function-transformation approach. +- Write small and general functions, compose them to create complex methods, reuse the same building blocks for similar algorithms. +- Consider compatibility with the broader JAX ecosystem (Flax, Optax, GPJax). +- Write code that is easy to read and understand. +- Write code that is well documented, describe in detail the inner mechanism of the algorithm and its use. ## Core implementation -There are three types of sampling algorithms blackjax currently supports: Markov Chain Monte Carlo (MCMC), Sequential Monte Carlo (SMC), and Stochastic Gradient MCMC (SGMCMC); and one type of approximate inference algorithm: Variational Inference (VI). Additionally, blackjax supports adaptation algorithms that efficiently tune the hyperparameters of sampling algorithms, usually aimed at reducing autocorrelation between sequential samples. +There are three types of sampling algorithms BlackJAX currently supports: Markov Chain Monte Carlo (MCMC), Sequential Monte Carlo (SMC), and Stochastic Gradient MCMC (SGMCMC); and one type of approximate inference algorithm: Variational Inference (VI). Additionally, BlackJAX supports adaptation algorithms that efficiently tune the hyperparameters of sampling algorithms, usually aimed at reducing autocorrelation between sequential samples. -Basic components are functions, which do specific tasks but are generally applicable, used to build all inference algorithms. When implementing a new inference algorithm, you should first break it down to its basic components, then find and use all that are already implemented *before* writing your own. A recurrent example is the Metropolis-Hastings step, a basic component used by many MCMC algorithms to keep the target distribution invariant. In blackjax, this common accept/reject step done with two functions: first the Hastings ratio is calculated by creating a proposal using `mcmc.proposal.proposal_generator`, then the proposal is accepted or rejected using `mcmc.proposal.static_binomial_sampling`. +Basic components are functions which do specific tasks but are generally applicable, used to build all inference algorithms. When implementing a new inference algorithm you should first break it down to its basic components then find and use all that are already implemented *before* writing your own. A recurrent example is the [Metropolis-Hastings](https://en.wikipedia.org/wiki/Metropolis%E2%80%93Hastings_algorithm) step, a basic component used by many MCMC algorithms to keep the target distribution invariant. In BlackJAX there are two basic components that do a specific (but simpler) and a general version of this accept/reject step: -Because JAX operates on pure functions, inference algorithms always return a NamedTuple containing the necessary variables to generate the next sample. Arguably, abstracting the handling of these variables is the whole point of blackjax, so it must be done in a way that abstracts the uninteresting bookkeeping from the end user but allows her to access important variables at each step. The algorithms should also return a NamedTuple with important information of each iteration. +- Metropolis step: if the proposal transition kernel is symmetric, i.e. if the probability of going from the initial to the proposed position is always equal to the probability of going from the proposed to the initial position, the acceptance probability is calculated by creating a proposal using `mcmc.proposal.proposal_generator`, then the proposal is accepted or rejected using `mcmc.proposal.static_binomial_sampling`. +- Metropolis-Hastings step: for the more general case of an asymmetric proposal transition kernel, the acceptance probability is calculated by creating a proposal using `mcmc.proposal.asymmetric_proposal_generator`, then the proposal is accepted or rejected using `mcmc.proposal.static_binomial_sampling`. + +When implementing an algorithm you could choose to replace the classic, reversible Metropolis-Hastings step with Neal's [non-reversible slice sampling](https://arxiv.org/abs/2001.11950) step by simply replacing `mcmc.proposal.static_binomial_sampling` with `mcmc.proposal.nonreversible_slice_sampling` on either of the previous implementations. Just make sure to carry over to the next iteration an updated slice, instead of passing a pseudo-random number generating key, for the slice sampling step! + +The previous example illustrates the power of basic components, useful not only to avoid rewriting the same methods for each new algorithm but also useful to personalize and test new algorithms which replace some steps of common efficient algorithms. Like how `blackjax.mcmc.ghmc` is `blackjax.mcmc.hmc` with a persistent momentum and a non-reversible slice sampling step instead of the Metropolis-Hastings step. + +Because JAX operates on pure functions, inference algorithms always return a `typing.NamedTuple` containing the necessary variables to generate the next sample. Arguably, abstracting the handling of these variables is the whole point of BlackJAX, so it must be done in a way that abstracts the uninteresting bookkeeping from the end user but allows her to access important variables at each step. The algorithms should also return a `typing.NamedTuple` with important information about each iteration. The user-facing interface of a **sampling algorithm** should work like this: ```python -import blackjax sampling_algorithm = blackjax.sampling_algorithm(logdensity_fn, *args, **kwargs) state = sampling_algorithm.init(initial_position) new_state, info = sampling_algorithm.step(rng_key, state) ``` -Achieve this by building from the basic skeleton of a sampling algorithm (here)[https://github.com/blackjax-devs/blackjax/tree/main/docs/developer/sampling_algorithm.py]. Only the `sampling_algorithm` class and the `init` and `build_kernel` functions need to be in the final version of your algorithm, the rest might become useful but are not necessary. +Achieve this by building from the basic skeleton of a sampling algorithm [here](https://github.com/blackjax-devs/blackjax/tree/main/docs/developer/sampling_algorithm.py). Only the `sampling_algorithm` class and the `init` and `build_kernel` functions need to be in the final version of your algorithm, the rest might become useful but are not necessary. The user-facing interface of an **approximate inference algorithm** should work like this: ```python -import blackjax approx_inf_algorithm = blackjax.approx_inf_algorithm(logdensity_fn, optimizer, *args, **kwargs) state = approx_inf_algorithm.init(initial_position) new_state, info = approx_inf_algorithm.step(rng_key, state) #user is able to build the approximate distribution using the state, or generate samples: position_samples = approx_inf_algorithm.sample(rng_key, state, num_samples) ``` -Achieve this by building from the basic skeleton of an approximate inference algorithm (here)[https://github.com/blackjax-devs/blackjax/tree/main/docs/developer/approximate_inf_algorithm.py]. Only the `approx_inf_algorithm` class and the `init`, `step` and `sample` functions need to be in the final version of your algorithm, the rest might become useful but are not necessary. - -Well documented code is essential for a useful library. Start by decomposing your algorithm into basic components, finding those that are already implemented, then implement your own and build the high-level API from basic components. +Achieve this by building from the basic skeleton of an approximate inference algorithm [here](https://github.com/blackjax-devs/blackjax/tree/main/docs/developer/approximate_inf_algorithm.py). Only the `approx_inf_algorithm` class and the `init`, `step` and `sample` functions need to be in the final version of your algorithm, the rest might become useful but are not necessary. diff --git a/docs/developer/sampling_algorithm.py b/docs/developer/sampling_algorithm.py index df9a057dd..747c9e23b 100644 --- a/docs/developer/sampling_algorithm.py +++ b/docs/developer/sampling_algorithm.py @@ -32,7 +32,7 @@ class SamplingAlgoState(NamedTuple): - """State of the sampling algorithm. + """State of your sampling algorithm. Give an overview of the variables needed at each iteration of the model. """ @@ -41,7 +41,7 @@ class SamplingAlgoState(NamedTuple): class SamplingAlgoInfo(NamedTuple): - """Additional information on the algorithm transition. + """Additional information on your algorithm transition. Given an overview of the collected values at each iteration of the model. """ @@ -56,7 +56,7 @@ def init(position: PyTree, logdensity_fn: Callable, *args, **kwargs): def build_kernel(*args, **kwargs): - """Build a HMC kernel. + """Build a your kernel. Parameters ---------- @@ -92,7 +92,7 @@ def kernel( class sampling_algorithm: - """Implements the (basic) user interface for the sampling kernel. + """Implements the (basic) user interface for your sampling kernel. Describe in detail the inner mechanism of the algorithm and its use. @@ -148,10 +148,20 @@ def sampling_algorithm_proposal(*args, **kwags) -> Callable: ------- Describe what is returned. """ - ... + # as an example, a Metropolis-Hastings step would look like this: + init_proposal, generate_proposal = proposal.proposal_generator(...) + sample_proposal = proposal.static_binomial_sampling(...) + + def generate(rng_key, state): + # propose a new sample + proposal_state = ... + + # accept or reject the proposed sample + proposal = init_proposal(state) + new_proposal, is_diverging = generate_proposal(proposal.energy, proposal_state) + sampled_proposal, *info = sample_proposal(rng_key, proposal, new_proposal) - def generate(*args, **kwargs): - """Generate a new chain state.""" + # build a new state and collect useful information sampled_state, info = ... return sampled_state, info diff --git a/docs/index.md b/docs/index.md index c0f9251c1..a6c801e67 100644 --- a/docs/index.md +++ b/docs/index.md @@ -141,5 +141,5 @@ maxdepth: 1 caption: DEVELOPER DOCUMENTATION hidden: --- -Guidelines +Guidelines ```