Merge pull request #430 from choderalab/bar-compute-overlap

Fix #427 : Add BAR compute overlap method

pymbar/__init__.py

@@ -30,7 +30,7 @@
 
 from pymbar import timeseries, testsystems, confidenceintervals, version
 from pymbar.mbar import MBAR
-from pymbar.bar import BAR, BARzero
+from pymbar.bar import BAR, BARoverlap, BARzero
 from pymbar.exp import EXP, EXPGauss
 import pymbar.old_mbar
 

pymbar/bar.py

@@ -147,7 +147,7 @@ def BARzero(w_F, w_R, DeltaF):
     return fzero
 
 
-def BAR(w_F, w_R, DeltaF=0.0, compute_uncertainty=True, uncertainty_method='BAR',maximum_iterations=500, relative_tolerance=1.0e-12, verbose=False, method='false-position', iterated_solution=True, return_dict=False):
+def BAR(w_F, w_R, DeltaF=0.0, compute_uncertainty=True, uncertainty_method='BAR', maximum_iterations=500, relative_tolerance=1.0e-12, verbose=False, method='false-position', iterated_solution=True, return_dict=False):
     """Compute free energy difference using the Bennett acceptance ratio (BAR) method.
 
     Parameters
@@ -174,7 +174,7 @@ def BAR(w_F, w_R, DeltaF=0.0, compute_uncertainty=True, uncertainty_method='BAR'
     method : str, optional, defualt='false-position'
         choice of method to solve BAR nonlinear equations, one of 'self-consistent-iteration' or 'false-position' (default: 'false-position')
     iterated_solution : bool, optional, default=True
-        whether to fully solve the optimized BAR equation to consistency, or to stop after one step, to be 
+        whether to fully solve the optimized BAR equation to consistency, or to stop after one step, to be
         equivalent to transition matrix sampling.
     return_dict : bool, default False
         If true, returns are a dict, else they are a tuple
@@ -242,7 +242,7 @@ def BAR(w_F, w_R, DeltaF=0.0, compute_uncertainty=True, uncertainty_method='BAR'
             # consider returning more information about failure
             print("Warning: BAR is likely to be inaccurate because of poor overlap. Improve the sampling, or decrease the spacing betweeen states.  For now, guessing that the free energy difference is 0 with no uncertainty.")
             if compute_uncertainty:
-                result_vals['Delta_f'] = 0.0 
+                result_vals['Delta_f'] = 0.0
                 result_vals['dDelta_f'] = 0.0
                 if return_dict:
                     return result_vals
@@ -504,6 +504,38 @@ def BAR(w_F, w_R, DeltaF=0.0, compute_uncertainty=True, uncertainty_method='BAR'
             return result_vals
         return DeltaF
 
+def BARoverlap(w_F, w_R):
+    """Compute overlap between foward and backward ensembles (using MBAR definition of overlap)
+
+    Parameters
+    ----------
+    w_F : np.ndarray
+        w_F[t] is the forward work value from snapshot t.
+        t = 0...(T_F-1)  Length T_F is deduced from vector.
+    w_R : np.ndarray
+        w_R[t] is the reverse work value from snapshot t.
+        t = 0...(T_R-1)  Length T_R is deduced from vector.
+
+    Returns
+    -------
+    overlap : float
+        The overlap: 0 denotes no overlap, 1 denotes complete overlap
+
+    """
+    from pymbar import MBAR
+    N_k = np.array( [len(w_F), len(w_R)] )
+    N = N_k.sum()
+    u_kn = np.zeros([2,N], np.float32)
+    u_kn[1,0:N_k[0]] = w_F[:]
+    u_kn[0,N_k[0]:N] = w_R[:]
+    mbar = MBAR(u_kn, N_k)
+
+    # Check to make sure u_kn has been correctly formed
+    BAR_DF, BAR_dDF = BAR(w_F, w_R, return_dict=False)
+    assert numpy.isclose(mbar.f_k[1] - mbar.f_k[0], BAR_DF), f'BAR: {BAR_DF} +- {BAR_dDF} | MBAR: {mbar.f_k[1] - mbar.f_k[0]}'
+
+    return mbar.computeOverlap()['scalar']
+
 #=============================================================================================
 # For compatibility with 2.0.1-beta
 #=============================================================================================

pymbar/tests/test_bar.py

@@ -0,0 +1,84 @@
+"""Test BAR by performing statistical tests on a set of model systems
+for which the true free energy differences can be computed analytically.
+"""
+
+import numpy as np
+from pymbar import BAR, BARoverlap, MBAR
+from pymbar.testsystems import harmonic_oscillators, exponential_distributions
+from pymbar.utils import ensure_type
+from pymbar.utils_for_testing import eq
+
+precision = 5 # the precision for systems that do have analytical results that should be matched.
+z_scale_factor = 12.0  # Scales the z_scores so that we can reject things that differ at the ones decimal place.  TEMPORARY HACK
+#0.5 is rounded to 1, so this says they must be within 3.0 sigma
+N_k = np.array([1000, 500])
+
+def generate_ho(O_k = np.array([1.0, 2.0]), K_k = np.array([0.5, 1.0])):
+    return "Harmonic Oscillators", harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, K_k)
+
+def generate_exp(rates = np.array([1.0, 2.0])): # Rates, e.g. Lambda
+    return "Exponentials", exponential_distributions.ExponentialTestCase(rates)
+
+def convert_to_differences(x_ij,dx_ij,xa):
+    xa_ij = np.array(np.matrix(xa) - np.matrix(xa).transpose())
+
+    # add ones to the diagonal of the uncertainties, because they are zero
+    for i in range(len(N_k)):
+        dx_ij[i,i] += 1
+    z = (x_ij - xa_ij) / dx_ij
+    for i in range(len(N_k)):
+        z[i,i] = x_ij[i,i]-xa_ij[i,i]  # these terms should be zero; so we only throw an error if they aren't
+    return z
+
+system_generators = [generate_ho, generate_exp]
+observables = ['position', 'position^2', 'RMS deviation', 'potential energy']
+
+def test_bar_free_energies():
+
+    """Can BAR calculate moderately correct free energy differences?"""
+
+    for system_generator in system_generators:
+        name, test = system_generator()
+        x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
+        eq(N_k, N_k_output)
+
+        i = 0 ; j = 1 ; i_indices = np.arange(0, N_k[0]) ; j_indices = np.arange(N_k[0], N_k[0]+N_k[1])
+        w_f = u_kn[j,i_indices] - u_kn[i,i_indices]
+        w_r = u_kn[i,j_indices] - u_kn[j,j_indices]
+        fe, fe_sigma = BAR(w_f, w_r)
+
+        # Compare with analytical
+        fe0 = test.analytical_free_energies()
+        fe0 = fe0[1] - fe0[0]
+
+        z = (fe - fe0) / fe_sigma
+        eq(z / z_scale_factor, 0*z, decimal=0)
+
+        # Compare with MBAR
+        mbar = MBAR(u_kn, N_k)
+        results = mbar.getFreeEnergyDifferences(return_dict=True)
+        fe_t, dfe_t = mbar.getFreeEnergyDifferences(return_dict=False)
+
+        eq(fe, fe_t[0,1])
+        eq(fe_sigma, dfe_t[0,1])
+
+
+def test_bar_computeOverlap():
+
+    for system_generator in system_generators:
+        name, test = system_generator()
+        x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
+        eq(N_k, N_k_output)
+
+        i = 0 ; j = 1 ; i_indices = np.arange(0, N_k[0]) ; j_indices = np.arange(N_k[0], N_k[0]+N_k[1])
+        w_f = u_kn[j,i_indices] - u_kn[i,i_indices]
+        w_r = u_kn[i,j_indices] - u_kn[j,j_indices]
+
+        # Compute overlap
+        bar_overlap = BARoverlap(w_f, w_r)
+
+        # Compare with MBAR
+        mbar = MBAR(u_kn, N_k)
+        mbar_overlap = mbar.computeOverlap()['scalar']
+
+        eq(bar_overlap, mbar_overlap)

pymbar/tests/test_mbar.py

@@ -119,8 +119,8 @@ def test_mbar_computeExpectations_position_differences():
         eq(N_k, N_k_output)
         mbar = MBAR(u_kn, N_k)
         results = mbar.computeExpectations(x_n, output = 'differences', return_dict=True)
-        mu_ij = results['mu'] 
-        sigma_ij = results['sigma'] 
+        mu_ij = results['mu']
+        sigma_ij = results['sigma']
 
         mu0 = test.analytical_observable(observable = 'position')
         z = convert_to_differences(mu_ij, sigma_ij, mu0)
@@ -137,7 +137,7 @@ def test_mbar_computeExpectations_position2():
         mbar = MBAR(u_kn, N_k)
         results = mbar.computeExpectations(x_n ** 2, return_dict=True)
         mu = results['mu']
-        sigma = results['sigma'] 
+        sigma = results['sigma']
         mu0 = test.analytical_observable(observable = 'position^2')
 
         z = (mu0 - mu) / sigma
@@ -154,7 +154,7 @@ def test_mbar_computeExpectations_potential():
         mbar = MBAR(u_kn, N_k)
         results = mbar.computeExpectations(u_kn, state_dependent = True, return_dict=True)
         mu = results['mu']
-        sigma = results['sigma'] 
+        sigma = results['sigma']
         mu0 = test.analytical_observable(observable = 'potential energy')
         print(mu)
         print(mu0)
@@ -198,11 +198,11 @@ def test_mbar_computeEntropyAndEnthalpy():
         mbar = MBAR(u_kn, N_k)
         results =  mbar.computeEntropyAndEnthalpy(u_kn, return_dict=True)
         f_t, df_t, u_t, du_t, s_t, ds_t = mbar.computeEntropyAndEnthalpy(u_kn, return_dict=False)
-        f_ij = results['Delta_f'] 
-        df_ij = results['dDelta_f'] 
-        u_ij = results['Delta_u'] 
-        du_ij = results['dDelta_u']  
-        s_ij = results['Delta_s'] 
+        f_ij = results['Delta_f']
+        df_ij = results['dDelta_f']
+        u_ij = results['Delta_u']
+        du_ij = results['dDelta_u']
+        s_ij = results['Delta_s']
         ds_ij = results['dDelta_s']
 
         eq(f_ij, f_t)