Repository: kazewong/flowMC (language: python)

Issue:
Sync Documentation up to main



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 3d7385d..fef648e 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,7 @@
 <img src="https://badgen.net/badge/License/MIT/blue" alt="doc"/>
 </a>
 
+![flowMC_logo](./docs//source/logo_0810.png)
 
 FlowMC is a Jax-based python package for normalizing-flow enhanced Markov chain Monte Carlo (MCMC) sampling.
 The code is open source under MIT license, and it is under active development.
@@ -26,8 +27,6 @@ The simplest way to install the package is to do it through pip
 pip install flowMC
 ```
 
-
-
 ## Requirements
 
 Here is a list of packages we use in the main library
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 0000000..11aaade
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,7 @@
+sphinx-book-theme
+myst-nb
+jax
+jaxlib
+flax
+distrax
+tqdm
\ No newline at end of file
diff --git a/docs/source/FAQ.rst b/docs/source/FAQ.rst
new file mode 100644
index 0000000..8040d81
--- /dev/null
+++ b/docs/source/FAQ.rst
@@ -0,0 +1,11 @@
+.. _FAQ-section-top:
+
+FAQ
+===
+
+**My local sampler is not accepting**
+
+**My global sample's loss is exploding/not decresing**
+
+**The sampler is stuck a bit until it starts sampling**
+
diff --git a/docs/source/analysis.rst b/docs/source/analysis.rst
index 0c4901e..55c12a4 100644
--- a/docs/source/analysis.rst
+++ b/docs/source/analysis.rst
@@ -1,4 +1,6 @@
+.. _analysis-section-top:
+
 Analysis
 ========
 
-Detail about analying a run
\ No newline at end of file
+Under construction
\ No newline at end of file
diff --git a/docs/source/api/nfsampler.nfmodel.rst b/docs/source/api/flowMC.nfmodel.rst
similarity index 70%
rename from docs/source/api/nfsampler.nfmodel.rst
rename to docs/source/api/flowMC.nfmodel.rst
index 1a988c3..d6a7d05 100644
--- a/docs/source/api/nfsampler.nfmodel.rst
+++ b/docs/source/api/flowMC.nfmodel.rst
@@ -1,11 +1,11 @@
 flowMC.nfmodel package
-=========================
+======================
 
 Submodules
 ----------
 
 flowMC.nfmodel.maf module
-----------------------------
+-------------------------
 
 .. automodule:: flowMC.nfmodel.maf
    :members:
@@ -13,16 +13,24 @@ flowMC.nfmodel.maf module
    :show-inheritance:
 
 flowMC.nfmodel.realNVP module
---------------------------------
+-----------------------------
 
 .. automodule:: flowMC.nfmodel.realNVP
    :members:
    :undoc-members:
    :show-inheritance:
 
-flowMC.nfmodel.utils module
+flowMC.nfmodel.rqSpline module
 ------------------------------
 
+.. automodule:: flowMC.nfmodel.rqSpline
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+flowMC.nfmodel.utils module
+---------------------------
+
 .. automodule:: flowMC.nfmodel.utils
    :members:
    :undoc-members:
diff --git a/docs/source/api/nfsampler.rst b/docs/source/api/flowMC.rst
similarity index 92%
rename from docs/source/api/nfsampler.rst
rename to docs/source/api/flowMC.rst
index aed8dbc..11af931 100644
--- a/docs/source/api/nfsampler.rst
+++ b/docs/source/api/flowMC.rst
@@ -1,5 +1,5 @@
 flowMC package
-=================
+==============
 
 Subpackages
 -----------
diff --git a/docs/source/api/nfsampler.sampler.rst b/docs/source/api/flowMC.sampler.rst
similarity index 72%
rename from docs/source/api/nfsampler.sampler.rst
rename to docs/source/api/flowMC.sampler.rst
index 938af51..9814eda 100644
--- a/docs/source/api/nfsampler.sampler.rst
+++ b/docs/source/api/flowMC.sampler.rst
@@ -1,11 +1,11 @@
 flowMC.sampler package
-=========================
+======================
 
 Submodules
 ----------
 
 flowMC.sampler.Gaussian\_random\_walk module
------------------------------------------------
+--------------------------------------------
 
 .. automodule:: flowMC.sampler.Gaussian_random_walk
    :members:
@@ -13,7 +13,7 @@ flowMC.sampler.Gaussian\_random\_walk module
    :show-inheritance:
 
 flowMC.sampler.MALA module
------------------------------
+--------------------------
 
 .. automodule:: flowMC.sampler.MALA
    :members:
@@ -21,7 +21,7 @@ flowMC.sampler.MALA module
    :show-inheritance:
 
 flowMC.sampler.NF\_proposal module
--------------------------------------
+----------------------------------
 
 .. automodule:: flowMC.sampler.NF_proposal
    :members:
@@ -29,13 +29,21 @@ flowMC.sampler.NF\_proposal module
    :show-inheritance:
 
 flowMC.sampler.Sampler module
---------------------------------
+-----------------------------
 
 .. automodule:: flowMC.sampler.Sampler
    :members:
    :undoc-members:
    :show-inheritance:
 
+flowMC.sampler.mMALA module
+---------------------------
+
+.. automodule:: flowMC.sampler.mMALA
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Module contents
 ---------------
 
diff --git a/docs/source/api/nfsampler.utils.rst b/docs/source/api/flowMC.utils.rst
similarity index 52%
rename from docs/source/api/nfsampler.utils.rst
rename to docs/source/api/flowMC.utils.rst
index 2f8da99..846fcc9 100644
--- a/docs/source/api/nfsampler.utils.rst
+++ b/docs/source/api/flowMC.utils.rst
@@ -1,11 +1,11 @@
 flowMC.utils package
-=======================
+====================
 
 Submodules
 ----------
 
 flowMC.utils.Analysis module
--------------------------------
+----------------------------
 
 .. automodule:: flowMC.utils.Analysis
    :members:
@@ -13,13 +13,29 @@ flowMC.utils.Analysis module
    :show-inheritance:
 
 flowMC.utils.PRNG\_keys module
----------------------------------
+------------------------------
 
 .. automodule:: flowMC.utils.PRNG_keys
    :members:
    :undoc-members:
    :show-inheritance:
 
+flowMC.utils.PythonFunctionWrap module
+--------------------------------------
+
+.. automodule:: flowMC.utils.PythonFunctionWrap
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+flowMC.utils.progressBar module
+-------------------------------
+
+.. automodule:: flowMC.utils.progressBar
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 Module contents
 ---------------
 
diff --git a/docs/source/api/modules.rst b/docs/source/api/modules.rst
index 28219c4..f54f65b 100644
--- a/docs/source/api/modules.rst
+++ b/docs/source/api/modules.rst
@@ -1,5 +1,7 @@
-API
-===
+.. _api-section-top:
+
+flowMC
+======
 
 .. toctree::
    :maxdepth: 4
diff --git a/docs/source/conf.py b/docs/source/conf.py
index bf81b0e..25a2761 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -12,47 +12,59 @@
 #
 import os
 import sys
-sys.path.append(os.path.abspath('../..'))
-sys.path.append(os.path.abspath('..'))
+
+sys.path.append(os.path.abspath("../../src"))
+sys.path.append(os.path.abspath(".."))
 
 
 # -- Project information -----------------------------------------------------
 
-project = 'flowMC'
-copyright = '2022, Kaze Wong, Marylou Gabrié'
-author = 'Kaze Wong, Marylou Gabrié'
+project = "flowMC"
+copyright = "2022, Kaze Wong, Marylou Gabrié"
+author = "Kaze Wong, Marylou Gabrié & contributors"
 
 # The full version, including alpha/beta/rc tags
-release = '0.0.2'
-master_doc = 'index'
+release = "0.0.4"
+master_doc = "index"
 
 # -- General configuration ---------------------------------------------------
 
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = [ 'sphinx.ext.autodoc',
-               'sphinx.ext.napoleon',
-               'sphinx.ext.autosummary',
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.autosummary",
+    "myst_parser"
 ]
 
+
+
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = []
+exclude_patterns = ["_build"]
 
+# HTML theme
+html_theme = "sphinx_book_theme"
+html_copy_source = True
+html_show_sourcelink = True
+html_sourcelink_suffix = ""
+html_title = "flowMC"
+html_theme_options = {
+    "repository_url": "https://github.com/kazewong/flowMC",
+    "repository_branch": "main",
+    "use_edit_page_button": True,
+    "use_issues_button": True,
+    "use_repository_button": True,
+    "use_download_button": True,
+}
 
-# -- Options for HTML output -------------------------------------------------
-
-# The theme to use for HTML and HTML Help pages.  See the documentation for
-# a list of builtin themes.
-#
-html_theme = 'alabaster'
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
\ No newline at end of file
diff --git a/docs/source/configuration.rst b/docs/source/configuration.rst
index 78fc376..7079222 100644
--- a/docs/source/configuration.rst
+++ b/docs/source/configuration.rst
@@ -1,8 +1,9 @@
+.. _configuration-section-top:
+
 Configuration
 =============
 
-
-# Parameters
+This page contains information about what knobs can you tune when the sampler doesn't work as expected.
 
 
 
@@ -27,3 +28,23 @@ momentum                           Momentum to use when training the normalizing
 batch_size                         Batch size to use when training the normalizing flow model.
 logging                            Whether we log or not.
 ===============================    ===================================
+
+Sampler class arguments
+-----------------------
+
+likelihood 
+
+
+Local sampler parameters
+------------------------
+
+Normalizing flow parameters
+---------------------------
+
+flowMC.sampler.Sampler module
+-----------------------------
+
+.. automodule:: flowMC.sampler.Sampler
+   :members:
+   :noindex:
+
diff --git a/docs/source/examples.rst b/docs/source/examples.rst
index ce68152..10c2969 100644
--- a/docs/source/examples.rst
+++ b/docs/source/examples.rst
@@ -1,3 +1,6 @@
+.. _example-section-top:
+
 Examples
 ========
 
+Under construction
\ No newline at end of file
diff --git a/docs/source/index.rst b/docs/source/index.rst
index d41d3cf..0488487 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -1,25 +1,46 @@
-Welcome to flowMC's documentation!
+flowMC
 =====================================
 
+**Normalizing-flow enhanced sampling package for probabilistic inference**
+
+
+.. image:: logo_0810.png
+
+.. image:: https://badgen.net/badge/Read/the doc/blue
+   :target: https://flowMC.readthedocs.io/en/latest/
+.. image:: https://badgen.net/badge/License/MIT/blue
+   :target: https//github.com/kazewong/flowMC/blob/Packaging/LICENSE
+
+
+
+flowMC is a Jax-based python package for normalizing-flow enhanced Markov chain Monte Carlo (MCMC) sampling.
+The code is open source under MIT license, and it is under active development.
+
+- Just-in-time compilation is supported.
+- Native support for GPU acceleration.
+- Suit for problems with multi-modality.
+
+
+Five steps to use flowMC
+========================
+
+#. You will find basic information such as installation and a quick guide in the :ref:`quickstart-section-top`.
+#. We give more information about tuning parameters of our sampler in the :ref:`configuration-section-top`.
+#. In :ref:`analysis-section-top`, we give some examples of common diagnostics one can use in understanding the performance of our sampler.
+#. We list some examples in :ref:`example-section-top` so users can see whether there is a similar use case they can adopt their code quickly.
+#. Finally, the auto-generated API doc is available in the :ref:`api-section-top`.
+
+User guide
+----------
+
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
    :caption: Contents
    
-   intro
+   quickstart
    configuration
    analysis
-   api/modules
-
-
-.. toctree::
-   :maxdepth: 1
-   :caption: Tutorial
-
    examples
-
-Indices and tables
-==================
-
-* :ref:`genindex`
-* :ref:`modindex`
-* :ref:`search`
+   FAQ
+   api/modules
+   
diff --git a/docs/source/intro.rst b/docs/source/intro.rst
deleted file mode 100644
index 61415c3..0000000
--- a/docs/source/intro.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-Introduction
-============
-
-Installation
-------------
-
-To install flowMC, first install it using pip:
diff --git a/docs/source/logo_0810.png b/docs/source/logo_0810.png
new file mode 100644
index 0000000..e165063
Binary files /dev/null and b/docs/source/logo_0810.png differ
diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst
new file mode 100644
index 0000000..651e14c
--- /dev/null
+++ b/docs/source/quickstart.rst
@@ -0,0 +1,70 @@
+.. _quickstart-section-top:
+
+Quick Start
+============
+
+Installation
+------------
+
+The recommended way to install flowMC is using pip
+
+.. code-block::
+
+    pip install flowMC
+
+This will install the latest stable release and its dependencies.
+flowMC is based on `Jax <https://github.com/google/jax>`_ and `Flax <https://github.com/google/flax>`_.
+By default, installing flowMC will automatically install Jax and Flax available on `PyPI <https://pypi.org/>`_.
+Jax does not install GPU support by default.
+If you want to use GPU with Jax, you need to install Jax with GPU support according to `their document <pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html>`_.
+At the time of writing this documentation page, this is the command to install Jax with GPU support:
+
+.. code-block::
+
+    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+
+Basic Usage
+-----------
+
+A minimum example to sample a N dimensional Gaussian, you would do something like:
+
+.. code-block:: python
+
+    import jax
+    import jax.numpy as jnp
+    from flowMC.nfmodel.realNVP import RealNVP
+    from flowMC.sampler.MALA import mala_sampler
+    from flowMC.sampler.Sampler import Sampler
+    from flowMC.utils.PRNG_keys import initialize_rng_keys
+    from flowMC.nfmodel.utils import *
+
+    def log_prob(x):
+        return -0.5 * jnp.sum(x ** 2)
+
+    d_logp = jax.grad(log_prob)
+
+    n_dim = 5
+    n_chains = 10
+
+    rng_key_set = initialize_rng_keys(n_chains, seed=42)
+    initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
+
+    model = RealNVP(10, n_dim, 64, 1)
+    run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None), out_axes=0)
+
+    nf_sampler = Sampler(n_dim, rng_key_set, model, run_mcmc,
+                        log_prob,
+                        n_chains=n_chains,
+                        d_likelihood=d_logp,)
+
+    nf_sampler.sample(initial_position)
+
+For more realistic test case, see the examples on `github <https://github.com/kazewong/flowMC/tree/main/example>`_.
+In particular, currently the best engineered test case is `dualmoon.py <https://github.com/kazewong/flowMC/blob/main/example/dualmoon.py>`_.
+
+In the ideal case, the only two things you will have to do is:
+
+#. Write down the function you want to sample in the form of :code:`p(x)`, where :code:`p` is the log probability density function and :code:`x` is the vector of variables of interest.
+#. Configure the sampler.
+
diff --git a/example/NonJaxLikelihood.py b/example/NonJaxLikelihood.py
index c6ea20e..58c72ef 100644
--- a/example/NonJaxLikelihood.py
+++ b/example/NonJaxLikelihood.py
@@ -2,17 +2,18 @@
 from flowMC.sampler.Gaussian_random_walk import rw_metropolis_sampler
 from flowMC.sampler.MALA import mala_sampler
 import jax
-import jax.numpy as jnp                # JAX NumPy
-import numpy as np  
+import jax.numpy as jnp  # JAX NumPy
+import numpy as np
 from scipy.stats import norm
 from flowMC.sampler.Sampler import Sampler
 from flowMC.utils.PRNG_keys import initialize_rng_keys
 from flowMC.utils.PythonFunctionWrap import wrap_python_log_prob_fn
 
+
 @wrap_python_log_prob_fn
 def neal_funnel(x):
-    y_pdf = norm.logpdf(x[0],loc=0,scale=3)
-    x_pdf = norm.logpdf(x[1:],loc=0,scale=np.exp(x[0]/2))
+    y_pdf = norm.logpdf(x[0], loc=0, scale=3)
+    x_pdf = norm.logpdf(x[1:], loc=0, scale=np.exp(x[0] / 2))
     return y_pdf + np.sum(x_pdf)
 
 
@@ -29,31 +30,37 @@ def neal_funnel(x):
 logging = True
 
 print("Preparing RNG keys")
-rng_key_set = initialize_rng_keys(n_chains,seed=42)
+rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 print("Initializing MCMC model and normalizing flow model.")
 
-initial_position = jax.random.normal(rng_key_set[0],shape=(n_chains,n_dim)) #(n_dim, n_chains)
+initial_position = jax.random.normal(
+    rng_key_set[0], shape=(n_chains, n_dim)
+)  # (n_dim, n_chains)
 
 
-model = RealNVP(10,n_dim,64, 1)
-run_mcmc = jax.vmap(rw_metropolis_sampler, in_axes=(0, None, None, 0),
-                    out_axes=0)
+model = RealNVP(10, n_dim, 64, 1)
+run_mcmc = jax.vmap(rw_metropolis_sampler, in_axes=(0, None, None, 0), out_axes=0)
 
 print("Initializing sampler class")
 
-nf_sampler = Sampler(n_dim, rng_key_set, model, run_mcmc,
-                    neal_funnel,
-                    n_loop=n_loop,
-                    n_local_steps=n_local_steps,
-                    n_global_steps=n_global_steps,
-                    n_chains=n_chains,
-                    n_epochs=num_epochs,
-                    n_nf_samples=100,
-                    learning_rate=learning_rate,
-                    momentum=momentum,
-                    batch_size=batch_size,
-                    stepsize=stepsize)
+nf_sampler = Sampler(
+    n_dim,
+    rng_key_set,
+    model,
+    run_mcmc,
+    neal_funnel,
+    n_loop=n_loop,
+    n_local_steps=n_local_steps,
+    n_global_steps=n_global_steps,
+    n_chains=n_chains,
+    n_epochs=num_epochs,
+    n_nf_samples=100,
+    learning_rate=learning_rate,
+    momentum=momentum,
+    batch_size=batch_size,
+    stepsize=stepsize,
+)
 
 print("Sampling")
 
@@ -65,9 +72,11 @@ def neal_funnel(x):
 import matplotlib.pyplot as plt
 
 # Plot one chain to show the jump
-plt.plot(chains[70,:,0],chains[70,:,1])
+plt.plot(chains[70, :, 0], chains[70, :, 1])
 plt.show()
 plt.close()
 
 # Plot all chains
-corner.corner(chains.reshape(-1,n_dim), labels=["$x_1$", "$x_2$", "$x_3$", "$x_4$", "$x_5$"])
+corner.corner(
+    chains.reshape(-1, n_dim), labels=["$x_1$", "$x_2$", "$x_3$", "$x_4$", "$x_5$"]
+)
diff --git a/example/dualmoon.py b/example/dualmoon.py
index 067b5ce..bb08dd3 100644
--- a/example/dualmoon.py
+++ b/example/dualmoon.py
@@ -1,12 +1,16 @@
 from flowMC.nfmodel.realNVP import RealNVP
 from flowMC.nfmodel.rqSpline import RQSpline
 from flowMC.sampler.MALA import make_mala_sampler
