blackjax.ns.utils

blackjax.ns.utils#

Utility functions for setting up and post-processing Nested Sampling runs.

References

Functions#

log1mexp(→ blackjax.types.Array)

Computes log(1 - exp(x)) in a numerically stable way.

compute_num_live(→ blackjax.types.Array)

Compute the effective number of live points at each death contour (Fowlie, Handley & Su, 2021).

logX(→ tuple[blackjax.types.Array, blackjax.types.Array])

Simulate the stochastic evolution of log prior volumes (Skilling, 2006).

log_weights(→ blackjax.types.Array)

Calculate the log importance weights for Nested Sampling results.

finalise(→ blackjax.ns.base.NSInfo)

Combines the history of dead particle information with the final live points.

ess(→ blackjax.types.Array)

Computes the Effective Sample Size (ESS) from log-weights.

sample(→ blackjax.types.ArrayTree)

Resamples particles according to their importance weights.

get_first_row(→ blackjax.types.ArrayTree)

Extracts the first "row" (element along the leading axis) of each leaf in a PyTree.

uniform_prior(→ tuple[blackjax.types.ArrayTree, Callable])

Sample initial particles and build a log-prior for a box-uniform prior.

Module Contents#

log1mexp(x: blackjax.types.Array) → blackjax.types.Array[source]#

Computes log(1 - exp(x)) in a numerically stable way.

compute_num_live(info: blackjax.ns.base.NSInfo) → blackjax.types.Array[source]#

Compute the effective number of live points at each death contour (Fowlie, Handley & Su, 2021).

When doing batch deletions, the jump in energy level can be smoothed by transforming 1 jump of size k into k jumps of size 1. This function computes the effective population size associated with this transformation.

Expects the complete finalised output – the dead points together with the final live particles (e.g. from finalise()): it relies on every particle’s birth event being present. Called on a dead-only subset that omits the initial live particles’ births, the live counts are wrong (1 instead of N for a standard run).

Returns:

An array where each element num_live[j] is the effective number of live points m*_i when the j-th particle (in the sorted list of dead particles) was considered “dead”.

Return type:

Array

logX(rng_key: blackjax.types.PRNGKey, dead_info: blackjax.ns.base.NSInfo, shape: int = 100) → tuple[blackjax.types.Array, blackjax.types.Array][source]#

Simulate the stochastic evolution of log prior volumes (Skilling, 2006).

Wraps the effective population size in compute_num_live, along with stochastic simulation of the log prior shrinkage associated with each deleted particle.

Parameters:

  • rng_key – A JAX PRNG key for generating uniform random variates.

  • dead_info – An NSInfo object (or compatible PyTree) containing loglikelihood_birth and loglikelihood for all dead particles accumulated during an NS run. It’s assumed these particles are already sorted by their death log-likelihood.

  • shape – The shape of Monte Carlo samples to generate for the stochastic log-volume sequence. Each sample represents one possible path of volume shrinkage. Default is 100.

Returns:

  • logX_cumulative: An array of shape (num_dead_particles, shape) containing shape simulated sequences of cumulative log prior volumes log(X_i).

  • log_dX_elements: An array of shape (num_dead_particles, shape) containing shape simulated sequences of log prior volume elements log(dX_i). dX_i is the trapezoidal volume element (X_{i-1} - X_{i+1}) / 2.

Return type:

tuple[Array, Array]

log_weights(rng_key: blackjax.types.PRNGKey, dead_info: blackjax.ns.base.NSInfo, shape: int = 100, beta: float = 1.0) → blackjax.types.Array[source]#

Calculate the log importance weights for Nested Sampling results.

Parameters:

  • rng_key – A JAX PRNG key for simulating log(dX_i).

  • dead_info – An NSInfo object (or compatible PyTree) containing loglikelihood_birth and loglikelihood for all dead particles.

  • shape – The shape of Monte Carlo samples to use for simulating log(dX_i). Default is 100.

  • beta – The inverse temperature. Typically 1.0 for standard evidence calculation. Allows for reweighting to different temperatures.

Returns:

An array of log importance weights, shape (num_dead_particles, *shape). The original order of particles in dead_info is preserved.

Return type:

Array

finalise(live: blackjax.ns.base.NSState, dead: list[blackjax.ns.base.NSInfo], update_info: bool = True) → blackjax.ns.base.NSInfo[source]#

Combines the history of dead particle information with the final live points.

Parameters:

  • live – The final NSState of the Nested Sampler, containing the live particles.

  • dead – A list of NSInfo objects, where each object contains information about the particles that “died” at one step of the NS algorithm.

  • update_info – Whether to concatenate the update_info from each element of dead. If False, the returned update_info is None. Default is True.

Returns:

A single NSInfo whose particles field concatenates all dead particles with the final live particles. When update_info=True the update_info field concatenates the update_info from each element of dead only – no entry is added for the final live points, so it is shorter than particles by the number of live points. When update_info=False the update_info field is None.

Return type:

NSInfo

ess(rng_key: blackjax.types.PRNGKey, dead: blackjax.ns.base.NSInfo) → blackjax.types.Array[source]#

Computes the Effective Sample Size (ESS) from log-weights.

Parameters:

  • rng_key – A JAX PRNG key, used by log_weights.

  • dead – An NSInfo object containing the full set of dead (and final live) particles, typically the output of finalise.

Returns:

The mean Effective Sample Size, a scalar float.

Return type:

Array

sample(rng_key: blackjax.types.PRNGKey, dead: blackjax.ns.base.NSInfo, shape: int = 1000) → blackjax.types.ArrayTree[source]#

Resamples particles according to their importance weights.

Parameters:

  • rng_key – A JAX PRNG key, used by log_weights and for resampling.

  • dead – An NSInfo object containing the full set of dead (and final live) particles, typically the output of finalise.

  • shape – The number of resampled particles to draw. Default is 1000.

Returns:

A PyTree of resampled particles, where each leaf has shape.

Return type:

ArrayTree

get_first_row(x: blackjax.types.ArrayTree) → blackjax.types.ArrayTree[source]#

Extracts the first “row” (element along the leading axis) of each leaf in a PyTree.

This is typically used to get a single particle’s structure or values from a PyTree representing a collection of particles, where the leading dimension of each leaf array corresponds to the particle index.

Parameters:

x – A PyTree of arrays, where each leaf array has a leading dimension.

Returns:

A PyTree with the same structure as x, but where each leaf is the first slice leaf[0] of the corresponding leaf in x.

Return type:

ArrayTree

uniform_prior(rng_key: blackjax.types.PRNGKey, num_live: int, bounds: dict[str, tuple[float, float]]) → tuple[blackjax.types.ArrayTree, Callable][source]#

Sample initial particles and build a log-prior for a box-uniform prior.

Parameters:

  • rng_key – A JAX PRNG key for random number generation.

  • num_live – The number of live particles to sample.

  • bounds – A dictionary mapping parameter names to their bounds (tuples of min and max). Each parameter will be sampled uniformly within these bounds. Example: {‘param1’: (0.0, 1.0), ‘param2’: (-5.0, 5.0)}

Returns:

  • particles: A PyTree of sampled parameters, where each leaf has shape (num_live,).

  • logprior_fn: A function that computes the log-prior probability for a given set of parameters.

Return type:

tuple