pespace.inference.sampler.base

Base classes and shared utilities for sampler interfaces.

Functions

_build_distribution(spec)	Instantiate a numpyro distribution from a dict specification.
_validate_prior_dict(prior_dict, required_params)	Validate the prior declarations against the likelihood.
map_fn(fn[, batch_size])	Map over the leading sample axis.

Classes

Cosine([validate_args])	Prior density p(x) = cos(x) / 2 on [-pi/2, pi/2] (latitude).
NumpyroPriorSamplerBase(likelihood, prior_dict)	Base class for samplers using priors specified with build-in distributions in numpyro.
SamplerBase(likelihood, prior_dict[, label, ...])	Framework-agnostic abstract base class for samplers.
Sine([validate_args])	Prior density p(x) = sin(x) / 2 on [0, pi] (inclination).

pespace.inference.sampler.base.map_fn(fn, batch_size=0)

Map over the leading sample axis.

Use vmap for batch_size=0 and sequential lax.map batches for positive sizes.

Parameters:

batch_size (int)

class pespace.inference.sampler.base.Sine(validate_args=None)

Bases: Distribution

Prior density p(x) = sin(x) / 2 on [0, pi] (inclination).

support = Interval(lower_bound=0.0, upper_bound=3.141592653589793)

Parameters:

x (ndarray | Array | number | int | float | complex)

Return type:

Array | ndarray | bool | number | bool | int | float | complex

arg_constraints: dict[str, Any] = {}

sample(key, sample_shape=())

log_prob(value)

class pespace.inference.sampler.base.Cosine(validate_args=None)

Bases: Distribution

Prior density p(x) = cos(x) / 2 on [-pi/2, pi/2] (latitude).

support = Interval(lower_bound=-1.5707963267948966, upper_bound=1.5707963267948966)

Parameters:

x (ndarray | Array | number | int | float | complex)

Return type:

Array | ndarray | bool | number | bool | int | float | complex

arg_constraints: dict[str, Any] = {}

sample(key, sample_shape=())

log_prob(value)

pespace.inference.sampler.base._build_distribution(spec)

Instantiate a numpyro distribution from a dict specification.

Parameters:

spec (dict)

Return type:

Distribution

pespace.inference.sampler.base._validate_prior_dict(prior_dict, required_params)

Validate the prior declarations against the likelihood.

Iterates the dictionary in insertion order and checks that

1. every 'from' input of a derived parameter is declared earlier (as a sampled, derived or fixed parameter);

2. no name is declared twice (a fixed parameter must not conflict with an earlier sampled/derived declaration, and vice versa);

3. every parameter required by the likelihood is declared.

Parameters:

• prior_dict (dict) – Prior specification: parameter names (or a tuple of names, for a transform returning several parameters at once) mapped to their declaration dicts.

• required_params (iterable of str) – Parameter names the likelihood needs, e.g. the fields of its likelihood_param_type namedtuple.

Raises:

ValueError – If a declaration is missing, declared twice, or depends on a parameter declared after it.

Return type:

None

class pespace.inference.sampler.base.SamplerBase(likelihood, prior_dict, label=None, outdir=None, rkey_seed=0)

Bases: ABC

Framework-agnostic abstract base class for samplers.

Parameters:

• prior_dict (dict)

• rkey_seed (int)

__init__(likelihood, prior_dict, label=None, outdir=None, rkey_seed=0)

Parameters:

• likelihood (object) – Likelihood instance providing a pure log_likelihood(params) -> scalar method.

• prior_dict (dict) – Input prior specification.

• label (str, optional) – Label attached to the sampling run, converted to str. The default None becomes the string "None".

• outdir (str, optional) – Output directory, created here and required not to exist yet, so that the result files of an earlier run can never be overwritten silently. The default None resolves to an outdir_{label} directory under the current working directory where the calling script runs (os.getcwd()).

• rkey_seed (int, optional) – Seed for the random number generator key. Default is 0.

Raises:

FileExistsError – If outdir already exists.

Return type:

None

logprior_fn(position)

Joint log-prior density (pure JAX, return a scalar).

Generic interface: concrete samplers should bind this name to the space-specific implementation (unbounded or physical) on the instance before sampling.

Return type:

Array

loglikelihood_fn(position)

Log-likelihood (pure JAX, return a scalar).

Generic interface: concrete samplers should bind this name to the space-specific implementation (unbounded or physical) on the instance before sampling.

Return type:

Array

logdensity_fn(position)

Log-posterior density up to a normalization constant (pure JAX, return a scalar).

Return type:

Array

format_output(posterior_samples, log_likelihoods, log_evidence, runing_time, weights=None)

Assemble the run result into a dictionary.

Parameters:

• posterior_samples (dict) – Categorized posterior samples with keys sampled_params, derived_params and fixed_params, each mapping parameter names to their physical parameter arrays.

• log_likelihoods (array-like) – Log-likelihood of each sample.

• log_evidence (float) – Log-evidence of the run (for SMC samplers).

• runing_time (float) – Wall-clock time of the sampling run, in seconds.

• weights (array-like, optional) – Importance weight of each sample (for SMC samplers). When None, an empty array is stored.

Returns:

