PolyChord

Bases: AbstractHypercubeSampler

The PolyChord Nested Sampler wrapped in a JAX interface.

Acts as an adapter layer between JAX PyTrees and PolyChord’s required flat 1D NumPy arrays. It automatically handles flattening and unflattening of complex parameter structures, JIT-compiles the likelihood and prior transforms for performance, and bridges JAX operations with PolyChord’s host-based MPI sampling routines.

All parameters default to None, deferring to PolyChord’s internal defaults.

Parameters:

nlive (int | None) – The number of live points (PolyChord default: 25 * ndims).
num_repeats (int | None) – The length of the slice sampling chain (PolyChord default: 5 * ndims).
nprior (int | None) – The number of prior samples to draw before clustering begins.
nfail (int | None) – The number of failed slice sampling steps before giving up.
do_clustering (bool | None) – Whether to use k-means clustering to handle multimodal posteriors.
feedback (int | None) – The level of output written to stdout. 0=none, 1=standard, 2=detailed.
precision_criterion (float | None) – The stopping criterion based on the estimated evidence precision.
logzero (float | None) – The numerical value used to represent log(0).
boost_posterior (float | None) – Boost the number of live points near the peak to improve posterior samples.
posteriors (bool | None) – Whether to produce standard posterior output files.
equals (bool | None) – Whether to output equally weighted posterior samples.
cluster_posteriors (bool | None) – Whether to produce posterior output files for individual clusters.
write_resume (bool | None) – Whether to continuously write resume files during the run.
write_paramnames (bool | None) – Whether to generate a .paramnames file for post-processing tools.
read_resume (bool | None) – Whether to attempt resuming from a previous partially completed run.
write_stats (bool | None) – Whether to write run statistics to a .stats file.
write_live (bool | None) – Whether to dump the current live points to disk.
write_dead (bool | None) – Whether to record the dead points (the core nested sampling output) to disk.
write_prior (bool | None) – Whether to write prior samples to disk.
maximise (bool | None) – Whether to perform a maximization phase to find the exact MAP estimate.
compression_factor (float | None) – The compression factor used for slice sampling.
synchronous (bool | None) – Whether to run MPI operations synchronously.
base_dir (str | None) – The base directory path where all output files will be saved.
file_root (str | None) – The root naming convention for all generated output files.
cluster_dir (str | None) – The directory name for cluster-specific outputs.
seed (int | None) – Random seed for the sampler. Uses time if set to -1.
nlives (dict | None) – A dictionary mapping log-likelihood contours to the number of live points.
paramnames (list of tuple | None) – A list of parameter names and LaTeX formatted names, e.g., [(“p1”, r” heta_1”)].

run(loglikelihood_fn: Callable[[PyTree, Any], Shaped[jaxlib._jax.Array, '']], prior_transform_fn: Callable[[PyTree, Any], PyTree], u0: PyTree, args: PyTree[Any], key: Array, init_cube_samples: PyTree | None = None, max_steps: int | None = None, **kwargs) → SampleResult

Execute the sampling algorithm.

Parameters:

loglikelihood_fn (callable) – A function taking the physical parameters and args as input and returning the log-likelihood.
prior_transform_fn (callable) – A function taking the hypercube parameters and args as input and returning the physical parameters.
u0 (PyTree) – The initial parameters in the unit hypercube, either for shape reference or as a starting point.
args (Any) – Args to pass to fn.
key (Array) – A random JAX key.
init_cube_samples (PyTree, optional) – An optional batched PyTree the same structure as u0 with initial hypercube samples to warm-start the algorithm.
max_steps (int, optional) – The maximum number of sampling steps to take. If None, implies there should be no limit.
**kwargs – Runtime arguments forward to the solver backend.

Returns:

An instance of pmrf.infer.SampleResult.

Return type:

results

base_dir: str | None = None

boost_posterior: float | None = None

cluster_dir: str | None = None

cluster_posteriors: bool | None = None

compression_factor: float | None = None

do_clustering: bool | None = None

equals: bool | None = None

feedback: int | None = None

file_root: str | None = None

logzero: float | None = None

maximise: bool | None = None

nfail: int | None = None

nlive: int | None = None

nlives: Dict[float, int] | None = None

nprior: int | None = None

num_repeats: int | None = None

paramnames: list[tuple[str, str]] | None = None

posteriors: bool | None = None

precision_criterion: float | None = None

read_resume: bool | None = None

property requires_hypercube

seed: int | None = None

synchronous: bool | None = None

write_dead: bool | None = None

write_live: bool | None = None

write_paramnames: bool | None = None

write_prior: bool | None = None

write_resume: bool | None = None

write_stats: bool | None = None