+
 import jax
 import jax.numpy as jnp  # JAX NumPy
+from jax.scipy.special import logsumexp
+
+
+from flowMC.nfmodel.realNVP import RealNVP
+from flowMC.sampler.MALA import mala_sampler
 from flowMC.sampler.Sampler import Sampler
 from flowMC.utils.PRNG_keys import initialize_rng_keys
-from jax.scipy.special import logsumexp
-import numpy as np
 
 from flowMC.nfmodel.utils import *
 
@@ -26,15 +30,14 @@ def dual_moon_pe(x):
 ### Demo config
 
 n_dim = 5
-n_chains = 100
+n_chains = 10
 n_loop = 5
 n_local_steps = 100
 n_global_steps = 100
-learning_rate = 0.01
-max_samples = 10000
+learning_rate = 0.1
 momentum = 0.9
-num_epochs = 100
-batch_size = 10000
+num_epochs = 5
+batch_size = 50
 stepsize = 0.01
 
 print("Preparing RNG keys")
@@ -46,35 +49,41 @@ def dual_moon_pe(x):
 
 
 # model = RealNVP(10, n_dim, 64, 1)
-model = RQSpline(n_dim, 10, [128,128], 8)
+model = RQSpline(n_dim, 10, [128, 128], 8)
 
-local_sampler,updater, kernel,logp,dlogp = make_mala_sampler(dual_moon_pe, d_dual_moon,1e-1, jit=True, M=jnp.eye(n_dim))
+local_sampler, updater, kernel, logp, dlogp = make_mala_sampler(
+    dual_moon_pe, d_dual_moon, 1e-1, jit=True, M=jnp.eye(n_dim)
+)
 
 
 print("Initializing sampler class")
 
-nf_sampler = Sampler(n_dim, rng_key_set, model, local_sampler,
-                    dual_moon_pe,
-                    d_likelihood=d_dual_moon,
-                    n_loop=n_loop,
-                    n_local_steps=n_local_steps,
-                    n_global_steps=n_global_steps,
-                    n_chains=n_chains,
-                    n_epochs=num_epochs,
-                    max_samples=max_samples,
-                    n_nf_samples=100,
-                    learning_rate=learning_rate,
-                    momentum=momentum,
-                    batch_size=batch_size,
-                    stepsize=stepsize,
-                    use_global=True,)
+nf_sampler = Sampler(
+    n_dim,
+    rng_key_set,
+    model,
+    run_mcmc,
+    dual_moon_pe,
+    d_likelihood=d_dual_moon,
+    n_loop=n_loop,
+    n_local_steps=n_local_steps,
+    n_global_steps=n_global_steps,
+    n_chains=n_chains,
+    n_epochs=num_epochs,
+    n_nf_samples=100,
+    learning_rate=learning_rate,
+    momentum=momentum,
+    batch_size=batch_size,
+    stepsize=stepsize,
+    use_global=True,
+)
 
 print("Sampling")
-    
+
 nf_sampler.sample(initial_position)
 
 chains, log_prob, local_accs, global_accs, loss_vals = nf_sampler.get_sampler_state()
-nf_samples = nf_sampler.sample_flow(10000)
+nf_samples = nf_sampler.sample_flow()
 
 print(
     "chains shape: ",
@@ -86,7 +95,7 @@ def dual_moon_pe(x):
 )
 
 chains = np.array(chains)
-nf_samples = np.array(nf_samples[1][0])
+nf_samples = np.array(nf_samples[1])
 loss_vals = np.array(loss_vals)
 
 import corner
diff --git a/example/kepler/kepler.py b/example/kepler/kepler.py
deleted file mode 100644
index bd5404a..0000000
--- a/example/kepler/kepler.py
+++ /dev/null
@@ -1,221 +0,0 @@
-import jax
-import jax.numpy as jnp                # JAX NumPy
-from jax.scipy.special import logsumexp
-import numpy as np  
-import matplotlib.pyplot as plt
-from scipy.optimize import minimize
-import corner
-import tqdm
-import time
-import pickle
-
-from flowMC.nfmodel.realNVP import RealNVP
-from flowMC.sampler.MALA import mala_sampler
-from flowMC.nfmodel.utils import *
-
-from utils import rv_model, log_likelihood, log_prior, sample_prior, get_kepler_params_and_log_jac
-from flowMC.sampler.Sampler import Sampler
-from flowMC.utils.PRNG_keys import initialize_rng_keys
-
-jax.config.update("jax_enable_x64", True)
-
-## Generate probelm
-true_params = jnp.array([
-    12.0, # v0
-    np.log(0.5), # log_s2
-    np.log(14.5), # log_P
-    np.log(2.3), # log_k
-    np.sin(1.5), # phi
-    np.cos(1.5),
-    0.4, # ecc
-    np.sin(-0.7), # w
-    np.cos(-0.7)
-])
-
-prior_kwargs = {
-    'ecc_alpha': 2, 'ecc_beta': 2,
-    'log_k_mean': 1, 'log_k_var': 1,
-    'v0_mean': 10, 'v0_var': 2,
-    'log_period_mean': 2.5, 'log_period_var': 0.5,
-    'log_s2_mean': -0.5, 'log_s2_var': 0.1,
-}
-
-random = np.random.default_rng(12345)
-t = np.sort(random.uniform(0, 100, 50))
-rv_err = 0.3
-sigma2 = rv_err ** 2 + jnp.exp(2 * true_params[1])
-rv_obs = rv_model(true_params, t) + random.normal(0, sigma2, len(t))
-# plt.plot(t, rv_obs, ".k")
-# x = np.linspace(0, 100, 500)
-# plt.plot(x, rv_model(true_params, x), "C0")
-# plt.show(block=False)
-
-## Setting up sampling -- takes one input at the time.
-def log_posterior(x):
-    return  log_likelihood(x.T, t, rv_err, rv_obs) + log_prior(x,**prior_kwargs)
-    
-d_log_posterior = jax.grad(log_posterior)
-
-config = {}
-n_dim = 9
-n_loop = 2
-n_local_steps = 5
-n_global_steps = 2
-n_chains = 5
-learning_rate = 0.01
-momentum = 0.9
-num_epochs = 5
-batch_size = 10
-stepsize = 1e-3
-logging = True
-
-print("Preparing RNG keys")
-rng_key_set = initialize_rng_keys(n_chains,seed=42)
-
-print("Initializing MCMC model and normalizing flow model.")
-
-# initial_position = jax.random.normal(rng_key_set[0],shape=(n_chains,n_dim)) #(n_chains, n_dim)
-
-kepler_params_ini = sample_prior(rng_key_set[0], n_chains,
-                                 **prior_kwargs)
-neg_logp_and_grad = jax.jit(jax.value_and_grad(lambda p: -log_posterior(p)))
-optimized = []
-for i in tqdm.tqdm(range(n_chains)):
-    soln = minimize(neg_logp_and_grad, kepler_params_ini.T[i].T, jac=True)
-    optimized.append(jnp.asarray(get_kepler_params_and_log_jac(soln.x)[0]))
-
-initial_position = jnp.stack(optimized) #(n_chains, n_dim)
-# initial_position = kepler_params_ini.T
-
-mean = initial_position.mean(0)
-init_centered = (initial_position - mean)
-cov = init_centered.T @ init_centered / n_chains
-
-model = RealNVP(10, n_dim, 64, 1)
-
-run_mala = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None),
-                    out_axes=0)
-
-print("Initializing sampler class")
-
-nf_sampler = Sampler(n_dim, rng_key_set, model, run_mala,
-                    log_posterior,
-                    d_likelihood=d_log_posterior,
-                    n_loop=n_loop,
-                    n_local_steps=n_local_steps,
-                    n_global_steps=n_global_steps,
-                    n_chains=n_chains,
-                    n_epochs=num_epochs,
-                    n_nf_samples=100,
-                    learning_rate=learning_rate,
-                    momentum=momentum,
-                    batch_size=batch_size,
-                    stepsize=stepsize)
-
-
-print("Sampling")
-
-start = time.time()
-_ = nf_sampler.sample(initial_position)
-chains, nf_samples, local_accs, global_accs, loss_vals = _
-
-print('Elapsed: ', time.time()-start, 's')
-
-print("Make plots")
-
-chains = np.array(chains)
-nf_samples = np.array(nf_samples)
-
-value1 = true_params
-n_dim = n_dim
-# Make the base corner plot
-
-figure = corner.corner(chains.reshape(-1,n_dim),
-                labels=['v0', 'log_s2', 'log_period', 'log_k', 'sin_phi_', 'cos_phi_', 'ecc_', 'sin_w_', 'cos_w_'], title_kwargs={"fontsize": 12})
-figure.set_size_inches(7, 7)
-figure.suptitle('Visualize chains')
-# Extract the axes
-axes = np.array(figure.axes).reshape((n_dim, n_dim))
-# Loop over the diagonal
-for i in range(n_dim):
-    ax = axes[i, i]
-    ax.axvline(value1[i], color="g")
-# Loop over the histograms
-for yi in range(n_dim):
-    for xi in range(yi):
-        ax = axes[yi, xi]
-        ax.plot(value1[xi], value1[yi], "sg")
-        ax.plot(chains[1, -1000:, xi],chains[1, -1000:, yi], alpha=0.75, lw=0.75)
-        ax.plot(chains[0, -1000:, xi],chains[0, -1000:, yi], alpha=0.75, lw=0.75)
-        ax.axvline(value1[xi], color="g")
-        ax.axhline(value1[yi], color="g")
-# plt.tight_layout()
-plt.show(block=False)
-
-plt.figure(figsize=(10,5))
-axs = [plt.subplot(121), plt.subplot(122)]
-plt.sca(axs[0])
-plt.plot(t, rv_obs, ".k", label='observations')
-x = np.linspace(0, 100, 500)
-plt.plot(x, rv_model(true_params, x), "C0", label='ground truth')
-for i in range(np.minimum(n_chains,10)):
-    params, log_jac = get_kepler_params_and_log_jac(chains[i,-1,:])
-    if i == 0:
-        plt.plot(x, rv_model(params, x), c='gray', alpha=0.5, label='final samples')
-    else:
-        plt.plot(x, rv_model(params, x), c='gray', alpha=0.5)
-plt.xlabel('t')
-plt.ylabel('radial velocity')
-plt.legend()
-
-plt.sca(axs[1])
-posterior_evolution = jax.vmap(log_posterior)(chains[-10:,:,:].reshape(-1, 9)).reshape(10, -1)
-shift = max(jnp.max(posterior_evolution),0)
-plt.plot(- (posterior_evolution.T - shift))
-plt.yscale('log')
-plt.ylabel('walker negative log-likelihood')
-plt.xlabel('iteration')
-
-plt.show(block=False)
-
-value1 = true_params
-n_dim = n_dim
-# Make the base corner plot
-figure = corner.corner(chains.reshape(-1,n_dim),
-                labels=['v0', 'log_s2', 'log_period', 'log_k', 'sin_phi_', 'cos_phi_', 'ecc_', 'sin_w_', 'cos_w_'], title_kwargs={"fontsize": 12})
-figure.set_size_inches(7, 7)
-figure.suptitle('Visualize initializations')
-# Extract the axes
-axes = np.array(figure.axes).reshape((n_dim, n_dim))
-# Loop over the diagonal
-for i in range(n_dim):
-    ax = axes[i, i]
-    ax.axvline(value1[i], color="g")
-# Loop over the histograms
-for yi in range(n_dim):
-    for xi in range(yi):
-        ax = axes[yi, xi]
-        ax.plot(value1[xi], value1[yi], "sg")
-        ax.plot(initial_position[:,xi], initial_position[:,yi],'+',ms=2)
-        ax.axvline(value1[xi], color="g")
-        ax.axhline(value1[yi], color="g")
-# plt.tight_layout()
-plt.show(block=False)
-
-
-results = {
-    'chains': chains,
-    'nf_samples': nf_samples,
-    'prior_samples': kepler_params_ini,
-    'optimized_init': initial_position,
-    'config': config,
-    'true_params': true_params,
-    'rv_obs': rv_obs,
-    't': t,
-    'prior_kwargs': prior_kwargs,
-    'rv_err': rv_err
-}
-
-random_id = np.random.randint(10000)
-with open('results_{:d}.pkl'.format(random_id), 'wb') as f:
-    pickle.dump(results, f)
\ No newline at end of file
diff --git a/example/kepler/utils.py b/example/kepler/utils.py
deleted file mode 100644
index a7b21e1..0000000
--- a/example/kepler/utils.py
+++ /dev/null
@@ -1,192 +0,0 @@
-import jax
-import jax.numpy as jnp
-import numpy as np
-import matplotlib.pyplot as plt
-
-@jax.custom_jvp
-def kepler(M, ecc):
-    """
-    Given a mean anomaly 'M' and eccentricity 'ecc', solve for the sin and cos
-    of the true anomaly 'f'.
-    """
-    # Wrap into the right range
-    M = M % (2 * jnp.pi)
-
-    # We can restrict to the range [0, pi)
-    high = M > jnp.pi
-    M = jnp.where(high, 2 * jnp.pi - M, M)
-
-    # Solve
-    ome = 1 - ecc
-    E = starter(M, ecc, ome)
-    E = refine(M, ecc, ome, E)
-
-    # Re-wrap back into the full range
-    E = jnp.where(high, 2 * jnp.pi - E, E)
-
-    # Convert to true anomaly; tan(0.5 * f)
-    tan_half_f = jnp.sqrt((1 + ecc) / (1 - ecc)) * jnp.tan(0.5 * E)
-    tan2_half_f = jnp.square(tan_half_f)
-
-    # Then we compute sin(f) and cos(f) using:
-    #  sin(f) = 2*tan(0.5*f)/(1 + tan(0.5*f)^2), and
-    #  cos(f) = (1 - tan(0.5*f)^2)/(1 + tan(0.5*f)^2)
-    denom = 1 / (1 + tan2_half_f)
-    sinf = 2 * tan_half_f * denom
-    cosf = (1 - tan2_half_f) * denom
-
-    return sinf, cosf
-
-
-# This defines the differentation rule
-@kepler.defjvp
-def kepler_jvp(primals, tangents):
-    M, e = primals
-    dM, de = tangents
-    sinf, cosf = kepler(M, e)
-
-    # Pre-compute some things
-    ecosf = e * cosf
-    ome2 = 1 - e ** 2
-
-    # Propagate the derivatives
-    df = 0.0
-    if type(dM) is not jax.interpreters.ad.Zero:
-        df += dM * (1 + ecosf) ** 2 / ome2 ** 1.5
-    if type(de) is not jax.interpreters.ad.Zero:
-        df += de * (2 + ecosf) * sinf / ome2
-
-    return (sinf, cosf), (cosf * df, -sinf * df)
-
-
-# The following two functions are helpers for the Kepler solver
-def starter(M, ecc, ome):
-    M2 = jnp.square(M)
-    alpha = 3 * jnp.pi / (jnp.pi - 6 / jnp.pi)
-    alpha += 1.6 / (jnp.pi - 6 / jnp.pi) * (jnp.pi - M) / (1 + ecc)
-    d = 3 * ome + alpha * ecc
-    alphad = alpha * d
-    r = (3 * alphad * (d - ome) + M2) * M
-    q = 2 * alphad * ome - M2
-    q2 = jnp.square(q)
-    w = (jnp.abs(r) + jnp.sqrt(q2 * q + r * r)) ** (2.0 / 3)
-    return (2 * r * w / (jnp.square(w) + w * q + q2) + M) / d
-
-
-def refine(M, ecc, ome, E):
-    sE = E - jnp.sin(E)
-    cE = 1 - jnp.cos(E)
-
-    f_0 = ecc * sE + E * ome - M
-    f_1 = ecc * cE + ome
-    f_2 = ecc * (E - sE)
-    f_3 = 1 - f_1
-    d_3 = -f_0 / (f_1 - 0.5 * f_0 * f_2 / f_1)
-    d_4 = -f_0 / (f_1 + 0.5 * d_3 * f_2 + (d_3 * d_3) * f_3 / 6)
-    d_42 = d_4 * d_4
-    dE = -f_0 / (
-        f_1 + 0.5 * d_4 * f_2 + d_4 * d_4 * f_3 / 6 - d_42 * d_4 * f_2 / 24
-    )
-
-    return E + dE
-
-
-def rv_model(kepler_params_, t):
-    v0, log_s2, log_period, log_k, sin_phi, cos_phi, ecc, sin_w, cos_w = kepler_params_
-    phi = jnp.arctan2(sin_phi, cos_phi)
-    sin_f, cos_f = kepler(2 * np.pi * t * jnp.exp(-log_period) + phi, ecc)
-    return v0 + jnp.exp(log_k) * (cos_w * cos_f - sin_w * sin_f + ecc * cos_w)
-
-
-def get_kepler_params_and_log_jac(kepler_params):
-    v0, log_s2, log_period, log_k, sin_phi_, cos_phi_, ecc_, sin_w_, cos_w_ = kepler_params
-    
-    # Log odds transform for eccentricity, since it needs to be in the range [0, 1)
-    ecc = 1.0 / (1.0 + jnp.exp(-ecc_))
-    log_jac = -2 * jnp.log(jnp.cosh(0.5 * ecc_))
-    # -2 * jnp.log(2 * jnp.cosh(0.5 * ecc_))? factor 2? to keep exact logjac with NF
-
-    # Normalize the unit vector parameterizations of the angles
-    phi_norm = sin_phi_ ** 2 + cos_phi_ ** 2
-    log_jac += -0.5 * phi_norm
-    phi_norm = jnp.sqrt(phi_norm)
-
-    w_norm = sin_w_ ** 2 + cos_w_ ** 2
-    log_jac += -0.5 * w_norm
-    w_norm = jnp.sqrt(w_norm)
-
-    return (
-        (v0, log_s2, log_period, log_k, sin_phi_ / phi_norm, cos_phi_ / phi_norm, ecc, sin_w_ / w_norm, cos_w_ / w_norm),
-        log_jac
-    )
-
-
-def log_likelihood(kepler_params, t, rv_err, rv_obs):
-    params, log_jac = get_kepler_params_and_log_jac(kepler_params)
-    v = rv_model(params, t)
-    log_s2 = params[1]
-    sigma2 = rv_err ** 2 + jnp.exp(2 * log_s2)
-    loglike = -0.5 * jnp.square(rv_obs - v) / sigma2 - 0.5 * jnp.log(sigma2)
-    return log_jac + jnp.sum(loglike)
-
-
-## we also need priors
-def log_prior(kepler_params,
-              ecc_alpha=2, ecc_beta=2,
-              log_k_mean=1, log_k_var=1,
-              v0_mean=0, v0_var=1,
-              log_period_mean=1, log_period_var=1,
-              log_s2_mean=1, log_s2_var=1e-2,
-              ):
-    v0, log_s2, log_period, log_k, sin_phi_, cos_phi_, ecc_, sin_w_, cos_w_ = kepler_params
-
-    logp_phis_ = - 0.5 * (sin_phi_ ** 2 + cos_phi_ ** 2)
-    logp_ws_ = - 0.5 * (sin_w_ ** 2 + cos_w_ ** 2)
-
-    ecc = 1.0 / (1.0 + jnp.exp(- ecc_))
-    logp_ecc = -2 * jnp.log(jnp.cosh(0.5 * ecc_)) 
-    logp_ecc += (ecc_alpha - 1) * jnp.log(ecc) 
-    logp_ecc += (ecc_beta - 1) * jnp.log(1 - ecc)
-
-    logp_log_k = - 0.5 * (log_k - log_k_mean) ** 2 / log_k_var
-    logp_v0 = - 0.5 * (v0 - v0_mean) ** 2 / v0_var
-    logp_log_period = - 0.5 * (log_period - log_period_mean) ** 2 / log_period_var
-    logp_log_s2 = - 0.5 * (log_s2 - log_s2_mean) ** 2 / log_s2_var
-
-    return logp_phis_ + logp_ws_ + logp_ecc + logp_log_k + logp_v0 + logp_log_period + logp_log_s2
-
-
-def sample_prior(rng_key, n_samples, ecc_alpha=2, ecc_beta=2,
-                 log_k_mean=1, log_k_var=1,
-                 v0_mean=0, v0_var=1,
-                 log_period_mean=1, log_period_var=1,
-                 log_s2_mean=1, log_s2_var=1e-2):
-    """
-    returns kepler_params
-    """
-    rng_key, sub_key = jax.random.split(rng_key)
-    v0 = jax.random.normal(sub_key, (n_samples,)) * jnp.sqrt(v0_var) + v0_mean
-    
-    rng_key, sub_key = jax.random.split(rng_key)
-    log_s2 = jax.random.normal(sub_key, (n_samples,)) * jnp.sqrt(log_s2_var)
-    log_s2 += log_s2_mean
-    
-    rng_key, sub_key = jax.random.split(rng_key)
-    log_k = jax.random.normal(sub_key, (n_samples,)) * jnp.sqrt(log_k_var)
-    log_k += log_k_mean
-    
-    rng_key, sub_key = jax.random.split(rng_key)
-    log_period = jax.random.normal(sub_key, (n_samples,)) * jnp.sqrt(log_period_var)
-    log_period += log_period_mean
-    
-    rng_key, sub_key = jax.random.split(rng_key)
-    ecc = jax.random.beta(sub_key, ecc_alpha, ecc_beta, (n_samples,))
-    ecc_ = jnp.log(ecc / (1 - ecc))
-
-    rng_key, sub_key = jax.random.split(rng_key)
-    sin_phi_, cos_phi_, sin_w_, cos_w_ = jax.random.multivariate_normal(sub_key, jnp.zeros(4), jnp.eye(4), (n_samples,) ).T
-    
-    kepler_params = v0, log_s2, log_period, log_k, sin_phi_, cos_phi_, ecc_, sin_w_, cos_w_
-
-    return jnp.stack(kepler_params) # shape = (9, n_samples)
-
diff --git a/example/mixture_gaussian.py b/example/mixture_gaussian.py
index b9fe0a4..92e62a8 100644
--- a/example/mixture_gaussian.py
+++ b/example/mixture_gaussian.py
@@ -17,7 +17,7 @@
 from jax.scipy.stats import norm, multivariate_normal
 from jax.scipy.special import logsumexp
 from functools import partial