Dictionary with keys label, outdir, rkey_seed, prior_dict, likelihood_class, runing_time, sampler_class, sampler_config, posteriors, log_likelihoods, log_evidence and weights.

Return type:

dict

_save_output(output, filename=None)

Save results to outdir/{label}_results.json and return the path.

numpy/JAX arrays and numpy scalars are converted to JSON-native types recursively before dumping. Anything else the JSON encoder cannot handle – e.g. the jax.Device objects stored under sampler_config['devices'] when the devices were passed explicitly – falls back to its string representation, so the result of a run is always written.

Parameters:

output (dict)

Return type:

str

abstractmethod run(*args, **kwargs)

Performing the sampling run.

Return type:

dict

_abc_impl = <_abc._abc_data object>

class pespace.inference.sampler.base.NumpyroPriorSamplerBase(likelihood, prior_dict, label=None, outdir=None, rkey_seed=0)

Bases: SamplerBase

Base class for samplers using priors specified with build-in distributions in numpyro.

The input prior dictionary is parsed into numpyro distributions which provide facilities to map parameters in limited ranges to the unbounded real domain for better supporting gradient-based (HMC/NUTS) kernels in samplers.

Notes

The prior dictionary maps parameter names to specifications. Three kinds of entries are supported:

(1) independent sampled priors (enter the sampling dimension) “param_name”: {“distribution”: “dist_name”, **dist_params}, e.g., “chirp_mass”: {“distribution”: “Uniform”, “low”: 5e5, “high”: 2e6}, “chi_1”: {“distribution”: “Normal”, “loc”: 0.0, “scale”: 0.5},

(2) derived parameters (computed from sampled/derived ones, not sampled; maps the sampled parameters to those required by the likelihood) “param_name”: {“transform”: <callable>, “from”: (“param_1”, …)}, (“param_name_1”, “param_name_2”): {“transform”: <callable>, “from”: (“param_1”, …)}, e.g. “symmetric_mass_ratio”: {“transform”: eta_from_mass1_mass2, “from”: (“mass_1”, “mass_2”)}, (“mass_1”, “mass_2”): {“transform”: component_mass_from_Mtotal_q, “from”: (“total_mass”, “mass_ratio”)},

(3) fixed parameters “param_name”: {“distribution”: “Delta”, “v”: 0.0}, “param_name”: {“fixed”: 1.0}, e.g., “coalescence_time”: {“distribution”: “Delta”, “v”: 0.0}, “coalescence_time”: {“fixed”: 0.0},

Dictionary insertion order defines the parameter order. The parameters obtained by transform must be specified after the original parameters from which they are derived.

Parameters specified in prior dictionary need to satisfy:

• All parameters required by likelihood computation are covered.

• Unused parameters that required by neither the likelihood nor transform will be discarded, and a warning will be thrown.

Parameters:

• prior_dict (dict)

• rkey_seed (int)

static _require_numpyro()

Raise if numpyro is unavailable.

Return type:

None

_parse_prior_dict()

Split the prior dictionary into sampled / derived / fixed entries.

Derived entries must provide a 'from' key naming their input parameters; a tuple key declares a transform returning several parameters at once.

Return type:

None

logprior_unbounded_fn(z)

Joint log-prior in the unbounded (unconstrained) space (pure JAX).

The change-of-variables log-abs-det-Jacobian is included in the prior term so the tempered target stays self-consistent: the Jacobian must NOT be rescaled with the tempering temperature.

Return type:

Array

_to_physical_params(z)

Restore constrained sampled parameters from unconstrained z.

_likelihood_params_from_unbounded(z)

Assemble the likelihood-input namedtuple from unbounded z.

_likelihood_params_from_physical(x)

Assemble the likelihood-input namedtuple from physical x.

x must provide the sampled parameters in their constrained (physical) domain; derived and fixed parameters are added here.

A transform declared under a tuple key returns several parameters at once (a tuple/list matching the key order); a single-name transform returns the parameter value directly.

loglikelihood_unbounded_fn(z)

Log-likelihood in the unbounded space (pure JAX).

Return type:

Array

logdensity_unbounded_fn(z)

Log-posterior density up to a constant in the unbounded space.

Return type:

Array

logprior_physical_fn(x)

Joint log-prior in the physical space (no Jacobian correction).

Return type:

Array

loglikelihood_physical_fn(x)

Log-likelihood for parameters given in the physical space.

Return type:

Array

logdensity_physical_fn(x)

Log-posterior density up to a constant in the physical space.

Return type:

Array

_sample_priors_unbounded(key, shape=())

Draw initial positions in the unconstrained space from a given key.

Pure-JAX helper (no mutable state), usable inside jax.pmap.

Parameters:

shape (tuple)

initial_position_unbounded(shape=())

Draw initial position(s) in the unbounded space.

Parameters:

shape (tuple)

_sample_priors_physical(key, shape=())

Draw initial positions in the physical space from a given key.

Pure-JAX helper (no mutable state), usable inside jax.pmap.

Parameters:

shape (tuple)

initial_position_physical(shape=())

Draw initial position(s) in the physical space.

Parameters:

shape (tuple)

_abc_impl = <_abc._abc_data object>