Documentation clean up

CONTRIBUTING.md

@@ -0,0 +1,59 @@
+# Guidelines for Contributing
+
+Thank you for interested in contributing to Blackjax! We value the following contributions:
+
+- Bug fixes
+- Documentation
+- High-level sampling algorithms from any family of algorithms: random walk,
+  hamiltonian monte carlo, sequential monte carlo, variational inference,
+  inference compilation, etc.
+- New building blocks, e.g. new metrics for HMC, integrators, etc.
+
+## How to contribute?
+
+1. Run `pip install -r requirements.txt` to install all the dev
+   dependencies.
+2. Run `pre-commit run --all-files` and `make test` before pushing on the repo; CI should pass if
+   these pass locally.
+
+## Editing documentations
+
+The Blackjax repository (and [sampling-book](/~https://github.com/blackjax-devs/sampling-book)) provides examples in the form of Markdown documents. [Jupytext](/~https://github.com/mwouts/jupytext) can be used by the users to convert their Jupyter notebooks to this format, or convert these documents to Jupyter notebooks. Examples are rendered in the [documentation](https://blackjax-devs.github.io/blackjax/).
+
+### Load examples in a Jupyter notebook
+
+To convert any example file to a Jupyter notebook you can use:
+
+```shell
+jupytext docs/examples/your_example_file.md --to notebook
+```
+
+you can then interact with the resulting notebook just like with any notebook.
+
+### Convert my Jupyter notebook to markdown
+
+If you implemented your example in a Jupyter notebook you can convert your `.ipynb` file to Markdown using the command below:
+
+```shell
+jupytext docs/examples/your_example_notebook.ipynb --to myst
+```
+
+Once the example file is converted to a Markdown file, you have two options for editing:
+
+1. Edit the Markdown version as it is a regular Markdown file.
+2. Edit the Notebook version, then convert it to a Markdown file once you finish editing with the command above. Jupytext can handle the change if the example has the same file name.
+
+**Please make sure to only commit the Markdown file.**
+
+### Composing Documentation on Sphinx-Doc
+
+We use `Sphinx` to generate documents for this repo. We highly encourage you to check how your changes to the examples are rendered in the documentation:
+
+1. Add your documentation to `docs/examples.rst`
+2. Run the command below:
+
+```shell
+ make build-docs
+```
+
+3. Check the generated HTML documentation in `docs/_build`

README.md

@@ -128,22 +128,7 @@ passing parameters.
 
 ## Contributions
 
-### What contributions?
-
-We value the following contributions:
-- Bug fixes
-- Documentation
-- High-level sampling algorithms from any family of algorithms: random walk,
-  hamiltonian monte carlo, sequential monte carlo, variational inference,
-  inference compilation, etc.
-- New building blocks, e.g. new metrics for HMC, integrators, etc.
-
-### How to contribute?
-
-1. Run `pip install -r requirements.txt` to install all the dev
-   dependencies.
-2. Run `pre-commit run --all-files` and `make test` before pushing on the repo; CI should pass if
-   these pass locally.
+Please follow our [short guide](/~https://github.com/blackjax-devs/blackjax/blob/main/CONTRIBUTING.md).
 
 ## Citing Blackjax
 

blackjax/mcmc/marginal_latent_gaussian.py

@@ -37,6 +37,7 @@ class MarginalState(NamedTuple):
     U_x
         Auxiliary attributes
     U_grad_x
+        Gradient of the auxiliary attributes
 
     """
 

blackjax/mcmc/proposal.py

@@ -97,8 +97,10 @@ def proposal_from_energy_diff(
     state: TrajectoryState,
 ) -> Tuple[Proposal, bool]:
     """Computes a new proposal from the energy difference between two states.
+
     It also verifies whether this difference is a divergence, if the
     energy diff is above divergence_threshold.
+
     Parameters
     ----------
     initial_energy
@@ -141,10 +143,13 @@ def asymmetric_proposal_generator(
     proposal_factory: Callable = proposal_from_energy_diff,
 ) -> Tuple[Callable, Callable]:
     """A proposal generator that takes into account the transition between
-    two states to compute a new proposal. In particular, both states are
-    used to compute the energies to consider in weighting the proposal,
+    two states to compute a new proposal.
+
+    In particular, both states are used to compute the energies to consider in weighting the proposal,
     to account for asymmetries.
-     ----------
+
+    Parameters
+    ----------
     transition_energy_fn
         A function that computes the energy of a transition from an initial state
         to a new state, given some optional keyword arguments.

blackjax/mcmc/random_walk.py

@@ -18,20 +18,27 @@
 with simpler versions of the algorithms, but in all cases they are particular instantiations
 of the Random Walk Rosenbluth-Metropolis-Hastings.
 
-Let's note x_{t-1} to the previous position and x_t to the newly sampled one.
+Let's note $x_{t-1}$ to the previous position and $x_t$ to the newly sampled one.
 
 The variants offered are:
 
 1. Proposal distribution as addition of random noice from previous position. This means
-x_t = x_{t-1} + step. Function: `additive_step`
+   $x_t = x_{t-1} + step$.
 
-2. Independent proposal distribution: P(x_t) doesn't depend on x_{t_1}. Function: `irmh`
+    Function: `additive_step`
 
-3. Proposal distribution using a symmetric function. That means P(x_t|x_{t-1}) = P(x_{t-1}|x_t).
- Function: `rmh` without proposal_logdensity_fn. See 'Metropolis Algorithm' in [1]
+2. Independent proposal distribution: $P(x_t)$ doesn't depend on $x_{t_1}$.
 
-4. Asymmetric proposal distribution. Function: `rmh` with proposal_logdensity_fn.
- See 'Metropolis-Hastings' Algorithm in [1]
+    Function: `irmh`
+
+3. Proposal distribution using a symmetric function. That means $P(x_t|x_{t-1}) = P(x_{t-1}|x_t)$.
+   See 'Metropolis Algorithm' in [1].
+
+    Function: `rmh` without proposal_logdensity_fn.
+
+4. Asymmetric proposal distribution. See 'Metropolis-Hastings' Algorithm in [1].
+
+    Function: `rmh` with proposal_logdensity_fn.
 
 Reference: :cite:p:`gelman2014bayesian` Section 11.2
 
@@ -225,6 +232,7 @@ def normal_random_walk(cls, logdensity_fn: Callable, sigma):
             The log density probability density function from which we wish to sample.
         sigma
             The value of the covariance matrix of the gaussian proposal distribution.
+
         Returns
         -------
              A ``SamplingAlgorithm``.
@@ -265,6 +273,7 @@ def kernel(
         proposal_distribution: Callable,
     ) -> Tuple[RWState, RWInfo]:
         """
+
         Parameters
         ----------
         proposal_distribution
@@ -338,6 +347,7 @@ def step_fn(rng_key: PRNGKey, state):
 
 def build_rmh():
     """Build a Rosenbluth-Metropolis-Hastings kernel.
+
     Returns
     -------
     A kernel that takes a rng_key and a Pytree that contains the current state