-from flowMC.sampler.MALA import mala_kernel,mala_sampler
+from flowMC.sampler.MALA import mala_kernel, mala_sampler
 
 
 def mixture_logpdf(x):
@@ -31,16 +31,17 @@ def mixture_logpdf(x):
     float
         The value of the log probability density function at x.
     """
-    cov = jnp.repeat(jnp.eye(2)[None,:],1,axis=0)
+    cov = jnp.repeat(jnp.eye(2)[None, :], 1, axis=0)
     mean = jnp.array([[0.0, 0.0]])
-    dist_1 = partial(multivariate_normal.logpdf, mean=mean+5.0, cov=cov)
-    dist_2 = partial(multivariate_normal.logpdf, mean=mean-5.0, cov=cov)
+    dist_1 = partial(multivariate_normal.logpdf, mean=mean + 5.0, cov=cov)
+    dist_2 = partial(multivariate_normal.logpdf, mean=mean - 5.0, cov=cov)
     # dist_3 = partial(norm.logpdf, loc=3.2, scale=5)
     # dist_4 = partial(norm.logpdf, loc=2.5, scale=2.8)
-    log_probs = jnp.array([dist_1(x), dist_2(x)])#, dist_3(x), dist_4(x)])
-    weights = jnp.array([0.3, 0.7])#, 0.1, 0.4])
+    log_probs = jnp.array([dist_1(x), dist_2(x)])  # , dist_3(x), dist_4(x)])
+    weights = jnp.array([0.3, 0.7])  # , 0.1, 0.4])
     return logsumexp(jnp.log(weights) + log_probs)
 
+
 d_pdf = jax.grad(mixture_logpdf)
 
 samples = 30000
@@ -55,15 +56,17 @@ def mixture_logpdf(x):
 rng_keys = jax.random.split(rng_key, n_chains)  # (nchains,)
 initial_position = jnp.zeros((n_dim, n_chains))  # (n_dim, n_chains)
 
-run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 1, None),
-                    out_axes=0)
-rng_key, positions, log_prob = run_mcmc(rng_keys, n_samples, mixture_logpdf, d_pdf, initial_position, 0.1)
+run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 1, None), out_axes=0)
+rng_key, positions, log_prob = run_mcmc(
+    rng_keys, n_samples, mixture_logpdf, d_pdf, initial_position, 0.1
+)
 assert positions.shape == (n_chains, n_samples, n_dim)
 positions.block_until_ready()
 
-flat_chain = np.concatenate(positions,axis=0)#[:,1000:],axis=0)
+flat_chain = np.concatenate(positions, axis=0)  # [:,1000:],axis=0)
 
 import matplotlib.pyplot as plt
-plt.figure(figsize=(10,10))
-plt.scatter(flat_chain[:,0],flat_chain[:,1],s=0.1)
-plt.show()
\ No newline at end of file
+
+plt.figure(figsize=(10, 10))
+plt.scatter(flat_chain[:, 0], flat_chain[:, 1], s=0.1)
+plt.show()
diff --git a/example/neal_funnel.py b/example/neal_funnel.py
index 717ab89..db6dc7c 100644
--- a/example/neal_funnel.py
+++ b/example/neal_funnel.py
@@ -1,4 +1,3 @@
-
 """
 Working Neal Funnel.
 We use a mapping to and from a normal distribution to efficiently sample from the Neal funnel.
@@ -11,17 +10,19 @@
 from flowMC.nfmodel.realNVP import RealNVP
 from flowMC.sampler.MALA import mala_sampler
 import jax
-import jax.numpy as jnp                # JAX NumPy
-import numpy as np  
+import jax.numpy as jnp  # JAX NumPy
+import numpy as np
 from jax.scipy.stats import norm
 from flowMC.sampler.Sampler import Sampler
 from flowMC.utils.PRNG_keys import initialize_rng_keys
 
+
 def neal_funnel(x):
-    y_pdf = norm.logpdf(x[0],loc=0,scale=3)
-    x_pdf = norm.logpdf(x[1:],loc=0,scale=jnp.exp(x[0]/2))
+    y_pdf = norm.logpdf(x[0], loc=0, scale=3)
+    x_pdf = norm.logpdf(x[1:], loc=0, scale=jnp.exp(x[0] / 2))
     return y_pdf + jnp.sum(x_pdf)
 
+
 d_neal_funnel = jax.grad(neal_funnel)
 
 n_dim = 5
@@ -37,32 +38,38 @@ def neal_funnel(x):
 logging = True
 
 print("Preparing RNG keys")
-rng_key_set = initialize_rng_keys(n_chains,seed=42)
+rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 print("Initializing MCMC model and normalizing flow model.")
 
-initial_position = jax.random.normal(rng_key_set[0],shape=(n_chains,n_dim)) #(n_dim, n_chains)
+initial_position = jax.random.normal(
+    rng_key_set[0], shape=(n_chains, n_dim)
+)  # (n_dim, n_chains)
 
 
-model = RealNVP(10,n_dim,64, 1)
-run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None),
-                    out_axes=0)
+model = RealNVP(10, n_dim, 64, 1)
+run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None), out_axes=0)
 
 print("Initializing sampler class")
 
-nf_sampler = Sampler(n_dim, rng_key_set, model, run_mcmc,
-                    neal_funnel,
-                    d_likelihood=d_neal_funnel,
-                    n_loop=n_loop,
-                    n_local_steps=n_local_steps,
-                    n_global_steps=n_global_steps,
-                    n_chains=n_chains,
-                    n_epochs=num_epochs,
-                    n_nf_samples=100,
-                    learning_rate=learning_rate,
-                    momentum=momentum,
-                    batch_size=batch_size,
-                    stepsize=stepsize)
+nf_sampler = Sampler(
+    n_dim,
+    rng_key_set,
+    model,
+    run_mcmc,
+    neal_funnel,
+    d_likelihood=d_neal_funnel,
+    n_loop=n_loop,
+    n_local_steps=n_local_steps,
+    n_global_steps=n_global_steps,
+    n_chains=n_chains,
+    n_epochs=num_epochs,
+    n_nf_samples=100,
+    learning_rate=learning_rate,
+    momentum=momentum,
+    batch_size=batch_size,
+    stepsize=stepsize,
+)
 
 print("Sampling")
 
@@ -72,9 +79,11 @@ def neal_funnel(x):
 import matplotlib.pyplot as plt
 
 # Plot one chain to show the jump
-plt.plot(chains[70,:,0],chains[70,:,1])
+plt.plot(chains[70, :, 0], chains[70, :, 1])
 plt.show()
 plt.close()
 
 # Plot all chains
-corner.corner(chains.reshape(-1,n_dim), labels=["$x_1$", "$x_2$", "$x_3$", "$x_4$", "$x_5$"])
+corner.corner(
+    chains.reshape(-1, n_dim), labels=["$x_1$", "$x_2$", "$x_3$", "$x_4$", "$x_5$"]
+)
diff --git a/example/train_RealNVP.py b/example/train_RealNVP.py
index a887755..da8cf2f 100644
--- a/example/train_RealNVP.py
+++ b/example/train_RealNVP.py
@@ -1,14 +1,13 @@
-
 from flowMC.nfmodel.realNVP import RealNVP
 import jax
-import jax.numpy as jnp                # JAX NumPy
-import jax.random as random            # JAX random
+import jax.numpy as jnp  # JAX NumPy
+import jax.random as random  # JAX random
 from jax.scipy.stats import multivariate_normal
 
 from flowMC.nfmodel.utils import *
-from flax import linen as nn           # The Linen API
+from flax import linen as nn  # The Linen API
 from flax.training import train_state  # Useful dataclass to keep train state
-import optax                           # Optimizers
+import optax  # Optimizers
 
 from sklearn.datasets import make_moons
 import matplotlib.pyplot as plt
@@ -28,26 +27,30 @@
 model = RealNVP(10, 2, n_hidden, 1)
 data = make_moons(n_samples=100000, noise=0.05)
 
-key1, rng, init_rng = jax.random.split(jax.random.PRNGKey(0),3)
+key1, rng, init_rng = jax.random.split(jax.random.PRNGKey(0), 3)
+
 
 def create_train_state(rng, learning_rate, momentum):
-    params = model.init(rng, jnp.ones((1,2)))['params']
+    params = model.init(rng, jnp.ones((1, 2)))["params"]
     tx = optax.adam(learning_rate, momentum)
     return train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)
 
+
 state = create_train_state(init_rng, learning_rate, momentum)
 
-variables = model.init(rng, jnp.ones((1,2)))['variables']
+variables = model.init(rng, jnp.ones((1, 2)))["variables"]
 
 
 data = jnp.array(data[0])
 
-rng, state, loss_values = train_flow(rng, model, state, data, num_epochs, batch_size, variables)
+rng, state, loss_values = train_flow(
+    rng, model, state, data, num_epochs, batch_size, variables
+)
 
 
 rng_key_nf = jax.random.PRNGKey(124098)
 nf_samples = sample_nf(model, state.params, rng_key_nf, 10000, variables)
 plt.figure(figsize=(10, 9))
-plt.scatter(data[:,0], data[:,1], s=0.1, c='r',label='Data')
-plt.scatter(nf_samples[1][0][:,0], nf_samples[1][0][:,1], s=0.1, c='b',label='NF')
-plt.show()
\ No newline at end of file
+plt.scatter(data[:, 0], data[:, 1], s=0.1, c="r", label="Data")
+plt.scatter(nf_samples[1][0][:, 0], nf_samples[1][0][:, 1], s=0.1, c="b", label="NF")
+plt.show()
diff --git a/example/wiggle.py b/example/wiggle.py
index f2b9faa..9c67586 100644
--- a/example/wiggle.py
+++ b/example/wiggle.py
@@ -1,101 +1,79 @@
-from flowMC.nfmodel.realNVP import RealNVP
-from flowMC.sampler.MALA import mala_sampler
+import numpy as np
+
 import jax
-import jax.numpy as jnp                # JAX NumPy
-from flowMC.utils import Sampler, initialize_rng_keys
+import jax.numpy as jnp  # JAX NumPy
 from jax.scipy.special import logsumexp
-import numpy as np  
+
+from flowMC.nfmodel.realNVP import RealNVP
+from flowMC.sampler.MALA import mala_sampler
+from flowMC.sampler.Sampler import Sampler
+from flowMC.utils.PRNG_keys import initialize_rng_keys
 
 from flowMC.nfmodel.utils import *
+
 import matplotlib.pyplot as plt
 
 
-def log_prob(x):
+def wiggle(x):
     """
-    Wiggle-distribution 
+    Wiggle-distribution
     """
     mean = jnp.array([0, 6])
     centered_x = x[0] - mean - jnp.sin(5 * x[1] / 5)
-    log_prob = - 0.5 * centered_x @ jnp.eye(2) @ centered_x.T
+    log_prob = -0.5 * centered_x @ jnp.eye(2) @ centered_x.T
     log_prob -= 0.5 * (jnp.linalg.norm(x) - 5) ** 2 / 8
     return log_prob
 
-def plot_2d_level(log_prob, dim=2, x_min=-10, x_max=10,
-                  y_min=None, y_max=None,
-                  n_points=1000, ax=None, title=''):
-    """
-    Args:
-    model (RealNVP_MLP or MoG): must have a .sample and .U method
-    """
-    # if 
-    x_range = np.linspace(x_min, x_max, n_points)
-    if y_min is None:
-        y_range = x_range
-    else:
-        y_range = np.linspace(y_min, y_max, n_points)
-
-    grid = np.meshgrid(x_range, y_range)
-    xys = np.stack(grid).reshape(2, n_points ** 2).T
-    if dim > 2:
-        blu = np.zeros(n_points ** 2, dim)
-        blu[:, 0:2] = xys
-        xys = blu
-
-    log_probs = jax.vmap(log_prob)(xys).reshape(n_points, n_points)
-    
-    if ax is None:
-        plt.figure()
-    else:
-        plt.sca(ax)
-    plt.contourf(x_range, y_range, np.exp(log_probs), 10)
-    # plt.axis('off')
-    plt.colorbar()
-    plt.title(title)
-    return log_probs
-
-d_log_prob = jax.grad(log_prob)
-
-config = {}
-config['n_dim'] = 2
-config['n_loop'] = 10
-config['n_local_steps'] = 20
-config['n_global_steps'] = 100
-config['n_chains'] = 100
-config['learning_rate'] = 0.01
-config['momentum'] = 0.9
-config['num_epochs'] = 10
-config['batch_size'] = 1000
-config['stepsize'] = 0.01
-config['logging'] = True
+
+d_wiggle = jax.grad(wiggle)
+
+### Demo config
+
+n_dim = 5
+n_chains = 10
+n_loop = 5
+n_local_steps = 100
+n_global_steps = 100
+learning_rate = 0.1
+momentum = 0.9
+num_epochs = 5
+batch_size = 50
+stepsize = 0.01
 
 print("Preparing RNG keys")
-rng_key_set = initialize_rng_keys(config['n_chains'],seed=42)
+rng_key_set = initialize_rng_keys(n_chains, seed=42)
 
 print("Initializing MCMC model and normalizing flow model.")
 
-initial_position = jax.random.normal(rng_key_set[0],shape=(config['n_chains'],config['n_dim'])) #(n_dim, n_chains)
+initial_position = jax.random.normal(rng_key_set[0], shape=(n_chains, n_dim)) * 1
 
 
-model = RealNVP(10, config['n_dim'], 64, 1)
-run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None),
-                    out_axes=0)
+model = RealNVP(10, n_dim, 64, 1)
+run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None), out_axes=0)
 
 print("Initializing sampler class")
 
-nf_sampler = Sampler(rng_key_set, config, model, run_mcmc, log_prob, d_log_prob)
-
+nf_sampler = Sampler(
+    n_dim,
+    rng_key_set,
+    model,
+    run_mcmc,
+    wiggle,
+    d_likelihood=d_wiggle,
+    n_loop=n_loop,
+    n_local_steps=n_local_steps,
+    n_global_steps=n_global_steps,
+    n_chains=n_chains,
+    n_epochs=num_epochs,
+    n_nf_samples=100,
+    learning_rate=learning_rate,
+    momentum=momentum,
+    batch_size=batch_size,
+    stepsize=stepsize,
+    use_global=True,
+)
 print("Sampling")
 
-chains, nf_samples = nf_sampler.sample(initial_position)
-
-import corner
-import matplotlib.pyplot as plt
-
-# Plot one chain to show the jump
-plot_2d_level(log_prob)
-plt.plot(chains[70,:,0],chains[70,:,1])
-plt.show(block=False)
+nf_sampler.sample(initial_position)
 
-# Plot all chains
-corner.corner(chains.reshape(-1, config['n_dim']), labels=["$x_1$", "$x_2$"])
-plt.show(block=False)
+chains, log_prob, local_accs, global_accs, loss_vals = nf_sampler.get_sampler_state()
diff --git a/readthedocs.yml b/readthedocs.yml
new file mode 100644
index 0000000..3ad6b15
--- /dev/null
+++ b/readthedocs.yml
@@ -0,0 +1,8 @@
+version: 2
+
+python:
+  install:
+    - requirements: docs/requirements.txt
+
+sphinx:
+  configuration: docs/source/conf.py
diff --git a/setup.cfg b/setup.cfg
index fe7de6d..e534490 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -16,7 +16,9 @@ install_requires =
     jax
     jaxlib
     flax
