blackjax.mcmc.metrics

blackjax.mcmc.metrics#

Metric space in which the Hamiltonian dynamic is embedded.

An important particular case (and the most used in practice) of metric for the position space in the Euclidean metric. It is defined by a definite positive matrix \(M\) with fixed value so that the kinetic energy of the hamiltonian dynamic is independent of the position and only depends on the momentum \(p\) [BBLG17].

For a Newtonian hamiltonian dynamic the kinetic energy is given by:

\[K(p) = \frac{1}{2} p^T M^{-1} p\]

We can also generate a relativistic dynamic [LPH+17].

Classes#

LowRankInverseMassMatrix

Pure-array description of a low-rank inverse mass matrix.

Functions#

`default_metric`(→ Metric)	Convert an input metric into a `Metric` object following sensible default rules.
`gaussian_euclidean`(→ Metric)	Hamiltonian dynamic on euclidean manifold with normally-distributed momentum
`gaussian_euclidean_low_rank`(→ Metric)	Euclidean metric with low-rank-modified mass matrix [SCC26].
`gaussian_riemannian`(→ Metric)	Hamiltonian dynamic on Riemannian manifold with normally-distributed momentum.
`lbfgs_inverse_hessian_to_low_rank_metric`(...)	Convert an L-BFGS factored inverse-Hessian to a `LowRankInverseMassMatrix`.

Module Contents#

class LowRankInverseMassMatrix[source]#

Pure-array description of a low-rank inverse mass matrix.

The inverse mass matrix has the form

\[M^{-1} = \operatorname{diag}(\sigma) \bigl(I + U(\Lambda - I)U^\top\bigr) \operatorname{diag}(\sigma)\]

where \(\sigma \in \mathbb{R}^d_{>0}\), \(U \in \mathbb{R}^{d \times k}\) has orthonormal columns and \(\Lambda = \operatorname{diag}(\lambda)\).

This is the array-only payload produced by window_adaptation_low_rank(). Unlike a fully-constructed Metric (whose fields are Python closures that capture these arrays), this NamedTuple is a pure JAX pytree and can be safely transported across jax.vmap / jax.pmap boundaries.

default_metric() expands this into a Metric at the kernel call site via gaussian_euclidean_low_rank().

sigma[source]#: Shape (d,). Positive diagonal scaling.

U[source]#: Shape (d, k). Matrix with orthonormal columns.

lam[source]#: Shape (k,). Positive eigenvalues.

sigma: blackjax.types.Array[source]#

U: blackjax.types.Array[source]#

lam: blackjax.types.Array[source]#

default_metric(metric: MetricTypes) → Metric[source]#

Convert an input metric into a Metric object following sensible default rules.

The metric can be specified in four different ways:

A Metric object that implements the full interface
A LowRankInverseMassMatrix NamedTuple holding (sigma, U, lam), which is expanded to a full Metric via gaussian_euclidean_low_rank(). This is the form returned by window_adaptation_low_rank() and is safe to transport across jax.vmap boundaries.
An Array which is assumed to specify the inverse mass matrix of a static metric
A function that takes a coordinate position and returns the mass matrix at that location

Returns:

A Metric object with sample_momentum, kinetic_energy,
check_turning, and scale fields.

gaussian_euclidean(inverse_mass_matrix: blackjax.types.Array) → Metric[source]#

Hamiltonian dynamic on euclidean manifold with normally-distributed momentum [Bet13].

The gaussian euclidean metric is a euclidean metric further characterized by setting the conditional probability density \(\pi(momentum|position)\) to follow a standard gaussian distribution. A Newtonian hamiltonian dynamics is assumed.

Parameters:

inverse_mass_matrix – One or two-dimensional array corresponding respectively to a diagonal or dense mass matrix. The inverse mass matrix is multiplied to a flattened version of the Pytree in which the chain position is stored (the current value of the random variables). The order of the variables should thus match JAX’s tree flattening order, and more specifically that of ravel_pytree. In particular, JAX sorts dictionaries by key when flattening them. The value of each variables will appear in the flattened Pytree following the order given by sort(keys).

Returns:

momentum_generator – A function that generates a value for the momentum at random.
kinetic_energy – A function that returns the kinetic energy given the momentum.
is_turning – A function that determines whether a trajectory is turning back on itself given the values of the momentum along the trajectory.

gaussian_euclidean_low_rank(sigma: blackjax.types.Array, U: blackjax.types.Array, lam: blackjax.types.Array) → Metric[source]#

Euclidean metric with low-rank-modified mass matrix [SCC26].

The inverse mass matrix has the form

\[M^{-1} = \operatorname{diag}(\sigma) \bigl(I + U(\Lambda - I)U^\top\bigr) \operatorname{diag}(\sigma)\]

