nifty8.re.hmc module

class AcceptedAndRejected(accepted_qp, rejected_qp, accepted, diverging)

Bases: NamedTuple

accepted: Array | bool

Alias for field number 2

accepted_qp: QP

Alias for field number 0

diverging: Array | bool

Alias for field number 3

rejected_qp: QP

Alias for field number 1

class QP(position: Q, momentum: Q)

Bases: NamedTuple

Object holding a pair of position and momentum.

Variables:

• position (Q) – Position.

• momentum (Q) – Momentum.

momentum: Q

Alias for field number 1

position: Q

Alias for field number 0

class Tree(left: QP, right: QP, logweight: Array | float, proposal_candidate: QP, turning: Array | bool, diverging: Array | bool, depth: Array | int, cumulative_acceptance: Array | float)

Bases: NamedTuple

Object carrying tree metadata.

Variables:

• right (left,) – Respective endpoints of the trees path.

• logweight (Union[jnp.ndarray, float]) – Sum over all -H(q, p) in the tree’s path.

• proposal_candidate (QP) – Sample from the trees path, distributed as exp(-H(q, p)).

• turning (Union[jnp.ndarray, bool]) – Indicator for either the left or right endpoint are a uturn or any subtree is a uturn.

• diverging (Union[jnp.ndarray, bool]) – Indicator for a large increase in energy in the next larger tree.

• depth (Union[jnp.ndarray, int]) – Levels of the tree.

• cumulative_acceptance (Union[jnp.ndarray, float]) – Sum of all acceptance probabilities relative to some initial energy value. This value is distinct from logweight as its absolute value is only well defined for the very final tree of NUTS.

cumulative_acceptance: Array | float

Alias for field number 7

depth: Array | int

Alias for field number 6

diverging: Array | bool

Alias for field number 5

left: QP

Alias for field number 0

logweight: Array | float

Alias for field number 2

proposal_candidate: QP

Alias for field number 3

right: QP

Alias for field number 1

turning: Array | bool

Alias for field number 4

add_single_qp_to_tree(key, tree, qp, go_right, potential_energy, kinetic_energy, inverse_mass_matrix, initial_neg_energy, max_energy_difference)

Helper function for progressive sampling. Takes a tree with a sample, and a new endpoint, propagates sample.

count_trailing_ones(n)

Count the number of trailing, consecutive ones in the binary representation of n.

Warning

n must be positive and strictly smaller than 2**64

Examples

>>> print(bin(23), count_trailing_one_bits(23))
0b10111 3

flip_momentum(qp: QP) → QP

generate_hmc_acc_rej(*, key, initial_qp, potential_energy, kinetic_energy, inverse_mass_matrix, stepper, num_steps, step_size, max_energy_difference) → AcceptedAndRejected

Generate a sample given the initial position.

Parameters:

• key (ndarray) – a PRNGKey used as the random key

• position (ndarray) – The the starting position of this step of the markov chain.

• potential_energy (Callable[[ndarray], float]) – The potential energy, which is the distribution to be sampled from.

• mass_matrix (ndarray) – The mass matrix used in the kinetic energy

• num_steps (int) – The number of steps the leapfrog integrator should perform.

• step_size (float) – The step size (usually epsilon) for the leapfrog integrator.

generate_nuts_tree(initial_qp, key, step_size, max_tree_depth, stepper: Callable[[Array | float, Q, QP], QP], potential_energy, kinetic_energy: Callable[[Q, Q], float], inverse_mass_matrix: Q, bias_transition: bool = True, max_energy_difference: Array | float = inf) → Tree

Generate a sample given the initial position.

This call implements a No-U-Turn-Sampler.

Parameters:

• initial_qp (QP) – Starting pair of (position, momentum). NOTE, the momentum must be resampled from conditional distribution BEFORE passing it into this function!

• key (ndarray) – PRNGKey used as the random key.

• step_size (float) – Step size (usually called epsilon) for the leapfrog integrator.

• max_tree_depth (int) – The maximum depth of the trajectory tree before the expansion is terminated. At the maximum iteration depth, the current value is returned even if the U-turn condition is not met. The maximum number of points (/integration steps) per trajectory is . This function requires memory linear in max_tree_depth, i.e. logarithmic in trajectory length. It is used to statically allocate memory in advance.

• stepper (Callable[[float, Q, QP], QP]) – The function that performs (Leapfrog) steps. Takes as arguments (in order) 1) step size (containing the direction): float , 2) inverse mass matrix: Q , 3) starting point: QP .

• potential_energy (Callable[[Q], float]) – The potential energy, of the distribution to be sampled from. Takes only the position part (QP.position) as argument.

• kinetic_energy (Callable[[Q, Q], float], optional) – Mapping of the momentum to its corresponding kinetic energy. As argument the function takes the inverse mass matrix and the momentum.

Returns:

current_tree – The final tree, carrying a sample from the target distribution.

Return type:

Tree

See also

No-U-Turn, NumPyro, Combination, Sampling

is_euclidean_uturn(qp_left, qp_right)

See also

Betancourt

iterative_build_tree(key, initial_tree, step_size, go_right, stepper, potential_energy, kinetic_energy, inverse_mass_matrix, max_tree_depth, initial_neg_energy, max_energy_difference)

Starting from either the left or right endpoint of a given tree, builds a new adjacent tree of the same size.

Parameters:

• key (ndarray) – PRNGKey to choose a sample when adding QPs to the tree.

• initial_tree (Tree) – Tree to be extended (doubled) on the left or right.

• step_size (float) – The step size (usually called epsilon) for the leapfrog integrator.

• go_right (bool) – If go_right start at the right end, going right else start at the left end, going left.

• stepper (Callable[[float, Q, QP], QP]) – The function that performs (Leapfrog) steps. Takes as arguments (in order) 1) step size (containing the direction): float , 2) inverse mass matrix: Q , 3) starting point: QP .

• potential_energy (Callable[[Q], float]) – Potential energy, of the distribution to be sampled from. Takes only the position part (QP.position) as argument.

• kinetic_energy (Callable[[Q, Q], float], optional) – Mapping of the momentum to its corresponding kinetic energy. As argument the function takes the inverse mass matrix and the momentum.

• max_tree_depth (int) – An upper bound on the ‘depth’ argument, but has no effect on the functions behaviour. It’s only required to statically set the size of the S array (Q).

leapfrog_step(potential_energy_gradient, kinetic_energy_gradient, step_size, inverse_mass_matrix, qp: QP)

Perform one iteration of the leapfrog integrator forwards in time.

Parameters:

• potential_energy_gradient (Callable[[ndarray], float]) – Potential energy gradient part of the hamiltonian (V). Depends on position only.

• qp (QP) – Point in position and momentum space from which to start integration.

• step_size (float) – Step length (usually called epsilon) of the leapfrog integrator.

merge_trees(key, current_subtree, new_subtree, go_right, bias_transition)

Merges two trees, propagating the proposal_candidate

sample_momentum_from_diagonal(*, key, mass_matrix_sqrt)

Draw a momentum sample from the kinetic energy of the hamiltonian.

Parameters:

• key (ndarray) – PRNGKey used as the random key.

• mass_matrix_sqrt (ndarray) – The left square-root mass matrix (i.e. square-root of the inverse diagonal covariance) to use for sampling. Diagonal matrix represented as (possibly pytree of) ndarray vector containing the entries of the diagonal.

select(pred, on_true, on_false)

total_energy_of_qp(qp, potential_energy, kinetic_energy_w_inv_mass)

tree_index_get(ptree, idx)

tree_index_update(x, idx, y)