+    tqdm
     distrax
+    
 python_requires =  >=3.7
 
 [options.packages.find]
diff --git a/src/flowMC/nfmodel/maf.py b/src/flowMC/nfmodel/maf.py
index 813fba7..9945c91 100644
--- a/src/flowMC/nfmodel/maf.py
+++ b/src/flowMC/nfmodel/maf.py
@@ -1,4 +1,4 @@
-from typing import (Any, Callable, Optional, Tuple, Iterable, Union)
+from typing import Any, Callable, Optional, Tuple, Iterable, Union
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
@@ -7,6 +7,7 @@
 from functools import reduce
 from jax import lax
 
+
 def get_masks(input_dim, hidden_dim=64, num_hidden=1):
     masks = []
     input_degrees = jnp.arange(input_dim)
@@ -17,9 +18,14 @@ def get_masks(input_dim, hidden_dim=64, num_hidden=1):
     degrees += [input_degrees % input_dim - 1]
 
     for (d0, d1) in zip(degrees[:-1], degrees[1:]):
-        masks += [jnp.transpose(jnp.expand_dims(d1, -1) >= jnp.expand_dims(d0, 0)).astype(jnp.float32)]
+        masks += [
+            jnp.transpose(jnp.expand_dims(d1, -1) >= jnp.expand_dims(d0, 0)).astype(
+                jnp.float32
+            )
+        ]
     return masks
 
+
 class MaskedDense(nn.Module):
     n_dim: int
     n_hidden: int
@@ -28,10 +34,11 @@ class MaskedDense(nn.Module):
 
     @nn.compact
     def __call__(self, x, mask):
-        weight = self.param('weights', self.kernel_init, (self.n_dim, self.n_hidden))
-        bias = self.param('bias', self.bias_init, (self.n_hidden,))
+        weight = self.param("weights", self.kernel_init, (self.n_dim, self.n_hidden))
+        bias = self.param("bias", self.bias_init, (self.n_hidden,))
         return jnp.dot(x, weight * mask) + bias
 
+
 class MaskedAutoEncoder(nn.Module):
     n_dim: int
     n_hidden: int
@@ -40,11 +47,11 @@ def setup(self):
         self.mask = get_masks(self.n_dim, self.n_hidden)
         self.up = MaskedDense(self.n_dim, self.n_hidden)
         self.mid = MaskedDense(self.n_hidden, self.n_hidden)
-        self.down = MaskedDense(self.n_hidden, 2*self.n_dim)
+        self.down = MaskedDense(self.n_hidden, 2 * self.n_dim)
 
     def __call__(self, inputs):
         log_weight, bias = self.forward(inputs)
-        outputs = (inputs - bias)*jnp.exp(-log_weight)
+        outputs = (inputs - bias) * jnp.exp(-log_weight)
         log_jacobian = -jnp.sum(log_weight, axis=-1)
         return outputs, log_jacobian
 
@@ -61,11 +68,14 @@ def inverse(self, inputs):
         for i_col in range(inputs.shape[1]):
             log_weight, bias = self.forward(outputs)
             outputs = jax.ops.index_update(
-                outputs, jax.ops.index[:, i_col], inputs[:, i_col] * jnp.exp(log_weight[:, i_col]) + bias[:, i_col]
+                outputs,
+                jax.ops.index[:, i_col],
+                inputs[:, i_col] * jnp.exp(log_weight[:, i_col]) + bias[:, i_col],
             )
         log_det_jacobian = -log_weight.sum(-1)
         return outputs, log_det_jacobian
 
+
 class MaskedAutoregressiveFlow(nn.Module):
     """
     Masked Autoregressive Flow model following https://arxiv.org/abs/1705.07057
@@ -74,14 +84,17 @@ class MaskedAutoregressiveFlow(nn.Module):
     Without the BatchNormFlow module, the network cannot be too deep or trained for too long.
     Otherwise some of the weights will start turning into nans.
     """
+
     n_dim: int
     n_hidden: int
     n_layer: int
 
     def setup(self):
-        self.layers = [MaskedAutoEncoder(self.n_dim, self.n_hidden) for _ in range(self.n_layer)]
+        self.layers = [
+            MaskedAutoEncoder(self.n_dim, self.n_hidden) for _ in range(self.n_layer)
+        ]
         self.norms = [BatchNormFlow(self.n_dim) for _ in range(self.n_layer)]
-    
+
     def __call__(self, inputs):
         log_jacobian = 0
         for layer, norm in zip(self.layers, self.norms):
@@ -89,27 +102,28 @@ def __call__(self, inputs):
             log_jacobian += log_jacobian_
             inputs, log_jacobian_ = layer(inputs)
             log_jacobian += log_jacobian_
-            inputs = inputs[:,::-1]
+            inputs = inputs[:, ::-1]
         return inputs, log_jacobian
 
     def inverse(self, inputs):
         # Be careful about flipping the inputs when inverting the flow.
         log_jacobian = 0
-        for layer,norm in zip(reversed(self.layers), reversed(self.norms)):
-            inputs = inputs[:,::-1]
+        for layer, norm in zip(reversed(self.layers), reversed(self.norms)):
+            inputs = inputs[:, ::-1]
             inputs, log_jacobian_ = layer.inverse(inputs)
             log_jacobian += log_jacobian_
             inputs, log_jacobian_ = norm(inputs, reverse=True)
             log_jacobian += log_jacobian_
         return inputs, log_jacobian
-    
+
     def sample(self, rng_key, n_samples, params):
-        mean = jnp.zeros((n_samples,self.n_dim))
-        cov = jnp.repeat(jnp.eye(self.n_dim)[None,:],n_samples,axis=0)
+        mean = jnp.zeros((n_samples, self.n_dim))
+        cov = jnp.repeat(jnp.eye(self.n_dim)[None, :], n_samples, axis=0)
         gaussian = jax.random.multivariate_normal(rng_key, mean, cov)
-        samples = self.apply({'params': params},gaussian,method=self.inverse)
+        samples = self.apply({"params": params}, gaussian, method=self.inverse)
         return samples
 
+
 from flax.linen.module import Module, compact, merge_param
 
 PRNGKey = Any
@@ -118,4 +132,3 @@ def sample(self, rng_key, n_samples, params):
 Dtype = Any  # this could be a real type?
 
 Axes = Union[int, Iterable[int]]
-
diff --git a/src/flowMC/nfmodel/realNVP.py b/src/flowMC/nfmodel/realNVP.py
index 87450d9..749899a 100644
--- a/src/flowMC/nfmodel/realNVP.py
+++ b/src/flowMC/nfmodel/realNVP.py
@@ -4,6 +4,7 @@
 from flax import linen as nn
 import numpy as np
 
+
 class MLP(nn.Module):
 
     """
@@ -31,7 +32,14 @@ class MLP(nn.Module):
     kernel_i: Callable = jax.nn.initializers.variance_scaling
 
     def setup(self):
-        self.layers = [nn.Dense(feat, use_bias=self.use_bias, kernel_init=self.kernel_i(self.init_weight_scale, "fan_in", "normal")) for feat in self.features]
+        self.layers = [
+            nn.Dense(
+                feat,
+                use_bias=self.use_bias,
+                kernel_init=self.kernel_i(self.init_weight_scale, "fan_in", "normal"),
+            )
+            for feat in self.features
+        ]
 
     def __call__(self, x):
         for l, layer in enumerate(self.layers[:-1]):
@@ -68,9 +76,9 @@ def setup(self):
         self.translate_MLP = MLP([self.n_features, self.n_hidden, self.n_features])
 
     def __call__(self, x):
-        s = self.mask * self.scale_MLP(x * (1-self.mask))
+        s = self.mask * self.scale_MLP(x * (1 - self.mask))
         s = jnp.tanh(s)
-        t = self.mask * self.translate_MLP(x * (1-self.mask))
+        t = self.mask * self.translate_MLP(x * (1 - self.mask))
         s = self.dt * s
         t = self.dt * t
         log_det = s.reshape(s.shape[0], -1).sum(axis=-1)
@@ -78,9 +86,9 @@ def __call__(self, x):
         return outputs, log_det
 
     def inverse(self, x):
-        s = self.mask * self.scale_MLP(x * (1-self.mask))
+        s = self.mask * self.scale_MLP(x * (1 - self.mask))
         s = jnp.tanh(s)
-        t = self.mask * self.translate_MLP(x * (1-self.mask))
+        t = self.mask * self.translate_MLP(x * (1 - self.mask))
         s = self.dt * s
         t = self.dt * t
         log_det = -s.reshape(s.shape[0], -1).sum(axis=-1)
@@ -88,7 +96,6 @@ def inverse(self, x):
         return outputs, log_det
 
 
-
 class RealNVP(nn.Module):
 
     """
@@ -106,7 +113,7 @@ class RealNVP(nn.Module):
         The number of affine coupling layers.
     dt : float
         Scaling factor for the affine coupling layer.
-    base_mean : ndarray 
+    base_mean : ndarray
         Mean of Gaussian base distribution
     base_cov : ndarray
         Covariance of Gaussian base distribution
@@ -116,21 +123,26 @@ class RealNVP(nn.Module):
     n_features: int
     n_hidden: int
     dt: float = 1
-    
 
     def setup(self):
         affine_coupling = []
         for i in range(self.n_layer):
             mask = np.ones(self.n_features)
-            mask[int(self.n_features/2):] = 0
+            mask[int(self.n_features / 2) :] = 0
             if i % 2 == 0:
                 mask = 1 - mask
             mask = jnp.array(mask)
-            affine_coupling.append(AffineCoupling(self.n_features, self.n_hidden, mask, dt=self.dt))
+            affine_coupling.append(
+                AffineCoupling(self.n_features, self.n_hidden, mask, dt=self.dt)
+            )
         self.affine_coupling = affine_coupling
 
-        self.base_mean = self.variable('variables','base_mean',jnp.zeros,((self.n_features)))
-        self.base_cov = self.variable('variables','base_cov',jnp.eye,(self.n_features))
+        self.base_mean = self.variable(
+            "variables", "base_mean", jnp.zeros, ((self.n_features))
+        )
+        self.base_cov = self.variable(
+            "variables", "base_cov", jnp.eye, (self.n_features)
+        )
 
     def __call__(self, x):
         log_det = jnp.zeros(x.shape[0])
@@ -142,16 +154,20 @@ def __call__(self, x):
     def inverse(self, x):
         log_det = jnp.zeros(x.shape[0])
         for i in range(self.n_layer):
-            x, log_det_i = self.affine_coupling[self.n_layer-1-i].inverse(x)
+            x, log_det_i = self.affine_coupling[self.n_layer - 1 - i].inverse(x)
             log_det += log_det_i
         return x, log_det
 
     def sample(self, rng_key, n_samples):
-        gaussian = jax.random.multivariate_normal(rng_key, self.base_mean.value, self.base_cov.value,shape=(n_samples,))
+        gaussian = jax.random.multivariate_normal(
+            rng_key, self.base_mean.value, self.base_cov.value, shape=(n_samples,)
+        )
         samples = self.inverse(gaussian)
         return samples
 
     def log_prob(self, x):
         y, log_det = self.__call__(x)
-        log_det = log_det + jax.scipy.stats.multivariate_normal.logpdf(y, self.base_mean.value, self.base_cov.value)
-        return log_det
\ No newline at end of file
+        log_det = log_det + jax.scipy.stats.multivariate_normal.logpdf(
+            y, self.base_mean.value, self.base_cov.value
+        )
+        return log_det
diff --git a/src/flowMC/nfmodel/rqSpline.py b/src/flowMC/nfmodel/rqSpline.py
index dba5289..e3a5eba 100644
--- a/src/flowMC/nfmodel/rqSpline.py
+++ b/src/flowMC/nfmodel/rqSpline.py
@@ -4,7 +4,7 @@
 import jax.numpy as jnp
 import distrax
 
-from flax import linen as nn           # The Linen API
+from flax import linen as nn  # The Linen API
 
 
 class MLP(nn.Module):
@@ -15,8 +15,14 @@ class MLP(nn.Module):
     kernel_i: Callable = jax.nn.initializers.variance_scaling
 
     def setup(self):
-        self.layers = [nn.Dense(feat, use_bias=self.use_bias, kernel_init=self.kernel_i(
-            self.init_weight_scale, "fan_in", "normal")) for feat in self.features]
+        self.layers = [
+            nn.Dense(
+                feat,
+                use_bias=self.use_bias,
+                kernel_init=self.kernel_i(self.init_weight_scale, "fan_in", "normal"),
+            )
+            for feat in self.features
+        ]
 
     def __call__(self, x):
         for l, layer in enumerate(self.layers[:-1]):
@@ -27,6 +33,7 @@ def __call__(self, x):
 
 class Reshape(nn.Module):
     shape: Sequence[int]
+
     def __call__(self, x):
         return jnp.reshape(x.T, self.shape)
 
@@ -37,50 +44,92 @@ class Conditioner(nn.Module):
     num_bijector_params: int
 
     def setup(self):
-        self.conditioner = nn.Sequential([MLP([self.n_features]+list(self.hidden_size)), nn.tanh, nn.Dense(
-            self.n_features * self.num_bijector_params, kernel_init=jax.nn.initializers.zeros, bias_init=jax.nn.initializers.zeros), Reshape((self.n_features, self.num_bijector_params))])
+        self.conditioner = nn.Sequential(
+            [
+                MLP([self.n_features] + list(self.hidden_size)),
+                nn.tanh,
+                nn.Dense(
+                    self.n_features * self.num_bijector_params,
+                    kernel_init=jax.nn.initializers.zeros,
+                    bias_init=jax.nn.initializers.zeros,
+                ),
+                Reshape((self.n_features, self.num_bijector_params)),
+            ]
+        )
 
     def __call__(self, x):
         return self.conditioner(x)
 
+
 class Scalar(nn.Module):
     n_features: int
 
     def setup(self):
-        self.shift = self.param('shifts',lambda rng, shape: jnp.zeros(shape), (self.n_features))
-        self.scale = self.param('scales',lambda rng, shape: jnp.ones(shape), (self.n_features))
+        self.shift = self.param(
+            "shifts", lambda rng, shape: jnp.zeros(shape), (self.n_features)
+        )
+        self.scale = self.param(
+            "scales", lambda rng, shape: jnp.ones(shape), (self.n_features)
+        )
 
     def __call__(self, x):
         return self.scale, self.shift
 
+
 def scalar_affine(params: jnp.ndarray):
-    return distrax.ScalarAffine(scale=params[0],shift=params[1])
+    return distrax.ScalarAffine(scale=params[0], shift=params[1])
 
 
 class RQSpline(nn.Module):
+
+    """
+    Rational quadratic spline normalizing flow model using distrax.
+
+    Parameters
+    ----------
+    n_features : int
+        Number of features in the data.
+    num_layers : int
+        Number of layers in the flow.
+    num_bins : int
+        Number of bins in the spline.
+    hidden_size : Sequence[int]
+        Size of the hidden layers in the conditioner.
+    spline_range : Sequence[float]
+        Range of the spline.
+    """
+
     n_features: int
     num_layers: int
     hidden_size: Sequence[int]
     num_bins: int
-    spline_range: Sequence[float] = (-10., 10.)
+    spline_range: Sequence[float] = (-10.0, 10.0)
 
     def setup(self):
         conditioner = []
         scalar = []
         for i in range(self.num_layers):
-            conditioner.append(Conditioner(self.n_features, self.hidden_size, 3*self.num_bins+1))
+            conditioner.append(
+                Conditioner(self.n_features, self.hidden_size, 3 * self.num_bins + 1)
+            )
             scalar.append(Scalar(self.n_features))
 
         self.conditioner = conditioner
         self.scalar = scalar
 
-        self.base_mean = self.variable('variables','base_mean',jnp.zeros,((self.n_features)))
-        self.base_cov = self.variable('variables','base_cov',jnp.eye,(self.n_features))
+        self.base_mean = self.variable(
+            "variables", "base_mean", jnp.zeros, ((self.n_features))
+        )
+        self.base_cov = self.variable(
+            "variables", "base_cov", jnp.eye, (self.n_features)
+        )
 
         self.vmap_call = jax.jit(jax.vmap(self.__call__))
-        
+
         def bijector_fn(params: jnp.ndarray):
-            return distrax.RationalQuadraticSpline(params, range_min=self.spline_range[0], range_max=self.spline_range[1])
+            return distrax.RationalQuadraticSpline(
+                params, range_min=self.spline_range[0], range_max=self.spline_range[1]
+            )
 
         self.bijector_fn = bijector_fn
 
@@ -89,15 +138,27 @@ def make_flow(self):
         mask_all = (jnp.zeros(self.n_features)).astype(bool)
         layers = []
         for i in range(self.num_layers):
-            layers.append(distrax.MaskedCoupling(
-                mask=mask_all, bijector=scalar_affine, conditioner=self.scalar[i]))
-            layers.append(distrax.MaskedCoupling(
-                mask=mask, bijector=self.bijector_fn, conditioner=self.conditioner[i]))
+            layers.append(
+                distrax.MaskedCoupling(
+                    mask=mask_all, bijector=scalar_affine, conditioner=self.scalar[i]
+                )
+            )
+            layers.append(
+                distrax.MaskedCoupling(
+                    mask=mask,
+                    bijector=self.bijector_fn,
+                    conditioner=self.conditioner[i],
+                )
+            )
             mask = jnp.logical_not(mask)
 
         flow = distrax.Inverse(distrax.Chain(layers))
-        base_dist = distrax.Independent(distrax.MultivariateNormalFullCovariance(
-            loc=jnp.zeros(self.n_features), covariance_matrix=jnp.eye(self.n_features)))
+        base_dist = distrax.Independent(
+            distrax.MultivariateNormalFullCovariance(
+                loc=jnp.zeros(self.n_features),
+                covariance_matrix=jnp.eye(self.n_features),
+            )
+        )
 
         return base_dist, flow
 
@@ -107,9 +168,10 @@ def __call__(self, x):
 
     def sample(self, rng, num_samples):
         base_dist, flow = self.make_flow()
-        samples = distrax.Transformed(base_dist, flow).sample(seed=rng, sample_shape=(num_samples))
-        return samples*jnp.sqrt(jnp.diag(self.base_cov.value))+self.base_mean.value
-
+        samples = distrax.Transformed(base_dist, flow).sample(
+            seed=rng, sample_shape=(num_samples)
+        )
+        return samples * jnp.sqrt(jnp.diag(self.base_cov.value)) + self.base_mean.value
 
-    def log_prob(self,x):
-        return self.vmap_call(x)
\ No newline at end of file
+    def log_prob(self, x):
+        return self.vmap_call(x)
diff --git a/src/flowMC/nfmodel/utils.py b/src/flowMC/nfmodel/utils.py
index 5cdb6d8..214352d 100644
--- a/src/flowMC/nfmodel/utils.py
+++ b/src/flowMC/nfmodel/utils.py
@@ -1,13 +1,26 @@
 import jax
