Nessai Interface

This module provides the bridge interface connecting Discovery models to the Nessai nested sampler. Nessai uses normalizing flows for efficient proposal generation and supports importance nested sampling.

Discovery ↔︎ nessai Interface

This module provides a light wrapper that adapts a Discovery-style model (a callable that returns a log-probability/log-likelihood given a parameter dictionary) to the API expected by nessai’s FlowSampler.

Design goals

• Keep input style similar to the (intended) Eryn interface: pass a discovery_model and a prior dictionary.

• Minimise hard-coding so this can be refactored later into a common abstract SamplerInterface.

• Clean separation between: parsing priors, converting (dict ↔︎ structured array), and running the sampler.

Notes on nessai model API

nessai expects a subclass of nessai.model.Model that defines: - names: list[str] – parameter names - bounds: dict[str, tuple[float, float]] – hard bounds used for sampling - log_prior(x) -> float – log prior for a single live point x - log_likelihood(x) -> float – log likelihood for a single live point

x is a structured numpy array with fields equal to names.

This wrapper constructs such a model on the fly, based on a prior-specification dictionary.

Prior specification

The prior dictionary maps parameter names to specs. Supported forms:

1. Distribution dicts: {

   ‘dist’: ‘uniform’ | ‘loguniform’ | ‘normal’ | ‘fixed’, # parameters depend on the dist (see _make_prior)

   }

2. Shorthand tuples for common cases: - (‘uniform’, min, max) - (‘loguniform’, a, b) - (‘normal’, mean, sigma) - (‘fixed’, value)

3. A callable prior: any object with logpdf(value) and, for non-fixed parameters, hard bounds provided via ‘bounds’: (min, max).

Fixed parameters are separated out and always injected into the model inputs before calling the Discovery model.

Example

>>> bridge = DiscoveryNessaiBridge(
...     discovery_model=my_model,  # callable or object with log_prob/log_likelihood
...     prior={
...         'm1': {'dist': 'uniform', 'min': 5, 'max': 50},
...         'm2': ('loguniform', 1e-1, 10.0),
...         'z':  ('fixed', 0.2),
...     },
... )
>>> # Run the sampler
>>> results = bridge.run_sampler(nlive=1000, max_iterations=50_000, output='./out')

class discoverysamplers.nessai_interface.DiscoveryNessaiBridge(discovery_model, priors, latex_labels=None, jit=True)

Bases: object

Bridge between a Discovery-style model and nessai’s FlowSampler.

Parameters:

• discovery_model (callable | object) – A callable or object with log_prob or log_likelihood.

• priors (Mapping[str, PriorSpec]) – Dictionary describing priors (see module docstring). Includes fixed parameters.

• labels (Optional[Mapping[str, str]]) – Optional display labels per parameter (not used internally yet).

sampled_names

Type:

list[str]

fixed_params

Type:

dict[str, float]

bounds

Type:

dict[str, tuple[float, float]]

model

Type:

DiscoveryNessaiModel

__init__(discovery_model, priors, latex_labels=None, jit=True)

dict_to_livepoint(d)

Convert a parameter dict to a nessai live point (structured array).

Return type:

ndarray

livepoint_to_dict(x)

Return type:

Dict[str, float]

run_sampler(*, nlive=1000, output='./nessai_out', resume=False, **kwargs)

Run nessai.FlowSampler with this model.

Parameters:

• nlive (int) – Number of live points.

• output (str) – Output directory for nessai (checkpoints, samples, plots).

• max_iterations (Optional[int]) – Maximum number of iterations. If None, uses nessai default.

• seed (Optional[int]) – Random seed passed to the sampler.

• resume (bool) – Resume from previous run if possible.

• **kwargs (Any) – Forwarded directly to FlowSampler.

Return type:

Any

Returns:

• The object returned by FlowSampler.run(), typically a results

• dictionary including posterior samples and evidences.

return_logZ(*, results=None)

Return the log evidence and its uncertainty from nested sampling.

Parameters:

results (dict, optional) – Results dict from run_sampler(). If None, uses stored results.

Returns:

Dictionary containing: - ‘logZ’: float - the log evidence estimate - ‘logZ_err’: float - uncertainty on logZ

Return type:

dict

Raises:

RuntimeError – If no results are available (run_sampler not called)

return_sampled_samples(*, results=None)

Returns sampled parameters only.

Returns:

• ‘names’ : list[str]

• ’labels’: list[str] (LaTeX or names)

• ’chain’ : ndarray (nsamples, n_sampled)

Return type:

dict with keys

return_all_samples(*, results=None)

Returns sampled + fixed parameters, arranged in the original priors order. Fixed parameters are filled with their constant values.

Return type:

Dict[str, Any]

plot_trace(*, burn=0, plot_fixed=False, results=None, **kwargs)

Plot trace of samples vs sample index.

Parameters:

• burn (int, optional) – Number of initial samples to discard, by default 0.

• plot_fixed (bool, optional) – If True, includes fixed parameters in the plot, by default False.

• results (dict, optional) – Results dict from run_sampler(). If None, uses stored results.

• **kwargs – Additional keyword arguments passed to plots.plot_trace().

Returns:

Figure containing the trace plots.

Return type:

matplotlib.figure.Figure

plot_corner(*, burn=0, results=None, **kwargs)

Corner plot of sampled parameters.

Parameters:

• burn (int, optional) – Number of initial samples to discard, by default 0.

• results (dict, optional) – Results dict from run_sampler(). If None, uses stored results.

• **kwargs – Additional keyword arguments passed to corner.corner().

Returns:

Corner plot figure.

Return type:

matplotlib.figure.Figure

class discoverysamplers.nessai_interface.DiscoveryNessaiModel(names, bounds, logprior_fns, fixed_params, discovery_adapter)

Bases: Model

A nessai Model that wraps a Discovery-style model.

Parameters:

• names (list[str]) – Names of sampled parameters (fixed parameters are injected internally).

• bounds (dict[str, tuple[float, float]]) – Sampling bounds for each sampled parameter.

• logprior_fns (dict[str, Callable[[float], float]]) – Per-parameter log-prior functions.

• fixed_params (dict[str, float]) – Parameters that are not sampled.

• discovery_adapter (LikelihoodWrapper) – Adapter to call the Discovery model with a parameter dict.

__init__(names, bounds, logprior_fns, fixed_params, discovery_adapter)

log_prior(x)

Return type:

ndarray

log_likelihood(x)

Return type:

ndarray

See Also

• Nessai Nested Sampler - User guide for Nessai

• Performance Optimization - Performance optimization