Parallel tempering
ParallelTempering
¤
Bases: Strategy
Sample a tempered PDF with one exchange step. This is in essence closer to TakeSteps than global tuning. Considering the tempered version of the PDF is only there to help with convergence, by default the extra information in temperature not equal to 1 is not saved.
There should be a version of this class that saves the extra information in the temperature not equal to 1, which could be used for other purposes such as diagnostics or training.
Methods:
| Name | Description |
|---|---|
__call__ |
Resources must contain: |
__init__ |
Args: |
__call__(rng_key: Key, resources: dict[str, Resource], initial_position: Float[Array, 'n_chains n_dims'], data: dict) -> tuple[Key, dict[str, Resource], Float[Array, 'n_chains n_dim']]
¤
Resources must contain
- TemperedPDF
- Local kernel
- A buffer holding the tempered positions
- A buffer holding the temperatures
This strategy has 3 main steps: 1. Sample from the tempered PDF using the local kernel for n_steps 2. Exchange the samples between the temperatures 3. Adapt the temperatures based on the acceptance rate
TODO: Add way to turn of temperature adaptation to maintain detail balance.
__init__(n_steps: int, tempered_logpdf_name: str, kernel_name: str, tempered_buffer_names: list[str], state_name: str, verbose: bool = False) -> None
¤
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
n_steps
|
int
|
Number of local kernel steps per temperature per call. |
required |
tempered_logpdf_name
|
str
|
Resource key for the
|
required |
kernel_name
|
str
|
Resource key for the local proposal kernel. |
required |
tempered_buffer_names
|
list[str]
|
Resource keys for the tempered-position buffer and the temperature buffer (in that order). |
required |
state_name
|
str
|
Resource key for the sampler
|
required |
verbose
|
bool
|
Enable debug logging. Defaults to False. |
False
|