-import jax.numpy as jnp                # JAX NumPy
-import jax.random as random            # JAX random
+import jax.numpy as jnp  # JAX NumPy
+import jax.random as random  # JAX random
 from tqdm import trange
 
 
 def make_training_loop(model):
+    """
+    Create a function that trains an NF model.
+
+    Args:
+        model: a neural network model with a `loss` method.
+
+    Returns:
+        train_flow (Callable): wrapper function that trains the model.
+
+    """
+
     def train_step(batch, state, variables):
         def loss(params):
-            log_det = model.apply({'params': params,'variables': variables}, batch, method=model.log_prob)
+            log_det = model.apply(
+                {"params": params, "variables": variables}, batch, method=model.log_prob
+            )
             return -jnp.mean(log_det)
 
         grad_fn = jax.value_and_grad(loss)
@@ -21,10 +34,10 @@ def train_epoch(rng, state, variables, train_ds, batch_size):
         """Train for a single epoch."""
         train_ds_size = len(train_ds)
         steps_per_epoch = train_ds_size // batch_size
-        if steps_per_epoch >0:
+        if steps_per_epoch > 0:
             perms = jax.random.permutation(rng, train_ds_size)
 
-            perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
+            perms = perms[: steps_per_epoch * batch_size]  # skip incomplete batch
             perms = perms.reshape((steps_per_epoch, batch_size))
             for perm in perms:
                 batch = train_ds[perm, ...]
@@ -38,20 +51,20 @@ def train_flow(rng, state, variables, data, num_epochs, batch_size):
         loss_values = jnp.zeros(num_epochs)
         best_state = state
         best_loss = 1e9
-        pbar = trange(num_epochs, desc="Training NF",miniters=int(num_epochs/10))
+        pbar = trange(num_epochs, desc="Training NF", miniters=int(num_epochs / 10))
         for epoch in pbar:
             # Use a separate PRNG key to permute image data during shuffling
             rng, input_rng = jax.random.split(rng)
             # Run an optimization step over a training batch
             value, state = train_epoch(input_rng, state, variables, data, batch_size)
-            #print('Train loss: %.3f' % value)
+            # print('Train loss: %.3f' % value)
             loss_values = loss_values.at[epoch].set(value)
             if loss_values[epoch] < best_loss:
                 best_state = state
                 best_loss = loss_values[epoch]
-            if epoch % int(num_epochs/10) == 0:
+            if epoch % int(num_epochs / 10) == 0:
                 pbar.set_description(f"Training NF, current loss: {value:.3f}")
-        
+
         return rng, best_state, loss_values
 
     return train_flow, train_epoch, train_step
@@ -59,6 +72,8 @@ def train_flow(rng, state, variables, data, num_epochs, batch_size):
 
 def sample_nf(model, param, rng_key, n_sample, variables):
     rng_key, subkey = random.split(rng_key)
-    samples = model.apply({'params': param,'variables': variables}, subkey, n_sample, method=model.sample)
+    samples = model.apply(
+        {"params": param, "variables": variables}, subkey, n_sample, method=model.sample
+    )
     # samples = jnp.flip(samples[0],axis=1)
-    return rng_key,samples
\ No newline at end of file
+    return rng_key, samples
diff --git a/src/flowMC/sampler/Gaussian_random_walk.py b/src/flowMC/sampler/Gaussian_random_walk.py
index 70a5804..d986937 100644
--- a/src/flowMC/sampler/Gaussian_random_walk.py
+++ b/src/flowMC/sampler/Gaussian_random_walk.py
@@ -5,7 +5,9 @@
 
 # @partial(jax.jit, static_argnums=(1,))
 def rw_metropolis_kernel(rng_key, logpdf, position, log_prob, dt=0.1):
+
     """Moves the chains by one step using the Random Walk Metropolis algorithm.
+
     Attributes
     ----------
     rng_key: jax.random.PRNGKey
@@ -36,26 +38,24 @@ def rw_metropolis_kernel(rng_key, logpdf, position, log_prob, dt=0.1):
 
 # @partial(jax.jit, static_argnums=(1, 2))
 def rw_metropolis_sampler(rng_key, n_steps, logpdf, initial_position):
-
     def mh_update_sol2(i, state):
         key, positions, log_prob, acceptance = state
         _, key = jax.random.split(key)
-        new_position, new_log_prob, do_accept = rw_metropolis_kernel(key, logpdf, positions[i-1], log_prob[i-1])
+        new_position, new_log_prob, do_accept = rw_metropolis_kernel(
+            key, logpdf, positions[i - 1], log_prob[i - 1]
+        )
         positions = positions.at[i].set(new_position)
         log_prob = log_prob.at[i].set(new_log_prob)
         acceptance = acceptance.at[i].set(do_accept)
         return (key, positions, log_prob, acceptance)
 
-
     logp = logpdf(initial_position)
     acceptance = jnp.zeros((n_steps,))
-    all_positions = jnp.zeros((n_steps,)+initial_position.shape) + initial_position
+    all_positions = jnp.zeros((n_steps,) + initial_position.shape) + initial_position
     all_logp = jnp.zeros((n_steps,)) + logp
     initial_state = (rng_key, all_positions, all_logp, acceptance)
-    rng_key, all_positions, all_logp, acceptance = jax.lax.fori_loop(1, n_steps, 
-                                                   mh_update_sol2, 
-                                                   initial_state)
-    
-    
-    return rng_key, all_positions, all_logp, acceptance
+    rng_key, all_positions, all_logp, acceptance = jax.lax.fori_loop(
+        1, n_steps, mh_update_sol2, initial_state
+    )
 
+    return rng_key, all_positions, all_logp, acceptance
diff --git a/src/flowMC/sampler/MALA.py b/src/flowMC/sampler/MALA.py
index 11c32a4..9bcc1bd 100644
--- a/src/flowMC/sampler/MALA.py
+++ b/src/flowMC/sampler/MALA.py
@@ -1,19 +1,33 @@
+from typing import Callable
 import jax
 import jax.numpy as jnp
 from jax.scipy.stats import multivariate_normal
 from tqdm import tqdm
 
-def make_mala_kernel(logpdf, d_logpdf, dt, M=None):
+
+def make_mala_kernel(logpdf: Callable, dt: float, M: jnp.array = None):
+    """
+    Making MALA kernel for a single step
+
+    Args:
+        logpdf (Callable): log-probability density function
+        dt (float): step size of the MALA step
+        M (jnp.array): mass matrix. Currently we only support diagonal mass matrix
+
+    Returns:
+        mala_kernel (Callable): MALA kernel for a single step
+
+    """
 
     if M != None:
-        dt  = dt*jnp.sqrt(M)
+        dt = dt * jnp.sqrt(M)
 
-    dt2 = dt*dt
+    dt2 = dt * dt
 
     def body(this_position, this_key):
-        this_log_prob,this_d_log = jax.value_and_grad(logpdf)(this_position)
-        proposal = this_position + jnp.dot(dt2, this_d_log)/2
-        proposal +=  jnp.dot(dt, jax.random.normal(this_key, shape=this_position.shape))
+        this_log_prob, this_d_log = jax.value_and_grad(logpdf)(this_position)
+        proposal = this_position + jnp.dot(dt2, this_d_log) / 2
+        proposal += jnp.dot(dt, jax.random.normal(this_key, shape=this_position.shape))
         return proposal, (proposal, this_log_prob, this_d_log)
 
     def mala_kernel(rng_key, position, log_prob):
@@ -37,29 +51,52 @@ def mala_kernel(rng_key, position, log_prob):
 
         """
         key1, key2 = jax.random.split(rng_key)
-        
-        _, (proposal, logprob, d_logprob) = jax.lax.scan(body, position, jnp.array([key1, key1]))
+
+        _, (proposal, logprob, d_logprob) = jax.lax.scan(
+            body, position, jnp.array([key1, key1])
+        )
 
         ratio = logprob[1] - logprob[0]
-        ratio -= multivariate_normal.logpdf(proposal[0], position+jnp.dot(dt2,d_logprob[0])/2,dt2)
-        ratio += multivariate_normal.logpdf(position, proposal[0]+jnp.dot(dt2,d_logprob[1])/2,dt2)
-        
+        ratio -= multivariate_normal.logpdf(
+            proposal[0], position + jnp.dot(dt2, d_logprob[0]) / 2, dt2
+        )
+        ratio += multivariate_normal.logpdf(
+            position, proposal[0] + jnp.dot(dt2, d_logprob[1]) / 2, dt2
+        )
+
         log_uniform = jnp.log(jax.random.uniform(key2))
         do_accept = log_uniform < ratio
 
         position = jnp.where(do_accept, proposal[0], position)
         log_prob = jnp.where(do_accept, logprob[1], logprob[0])
         return position, log_prob, do_accept
+
     return mala_kernel
 
 
 def make_mala_update(logpdf, d_logpdf, dt, M=None):
-    mala_kernel = make_mala_kernel(logpdf, d_logpdf, dt, M)
-    def mala_update(i,state):
+    """
+
+    Making MALA update function for multiple steps
+
+    Args:
+        logpdf (Callable): log-probability density function
+        d_logpdf (Callable): gradient of log-probability density function
+        dt (float): step size of the MALA step
+        M (jnp.array): mass matrix. Currently we only support diagonal mass matrix
+
+    Returns:
+        mala_update (Callable): MALA update function for multiple steps
+
+    """
+    mala_kernel = make_mala_kernel(logpdf, dt, M)
+
+    def mala_update(i, state):
         key, positions, log_prob, acceptance = state
         _, key = jax.random.split(key)
-        new_position, new_log_prob, do_accept = mala_kernel(key, positions[i-1],
-                                                                log_prob[i-1])
+        new_position, new_log_prob, do_accept = mala_kernel(
+            key, positions[i - 1], log_prob[i - 1]
+        )
         positions = positions.at[i].set(new_position)
         log_prob = log_prob.at[i].set(new_log_prob)
         acceptance = acceptance.at[i].set(do_accept)
@@ -72,6 +109,7 @@ def mala_update(i,state):
 
     return mala_update, mala_kernel, logpdf, d_logpdf
 
+
 def make_mala_sampler(logpdf, d_logpdf, dt=1e-5, jit=False, M=None):
     mala_update, mk, lp, dlp = make_mala_update(logpdf, d_logpdf, dt, M)
     # Somehow if I define the function inside the other function,
@@ -82,7 +120,7 @@ def make_mala_sampler(logpdf, d_logpdf, dt=1e-5, jit=False, M=None):
         lp = jax.jit(lp)
         dlp = jax.jit(dlp)
 
-    mala_update = jax.vmap(mala_update, in_axes=(None,(0,0,0,0)))
+    mala_update = jax.vmap(mala_update, in_axes=(None, (0, 0, 0, 0)))
     mk = jax.vmap(mk, in_axes=(0, 0, 0))
     lp = jax.vmap(lp)
     dlp = jax.vmap(dlp)
@@ -94,8 +132,8 @@ def mala_sampler(rng_key, n_steps, initial_position):
         This function do n step with the MALA kernel.
 
         Args:
-            rng_key (n_chains, 2): random key for the sampler 
-            n_steps (int): number of local steps 
+            rng_key (n_chains, 2): random key for the sampler
+            n_steps (int): number of local steps
             logpdf (function): log-density function
             d_logpdf (function): gradient of log-density function
             initial_position (n_chains, n_dim): initial position of the chain
@@ -105,21 +143,46 @@ def mala_sampler(rng_key, n_steps, initial_position):
             rng_key (n_chains, 2): random key for the sampler after the sampling
             all_positions (n_chains, n_steps, n_dim): all the positions of the chain
             log_probs (n_chains, ): log probability at the end of the chain
