aspire.samples

Attributes

logger

Classes

BaseSamples	Class for storing samples and corresponding weights.
Samples	Class for storing samples and corresponding weights.
MCMCSamples	Class for storing MCMC samples and chain metadata.
PTMCMCSamples	Class for storing parallel-tempered MCMC samples.
SMCSamples	Class for storing samples and corresponding weights.

Module Contents

aspire.samples.logger

class aspire.samples.BaseSamples

Class for storing samples and corresponding weights.

If xp is not specified, all inputs will be converted to match the array type of x.

x: array_api_compat.common._typing.Array

Array of samples, shape (n_samples, n_dims).

log_likelihood: array_api_compat.common._typing.Array | None = None

Log-likelihood values for the samples.

log_prior: array_api_compat.common._typing.Array | None = None

Log-prior values for the samples.

log_q: array_api_compat.common._typing.Array | None = None

Log-probability values under the proposal distribution.

parameters: list[str] | None = None

List of parameter names.

dtype: Any | str | None = None

Data type of the samples.

If None, the default dtype for the array namespace will be used.

xp: Callable | None = None

The array namespace to use for the samples.

If None, the array namespace will be inferred from the type of x.

device: Any = None

Device to store the samples on.

If None, the device will be inferred from the array namespace of x.

property dims

Number of dimensions (parameters) in the samples.

to_numpy(dtype=None)

Parameters:

dtype (Any | str | None)

to_namespace(xp, dtype=None)

Parameters:

dtype (Any | str | None)

array_to_namespace(x, dtype=None)

Convert an array to the same namespace as the samples

to_dict(flat=True, copy=True)

Convert the samples to a dictionary.

Parameters:

• flat (bool) – If True, the samples are stored as separate keys for each parameter. If False, the samples are stored in a “samples” key as a dictionary of parameter arrays.

• copy (bool) – If True, the arrays in the dictionary are deep-copied. If False, they are not copied and may share memory with the original samples.

Returns:

A dictionary representation of the samples.

Return type:

dict

classmethod from_dict(dictionary)

Load samples from a dictionary.

The dictionary can either be in a flat format, where the samples are stored as separate keys for each parameter, or in a nested format, where the samples are stored in a “samples” key as a dictionary of parameter.

to_dataframe(include=None)

Convert the samples to a pandas DataFrame.

Only includes samples, log_likelihood, log_prior, and log_q by default, since additional fields have varying shapes and may not be compatible with a DataFrame format.

Parameters:

include (list[str] | None) – List of fields to include in the DataFrame. If None, includes x, log_likelihood, log_prior, and log_q. x is always included irrespective of the value of include.

Returns:

A DataFrame representation of the samples.

Return type:

pd.DataFrame

plot_corner(parameters=None, fig=None, **kwargs)

Plot a corner plot of the samples.

Parameters:

• parameters (list[str] | None) – List of parameters to plot. If None, all parameters are plotted. Figure to plot on. If None, a new figure is created.

• **kwargs (dict) – Additional keyword arguments to pass to corner.corner(). Kwargs are deep-copied before use.

• fig (matplotlib.figure.Figure | None)

save(h5_file, path='samples', flat=False)

Save the samples to an HDF5 file.

This converts the samples to numpy and then to a dictionary.

Parameters:

• h5_file (h5py.File) – The HDF5 file to save to.

• path (str) – The path in the HDF5 file to save to.

• flat (bool) – If True, save the samples as a flat dictionary. If False, save the samples as a nested dictionary.

classmethod load(h5_file, path='samples')

Load the samples from an HDF5 file.

classmethod concatenate(samples)

Concatenate multiple Samples objects.

Parameters:

samples (list[BaseSamples])

Return type:

BaseSamples

classmethod from_samples(samples, **kwargs)

Create a Samples object from a BaseSamples object.

Parameters:

samples (BaseSamples)

Return type:

BaseSamples

class aspire.samples.Samples

Bases: BaseSamples

Class for storing samples and corresponding weights.

If xp is not specified, all inputs will be converted to match the array type of x.

log_w: array_api_compat.common._typing.Array

weights: array_api_compat.common._typing.Array

evidence: float

evidence_error: float

log_evidence: float | None = None

log_evidence_error: float | None = None

effective_sample_size: float

property efficiency

Efficiency of the weighted samples.

Defined as ESS / number of samples.

