bilby.core.utils.parallel.bilby_pool

bilby.core.utils.parallel.bilby_pool(likelihood=None, priors=None, use_ratio=None, search_parameter_keys=None, npool=None, pool=None, parameters=None)

Yield a parallel pool object that is initialized with variables typically needed by Bilby for parallel tasks that is automatically close when closing the context.

Parameters:

likelihood: bilby.core.likelihood.Likelihood, None

The likelihood to copy into each process

priors: bilby.core.prior.PriorDict, None

The Bilby prior dictionary to copy into each process

use_ratio: bool, None

Whether to evaluate the log_likelihood_ratio

search_parameter_keys: list[str], None

The names for parameters being sampled over

npool: int, None

The number of processes to use for multiprocessing. If a user pool is not provided and this is either 1 or None, this functions returns None.

pool: pool-like, str, None

Either a premade pool object, or the pool kind (mpi, multiprocessing). If a pre-made pool is passed, it is returned directly with no checks performed.

parameters: dict, None

Parameters to pass through to the new processes, e.g., if default parameters are to be passed.

Yields:

pool: schwimmbad.MPIPool, multiprocessing.Pool, None

Returns either a pool that can be used for mapping function calls. Each process attached to the pool has been initialized with the bilby.core.utils.parallel.sampling_convenience_dump.

Examples

>>> import numpy as np
>>> from bilby.core.likelihood import AnalyticalMultidimensionalCovariantGaussian
>>> from bilby.core.prior import Normal, PriorDict
>>> from bilby.core.utils.parallel import bilby_pool

>>> likelihood = AnalyticalMultidimensionalCovariantGaussian(
...     mean=np.zeros(4), cov=np.eye(4)
... )
>>> priors = PriorDict({f"x{ii}": Normal(0, 1) for ii in range(4)})
>>> parameters = [priors.sample() for _ in range(10)]
>>> with bilby_pool(likelihood, priors, npool=4) as pool:
...     log_ls = list(pool.map(likelihood.log_likelihood, parameters))

Note

The above example passes the likelihood.log_likelihood method directly to the pool. This is possible, but will lead to the likelihood object being pickled and sent to each process for every call to the pool. This can add significant overhead if the likelihood carries a lot of data, e.g., for the gravitational-wave transient likelihoods. In this case, we recommend creating a wrapper function that uses the sampling_convenience_dump to access the likelihood object, e.g.,

>>> from bilby.core.utils.parallel import sampling_convenience_dump

>>> def parallel_likelihood_eval(parameters):
...     return sampling_convenience_dump.likelihood.log_likelihood(parameters)