aspire

aspire: Accelerated Sequential Posterior Inference via REuse

Submodules

• aspire.aspire
• aspire.flows
• aspire.history
• aspire.plot
• aspire.samplers
• aspire.samples
• aspire.transforms
• aspire.utils

Classes

Aspire	Accelerated Sequential Posterior Inference via REuse (aspire).
Samples	Class for storing samples and corresponding weights.

Package Contents

class aspire.Aspire(*, log_likelihood, log_prior, dims, parameters=None, periodic_parameters=None, prior_bounds=None, bounded_to_unbounded=True, bounded_transform='logit', device=None, xp=None, flow=None, flow_backend='zuko', flow_matching=False, eps=1e-06, dtype=None, **kwargs)

Accelerated Sequential Posterior Inference via REuse (aspire).

Parameters:

• log_likelihood (Callable) – The log likelihood function.

• log_prior (Callable) – The log prior function.

• dims (int) – The number of dimensions.

• parameters (list[str] | None) – The list of parameter names. If None, any samples objects will not have the parameters names specified.

• periodic_parameters (list[str] | None) – The list of periodic parameters.

• prior_bounds (dict[str, tuple[float, float]] | None) – The bounds for the prior. If None, some parameter transforms cannot be applied.

• bounded_to_unbounded (bool) – Whether to transform bounded parameters to unbounded ones.

• bounded_transform (str) – The transformation to use for bounded parameters. Options are ‘logit’, ‘exp’, or ‘tanh’.

• device (str | None) – The device to use for the flow. If None, the default device will be used. This is only used when using the PyTorch backend.

• xp (Callable | None) – The array backend to use. If None, the default backend will be used.

• flow (Flow | None) – The flow object, if it already exists. If None, a new flow will be created.

• flow_backend (str) – The backend to use for the flow. Options are ‘zuko’ or ‘flowjax’.

• flow_matching (bool) – Whether to use flow matching.

• eps (float) – The epsilon value to use for data transforms.

• dtype (Any | str | None) – The data type to use for the samples, flow and transforms.

• **kwargs – Keyword arguments to pass to the flow.

log_likelihood

log_prior

dims

parameters = None

device = None

eps = 1e-06

periodic_parameters = None

prior_bounds = None

bounded_to_unbounded = True

bounded_transform = 'logit'

flow_matching = False

flow_backend = 'zuko'

flow_kwargs

xp = None

dtype = None

property flow

The normalizing flow object.

property sampler: aspire.samplers.base.Sampler | None

The sampler object.

Return type:

aspire.samplers.base.Sampler | None

property n_likelihood_evaluations

The number of likelihood evaluations.

convert_to_samples(x, log_likelihood=None, log_prior=None, log_q=None, evaluate=True, xp=None)

Parameters:

evaluate (bool)

Return type:

aspire.samples.Samples

init_flow()

fit(samples, checkpoint_path=None, checkpoint_save_config=True, overwrite=False, **kwargs)

Fit the normalizing flow to the provided samples.

Parameters:

• samples (Samples) – The samples to fit the flow to.

• checkpoint_path (str | None) – Path to save the checkpoint. If None, no checkpoint is saved.

• checkpoint_save_config (bool) – Whether to save the Aspire configuration to the checkpoint.

• overwrite (bool) – Whether to overwrite an existing flow in the checkpoint file.

• kwargs (dict) – Keyword arguments to pass to the flow’s fit method.

Return type:

aspire.history.History

get_sampler_class(sampler_type)

Get the sampler class based on the sampler type.

Parameters:

sampler_type (str) – The type of sampler to use. Options are ‘importance’, ‘emcee’, or ‘smc’.

Return type:

Callable

init_sampler(sampler_type, preconditioning=None, preconditioning_kwargs=None, **kwargs)

Initialize the sampler for posterior sampling.

Parameters:

• sampler_type (str) – The type of sampler to use. Options are ‘importance’, ‘emcee’, or ‘smc’.

• preconditioning (str) – Type of preconditioning to apply in the sampler. Options are ‘default’, ‘flow’, or ‘none’.

• preconditioning_kwargs (dict) – Keyword arguments to pass to the preconditioning transform.

• kwargs (dict) – Keyword arguments to pass to the sampler.

Return type:

Callable

sample_posterior(n_samples=None, sampler='importance', xp=None, return_history=False, preconditioning=None, preconditioning_kwargs=None, checkpoint_path=None, checkpoint_every=1, checkpoint_save_config=True, **kwargs)

Draw samples from the posterior distribution.

If using a sampler that calls an external sampler, e.g. minipcn then keyword arguments for this sampler should be specified in sampler_kwargs. For example:

aspire = aspire(...)
aspire.sample_posterior(
    n_samples=1000,
    sampler="minipcn_smc",
    adaptive=True,
    sampler_kwargs=dict(
        n_steps=100,
        step_fn="tpcn",
    )
)

