Python API#

Sampling functions#

walnutpie.walnuts_pyfunc(logp: Callable[[ndarray], tuple[float, ndarray]] | numba.core.ccallback.CFunc | tuple[CFUNCTYPE, Any], *, num_params: int | None = None, inits: ndarray | None = None, num_chains: int = 4, seed: int | None = None, id: int = 1, init_radius: float = 2.0, init_inv_metric: ndarray | None = None, save_inv_metric: bool = False, min_warmup_iter: int = 50, max_warmup_iter: int = 1000, min_sampling_iter: int = 50, max_sampling_iter: int = 1000, max_trajectory_doublings: int = 5, max_step_halvings: int = 5, min_micro_steps: int = 1, max_hamiltonian_error: float = 0.5, step_size_converge_tol: float = 0.1, mass_converge_tol: float = 1.0, rhat_converge_tol: float = 1.01, mass_init_count: float = 4.0, mass_additive_smoothing: float = 1e-05, max_macro_steps_target: float = 15.0, step_size_init: float = 1.0, step_accept_rate_target: float = 0.8, step_learning_rate: float = 0.05, step_gradient_decay: float = 0.8, step_sq_gradient_decay: float = 0.9, step_stabilization: float = 0.0001, step_learn_rate_decay: float = 0.5, save_warmup: bool = False, refresh: int = 0) → list[WalnutsOutputArray][source]#

Sample from the specified model coded in Python using the specified configuration.

Parameters:

logp (Union[ Callable[[np.ndarray], tuple[float, np.ndarray]], numba.core.ccallback.CFunc, tuple[ctypes.CFUNCTYPE, Any], ]) –
The target log density and gradient function. Acceptable formats include:
- A function that accepts an numpy.ndarray and returns a tuple of the log density and gradient.
- A numba.cfunc() of type (size_t, doubleptr, doubleptr, doubleptr, voidptr) which accepts the position as the second argument and sets the gradient in argument 3 and log density in argument 4.
- A tuple of a ctypes callback of the same signature and an arbitrary data argument passed through the final pointer.
num_params (Optional[int], optional) – The dimensionality of the problem, by default None. At least one of num_params or inits must be specified.
inits (Optional[np.ndarray], optional) – The constrained initialization to use for all chains, or a list of constrained initializations, one for each chain, or None to indicate fully random initialization, by default None. At least one of num_params or inits must be specified.
num_chains (int, optional) – The number of Markov chains to run, positive, by default 4.
seed (Optional[int], optional) – The pseudo-random number generator seed, non-negative, or None to use a random seed, by default None.
id (int, optional) – Numeric id for the first chain, by default 1. The remaining chains are given consecutive ids following this one. This controls the random number generation, along with the seed.
init_radius (float, optional) – The bounds of uniform random initialization (-init_radius, init_radius), positive, by default 2.0.
init_inv_metric (Optional[np.ndarray], optional) – The diagonal of the initial diagonal inverse metric, positive entries and size equal to transformed (unconstrained) dimension, or None, in which case the mass matrix is initialized with a smoothed. negative outer product of gradients at the initial position, by default None.
save_inv_metric (bool, optional) – Set to True to save the inverse metric after adaptation, by default False.
min_warmup_iter (int, optional) – The minimum number of warmup iterations, greater than or equal to 0, by default 50.
max_warmup_iter (int, optional) – The maximum number of warmup iterations, greater than or equal to min_warmup_iter, by default 1000.
min_sampling_iter (int, optional) – The minimum number of sampling iterations, greater than or equal to 0, by default 50.
max_sampling_iter (int, optional) – The maximum number of sampling iterations, greater than or equal to min_sampling_iter, by default 1000.
max_trajectory_doublings (int, optional) – The maximum number of trajectory doublings for the no-U-turn sampler, positive, by default 5.
max_step_halvings (int, optional) – The maximum number of step size halvings in Walnuts, non-negative, by default 5.
min_micro_steps (int, optional) – The minimum number of micro steps per macro step, positive, by default 1.
max_hamiltonian_error (float, optional) – The maximum error allowed in the Hamiltonian, positive, by default 0.5.
step_size_converge_tol (float, optional) – The relative converge tolerance for difference in step sizes from the geometric mean across chains, positive, by default 0.1.
mass_converge_tol (float, optional) – The relative mass matrix norm convergence tolerance from the geometric mean across chains, by default 1.0.
rhat_converge_tol (float, optional) – The convergence tolernace for R-hat, greater than 1, by default 1.01.
mass_init_count (float, optional) – The pseudo-observation count for the initial mass matrix, positive, by default 4.0.
mass_additive_smoothing (float, optional) – The amount to add to the mass matrix estimators for smoothing, non-negative, by default 1e-5.
max_macro_steps_target (float, optional) – The target maximum number of macro steps for adaptation, positive, by default 15.0.
step_size_init (float, optional) – The initial step size, positive, by default 1.0.
step_accept_rate_target (float, optional) – The acceptance rate target for step size adaptation, in (0, 1), by default 0.8.
step_learning_rate (float, optional) – The learning rate for step size in Adam, positive, by default 0.05.
step_gradient_decay (float, optional) – The step size gradient decay in Adam, positive, by default 0.8.
step_sq_gradient_decay (float, optional) – The step size square gradient decay in Adam, positive, by default 0.9.
step_stabilization (float, optional) – The additive step stabilization factor for Adam, non-negative, by default 1e-4.
step_learn_rate_decay (float, optional) – The learning rate decay for Adam, non-negative, by default 0.5.
save_warmup (bool, optional) – Set to True to save warmup iterations, by default False.
refresh (int, optional) – Period between iteration console feedback, with 0 indicating no feedback, non-netative, by default 0.