-            acceptance: acceptance rate of the chain 
+            acceptance: acceptance rate of the chain
         """
 
         logp = lp(initial_position)
         n_chains = rng_key.shape[0]
-        acceptance = jnp.zeros((n_chains,n_steps,))
-        all_positions = jnp.zeros((n_chains, n_steps,)+initial_position.shape[-1:]) + initial_position[:,None]
-        all_logp = jnp.zeros((n_chains,n_steps,)) + logp[:,None]
+        acceptance = jnp.zeros(
+            (
+                n_chains,
+                n_steps,
+            )
+        )
+        all_positions = (
+            jnp.zeros(
+                (
+                    n_chains,
+                    n_steps,
+                )
+                + initial_position.shape[-1:]
+            )
+            + initial_position[:, None]
+        )
+        all_logp = (
+            jnp.zeros(
+                (
+                    n_chains,
+                    n_steps,
+                )
+            )
+            + logp[:, None]
+        )
         state = (rng_key, all_positions, all_logp, acceptance)
-        for i in tqdm(range(1, n_steps),desc='Sampling Locally',miniters=int(n_steps/10)):
+        for i in tqdm(
+            range(1, n_steps), desc="Sampling Locally", miniters=int(n_steps / 10)
+        ):
             state = mala_update(i, state)
         return state
 
     return mala_sampler, mala_update, mk, lp, dlp
 
+
 ################### Scan API ##############################
 
 # def make_mala_kernel(logpdf, d_logpdf, dt):
@@ -150,5 +213,5 @@ def mala_sampler(rng_key, n_steps, initial_position):
 #     mala_update = jax.vmap(mala_update, in_axes=(0,0,0,None))
 #     mala_kernel_vec = jax.vmap(mala_kernel, in_axes=((0,0,0,0),None))
 #     logpdf = jax.vmap(logpdf)
-#     d_logpdf = jax.vmap(d_logpdf)    
+#     d_logpdf = jax.vmap(d_logpdf)
 #     return mala_update, mala_kernel_vec, logpdf, d_logpdf
diff --git a/src/flowMC/sampler/NF_proposal.py b/src/flowMC/sampler/NF_proposal.py
index ad004da..a047a73 100644
--- a/src/flowMC/sampler/NF_proposal.py
+++ b/src/flowMC/sampler/NF_proposal.py
@@ -1,3 +1,4 @@
+from typing import Tuple
 import jax
 import jax.numpy as jnp
 from jax import random, jit, vmap
@@ -6,8 +7,34 @@
 n_sample_max = 100000
 
 
-def nf_metropolis_kernel(rng_key, proposal_position, initial_position,
-                         proposal_pdf, proposal_nf_pdf, initial_pdf, initial_nf_pdf):
+def nf_metropolis_kernel(
+    rng_key: jax.random.PRNGKey,
+    proposal_position: jnp.ndarray,
+    initial_position: jnp.ndarray,
+    proposal_pdf: jnp.ndarray,
+    proposal_nf_pdf: jnp.ndarray,
+    initial_pdf: jnp.ndarray,
+    initial_nf_pdf: jnp.ndarray,
+):
+
+    """
+    A one-step global sampling kernel for a given normalizing flow model.
+
+    Args:
+        rng_key: Jax PRNGKey.
+        proposal_position: Proposed positions, shape (Ndim).
+        initial_position: Initial positions, shape (Ndim).
+        proposal_pdf: Pdf value evaluate using the target function at the proposal position, shape (Ndim).
+        proposal_nf_pdf: Pdf value evaluate using the normalizing flow model at the proposal position, shape (Ndim).
+        initial_pdf: Pdf value evaluate using the target function at the initial position, shape (Ndim).
+        initial_nf_pdf: Pdf value evaluate using the normalizing flow model at the initial position, shape (Ndim).
+
+    Returns:
+        position: New positions, shape (Ndim).
+        log_prob: Pdf value evaluate using the target function at the new position, shape (Ndim).
+        log_prob_nf: Pdf value evaluate using the normalizing flow model at the new position, shape (Ndim).
+        do_accept: Acceptance boolean, shape (Ndim).
+    """
 
     rng_key, subkeys = random.split(rng_key, 2)
     ratio = (proposal_pdf - initial_pdf) - (proposal_nf_pdf - initial_nf_pdf)
@@ -23,13 +50,37 @@ def nf_metropolis_kernel(rng_key, proposal_position, initial_position,
 
 
 @jax.jit
-def nf_metropolis_update(i, state):
-    key, positions, proposal, log_prob, log_prob_nf, log_prob_proposal, log_prob_nf_proposal, acceptance = state
-    key, *sub_key = jax.random.split(key, positions.shape[1]+1)
+def nf_metropolis_update(i: int, state: Tuple):
+
+    """
+    A multistep global sampling kernel for a given normalizing flow model.
+
+    Args:
+        i: Number of iteration.
+        state: A tuple containing the current state of the sampler.
+
+    """
+
+    (
+        key,
+        positions,
+        proposal,
+        log_prob,
+        log_prob_nf,
+        log_prob_proposal,
+        log_prob_nf_proposal,
+        acceptance,
+    ) = state
+    key, *sub_key = jax.random.split(key, positions.shape[1] + 1)
     sub_key = jnp.array(sub_key)
     new_position, new_log_prob, new_log_prob_nf, do_accept = nf_metropolis_kernel(
-        sub_key, proposal[i], positions[i-1], log_prob_proposal[i],
-        log_prob_nf_proposal[i], log_prob[i-1], log_prob_nf[i-1]
+        sub_key,
+        proposal[i],
+        positions[i - 1],
+        log_prob_proposal[i],
+        log_prob_nf_proposal[i],
+        log_prob[i - 1],
+        log_prob_nf[i - 1],
     )
     positions = positions.at[i].set(new_position)
     log_prob = log_prob.at[i].set(new_log_prob)
@@ -37,82 +88,153 @@ def nf_metropolis_update(i, state):
 
     acceptance = acceptance.at[i].set(do_accept)
 
-    return (key, positions, proposal, log_prob, log_prob_nf, log_prob_proposal, log_prob_nf_proposal, acceptance)
+    return (
+        key,
+        positions,
+        proposal,
+        log_prob,
+        log_prob_nf,
+        log_prob_proposal,
+        log_prob_nf_proposal,
+        acceptance,
+    )
 
 
 def make_nf_metropolis_sampler(nf_model):
 
+    """
+    A function make the normalizing flow sampler for a given normalizing flow model.
+
+    Args:
+        nf_model: A normalizing flow model.
+        
+    Returns:
+        nf_metropolis_sampler: A normalizing flow sampler that has the following signature:
+
+            Args:
+                rng_key: Jax PRNGKey.
+                n_steps: Number of steps.
+                nf_param: Normalizing flow parameters.
+                nf_variables: Normalizing flow variables.
+                target_pdf: Target pdf.
+                initial_position: Initial position.
+
+            Returns:
+                rng_key: Current state of random key
+                all_positions: All positions of the chain
+                all_logp: Log probability of the target function at all positions
+                all_logp_nf: Log probability of the normalizing flow model at all positions
+                acceptance: Acceptance boolean
+
+    """
+
     @jax.jit
     def eval_nf_logprob(position, nf_params, nf_variables):
-        return nf_model.apply({'params': nf_params, 'variables': nf_variables},
-                              position,
-                              method=nf_model.log_prob)
+        return nf_model.apply(
+            {"params": nf_params, "variables": nf_variables},
+            position,
+            method=nf_model.log_prob,
+        )
 
     def sample_nf(rng_key, n_samples, nf_params, nf_variables):
-        return nf_model.apply({'params': nf_params, 'variables': nf_variables}, rng_key,
-                              n_samples,
-                              method=nf_model.sample)
+        return nf_model.apply(
+            {"params": nf_params, "variables": nf_variables},
+            rng_key,
+            n_samples,
+            method=nf_model.sample,
+        )
 
     sample_nf = jax.jit(sample_nf, static_argnums=(1))
 
-    def nf_metropolis_sampler(rng_key, n_steps, nf_param, nf_variables, target_pdf,
-                              initial_position):
-        """
-        Returns:
-            rng_key: current state of random key
-            all_positions (n_steps, dim): all the positions of the chain
-            log_prob (): log probability at the end of the chain
-            log_prob_nf (): log probability at the end of the chain
-            acceptance (n_steps, ): acceptance table of the chains
-        """
-
+    def nf_metropolis_sampler(
+        rng_key, n_steps, nf_param, nf_variables, target_pdf, initial_position
+    ):
         rng_key, *subkeys = random.split(rng_key, 3)
 
-        total_sample = initial_position.shape[0]*n_steps
+        total_sample = initial_position.shape[0] * n_steps
 
-        initial_position_trans = (initial_position-nf_variables['base_mean'])/jnp.sqrt(jnp.diag(nf_variables['base_cov']))
-        log_pdf_nf_initial = eval_nf_logprob(initial_position_trans, nf_param, nf_variables)
+        initial_position_trans = (
+            initial_position - nf_variables["base_mean"]
+        ) / jnp.sqrt(jnp.diag(nf_variables["base_cov"]))
+        log_pdf_nf_initial = eval_nf_logprob(
+            initial_position_trans, nf_param, nf_variables
+        )
         log_pdf_initial = target_pdf(initial_position)
 
-
         if total_sample > n_sample_max:
             proposal_position = jnp.zeros((total_sample, initial_position.shape[-1]))
             log_pdf_nf_proposal = jnp.zeros((total_sample,))
             log_pdf_proposal = jnp.zeros((total_sample,))
             local_key, subkey = random.split(subkeys[0], 2)
-            for i in tqdm(range(total_sample//n_sample_max), desc='Sampling Globally',miniters=(total_sample//n_sample_max)//10):
+            for i in tqdm(
+                range(total_sample // n_sample_max),
+                desc="Sampling Globally",
+                miniters=(total_sample // n_sample_max) // 10,
+            ):
                 local_samples = sample_nf(subkey, n_sample_max, nf_param, nf_variables)
-                local_samples_trans = (local_samples-nf_variables['base_mean'])/jnp.sqrt(jnp.diag(nf_variables['base_cov']))
-                proposal_position = proposal_position.at[i*n_sample_max:(i+1)*n_sample_max].set(local_samples)
-                log_pdf_nf_proposal = log_pdf_nf_proposal.at[i*n_sample_max:(i+1)*n_sample_max].set(eval_nf_logprob(local_samples_trans, nf_param, nf_variables))
-                log_pdf_proposal = log_pdf_proposal.at[i*n_sample_max:(i+1)*n_sample_max].set(target_pdf(local_samples))
+                local_samples_trans = (
+                    local_samples - nf_variables["base_mean"]
+                ) / jnp.sqrt(jnp.diag(nf_variables["base_cov"]))
+                proposal_position = proposal_position.at[
+                    i * n_sample_max : (i + 1) * n_sample_max
+                ].set(local_samples)
+                log_pdf_nf_proposal = log_pdf_nf_proposal.at[
+                    i * n_sample_max : (i + 1) * n_sample_max
+                ].set(eval_nf_logprob(local_samples_trans, nf_param, nf_variables))
+                log_pdf_proposal = log_pdf_proposal.at[
+                    i * n_sample_max : (i + 1) * n_sample_max
+                ].set(target_pdf(local_samples))
                 local_key, subkey = random.split(local_key, 2)
 
         else:
-            proposal_position = sample_nf(subkeys[0], total_sample, nf_param, nf_variables)
-            proposal_position_trans = (proposal_position-nf_variables['base_mean'])/jnp.sqrt(jnp.diag(nf_variables['base_cov']))
-            log_pdf_nf_proposal = eval_nf_logprob(proposal_position_trans, nf_param, nf_variables)
+            proposal_position = sample_nf(
+                subkeys[0], total_sample, nf_param, nf_variables
+            )
+            proposal_position_trans = (
+                proposal_position - nf_variables["base_mean"]
+            ) / jnp.sqrt(jnp.diag(nf_variables["base_cov"]))
+            log_pdf_nf_proposal = eval_nf_logprob(
+                proposal_position_trans, nf_param, nf_variables
+            )
             log_pdf_proposal = target_pdf(proposal_position)
 
-        proposal_position = proposal_position.reshape(n_steps,
-                                                      initial_position.shape[0],
-                                                      initial_position.shape[1])
-        log_pdf_nf_proposal = log_pdf_nf_proposal.reshape(n_steps,
-                                                          initial_position.shape[0])
-        log_pdf_proposal = log_pdf_proposal.reshape(
-            n_steps, initial_position.shape[0])
-
-        all_positions = jnp.zeros((n_steps,) + initial_position.shape) + \
-            initial_position
+        proposal_position = proposal_position.reshape(
+            n_steps, initial_position.shape[0], initial_position.shape[1]
+        )
+        log_pdf_nf_proposal = log_pdf_nf_proposal.reshape(
+            n_steps, initial_position.shape[0]
+        )
+        log_pdf_proposal = log_pdf_proposal.reshape(n_steps, initial_position.shape[0])
+
+        all_positions = (
+            jnp.zeros((n_steps,) + initial_position.shape) + initial_position
+        )
         all_logp = jnp.zeros((n_steps, initial_position.shape[0])) + log_pdf_initial
-        all_logp_nf = jnp.zeros((n_steps, initial_position.shape[0])) + log_pdf_nf_initial
+        all_logp_nf = (
+            jnp.zeros((n_steps, initial_position.shape[0])) + log_pdf_nf_initial
+        )
         acceptance = jnp.zeros((n_steps, initial_position.shape[0]))
 
-        initial_state = (subkeys[1], all_positions, proposal_position, all_logp,
-                         all_logp_nf, log_pdf_proposal, log_pdf_nf_proposal, acceptance)
-        rng_key, all_positions, proposal_position, all_logp, all_logp_nf, log_pdf_proposal, log_pdf_nf_proposal, acceptance = jax.lax.fori_loop(1, n_steps,
-                                                                                                                                                nf_metropolis_update,
-                                                                                                                                                initial_state)
+        initial_state = (
+            subkeys[1],
+            all_positions,
+            proposal_position,
+            all_logp,
+            all_logp_nf,
+            log_pdf_proposal,
+            log_pdf_nf_proposal,
+            acceptance,
+        )
+        (
+            rng_key,
+            all_positions,
+            proposal_position,
+            all_logp,
+            all_logp_nf,
+            log_pdf_proposal,
+            log_pdf_nf_proposal,
+            acceptance,
+        ) = jax.lax.fori_loop(1, n_steps, nf_metropolis_update, initial_state)
         all_positions = all_positions.swapaxes(0, 1)
         all_logp = all_logp.swapaxes(0, 1)
         all_logp_nf = all_logp_nf.swapaxes(0, 1)
diff --git a/src/flowMC/sampler/Sampler.py b/src/flowMC/sampler/Sampler.py
index cc44721..59e0a0d 100644
--- a/src/flowMC/sampler/Sampler.py
+++ b/src/flowMC/sampler/Sampler.py
@@ -6,8 +6,9 @@
 from flowMC.sampler.NF_proposal import make_nf_metropolis_sampler
 from flax.training import train_state  # Useful dataclass to keep train state
 import flax
-import optax          
-                 # Optimizers
+import optax
+
+# Optimizers
 class Sampler(object):
     """
     Sampler class that host configuration parameters, NF model, and local sampler
@@ -21,27 +22,31 @@ class Sampler(object):
         d_likelihood (Device Array): Derivative of the likelihood function.
     """
 
-
-
-    def __init__(self, n_dim: int, rng_key_set, nf_model, local_sampler,
-                likelihood, 
-                d_likelihood = None,
-                n_loop: int = 2,
-                n_local_steps: int = 5,
-                n_global_steps: int = 5,
-                n_chains: int = 5,
-                n_epochs: int = 5,
-                n_nf_samples: int = 100,
-                learning_rate: float = 0.01,
-                max_samples: int = 10000,
-                momentum: float = 0.9,
-                batch_size: int = 10,
-                stepsize: float = 1e-3,
-                use_global: bool = True,
-                logging: bool = True,
-                nf_variable = None,
-                keep_quantile=0):
-        rng_key_init ,rng_keys_mcmc, rng_keys_nf, init_rng_keys_nf = rng_key_set
+    def __init__(
+        self,
+        n_dim: int,
+        rng_key_set,
+        nf_model,
+        local_sampler,
+        likelihood,
+        d_likelihood=None,
+        n_loop: int = 2,
+        n_local_steps: int = 5,
+        n_global_steps: int = 5,
+        n_chains: int = 5,
+        n_epochs: int = 5,
+        n_nf_samples: int = 100,
+        learning_rate: float = 0.01,
+        max_samples: int = 10000,
+        momentum: float = 0.9,
+        batch_size: int = 10,
+        stepsize: float = 1e-3,
+        use_global: bool = True,
+        logging: bool = True,
+        nf_variable=None,
+        keep_quantile=0,
+    ):
+        rng_key_init, rng_keys_mcmc, rng_keys_nf, init_rng_keys_nf = rng_key_set
 
         self.local_sampler = local_sampler
         self.likelihood = likelihood
@@ -66,27 +71,30 @@ def __init__(self, n_dim: int, rng_key_set, nf_model, local_sampler,
 
         self.nf_model = nf_model
         # params = nf_model.init(init_rng_keys_nf, jnp.ones((self.batch_size,self.n_dim)))['params']
-        model_init = nf_model.init(init_rng_keys_nf, jnp.ones((1,self.n_dim)))
-        params = model_init['params']
-        self.variables = model_init['variables']
+        model_init = nf_model.init(init_rng_keys_nf, jnp.ones((1, self.n_dim)))
+        params = model_init["params"]
+        self.variables = model_init["variables"]
         if nf_variable is not None:
             self.variables = self.variables
 
         self.keep_quantile = keep_quantile
-        self.nf_training_loop, train_epoch, train_step = make_training_loop(self.nf_model) 
+        self.nf_training_loop, train_epoch, train_step = make_training_loop(
+            self.nf_model
+        )
         self.global_sampler = make_nf_metropolis_sampler(self.nf_model)
 
         tx = optax.adam(self.learning_rate, self.momentum)
-        self.state = train_state.TrainState.create(apply_fn=nf_model.apply,
-                                                   params=params, tx=tx)
+        self.state = train_state.TrainState.create(
+            apply_fn=nf_model.apply, params=params, tx=tx
+        )
 
-        self.chains = jnp.empty((self.n_chains, 0 , self.n_dim))
+        self.chains = jnp.empty((self.n_chains, 0, self.n_dim))
         self.log_prob = jnp.empty((self.n_chains, 0))
         self.local_accs = jnp.empty((self.n_chains, 0))
         self.global_accs = jnp.empty((self.n_chains, 0))
         self.loss_vals = jnp.empty((0, self.n_epochs))
 
-    def sample(self,initial_position):
+    def sample(self, initial_position):
         """
         Sample from the posterior using the local sampler.
 
@@ -95,7 +103,7 @@ def sample(self,initial_position):
 
         Returns:
             chains (Device Array): Samples from the posterior.
-            nf_samples (Device Array): (n_nf_samples, n_dim) 
+            nf_samples (Device Array): (n_nf_samples, n_dim)
             local_accs (Device Array): (n_chains, n_local_steps * n_loop)
             global_accs (Device Array): (n_chains, n_global_steps * n_loop)
             loss_vals (Device Array): (n_epoch * n_loop,)
@@ -104,16 +112,18 @@ def sample(self,initial_position):
         last_step = initial_position
 
         for i in range(self.n_loop):
-            self.rng_keys_nf, self.rng_keys_mcmc, self.state, last_step = self.sampling_loop(
-                self.rng_keys_nf, self.rng_keys_mcmc, self.state, last_step)
+            (
+                self.rng_keys_nf,
+                self.rng_keys_mcmc,
+                self.state,
+                last_step,
+            ) = self.sampling_loop(
+                self.rng_keys_nf, self.rng_keys_mcmc, self.state, last_step
+            )
 
-  
         # return chains, log_prob, nf_samples, self.local_accs, global_accs, loss_vals
 
-    def sampling_loop(self,rng_keys_nf,
-                rng_keys_mcmc,
-                state,
-                initial_position):
+    def sampling_loop(self, rng_keys_nf, rng_keys_mcmc, state, initial_position):
 
         """
         Sampling loop for both the global sampler and the local sampler.
@@ -123,21 +133,21 @@ def sampling_loop(self,rng_keys_nf,
             rng_keys_mcmc (Device Array): RNG keys for the local sampler.
             d_likelihood ?
             TODO: likelihood vs posterior?
-            TODO: nf_samples - sometime int, sometimes samples 
+            TODO: nf_samples - sometime int, sometimes samples
 
         """
 
         if self.d_likelihood is None:
             rng_keys_mcmc, positions, log_prob, local_acceptance = self.local_sampler(
-                rng_keys_mcmc, self.n_local_steps, self.likelihood, initial_position)
+                rng_keys_mcmc, self.n_local_steps, self.likelihood, initial_position
+            )
         else:
             rng_keys_mcmc, positions, log_prob, local_acceptance = self.local_sampler(
-                rng_keys_mcmc, self.n_local_steps, initial_position)
+                rng_keys_mcmc, self.n_local_steps, initial_position
+            )
 
         log_prob_output = np.copy(log_prob)
 
-
-            
         if self.use_global == True:
             if self.keep_quantile > 0:
                 max_log_prob = jnp.max(log_prob_output, axis=1)
