API Reference

TMGSampler

class tmg_hmc.TMGSampler(mu=None, Sigma=None, T=None, gpu=False, *, Sigma_half=None)

Bases: object

Hamiltonian Monte Carlo sampler for Multivariate Gaussian distributions with linear and quadratic constraints.

Parameters:

• mu (Array | None)

• Sigma (Array | None)

• T (float | None)

• gpu (bool)

• Sigma_half (Array | None)

add_constraint(*, A=None, f=None, c=0.0, sparse=True, compiled=True)

Adds a constraint to the sampler.

\[\mathbf{x}^T A \mathbf{x} + \mathbf{f}^T \mathbf{x} + c \geq 0\]

Parameters:

• A (Array, optional) – Quadratic term matrix, defaults to the zero matrix if not provided.

• f (Array, optional) – Linear term vector, defaults to the zero vector if not provided.

• c (float, optional) – Constant term. Default is 0.0.

• sparse (bool, optional) – Whether to store A and f in sparse format. Default is True.

• compiled (bool, optional) – Whether to use compiled constraint solutions for full quadratic constraints. Default is True.

Raises:

ValueError – If A is not symmetric when provided, or if neither A nor f is provided.

Return type:

None

Notes

The constraint is automatically transformed to account for the Gaussian’s mean and covariance.

\[\mathbf{y}^T (S A S) \mathbf{y} + (2 S A \pmb{\mu} + S \mathbf{f})^T \mathbf{y} + (\pmb{\mu}^T A \pmb{\mu} + \pmb{\mu}^T \mathbf{f} + c) \geq 0\]

where \(\mathbf{y} = S^{-1}(\mathbf{x} - \pmb{\mu})\) and \(S = \Sigma^{1/2}\). Depending on whether \(A\) and \(\mathbf{f}\) are non-zero after transformation, the appropriate constraint type is chosen.

add_product_constraint(*, parameters, sparse=True, compiled=True)

Adds a constraint to the sampler.

\[\mathbf{x}^T A \mathbf{x} + \mathbf{f}^T \mathbf{x} + c \geq 0\]

Parameters:

• parameters (list[list[Array]] | list[dict[str,Array]]) – List of constraint parameters as either lists [A, f, c] or dictionaries {‘A’: A, ‘f’: f, ‘c’: c}. If list, each element must be of length 3 corresponding to A, f, and c. If dictionary, missing keys ‘A’, ‘f’, and ‘c’ default to None, None, and 0.0 respectively.

• sparse (bool, optional) – Whether to store A and f in sparse format. Default is True.

• compiled (bool, optional) – Whether to use compiled constraint solutions for full quadratic constraints. Default is True.

Raises:

ValueError – If A is not symmetric when provided, or if neither A nor f is provided.

Return type:

None

Notes

For product constraints, you must provide lists of each component (A, f, c). The constraint is automatically transformed to account for the Gaussian’s mean and covariance.

\[\mathbf{y}^T (S A S) \mathbf{y} + (2 S A \pmb{\mu} + S \mathbf{f})^T \mathbf{y} + (\pmb{\mu}^T A \pmb{\mu} + \pmb{\mu}^T \mathbf{f} + c) \geq 0\]

where \(\mathbf{y} = S^{-1}(\mathbf{x} - \pmb{\mu})\) and \(S = \Sigma^{1/2}\). Depending on whether \(A\) and \(\mathbf{f}\) are non-zero after transformation, the appropriate constraint type is chosen.

sample_xdot()

Samples a new momentum vector xdot from the standard normal distribution handling GPU if necessary.

Return type:

ndarray | Tensor

sample(x0=None, n_samples=100, burn_in=100, verbose=False, cont=False)

Generates samples from the truncated multivariate Gaussian distribution.

Parameters:

• x0 (Array | None) – Initial point for the sampler. Optional if cont is True.

• n_samples (int, optional) – Number of samples to generate, default is 100.