Returns:

A list of Markov chain of length num_chains, which may not all have the same number of draws.

Return type:

list[WalnutsOutputArray]

Raises:

ValueError – If any argument is out of its valid range (documented above) or has inconsistent dimensionality.

walnutpie.walnuts_stan(model: StanModel, *, num_chains: int = 4, inits: str | PathLike | Mapping[str, Any] | List[str | PathLike | Mapping[str, Any]] | None = None, seed: int | None = None, id: int = 1, init_radius: float = 2.0, init_inv_metric: ndarray | None = None, save_inv_metric: bool = False, min_warmup_iter: int = 50, max_warmup_iter: int = 1000, min_sampling_iter: int = 50, max_sampling_iter: int = 1000, max_trajectory_doublings: int = 5, max_step_halvings: int = 5, min_micro_steps: int = 1, max_hamiltonian_error: float = 0.5, step_size_converge_tol: float = 0.1, mass_converge_tol: float = 1.0, rhat_converge_tol: float = 1.01, mass_init_count: float = 4.0, mass_additive_smoothing: float = 1e-05, max_macro_steps_target: float = 15.0, step_size_init: float = 1.0, step_accept_rate_target: float = 0.8, step_learning_rate: float = 0.05, step_gradient_decay: float = 0.8, step_sq_gradient_decay: float = 0.9, step_stabilization: float = 0.0001, step_learn_rate_decay: float = 0.5, save_warmup: bool = False, refresh: int = 0) → list[StanOutput][source]#

Sample from the specified Stan model using the specified configuration.

Parameters:

model (bridgestan.StanModel) – The BridgeStan model to fit.
num_chains (int, optional) – The number of Markov chains to run, positive, by default 4.
inits (Union[StanData, List[StanData], None], optional) – The constrained initialization to use for all chains, or a list of constrained initializations, one for each chain, or None to indicate fully random initialization, by default None.
seed (Optional[int], optional) – The pseudo-random number generator seed, non-negative, or None to use a random seed, by default None.
id (int, optional) – Numeric id for the first chain, by default 1. The remaining chains are given consecutive ids following this one. This controls the random number generation, along with the seed.
init_radius (float, optional) – The bounds of uniform random initialization (-init_radius, init_radius), positive, by default 2.0.
init_inv_metric (Optional[np.ndarray], optional) – The diagonal of the initial diagonal inverse metric, positive entries and size equal to transformed (unconstrained) dimension, or None, in which case the mass matrix is initialized with a smoothed. negative outer product of gradients at the initial position, by default None.
save_inv_metric (bool, optional) – Set to True to save the inverse metric after adaptation, by default False.
min_warmup_iter (int, optional) – The minimum number of warmup iterations, greater than or equal to 0, by default 50.
max_warmup_iter (int, optional) – The maximum number of warmup iterations, greater than or equal to min_warmup_iter, by default 1000.
min_sampling_iter (int, optional) – The minimum number of sampling iterations, greater than or equal to 0, by default 50.
max_sampling_iter (int, optional) – The maximum number of sampling iterations, greater than or equal to min_sampling_iter, by default 1000.
max_trajectory_doublings (int, optional) – The maximum number of trajectory doublings for the no-U-turn sampler, positive, by default 5.
max_step_halvings (int, optional) – The maximum number of step size halvings in Walnuts, non-negative, by default 5.
min_micro_steps (int, optional) – The minimum number of micro steps per macro step, positive, by default 1.
max_hamiltonian_error (float, optional) – The maximum error allowed in the Hamiltonian, positive, by default 0.5.
step_size_converge_tol (float, optional) – The relative converge tolerance for difference in step sizes from the geometric mean across chains, positive, by default 0.1.
mass_converge_tol (float, optional) – The relative mass matrix norm convergence tolerance from the geometric mean across chains, by default 1.0.
rhat_converge_tol (float, optional) – The convergence tolernace for R-hat, greater than 1, by default 1.01.
mass_init_count (float, optional) – The pseudo-observation count for the initial mass matrix, positive, by default 4.0.
mass_additive_smoothing (float, optional) – The amount to add to the mass matrix estimators for smoothing, non-negative, by default 1e-5.
max_macro_steps_target (float, optional) – The target maximum number of macro steps for adaptation, positive, by default 15.0.
step_size_init (float, optional) – The initial step size, positive, by default 1.0.
step_accept_rate_target (float, optional) – The acceptance rate target for step size adaptation, in (0, 1), by default 0.8.
step_learning_rate (float, optional) – The learning rate for step size in Adam, positive, by default 0.05.
step_gradient_decay (float, optional) – The step size gradient decay in Adam, positive, by default 0.8.
step_sq_gradient_decay (float, optional) – The step size square gradient decay in Adam, positive, by default 0.9.
step_stabilization (float, optional) – The additive step stabilization factor for Adam, non-negative, by default 1e-4.
step_learn_rate_decay (float, optional) – The learning rate decay for Adam, non-negative, by default 0.5.
save_warmup (bool, optional) – Set to True to save warmup iterations, by default False.
refresh (int, optional) – Period between iteration console feedback, with 0 indicating no feedback, non-netative, by default 0.

