nessai.utils.multiprocessing

Utilities related to multiprocessing.

Module Contents

Functions

get_n_pool(pool)	Determine the number of processes in a multiprocessing pool.
check_multiprocessing_start_method()	Check the multiprocessing start method.
initialise_pool_variables(model)	Prepare the model for use with a multiprocessing pool.
log_likelihood_wrapper(x)	Wrapper for the log-likelihood for use with multiprocessing.
log_prior_wrapper(x)	Wrapper for the log-prior for use with multiprocessing.
batch_evaluate_function(func, x, vectorised[, ...])	Evaluate a function over a batch of inputs.
check_vectorised_function(func, x[, dtype, atol, rtol])	Check if a function is vectorised given a set of inputs.

nessai.utils.multiprocessing.get_n_pool(pool)

Determine the number of processes in a multiprocessing pool.

Parameters:

poolobject

Multiprocessing pool or similar.

Returns:

int or None

Number of processes. Returns None if number could not be determined.

nessai.utils.multiprocessing.check_multiprocessing_start_method()

Check the multiprocessing start method.

Print a warning if the start method is not fork.

nessai.utils.multiprocessing.initialise_pool_variables(model)

Prepare the model for use with a multiprocessing pool.

Makes a global copy of the model. Should be called before initialising a pool or passed to the initializer argument with the model as one of the initargs.

Parameters:

modelnessai.model.Model

Model to be copied to a global variable.

nessai.utils.multiprocessing.log_likelihood_wrapper(x)

Wrapper for the log-likelihood for use with multiprocessing.

Should be used alongside nessai.utils.multiprocessing.initialise_pool_variables()

Parameters:

xnumpy.ndarray

Array of samples.

Returns:

numpy.ndarray

Array of log-likelihoods.

nessai.utils.multiprocessing.log_prior_wrapper(x)

Wrapper for the log-prior for use with multiprocessing.

Should be used alongside nessai.utils.multiprocessing.initialise_pool_variables()

Parameters:

xnumpy.ndarray

Array of samples.

Returns:

numpy.ndarray

Array of log-prior values.

nessai.utils.multiprocessing.batch_evaluate_function(func, x, vectorised, chunksize=None, pool=None, n_pool=None, func_wrapper=None)

Evaluate a function over a batch of inputs.

Parameters:

funcCallable

The function to evaluate.

xnumpy.ndarray

The values over which to evaluate the function.

vectorisedbool

Boolean to indicate if the function is vectorised or not.

chunksizeOptional[int]

Maximum number of inputs that will be passed to the function at once. Only applies if the function is vectorised.

poolOptional[multiprocessing.Pool]

Multiprocessing pool used to evaluate the function.

n_poolOptional[int]

Number of cores in the multiprocessing. Determines how values are split when calling Pool.map.

func_wrapperOptional[Callable]

Wrapper to the function to use instead of the function when using the pool.

Returns:

numpy.ndarray

Array of outputs

nessai.utils.multiprocessing.check_vectorised_function(func, x, dtype=None, atol=1e-15, rtol=1e-15)

Check if a function is vectorised given a set of inputs.

Parameters:

funcCallable

Function to test

xnumpy.ndarray

Inputs over which to evaluate the function

dtypeOptional[Union[str, numpy.dtype]]

Dtype to cast the results to. If not specified, the default dtype for livepoints is used.

rtolfloat

Relative tolerance, see numpy documentation for allclose.

atolfloat

Abosoulte tolerance, see numpy documentation for allclose.

Returns:

numpy.ndarray

Array of outputs