• burn_in (int, optional) – Number of burn-in iterations, default is 100.

• verbose (bool, optional) – Whether to print verbose output, default is False.

• cont (bool, optional) – Whether to continue from the last sampled point. Default is False.

Returns:

Generated samples.

Return type:

Array

Raises:

ValueError – If cont is False and x0 is not provided, or if x0 does not satisfy constraints.

save(filename)

Saves the sampler state to a pickled file.

Return type:

None

Parameters:

filename (str)

classmethod load(filename)

Loads the sampler state from a pickled file.

Return type:

TMGSampler

Parameters:

filename (str)

Constraints

class tmg_hmc.constraints.LinearConstraint(f, c)

Bases: Constraint

Constraint of the form \(\mathbf{f}^T \mathbf{x} + c \geq 0\)

Parameters:

• f (Array)

• c (float)

classmethod build_from_dict(d, gpu)

Build a LinearConstraint from a dictionary representation

Parameters:

• d (dict) – Dictionary representation of the constraint

• gpu (bool) – Whether to load tensors onto the GPU

Returns:

The constructed constraint

Return type:

Linear Constraint

value(x)

Compute the value of the constraint at x

Parameters:

x (Array) – Point to evaluate the constraint at

Returns:

Value of the constraint at \(\mathbf{x}\) given by \(\mathbf{f}^T \mathbf{x} + c\)

Return type:

float

normal(x)

Compute the normal vector of the constraint at x

Parameters:

x (Array) – Point to evaluate the normal vector at

Returns:

Normal vector of the constraint at \(\mathbf{x}\) given by \(\mathbf{f}\)

Return type:

Array

compute_q(a, b)

Compute the 2 q terms for the linear constraint

Parameters:

• a (Array) – The velocity of the point in the HMC trajectory

• b (Array) – The position of the point in the HMC trajectory

Returns:

q terms for the constraint

Return type:

Tuple[float, float]

Notes

These expressions are defined such that Eqn 2.22 in Pakman and Paninski (2014) simplifies to:

\[q_1 \sin(t) + q_2 \cos(t) + c = 0\]

hit_time(x, xdot)

Compute the hit time of the constraint along the trajectory defined by x and xdot

Parameters:

• x (Array) – The position of the point in the HMC trajectory

• xdot (Array) – The velocity of the point in the HMC trajectory

Returns:

Hit time of the constraint along the trajectory

Return type:

Array

Notes

Hit time is computed by solving Eqn 2.26 in Pakman and Paninski (2014) See resources/HMC_exact_soln.nb for derivation Due to the sum of inverse trig functions, we check the solution and the solution \(\pm \pi\) to ensure we capture all hit times.

Only positive hit times are returned and any ghost solutions are filtered out at a later stage.

class tmg_hmc.constraints.SimpleQuadraticConstraint(A, c, S, sparse=False)

Bases: BaseQuadraticConstraint

Constraint of the form \(\mathbf{x}^T A \mathbf{x} + c \geq 0\)

Parameters:

• A (Array)

• c (float)

• S (Array)

• sparse (bool)

classmethod build_from_dict(d, gpu)

Build a SimpleQuadraticConstraint from a dictionary representation

Parameters:

• d (dict) – Dictionary representation of the constraint

• gpu (bool) – Whether to load tensors onto the GPU

Returns:

The constructed constraint

Return type:

SimpleQuadraticConstraint

value_(x)

Compute the value of the constraint at x using dense matrix computations

Parameters:

x (Array) – Point to evaluate the constraint at

Returns:

Value of the constraint at \(\mathbf{x}\) given by \(\mathbf{x}^T A \mathbf{x} + c\)

Return type:

float

value_sparse(x)

Compute the value of the constraint at x using sparse matrix computations

Parameters:

x (Array) – Point to evaluate the constraint at

Returns:

Value of the constraint at \(\mathbf{x}\) given by \(\mathbf{x}^T A \mathbf{x} + c\)

Return type:

