public

jaxns.experimental.public

Module Contents

class GlobalOptimisation

A global optimiser using nested sampling as the core algorithm. Can easily globally optimise complex models, with curving degeneracies and multimodal structure. Highly parallelisable. Recommended to use gradient information by setting gradient_slice=True.

Note, the log-likelihood over the model is maximised NOT the posterior. The prior acts as the search space prior, by constraining the search space and giving search preference to regions of high prior probability. Thus, the prior should encode your prior belief about where you think the global maximum is located.

Parameters:

• model – a model to perform global optimisation on over the sample space.

• num_search_chains – number of search chains to use.

• s – number of slices to use per dimension.

• k – number of phantom samples to use.

• gradient_slice – if true use gradient information to improve. Default True.

• shell_frac – fraction of the shell to discard in parallel.

• devices – devices to use for parallel sharded computation. Default all available devices.

• verbose – whether to print verbose output. Default False.

model: jaxns.framework.bases.BaseAbstractModel

num_search_chains: int | None = None

s: int | None = None

k: int | None = None

gradient_slice: bool = True

shell_frac: float | None = None

devices: jaxlib.xla_client.Device | None = None

verbose: bool = False

__post_init__()

__call__(key, term_cond=None, finetune=False)

Runs the global optimisation.

Parameters:

• key (jaxns.internals.types.PRNGKey) – PRNGKey

• term_cond (Optional[jaxns.experimental.GlobalOptimisationTerminationCondition]) – termination condition

• finetune (bool) – whether to use gradient-based fine-tune. Default False because not all models have gradients.

Returns:

results of the global optimisation

Return type:

jaxns.experimental.GlobalOptimisationResults

DefaultGlobalOptimisation