nessai.utils.multiprocessing#

Utilities related to multiprocessing.

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.
`log_prior_unit_hypercube_wrapper`(x)	Wrapper for the log-prior in the unit-hypercube for use with
`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.

Module Contents#

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.log_prior_unit_hypercube_wrapper(x)#

Wrapper for the log-prior in the unit-hypercube 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 in the unit-hypercube.

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