float

normal_(x)

Compute the normal vector at x using dense matrix computations

Parameters:

x (Array) – Point to evaluate the normal vector at

Returns:

Normal vector at \(\mathbf{x}\) given by \(2 A \mathbf{x}\)

Return type:

Array

normal_sparse(x)

Compute the normal vector at x using sparse matrix computations

Parameters:

x (Array) – Point to evaluate the normal vector at

Returns:

Normal vector at \(\mathbf{x}\) given by \(2 A \mathbf{x}\)

Return type:

Array

compute_q_(a, b)

Compute the 3 q terms for the simple quadratic constraint using dense matrix computations

Parameters:

• a (Array) – The velocity of the point in the HMC trajectory

• b (Array) – The position of the point in the HMC trajectory

Returns:

q terms for the constraint

Return type:

Tuple[float, float, float]

Notes

These expressions are the nonzero q terms defined in equation 2.45 in Pakman and Paninski (2014)

compute_q_sparse(a, b)

Compute the 3 q terms for the simple quadratic constraint using sparse matrix computations

Parameters:

• a (Array) – The velocity of the point in the HMC trajectory

• b (Array) – The position of the point in the HMC trajectory

Returns:

q terms for the constraint

Return type:

Tuple[float, float, float]

Notes

These expressions are the nonzero q terms defined in equation 2.45 in Pakman and Paninski (2014)

hit_time(x, xdot)

Compute the hit time for the simple quadratic constraint

Parameters:

• x (Array) – The position of the point in the HMC trajectory

• xdot (Array) – The velocity of the point in the HMC trajectory

Returns:

The hit time for the constraint

Return type:

Array

Notes

Hit time is computed by solving Eqn 2.45 in Pakman and Paninski (2014) See resources/HMC_exact_soln.nb for derivation Only positive hit times are returned and any ghost solutions are filtered out at a later stage.

class tmg_hmc.constraints.QuadraticConstraint(A, f, c, S, sparse=True, compiled=True)

Bases: BaseQuadraticConstraint

Constraint of the form \(\mathbf{x}^T A \mathbf{x} + \mathbf{f}^T \mathbf{x} + c \geq 0\)

Parameters:

• A (Array)

• f (Array)

• c (float)

• S (Array)

• sparse (bool)

• compiled (bool)

classmethod build_from_dict(d, gpu)

Build a QuadraticConstraint from a dictionary representation

Parameters:

• d (dict) – Dictionary representation of the constraint

• gpu (bool) – Whether to load tensors onto the GPU

Returns:

The constructed constraint

Return type:

QuadraticConstraint

value_(x)

Compute the value of the constraint at x using dense matrix computations

Parameters:

x (Array) – Point to evaluate the constraint at

Returns:

The value of the constraint at \(\mathbf{x}\) given by \(\mathbf{x}^T A \mathbf{x} + \mathbf{f}^T \mathbf{x} + c\)

Return type:

float

value_sparse(x)

Compute the value of the constraint at x using sparse matrix computations

Parameters:

x (Array) – Point to evaluate the constraint at

Returns:

The value of the constraint at \(\mathbf{x}\) given by \(\mathbf{x}^T A \mathbf{x} + \mathbf{f}^T \mathbf{x} + c\)

Return type:

float

normal_(x)

Compute the normal vector at x using dense matrix computations :type x: ndarray | Tensor :param x: Point to evaluate the normal vector at :type x: Array

Returns:

Normal vector at \(\mathbf{x}\) given by \(2 A \mathbf{x} + \mathbf{f}\)

Return type:

Array

Parameters:

x (ndarray | Tensor)

normal_sparse(x)

Compute the normal vector at x using sparse matrix computations :type x: ndarray | Tensor :param x: Point to evaluate the normal vector at :type x: Array

Returns:

Normal vector at \(\mathbf{x}\) given by \(2 A \mathbf{x} + \mathbf{f}\)

Return type:

Array

Parameters:

