Methods

This page describes how JNesty implements nested sampling. The algorithms closely follow dynesty; we document the key ideas and JNesty-specific details here.

Nested Sampling Overview

Nested sampling (Skilling 2006) is a Monte Carlo method for computing the Bayesian evidence (marginal likelihood)

\[Z = \int \mathcal{L}(\theta)\,\pi(\theta)\,d\theta\]

by reparameterizing the integral in terms of the prior mass $X$. At each iteration the live point with the lowest likelihood is removed (“dead point”), the evidence contribution $\Delta Z = \mathcal{L}_i\,\Delta X_i$ is accumulated, and a new live point is drawn above the current likelihood threshold. The process terminates when the remaining evidence estimate $\Delta\log Z$ falls below a user-specified threshold.

Two-Phase Sampling

JNesty uses a two-phase approach matching dynesty:

Phase 1 — Uniform rejection sampling Draws batches of uniform random points from the unit cube and accepts the first one above the current likelihood threshold. This is fast when the prior volume is large relative to the posterior (early iterations). Phase 1 runs until sampling efficiency drops below min_eff (default 10%) after at least min_ncall (default $2 \times \text{nlive}$) likelihood calls.

Phase 2 — Random walk with adaptive bounds Uses the RWalkSampler with a chosen bounding method. Runs until convergence ($\Delta\log Z < \text{threshold}$) or the iteration limit is reached.

Random Walk Sampler

The random walk sampler (rwalk) follows dynesty’s generic_random_walk / propose_ball_point implementation:

Proposal: Generate a point uniformly inside an n-ball via $\hat{x} = z \cdot (U^{1/n}/\|z\|)$ where $z \sim \mathcal{N}(0, I)$ and $U \sim \text{Uniform}(0,1)$. Transform by the axes matrix from the current bound and add to the current position.
Acceptance: Accept if the proposal is inside the unit cube $[0,1]^{\text{ndim}}$ and has $\log\mathcal{L} > \log\mathcal{L}^*$ (the current worst live point). No Metropolis step — pure constraint enforcement matching dynesty.
Scale adaptation: After each walk, update the proposal scale via Robbins-Monro: $$\sigma \leftarrow \sigma \cdot \exp\!\left(\frac{f_{\text{acc}} - f_{\text{target}}}{n_{\text{cdim}} \cdot f_{\text{target}}}\right)$$ where $f_{\text{acc}}$ is the measured acceptance rate and $n_{\text{cdim}}$ is the number of clustered dimensions (default: ndim).
Clustered vs non-clustered dims: The first ncdim dimensions are perturbed by the axes transform. Remaining dimensions are resampled uniformly from $[0,1]$, matching dynesty’s ncdim split.
Retry loop: Each iteration retries the walk with fresh randomness until a valid replacement ($\log\mathcal{L} > \log\mathcal{L}^*$) is found, matching dynesty’s _new_point() while-True behavior.

Bounding Methods

JNesty implements three bounding methods, selected by the bound parameter:

UnitCube (`bound='none'`)

No bounding — proposals are generated within the full unit cube. The axes matrix is the identity. Simple and robust, but less efficient for concentrated or elongated posteriors.

SingleEllipsoid (`bound='single'`)

Fits a single ellipsoid to the live points using covariance eigenvalue decomposition. The axes matrix is $\mathbf{A} = \mathbf{V}\sqrt{\boldsymbol{\Lambda}}$ where $\mathbf{V}$ are eigenvectors and $\boldsymbol{\Lambda}$ are eigenvalues. Good for unimodal, elongated posteriors.

MultiEllipsoidBound (`bound='multi'`)

Fits multiple ellipsoids via recursive k-means splitting with a BIC stopping criterion, matching dynesty’s _bounding_ellipsoids:

Fit one bounding ellipsoid to all live points
Split into 2 clusters via k-means along the major axis
Accept the split if the combined volume decreases significantly (BIC)
Recurse on each sub-cluster (bounded depth: max_ellipsoids)
Enlarge all ellipsoids by a factor (default 1.25) for numerical safety

For rwalk proposals, one ellipsoid is selected proportional to volume and its axes are used to shape the n-ball proposal. This lets the walker focus on one mode at a time.

Periodic refitting is controlled by bound_update_interval (measured in likelihood calls, default: rwalk_K * nlive for multi-ellipsoid).