@@ -145,53 +155,86 @@ def sampling_loop(self,rng_keys_nf,
                 cut_chains = positions[max_log_prob > cut]
             else:
                 cut_chains = positions
-            chain_size = cut_chains.shape[0]*cut_chains.shape[1]
+            chain_size = cut_chains.shape[0] * cut_chains.shape[1]
             if chain_size > self.max_samples:
-                flat_chain = cut_chains[:,-int(self.max_samples/self.n_chains):].reshape(-1,self.n_dim)
+                flat_chain = cut_chains[
+                    :, -int(self.max_samples / self.n_chains) :
+                ].reshape(-1, self.n_dim)
             else:
                 flat_chain = cut_chains.reshape(-1, self.n_dim)
 
             variables = self.variables.unfreeze()
-            variables['base_mean'] = jnp.mean(flat_chain, axis=0)
-            variables['base_cov'] = jnp.cov(flat_chain.T)
+            variables["base_mean"] = jnp.mean(flat_chain, axis=0)
+            variables["base_cov"] = jnp.cov(flat_chain.T)
             self.variables = flax.core.freeze(variables)
 
-            flat_chain = (flat_chain-variables['base_mean'])/jnp.sqrt(jnp.diag(variables['base_cov']))
+            flat_chain = (flat_chain - variables["base_mean"]) / jnp.sqrt(
+                jnp.diag(variables["base_cov"])
+            )
 
-            rng_keys_nf, state, loss_values = self.nf_training_loop(rng_keys_nf, state, self.variables, flat_chain,
-                                                                    self.n_epochs, self.batch_size)
-            rng_keys_nf, nf_chain, log_prob, log_prob_nf, global_acceptance = self.global_sampler(
-                rng_keys_nf, self.n_global_steps, state.params, self.variables, self.likelihood_vec,
-                positions[:,-1])
+            rng_keys_nf, state, loss_values = self.nf_training_loop(
+                rng_keys_nf,
+                state,
+                self.variables,
+                flat_chain,
+                self.n_epochs,
+                self.batch_size,
+            )
+            (
+                rng_keys_nf,
+                nf_chain,
+                log_prob,
+                log_prob_nf,
+                global_acceptance,
+            ) = self.global_sampler(
+                rng_keys_nf,
+                self.n_global_steps,
+                state.params,
+                self.variables,
+                self.likelihood_vec,
+                positions[:, -1],
+            )
+
+            positions = jnp.concatenate((positions, nf_chain), axis=1)
+            log_prob_output = jnp.concatenate((log_prob_output, log_prob), axis=1)
 
-            positions = jnp.concatenate((positions,nf_chain),axis=1)
-            log_prob_output = jnp.concatenate((log_prob_output,log_prob),axis=1)
-        
         self.chains = jnp.append(self.chains, positions, axis=1)
         self.log_prob = jnp.append(self.log_prob, log_prob_output, axis=1)
         self.local_accs = jnp.append(self.local_accs, local_acceptance, axis=1)
         if self.use_global == True:
             self.global_accs = jnp.append(self.global_accs, global_acceptance, axis=1)
-            self.loss_vals = jnp.append(self.loss_vals, loss_values.reshape(1,-1), axis=0)
-
+            self.loss_vals = jnp.append(
+                self.loss_vals, loss_values.reshape(1, -1), axis=0
+            )
 
         last_step = positions[:, -1]
 
         return rng_keys_nf, rng_keys_mcmc, state, last_step
 
-
     def get_sampler_state(self):
-        return self.chains, self.log_prob, self.local_accs, self.global_accs, self.loss_vals
+        return (
+            self.chains,
+            self.log_prob,
+            self.local_accs,
+            self.global_accs,
+            self.loss_vals,
+        )
 
     def sample_flow(self, n_samples=None):
         if n_samples is None:
             n_samples = self.n_nf_samples
-        nf_samples = sample_nf(self.nf_model, self.state.params, self.rng_keys_nf, n_samples, self.variables)
+        nf_samples = sample_nf(
+            self.nf_model,
+            self.state.params,
+            self.rng_keys_nf,
+            n_samples,
+            self.variables,
+        )
         return nf_samples
 
     def reset(self):
-        self.chains = jnp.empty((self.n_chains, 0 , self.n_dim))
+        self.chains = jnp.empty((self.n_chains, 0, self.n_dim))
         self.log_prob = jnp.empty((self.n_chains, 0))
         self.local_accs = jnp.empty((self.n_chains, 0))
         self.global_accs = jnp.empty((self.n_chains, 0))
-        self.loss_vals = jnp.empty((0, self.n_epochs))
\ No newline at end of file
+        self.loss_vals = jnp.empty((0, self.n_epochs))
diff --git a/src/flowMC/sampler/mMALA.py b/src/flowMC/sampler/mMALA.py
index cf4dde8..d859a82 100644
--- a/src/flowMC/sampler/mMALA.py
+++ b/src/flowMC/sampler/mMALA.py
@@ -3,6 +3,7 @@
 from jax.scipy.stats import multivariate_normal
 from tqdm import tqdm
 
+
 def make_mMALA_kernel(logpdf, d_logpdf, d2_logpdf, dt):
     def mMALA_kernel(rng_key, position, log_prob):
 
@@ -29,45 +30,61 @@ def mMALA_kernel(rng_key, position, log_prob):
         n_dim = position.shape[-1]
         d_log_current = d_logpdf(position)
         d2_current = d2_logpdf(position)
-        d2_current = jax.lax.cond(jnp.any(jnp.isnan(d2_current)), lambda:jnp.eye(n_dim), lambda:d2_current)
+        d2_current = jax.lax.cond(
+            jnp.any(jnp.isnan(d2_current)), lambda: jnp.eye(n_dim), lambda: d2_current
+        )
         d2_current = jnp.linalg.inv(d2_current)
-        
 
-        dt_current_sqrt = dt*jnp.sqrt(d2_current)
+        dt_current_sqrt = dt * jnp.sqrt(d2_current)
 
-        proposal = position + jnp.dot((dt*d2_current)**2,d_log_current)/2
-        proposal += jnp.dot(dt_current_sqrt,jax.random.normal(key1, shape=position.shape))
+        proposal = position + jnp.dot((dt * d2_current) ** 2, d_log_current) / 2
+        proposal += jnp.dot(
+            dt_current_sqrt, jax.random.normal(key1, shape=position.shape)
+        )
         proposal_log_prob = logpdf(proposal)
         d2_proposal = d2_logpdf(proposal)
-        d2_proposal = jax.lax.cond(jnp.any(jnp.isnan(d2_proposal)), lambda:jnp.eye(n_dim), lambda:d2_proposal)
+        d2_proposal = jax.lax.cond(
+            jnp.any(jnp.isnan(d2_proposal)), lambda: jnp.eye(n_dim), lambda: d2_proposal
+        )
         d2_proposal = jnp.linalg.inv(d2_proposal)
 
         ratio = proposal_log_prob - logpdf(position)
-        ratio -= multivariate_normal.logpdf(position, proposal+jnp.dot(d2_proposal*dt*dt,d_logpdf(proposal))/2,d2_proposal*dt*dt)
-        ratio += multivariate_normal.logpdf(proposal, position+jnp.dot(d2_current*dt*dt,d_log_current)/2,d2_current*dt*dt)
-        
+        ratio -= multivariate_normal.logpdf(
+            position,
+            proposal + jnp.dot(d2_proposal * dt * dt, d_logpdf(proposal)) / 2,
+            d2_proposal * dt * dt,
+        )
+        ratio += multivariate_normal.logpdf(
+            proposal,
+            position + jnp.dot(d2_current * dt * dt, d_log_current) / 2,
+            d2_current * dt * dt,
+        )
+
         log_uniform = jnp.log(jax.random.uniform(key2))
         do_accept = log_uniform < ratio
 
         position = jnp.where(do_accept, proposal, position)
         log_prob = jnp.where(do_accept, proposal_log_prob, log_prob)
         return position, log_prob, do_accept
+
     return mMALA_kernel
 
 
 def make_mMALA_update(logpdf, d_logpdf, d2_logpdf, dt):
     mMALA_kernel = make_mMALA_kernel(logpdf, d_logpdf, d2_logpdf, dt)
-    def mMALA_update(i,state):
+
+    def mMALA_update(i, state):
         key, positions, log_prob, acceptance = state
         _, key = jax.random.split(key)
-        new_position, new_log_prob, do_accept = mMALA_kernel(key, positions[i-1],
-                                                                log_prob[i-1])
+        new_position, new_log_prob, do_accept = mMALA_kernel(
+            key, positions[i - 1], log_prob[i - 1]
+        )
         positions = positions.at[i].set(new_position)
         log_prob = log_prob.at[i].set(new_log_prob)
         acceptance = acceptance.at[i].set(do_accept)
         return (key, positions, log_prob, acceptance)
 
-    mMALA_update = jax.vmap(mMALA_update, in_axes=(None,(0,0,0,0)))
+    mMALA_update = jax.vmap(mMALA_update, in_axes=(None, (0, 0, 0, 0)))
     mMALA_kernel_vec = jax.vmap(mMALA_kernel, in_axes=(0, 0, 0))
     logpdf = jax.vmap(logpdf)
     d_logpdf = jax.vmap(d_logpdf)
@@ -79,6 +96,7 @@ def mMALA_update(i,state):
 
     return mMALA_update, mMALA_kernel_vec, logpdf, d_logpdf
 
+
 def make_MALA_sampler(logpdf, dt=1e-5, jit=False):
 
     d_logpdf = jax.grad(logpdf)
@@ -100,8 +118,8 @@ def mMALA_sampler(rng_key, n_steps, initial_position):
         This function do n step with the mMALA kernel.
 
         Args:
-            rng_key (n_chains, 2): random key for the sampler 
-            n_steps (int): number of local steps 
+            rng_key (n_chains, 2): random key for the sampler
+            n_steps (int): number of local steps
             logpdf (function): log-density function
             d_logpdf (function): gradient of log-density function
             initial_position (n_chains, n_dim): initial position of the chain
@@ -111,16 +129,40 @@ def mMALA_sampler(rng_key, n_steps, initial_position):
             rng_key (n_chains, 2): random key for the sampler after the sampling
             all_positions (n_chains, n_steps, n_dim): all the positions of the chain
             log_probs (n_chains, ): log probability at the end of the chain
-            acceptance: acceptance rate of the chain 
+            acceptance: acceptance rate of the chain
         """
 
         logp = lp(initial_position)
         n_chains = rng_key.shape[0]
-        acceptance = jnp.zeros((n_chains,n_steps,))
-        all_positions = jnp.zeros((n_chains, n_steps,)+initial_position.shape[-1:]) + initial_position[:,None]
-        all_logp = jnp.zeros((n_chains,n_steps,)) + logp[:,None]
+        acceptance = jnp.zeros(
+            (
+                n_chains,
+                n_steps,
+            )
+        )
+        all_positions = (
+            jnp.zeros(
+                (
+                    n_chains,
+                    n_steps,
+                )
+                + initial_position.shape[-1:]
+            )
+            + initial_position[:, None]
+        )
+        all_logp = (
+            jnp.zeros(
+                (
+                    n_chains,
+                    n_steps,
+                )
+            )
+            + logp[:, None]
+        )
         state = (rng_key, all_positions, all_logp, acceptance)
-        for i in tqdm(range(1, n_steps),desc='Sampling Locally',miniters=int(n_steps/10)):
+        for i in tqdm(
+            range(1, n_steps), desc="Sampling Locally", miniters=int(n_steps / 10)
+        ):
             state = mMALA_update(i, state)
         return state
 
diff --git a/src/flowMC/utils/Analysis.py b/src/flowMC/utils/Analysis.py
index 9e148d9..bd704b1 100644
--- a/src/flowMC/utils/Analysis.py
+++ b/src/flowMC/utils/Analysis.py
@@ -3,6 +3,7 @@
 
 import jax.numpy as jnp
 
+
 def ESS(x):
     """
     Compute the effective sample size of estimand of interest.
@@ -14,15 +15,18 @@ def ESS(x):
     """
     if x.shape < (2,):
         raise ValueError(
-            'Calculation of effective sample size'
-            'requires multiple chains of the same length.')
+            "Calculation of effective sample size"
+            "requires multiple chains of the same length."
+        )
     try:
-        m_chains, n_iter = x.shape # 1d variable
+        m_chains, n_iter = x.shape  # 1d variable
     except ValueError:
-        return [ESS(y.T) for y in x.T] # apply sequentially to each dimension
+        return [ESS(y.T) for y in x.T]  # apply sequentially to each dimension
 
-    def variogram(t): return (
-        (x[:, t:] - x[:, :(n_iter - t)])**2).sum() / (m_chains * (n_iter - t))
+    def variogram(t):
+        return ((x[:, t:] - x[:, : (n_iter - t)]) ** 2).sum() / (
+            m_chains * (n_iter - t)
+        )
 
     post_var = gelman_rubin(x)
     assert post_var > 0
@@ -36,7 +40,7 @@ def variogram(t): return (
         rho = rho.at[t].set(1 - variogram(t) / (2 * post_var))
 
         if not t % 2:
-            negative_autocorr = sum(rho[t - 1:t + 1]) < 0
+            negative_autocorr = sum(rho[t - 1 : t + 1]) < 0
 
         t += 1
 
@@ -51,11 +55,10 @@ def gelman_rubin(x):
     m_chains, n_iter = x.shape
 
     # Calculate between-chain variance
-    B_over_n = ((jnp.mean(x, axis=1) - jnp.mean(x))**2).sum() / (m_chains - 1)
+    B_over_n = ((jnp.mean(x, axis=1) - jnp.mean(x)) ** 2).sum() / (m_chains - 1)
 
     # Calculate within-chain variances
-    W = ((x - x.mean(axis=1, keepdims=True)) **
-         2).sum() / (m_chains * (n_iter - 1))
+    W = ((x - x.mean(axis=1, keepdims=True)) ** 2).sum() / (m_chains * (n_iter - 1))
 
     # (over) estimate of variance
     s2 = W * (n_iter - 1) / n_iter + B_over_n
diff --git a/src/flowMC/utils/PRNG_keys.py b/src/flowMC/utils/PRNG_keys.py
index 1bd65be..ba8bedf 100644
--- a/src/flowMC/utils/PRNG_keys.py
+++ b/src/flowMC/utils/PRNG_keys.py
@@ -18,10 +18,9 @@ def initialize_rng_keys(n_chains, seed=42):
         init_rng_keys_nf (Device Array): RNG keys for initializing wieght of the normalizing flow model.
     """
     rng_key = jax.random.PRNGKey(seed)
-    rng_key_init, rng_key_mcmc, rng_key_nf = jax.random.split(rng_key,3)
+    rng_key_init, rng_key_mcmc, rng_key_nf = jax.random.split(rng_key, 3)
 
     rng_keys_mcmc = jax.random.split(rng_key_mcmc, n_chains)  # (nchains,)
-    rng_keys_nf, init_rng_keys_nf = jax.random.split(rng_key_nf,2)
-    
-    return rng_key_init ,rng_keys_mcmc, rng_keys_nf, init_rng_keys_nf
+    rng_keys_nf, init_rng_keys_nf = jax.random.split(rng_key_nf, 2)
 
+    return rng_key_init, rng_keys_mcmc, rng_keys_nf, init_rng_keys_nf
diff --git a/src/flowMC/utils/PythonFunctionWrap.py b/src/flowMC/utils/PythonFunctionWrap.py
index b340096..7a1eb32 100644
--- a/src/flowMC/utils/PythonFunctionWrap.py
+++ b/src/flowMC/utils/PythonFunctionWrap.py
@@ -18,10 +18,7 @@
 from jax.tree_util import tree_flatten, tree_unflatten
 
 
-
-def wrap_python_log_prob_fn(
-    python_log_prob_fn: Callable[..., Array]
-):
+def wrap_python_log_prob_fn(python_log_prob_fn: Callable[..., Array]):
     @custom_vmap
     @wraps(python_log_prob_fn)
     def log_prob_fn(params: Array) -> Array:
@@ -33,9 +30,7 @@ def log_prob_fn(params: Array) -> Array:
         )
 
     @log_prob_fn.def_vmap
-    def _(
-        axis_size: int, in_batched: List[bool], params: Array
-    ) -> Tuple[Array, bool]:
+    def _(axis_size: int, in_batched: List[bool], params: Array) -> Tuple[Array, bool]:
         del axis_size, in_batched
 
         if _arraylike(params):
@@ -45,15 +40,13 @@ def _(
             flat_params, unravel = ravel_ensemble(params)
             eval_one = lambda x: python_log_prob_fn(unravel(x))
 
-        result_shape = jax.ShapeDtypeStruct(
-            (flat_params.shape[0],), flat_params.dtype
-        )
+        result_shape = jax.ShapeDtypeStruct((flat_params.shape[0],), flat_params.dtype)
 
         result = host_callback.call(
-                lambda y: np.stack([eval_one(x) for x in y]),
-                flat_params,
-                result_shape=result_shape,
-            )
+            lambda y: np.stack([eval_one(x) for x in y]),
+            flat_params,
+            result_shape=result_shape,
+        )
         return (
             result,
             True,
@@ -102,9 +95,7 @@ def _ravel_inner(lst: List[Array]) -> Tuple[Array, UnravelFn]:
 
         def unravel(arr: Array) -> PyTree:
             chunks = jnp.split(arr, indices[:-1])
-            return [
-                chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)
-            ]
+            return [chunk.reshape(shape) for chunk, shape in zip(chunks, shapes)]
 
         ravel = lambda arg: jnp.concatenate([jnp.ravel(e) for e in arg])
         raveled = jax.vmap(ravel)(lst)
@@ -131,4 +122,4 @@ def unravel(arr: Array) -> PyTree:
             [jnp.ravel(lax.convert_element_type(e, to_dtype)) for e in arg]
         )
         raveled = jax.vmap(ravel)(lst)
-        return raveled, unravel
\ No newline at end of file
+        return raveled, unravel
diff --git a/src/flowMC/utils/progressBar.py b/src/flowMC/utils/progressBar.py
index d4abe7c..8329029 100644
--- a/src/flowMC/utils/progressBar.py
+++ b/src/flowMC/utils/progressBar.py
@@ -7,15 +7,14 @@
 def progress_bar_scan(num_samples, message=None):
     "Progress bar for a JAX scan"
     if message is None:
-            message = f"Running for {num_samples:,} iterations"
+        message = f"Running for {num_samples:,} iterations"
     tqdm_bars = tqdm(range(num_samples))
     tqdm_bars.set_description(message)
 
-
     if num_samples > 20:
         print_rate = int(num_samples / 20)
     else:
-        print_rate = 1 # if you run the sampler for less than 20 iterations
+        print_rate = 1  # if you run the sampler for less than 20 iterations
 
     print_rate = 1
     remainder = num_samples % print_rate