x (ndarray | Tensor)

compute_q_(a, b)

Compute the 5 q terms for the quadratic constraint using dense matrix computations

Parameters:

• a (Array) – The velocity of the point in the HMC trajectory

• b (Array) – The position of the point in the HMC trajectory

Returns:

q terms for the quadratic constraint

Return type:

Tuple[float, float, float, float, float]

Notes

These expressions are defined in Eqns 2.40-2.44 in Pakman and Paninski (2014)

compute_q_sparse(a, b)

Compute the 5 q terms for the quadratic constraint using sparse matrix computations

Parameters:

• a (Array) – The velocity of the point in the HMC trajectory

• b (Array) – The position of the point in the HMC trajectory

Returns:

q terms for the quadratic constraint

Return type:

Tuple[float, float, float, float, float]

Notes

These expressions are defined in Eqns 2.40-2.44 in Pakman and Paninski (2014)

hit_time_cpp(x, xdot)

Compute the hit time for the quadratic constraint using compiled code

Parameters:

• x (Array) – The position of the point in the HMC trajectory

• xdot (Array) – The velocity of the point in the HMC trajectory

Returns:

The hit time for the constraint

Return type:

Array

Notes

Hit time is computed by solving Eqn 2.48 in Pakman and Paninski (2014) See resources/HMC_exact_soln.nb for derivation Only positive hit times are returned and any ghost solutions are filtered out at a later stage.

Compiled code is both written in C++ and optimized to remove all redundant computations see paper for details.

hit_time_py(x, xdot)

Compute the hit time for the quadratic constraint using Python code

Parameters:

• x (Array) – The position of the point in the HMC trajectory

• xdot (Array) – The velocity of the point in the HMC trajectory

Returns:

The hit time for the constraint

Return type:

Array

Notes

Hit time is computed by solving Eqn 2.48 in Pakman and Paninski (2014) See resources/HMC_exact_soln.nb for derivation Only positive hit times are returned and any ghost solutions are filtered out at a later stage.

It is highly recommended to use the compiled version for performance reasons. This Python version is maintained for testing and validation purposes.

hit_time(x, xdot)

Dispatch method for the hit time for the quadratic constraint

Parameters:

• x (Array) – The position of the point in the HMC trajectory

• xdot (Array) – The velocity of the point in the HMC trajectory

Returns:

The hit time for the constraint

Return type:

Array

Notes

Hit time is computed by solving Eqn 2.48 in Pakman and Paninski (2014) See resources/HMC_exact_soln.nb for derivation Only positive hit times are returned and any ghost solutions are filtered out at a later stage.

This method dispatches to either the compiled or Python version based on the compiled attribute. It is highly recommended to use the compiled version for performance reasons.

class tmg_hmc.constraints.ProductConstraint(constraints)

Bases: Constraint

Constraint that is the product of multiple linear or quadratic constraints

Parameters:

constraints (Tuple[Constraint, ...])

value(x)

Compute the value of the product constraint at x

Parameters:

x (Array) – Point to evaluate the constraint at

Returns:

Value of the product constraint at x

Return type:

float

normal(x)

Compute the normal vector of the product constraint at x

Parameters:

x (Array) – Point to evaluate the normal vector at

Returns:

Normal vector of the product constraint at x

Return type:

Array

hit_time(x, xdot)

Compute the hit time of the product constraint along the trajectory defined by \(\mathbf{x}\) and \(\dot{\mathbf{x}}\).

Parameters:

• x (Array) – The position of the point in the HMC trajectory

• xdot (Array) – The velocity of the point in the HMC trajectory

Returns:

Hit times of the product constraint along the trajectory

Return type:

Array

compute_q(a, b)

Compute the coefficients of the constraint equation along the trajectory defined by \(\mathbf{a}\) and \(\mathbf{b}\).

Return type:

Tuple[float, ...]

Parameters:

• a (ndarray | Tensor)

• b (ndarray | Tensor)