quicksampler.samplers package

Submodules

quicksampler.samplers.DomainChanger module

class quicksampler.samplers.DomainChanger.DomainChanger(ranges, backend='numpy', epsilon=1e-05)[source]

Bases: LogisticTransform

A class for transforming variables between different bounded domains using the Logistic Transformation.

If one wants to transform points using a logistic transformation from a closed interval \([a,b]\) to \((-\infty, +\infty)\). The transformation is defined as:

\[\begin{split}\begin{align*} f &: [a,b] \to (-\infty, \infty) \\ f(x) &= \log\left(\frac{\frac{x-a}{b-a}}{1 - \frac{x-a}{b-a}}\right) \\ f^{-1}(y) &= \frac{b-a}{1+e^{-y}} + a \end{align*}\end{split}\]
Parameters:

  • ranges (dict) – A dictionary specifying the variable ranges, where keys represent variable names and values can be either ‘infinite’ for unbounded variables or a list [a, b] representing the closed interval :obj`[a, b]`.

  • backend (str, optional) – The backend for computation (‘JAX’ or ‘numpy’). Defaults to ‘numpy’.

  • epsilon (float, optional) – A small value for numerical stability, defaults to 1e-5.

compute_transforms()[source]

Compute the transformations and inverse transformations for each variable range.

inverse_log_jacobian(x)[source]

Compute the sum of the logarithm of Jacobian determinants for each variable.

Parameters:

x (dict) – A dictionary containing variable names as keys and their corresponding transformed values.

Returns:

log_det – The sum of the logarithm of Jacobian determinants.

Return type:

float

inverse_transform(x, suffix='')[source]

Inverse transform variables in the given dictionary from the infinite domain.

Parameters:

  • x (dict) – A dictionary containing variable names as keys and their corresponding transformed values.

  • suffix (str, optional) – A suffix to append to the original variable names. Defaults to an empty string.

Returns:

new_x – A dictionary containing inverse transformed variables with original names.

Return type:

dict

inverse_transform_from_infinite_from_range(x, ranges)[source]

Inverse transform variables from (-∞, +∞) to a specified range.

Parameters:

  • x (float or array-like) – The input variable(s) in the infinite domain.

  • ranges (str or list) – The specified variable range.

Returns:

transformed_x – The inverse-transformed variable(s).

Return type:

float or array-like

logprob_wrapped(logprob_func)[source]

Create a wrapped likelihood function that includes the inverse log Jacobian term.

Parameters:

logprob_func (function) – A log probability function that takes a dictionary of variables as input.

Returns:

likelihood_func – A wrapped likelihood function that includes the inverse log Jacobian term.

Return type:

function

transform(x, suffix='')[source]

Transform variables in the given dictionary to the infinite domain.

Parameters:

  • x (dict) – A dictionary containing variable names as keys and their corresponding values.

  • suffix (str, optional) – A suffix to append to the transformed variable names. Defaults to an empty string.

Returns:

new_x – A dictionary containing transformed variables with updated names.

Return type:

dict

transform_to_infinite_from_range(x, ranges)[source]

Transform variables from a specified range to (-∞, +∞).

Parameters:

  • x (float or array-like) – The input variable(s) to be transformed.

  • ranges (str or list) – The specified variable range.

Returns:

transformed_x – The transformed variable(s).

Return type:

float or array-like

class quicksampler.samplers.DomainChanger.LogisticTransform(backend, epsilon=1e-05)[source]

Bases: object

Implementation of the Logistic Transformation between spaces

If one wants to transform points using a logistic transformation from a closed interval \([a,b]\) to \((-\infty, +\infty)\). The transformation is defined as:

\[\begin{split}\begin{align*} f &: [a,b] \to (-\infty, \infty) \\ f(x) &= \log\left(\frac{\frac{x-a}{b-a}}{1 - \frac{x-a}{b-a}}\right) \\ f^{-1}(y) &= \frac{b-a}{1+e^{-y}} + a \end{align*}\end{split}\]
epsilon

