Enhance stability of minimization

[jfennick]

Description

I have been experiencing some instabilities which, after much difficulty, I was able to track down to the minimizer. The observed behavior is that the minimizer will finish 'successfully', but some time later, during equilibration and/or propagating replicas, the simulation may or may not randomly crash with NaNs or even throw CUDA_ERROR_ILLEGAL_ADDRESS (700).

Nearly identical behavior can be found here https://github.com/openmm/openmm/issues/3414. In that ticket, the issue was that the box vectors were not being updated correctly. For stability, during minimization the box vectors should not be allowed to change at all. Subsequent equilibration at NPT should be used to adjust the box vectors.

In my case(s), the box vectors were changing only at the fourth significant figure, and yet that was sufficient to (usually) crash the simulations. Inserting Gradient Descent before FIRE seems to greatly ameliorate but not completely eliminate the instabilities. Replacing FIRE with L-BFGS (with or without Gradient Descent) also seems to fix the instabilities.

This PR adds three changes to improve the stability of minimization:

1. First and foremost, it temporarily disables the barostat during minimization to keep the box vectors fixed. This is probably the only change that is absolutely necessary to fix the instabilities.
2. It inserts a very short (24 timestep) Gradient Descent minimization before the main FIRE minimization. Gradient Descent is not asymptotically fast, but it is excellent for preconditioning a 'better' minimizer.
3. I have replaced FIRE with L-BFGS. Unlike FIRE, L-BFGS does not modify the box vectors (even when the pressure is not None) and it's a fine minimizer anyway.

Todos

• [ ] Implement feature / fix bug
• [ ] Add tests
• [ ] Update documentation as needed
• [ ] Update changelogNotable points that this PR has either accomplished or will accomplish.

Status

• [ ] Ready to go

[jfennick]

See also https://github.com/choderalab/yank/pull/1267

[mikemhenry]

@jfennick Thanks for this! I'll have @jchodera take a look!

[jfennick]

FYI I just realized that I accidentally named this branch stable_equilibration instead of stable_minimization.

[jfennick]

I'm not sure why this is failing CI. pytest openmmtools/tests/test_sampling.py -k 'test_minimize' works on my machine.

[jfennick]

Are there any thoughts on this PR? Again, disabling the barostat is really the only thing that is critical, which is why I originally separated it out into it's own 3-line commit. Perhaps it was a mistake to add additional enhancements. Can we at least merge the pressure bug in the first commit?

[ijpulidos]

@jfennick Thanks for your contributions. I'm tagging @jchodera for him to take another look into these, sine we need to see if the previously raised concerns are solved. I will also review this in detail, since this can potentially change many things in our calculations. For now I'm marking this to be released on our 0.22 milestone.

[ijpulidos]

Resolves #668