jexplore.sampler
================

.. py:module:: jexplore.sampler

.. autoapi-nested-parse::

   Jexplore: Jax based markov sampler



Attributes
----------

.. autoapisummary::

   jexplore.sampler.logger
   jexplore.sampler.Tepoch
   jexplore.sampler.Tsampling
   jexplore.sampler.Tbackend
   jexplore.sampler.Tstate
   jexplore.sampler.StepBuilder


Classes
-------

.. autoapisummary::

   jexplore.sampler.Steps
   jexplore.sampler.JaxSampler


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

.. py:data:: logger

   logging instance

.. py:data:: Tepoch

   placeholder for epoch class

.. py:data:: Tsampling

   placeholder for samping definition class

.. py:data:: Tbackend

   placeholder for backend class

.. py:data:: Tstate

   placholder for state class

.. py:type:: StepBuilder
   :canonical: Callable[[Tepoch], StepSample[Tstate]]


   Prototype definition of a step builder.
   This is a function getting an epoch as argument and
   returning a :py:attr:`jexplore.sampling.epoch.StepSample` function.

.. py:class:: Steps

   Bases: :py:obj:`list`\ [\ :py:obj:`dict`\ [\ :py:obj:`StepBuilder`\ [\ :py:obj:`Tepoch`\ , :py:obj:`Tstate`\ ]\ , :py:obj:`float`\ ]\ ]


   Sampler steps class. This is a list of dictionnaries.
   Each element of the list is a sequential substep and the corresponding
   dictionnary keys are the builders of the alternative steps for that
   substep and the values are the weights.


   .. py:method:: epoch_cycle(epoch)

      Defines the epoch cycle. Each one of the `StepBuilder` dictionnary keys
      is called on with argument `epoch`, and replaced by the returned
      :py:attr:`jexplore.sampling.epoch.StepSample` instance.

      :param epoch: this epoch

      :return: the epoch cycle list of dictionnaries.



.. py:class:: JaxSampler

   Bases: :py:obj:`Generic`\ [\ :py:obj:`Tsampling`\ , :py:obj:`Tepoch`\ , :py:obj:`Tbackend`\ ]


   Main jexplore sampler class.

   :param sampling: sampling definition
   :param steps: steps definition
   :param backend: backend instance


   .. py:attribute:: sampling
      :type:  Tsampling

      sampling parameters definition


   .. py:attribute:: steps
      :type:  Steps

      steps definition


   .. py:attribute:: backend
      :type:  Tbackend

      backend instance


   .. py:method:: run(epoch, niters, nepoch = 1, seed = None)

      Runs MCMC, stores results in the backend and returns last epoch.

      :param epoch: starting epoch
      :param niters: number of iterations per epoch
      :param nepoch: number of epochs
      :param seed: seed. If None a random seed is generated.

      :return: last epoch.



   .. py:method:: run_epoch(epoch, niters, key)
      :staticmethod:


      Runs one MCMC epoch iterations  and returns the resulting Epoch object.

      :param epoch: last epoch
      :param int nstep: number of MCMC iterations.
      :param PRNGKey key: PRNGKey number

      :return: key and the new Epoch