Parameters:

• n_samples (int | None) – The number of sample to draw. If None, the behavior will depend on the sampler.

• sampler (str) – Sampling algorithm to use for drawing the posterior samples.

• xp (Any) – Array API for the final samples.

• return_history (bool) – Whether to return the history of the sampler.

• preconditioning (str) – Type of preconditioning to apply in the sampler. Options are ‘default’, ‘flow’, or ‘none’. If not specified, the default will depend on the sampler being used. The importance sampler will default to ‘none’ and the other samplers to ‘default’

• preconditioning_kwargs (dict) – Keyword arguments to pass to the preconditioning transform.

• checkpoint_path (str | None) – Path to save the checkpoint. If None, no checkpoint is saved unless within an auto_checkpoint() context or a custom callback is provided.

• checkpoint_every (int) – Frequency (in number of sampler iterations) to save the checkpoint.

• checkpoint_save_config (bool) – Whether to save the Aspire configuration to the checkpoint.

• kwargs (dict) – Keyword arguments to pass to the sampler. These are passed automatically to the init method of the sampler or to the sample method.

Returns:

samples – Samples object contain samples and their corresponding weights.

Return type:

Samples

classmethod resume_from_file(file_path, *, log_likelihood, log_prior, sampler=None, checkpoint_path='checkpoint', checkpoint_dset='state', flow_path='flow', config_path='aspire_config', resume_kwargs=None)

Recreate an Aspire object from a single file and prepare to resume sampling.

Parameters:

• file_path (str) – Path to the HDF5 file containing config, flow, and checkpoint.

• log_likelihood (Callable) – Log-likelihood function (required, not pickled).

• log_prior (Callable) – Log-prior function (required, not pickled).

• sampler (str) – Sampler type to use (e.g., ‘smc’, ‘minipcn_smc’, ‘emcee_smc’). If None, will attempt to infer from saved config or checkpoint metadata.

• checkpoint_path (str) – HDF5 group path where the checkpoint is stored.

• checkpoint_dset (str) – Dataset name within the checkpoint group.

• flow_path (str) – HDF5 path to the saved flow.

• config_path (str) – HDF5 path to the saved Aspire config.

• resume_kwargs (dict | None) – Optional overrides to apply when resuming (e.g., checkpoint_every).

auto_checkpoint(path, every=1, save_config=True, save_flow=True, resume=False)

Context manager to auto-save checkpoints, config, and flow to a file.

Within the context, sample_posterior will default to writing checkpoints to the given path with the specified frequency, and will append config/flow after sampling.

Parameters:

• path (str) – Path to save the checkpoint file.

• every (int) – Frequency (in number of sampler iterations) to save the checkpoint.

• save_config (bool) – Whether to save the Aspire configuration to the checkpoint file.

• save_flow (bool) – Whether to save the flow to the checkpoint file.

• resume (bool) – Whether to attempt to resume from an existing checkpoint at the path.

enable_pool(pool, **kwargs)

Context manager to temporarily replace the log_likelihood method with a version that uses a multiprocessing pool to parallelize computation.

Parameters:

pool (multiprocessing.Pool) – The pool to use for parallel computation.

config_dict(include_sampler_config=False, **kwargs)

Return a dictionary with the configuration of the aspire object.

Parameters:

• include_sampler_config (bool) – Whether to include the configuration of the sampler. Default is False.

• kwargs (dict) – Additional keyword arguments to pass to the config_dict() method of the sampler.

Return type:

dict

save_config(h5_file, path='aspire_config', **kwargs)

Save the configuration to an HDF5 file.

Parameters:

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

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

• kwargs (dict) – Additional keyword arguments to pass to the config_dict() method.

Return type:

None

save_sampler_config(h5_file, path='sampler_config', **kwargs)

Save the sampler configuration to an HDF5 file.

By default, this will save the configuration of the last sampler used in a call to sample_posterior(). Set include_sample_class='all' to include all the calls to :py:meth:`sample_posterior and their corresponding sampler configurations.

Parameters:

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

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

• kwargs (dict) – Additional keyword arguments to pass to the config_dict() method of the sampler.

Return type:

None

save_flow(h5_file, path='flow')

Save the flow to an HDF5 file.

Parameters:

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

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

Return type:

None

load_flow(h5_file, path='flow')

Load the flow from an HDF5 file.

Parameters:

• h5_file (h5py.File) – The HDF5 file to load the flow from.

• path (str) – The path in the HDF5 file to load the flow from.

Return type:

None

save_config_to_json(filename)

Save the configuration to a JSON file.

Parameters:

filename (str)

Return type:

None

sample_flow(n_samples=1, xp=None)

Sample from the flow directly.

Includes the data transform, but does not compute log likelihood or log prior.

Parameters:

n_samples (int)

Return type:

aspire.samples.Samples

class aspire.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