Fast Solver for Large Scale Multistate Bennett Acceptance Ratio Equations

Abstract

The multistate Bennett acceptance ratio method (MBAR) and unbinned weighted histogram analysis method (UWHAM) are widely employed approaches to calculate relative free energies of multiple thermodynamic states that gain statistical precision by employing free energy contributions from configurations sampled at each of the simulated λ states. With the increasing availability of high throughput computing resources, a large number of configurations can be sampled from hundreds or even thousands of states. Combining sampled configurations from all states to calculate relative free energies requires the iterative solution of large scale MBAR/UWHAM equations. In the current work, we describe the development of a fast solver to iteratively solve these large scale MBAR/UWHAM equations utilizing our previous findings that the MBAR/UWHAM equations can be derived as a Rao-Blackwell estimator. The solver is implemented and distributed as a Python module called FastMBAR. Our benchmark results show that FastMBAR is more than 2 times faster than the widely used solver pymbar, when it runs on a central processing unit (CPU) and more than 100 times faster than pymbar when it runs on a graphical processing unit (GPU). The significant speedup achieved by FastMBAR running on a GPU is useful not only for solving large scale MBAR/UWHAM equations but also for estimating uncertainty of calculated free energies using bootstrapping where the MBAR/UWHAM equations need to be solved multiple times.

Figure 1:

(A) Comparison of wall times required to solve MBAR/UWHAM equations of different sizes by pymbar and FastMBAR. Both pymbar and FastMBAR (CPU) were run on one Intel^® Xeon^® Processor E5520. The GPU used was a NVIDIA^® GEFORCE^® GTX 1080. (B) Number of iterations required by pymbar and FastMBAR (running on CPUs and GPUs) to solve MBAR/UWHAM equations to the same estimator precision (the relative reduction of objective functions is smaller than 1e-12). (C) Wall times required by pymbar, FastMBAR (CPU), and FastMBAR (GPU) to evaluate objective and gradient functions once. (D) Estimator precision changes with the number of iterations for pymbar, FastMBAR (CPU), and FastMBAR (GPU) when solving the MBAR/UWHAM equation with the energy matrix size of 100×12421. Similar plots for solving MBAR/UWHAM equations of other energy matrix sizes can be found in Fig. S1-S7.