jexplore.steps.single_masked

jexplore.steps.single_masked#

This modules defines the classes for Metropolis-Hastings steps.

Classes#

SingleMasked

Full parallel all chains MH or modsel step.

Module Contents#

class SingleMasked[Tepoch: jexplore.sampling.EpochMH, Tstate: jexplore.sampling.StateMH, Tsampling: jexplore.sampling.SamplingMH](mask=None)[source]#

Bases: jexplore.steps.mh.AllChains[Tepoch, Tstate, Tsampling]

Full parallel all chains MH or modsel step.

Parameters:

mask (jax.Array | None) – mask of the subspace on which this step si defined.

abstract sc_masked_proposal(key, p, ch)[source]#

Single chain masked proposal. This is a prototype.

Parameters:

  • key (jax.Array) – PRNG key

  • p (jax.Array) – single chain point

  • ch (jax.Array) – chain index

Returns:

the proposed chain point, the transition log probability.

Return type:

tuple[jax.Array, jax.Array]

proposal(key, state)[source]#

All chain proposal. This is a prototype.

Parameters:

  • key (jax.Array) – PRNG key

  • state (Tstate) – current state

Returns:

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

Return type:

tuple[jax.Array, Tstate, jax.Array]