blackjax.ns.base

blackjax.ns.base#

Base components for Nested Sampling.

Defines the particle state carrying loglikelihood information, a generic kernel builder that deletes the lowest-likelihood particles and replaces them with an inner kernel, and the default deletion strategy selecting the num_delete particles with the lowest loglikelihoods.

References

Classes#

`NSState`	State of the Nested Sampler.
`NSInfo`	Additional information returned at each step of the Nested Sampling algorithm.

Functions#

`init`(→ NSState)	Initializes the Nested Sampler state.
`build_kernel`(→ Callable)	Build a generic Nested Sampling kernel.
`delete_fn`(→ tuple[blackjax.types.Array, ...)	Identifies particles to be deleted.

Module Contents#

class NSState[source]#

State of the Nested Sampler.

At the most basic level, this is just a wrapper around a StateWithLogLikelihood; richer NS implementations (e.g. AdaptiveNSState) carry extra fields.

particles[source]#: The StateWithLogLikelihood of the current live particles.

particles: StateWithLogLikelihood[source]#

class NSInfo[source]#

Additional information returned at each step of the Nested Sampling algorithm.

particles[source]#: The StateWithLogLikelihood of particles that were marked as “dead” (replaced).

update_info[source]#: A NamedTuple (or any PyTree) containing information from the update step (inner kernel) used to generate new live particles.

particles: StateWithLogLikelihood[source]#

update_info: NamedTuple[source]#

init(positions: blackjax.types.ArrayLikeTree, init_state_fn: Callable, loglikelihood_birth: float = jnp.nan) → NSState[source]#

Initializes the Nested Sampler state.

Parameters:

positions – An initial set of positions (PyTree of arrays) drawn from the prior distribution. The leading dimension of each leaf array must be equal to the number of positions.
init_state_fn – A function that builds the particle state (StateWithLogLikelihood) from positions; init wraps the result in an NSState. Typically vmapped over the live set.
loglikelihood_birth – The initial log-likelihood birth threshold. Defaults to NaN, which implies no initial likelihood constraint beyond the prior.

Returns:

The initial state of the Nested Sampler.

Return type:

NSState

build_kernel(delete_fn: Callable, inner_kernel: Callable) → Callable[source]#

Build a generic Nested Sampling kernel.

This function creates a kernel for the Nested Sampling algorithm by combining a particle deletion function and an inner kernel for generating new particles.

Parameters:

delete_fn – A deletion function, typically partially applied with num_delete, with effective signature (state) -> (dead_idx, target_update_idx). Receives the full NS state (duck-typed) and identifies particles to be deleted and the indices to update.
inner_kernel – A kernel function with the signature (rng_key, state, loglikelihood_0) -> (new_particles, info) that generates replacement particles. Receives the full NS state (duck-typed) and a single PRNG key; returns a StateWithLogLikelihood with leading dimension num_delete. The number of particles to produce is known at construction time.

Returns:

A kernel function for Nested Sampling: (rng_key, state) -> (new_state, ns_info).

Return type:

Callable

delete_fn(state: NSState, num_delete: int) → tuple[blackjax.types.Array, blackjax.types.Array][source]#

Identifies particles to be deleted.

Selects the num_delete particles with the lowest log-likelihoods and marks them as “dead”.

Parameters:

state – The current NS state (duck-typed; must have .particles.loglikelihood).
num_delete – The number of particles to delete and subsequently replace.

Returns:

A tuple (dead_idx, target_update_idx) of indices marked for deletion
and of slots to overwrite (identical here).