Brainstorming ideas for speeding up YANK

[jchodera]

I'm just collecting some ideas to consider after the 1.0 release.

Overall optimization

• Optimize code based on profiling
• Non-cubic/rectangular boxes (truncated octahedron, dodecahedron)
• Tune nonbonded cutoffs to balance performance and overlap
• Optimize PME parameters to maximize performance given error threshold
• Use crude PME tolerances in intermediate states
• optimize alchemical pathway to minimize number of thermodynamic states

Startup time

• Use FIRE minimizer for minimizations instead of very slow LocalEnergyMinimizer
• Pre-equilibrate fully-interacting system

Propagation

• Adjust number of steps per alchemical state to minimize discrepancy between end times for propagation steps
• Use better integrators that permit larger timestep, such as g-BAOAB and other multiple-timestep algorithms that make use of hydrogen mass repartitioning, solute/solvent splitting, force splitting, etc. This may give a 2-6x speedup of propagation.
• Speed up ligand Monte Carlo using a CustomIntegrator or applying it only to weakly interacting replicas
• Use distributed computing framework that allows multiple YANK jobs to share same set of worker nodes and interleave free energy calculations
• Allow Monte Carlo rotation/translation to be disabled for Boresch restraints, since none will be accepted. (Perhaps ring flips could be used instead?) Would save ~16 s / 124s per iteration for Abl:imatinib (~13%)
• Consolidate sanity checks to shave off some time from _propagate_replicas.
• Have OpenMM use pairlists for interaction groups (Efficiency considerations in CustomNonbondedForce.addInteractionGroup? openmm/openmm#1765)

Energy computation

• Only compute energies for K neighbor states and use local methods from Z. Tan for state mixing and free energy estimates. Currently, energy computation takes 29s / 130s ~ 22%.
• Only compute energies for Forces with alchemical parameters (Possible optimization to _compute_replica_energies #678).

Writing configurations

• Have separate thread write data to disk
• Only write full system positions checkpoints to disk every N iterations; only resume from these points (for Abl:imatinib, writing currently consumes 42s / 141s ~ 30%). (Implemented in Checkpointing Feature #675)

[jlerche]

Re: separate thread writing to disk. Just need to be aware that due to the global interpreter lock threads created with the threading library aren't actually run in a separate os process, the multiprocessing library would have to be used to spawn a separate process, with the overhead of loading a copy of the yank python files into memory.

edit: well i'm fuzzy on what exactly gets copied into memory, whether it's the entire application or just the module that spawns the process.

[jchodera]

Thanks! I've used MPI processes in the past to do this kind of thing with Python, but there are now plenty of options for truly parallel processing that avoids the GIL.

[jchodera]

@Lnaden has implemented infrequent system checkpoints in #675!

[jchodera]

Switching to g-BAOAB with hydrogen mass repartitioning is probably our easiest win in terms of getting a big speed boost for little effort. Once @maxentile's assessment of HMR is complete in a couple of weeks, we'll have an idea of what the optimal settings are here.