jexplore.steps.stretch

jexplore.steps.stretch#

This module define the class for a MH step based on stretch proposal.

Classes#

Stretch

Class implementing a MH steps based on stretch proposal

Module Contents#

class Stretch[Tepoch: jexplore.sampling.EpochMH, Tstate: jexplore.sampling.StateMH, Tsampling: jexplore.sampling.SamplingMH](a=2.0, ngroups=2, permute=False)[source]#

Bases: jexplore.steps.colored.ColoredSC[Tepoch, Tstate, Tsampling]

Class implementing a MH steps based on stretch proposal

Parameters:

  • a (float) – stretch proposal a parameter

  • ngroups (int) – number of groups. Default 2.

  • permute (bool) – if true walkers are permuted at each iteration.

npart: int = 1#

number of partners needed to build the proposal

a: float#

stretch proposal a parameter

sa#
sample_z(key, size)[source]#

Sample the z distribution

Parameters:

  • key (jax.Array) – PRNG key

  • size (int) – output size

Returns:

samples

Return type:

jax.Array

proposal(key, state, group, cgroup)[source]#

Propose a new state according to the stretch proposal algorithm.

Parameters:

  • key (jax.Array) – PRNG key

  • state (Tstate) – current state

  • group (jax.Array) – indexes of the color chains

  • cgroup (jax.Array) – indexes of the complementary chains for this color.

Returns:

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

Return type:

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