Fast Solver for Large Scale Multistate Bennett Acceptance Ratio Equations

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.