jexplore.steps.mh
=================

.. py:module:: jexplore.steps.mh

.. autoapi-nested-parse::

   This modules defines the classes for Metropolis-Hastings
   steps.



Classes
-------

.. autoapisummary::

   jexplore.steps.mh.MHStep
   jexplore.steps.mh.AllChains


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

.. py:class:: MHStep[Tepoch: jexplore.sampling.EpochMH, Tstate: jexplore.sampling.StateMH, Tsampling: jexplore.sampling.SamplingMH]

   Bases: :py:obj:`jexplore.steps.step.Step`\ [\ :py:obj:`Tepoch`\ , :py:obj:`Tstate`\ , :py:obj:`Tsampling`\ ]


   Base class for MH and Model selection step. It specializes
   `build` method to define a suitable array of :math:`\beta_i = 1/T_i`.
   It also defines a generic MH acceptance step (the `mh` method).


   .. py:attribute:: beta
      :type:  jax.Array

      Array of temperatures inverses for all chains. Shape is (nchains, 1)


   .. py:method:: build(epoch)

      Step epoch initialisation method. Extend :py:attr:`jexplore.steps.step.Step.build`
      by populating the `betas` attribute.

      :param epoch: current epoch.



   .. py:method:: compute(state)

      Compute the loglik and logprior values of a state
      :param state: input state

      :return: new state with populated loglik and logprior values.



   .. py:method:: mh(key, state, prop, betas, qxy)

      Metropolis-Hastings acceptance step

      :param key: PRNG key
      :param state: current state
      :param prop: proposed state
      :param betas: (nchains, 1) betas array
      :param qxy: transition probabilities (nchains, 1)

      :return: new state with accepted changes and boolean mask of the changed chains.



   .. py:method:: step(key, state)
      :abstractmethod:


      Step sampling method. This is just a prototype.

      :param key: PRNG key
      :param state: current state

      :return: new state and the boolean mask of the chains modified by the step.



.. py:class:: AllChains[Tepoch: jexplore.sampling.EpochMH, Tstate: jexplore.sampling.StateMH, Tsampling: jexplore.sampling.SamplingMH](mask = None)

   Bases: :py:obj:`MHStep`\ [\ :py:obj:`Tepoch`\ , :py:obj:`Tstate`\ , :py:obj:`Tsampling`\ ]


   Full parallel all chains MH or modsel step.

   :param mask: proposal dimensions mask (default all space)


   .. py:attribute:: mask
      :type:  jax.Array

      List of indices that the proposal will act on


   .. py:attribute:: masked_covs
      :type:  jax.Array

      Masked version of the covariance


   .. py:method:: build(epoch)

      Step epoch initialisation method. Extend :py:attr:`jexplore.steps.step.Step.build`
      by populating the `betas` attribute.

      :param epoch: current epoch.



   .. py:method:: proposal(key, state)
      :abstractmethod:


      All chain proposal.

      :param key: PRNG key
      :param state: current state

      :return: the new PRNG key, the proposed state, the transition log probability.



   .. py:method:: step(key, state)

      Metropolis Hasting step sampling method. It proposes a new state
      calling the `proposal` method and then performs a MH acceptance calling
      :py:attr:`jexplore.steps.mh.MHState.mh` method.

      :param key: PRNG key
      :param state: current state

      :return: new state and the boolean mask of the chains modified by the step.