Any input within \(\epsilon\) of \([a,b]\) will be clipped onto the value of the transformation at \(f(a + \epsilon)\) or \(f(b - \epsilon)\)

Type:

float, default = 1e-5

backend

Will compute the transformation for JAX arrays if 'JAX' and compute it for numpy arrays if 'numpy'

Type:

str, optional

inverse_transform_from_infinite(y, a, b)[source]

Inverse transforms from \((-\infty, +\infty)\) to \([a,b]\)

Parameters:

  • y (array or float from either JAX or numpy) – The points in the infinite domain to be transformed back.

  • a (array or float from either JAX or numpy) – The lower end of the bounded domain

  • b (array or float from either JAX or numpy) – The upper end of the bounded domain

Returns:

x – The transformed points back into the closed interval [a, b].

Return type:

array or float of size and type similar to y

log_jacobian_determinant(variable_range, y)[source]

Compute the logarithm of the Jacobian determinant of the transformation.

Parameters:

  • variable_range (str or list) – If a string, should be ‘infinite’ or a domain range represented as [a, b]. If a list, it represents the domain range [a, b].

  • y (array or float from either JAX or numpy) – The points in the infinite domain.

Returns:

log_det – The logarithm of the Jacobian determinant of the transformation.

Return type:

float

transform_to_infinite(x, a, b)[source]

Transforms from \([a,b]\) to \((-\infty, +\infty)\)

Parameters:

  • x (array or float of from either JAX or numpy) – The points to be transformed

  • a (array or float of from either JAX or numpy) – The lower end of the bounded domain

  • b (array or float of from either JAX or numpy) – The upper end of the bounded domain

Returns:

y – The transformed points that now live in the infinite domain

Return type:

array or float of size and type similar to x

quicksampler.samplers.MHSampler module

class quicksampler.samplers.MHSampler.AbstractProposal(step_size, backend='numpy', rng_key=None)[source]

Bases: object

Abstract proposal for generating new states based on the current state.

This class serves as a base for specific proposal distributions. It defines the common functionality and interface for proposing new values for variables in the state dictionary.

Parameters:

  • step_size (float) – The step size for the proposal.

  • backend (str, optional) – The backend for computation (‘numpy’ or ‘JAX’), defaults to ‘numpy’.

  • rng_key (array, optional) – The random number generator key, defaults to None.

__call__(x)[source]

Generate a new state based on the current state ‘x’.

class quicksampler.samplers.MHSampler.MHSampler(likelihood, init_position, step_size=1, limits=None, rng_key=None, backend='numpy')[source]

Bases: object

Metropolis-Hastings sampler for generating samples from a distribution.

This class implements the Metropolis-Hastings algorithm for generating samples from a target distribution specified by a likelihood function. It works with both ‘numpy’ and ‘JAX’ backends.

Samplers will accept an object with a logpdf method like the following:

Parameters:

  • likelihood (object) – The likelihood object representing the target distribution.

  • init_position (dict) – The initial position in the state space.

  • step_size (float, optional) – The step size for the Metropolis-Hastings proposal, defaults to 1.

  • limits (dict, optional) – The limits for variables in the state space, defaults to None.

  • rng_key (array, optional) – The random number generator key, defaults to None.

  • backend (str, optional) – The backend for computation (‘numpy’ or ‘JAX’), defaults to ‘numpy’.

accept_reject(p_accept)[source]

Accept or reject a proposed state based on the acceptance probability.

Parameters:

p_accept (float) – The probability of accepting the proposed state.

Returns:

True if the proposed state is accepted, False otherwise.

Return type:

bool

property result

Get the result of the sampler in a dictionary format.

Returns:

The result of the sampler in a dictionary format.

Return type:

dict

run(n_steps=1000, max_rejects=10000)[source]

Run the Metropolis-Hastings sampler for a specified number of steps.

