Posterior sampling

ds.inference.sample draws from the posterior of a MAP fit with MCMC, and ds.inference.mcmc_report checks the draws. See posterior sampling.

diff_soup.inference.sample

sample(problem: Any, fit: Any = None, *, kernel: str = 'nuts', chains: int = 4, draws: int = 1000, warmup: int = 200, adapt: Optional[str] = 'pooled', mass: str = 'diag', step_size: Optional[float] = None, max_tree_depth: int = 10, key: Array, preconditioner: Any = 'laplace', curvature: str = 'observed', chain_map: Optional[str] = None, **kernel_options) -> Draws

Draw from the posterior of a MAP problem with MCMC, with the chains started near the fit.

Each chain starts from the fit plus random noise of half a standard deviation per coordinate, and every draw is in the parameters' own coordinates. It is refused in a run of several processes.

Parameters:

Name	Type	Description	Default
problem	Any	An InferenceProblem with a MAP objective, or a target from posterior_target, which takes no fit.	required
fit	Any	The MAP fit the chains start from.	None
kernel	str	"nuts" for the No-U-Turn sampler, which first audits the gradient at a perturbed chain start, or "rmh" for random-walk Metropolis, which needs no gradient.	'nuts'
chains	int	The number of chains.	4
draws	int	The number of draws kept per chain.	1000
warmup	int	The number of steps before the kept draws, which random-walk chains discard and NUTS chains spend adapting, keeping their last state.	200
adapt	Optional[str]	How NUTS adapts, "pooled" from all chains at once, "window" each chain alone, "step_size" the step size alone, or None not at all.	'pooled'
mass	str	The mass matrix NUTS adapts, "diag" or "dense", which needs adapt="window" and at most 50 parameters.	'diag'
step_size	Optional[float]	The NUTS step size adaptation starts from, 1.0 when None, and the fixed one with adapt=None.	None
max_tree_depth	int	The most times a NUTS trajectory doubles in one draw.	10
key	Array	The JAX random key that fixes the draws.	required
preconditioner	Any	"laplace" to scale the proposals by the Laplace approximation at the fit, or None to run without it.	'laplace'
curvature	str	The curvature of that Laplace approximation, "observed" or "expected".	'observed'
chain_map	Optional[str]	How the chains run, "vmap" as one batched call, "sequential" one after another, "devices" on the local devices, or None to choose per problem.	None
**kernel_options		nuts_target_acceptance, the acceptance rate NUTS adapts its step size toward, by default 0.8.	{}

Returns:

Type	Description
Draws	The Draws.

diff_soup.inference.Draws dataclass

The MCMC draws of sample, in the parameters' own coordinates.

draws[handle] returns the draws of one parameter, or a dictionary by member for a parameter that varies by member.

Attributes:

Name	Type	Description
values	Array	The draws, shaped (chains, draws, parameters).
names	tuple	The label of each parameter column of values.
info	dict	Sampler statistics per draw, acceptance_rate, and for NUTS is_divergent, num_integration_steps and the adaptation record.
chain_map	str	How the chains ran, "vmap", "sequential" or "devices".

values instance-attribute

values: Array

names instance-attribute

names: tuple

chain_map class-attribute instance-attribute

chain_map: str = 'vmap'

positions property

positions

The draws, the same array as values.

diff_soup.inference.mcmc_report

mcmc_report(draws: Any, *, rhat_max: float = 1.01, ess_floor: float = 400.0) -> McmcReport

Check MCMC draws for convergence and effective sample size, per parameter.

The diagnostics are the rank-normalized split R-hat and the bulk and tail effective sample sizes. Each chain is split in half, so one chain is enough, and every chain needs at least eight draws.

Parameters:

Name	Type	Description	Default
draws	Any	The Draws of ds.inference.sample, or an array shaped (chain, draw, dim).	required
rhat_max	float	The largest split R-hat that passes.	1.01
ess_floor	float	The smallest bulk or tail effective sample size that passes.	400.0

Returns:

Type	Description
McmcReport	An McmcReport with the diagnostics per parameter, the sampler's statistics when draws carries them, and the failed checks in flags.