Architectural overhaul to simplify handling of chains and swapping

[HarrisonWilde]

The main contributions of this PR are:

• Moving from the old chain_to_process and process_to_chain maps to instead keep chain i's states and samplers easily accessible via state.internal[i], sampler.internal[i] and only manage mappings for the inverse temperature (inverse_temperatures[j]) currently acting to temper chain i via chain_order[j] = i and the temperature itself being accessible via chain_to_inverse_temperature_map[i]. This simplifies the structure of the TemperedState significantly and should make parallelised and distributed work more simple in the future.
• In doing the above, we fix an old bug with the state to process maps, meaning incorrect samples were being recorded throughout the sampling process.
• Additionally, we filter "stale" samples that arise in the situation where chains are swapped and subsequently reject the first proposals made under their new temperatures. This leads to transitions from their previously associated inverse temperature steps being returned and associated with their new temperature, meaning crucially that our eventual target sample contained samples from tempered target steps, leading to an implied poor performance of the package due to this leakage. The work to identify and track "stale" samples mainly happens through the maintenance of an array of length inverse_temperatures (is_stale) keeping track of when this happens during a step.
• Merging the swap phase of a step with the step itself: rather than having some steps returned to AbstractMCMC.sample that correspond only to the result of swapping temperatures between chains, we carry out the check for whether we should_swap and do the swap itself as part of the AbstractMCMC.step and always progress the chains using the underlying exploratory samplers / kernels as well, this leads to a 1:1 correspondence between nsamples specified at sample time and the total number of steps taken by the samplers, which is desirable.
• Correct implementations of standard swap strategies SingleSwap is a more clear renaming of the previous StandardSwap, we now provide correct implementations of SingleRandomSwap, RandomSwap, ReversibleSwap and NonReversibleSwap in line with their definitions in the referenced literature. Reversible and Non-Reversible schemes are particularly important for reasonable baseline performance of a functional PT package, as they increase our robustness to untuned temperature ladders.
• Numerous TODO markers in the code have been resolved, including the handling of adaptation configuration in the sampler, the top level stuff is now in the sampler and the AdaptationState is only initiated and stored when the TemperedState is first initialised. Additionally, the flag to turn adaptation has been removed in favour of a NoAdapt() schedule to turn off all adaptation.
• Simplify the swapping.jl file to only support LogDensityProblems and equivalent implementations of MCMCTempering for now, mainly through removing the versions of functions like compute_tempered_likelihoods with extra unused arguments, this stuff may now be fully redundant given Turing.jl's support of LogDensityProblems. Additionally, update from k to i and j indexing to support a more general swap not necessarily between two neighbouring chains, in general this allows for clearer variable naming and reference.
• The README has been reworked, focussing mainly on the example to be in line with most up to date syntax and approaches, including a simple comparison of MH sampling on a GMM with and without PT, with visual aids.
• We now support the chain_type arg to directly return an MCMCChains.Chains object for example rather than making a user bundle_samples manually
• Consistency in some variable naming has been improved, documentation improved and made more clear in line with the new approach detailed at the start of this list.

PENDING TODO:

• Old testing suite does not work, I will fix this in the coming days but wanted to make the PR available now
• Adaptation is currently broken under this new approach, mainly due to use explicitly allowing swaps between non-neighbouring chains under the random permutation family of swap strategies. This could be fixed fairly easily.
• Possibly need some further documentation in the code to make get_fresh_transition as clear as possible, didn't want to over-explain that part but would appreciate second opinions.

FUTURE:

• Parallelisation (multi threading) may be trivially possible now with some changes to the underling data structures
• More QoL changes are possible, for example removing the dependency on ConcreteStructs

[yebai]

@HarrisonWilde Is this still needed, given recent changes?

[HarrisonWilde]

Yeah definitely not, just keeping around the draft as we incorporate the changes that it proposed (still not fully there)

On 21 Mar 2023, at 20:05, Hong Ge @.***> wrote:



@HarrisonWildehttps://github.com/HarrisonWilde Is this still needed, given recent changes?

— Reply to this email directly, view it on GitHubhttps://github.com/TuringLang/MCMCTempering.jl/pull/140#issuecomment-1478512293, or unsubscribehttps://github.com/notifications/unsubscribe-auth/AIQUYBAZPH6CCXCKXJ4YPQTW5ICZFANCNFSM6AAAAAAU6F2I4A. You are receiving this because you were mentioned.Message ID: @.***>