Repository: minaskar/pocomc (language: python)

Issue:
fix bridge sampling and call counter



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 80e3743..868e7c3 100644
--- a/README.md
+++ b/README.md
@@ -47,7 +47,7 @@ def log_prior(x):
     if np.any((x < -10.0) | (x > 10.0)):  # If any dimension is out of bounds, the log prior is -infinity
         return -np.inf 
     else:
-        return -const
+        return 0.0
 
 def log_likelihood(x):
     return -np.sum(10.0*(x[:,::2]**2.0 - x[:,1::2])**2.0 \
@@ -55,15 +55,16 @@ def log_likelihood(x):
 
 
 n_particles = 1000
-prior_samples = np.random.uniform(size=(nwalkers, ndim), low=-10.0, high=10.0)
-
-sampler = pc.Sampler(n_particles,
-                     n_dim,
-                     log_likelihood,
-                     log_prior,
-                     vectorize_likelihood=True,
-                     bounds=(-10.0, 10.0)
-                    )
+prior_samples = np.random.uniform(size=(n_particles, n_dim), low=-10.0, high=10.0)
+
+sampler = pc.Sampler(
+    n_particles,
+    n_dim,
+    log_likelihood,
+    log_prior,
+    vectorize_likelihood=True,
+    bounds=(-10.0, 10.0)
+)
 sampler.run(prior_samples)
 
 results = sampler.results # Dictionary with results
diff --git a/docs/source/advanced.rst b/docs/source/advanced.rst
index d7ef7d9..e585573 100644
--- a/docs/source/advanced.rst
+++ b/docs/source/advanced.rst
@@ -17,7 +17,7 @@ If you are not familiar with these terms I encourage you to visit the :doc:`back
 
 Log-likelihood function
 -----------------------
-Suppose that we want our *likelihood* function to be a *Gaussian density* with 10 paramerers or in 10-D, we would do 
+Suppose that we want our *likelihood* function to be a *Gaussian density* with 10 parameters or in 10-D, we would do
 something like::
 
     import numpy as np
@@ -59,7 +59,7 @@ e.g. :math:`x\sim\mathcal{N}(0,3^{2})`, we would do something like::
     def log_prior(x):
         return - 0.5 * np.dot(x, x) / 3.0**2.0
 
-Alternatively, we can have priors in which not all of the paramerers have the same density a priori. For instance, suppose that 
+Alternatively, we can have priors in which not all of the parameters have the same density a priori. For instance, suppose that
 we want the first five parameters to have a flat/uniform prior :math:`x_{i}\sim\mathcal{U}(-10,10)` for :math:`i=0,1,\dots,4` and
 the last five to have a Gaussian/normal prior  :math:`x_{i}\sim\mathcal{N}(0,3^{2})` for :math:`i=5,6,\dots,9`, we would do::
 
@@ -132,7 +132,7 @@ The sampler also accepts other arguments, for a full list see :doc:`api`. Those
   or to the log-prior using the arguments ``log_prior_args`` and ``log_prior_kwargs``.
 - The arguments ``vectorize_likelihood`` and ``vectorize_prior`` which accept boolean values allow the user to use vectorized
   log-likelihood and log-prior functions.
-- The ``periodic`` and  ``reflective`` arguments that accept a list of indeces corresponding to parameters of the model that
+- The ``periodic`` and  ``reflective`` arguments that accept a list of indices corresponding to parameters of the model that
   have *periodic* or *reflective* boundary conditions. The first kind include *phase* parameters that might be periodic e.g. 
   on a range :math:`[0,2\pi]`. The latter can arise in cases where parameters are ratios where :math:`a/b` and :math:`b/a`
   are equivalent.
@@ -151,7 +151,7 @@ The sampler also accepts other arguments, for a full list see :doc:`api`. Those
 
     train_config = dict(validation_split = 0.2, # Percentage of particles to use for validation
                         epochs = 1000, # Maximum number of epochs
-                        batch_size = nparticles, # Batch size used for training
+                        batch_size = n_particles, # Batch size used for training
                         patience = 30, # Number of iterations to wait with no improvement in the (monitor) loss until stopping.
                         monitor = 'val_loss', # Which loss to monitor for early stopping. Options are 'val_loss' and 'loss'.
                         shuffle = True, # Shuffle the particles
@@ -196,7 +196,7 @@ works well for most target distributions.
 Correlation coefficient threshold :math:`\gamma` and number of MCMC steps
 """""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
 
-The :math:`\gamma` factor ( default is ``gamma = 0.75``) determnines the correlation coefficient threshold of the MCMC trajectories
+The :math:`\gamma` factor ( default is ``gamma = 0.75``) determines the correlation coefficient threshold of the MCMC trajectories
 of the particles. The lower this threshold, the more MCMC steps are performed per iteration and thus less correlated the final positions
 of the particles are with respect to their initial ones. In other words, the value :math:`\gamma` determines the number of MCMC steps
 to perform. The default value works well for most problems. If you find that your particles are not decorrelated enough and you want
@@ -224,7 +224,7 @@ This is a dictionary which includes the following arrays:
 2. ``results['loglikelihood']`` - Array with the **values of the log-likelihood** for the posterior samples given by ``results['samples']``.
 3. ``results['logprior']`` - Array with the **values of the log-prior** for the posterior samples given by ``results['samples']``.
 4. ``results['logz']`` - Array with the evolution of the estimate of the **logarithm of the model evidence** :math:`\log\mathcal{Z}`. This is usually what you need for model comparison.
-5. ``results['iter']`` - Array with number iteration indeces (e.g. ``np.array([0, 1, 2, ...])``)
+5. ``results['iter']`` - Array with number iteration indices (e.g. ``np.array([0, 1, 2, ...])``)
 6. ``results['x']`` - Array with the final samples from all the intermediate distributions.
 7. ``results['logl']`` - Array with the values of the log-likelihood for the samples from all the intermediate distributions.
 8. ``results['logp']`` - Array with the values of the log-prior for the samples from all the intermediate distributions.
@@ -242,7 +242,7 @@ Visualising the results
 
 First of all, we can plot the *run-plot* that shows us various metrics calculated during the run using the command::
 
-    import matplotliib.pyplot as plt
+    import matplotlib.pyplot as plt
 
     pc.plotting.run(results)
     plt.show()
@@ -279,7 +279,7 @@ Parallelisation
 If you want to run computations in parallel, ``pocoMC`` can use a user-defined ``pool`` to execute a variety of expensive operations 
 in parallel rather than in serial. This can be done by passing the ``pool`` object to the sampler upon initialization::
 
-    sampler = pc.Sampler(n_particles = nparticles,
+    sampler = pc.Sampler(n_particles = n_particles,
                          n_dim = ndim,
                          log_likelihood = log_like,
                          log_prior = log_prior,
@@ -328,20 +328,20 @@ it is a good idea to specify how many CPUs should be used for that using::
 
     os.environ["OMP_NUM_THREADS"] = "1"
 
-at the beggining of the code. This can affect the speed of the normalising flow training.
+at the beginning of the code. This can affect the speed of the normalising flow training.
 
 Finally, other pools can also be used, particularly if you plan to use ``pocoMC`` is a supercomputing cluster you may want to use
 an ``mpi4py`` pool so that you can utilise multiple nodes.
 
-The speed-up offered by parallisation in ``pocoMC`` is expected to be linear in the number of particles ``n_particles``.
+The speed-up offered by parallelisation in ``pocoMC`` is expected to be linear in the number of particles ``n_particles``.
 
 
 Saving and resuming runs
 ------------------------
 
 A useful option, especially for long runs, is to be able to store the state of ``pocoMC`` in a file and also the to use
-that file in order to later continue the same run. This can help avoid disastrous situations in which a run is interupted
-or terminated prematurely (e.g. due to time limititation in computing clusters or possible crashes).
+that file in order to later continue the same run. This can help avoid disastrous situations in which a run is interrupted
+or terminated prematurely (e.g. due to time limitation in computing clusters or possible crashes).
 
 Fortunately, ``pocoMC`` offers both options to save and load a previous state of the sampler.
 
@@ -361,7 +361,7 @@ something like::
 The default directory in which the state files are saved is a folder named ``states`` in the current directory. One can change
 this using the ``output_dir`` argument when initialising the sampler (e.g. ``output_dir = "new_run"``). By default, the state
 files follow the naming convention ``pmc_{i}.state`` where ``i`` is the iteration index. For instance, if ``save_every=3`` was 
-specified then the ``output_dir`` directory will incliude the files ``pmc_3.state``, ``pmc_6.state``, etc. One can also change
+specified then the ``output_dir`` directory will include the files ``pmc_3.state``, ``pmc_6.state``, etc. One can also change
 the label from ``pmc`` to anything else by using the ``output_label`` argument when initialising the sampler (e.g. 
 ``output_label="grav_waves"``).
 
diff --git a/docs/source/background.rst b/docs/source/background.rst
index ed092ce..e9ab9ff 100644
--- a/docs/source/background.rst
+++ b/docs/source/background.rst
@@ -58,7 +58,7 @@ For the task of Bayesian model comparison, one is interested in the ratio of pos
     \frac{p(\mathcal{M}_{i}\vert d)}{p(\mathcal{M}_{j}\vert d)} = \frac{p(d\vert\mathcal{M}_{i})}{p(d\vert\mathcal{M}_{j})} \times \frac{p(\mathcal{M}_{i})}{p(\mathcal{M}_{j})}
 
 where the first term on the right-hand-side is the so called *Bayes factor* and the second term is the ratio of
-prior probabilities of the two models. The latter is often set to 1 (i.e. no model is prefered a priori). The 
+prior probabilities of the two models. The latter is often set to 1 (i.e. no model is preferred a priori). The
 Bayes factor on the other hand is simply the ratio of the model evidences of the two models, or
 
 .. math::
@@ -114,7 +114,7 @@ Normalising Flow Preconditioning
 
 The purpose of the *mutation* step in SMC is to diversify the particles and allow their distribution to reach the 
 equilibrium distribution :math:`p_{t}(\theta)`. To this end, common MCMC methods (e.g. Metropolis-Hastings) are
-often employed in practice. However, the sampling efficiency of MCMC is affected by the presense of correlation
+often employed in practice. However, the sampling efficiency of MCMC is affected by the presence of correlation
 between the parameters and multimodality. In order to maintain their efficiency, PMC and `pocoMC` utilise a 
 *Normalising Flow (NF)* (i.e. an invertible transformation parameterised by a neural network) to precondition or
 decorrelate the parameters of the target distribution. *Normalising Flow Preconditioning* works by approximately
@@ -135,4 +135,4 @@ with sampling procedure of PMC, and consists of:
 
 1. The number of particles that will be used to sample. The default value is 1000 but more might be needed for challenging or high dimensional target distributions.
 2. The *effective sample size (ESS)* that will be maintain constant during the run. The default value is :math:`ESS=0.95`. This effectively determines the schedule of the :math:`\beta_{t}` values as well, with higher ESS resulting in great number of :math:`\beta_{t}` levels and more conservative exploration.
-3. The *correlation coefficient (CC)* which takes values in the range :math:`(0,1)` determines how long the particles are propaged with MCMC in each iteration. As the value of CC expresses the mean correlation between the current and initial distribution of particles,  a lower value results in more MCMC steps per iteration and thus more careful and conservative exploration of the target distribution. The default value is :math:`CC=0.75`.
\ No newline at end of file
+3. The *correlation coefficient (CC)* which takes values in the range :math:`(0,1)` determines how long the particles are propaged with MCMC in each iteration. As the value of CC expresses the mean correlation between the current and initial distribution of particles, a lower value results in more MCMC steps per iteration and thus more careful and conservative exploration of the target distribution. The default value is :math:`CC=0.75`.
\ No newline at end of file
diff --git a/docs/source/examples.rst b/docs/source/examples.rst
index 706415b..7c86231 100644
--- a/docs/source/examples.rst
+++ b/docs/source/examples.rst
@@ -11,7 +11,7 @@ is required to be installed to generate some of the plots.
 Gaussian Shells
 ===============
 
-The ``Gaussian Shells`` are a clear example of a bimodal distribution. Allthough
+The ``Gaussian Shells`` are a clear example of a bimodal distribution. Although
 the example is only in 2D it is a difficult target for many MCMC methods due to 
 its peculiar geometry in each mode.
 
diff --git a/docs/source/tutorials/bridge_sampling.ipynb b/docs/source/tutorials/bridge_sampling.ipynb
index f336327..71310c1 100644
--- a/docs/source/tutorials/bridge_sampling.ipynb
+++ b/docs/source/tutorials/bridge_sampling.ipynb
@@ -2,21 +2,33 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "# Estimating the Model Evidence using Bridge Sampling"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Background"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "In the previous tutorial we showed how one can esstimate the **model evidence** $\\mathcal{Z}$ for two models in order to perform **Bayesian model comparison**. To this end, we computed the **Bayes factor** as the ratio of the two model evidences (i.e. one for each model). The estimates of $\\mathcal{Z}$ which we used were the ones provided by ``pocoMC`` at the end of the *PMC* run.\n",
     "\n",
@@ -35,7 +47,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Defining the likelihood and prior\n",
     "\n",
@@ -45,7 +61,11 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "import numpy as np\n",
@@ -80,7 +100,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Running PMC\n",
     "\n",
@@ -90,7 +114,11 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -137,7 +165,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Bridge Sampling\n",
     "\n",
@@ -147,7 +179,11 @@
   {
    "cell_type": "code",
    "execution_count": 3,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "name": "stdout",
@@ -165,7 +201,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "The Bridge Sampling estimate is much closer to the true value of $\\log\\mathcal{Z}_{a}-5.804$ and the error-bar looks reasonable. Although the difference between the two methods is small in this example, this is not generally the case. In particular, for higher dimensional problems,  Bridge Sampling should always be preferred!"
    ]
@@ -173,7 +213,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": []
   }
@@ -205,4 +249,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/docs/source/tutorials/fitting.ipynb b/docs/source/tutorials/fitting.ipynb
index 6d7bb96..47e3769 100644
--- a/docs/source/tutorials/fitting.ipynb
+++ b/docs/source/tutorials/fitting.ipynb
@@ -2,13 +2,17 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "# Fitting a model to data\n",
     "\n",
     "In this tutorial we will see how one can use ``pocoMC`` to fit a model to some data.\n",
     "\n",
-    "The model that we will use is a *cosine* oscilation of frequency $\\omega$ and phase $\\phi$ with a linearly-evolving amplitude:\n",
+    "The model that we will use is a *cosine* oscillation of frequency $\\omega$ and phase $\\phi$ with a linearly-evolving amplitude:\n",
     "\n",
     "$$\n",
     "x(t) = (A + B\\,t) \\, cos(\\omega\\,t + \\phi)\n",
@@ -20,7 +24,11 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "import numpy as np\n",
@@ -30,9 +38,13 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
-    "## Definiing the model and generating the data\n",
+    "## Defining the model and generating the data\n",
     "\n",
     "We define the model as function that accepts an array of parameters and an array of the time instances. We then use the same model to\n",
     "generate some data by specifying the *true* parameters (these will be what we will attempt to estimate), a noise level ``sigma`` and the\n",
@@ -42,7 +54,11 @@
   {
    "cell_type": "code",
    "execution_count": 17,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "data": {
@@ -94,7 +110,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Log-likelihood, log-prior and bounds\n",
     "\n",
@@ -104,7 +124,11 @@
   {
    "cell_type": "code",
    "execution_count": 18,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "# Log-likelihood functions\n",
@@ -129,7 +153,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Initial particle positions\n",
     "\n",
@@ -139,7 +167,11 @@
   {
    "cell_type": "code",
    "execution_count": 19,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "n_particles = 1000\n",
@@ -150,7 +182,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Initialising and running Preconditioned Monte Carlo\n",
     "\n",
@@ -169,7 +205,11 @@
   {
    "cell_type": "code",
    "execution_count": 20,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "# Sampler initialisation\n",
@@ -191,7 +231,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Results\n",
     "\n",
@@ -201,7 +245,11 @@
   {
    "cell_type": "code",
    "execution_count": 23,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "results = sampler.results"
@@ -210,7 +258,11 @@
   {
    "cell_type": "code",
    "execution_count": 24,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "data": {
@@ -232,7 +284,11 @@
   {
    "cell_type": "code",
    "execution_count": 45,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "data": {
@@ -254,7 +310,11 @@
   {
    "cell_type": "code",
    "execution_count": 44,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "data": {
@@ -275,7 +335,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "We can also compute any expectation values we want using the posterior samples."
    ]
@@ -283,7 +347,11 @@
   {
    "cell_type": "code",
    "execution_count": 47,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "name": "stdout",
@@ -301,7 +369,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Posterior predictive check\n",
     "\n",
@@ -320,7 +392,11 @@
   {
    "cell_type": "code",
    "execution_count": 28,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "sim_data = []\n",
@@ -334,7 +410,11 @@
   {
    "cell_type": "code",
    "execution_count": 42,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "data": {
@@ -365,7 +445,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": []
   }
@@ -397,4 +481,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/docs/source/tutorials/model_comparison.ipynb b/docs/source/tutorials/model_comparison.ipynb
index c2fdabc..d4b8465 100644
--- a/docs/source/tutorials/model_comparison.ipynb
+++ b/docs/source/tutorials/model_comparison.ipynb
@@ -2,7 +2,11 @@
  "cells": [
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "# Comparing two models\n",
     "\n",
@@ -26,7 +30,11 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "import numpy as np\n",
@@ -37,7 +45,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Defining the models and generating the data\n",
     "\n",
@@ -48,7 +60,11 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "data": {
@@ -105,17 +121,25 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Log-likelihood, log-prior and parameter bounds\n",
     "\n",
-    "We then define the log-likelihood and log-prior functions, as well as the low & high parameter bounds. We will use the same functions for both models and provide the model as an additional argument. The parameter bounds are identical for the common parameters of the two models. The bounds of the extended model include an extra parameter of course for the linear dependance of the amplitude."
+    "We then define the log-likelihood and log-prior functions, as well as the low & high parameter bounds. We will use the same functions for both models and provide the model as an additional argument. The parameter bounds are identical for the common parameters of the two models. The bounds of the extended model include an extra parameter of course for the linear dependence of the amplitude."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 3,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": [
     "# Log-likelihood function used for either model\n",
@@ -145,7 +169,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Initialising and running Preconditioned Monte Carlo\n",
     "\n",
@@ -157,7 +185,11 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -202,7 +234,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "### Simple model"
    ]
@@ -210,7 +246,11 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "name": "stderr",
@@ -228,14 +268,14 @@
     "# Prior samples of simple model\n",
     "prior_samples_simple = np.random.uniform(low = bounds_simple.T[0],\n",
     "                                         high = bounds_simple.T[1],\n",
-    "                                         size = (_particles, n_dim_simple))\n",
+    "                                         size = (n_particles, n_dim_simple))\n",
     "\n",
     "# Initialise sampler for simple model\n",
     "sampler_simple = pc.Sampler(n_particles = n_particles,\n",
     "                            n_dim = n_dim_simple,\n",
     "                            log_likelihood = log_like,\n",
     "                            log_prior = log_prior,\n",
-    "                            bounds_simple,\n",
+    "                            bounds=bounds_simple,\n",
     "                            loglikelihood_args = [t, data, sigma, model_simple],\n",
     "                            logprior_args = [bounds_simple]\n",
     "                            )\n",
@@ -250,7 +290,21 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "source": [],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "## Bayes factor and model comparison\n",
     "\n",
@@ -260,7 +314,11 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [
     {
      "name": "stdout",
@@ -279,7 +337,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "It is important to mention here that there some caveats in the use of Bayes factors for model comparison.\n",
     "- First of all, the estimates of the model evidences $\\mathcal{Z}$ that constitute the numerator and denominator of the Bayes factor are subject to noise.\n",
@@ -290,7 +352,11 @@
   },
   {
    "cell_type": "markdown",
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
    "source": [
     "**Finally, in order to get more accurate estimates of the model evidence, you can use the Bridge Sampling estimator implemented in pocoMC (see  the next tutorial for this)!**"
    ]
@@ -298,7 +364,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
    "outputs": [],
    "source": []
   }
@@ -330,4 +400,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/pocomc/sampler.py b/pocomc/sampler.py
index 4c871b9..2f0e8da 100644
--- a/pocomc/sampler.py
+++ b/pocomc/sampler.py
@@ -1,4 +1,3 @@
-import warnings
 from pathlib import Path
 from typing import Union
 
diff --git a/pocomc/scaler.py b/pocomc/scaler.py
index 9058a23..a347de4 100644
--- a/pocomc/scaler.py
+++ b/pocomc/scaler.py
@@ -1,4 +1,3 @@
-import imp
 from typing import Union, List
 
 import numpy as np


Test patch:
diff --git a/tests/test_flow.py b/tests/test_flow.py
index 8093ba8..ef9fb24 100644
--- a/tests/test_flow.py
+++ b/tests/test_flow.py
@@ -1,6 +1,5 @@
 import unittest
 import torch
-
 from pocomc.flow import Flow
 
 
@@ -49,6 +48,7 @@ def test_logprob(self):
 
         x = self.make_data()
         flow = Flow(ndim=x.shape[1])
+
         log_prob = flow.logprob(x)
 
         self.assertFalse(torch.any(torch.isnan(log_prob)))
@@ -109,7 +109,8 @@ def test_logprob_float64(self):
         x = self.make_data()
         x = x.double()
         flow = Flow(ndim=x.shape[1])
-        log_prob = flow.logprob(x)
+        with self.assertWarns(UserWarning):
+            log_prob = flow.logprob(x)
 
         self.assertFalse(torch.any(torch.isnan(log_prob)))
         self.assertFalse(torch.any(torch.isinf(log_prob)))