aspire.samplers.smc.blackjax
============================

.. py:module:: aspire.samplers.smc.blackjax


Attributes
----------

.. autoapisummary::

   aspire.samplers.smc.blackjax.logger


Classes
-------

.. autoapisummary::

   aspire.samplers.smc.blackjax.BlackJAXSMC


Module Contents
---------------

.. py:data:: logger

.. py:class:: BlackJAXSMC(log_likelihood, log_prior, dims, prior_flow, xp, dtype=None, parameters=None, preconditioning_transform=None, rng = None)

   Bases: :py:obj:`aspire.samplers.smc.base.SMCSampler`


   BlackJAX SMC sampler.


   .. py:attribute:: key
      :value: None



   .. py:attribute:: rng


   .. py:method:: log_prob(x, beta=None)

      Log probability function compatible with BlackJAX.



   .. py:method:: sample(n_samples, n_steps = None, adaptive = True, target_efficiency = 0.5, target_efficiency_rate = 1.0, n_final_samples = None, sampler_kwargs = None, rng_key=None, checkpoint_callback=None, checkpoint_every = None, checkpoint_file_path = None, resume_from = None)

      Sample using BlackJAX SMC.

      :param n_samples: Number of samples to draw.
      :type n_samples: :py:class:`int`
      :param n_steps: Number of SMC steps.
      :type n_steps: :py:class:`int`
      :param adaptive: Whether to use adaptive tempering.
      :type adaptive: :py:class:`bool`
      :param target_efficiency: Target efficiency for adaptive tempering.
      :type target_efficiency: :py:class:`float`
      :param n_final_samples: Number of final samples to return.
      :type n_final_samples: :py:class:`int | None`
      :param sampler_kwargs: Additional arguments for the BlackJAX sampler.
                             - algorithm: str, one of "nuts", "hmc", "rwmh", "random_walk"
                             - n_steps: int, number of MCMC steps per mutation
                             - step_size: float, step size for HMC/NUTS
                             - inverse_mass_matrix: array, inverse mass matrix
                             - sigma: float or array, proposal covariance for random walk MH
                             - num_integration_steps: int, integration steps for HMC
      :type sampler_kwargs: :py:class:`dict | None`
      :param rng_key: JAX random key for reproducibility.
      :type rng_key: :py:class:`jax.random.key| None`



   .. py:method:: mutate(particles, beta, n_steps=None)

      Mutate particles using BlackJAX MCMC.