compute_weights()

Compute the posterior weights.

property scaled_weights

rejection_sample(rng=None)

plot_corner(include_weights=True, **kwargs)

Plot a corner plot of the samples.

Parameters:

• parameters (list[str] | None) – List of parameters to plot. If None, all parameters are plotted. Figure to plot on. If None, a new figure is created.

• **kwargs (dict) – Additional keyword arguments to pass to corner.corner(). Kwargs are deep-copied before use.

• include_weights (bool)

to_namespace(xp)

to_numpy()

to_dataframe(include=None)

Convert the samples to a pandas DataFrame.

By default, includes log_likelihood, log_prior, log_q, and log_w. See parent class for more details.

Parameters:

include (list[str] | None) – List of fields to include in the DataFrame. If None, includes log_likelihood, log_prior, log_q, and log_w.

Returns:

A DataFrame representation of the samples.

Return type:

pd.DataFrame

class aspire.samples.MCMCSamples

Bases: BaseSamples

Class for storing MCMC samples and chain metadata.

Samples are stored flattened in x, with chain_shape capturing the original chain layout (excluding the parameter dimension).

chain_shape: tuple[int, Ellipsis] | None = None

Shape of the chain excluding the parameter dimension.

thin: int | None = None

Thinning factor used to produce the chain, if any.

burn_in: int | None = None

Number of burn-in steps removed, if any.

autocorrelation_time: array_api_compat.common._typing.Array | None = None

Autocorrelation time per dimension, if available.

minimum_chain_ndim: ClassVar[int] = 2

Minimum required dimensionality for the input chain.

classmethod from_chain(chain, log_likelihood=None, log_prior=None, log_q=None, parameters=None, xp=None, dtype=None, device=None, thin=None, burn_in=None, autocorrelation_time=None, **extra_kwargs)

Create samples from a chain array.

Parameters:

• chain – MCMC chain with shape (…, n_dims). For a single chain, this is typically (n_steps, n_dims). For ensemble samplers, use (n_steps, n_walkers, n_dims).

• log_likelihood – Optional arrays matching the chain shape without the last dimension. They will be flattened to align with x.

• log_prior – Optional arrays matching the chain shape without the last dimension. They will be flattened to align with x.

• log_q – Optional arrays matching the chain shape without the last dimension. They will be flattened to align with x.

• parameters – Optional list of parameter names.

• xp – Optional array namespace. If None, inferred from chain.

• dtype (Any | str | None) – Optional dtype to use for stored arrays.

• device (Any) – Optional device to place arrays on.

• thin (int | None) – Optional metadata describing how the chain was processed.

• burn_in (int | None) – Optional metadata describing how the chain was processed.

• autocorrelation_time (float | None)

Return type:

MCMCSamples

property chain

The chain reshaped to its original layout

property n_steps: int

Number of steps in the chain, excluding walkers

Return type:

int

property n_chains: int

Number of parallel chains (e.g. walkers), if applicable.

Return type:

int

post_process(burn_in=0, thin=1)

Return a new MCMCSamples object with burn-in discarded and/or thinned.

Parameters:

• burn_in (int)

• thin (int)

Return type:

MCMCSamples

to_numpy()

class aspire.samples.PTMCMCSamples

Bases: MCMCSamples

Class for storing parallel-tempered MCMC samples.

betas: array_api_compat.common._typing.Array | None = None

Inverse temperatures for the chains.

Should be a 1D array of shape (n_temps,) in decreasing order, starting at 1.

minimum_chain_ndim: ClassVar[int] = 3

Minimum required dimensionality for the PTMCMC input chain.

subsample(n_samples_per_temperature, rng=None)

Subsample the chain to a fixed number of samples per temperature.

Samples are drawn without replacement independently for each temperature.

Parameters:

• n_samples_per_temperature (int) – Number of samples to draw per temperature.

• rng – Optional numpy random generator. If None, a new one is created.

Return type:

PTMCMCSamples

classmethod from_chain(chain, betas=None, log_likelihood=None, log_prior=None, log_q=None, parameters=None, xp=None, dtype=None, device=None, thin=None, burn_in=None, autocorrelation_time=None)

Create samples from a parallel-tempered MCMC chain.

Parameters:

• chain – PTMCMC chain with shape (n_temps, …, n_dims). A common layout is (n_temps, n_steps, n_walkers, n_dims).