@@ -38,7 +37,7 @@ def _update_progress_bar(iter_num):
 
         _ = lax.cond(
             # update tqdm every multiple of `print_rate` except at the end
-            (iter_num % print_rate == 0) & (iter_num != num_samples-remainder),
+            (iter_num % print_rate == 0) & (iter_num != num_samples - remainder),
             lambda _: host_callback.id_tap(_update_tqdm, print_rate, result=iter_num),
             lambda _: iter_num,
             operand=None,
@@ -46,7 +45,7 @@ def _update_progress_bar(iter_num):
 
         _ = lax.cond(
             # update tqdm by `remainder`
-            iter_num == num_samples-remainder,
+            iter_num == num_samples - remainder,
             lambda _: host_callback.id_tap(_update_tqdm, remainder, result=iter_num),
             lambda _: iter_num,
             operand=None,
@@ -57,13 +56,12 @@ def _close_tqdm(arg, transform):
 
     def close_tqdm(result, iter_num):
         return lax.cond(
-            iter_num == num_samples-1,
+            iter_num == num_samples - 1,
             lambda _: host_callback.id_tap(_close_tqdm, None, result=result),
             lambda _: result,
             operand=None,
         )
 
-
     def _progress_bar_scan(func):
         """Decorator that adds a progress bar to `body_fun` used in `lax.scan`.
         Note that `body_fun` must either be looping over `np.arange(num_samples)`,


Test patch:
diff --git a/example/kepler/test.py b/example/kepler/test.py
deleted file mode 100644
index cc04b61..0000000
--- a/example/kepler/test.py
+++ /dev/null
@@ -1,39 +0,0 @@
-import jax
-import jax.numpy as jnp
-import numpy as np
-import matplotlib.pyplot as plt
-
-#import jaxopt
-from scipy.optimize import minimize
-from utils import rv_model, get_kepler_params_and_log_jac, log_likelihood
-
-jax.config.update("jax_enable_x64", True)
-
-
-true_params = jnp.array([
-    12.0, np.log(0.5), np.log(14.5), np.log(2.3), 
-    np.sin(1.5), np.cos(1.5), 0.4, np.sin(-0.7), np.cos(-0.7)
-])
-
-random = np.random.default_rng(12345)
-t = np.sort(random.uniform(0, 100, 50))
-rv_err = 0.3
-rv_obs = rv_model(true_params, t) + random.normal(0, rv_err, len(t))
-
-plt.plot(t, rv_obs, ".k")
-x = np.linspace(0, 100, 500)
-plt.plot(x, rv_model(true_params, x), "C0")
-plt.show(block=False)
-
-
-# neg_logp_and_grad = jax.jit(lambda p: -log_likelihood(p, t, rv_err, rv_obs))
-# solver = jaxopt.LBFGS(fun=neg_logp_and_grad, maxiter=1000)
-# res = solver.run(true_params)
-
-neg_logp_and_grad = jax.jit(jax.value_and_grad(lambda p: -log_likelihood(p, t, rv_err, rv_obs)))
-soln = minimize(neg_logp_and_grad, true_params, jac=True)
-pos_ini = jnp.asarray(get_kepler_params_and_log_jac(soln.x)[0])
-
-pos_ini = pos_ini.repeat(10, 0)
-
-jnp.asarray(get_kepler_params_and_log_jac(soln.x)[0])
diff --git a/test/kernel_test_old.py b/test/kernel_test_old.py
index 956106d..41a4831 100644
--- a/test/kernel_test_old.py
+++ b/test/kernel_test_old.py
@@ -1,9 +1,9 @@
 from flowMC.nfmodel.realNVP import RealNVP
 import jax
-import jax.numpy as jnp                # JAX NumPy
+import jax.numpy as jnp  # JAX NumPy
 from flowMC.utils import Sampler, initialize_rng_keys
 from jax.scipy.special import logsumexp
-import numpy as np  
+import numpy as np
 
 
 from flowMC.nfmodel.utils import *
@@ -13,10 +13,20 @@ def mala_kernel(rng_key, logpdf, d_logpdf, position, log_prob, kernal_size=0.1):
 
     key1, key2 = jax.random.split(rng_key)
     proposal = position + kernal_size * d_logpdf(position)
-    proposal += kernal_size * jnp.sqrt(2/kernal_size) * jax.random.normal(key1, shape=position.shape)
+    proposal += (
+        kernal_size
+        * jnp.sqrt(2 / kernal_size)
+        * jax.random.normal(key1, shape=position.shape)
+    )
     ratio = logpdf(proposal) - logpdf(position)
-    ratio -= ((position - proposal - kernal_size * d_logpdf(proposal)) ** 2 / (4 * kernal_size)).sum()
-    ratio += ((proposal - position - kernal_size * d_logpdf(position)) ** 2 / (4 * kernal_size)).sum()
+    ratio -= (
+        (position - proposal - kernal_size * d_logpdf(proposal)) ** 2
+        / (4 * kernal_size)
+    ).sum()
+    ratio += (
+        (proposal - position - kernal_size * d_logpdf(position)) ** 2
+        / (4 * kernal_size)
+    ).sum()
     proposal_log_prob = logpdf(proposal)
 
     log_uniform = jnp.log(jax.random.uniform(key2))
@@ -26,39 +36,44 @@ def mala_kernel(rng_key, logpdf, d_logpdf, position, log_prob, kernal_size=0.1):
     log_prob = jnp.where(do_accept, proposal_log_prob, log_prob)
     return position, log_prob, do_accept.astype(jnp.int8)
 
+
 mala_kernel_vec = jax.vmap(mala_kernel, in_axes=(0, None, None, 0, 0, None))
 
-def mala_sampler(rng_key, n_samples, logpdf, d_logpdf, initial_position, kernal_size=0.1):
 
+def mala_sampler(
+    rng_key, n_samples, logpdf, d_logpdf, initial_position, kernal_size=0.1
+):
     def mh_update_sol2(i, state):
         key, positions, log_prob, acceptance = state
         _, key = jax.random.split(key)
-        new_position, new_log_prob, accept_local = mala_kernel(key, logpdf, d_logpdf, positions[i-1], log_prob, kernal_size)
-        positions=positions.at[i].set(new_position)
+        new_position, new_log_prob, accept_local = mala_kernel(
+            key, logpdf, d_logpdf, positions[i - 1], log_prob, kernal_size
+        )
+        positions = positions.at[i].set(new_position)
         acceptance += accept_local
-        return (key, positions, new_log_prob,acceptance)
-
+        return (key, positions, new_log_prob, acceptance)
 
     logp = logpdf(initial_position)
     acceptance = jnp.zeros(logp.shape)
-    all_positions = jnp.zeros((n_samples,)+initial_position.shape) + initial_position
-    initial_state = (rng_key,all_positions, logp, acceptance)
-    rng_key, all_positions, log_prob, acceptance = jax.lax.fori_loop(1, n_samples, 
-                                                   mh_update_sol2, 
-                                                   initial_state)
-    
-    
-    return rng_key, all_positions, log_prob, acceptance/n_samples
+    all_positions = jnp.zeros((n_samples,) + initial_position.shape) + initial_position
+    initial_state = (rng_key, all_positions, logp, acceptance)
+    rng_key, all_positions, log_prob, acceptance = jax.lax.fori_loop(
+        1, n_samples, mh_update_sol2, initial_state
+    )
+
+    return rng_key, all_positions, log_prob, acceptance / n_samples
+
 
 def dual_moon_pe(x):
     """
     Term 2 and 3 separate the distriubiotn and smear it along the first and second dimension
     """
     term1 = 0.5 * ((jnp.linalg.norm(x) - 2) / 0.1) ** 2
-    term2 = -0.5 * ((x[:1] + jnp.array([-3., 3.])) / 0.8) ** 2
-    term3 = -0.5 * ((x[1:2] + jnp.array([-3., 3.])) / 0.6) ** 2
+    term2 = -0.5 * ((x[:1] + jnp.array([-3.0, 3.0])) / 0.8) ** 2
+    term3 = -0.5 * ((x[1:2] + jnp.array([-3.0, 3.0])) / 0.6) ** 2
     return -(term1 - logsumexp(term2) - logsumexp(term3))
 
+
 d_dual_moon = jax.jit(jax.grad(dual_moon_pe))
 d_dual_moon_vec = jax.jit(jax.vmap(jax.grad(dual_moon_pe)))
 dual_moon_pe_vec = jax.jit(jax.vmap(dual_moon_pe))
@@ -75,28 +90,29 @@ def dual_moon_pe(x):
 
 
 config = {}
-config['n_dim'] = 5
-config['n_loop'] = 5
-config['n_local_steps'] = 20
-config['n_global_steps'] = 100
-config['n_chains'] = 100
-config['learning_rate'] = 0.01
-config['momentum'] = 0.9
-config['num_epochs'] = 100
-config['batch_size'] = 1000
-config['stepsize'] = 0.01
+config["n_dim"] = 5
+config["n_loop"] = 5
+config["n_local_steps"] = 20
+config["n_global_steps"] = 100
+config["n_chains"] = 100
+config["learning_rate"] = 0.01
+config["momentum"] = 0.9
+config["num_epochs"] = 100
+config["batch_size"] = 1000
+config["stepsize"] = 0.01
 
 print("Preparing RNG keys")
-rng_key_set = initialize_rng_keys(config['n_chains'],seed=42)
+rng_key_set = initialize_rng_keys(config["n_chains"], seed=42)
 
 print("Initializing MCMC model and normalizing flow model.")
 
-initial_position = jax.random.normal(rng_key_set[0],shape=(config['n_chains'], config['n_dim']))
+initial_position = jax.random.normal(
+    rng_key_set[0], shape=(config["n_chains"], config["n_dim"])
+)
 
 
-model = RealNVP(10,config['n_dim'],64, 1)
-run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None),
-                    out_axes=0)
+model = RealNVP(10, config["n_dim"], 64, 1)
+run_mcmc = jax.vmap(mala_sampler, in_axes=(0, None, None, None, 0, None), out_axes=0)
 
 # print("Initializing sampler class")
 
diff --git a/test/old_test_proposal.py b/test/old_test_proposal.py
index f4aeb44..2ef41f7 100644
--- a/test/old_test_proposal.py
+++ b/test/old_test_proposal.py
@@ -1,18 +1,17 @@
-
 from flowMC.nfmodel.maf import MaskedAutoregressiveFlow
 from flowMC.nfmodel.realNVP import RealNVP
 from flowMC.sampler.Gaussian_random_walk import rw_metropolis_sampler
 from flowMC.sampler.NF_proposal import nf_metropolis_sampler, nf_metropolis_kernel
 import jax
-import jax.numpy as jnp                # JAX NumPy
-import jax.random as random            # JAX random
+import jax.numpy as jnp  # JAX NumPy
+import jax.random as random  # JAX random
 from jax.scipy.special import logsumexp
 from jax.scipy.stats import multivariate_normal
-import numpy as np  
+import numpy as np
 
-from flax import linen as nn           # The Linen API
+from flax import linen as nn  # The Linen API
 from flax.training import train_state  # Useful dataclass to keep train state
-import optax                           # Optimizers
+import optax  # Optimizers
 from functools import partial
 from jax.scipy.stats import norm
 
@@ -26,30 +25,31 @@
 """
 
 
-
 def train_step(model, state, batch):
-	def loss(params):
-		y, log_det = model.apply({'params': params},batch)
-		mean = jnp.zeros((batch.shape[0],model.n_features))
-		cov = jnp.repeat(jnp.eye(model.n_features)[None,:],batch.shape[0],axis=0)
-		log_det = log_det + multivariate_normal.logpdf(y,mean,cov)
-		return -jnp.mean(log_det)
-	grad_fn = jax.value_and_grad(loss)
-	value, grad = grad_fn(state.params)
-	state = state.apply_gradients(grads=grad)
-	return value,state
-
-train_step = jax.jit(train_step,static_argnums=(0,))
+    def loss(params):
+        y, log_det = model.apply({"params": params}, batch)
+        mean = jnp.zeros((batch.shape[0], model.n_features))
+        cov = jnp.repeat(jnp.eye(model.n_features)[None, :], batch.shape[0], axis=0)
+        log_det = log_det + multivariate_normal.logpdf(y, mean, cov)
+        return -jnp.mean(log_det)
+
+    grad_fn = jax.value_and_grad(loss)
+    value, grad = grad_fn(state.params)
+    state = state.apply_gradients(grads=grad)
+    return value, state
+
+
+train_step = jax.jit(train_step, static_argnums=(0,))
 
-def train_flow(rng, model, state, data):
 
+def train_flow(rng, model, state, data):
     def train_epoch(state, train_ds, batch_size, epoch, rng):
         """Train for a single epoch."""
         train_ds_size = len(train_ds)
         steps_per_epoch = train_ds_size // batch_size
 
         perms = jax.random.permutation(rng, train_ds_size)
-        perms = perms[:steps_per_epoch * batch_size]  # skip incomplete batch
+        perms = perms[: steps_per_epoch * batch_size]  # skip incomplete batch
         perms = perms.reshape((steps_per_epoch, batch_size))
         for perm in perms:
             batch = train_ds[perm, ...]
@@ -58,29 +58,34 @@ def train_epoch(state, train_ds, batch_size, epoch, rng):
         return value, state
 
     for epoch in range(1, num_epochs + 1):
-        print('Epoch %d' % epoch)
+        print("Epoch %d" % epoch)
         # Use a separate PRNG key to permute image data during shuffling
         rng, input_rng = jax.random.split(rng)
         # Run an optimization step over a training batch
         value, state = train_epoch(state, data, batch_size, epoch, input_rng)
-        print('Train loss: %.3f' % value)
+        print("Train loss: %.3f" % value)
 
     return rng, state
 
+
 def dual_moon_pe(x):
     """
     Term 2 and 3 separate the distriubiotn and smear it along the first and second dimension
     """
     term1 = 0.5 * ((jnp.linalg.norm(x, axis=-1) - 2) / 0.1) ** 2
-    term2 = -0.5 * ((x[..., :1] + jnp.array([-3., 3.])) / 0.8) ** 2
-    term3 = -0.5 * ((x[..., 1:2] + jnp.array([-3., 3.])) / 0.6) ** 2
+    term2 = -0.5 * ((x[..., :1] + jnp.array([-3.0, 3.0])) / 0.8) ** 2
+    term3 = -0.5 * ((x[..., 1:2] + jnp.array([-3.0, 3.0])) / 0.6) ** 2
     return -(term1 - logsumexp(term2, axis=-1) - logsumexp(term3, axis=-1))
 
-def sample_nf(model, param, rng_key,n_sample):
+
+def sample_nf(model, param, rng_key, n_sample):
     rng_key, subkey = random.split(rng_key)
-    samples = model.apply({'params': param}, subkey, n_sample,param, method=model.sample)
-    samples = jnp.flip(samples[0],axis=1)
-    return rng_key,samples
+    samples = model.apply(
+        {"params": param}, subkey, n_sample, param, method=model.sample
+    )
+    samples = jnp.flip(samples[0], axis=1)
+    return rng_key, samples
+
 
 n_dim = 5
 n_samples = 200
@@ -94,21 +99,20 @@ def sample_nf(model, param, rng_key,n_sample):
 
 print("Preparing RNG keys")
 rng_key = jax.random.PRNGKey(42)
-rng_key_mcmc, rng_key_nf = jax.random.split(rng_key,2)
+rng_key_mcmc, rng_key_nf = jax.random.split(rng_key, 2)
 
 rng_keys_mcmc = jax.random.split(rng_key_mcmc, n_chains)  # (nchains,)
-rng_keys_nf, init_rng_keys_nf = jax.random.split(rng_key_nf,2)
+rng_keys_nf, init_rng_keys_nf = jax.random.split(rng_key_nf, 2)
 
 print("Initializing MCMC model and normalizing flow model.")
 
-initial_position = jnp.zeros((n_dim, n_chains)) #(n_dim, n_chains)
+initial_position = jnp.zeros((n_dim, n_chains))  # (n_dim, n_chains)
 
-#model = MaskedAutoregressiveFlow(n_dim,64,4)
-model = RealNVP(10,n_dim,64, 1)
-params = model.init(init_rng_keys_nf, jnp.ones((batch_size,n_dim)))['params']
+# model = MaskedAutoregressiveFlow(n_dim,64,4)
+model = RealNVP(10, n_dim, 64, 1)
+params = model.init(init_rng_keys_nf, jnp.ones((batch_size, n_dim)))["params"]
 
-run_mcmc = jax.vmap(rw_metropolis_sampler, in_axes=(0, None, None, 1),
-                    out_axes=0)
+run_mcmc = jax.vmap(rw_metropolis_sampler, in_axes=(0, None, None, 1), out_axes=0)
 
 tx = optax.adam(learning_rate, momentum)
 state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)
@@ -117,38 +121,65 @@ def sample_nf(model, param, rng_key,n_sample):
 
 
 def sampling_loop(rng_keys_nf, rng_keys_mcmc, model, state, initial_position):
-	rng_keys_mcmc, positions, log_prob = run_mcmc(rng_keys_mcmc, n_samples, dual_moon_pe, initial_position)
-	flat_chain = positions.reshape(-1,n_dim)	
-	rng_keys_nf, state = train_flow(rng_key_nf, model, state, flat_chain)	
-	rng_keys_nf, nf_chain, log_prob, log_prob_nf = nf_metropolis_sampler(rng_keys_nf, nf_samples, model, state.params , dual_moon_pe, positions[:,-1])
+    rng_keys_mcmc, positions, log_prob = run_mcmc(
+        rng_keys_mcmc, n_samples, dual_moon_pe, initial_position
+    )
+    flat_chain = positions.reshape(-1, n_dim)
+    rng_keys_nf, state = train_flow(rng_key_nf, model, state, flat_chain)
+    rng_keys_nf, nf_chain, log_prob, log_prob_nf = nf_metropolis_sampler(
+        rng_keys_nf, nf_samples, model, state.params, dual_moon_pe, positions[:, -1]
+    )
+
+    positions = jnp.concatenate((positions, nf_chain), axis=1)
+    return rng_keys_nf, rng_keys_mcmc, state, positions
 
-	positions = jnp.concatenate((positions,nf_chain),axis=1)
-	return rng_keys_nf, rng_keys_mcmc, state, positions
 
 last_step = initial_position
 chains = []
 for i in range(3):
-	rng_keys_nf, rng_keys_mcmc, state, positions = sampling_loop(rng_keys_nf, rng_keys_mcmc, model, state, last_step)
-	last_step = positions[:,-1].T
-	# rng_keys_mcmc, positions, log_prob = run_mcmc(rng_keys_mcmc, n_samples, likelihood, initial_position)
-	# last_step = last_step.T
-	chains.append(positions)
-chains = np.concatenate(chains,axis=1)
-#nf_samples = sample_nf(model, state.params, rng_keys_nf, 10000)
-
-rng_key, subkey = random.split(rng_keys_nf,2)
-proposal_position = model.apply({'params': params}, subkey, initial_position.shape[1]*nf_samples, params, method=model.sample)[0]
-
-log_pdf_nf_proposal = model.apply({'params': params}, proposal_position, method=model.log_prob)
-log_pdf_nf_initial = model.apply({'params': params}, initial_position.T, method=model.log_prob)
+    rng_keys_nf, rng_keys_mcmc, state, positions = sampling_loop(
+        rng_keys_nf, rng_keys_mcmc, model, state, last_step
+    )
+    last_step = positions[:, -1].T
+    # rng_keys_mcmc, positions, log_prob = run_mcmc(rng_keys_mcmc, n_samples, likelihood, initial_position)
+    # last_step = last_step.T
+    chains.append(positions)
+chains = np.concatenate(chains, axis=1)
+# nf_samples = sample_nf(model, state.params, rng_keys_nf, 10000)
+
+rng_key, subkey = random.split(rng_keys_nf, 2)
+proposal_position = model.apply(
+    {"params": params},
+    subkey,
+    initial_position.shape[1] * nf_samples,
+    params,
+    method=model.sample,
+)[0]
+
+log_pdf_nf_proposal = model.apply(
+    {"params": params}, proposal_position, method=model.log_prob
+)
+log_pdf_nf_initial = model.apply(
+    {"params": params}, initial_position.T, method=model.log_prob
+)
 log_pdf_proposal = dual_moon_pe(proposal_position)
 log_pdf_initial = dual_moon_pe(initial_position.T)
 
-proposal_position = proposal_position.reshape(nf_samples, initial_position.shape[1], initial_position.shape[0])
+proposal_position = proposal_position.reshape(
+    nf_samples, initial_position.shape[1], initial_position.shape[0]
+)
 log_pdf_nf_proposal = log_pdf_nf_proposal.reshape(nf_samples, initial_position.shape[1])
 log_pdf_proposal = log_pdf_proposal.reshape(nf_samples, initial_position.shape[1])
 
-rng_key, *subkeys = random.split(rng_key,initial_position.shape[1]+1)
+rng_key, *subkeys = random.split(rng_key, initial_position.shape[1] + 1)
 subkeys = jnp.array(subkeys)
 
-nf_metropolis_kernel(subkeys, proposal_position[0], initial_position.T, log_pdf_proposal[0], log_pdf_nf_proposal[0], log_pdf_initial, log_pdf_nf_initial)
\ No newline at end of file
+nf_metropolis_kernel(
+    subkeys,
+    proposal_position[0],
+    initial_position.T,
+    log_pdf_proposal[0],
+    log_pdf_nf_proposal[0],
+    log_pdf_initial,
+    log_pdf_nf_initial,
+)
diff --git a/test/sampler_test.py b/test/sampler_test.py
index 1262dd5..381329f 100644
--- a/test/sampler_test.py
+++ b/test/sampler_test.py
@@ -3,20 +3,26 @@
 from flowMC.sampler.MALA import mala_kernel
 from flowMC.sampler.Gaussian_random_walk import rw_metropolis_kernel
 
+
 def test_mala_kernel():
     rng_key = jax.random.PRNGKey(0)
     logpdf = lambda x: x**2
-    d_logpdf = lambda x: 2*x
-    position = jnp.array([1., 2.])
+    d_logpdf = lambda x: 2 * x
+    position = jnp.array([1.0, 2.0])
     log_prob = logpdf(position)
     dt = 0.1
-    position, log_prob, do_accept = mala_kernel(rng_key, logpdf, d_logpdf, position, log_prob, dt)
+    position, log_prob, do_accept = mala_kernel(
+        rng_key, logpdf, d_logpdf, position, log_prob, dt
+    )
+
 
 def test_rw_metropolis_kernel():
     rng_key = jax.random.PRNGKey(0)
     logpdf = lambda x: x**2
-    d_logpdf = lambda x: 2*x
-    position = jnp.array([1., 2.])
+    d_logpdf = lambda x: 2 * x
+    position = jnp.array([1.0, 2.0])
     log_prob = logpdf(position)
     dt = 0.1
-    position, log_prob, do_accept = rw_metropolis_kernel(rng_key, logpdf, d_logpdf, position, log_prob, dt)
\ No newline at end of file
+    position, log_prob, do_accept = rw_metropolis_kernel(
+        rng_key, logpdf, d_logpdf, position, log_prob, dt
+    )