aspire ====== .. py:module:: aspire .. autoapi-nested-parse:: aspire: Accelerated Sequential Posterior Inference via REuse Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/aspire/aspire/index /autoapi/aspire/flows/index /autoapi/aspire/history/index /autoapi/aspire/plot/index /autoapi/aspire/samplers/index /autoapi/aspire/samples/index /autoapi/aspire/transforms/index /autoapi/aspire/utils/index Classes ------- .. autoapisummary:: aspire.Aspire aspire.Samples Package Contents ---------------- .. py:class:: 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). :param log_likelihood: The log likelihood function. :type log_likelihood: :py:class:`Callable` :param log_prior: The log prior function. :type log_prior: :py:class:`Callable` :param dims: The number of dimensions. :type dims: :py:class:`int` :param parameters: The list of parameter names. If None, any samples objects will not have the parameters names specified. :type parameters: :py:class:`list[str] | None` :param periodic_parameters: The list of periodic parameters. :type periodic_parameters: :py:class:`list[str] | None` :param prior_bounds: The bounds for the prior. If None, some parameter transforms cannot be applied. :type prior_bounds: :py:class:`dict[str`, :py:class:`tuple[float`, :py:class:`float]] | None` :param bounded_to_unbounded: Whether to transform bounded parameters to unbounded ones. :type bounded_to_unbounded: :py:class:`bool` :param bounded_transform: The transformation to use for bounded parameters. Options are 'logit', 'exp', or 'tanh'. :type bounded_transform: :py:class:`str` :param device: The device to use for the flow. If None, the default device will be used. This is only used when using the PyTorch backend. :type device: :py:class:`str | None` :param xp: The array backend to use. If None, the default backend will be used. :type xp: :py:class:`Callable | None` :param flow: The flow object, if it already exists. If None, a new flow will be created. :type flow: :py:class:`Flow | None` :param flow_backend: The backend to use for the flow. Options are 'zuko' or 'flowjax'. :type flow_backend: :py:class:`str` :param flow_matching: Whether to use flow matching. :type flow_matching: :py:class:`bool` :param eps: The epsilon value to use for data transforms. :type eps: :py:class:`float` :param dtype: The data type to use for the samples, flow and transforms. :type dtype: :py:class:`Any | str | None` :param \*\*kwargs: Keyword arguments to pass to the flow. .. py:attribute:: log_likelihood .. py:attribute:: log_prior .. py:attribute:: dims .. py:attribute:: parameters :value: None .. py:attribute:: device :value: None .. py:attribute:: eps :value: 1e-06 .. py:attribute:: periodic_parameters :value: None .. py:attribute:: prior_bounds :value: None .. py:attribute:: bounded_to_unbounded :value: True .. py:attribute:: bounded_transform :value: 'logit' .. py:attribute:: flow_matching :value: False .. py:attribute:: flow_backend :value: 'zuko' .. py:attribute:: flow_kwargs .. py:attribute:: xp :value: None .. py:attribute:: dtype :value: None .. py:property:: flow The normalizing flow object. .. py:property:: sampler :type: aspire.samplers.base.Sampler | None The sampler object. .. py:property:: n_likelihood_evaluations The number of likelihood evaluations. .. py:method:: convert_to_samples(x, log_likelihood=None, log_prior=None, log_q=None, evaluate = True, xp=None) .. py:method:: init_flow() .. py:method:: fit(samples, checkpoint_path = None, checkpoint_save_config = True, overwrite = False, **kwargs) Fit the normalizing flow to the provided samples. :param samples: The samples to fit the flow to. :type samples: :py:class:`Samples` :param checkpoint_path: Path to save the checkpoint. If None, no checkpoint is saved. :type checkpoint_path: :py:class:`str | None` :param checkpoint_save_config: Whether to save the Aspire configuration to the checkpoint. :type checkpoint_save_config: :py:class:`bool` :param overwrite: Whether to overwrite an existing flow in the checkpoint file. :type overwrite: :py:class:`bool` :param kwargs: Keyword arguments to pass to the flow's fit method. :type kwargs: :py:class:`dict` .. py:method:: get_sampler_class(sampler_type) Get the sampler class based on the sampler type. :param sampler_type: The type of sampler to use. Options are 'importance', 'emcee', or 'smc'. :type sampler_type: :py:class:`str` .. py:method:: init_sampler(sampler_type, preconditioning = None, preconditioning_kwargs = None, **kwargs) Initialize the sampler for posterior sampling. :param sampler_type: The type of sampler to use. Options are 'importance', 'emcee', or 'smc'. :type sampler_type: :py:class:`str` :param preconditioning: Type of preconditioning to apply in the sampler. Options are 'default', 'flow', or 'none'. :type preconditioning: :py:class:`str` :param preconditioning_kwargs: Keyword arguments to pass to the preconditioning transform. :type preconditioning_kwargs: :py:class:`dict` :param kwargs: Keyword arguments to pass to the sampler. :type kwargs: :py:class:`dict` .. py:method:: 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. :code:`minipcn` then keyword arguments for this sampler should be specified in :code:`sampler_kwargs`. For example: .. code-block:: python aspire = aspire(...) aspire.sample_posterior( n_samples=1000, sampler="minipcn_smc", adaptive=True, sampler_kwargs=dict( n_steps=100, step_fn="tpcn", ) ) :param n_samples: The number of sample to draw. If None, the behavior will depend on the sampler. :type n_samples: :py:class:`int | None` :param sampler: Sampling algorithm to use for drawing the posterior samples. :type sampler: :py:class:`str` :param xp: Array API for the final samples. :type xp: :py:class:`Any` :param return_history: Whether to return the history of the sampler. :type return_history: :py:class:`bool` :param preconditioning: 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' :type preconditioning: :py:class:`str` :param preconditioning_kwargs: Keyword arguments to pass to the preconditioning transform. :type preconditioning_kwargs: :py:class:`dict` :param checkpoint_path: Path to save the checkpoint. If None, no checkpoint is saved unless within an :py:meth:`auto_checkpoint` context or a custom callback is provided. :type checkpoint_path: :py:class:`str | None` :param checkpoint_every: Frequency (in number of sampler iterations) to save the checkpoint. :type checkpoint_every: :py:class:`int` :param checkpoint_save_config: Whether to save the Aspire configuration to the checkpoint. :type checkpoint_save_config: :py:class:`bool` :param kwargs: Keyword arguments to pass to the sampler. These are passed automatically to the init method of the sampler or to the sample method. :type kwargs: :py:class:`dict` :returns: **samples** -- Samples object contain samples and their corresponding weights. :rtype: :py:class:`Samples` .. py:method:: 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) :classmethod: Recreate an Aspire object from a single file and prepare to resume sampling. :param file_path: Path to the HDF5 file containing config, flow, and checkpoint. :type file_path: :py:class:`str` :param log_likelihood: Log-likelihood function (required, not pickled). :type log_likelihood: :py:class:`Callable` :param log_prior: Log-prior function (required, not pickled). :type log_prior: :py:class:`Callable` :param sampler: Sampler type to use (e.g., 'smc', 'minipcn_smc', 'emcee_smc'). If None, will attempt to infer from saved config or checkpoint metadata. :type sampler: :py:class:`str` :param checkpoint_path: HDF5 group path where the checkpoint is stored. :type checkpoint_path: :py:class:`str` :param checkpoint_dset: Dataset name within the checkpoint group. :type checkpoint_dset: :py:class:`str` :param flow_path: HDF5 path to the saved flow. :type flow_path: :py:class:`str` :param config_path: HDF5 path to the saved Aspire config. :type config_path: :py:class:`str` :param resume_kwargs: Optional overrides to apply when resuming (e.g., checkpoint_every). :type resume_kwargs: :py:class:`dict | None` .. py:method:: 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. :param path: Path to save the checkpoint file. :type path: :py:class:`str` :param every: Frequency (in number of sampler iterations) to save the checkpoint. :type every: :py:class:`int` :param save_config: Whether to save the Aspire configuration to the checkpoint file. :type save_config: :py:class:`bool` :param save_flow: Whether to save the flow to the checkpoint file. :type save_flow: :py:class:`bool` :param resume: Whether to attempt to resume from an existing checkpoint at the path. :type resume: :py:class:`bool` .. py:method:: enable_pool(pool, **kwargs) Context manager to temporarily replace the log_likelihood method with a version that uses a multiprocessing pool to parallelize computation. :param pool: The pool to use for parallel computation. :type pool: :py:class:`multiprocessing.Pool` .. py:method:: config_dict(include_sampler_config = False, **kwargs) Return a dictionary with the configuration of the aspire object. :param include_sampler_config: Whether to include the configuration of the sampler. Default is False. :type include_sampler_config: :py:class:`bool` :param kwargs: Additional keyword arguments to pass to the :py:meth:`config_dict` method of the sampler. :type kwargs: :py:class:`dict` .. py:method:: save_config(h5_file, path='aspire_config', **kwargs) Save the configuration to an HDF5 file. :param h5_file: The HDF5 file to save the configuration to. :type h5_file: :py:class:`h5py.File` :param path: The path in the HDF5 file to save the configuration to. :type path: :py:class:`str` :param kwargs: Additional keyword arguments to pass to the :py:meth:`config_dict` method. :type kwargs: :py:class:`dict` .. py:method:: 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 :py:meth:`sample_posterior`. Set :code:`include_sample_class='all' to include all the calls to :py:meth:`sample_posterior` and their corresponding sampler configurations. :param h5_file: The HDF5 file to save the sampler configuration to. :type h5_file: :py:class:`h5py.File` :param path: The path in the HDF5 file to save the sampler configuration to. :type path: :py:class:`str` :param kwargs: Additional keyword arguments to pass to the :py:meth:`config_dict` method of the sampler. :type kwargs: :py:class:`dict` .. py:method:: save_flow(h5_file, path='flow') Save the flow to an HDF5 file. :param h5_file: The HDF5 file to save the flow to. :type h5_file: :py:class:`h5py.File` :param path: The path in the HDF5 file to save the flow to. :type path: :py:class:`str` .. py:method:: load_flow(h5_file, path='flow') Load the flow from an HDF5 file. :param h5_file: The HDF5 file to load the flow from. :type h5_file: :py:class:`h5py.File` :param path: The path in the HDF5 file to load the flow from. :type path: :py:class:`str` .. py:method:: save_config_to_json(filename) Save the configuration to a JSON file. .. py:method:: 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. .. py:class:: Samples Bases: :py:obj:`BaseSamples` Class for storing samples and corresponding weights. If :code:`xp` is not specified, all inputs will be converted to match the array type of :code:`x`. .. py:attribute:: log_w :type: array_api_compat.common._typing.Array .. py:attribute:: weights :type: array_api_compat.common._typing.Array .. py:attribute:: evidence :type: float .. py:attribute:: evidence_error :type: float .. py:attribute:: log_evidence :type: float | None :value: None .. py:attribute:: log_evidence_error :type: float | None :value: None .. py:attribute:: effective_sample_size :type: float .. py:property:: efficiency Efficiency of the weighted samples. Defined as ESS / number of samples. .. py:method:: compute_weights() Compute the posterior weights. .. py:property:: scaled_weights .. py:method:: rejection_sample(rng=None) .. py:method:: plot_corner(include_weights = True, **kwargs) Plot a corner plot of the samples. :param parameters: List of parameters to plot. If None, all parameters are plotted. Figure to plot on. If None, a new figure is created. :type parameters: :py:class:`list[str] | None` :param \*\*kwargs: Additional keyword arguments to pass to corner.corner(). Kwargs are deep-copied before use. :type \*\*kwargs: :py:class:`dict` .. py:method:: to_namespace(xp) .. py:method:: to_numpy() .. py:method:: 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. :param include: List of fields to include in the DataFrame. If None, includes log_likelihood, log_prior, log_q, and log_w. :type include: :py:class:`list[str] | None` :returns: A DataFrame representation of the samples. :rtype: :py:class:`pd.DataFrame`