###########################
Parallelisation with nessai
###########################

One benefit of the proposal method used in ``nessai`` is it allows for simple parallelisation of the likelihood evaluation since new live points are drawn in batches and then stored. The likelihood can therefore be precomputed and stored later use.

************************
Enabling parallelisation
************************

There are two keyword arguments that must be set to enable parallelisation:

- :code:`n_pool`: The number of threads to use for evaluating the likelihood
- :code:`max_threads`: The maximum number of threads to use, this should be at least 1 larger than :code:`n_pool`. Extra threads are allocated to PyTorch's CPU parallelisation.

.. note::
    If running ``nessai`` via a job scheduler such as HTCondor, remember to set the number of requested threads accordingly. This should match :code:`max_threads`.


*****************
Specifying a pool
*****************

Alternatively, ``nessai`` can use a user-defined pool. This is specified by setting the :code:`pool` argument in :code:`NestedSampler` or :code:`FlowSampler`. Some variables must be initialised when creating the pool, this is done using :py:func:`~nessai.utils.multiprocessing.initialise_pool_variables`:

.. code-block:: python

    from multiprocessing import Pool
    from nessai.utils.multiprocessing import initialise_pool_variables

    model = GaussianModel()
    pool = Pool(
        processes=2,
        initializer=initialise_pool_variables,                                     
        initargs=(model,),
    )

:code:`pool` can then passed to the :code:`pool` keyword argument when setting up the sampler.

-------------
Using ``ray``
-------------

``ray`` includes a `distributed multiprocessing pool <https://docs.ray.io/en/latest/ray-more-libs/multiprocessing.html>`_ that can also be used with ``nessai``. Simply import :code:`ray.util.multiprocessing.Pool` instead of the standard pool and initialise using the method described above.


*************
Example usage
*************

.. literalinclude:: ../examples/parallelisation_example.py
    :language: python

********
See also
********

- :py:func:`nessai.utils.threading.configure_threads`
- :py:func:`nessai.utils.multiprocessing.initialise_pool_variables`
- :py:meth:`nessai.model.Model.configure_pool`
- :py:meth:`nessai.model.Model.close_pool`