where \(\sigma \in \mathbb{R}^d_{>0}\) is a diagonal scaling, \(U \in \mathbb{R}^{d \times k}\) has orthonormal columns, and \(\Lambda = \operatorname{diag}(\lambda)\) with \(\lambda > 0\). When \(\lambda = \mathbf{1}\) the metric reduces to a diagonal metric with scale \(\sigma\). All HMC operations are \(O(dk)\), making this efficient when \(k \ll d\).

Parameters:

sigma – Shape (d,). Positive diagonal scaling; plays the role of marginal standard deviations.
U – Shape (d, k). Matrix with orthonormal columns spanning the low-rank correction subspace.
lam – Shape (k,). Positive eigenvalues for the low-rank correction.

Return type:

A Metric object whose operations all run in \(O(dk)\).

gaussian_riemannian(mass_matrix_fn: Callable) → Metric[source]#

Hamiltonian dynamic on Riemannian manifold with normally-distributed momentum.

Parameters:

mass_matrix_fn – A callable that takes a position and returns the mass matrix at that location (positive definite, one or two-dimensional array).

Returns:

A Metric object with sample_momentum, kinetic_energy,
check_turning, and scale fields.

lbfgs_inverse_hessian_to_low_rank_metric(alpha: blackjax.types.Array, beta: blackjax.types.Array, gamma: blackjax.types.Array) → LowRankInverseMassMatrix[source]#

Convert an L-BFGS factored inverse-Hessian to a LowRankInverseMassMatrix.

The L-BFGS inverse Hessian is stored in the factored form

\[H^{-1} = \operatorname{diag}(\alpha) + \beta \Gamma \beta^\top\]

(formula II.1 / II.3 of [ZCGV22]). This adapter rewrites it as a LowRankInverseMassMatrix with

\[M^{-1} = \operatorname{diag}(\sigma) \bigl(I + U(\Lambda - I)U^\top\bigr) \operatorname{diag}(\sigma), \quad \sigma = \sqrt{\alpha}\]

via a compact \(O((2m)^3)\) inner eigendecomposition that avoids the \(O(d^3)\) full eigenproblem:

Set \(D = \operatorname{diag}(\sigma)\) and factor out to obtain \(H^{-1} = D(I + \tilde B \Gamma \tilde B^\top)D\) with \(\tilde B = D^{-1}\beta \in \mathbb{R}^{d \times 2m}\).
QR-decompose \(\tilde B = Q R\) (thin QR, \(Q\) orthonormal \(d \times r\), \(r = \min(d, 2m)\)).
The inner correction satisfies \(\tilde B \Gamma \tilde B^\top = Q (R \Gamma R^\top) Q^\top\), so its eigenvalues are those of the \(r \times r\) matrix \(R \Gamma R^\top\).
Eigendecompose \(R \Gamma R^\top = V \Lambda_c V^\top\) (eigh, \(r \times r\) only).
Return \(U = QV\) (orthonormal eigenvectors, shape \(d \times r\)) and \(\lambda = 1 + \Lambda_c\) (eigenvalues of \(I + \tilde B \Gamma \tilde B^\top\)).

When to use this. Pass the (alpha, beta, gamma) triple produced by lbfgs_inverse_hessian_factors() to obtain a JAX-pytree-safe LowRankInverseMassMatrix that can cross jax.vmap boundaries and feed any consumer that accepts the unified representation (HMC, MCLMC, etc.).

Note

This function is a pure adapter — it does not alter Pathfinder’s internal sampling path. Wiring this adapter into Pathfinder’s own sampling path is deliberate follow-up work; for now it ships as an adapter with parity tests only.

Warning

Positive-definiteness precondition. The triple (alpha, beta, gamma) must yield a positive-definite dense form diag(alpha) + beta @ gamma @ beta.T; this is guaranteed when the triple comes from lbfgs_inverse_hessian_factors() under a Wolfe-condition line search. Non-positive-definite inputs produce lam <= 0 silently here and surface as NaN at momentum sampling. Additionally, at float32 near-singular metrics (condition number ≳ 1e7) can resolve the smallest eigenvalue with unreliable sign; for such inputs prefer float64 factors.

Parameters:

alpha – Shape (d,). Positive diagonal of the inverse Hessian approximation.
beta – Shape (d, 2m). Left factor of the low-rank correction. When m = 0 (empty L-BFGS history) beta has shape (d, 0) and the adapter returns a pure diagonal metric.
gamma – Shape (2m, 2m). Symmetric inner factor of the low-rank correction.

Returns:

With sigma = sqrt(alpha), U the (d, r) orthonormal eigenvector matrix, and lam = 1 + eigenvalues(R Γ Rᵀ). Empty-history edge: U has shape (d, 0) and lam shape (0,), representing a pure diagonal metric with scale sigma.

Return type:

LowRankInverseMassMatrix