blackjax.ns.nss

blackjax.ns.nss#

Nested Slice Sampling (NSS) algorithm.

An example implementation of Nested Sampling with a slice sampler as the inner MCMC kernel (Yallup, Kroupa & Handley, 2026, arXiv:2601.23252). The default build_kernel() uses hit-and-run moves shaped by the live-point covariance; build_swig_kernel() offers an axis-aligned slice-within-Gibbs alternative.

Functions#

covariance_proposal(→ Callable)

Proposal generator for nested slice sampling.

coordinate_proposal(→ Callable)

Per-axis proposal generator for nested slice-within-Gibbs (SwiG).

live_covariance(→ dict[str, blackjax.types.ArrayTree])

Compute the live-point covariance for covariance-based custom proposals.

live_covariance_factor(→ dict[str, ...)

Factor the live-point covariance once per nested-sampling step.

live_widths(→ dict[str, blackjax.types.ArrayTree])

Per-axis live-point spread (std): the per-coordinate slice widths for SwiG.

slice_constrained_step(→ Callable)

The slice-family constrained inner step for nested sampling.

build_kernel(→ Callable)

Build the Nested Slice Sampling kernel.

coordinate_constrained_step(→ Callable)

The coordinate-sweep constrained inner step for nested sampling (SwiG).

build_swig_kernel(→ Callable)

Build the Nested Slice-within-Gibbs (SwiG) kernel.

as_top_level_api(→ blackjax.SamplingAlgorithm)

Creates a Nested Slice Sampling (NSS) algorithm, blackjax.nss.

swig_as_top_level_api(→ blackjax.SamplingAlgorithm)

Creates a Nested Slice-within-Gibbs (SwiG) sampling algorithm, blackjax.nsswig.

Module Contents#

covariance_proposal(init_state_fn: Callable, loglikelihood_0: blackjax.types.Array, cov: blackjax.types.Array | None = None, *, covariance_factor: blackjax.types.Array | None = None) → Callable[source]#

Proposal generator for nested slice sampling.

The nested-sampling analogue of direction_proposal(): it steps along a covariance-shaped direction and gates the hard likelihood constraint into is_valid. The returned slice_fn builds the candidate particle (recording its log-likelihood, computed once) and reports it admissible only when loglikelihood > loglikelihood_0. Override it to write a custom nested stepper.

The default NSS kernel supplies covariance_factor so the Cholesky factorization is shared by all inner steps. cov remains supported for covariance-based custom parameter callbacks and direct callers.

Parameters:

  • init_state_fn – Builds a particle state from a position and birth log-likelihood.

  • loglikelihood_0 – Hard lower likelihood threshold for valid proposals.

  • cov – Live-point covariance matrix. Used only when covariance_factor is not supplied.

  • covariance_factor – Precomputed lower-triangular Cholesky factor of the live covariance.

Return type:

A proposal generator consumed by the univariate slice kernel.

coordinate_proposal(init_state_fn: Callable, loglikelihood_0: blackjax.types.Array, i: blackjax.types.Array, width: blackjax.types.Array) → Callable[source]#

Per-axis proposal generator for nested slice-within-Gibbs (SwiG).

The coordinate counterpart of covariance_proposal(): it steps along axis i scaled by width (the direction width * e_i) and gates the hard likelihood constraint into is_valid. Like covariance_proposal(), the move’s scale lives in the direction, so the univariate slice always runs at unit width. The returned slice_fn builds the candidate particle (recording its log-likelihood) and reports it admissible only when loglikelihood > loglikelihood_0; it threads the full particle, so the recorded loglikelihood survives the sweep. Override it to write a custom nested coordinate stepper.

live_covariance(rng_key: blackjax.types.PRNGKey, state: blackjax.ns.base.NSState, info: blackjax.ns.base.NSInfo, params: dict[str, blackjax.types.ArrayTree] | None = None) → dict[str, blackjax.types.ArrayTree][source]#

Compute the live-point covariance for covariance-based custom proposals.

Parameters:

  • rng_key – Unused key required by the adaptive-kernel callback protocol.

  • state – Nested-sampling state containing the current live particles.

  • info – Unused transition information required by the callback protocol.

  • params – Unused previous parameters required by the callback protocol.

Return type:

A parameter dictionary containing the live-point covariance.

live_covariance_factor(rng_key: blackjax.types.PRNGKey, state: blackjax.ns.base.NSState, info: blackjax.ns.base.NSInfo, params: dict[str, blackjax.types.ArrayTree] | None = None) → dict[str, blackjax.types.ArrayTree][source]#

Factor the live-point covariance once per nested-sampling step.

Parameters:

  • rng_key – Unused key required by the adaptive-kernel callback protocol.

  • state – Nested-sampling state containing the current live particles.

  • info – Unused transition information required by the callback protocol.

  • params – Unused previous parameters required by the callback protocol.

Return type:

A parameter dictionary containing the lower-triangular Cholesky factor.

live_widths(rng_key: blackjax.types.PRNGKey, state: blackjax.ns.base.NSState, info: blackjax.ns.base.NSInfo, params: dict[str, blackjax.types.ArrayTree] | None = None) → dict[str, blackjax.types.ArrayTree][source]#

Per-axis live-point spread (std): the per-coordinate slice widths for SwiG.

The coordinate counterpart of live_covariance_factor(): only the marginal per-axis spread is used, so axis correlations are deliberately ignored – the defining trait of a coordinate (slice-within-Gibbs) move. Overridable via the inner_kernel_params seam of build_swig_kernel() and swig_as_top_level_api(), mirroring live_covariance_factor().

slice_constrained_step(init_state_fn: Callable, slice_kernel: Callable, proposal: Callable) → Callable[source]#

The slice-family constrained inner step for nested sampling.

Runs slice_kernel with a constrained proposal generator built by proposal(init_state_fn, loglikelihood_0, **params); the proposal’s slice_fn gates is_valid on the likelihood contour, so the slice shrinks until it lands inside it (no wasted steps). The slice counterpart to reject_constrained_step(), consumed by build_kernel().

build_kernel(init_state_fn: Callable, num_inner_steps: int, num_delete: int = 1, max_steps: int = 10, max_shrinkage: int = 100, proposal: Callable = covariance_proposal, inner_kernel_params: Callable | None = None, update_strategy: Callable = update_with_mcmc_take_last) → Callable[source]#

Build the Nested Slice Sampling kernel.

Parameters:

  • init_state_fn – Builds a particle state from a position and birth log-likelihood.

  • num_inner_steps – Number of slice steps per new particle. Prefer num_inner_steps >= max(5, 2 * dim) for reliable mixing (bare dim is the minimum; see as_top_level_api()).

  • num_delete – Number of particles deleted and replaced per step (default 1).

  • max_steps – Cap on stepping-out expansions per slice (default 10).

  • max_shrinkage – Cap on shrinkage evaluations per slice (default 100).

  • proposal – Proposal factory (init_state_fn, loglikelihood_0, **params) -> proposal_generator (covariance_proposal() by default). The default proposal consumes a precomputed covariance_factor. Override to write a custom nested stepper.

  • inner_kernel_params – Computes the inner-kernel parameters from the live points each step, (rng_key, state, info, params) -> params. When None, uses live_covariance_factor() with the default proposal and live_covariance() with a custom proposal, preserving the existing covariance-based extension seam.

  • update_strategy – Inner-kernel factory (default: update_with_mcmc_take_last()). See build_kernel() for the contract.

Return type:

A kernel kernel(rng_key, state) that returns (new_state, info).

coordinate_constrained_step(init_state_fn: Callable, slice_kernel: Callable, proposal: Callable = coordinate_proposal, coordinate_order: Callable = random_order) → Callable[source]#

The coordinate-sweep constrained inner step for nested sampling (SwiG).

The slice-within-Gibbs counterpart of slice_constrained_step(): one call sweeps every axis once – in the order set by coordinate_order – updating each by a univariate slice from the per-axis proposal generator proposal(init_state_fn, loglikelihood_0, i, width) (coordinate_proposal() by default, the axis analogue of covariance_proposal() passed to slice_constrained_step()), which steps along width * e_i and gates the likelihood contour into is_valid. As with the hit-and-run path the scale lives in the direction, so the univariate slice runs at unit width. Consumed by build_kernel() exactly like the hit-and-run step.

build_swig_kernel(init_state_fn: Callable, num_inner_steps: int, num_delete: int = 1, max_steps: int = 10, max_shrinkage: int = 100, proposal: Callable = coordinate_proposal, coordinate_order: Callable = random_order, inner_kernel_params: Callable = live_widths, update_strategy: Callable = update_with_mcmc_take_last) → Callable[source]#

Build the Nested Slice-within-Gibbs (SwiG) kernel.

The coordinate counterpart of build_kernel(): each inner step is a full coordinate sweep rather than a hit-and-run move. Axes are visited in the order set by coordinate_order (random_order() by default, or fixed_order()), and each is updated by a univariate slice gated on the likelihood contour and scaled by that axis’s live width (the per-axis spread of the live points; correlations are ignored). Prefer this when the target is close to axis-aligned, or when its correlations are unreliable to estimate. Pair with swig_as_top_level_api() for the bundled (init, step) algorithm.

Parameters:

  • init_state_fn – Builds a particle state from a position and birth log-likelihood.

  • num_inner_steps – Number of coordinate sweeps per new particle. Prefer num_inner_steps >= max(5, 2 * dim) for reliable mixing (bare dim is the minimum; see swig_as_top_level_api()).

  • num_delete – Number of particles deleted and replaced per step (default 1).

  • max_steps – Cap on stepping-out expansions per univariate slice (default 10).

  • max_shrinkage – Cap on shrinkage evaluations per univariate slice (default 100).

  • proposal – Per-axis proposal factory (init_state_fn, loglikelihood_0, i, width) -> proposal_generator (coordinate_proposal() by default). The coordinate analogue of the proposal seam on build_kernel().

  • coordinate_order – Sweep-order primitive (rng_key, d) -> indices (random_order() by default).

  • inner_kernel_params – Computes the inner-kernel parameters from the live points each step, (rng_key, state, info, params) -> params (live_widths() by default, the per-axis live-point spread).

  • update_strategy – Inner-kernel factory (default: update_with_mcmc_take_last()). See build_kernel() for the contract.

Return type:

A kernel kernel(rng_key, state) that returns (new_state, info).

as_top_level_api(logprior_fn: Callable, loglikelihood_fn: Callable, num_inner_steps: int, num_delete: int = 1, max_steps: int = 10, max_shrinkage: int = 100, proposal: Callable = covariance_proposal, inner_kernel_params: Callable | None = None, update_strategy: Callable = update_with_mcmc_take_last) → blackjax.SamplingAlgorithm[source]#

Creates a Nested Slice Sampling (NSS) algorithm, blackjax.nss.

Nested Sampling with a hit-and-run slice inner kernel: each particle replacement runs num_inner_steps constrained slice moves along directions shaped by the live-point covariance.

Parameters:

  • logprior_fn – Log-prior of a single particle.

  • loglikelihood_fn – Log-likelihood of a single particle.

  • num_inner_steps – Number of slice steps per new particle. Use num_inner_steps >= max(5, 2 * dim) for reliable mixing within the likelihood contour; bare dim is the minimum and can bias the evidence upward for dim > 10 (the inner chain must decorrelate the new particle from the deleted one, not merely satisfy the constraint).

  • num_delete – Number of particles deleted and replaced per step (default 1).

  • max_steps – Cap on stepping-out expansions per slice (default 10).

  • max_shrinkage – Cap on shrinkage evaluations per slice (default 100).

  • proposal – Proposal factory (init_state_fn, loglikelihood_0, **params) -> proposal_generator (covariance_proposal() by default). The default proposal consumes a precomputed covariance_factor. Override to write a custom nested stepper.

  • inner_kernel_params – Computes the inner-kernel parameters from the live points, (rng_key, state, info, params) -> params. When None, uses live_covariance_factor() with the default proposal and live_covariance() with a custom proposal. Used both to seed init and to update each step.

  • update_strategy – Inner-kernel factory (default: update_with_mcmc_take_last()). See build_kernel() for the contract.

Returns:

  • A SamplingAlgorithm whose step(rng_key, state) returns

  • (new_state, info).

Notes

The live particles in the run state (state.particles) are not posterior samples: they are the current likelihood shell and at termination collapse to the highest-likelihood mode. For correctly-weighted posterior draws, pass the dead points through finalise() and resample with sample().

The covariance-shaped proposal bridges between modes only up to moderate separation. For strongly multimodal targets, ensure the initial live points span every mode (and consider a clustering inner kernel); minor modes are still weighted correctly in the resampled posterior, but may be absent from the final live set.

swig_as_top_level_api(logprior_fn: Callable, loglikelihood_fn: Callable, num_inner_steps: int, num_delete: int = 1, max_steps: int = 10, max_shrinkage: int = 100, proposal: Callable = coordinate_proposal, coordinate_order: Callable = random_order, inner_kernel_params: Callable = live_widths, update_strategy: Callable = update_with_mcmc_take_last) → blackjax.SamplingAlgorithm[source]#

Creates a Nested Slice-within-Gibbs (SwiG) sampling algorithm, blackjax.nsswig.

Nested Sampling with an axis-aligned slice-within-Gibbs inner kernel: each particle replacement runs num_inner_steps constrained coordinate sweeps, each axis scaled by the live-point spread (correlations are ignored). The coordinate counterpart of as_top_level_api(); prefer it when the target is close to axis-aligned or its correlations are unreliable to estimate.

Parameters:

  • logprior_fn – Log-prior of a single particle.

  • loglikelihood_fn – Log-likelihood of a single particle.

  • num_inner_steps – Number of coordinate sweeps per new particle. Use num_inner_steps >= max(5, 2 * dim) for reliable mixing within the likelihood contour; bare dim is the minimum and can bias the evidence upward for dim > 10 (the inner chain must decorrelate the new particle from the deleted one, not merely satisfy the constraint).

  • num_delete – Number of particles deleted and replaced per step (default 1).

  • max_steps – Cap on stepping-out expansions per univariate slice (default 10).

  • max_shrinkage – Cap on shrinkage evaluations per univariate slice (default 100).

  • proposal – Per-axis proposal factory (init_state_fn, loglikelihood_0, i, width) -> proposal_generator (coordinate_proposal() by default), the coordinate analogue of the proposal seam on as_top_level_api().

  • coordinate_order – Sweep-order primitive (rng_key, d) -> indices (random_order() by default).

  • inner_kernel_params – Computes the inner-kernel parameters from the live points, (rng_key, state, info, params) -> params (live_widths() by default). Used both to seed init and to update each step.

  • update_strategy – Inner-kernel factory (default: update_with_mcmc_take_last()). See build_kernel() for the contract.

Returns:

  • A SamplingAlgorithm whose step(rng_key, state) returns

  • (new_state, info).

Notes

The live particles in the run state (state.particles) are not posterior samples: they are the current likelihood shell and at termination collapse to the highest-likelihood mode. For correctly-weighted posterior draws, pass the dead points through finalise() and resample with sample().

For strongly multimodal targets, ensure the initial live points span every mode (the axis-aligned per-particle proposal does not bridge well-separated modes); minor modes are still weighted correctly in the resampled posterior, but may be absent from the final live set.