Returns:

A list of Stan fits of length num_chains, which may not all have the same number of draws.

Return type:

list[StanOutput]

Raises:

ValueError – If any argument is out of its valid range (documented above) or has inconsistent dimensionality.

Note on reproducibility#

When the adaptive stopping criteria are used (i.e., when min_warmup_iter != max_warmup_iter and min_sampling_iter != max_sampling_iter), thread scheduling will result in different draws being produced for the same seed. When precise reproducibility is required, setting these numbers to the same value will turn off adaptive stopping and result in reproducible runs.

Output Classes#

class walnutpie.pyfunc.WalnutsOutputArray(input_array, warmup: WarmupInfo[ndarray])[source]#

Bases: ndarray

An adapter for ndarray to set extra atrributes.

See the numpy.org documentation: Adding extra attributes to ndarray.

warmup: WarmupInfo[ndarray]#: The saved adaptation and warmup draws, if requested.

class walnutpie.stan.StanOutput(parameters: List[str], data: ndarray, warmup: WarmupInfo[StanOutputBase])[source]#

Bases: StanOutputBase

A holder for the output of a Stan run.

The data attribute contains the raw output from Stan.

If a specific parameter is needed, it can be extracted using the get() method, or by using the object as a dictionary.

Parameters:

parameters (List[str]) – The names of the (constrained) model parameters, in output order.
data (np.ndarray) – The raw sampler output.
warmup (WarmupInfo[StanOutputBase]) – The warmup diagnostics and warmup draws for this chain.

__getitem__(key: str) → ndarray#: Extract a parameter from the Stan output.

create_inits(*, chains: int = 4, seed: int | None = None) → Dict[str, ndarray] | List[Dict[str, ndarray]][source]#

Create a dictionary of parameters suitable for initializing a new Stan run.

Parameters:

chains (int, optional) – The number of chains needed, by default 4.
seed (Optional[int], optional) – The seed to use for the random number generator. If not provided, a random seed will be generated.

Returns:

A dictionary of parameters, or a list of dictionaries if chains > 1.

Return type:

Union[Dict[str, np.ndarray], List[Dict[str, np.ndarray]]]

property data: ndarray#: The underlying draws from the Stan model.

get(key: str) → ndarray#

Extract a parameter from the Stan output. Synonym for obj[key].

Parameters:: key (str) – The name of the parameter to extract.
Returns:: The parameter values. Shape depends on the Stan type and algorithm used.
Return type:: np.ndarray

property parameters: List[str]#: The names of the parameters in the Stan model.

warmup: WarmupInfo[StanOutputBase]#: The saved adaptation and warmup draws, if requested.

class walnutpie.stan.StanOutputBase(parameters: List[str], data: ndarray)[source]#

A holder for the output of a Stan run.

The data attribute contains the raw output from Stan.

If a specific parameter is needed, it can be extracted using the get() method, or by using the object as a dictionary.

