| name | /usage |
| description | JNesty usage guide — public API, parameters, results access, FITS I/O, bounding methods, tuning, and plotting. |
Quick Start
from jnesty import NestedSampler, save_results, load_results
import jax.numpy as jnp
def loglikelihood(x):
return -0.5 * jnp.sum(x**2)
def prior_transform(u):
return (u - 0.5) * 10.0
sampler = NestedSampler(loglikelihood, prior_transform, ndim=5)
sampler.run_nested(max_iterations=10000, delta_logZ_threshold=0.01)
results = sampler.results
print(f"logZ = {results['logz']:.4f} +/- {results['logzerr']:.4f}")
save_results(results, 'output.fits')
loaded = load_results('output.fits')
NestedSampler Constructor
NestedSampler(
loglikelihood,
prior_transform,
ndim,
nlive=500,
rwalk_K=None,
rwalk_step_scale=None,
target_acceptance=0.5,
scale_adapt_interval=1,
device='gpu',
verbose=True,
bound='none',
bound_update_interval=None,
max_ellipsoids=20,
batch_size=None,
queue_size=None,
memory_frac=0.9,
unit_cube_batch_size=200,
min_eff=10.0,
min_ncall=None,
)
Sampling Modes
JNesty has two sampling modes, selected automatically:
Legacy mode (queue_size=0, default for bound='none'):
- Runs
batch_size independent walks in parallel via jax.vmap
- Each walk gets
rwalk_K // batch_size steps
- Scale adapts every iteration
Queue mode (queue_size>1, default for bound='multi'):
- Dynesty-style GPU parallelism
- Generates
queue_size candidates in parallel, each with full rwalk_K steps
- Tests one candidate per iteration against current loglstar
- Scale adapts only when queue drains (matching Dynesty's multiprocessing)
- Leftover candidates preserved across iterations
Auto-tuning Defaults
| Parameter | Auto value | Logic |
|---|
rwalk_K | max(25, ndim+20) | Dynesty's ndim+20 formula |
rwalk_step_scale | 1.0 | Dynesty default |
batch_size | max(1, rwalk_K // max(1, rwalk_K*10//nlive)) | ~5 steps/walk |
queue_size | 8 for multi, 0 otherwise | Match Dynesty multiprocessing |
bound_update_interval | rwalk_K*nlive calls for multi, 0 otherwise | In likelihood calls |
run_nested() Parameters
sampler.run_nested(
max_iterations=100000,
delta_logZ_threshold=0.01,
print_progress=True,
)
Two-Phase Sampling
The sampler uses a two-phase approach:
-
Phase 1 (uniform rejection): Draws batches of uniform random points from the unit cube, picks the first valid replacement. Fast when efficiency is high. Switches to Phase 2 when efficiency drops below min_eff% after min_ncall likelihood calls.
-
Phase 2 (random walk): Uses RWalkSampler with adaptive bounds. Runs until delta_logZ < threshold or max_iterations reached.
Progress bar shows: logZ, dlogZ (with threshold), logl* (current worst logL), eff(%) (overall efficiency).
Results Access
The results property returns a Results object supporting dict-style and attribute access:
results = sampler.results
results['logz']
results['logzerr']
results['information']
results['samples']
results['samples_u']
results['logl']
results['logwt']
results['logvol']
results['logz_trajectory']
results['logzerr_trajectory']
results['delta_logZ_trajectory']
results['scale_trajectory']
results['nlive']
results['niter']
results['eff']
results['acceptance_rate']
results['converged']
results['delta_logz']
results['delta_logZ_threshold']
results['rwalk_K']
Results Methods
results.summary()
results.samples_equal()
FITS I/O
from jnesty import save_results, load_results
save_results(results, 'output.fits')
loaded = load_results('output.fits')
FITS structure:
- HDU 0 (PrimaryHDU): Header with scalar metadata (LOGZ, LOGZERR, H, NLIVE, NITER, etc.)
- HDU 1 (BinTableHDU): Per-sample arrays (LOGL, LOGWT, LOGVOL, SAMPLES, SAMPLES_U, trajectories)
Bounding Methods
| Method | String | Use Case | Auto-updates |
|---|
| Unit cube | 'none' | Simple unimodal posteriors (default) | No |
| Single ellipsoid | 'single' | Elongated posteriors | No |
| Multi-ellipsoid | 'multi' | Multi-modal or complex geometries | Yes (periodic refit) |
Multi-ellipsoid uses recursive k-means splitting with BIC selection. Periodic refitting is measured in likelihood calls (default: rwalk_K * nlive calls, approximately nlive iterations). Queue mode (queue_size=8) is auto-enabled for multi-ellipsoid.
Parameter Tuning
- nlive: 500 is a good default. Use 1000+ for high-dimensional or multi-modal problems.
- rwalk_K: Auto-tuned to
max(25, ndim+20). Higher values improve mixing at cost of speed.
- rwalk_step_scale: Start at 1.0. Adapted automatically via Robbins-Munro.
- target_acceptance: 0.5 is Dynesty's default. Lower values explore more aggressively.
- delta_logZ_threshold: 0.01 is standard. Use 0.1 for quick tests, 0.001 for publication.
- batch_size (legacy mode only): Auto-tuned to ~5 steps/walk. Runs
batch_size independent walks in parallel via jax.vmap. Set to 1 to disable parallelism. Larger values benefit expensive likelihoods on GPU.
- queue_size (queue mode): 8 by default for
bound='multi'. Each queue entry gets full rwalk_K steps. Scale adapts only at queue drain.
- memory_frac: Fraction of GPU memory available for batch walks (default 0.9). Caps auto-tuned
batch_size if it would exceed this fraction of GPU memory. Uses compile-time memory estimation via XLA's memory_analysis(). Ignored on CPU.
- bound_update_interval: In likelihood calls (not iterations). Default
rwalk_K * nlive for multi-ellipsoid. Set to 0 to disable periodic updates. Float values in (0,1) are interpreted as fraction of rwalk_K * nlive.
- min_eff: Efficiency threshold (default 10%) to switch from Phase 1 to Phase 2.
Plotting
Built-in plotting (no Dynesty dependency)
from jnesty import plotting
fig, axes = plotting.runplot(sampler.results, lnz_truth=-10.0)
fig, axes = plotting.traceplot(sampler.results, truths=true_params)
fig, axes = plotting.cornerplot(sampler.results, truths=true_params)
fig, axes = plotting.cornerpoints(sampler.results)
fig, axes = plotting.diagnostics(sampler.results)
Dynesty plotting (requires dynesty installed)
sampler.plot_run()
sampler.plot_trace()
sampler.plot_corner()
sampler.plot_diagnostics()
These convenience methods call to_dynesty_results() internally to convert results for Dynesty's plotting functions.
Convenience Methods
sampler.print_summary()
sampler.get_samples()
sampler.get_logz()
sampler.to_dynesty_results()
Typical Usage Patterns
Simple problem (unimodal, low-dim)
sampler = NestedSampler(loglike, prior_transform, ndim=5, bound='none')
sampler.run_nested(delta_logZ_threshold=0.01)
Multi-modal problem
sampler = NestedSampler(loglike, prior_transform, ndim=5,
bound='multi', nlive=1000)
sampler.run_nested()
High-dimensional problem (>50D)
sampler = NestedSampler(loglike, prior_transform, ndim=100,
nlive=500, bound='none')
sampler.run_nested(delta_logZ_threshold=0.1)