GPU Parallelism

JNesty offers two GPU parallelism modes for the random walk phase:

Batch mode (`queue_size=0`, default for `bound='none'`)

Runs batch_size independent walks in parallel via jax.vmap, each with rwalk_K // batch_size steps. Each walk starts from a different live point with a different (randomly selected) ellipsoid. The first valid candidate becomes the replacement. Batch size is auto-tuned to ~5 steps/walk and capped by GPU memory.

Queue mode (`queue_size>1`, default for `bound='multi'`)

Dynesty-style GPU parallelism: generates queue_size candidates in parallel, each with a full rwalk_K steps. The main loop tests one candidate per iteration against the current $\log\mathcal{L}^*$. Scale adapts only when the queue drains, matching dynesty’s multiprocessing behavior. Leftover candidates are preserved across iterations.

Convergence

The convergence criterion is:

\[\Delta\log Z < \text{threshold}\]

where $\Delta\log Z$ estimates the maximum remaining evidence contribution from the live points:

\[\Delta\log Z = \log\!\left(1 + e^{\log\mathcal{L}_{\max} + \log X_{\text{current}} - \log Z}\right)\]

The default threshold is 0.01. Use 0.1 for quick tests, 0.001 for publication-quality results.

Evidence and Weight Computation

After sampling completes, format_results() computes posterior weights:

Dead points: Prior volume shrinks as $\log X_i = -i / \text{nlive}$. Trapezoidal weights combine adjacent likelihoods: $\log w_i = \log\text{addexp}(\log\mathcal{L}_i, \log\mathcal{L}_{i-1}) + \log\Delta V_i$
Live points: Remaining live points are added using dynesty’s add_live_points() formula with accelerating volume shrinkage: $\log V_{\text{live},j} = \log(1 - (j+1)/(n_{\text{live}}+1)) + \log V_{\text{last dead}}$
Cumulative evidence: $\log Z_i = \text{logaddexp.accumulate}(\log w_i)$

The evidence error is estimated as $\sigma_{\log Z} = \sqrt{|H| / n_{\text{live}}}$ where $H$ is the information.

Dynamic Nested Sampling

JNesty implements dynamic nested sampling following Higson et al. (2019) and dynesty. The key idea is to run a baseline static NS survey, then iteratively add batches of live points targeted at the regions that contribute most to posterior or evidence uncertainty.

Algorithm

Base run: A standard static NS run with nlive_init live points produces the initial set of dead points.
Batch loop (repeated until stopping_function returns True or maxbatch is reached):
- Weight function: computes where additional samples would be most valuable, returning (logl_min, logl_max) bounds based on a weighted combination of posterior and evidence importance: $$w_i = p_{\text{frac}} \cdot p_i + (1 - p_{\text{frac}}) \cdot z_i$$
- Seed live points: volume-weighted subset selection from saved dead points above logl_min, then fresh point generation via rejection sampling to produce nlive_batch independent live points.
- Batch run: constrained static NS between (logl_min, logl_max) with pre-loaded saved history for correct convergence.
- Combine: merge batch dead points into the saved run, recompute logvol with combined live-point counts, and recalculate evidence via compute_integrals.
Stopping: the run stops when $$\text{stop} = p_{\text{frac}} \cdot \frac{n_{\text{target}}}{n_{\text{eff}}}
- (1 - p_{\text{frac}}) \cdot \frac{\sigma_{\log Z}}{\text{evid_thresh}} \leq 1$$

GPU Acceleration

The fresh point generation step uses a single-shot GPU batch: samples nlive_batch * 8 points from the unit cube, evaluates logL via a JIT-compiled vmap, and filters above logl_min — all in one JAX dispatch. The compiled logL evaluator is cached across batches.

API

from jnesty import DynamicNestedSampler

dsampler = DynamicNestedSampler(loglikelihood, prior_transform, ndim,
                                 nlive=500, bound='single')
dsampler.run_nested(nlive_init=500, nlive_batch=500, maxbatch=10,
                     pfrac=0.8, n_effective=10000)

Key parameters:

pfrac: posterior/evidence split (0.8 default; 1.0 = posterior-only; 0.0 = evidence-only)
nlive_batch: live points per batch (default = nlive)
maxbatch: maximum number of batches
n_effective: target effective sample size for stopping