[ENH] Expose swap attempts for ReplicaExchangeSampler

[IAlibay]

Exchanges are quite expensive, how would folks feel about the possibility of adding a control parameter for the number of swap attempts?

In the back of my mind I'm thinking that eventually it would be great to build some kind of adaptive approach that minimises the number of swap attempts for a target mixing ratio.

[ijpulidos]

Thanks for the feedback! That could indeed be useful to have, but can you share an example where the replica exchange/swap is taking too long? Because in my experience this part is pretty fast, so I wonder if we may not be reporting the times correctly, even though I double checked some of the logs I have and they seem consistent with this idea of mixing of replicas not taking long at all.

[IAlibay]

@ijpulidos - I'm currently debugging this but things seem to have changed a bit since yesterday. What we've found is:

1. There is a large performance drop after a certain threshold of swaps has been done (still checking exactly when, but it looks a lot like we're hitting something similar https://github.com/choderalab/perses/issues/613
2. As part of the above we've found that there seems to be no substantial difference between the numba and non-numba replica mixing loop. This might not be a problem though.
3. There are cases where the number of replica mixes are probably insufficient (n_windows**3 is a lot lower than what Vytas did in his own assessments for ABFEs for example).

[jchodera]

The rationale for the number of swaps we currently do is to ensure that we swap as much as possible without it taking a significant fraction of iteration time. A better approach might be to have a setting that specifies how many seconds to spend on swapping, and then run batches of swaps and check if the timer has exceeded the allowed time.

There should be a theoretical bound on the number of swaps needed to ensure thorough mixing, and I'd even attempted to engage the world's foremost probabilist of Markov chains on permutation spaces while I was at Stanford, but was never able to crack the problem.

We might also consider alternatives to numba, such as jax, that may offer good or better performance and better robustness.