Parameters:

  • n_steps (int, optional) – The number of Metropolis-Hastings steps to run, defaults to 1000.

  • max_rejects (int, optional) – The maximum number of rejections before raising an error, defaults to MAX_REJECTS_DEFAULT.

Returns:

The result of the sampler in a dictionary format.

Return type:

dict

step(x, max_rejects=10000)[source]

Perform a single Metropolis-Hastings step to generate a new state.

Parameters:

  • x (dict) – The current state.

  • max_rejects (int, optional) – The maximum number of rejections before raising an error, defaults to 10000.

Returns:

The new proposed state.

Return type:

dict

class quicksampler.samplers.MHSampler.SphericalGaussianProposal(step_size, backend='numpy', rng_key=None)[source]

Bases: AbstractProposal

Spherical Gaussian proposal for generating new states.

This class defines a proposal distribution where the next step is sampled from a spherical Gaussian distribution with a specified step size.

Parameters:

  • step_size (float) – The step size for the proposal.

  • backend (str, optional) – The backend for computation (‘numpy’ or ‘JAX’), defaults to ‘numpy’.

  • rng_key (array, optional) – The random number generator key, defaults to None.

next_step(key, value)[source]

Compute the next proposed value for any variable inside the state dictionary x.

Within the state dictionary x (e.g., x = {‘mu’: [0.1, 0.2], ‘kappa’: 0.3}), each state variable (the key) has a value that parameterizes its current state. This function will take in a value and compute the change in position needed to go to the next step in the proposal.

e.g.

>>> x = {'mu': np.array([0.1, 0.2]), 'kappa': 0.3}
>>> SphericalGaussianProposal(step_size=0.1)
>>> SphericalGaussianProposal.next_step('mu', x['mu']) # random normal proposal with the right shape.
Parameters:

  • key (str) – The variable key.

  • value (float, array-like) – The current value of the variable.

Returns:

The next proposed value for the variable.

Return type:

array

quicksampler.samplers.MHSampler.convert_to_jax_array(dictionary)[source]

Convert selected values in a dictionary to JAX arrays.

Parameters:

dictionary (dict) – The input dictionary containing values to be converted.

Returns:

A new dictionary with selected values converted to JAX arrays.

Return type:

dict

quicksampler.samplers.MHSampler.create_rng_key(backend)[source]

Create a random number generator key based on the backend.

Parameters:

backend (str) – The backend for computation (‘numpy’ or ‘JAX’).

Returns:

A random number generator key.

Return type:

array

quicksampler.samplers.NUTS module

class quicksampler.samplers.NUTS.NUTS(likelihood, init_position, limits=None, step_size=0.001, inverse_mass_matrix=None, rng_key=None, warmup_steps=100)[source]

Bases: object

Solves a problem specified by a likelihood object using the NUTS sampler.

inference_loop(num_samples, sample_key=None)[source]

Perform the inference loop to obtain samples from the NUTS sampler.

Parameters:

  • num_samples (int) – The number of samples to generate.

  • sample_key (array, optional) – The random number generator key for sampling, defaults to None.

Returns:

A list of states generated by the sampler.

Return type:

states

property result

Get the result of the sampler.

Returns:

result – A dictionary containing the sampled positions.

Return type:

dict

run(num_samples=100)[source]

Run the NUTS sampler to obtain samples.

Parameters:

num_samples (int, optional) – The number of samples to generate, defaults to 100.

Returns:

result – A dictionary containing the sampled positions.

Return type:

dict

step(rng_key, x)[source]

Perform a single step of the NUTS sampler.

Parameters:

  • rng_key (array) – The random number generator key.

  • x (dict) – The current position in the parameter space.

Returns:

The updated state after a single step.

Return type:

state

quicksampler.samplers.NUTS.convert_to_jax_array(dictionary)[source]

Convert selected values in a dictionary to JAX arrays.

Parameters:

dictionary (dict) – The input dictionary containing values to be converted.

Returns:

A new dictionary with selected values converted to JAX arrays.

Return type:

dict

Module contents