[Draft] Nested sampling implementation

blackjax/__init__.py

@@ -37,6 +37,7 @@
 from .smc import inner_kernel_tuning as _inner_kernel_tuning
 from .smc import partial_posteriors_path as _partial_posteriors_smc
 from .smc import tempered
+from .ns import adaptive
 from .vi import meanfield_vi as _meanfield_vi
 from .vi import pathfinder as _pathfinder
 from .vi import schrodinger_follmer as _schrodinger_follmer
@@ -122,6 +123,8 @@ def generate_top_level_api_from(module):
 inner_kernel_tuning = generate_top_level_api_from(_inner_kernel_tuning)
 partial_posteriors_smc = generate_top_level_api_from(_partial_posteriors_smc)
 
+adaptive_ns = generate_top_level_api_from(adaptive)
+
 smc_family = [tempered_smc, adaptive_tempered_smc, partial_posteriors_smc]
 "Step_fn returning state has a .particles attribute"
 

blackjax/ns/__init__.py

@@ -0,0 +1,5 @@
+from . import adaptive
+
+__all__ = [
+    "adaptive",
+]

blackjax/ns/adaptive.py

@@ -0,0 +1,195 @@
+from functools import partial
+from typing import Callable, Dict
+
+import jax.numpy as jnp
+
+from blackjax import SamplingAlgorithm
+from blackjax.ns.base import NSInfo, NSState
+from blackjax.ns.base import build_kernel as base_ns
+from blackjax.ns.base import delete_fn
+from blackjax.ns.base import init as init_base
+from blackjax.ns.vectorized_slice import build_kernel as slice_kernel
+from blackjax.ns.vectorized_slice import init as slice_init
+from blackjax.smc.inner_kernel_tuning import StateWithParameterOverride
+from blackjax.smc.tuning.from_particles import particles_covariance_matrix
+from blackjax.types import ArrayLikeTree, ArrayTree, PRNGKey
+
+__all__ = ["init", "as_top_level_api", "build_kernel", "nss"]
+
+
+def init(position, loglikelihood_fn, parameter_update_function):
+    state = init_base(position, loglikelihood_fn)
+    initial_parameter_value = parameter_update_function(
+        state, NSInfo(state, state, state, None)
+    )
+    return StateWithParameterOverride(state, initial_parameter_value)
+
+
+def build_kernel(
+    logprior_fn: Callable,
+    loglikelihood_fn: Callable,
+    delete_fn: Callable,
+    mcmc_step_fn: Callable,
+    mcmc_init_fn: Callable,
+    mcmc_parameter_update_fn: Callable,
+    num_mcmc_steps: int,
+) -> Callable:
+    r"""Build an adaptive Nested Sampling kernel. Tunes the inner kernel parameters
+    at each iteration.
+
+    Parameters
+    ----------
+    logprior_fn : Callable
+        A function that computes the log prior probability.
+    loglikelihood_fn : Callable
+        A function that computes the log likelihood.
+    delete_fn : Callable
+        Function that takes an array of log likelihoods and marks particles for deletion and updates.
+    mcmc_step_fn:
+        The initialisation of the transition kernel, should take as parameters.
+        kernel = mcmc_step_fn(logprior, loglikelihood, logL0 (likelihood threshold), **mcmc_parameter_update_fn())
+    mcmc_init_fn
+        A callable that initializes the inner kernel
+    mcmc_parameter_update_fn : Callable[[NSState, NSInfo], Dict[str, ArrayTree]]
+        Function that updates the parameters of the inner kernel.
+    num_mcmc_steps: int
+        Number of MCMC steps to perform. Recommended is 5 times the dimension of the parameter space.
+
+    Returns
+    -------
+    Callable
+        A function that takes a rng_key and a NSState that contains the current state
+        of the chain and returns a new state of the chain along with
+        information about the transition.
+    """
+
+    def kernel(
+        rng_key: PRNGKey,
+        state: StateWithParameterOverride,
+    ) -> tuple[StateWithParameterOverride, NSInfo]:
+        step_fn = base_ns(
+            logprior_fn,
+            loglikelihood_fn,
+            delete_fn,
+            mcmc_step_fn,
+            mcmc_init_fn,
+            num_mcmc_steps,
+        )
+        new_state, info = step_fn(
+            rng_key, state.sampler_state, state.parameter_override
+        )
+        new_parameter_override = mcmc_parameter_update_fn(new_state, info)
+        return (
+            StateWithParameterOverride(new_state, new_parameter_override),
+            info,
+        )
+
+    return kernel
+
+
+def as_top_level_api(
+    logprior_fn: Callable,
+    loglikelihood_fn: Callable,
+    mcmc_step_fn: Callable,
+    mcmc_init_fn: Callable,
+    mcmc_parameter_update_fn: Callable[[NSState, NSInfo], Dict[str, ArrayTree]],
+    num_mcmc_steps: int,
+    n_delete: int = 1,
+) -> SamplingAlgorithm:
+    """Implements the (basic) user interface for the Adaptive Nested Sampling kernel.
+
+    Parameters
+    ----------
+    logprior_fn : Callable
+        A function that computes the log prior probability.
+    loglikelihood_fn : Callable
+        A function that computes the log likelihood.
+    mcmc_step_fn:
+        The initialisation of the transition kernel, should take as parameters.
+        kernel = mcmc_step_fn(logprior, loglikelihood, logL0 (likelihood threshold), **mcmc_parameter_update_fn())
+    mcmc_init_fn
+        A callable that initializes the inner kernel
+    mcmc_parameter_update_fn : Callable[[NSState, NSInfo], Dict[str, ArrayTree]]
+        A function that updates the parameters given the current state and info.
+    num_mcmc_steps: int
+        Number of MCMC steps to perform. Recommended is 5 times the dimension of the parameter space.
+    n_delete : int, optional
+        Number of particles to delete in each iteration. Default is 1.
+
+    Returns
+    -------
+    SamplingAlgorithm
+        A sampling algorithm object.
+    """
+    delete_func = partial(delete_fn, n_delete=n_delete)
+
+    kernel = build_kernel(
+        logprior_fn,
+        loglikelihood_fn,
+        delete_func,
+        mcmc_step_fn,
+        mcmc_init_fn,
+        mcmc_parameter_update_fn,
+        num_mcmc_steps,
+    )
+
+    def init_fn(position: ArrayLikeTree, rng_key=None):
+        del rng_key
+        return init(position, loglikelihood_fn, mcmc_parameter_update_fn)
+
+    def step_fn(rng_key: PRNGKey, state):
+        return kernel(rng_key, state)
+
+    return SamplingAlgorithm(init_fn, step_fn)
+
+
+def nss(
+    logprior_fn: Callable,
+    loglikelihood_fn: Callable,
+    num_mcmc_steps: int,
+    n_delete: int = 1,
+) -> SamplingAlgorithm:
+    """Implements the a baseline Nested Slice Sampling kernel.
+
+    Parameters
+    ----------
+    logprior_fn: Callable
+        A function that computes the log prior probability.
+    loglikelihood_fn: Callable
+        A function that computes the log likelihood.
+    num_mcmc_steps: int
+        Number of MCMC steps to perform. Recommended is 5 times the dimension of the parameter space.
+    n_delete: int, optional
+        Number of particles to delete in each iteration. Default is 1.
+
+    Returns
+    -------
+    SamplingAlgorithm
+        A sampling algorithm object.
+    """
+    delete_func = partial(delete_fn, n_delete=n_delete)
+    mcmc_step_fn = slice_kernel
+    mcmc_init_fn = slice_init
+
+    def parameter_update_fn(state, _):
+        cov = jnp.atleast_2d(particles_covariance_matrix(state.particles))
+        return {"cov": cov}
+
+    kernel = build_kernel(
+        logprior_fn,
+        loglikelihood_fn,
+        delete_func,
+        mcmc_step_fn,
+        mcmc_init_fn,
+        parameter_update_fn,
+        num_mcmc_steps,
+    )
+
+    def init_fn(position: ArrayLikeTree, rng_key=None):
+        del rng_key
+        return init(position, loglikelihood_fn, parameter_update_fn)
+
+    def step_fn(rng_key: PRNGKey, state):
+        return kernel(rng_key, state)
+
+    return SamplingAlgorithm(init_fn, step_fn)