class walnutpie.util.WarmupInfo(stepsize: float, inv_metric: ndarray | None, warmup_draws: T | None)[source]#

Warmup output from a single chain, parameterized by the array type.

stepsize#

The adapted step size from warmup.

Type:: float

inv_metric#

The diagonal inverse mass matrix estimated during warmup, or `None if not saved.

Type:: Optional[np.ndarray]

warmup_draws#

The warmup draws, or None if not saved.

Type:: Optional[T]

Posterior Analysis Functions#

When adaptive stopping is used, it is likely that each chain will end with a different number of draws, which may by a challenge to process using existing tools. The following common posterior analysis functions are implemented in such a manner to account for this.

class walnutpie.Summarizer(draws: List[ndarray] | List[StanOutputBase])[source]#

A class to hold multivariate Markov chain Monte Carlo draws and provide summary statistics for their variables.

Construct an instance from a list of multivariate Markov chains.

Parameters:: draws – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.

ess() → ndarray[source]#

Return the estimated effective sample size of the sampled variables.

The implementation uses initial monotonic sequence estimators for integrated autocorrelation. It also discounts ESS for non-convergence across chains.

Returns:: The estimated effective sample sizes.
Return type:: np.ndarray

mcse() → ndarray[source]#

Return an estimate of the Monte Carlo standard error for the sampled variables.

The MCSE is computed in the standard way as the estimated standard deviation divided by the square root of the estimated sample size.

Returns:: The Monte Carlo standard error estimates.
Return type:: np.ndarray

mean()[source]#

Compute the arithmetic mean of sampled variables across all draws.

Returns:: The posterior means.
Return type:: np.ndarray

r_hat() → ndarray[source]#

Return the potential scale reduction statistic R-hat for the sampled variables.

The definition of R-hat for ragged chains is conservative in that it (a) weighs each chain identically, not by chain length, and (b) replaces the factor of (N - 1) / N when all chains are of length N with 1.

Returns:: The R-hat statistics.
Return type:: np.ndarray

standard_deviation()[source]#

Compute the sample standard deviation (ddof = 1) of the sampled variables across all draws.

Returns:: The posterior sample standard deviations.
Return type:: np.ndarray

variance()[source]#

Compute the sample variance (ddof = 1) of the sampled variables across all draws.

Returns:: The posterior sample variances.
Return type:: np.ndarray

walnutpie.ess(draws: List[ndarray] | List[StanOutputBase]) → ndarray[source]#

Return the estimated effective sample size of the sampled variables.

The implementation uses initial monotonic sequence estimators for integrated autocorrelation. It also discounts ESS for non-convergence across chains.

Parameters:: draws (Union[List[np.ndarray], List["StanOutputBase"]]) – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.
Returns:: The estimated effective sample sizes.
Return type:: np.ndarray

walnutpie.r_hat(draws: List[ndarray] | List[StanOutputBase]) → ndarray[source]#

Return the potential scale reduction statistic R-hat for the sampled variables.

Parameters:: draws (Union[List[np.ndarray], List["StanOutputBase"]]) – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.
Returns:: The R-hat statistics.
Return type:: np.ndarray

walnutpie.mcse(draws: List[ndarray] | List[StanOutputBase]) → ndarray[source]#

Return an estimate of the Monte Carlo standard error for the sampled variables.

The MCSE is computed in the standard way as the estimated standard deviation divided by the square root of the estimated sample size.

Parameters:: draws (Union[List[np.ndarray], List["StanOutputBase"]]) – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.
Returns:: The Monte Carlo standard error estimates.
Return type:: np.ndarray

walnutpie.mean(draws: List[ndarray] | List[StanOutputBase]) → ndarray[source]#

Compute the arithmetic mean of sampled variables across all draws.

Parameters:: draws (Union[List[np.ndarray], List["StanOutputBase"]]) – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.
Returns:: The posterior means.
Return type:: np.ndarray

walnutpie.variance(draws: List[ndarray] | List[StanOutputBase]) → ndarray[source]#

Compute the sample variance (ddof = 1) of the sampled variables across all draws.

Parameters:: draws (Union[List[np.ndarray], List["StanOutputBase"]]) – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.
Returns:: The posterior sample variances.
Return type:: np.ndarray

walnutpie.standard_deviation(draws: List[ndarray] | List[StanOutputBase]) → ndarray[source]#

Compute the sample standard deviation (ddof = 1) of the sampled variables across all draws.

Parameters:: draws (Union[List[np.ndarray], List["StanOutputBase"]]) – A list of Markov chains represented as matrices with one row per draw or a list of Stan outputs.
Returns:: The posterior sample standard deviations.
Return type:: np.ndarray