• betas – Optional inverse temperatures with shape (n_temps,).

• log_likelihood – Optional arrays matching the chain shape without the last dimension. They will be flattened to align with x.

• log_prior – Optional arrays matching the chain shape without the last dimension. They will be flattened to align with x.

• log_q – Optional arrays matching the chain shape without the last dimension. They will be flattened to align with x.

• parameters – Optional list of parameter names.

• xp – Optional array namespace. If None, inferred from chain.

• dtype (Any | str | None) – Optional dtype to use for stored arrays.

• device (Any) – Optional device to place arrays on.

• thin (int | None) – Optional metadata describing how the chain was processed.

• burn_in (int | None) – Optional metadata describing how the chain was processed.

• autocorrelation_time (float | None)

Return type:

PTMCMCSamples

property n_temps: int

Number of temperatures in the parallel-tempered chain.

Return type:

int

at_temperature(index)

Parameters:

index (int)

Return type:

MCMCSamples

cold_chain()

Return type:

MCMCSamples

log_evidence_thermodynamic_integration(burn_in_fraction=0.1, method='variance')

Compute the log evidence using thermodynamic integration.

Notes

By default, follows the implementation outlined in Section 2.1.3 of Annis et al. (2019) [1]_. If method=”variance”, the error is estimated using the variance of the TI estimate across samples as in Eq. (37). If method=”coarse”, the error is estimated by comparing to a coarser integration using every other temperature, as in the original ptemcee implementation.

[1]

Annis, J., et al. https://doi.org/10.1016/j.jmp.2019.01.005

Parameters:

• burn_in_fraction (float | None) – Fraction of initial samples to discard as burn-in. If None, no burn-in is discarded. This is applied independently to each temperature chain before integration.

• method (str) – Method for estimating the uncertainty in the log evidence. Options are “variance” or “coarse”. See Notes for details.

Returns:

• log_evidence – The estimated log evidence from thermodynamic integration.

• log_evidence_error – An estimate of the uncertainty in the log evidence.

Return type:

tuple[float, float]

log_evidence_stepping_stone(burn_in_fraction=0.1)

Compute the log evidence using the stepping-stone estimator.

Notes

Follows the implementation outlined in Section 2.2.3 of Annis et al. (2019) [1]_

[1]

Annis, J., et al. https://doi.org/10.1016/j.jmp.2019.01.005

Parameters:

burn_in_fraction (float | None) – Fraction of initial samples to discard as burn-in. If None, no burn-in is discarded. This is applied independently to each temperature chain before integration.

Returns:

• log_evidence – The estimated log evidence from thermodynamic integration.

• log_evidence_error – An estimate of the uncertainty in the log evidence.

Return type:

float

plot_chain(beta_index, n_walkers=None, burn_in=0, parameters=None, fig=None, **kwargs)

Parameters:

• beta_index (int)

• n_walkers (int | None)

• burn_in (int)

• parameters (list[str] | None)

• fig (matplotlib.figure.Figure | None)

class aspire.samples.SMCSamples

Bases: BaseSamples

Class for storing samples and corresponding weights.

If xp is not specified, all inputs will be converted to match the array type of x.

beta: float | None = None

Temperature parameter for the current samples.

log_evidence: float | None = None

Log evidence estimate for the current samples.

log_evidence_error: float | None = None

Log evidence error estimate for the current samples.

log_p_t(beta)

unnormalized_log_weights(beta)

Parameters:

beta (float)

Return type:

array_api_compat.common._typing.Array

log_evidence_ratio(beta)

Parameters:

beta (float)

Return type:

float

log_evidence_ratio_variance(beta)

Estimate the variance of the log evidence ratio using the delta method.

Defined as Var(log Z) = Var(w) / (E[w])^2 where w are the unnormalized weights.

Parameters:

beta (float)

Return type:

float

log_weights(beta)

Parameters:

beta (float)

Return type:

array_api_compat.common._typing.Array

resample(beta, n_samples=None, rng=None)

Resample the samples to a new inverse temperature beta.

If beta is the same as the current beta and n_samples is None or equal to the current number of samples, returns the same samples with updated beta.

Parameters:

• n_samples (int | None)

• rng (numpy.random.Generator)

Return type:

SMCSamples

to_standard_samples()

Convert the samples to standard samples.

to_numpy()