gensbi.inference.posterior

NLE posterior target: build log-densities from a trained likelihood flow + prior.

The flow is NLE-trained (obs = x, cond = theta), so flow.log_prob(x, theta) is log q(x | theta). NLEPosterior turns (flow, prior, x_o) into a PosteriorTarget (separate log-prior / log-likelihood / log-posterior closures), which a Sampler consumes. The flow params are frozen constants inside the closures; only theta is traced/differentiated.

Classes

NLEPosterior	Amortized NLE posterior over a trained likelihood flow.
PosteriorTarget	Log-densities for one observation x_o.

Module Contents

class gensbi.inference.posterior.NLEPosterior(flow, prior, *, structured_obs=False)

Amortized NLE posterior over a trained likelihood flow.

Parameters:

• flow (object) – Exposes log_prob(x, cond) -> (B,) (an NLE-trained MAFlow/TarFlow).

• prior (numpyro.distributions.Distribution) – Prior over theta; prior.log_prob(theta) is a scalar and prior.sample(key, ()) returns (dim,).

• structured_obs (bool, optional) – If True, x_o keeps its (image/field) shape instead of being flattened. Default is False.

build_target(x_o)

Build a posterior target for a single observation.

Parameters:

x_o (Array) – Observed data. For non-structured: squeezed then promoted to (dim_x, 1) (channel-carrying). For structured_obs=True: kept as-is (image/field shape).

Returns:

Frozen log-density container for x_o.

Return type:

PosteriorTarget

sample(key, x_o, sampler=None, *, return_info=False)

Draw posterior samples for a single observation.

Parameters:

• key (jax.random.PRNGKey) – Random key.

• x_o (Array) – Observed data passed to build_target().

• sampler (Sampler or None, optional) – Sampler instance to use. If None, defaults to MCLMC. Default is None.

• return_info (bool, optional) – If True, return a (samples, info) tuple instead of just samples. Default is False.

Returns:

• samples (Array) – Posterior samples of shape (n, dim, 1). When return_info=False (the default), this is the only return value.

• info (object) – Sampler-specific info object (MclmcInfo, SmcInfo, or NestedSamplerInfo). Only present when return_info=True.

flow

prior

structured_obs = False

class gensbi.inference.posterior.PosteriorTarget

Log-densities for one observation x_o.

Frozen dataclass produced by NLEPosterior.build_target(). All callables accept a flat parameter vector theta of shape (dim,).

Parameters:

• log_prior (Callable) – Log-prior density. Signature: log_prior(theta) -> float.

• log_likelihood (Callable) – Log-likelihood log q(x_o | theta) from the NLE-trained flow. Signature: log_likelihood(theta) -> float.

• log_posterior (Callable) – Unnormalised log-posterior log_likelihood(theta) + log_prior(theta). Signature: log_posterior(theta) -> float.

• prior (object) – Prior distribution; must expose sample(key, shape) and log_prob(theta).

• dim (int) – Dimensionality of the parameter space.

dim: int

log_likelihood: Callable

log_posterior: Callable

log_prior: Callable

prior: object