JNesty Documentation

GPU-accelerated nested sampling with JAX, designed for expensive likelihood functions. JNesty parallelizes likelihood evaluations across GPU cores using JAX’s vmap, making it well-suited for problems where each likelihood call is computationally costly (forward models, simulations, large-data inferences).

The API follows dynesty, so switching is straightforward.

• Methods
  • Nested Sampling Overview
  • Two-Phase Sampling
  • Random Walk Sampler
  • Bounding Methods
  • GPU Parallelism
  • Convergence
  • Evidence and Weight Computation
  • Dynamic Nested Sampling
• Examples
  • 1. Multi-Modal Gaussian Mixture (5D)
  • 2. Rosenbrock Banana (2D)
  • 3. High-Dimensional Gaussian (20D)
  • 4. Gaussian Shells (2D)
  • 5. Dynamic Nested Sampling (3D Correlated Gaussian)
• API Reference
  • NestedSampler
  • DynamicNestedSampler
  • Results
  • Plotting
  • Core Sampling Loop
  • Bounding Methods
  • Internal Samplers
  • Utilities
  • Multi-Ellipsoid Fitting

Installation

pip install -e .

This installs JNesty with JAX CUDA 12 GPU support. You need a CUDA-compatible GPU and driver. For CPU-only use, set device='cpu' when constructing the sampler.

Optional dependencies for dynesty comparison and plotting interop:

pip install -e ".[dynesty]"

Quick Start

import jax.numpy as jnp
from jnesty import NestedSampler, plotting

def loglikelihood(x):
    return -0.5 * jnp.sum(x**2)

def prior_transform(u):
    return 20.0 * u - 10.0  # uniform [-10, 10] in each dim

sampler = NestedSampler(loglikelihood, prior_transform, ndim=5)
sampler.run_nested(max_iterations=50000, delta_logZ_threshold=0.01)

r = sampler.results
print(f"logZ = {r['logz']:.4f} +/- {r['logzerr']:.4f}")
print(f"Converged: {r['converged']}")

# Built-in plots
fig, axes = plotting.runplot(r)
fig, axes = plotting.traceplot(r)
fig, axes = plotting.cornerplot(r)

Next Steps

• Methods — How the sampling and bounding algorithms work

• Examples — Five worked examples (multi-modal, Rosenbrock, high-D, shells, dynamic)

• API Reference — Full API reference