def calcTension(energy_data, verbose=False):
dE1 = energy_data[:, 1] - energy_data[:, 0]
dE2 = energy_data[:, 2] - energy_data[:, 0]
BdE1 = dE1 / kTkJmol
BdE2 = dE2 / kTkJmol
nstates = 2
nframes = len(dE1)
u_kln = np.zeros([nstates, nstates, nframes], np.float64)
u_kln[0, 1, :] = BdE1
u_kln[1, 0, :] = BdE2
N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
for k in range(nstates):
[nequil, g, Neff_max] = timeseries.detectEquilibration(u_kln[k, k, :])
indices = timeseries.subsampleCorrelatedData(u_kln[k, k, :], g=g)
N_k[k] = len(indices)
u_kln[k, :, 0:N_k[k]] = u_kln[k, :, indices].T
if verbose:
print("...found {} uncorrelated samples out of {} total samples...".
format(N_k, nframes))
if verbose: print("=== Computing free energy differences ===")
mbar = MBAR(u_kln, N_k)
[DeltaF_ij, dDeltaF_ij, Theta_ij] = mbar.getFreeEnergyDifferences()
tension = DeltaF_ij[
0,
1] / da * 1e18 * kT #(in J/m^2). note da already has a factor of two for the two areas!
tensionError = dDeltaF_ij[0, 1] / da * 1e18 * kT
if verbose:
print('tension (pymbar): {} +/- {}N/m'.format(tension, tensionError))
return tension, tensionError
def test_mbar_computeMultipleExpectations():
"""Can MBAR calculate E(u_kn)??"""
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)
mbar = MBAR(u_kn, N_k)
A = np.zeros([2, len(x_n)])
A[0, :] = x_n
A[1, :] = x_n**2
state = 1
results = mbar.computeMultipleExpectations(A,
u_kn[state, :],
return_dict=True)
mu_t, sigma_t = mbar.computeMultipleExpectations(A,
u_kn[state, :],
return_dict=False)
mu = results['mu']
sigma = results['sigma']
eq(mu, mu_t)
eq(sigma, sigma_t)
mu0 = test.analytical_observable(observable='position')[state]
mu1 = test.analytical_observable(observable='position^2')[state]
z = (mu0 - mu[0]) / sigma[0]
eq(z / z_scale_factor, 0 * z, decimal=0)
z = (mu1 - mu[1]) / sigma[1]
eq(z / z_scale_factor, 0 * z, decimal=0)
def test_mbar_computePMF():
""" testing computePMF """
name, test = generate_ho()
x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
mbar = MBAR(u_kn,N_k)
#do a 1d PMF of the potential in the 3rd state:
refstate = 2
dx = 0.25
xmin = test.O_k[refstate] - 1
xmax = test.O_k[refstate] + 1
within_bounds = (x_n >= xmin) & (x_n < xmax)
bin_centers = dx*np.arange(np.int(xmin/dx),np.int(xmax/dx)) + dx/2
bin_n = np.zeros(len(x_n),int)
bin_n[within_bounds] = 1 + np.floor((x_n[within_bounds]-xmin)/dx)
# 0 is reserved for samples outside the domain. We will ignore this state
range = np.max(bin_n)+1
results = mbar.computePMF(u_kn[refstate,:], bin_n, range, uncertainties = 'from-specified', pmf_reference = 1)
f_i = results['f_i']
df_i = results['df_i']
f0_i = 0.5*test.K_k[refstate]*(bin_centers-test.O_k[refstate])**2
f_i, df_i = f_i[2:], df_i[2:] # first state is ignored, second is zero, with zero uncertainty
normf0_i = f0_i[1:] - f0_i[0] # normalize to first state
z = (f_i - normf0_i) / df_i
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def sqdeltaW(self, mu_VLE, eps_scaled):
'''
Computes the square difference between the sum of the weights in the
vapor and liquid phases.
Stores the optimal reduced free energy as f_k_guess for future iterations
Stores mbar, sumWliq, and sumWvap for computing VLE properties if converged
'''
nTsim, U_flat, Nmol_flat, Ncut, f_k_guess, Temp_VLE, u_kn_all, N_k_all = self.nTsim, self.U_flat, self.Nmol_flat, self.Ncut, self.f_k_guess, self.Temp_VLE, self.u_kn_all, self.N_k_all
for jT, (Temp, mu) in enumerate(zip(Temp_VLE, mu_VLE)):
u_kn_all[nTsim + jT, :] = self.U_to_u(eps_scaled * U_flat, Temp,
mu, Nmol_flat)
mbar = MBAR(u_kn_all, N_k_all, initial_f_k=f_k_guess)
sumWliq = np.sum(mbar.W_nk[:, nTsim:][Nmol_flat > Ncut], axis=0)
sumWvap = np.sum(mbar.W_nk[:, nTsim:][Nmol_flat <= Ncut], axis=0)
sqdeltaW_VLE = (sumWliq - sumWvap)**2
### Store previous solutions to speed-up future convergence of MBAR
Deltaf_ij = mbar.getFreeEnergyDifferences(return_theta=False)[0]
self.f_k_guess = Deltaf_ij[0, :]
self.mbar, self.sumWliq, self.sumWvap = mbar, sumWliq, sumWvap
return sqdeltaW_VLE
def calc_df(u_kln):
"""
u_kln should be (nstates) x (nstates) x (nframes)
note that u_kln should be normalized by kT already
where each element is
a config from frame `n` of a trajectory conducted with state `k`
with energy recalculated using parameters of state `l`
"""
dims = u_kln.shape
if dims[0] != dims[1]:
raise ValueError(
"dimensions {} of u_kln should be square in the first two indices".
format(dims))
nstates = dims[0]
N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
for k in range(nstates):
[nequil, g, Neff_max] = timeseries.detectEquilibration(u_kln[k, k, :])
indices = timeseries.subsampleCorrelatedData(u_kln[k, k, :], g=g)
N_k[k] = len(indices)
u_kln[k, :, 0:N_k[k]] = u_kln[k, :, indices].T
# Compute free energy differences and statistical uncertainties
mbar = MBAR(u_kln, N_k)
[DeltaF_ij, dDeltaF_ij, Theta_ij] = mbar.getFreeEnergyDifferences()
# save data?
return DeltaF_ij, dDeltaF_ij
def test_mbar_computePerturbedFreeEnergeies():
""" testing computePerturbedFreeEnergies """
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')
numN = np.sum(N_k[:2])
mbar = MBAR(u_kn[:2, :numN],
N_k[:2]) # only do MBAR with the first and last set
results = mbar.computePerturbedFreeEnergies(u_kn[2:, :numN],
return_dict=True)
f_t, df_t = mbar.computePerturbedFreeEnergies(u_kn[2:, :numN],
return_dict=False)
fe = results['Delta_f']
fe_sigma = results['dDelta_f']
eq(fe, f_t)
eq(fe_sigma, df_t)
fe, fe_sigma = fe[0, 1:], fe_sigma[0, 1:]
print(fe, fe_sigma)
fe0 = test.analytical_free_energies()[2:]
fe0 = fe0[1:] - fe0[0]
z = (fe - fe0) / fe_sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeEntropyAndEnthalpy():
"""Can MBAR calculate f_k, <u_k> and s_k ??"""
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)
mbar = MBAR(u_kn, N_k)
results = mbar.computeEntropyAndEnthalpy(u_kn)
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']
fa = test.analytical_free_energies()
ua = test.analytical_observable('potential energy')
sa = test.analytical_entropies()
fa_ij = np.array(np.matrix(fa) - np.matrix(fa).transpose())
ua_ij = np.array(np.matrix(ua) - np.matrix(ua).transpose())
sa_ij = np.array(np.matrix(sa) - np.matrix(sa).transpose())
z = convert_to_differences(f_ij, df_ij, fa)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
z = convert_to_differences(u_ij, du_ij, ua)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
z = convert_to_differences(s_ij, ds_ij, sa)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
def test_mbar_computePMF():
""" testing computePMF """
name, test = generate_ho()
x_n, u_kn, N_k_output, s_n = test.sample(N_k, mode='u_kn')
mbar = MBAR(u_kn, N_k)
#do a 1d PMF of the potential in the 3rd state:
refstate = 2
dx = 0.25
xmin = test.O_k[refstate] - 1
xmax = test.O_k[refstate] + 1
within_bounds = (x_n >= xmin) & (x_n < xmax)
bin_centers = dx * np.arange(np.int(xmin / dx), np.int(xmax / dx)) + dx / 2
bin_n = np.zeros(len(x_n), int)
bin_n[within_bounds] = 1 + np.floor((x_n[within_bounds] - xmin) / dx)
# 0 is reserved for samples outside the domain. We will ignore this state
range = np.max(bin_n) + 1
results = mbar.computePMF(u_kn[refstate, :],
bin_n,
range,
uncertainties='from-specified',
pmf_reference=1)
f_i = results['f_i']
df_i = results['df_i']
f0_i = 0.5 * test.K_k[refstate] * (bin_centers - test.O_k[refstate])**2
f_i, df_i = f_i[2:], df_i[
2:] # first state is ignored, second is zero, with zero uncertainty
normf0_i = f0_i[1:] - f0_i[0] # normalize to first state
z = (f_i - normf0_i) / df_i
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def run_mbar(self, test_overlap = True):
r"""Runs MBAR free energy estimate """
MBAR_obj = MBAR(self._u_kln, self._N_k, verbose=True)
self._f_k = MBAR_obj.f_k
try:
(deltaF_ij, dDeltaF_ij, theta_ij) = MBAR_obj.getFreeEnergyDifferences()
except:
(deltaF_ij, dDeltaF_ij, theta_ij) = MBAR_obj.getFreeEnergyDifferences(return_theta=True)
self._deltaF_mbar = deltaF_ij[0, self._lambda_array.shape[0]-1]
self._dDeltaF_mbar = dDeltaF_ij[0, self._lambda_array.shape[0]-1]
self._pmf_mbar = numpy.zeros(shape=(self._lambda_array.shape[0], 3))
self._pmf_mbar[:, 0] = self._lambda_array
self._pmf_mbar[:, 1] = self._f_k
self._pmf_mbar[:,2] = dDeltaF_ij[0]
self._pairwise_F = numpy.zeros(shape=(self._lambda_array.shape[0]-1,4))
self._pairwise_F[:,0] = self._lambda_array[:-1]
self._pairwise_F[:,1] = self._lambda_array[1:]
self._pairwise_F[:,2] = numpy.diag(deltaF_ij,1)
self._pairwise_F[:,3] = numpy.diag(dDeltaF_ij,1)
##testing data overlap:
if test_overlap:
overlap_matrix = MBAR_obj.computeOverlap()
self._overlap_matrix = overlap_matrix[2]
def compute_mbar(self):
if 'method' in self.kwargs:
method = self.kwargs['method']
else:
method = 'adaptive'
if self.mbar_f_ki is not None:
self.mbar = MBAR(self.u_kln,
self.N_k,
verbose=self.verbose,
method=method,
initial_f_k=self.mbar_f_ki,
subsampling_protocol=[{
'method': 'L-BFGS-B',
'options': {
'disp': self.verbose
}
}],
subsampling=1)
else:
self.mbar = MBAR(self.u_kln,
self.N_k,
verbose=self.verbose,
method=method,
subsampling_protocol=[{
'method': 'L-BFGS-B',
'options': {
'disp': self.verbose
}
}],
subsampling=1)
self.mbar_ready = True
def build_MBAR_sim(self):
'''
Creates an instance of the MBAR object for just the simulated state points
N_k: contains the number of snapshots from each state point simulated
Nmol_kn: contains all of the Number of molecules in 1-d array
u_kn_sim: contains all the reduced potential energies just for the simulated points
f_k_sim: the converged reduced free energies for each simulated state point (used as initial guess for non-simulated state points)
'''
Temp_sim, mu_sim, nSnapshots, Nmol_flat, U_flat = self.Temp_sim, self.mu_sim, self.K_sim, self.Nmol_flat, self.U_flat
N_k_sim = np.array(nSnapshots)
sumN_k = np.sum(N_k_sim)
# Nmol_flat = np.array(N_data_sim).flatten()
# U_flat = np.array(U_data_sim).flatten()
u_kn_sim = np.zeros([len(Temp_sim), sumN_k])
for iT, (Temp, mu) in enumerate(zip(Temp_sim, mu_sim)):
u_kn_sim[iT] = self.U_to_u(U_flat, Temp, mu, Nmol_flat)
mbar_sim = MBAR(u_kn_sim, N_k_sim)
Deltaf_ij = mbar_sim.getFreeEnergyDifferences(return_theta=False)[0]
f_k_sim = Deltaf_ij[0, :]
# print(f_k_sim)
self.u_kn_sim, self.f_k_sim, self.sumN_k, self.N_k_sim, self.mbar_sim = u_kn_sim, f_k_sim, sumN_k, N_k_sim, mbar_sim
def test_mbar_computeOverlap():
# tests with identical states, which gives analytical results.
d = len(N_k)
even_O_k = 2.0*np.ones(d)
even_K_k = 0.5*np.ones(d)
even_N_k = 100*np.ones(d)
name, test = generate_ho(O_k = even_O_k, K_k = even_K_k)
x_n, u_kn, N_k_output, s_n = test.sample(even_N_k, mode='u_kn')
mbar = MBAR(u_kn, even_N_k)
overlap_scalar, eigenval, O = mbar.computeOverlap()
reference_matrix = np.matrix((1.0/d)*np.ones([d,d]))
reference_eigenvalues = np.zeros(d)
reference_eigenvalues[0] = 1.0
reference_scalar = np.float64(1.0)
eq(O, reference_matrix, decimal=precision)
eq(eigenval, reference_eigenvalues, decimal=precision)
eq(overlap_scalar, reference_scalar, decimal=precision)
# test of more straightforward examples
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')
mbar = MBAR(u_kn, N_k)
overlap_scalar, eigenval, O = mbar.computeOverlap()
# rows of matrix should sum to one
sumrows = np.array(np.sum(O,axis=1))
eq(sumrows, np.ones(np.shape(sumrows)), decimal=precision)
eq(eigenval[0], np.float64(1.0), decimal=precision)
def update_logZ_with_mbar(self):
"""
Use MBAR to update logZ estimates.
"""
if not self.ncfile:
raise Exception("Cannot update logZ using MBAR since no NetCDF file is storing history.")
if not self.sampler.update_scheme == 'global-jump':
raise Exception("Only global jump is implemented right now.")
if not self.ncfile:
raise Exception("Must have a storage file attached to use MBAR updates")
# Extract relative energies.
if self.verbose:
print('Updating logZ estimate with MBAR...')
initial_time = time.time()
from pymbar import MBAR
#first = int(self.iteration / 2)
first = 0
u_kn = np.array(self.ncfile.variables['u_k'][first:,:]).T
[N_k, bins] = np.histogram(self.ncfile.variables['state_index'][first:], bins=(np.arange(self.sampler.nstates+1) - 0.5))
mbar = MBAR(u_kn, N_k)
Deltaf_ij, dDeltaf_ij, Theta_ij = mbar.getFreeEnergyDifferences(compute_uncertainty=True, uncertainty_method='approximate')
self.logZ[:] = -mbar.f_k[:]
self.logZ -= self.logZ[0]
final_time = time.time()
elapsed_time = final_time - initial_time
self._timing['MBAR time'] = elapsed_time
if self.verbose:
print('MBAR time %8.3f s' % elapsed_time)
def test_mbar_computeEntropyAndEnthalpy():
"""Can MBAR calculate f_k, <u_k> and s_k ??"""
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)
mbar = MBAR(u_kn, N_k)
f_ij, df_ij, u_ij, du_ij, s_ij, ds_ij = mbar.computeEntropyAndEnthalpy(u_kn)
fa = test.analytical_free_energies()
ua = test.analytical_observable('potential energy')
sa = test.analytical_entropies()
fa_ij = np.array(np.matrix(fa) - np.matrix(fa).transpose())
ua_ij = np.array(np.matrix(ua) - np.matrix(ua).transpose())
sa_ij = np.array(np.matrix(sa) - np.matrix(sa).transpose())
z = convert_to_differences(f_ij,df_ij,fa)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
z = convert_to_differences(u_ij,du_ij,ua)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
z = convert_to_differences(s_ij,ds_ij,sa)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
def calc_abs_press_int(self,show_plot=True):
'''
Fits ln(Xi) with respect to N for low-density vapor
'''
Temp_sim, u_kn_sim,f_k_sim,sumN_k = self.Temp_sim, self.u_kn_sim,self.f_k_sim,self.sumN_k
nTsim, U_flat, Nmol_flat,Ncut = self.nTsim, self.U_flat, self.Nmol_flat, self.Ncut
Temp_IG = np.min(Temp_sim[self.mu_sim == self.mu_sim.min()])
# print(Temp_IG)
mu_IG = np.linspace(2.*self.mu_opt[self.Temp_VLE==Temp_IG],5.*self.mu_opt[self.Temp_VLE==Temp_IG],10)
N_k_all = self.K_sim[:]
N_k_all.extend([0]*len(mu_IG))
u_kn_IG = np.zeros([len(mu_IG),sumN_k])
u_kn_all = np.concatenate((u_kn_sim,u_kn_IG))
f_k_guess = np.concatenate((f_k_sim,np.zeros(len(mu_IG))))
for jT, mu in enumerate(mu_IG):
u_kn_all[nTsim+jT,:] = self.U_to_u(U_flat,Temp_IG,mu,Nmol_flat)
mbar = MBAR(u_kn_all,N_k_all,initial_f_k=f_k_guess)
sumW_IG = np.sum(mbar.W_nk[:,nTsim:][Nmol_flat<Ncut],axis=0)
Nmol_IG = np.sum(mbar.W_nk[:,nTsim:][Nmol_flat<Ncut].T*Nmol_flat[Nmol_flat<Ncut],axis=1)/sumW_IG
# print(sumW_IG,Nmol_IG)
# print(mbar.W_nk[:,nTsim:][Nmol_flat<Ncut].T)
# print(mbar.W_nk[:,nTsim:][Nmol_flat<Ncut].T*Nmol_flat[Nmol_flat<Ncut])
### Store previous solutions to speed-up future convergence of MBAR
Deltaf_ij = mbar.getFreeEnergyDifferences(return_theta=False)[0]
f_k_IG = Deltaf_ij[nTsim:,0]
# print(f_k_sim,f_k_guess[:nTsim+1],Deltaf_ij[0,:nTsim],f_k_IG)#,Nmol_IG,press_IG,Psat)
fit=stats.linregress(Nmol_IG[mu_IG<2.*self.mu_sim.min()],f_k_IG[mu_IG<2.*self.mu_sim.min()])
if show_plot:
Nmol_plot = np.linspace(Nmol_IG.min(),Nmol_IG.max(),50)
lnXi_plot = fit.intercept + fit.slope*Nmol_plot
plt.figure(figsize=[6,6])
plt.plot(Nmol_IG,f_k_IG,'bo',mfc='None',label='MBAR-GCMC')
plt.plot(Nmol_plot,lnXi_plot,'k-',label='Linear fit')
plt.xlabel('Number of Molecules')
plt.ylabel(r'$\ln(\Xi)$')
plt.legend()
plt.show()
print('Slope for ideal gas is 1, actual slope is: '+str(fit.slope))
print('Intercept for absolute pressure is:'+str(fit.intercept))
self.abs_press_int, self.Temp_IG, self.f_k_IG, self.Nmol_IG = fit.intercept, Temp_IG, f_k_IG, Nmol_IG
def run_mbar(self):
r"""Runs MBAR free energy estimate """
MBAR_obj = MBAR(self._u_kln, self._N_k, verbose=True)
self._f_k = MBAR_obj.f_k
(deltaF_ij, dDeltaF_ij, theta_ij) = MBAR_obj.getFreeEnergyDifferences()
self._deltaF_mbar = deltaF_ij[0, self._lambda_array.shape[0]-1]
self._dDeltaF_mbar = dDeltaF_ij[0, self._lambda_array.shape[0]-1]
self._pmf_mbar = np.zeros(shape=(self._lambda_array.shape[0], 2))
self._pmf_mbar[:, 0] = self._lambda_array
self._pmf_mbar[:, 1] = self._f_k
def _calc_singleLP_exTI_pred_mbar(Es, dEs, LPs_pred, nfr, T=300):
assert Es.shape == dEs.shape
kT = kb * T
mbar = MBAR(Es / kT, nfr)
w = mbar.getWeights()
wt = w.T
dgdl_int = []
for i in range(len(LPs_pred)):
dgdl_int.append(np.dot(wt[i], dEs[i]))
return np.array(dgdl_int)
def test_mbar_getWeights():
""" testing getWeights """
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')
mbar = MBAR(u_kn, N_k)
# rows should be equal to zero
W = mbar.getWeights()
sumrows = np.sum(W, axis=0)
eq(sumrows, np.ones(len(sumrows)), decimal=precision)
def test_mbar_getWeights():
""" testing getWeights """
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')
mbar = MBAR(u_kn, N_k)
# rows should be equal to zero
W = mbar.getWeights()
sumrows = np.sum(W,axis=0)
eq(sumrows, np.ones(len(sumrows)), decimal=precision)
def test_mbar_computeExpectations_position_differences():
"""Can MBAR calculate E(x_n)??"""
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)
mbar = MBAR(u_kn, N_k)
mu_ij, sigma_ij = mbar.computeExpectations(x_n, output='differences')
mu0 = test.analytical_observable(observable='position')
z = convert_to_differences(mu_ij, sigma_ij, mu0)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
def test_mbar_computeExpectationsInner():
"""Can MBAR calculate general expectations inner code (note: this just tests completion)"""
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)
mbar = MBAR(u_kn, N_k)
A_in = np.array([x_n, x_n**2, x_n**3])
u_n = u_kn[:2, :]
state_map = np.array([[0, 0], [1, 0], [2, 0], [2, 1]], int)
_ = mbar.computeExpectationsInner(A_in, u_n, state_map)
def test_mbar_computeExpectationsInner():
"""Can MBAR calculate general expectations inner code (note: this just tests completion)"""
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)
mbar = MBAR(u_kn, N_k)
A_in = np.array([x_n, x_n ** 2, x_n ** 3])
u_n = u_kn[:2,:]
state_map = np.array([[0,0],[1,0],[2,0],[2,1]],int)
[A_i, d2A_ij] = mbar.computeExpectationsInner(A_in, u_n, state_map)
def test_mbar_computeExpectations_position_averages():
"""Can MBAR calculate E(x_n)??"""
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)
mbar = MBAR(u_kn, N_k)
mu, sigma = mbar.computeExpectations(x_n)
mu0 = test.analytical_observable(observable='position')
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeExpectations_position_differences():
"""Can MBAR calculate E(x_n)??"""
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)
mbar = MBAR(u_kn, N_k)
mu_ij, sigma_ij = mbar.computeExpectations(x_n, output = 'differences')
mu0 = test.analytical_observable(observable = 'position')
z = convert_to_differences(mu_ij, sigma_ij, mu0)
eq(z / z_scale_factor, np.zeros(np.shape(z)), decimal=0)
def test_exponential_mbar_xkn_squared():
"""Exponential Distribution Test: can MBAR calculate E(x_kn^2)"""
test = exponential_distributions.ExponentialTestCase(rates)
x_kn, u_kln, N_k_output = test.sample(N_k)
eq(N_k, N_k_output)
mbar = MBAR(u_kln, N_k)
mu, sigma = mbar.computeExpectations(x_kn ** 2)
mu0 = test.analytical_x_squared()
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeExpectations_position_averages():
"""Can MBAR calculate E(x_n)??"""
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)
mbar = MBAR(u_kn, N_k)
mu, sigma = mbar.computeExpectations(x_n)
mu0 = test.analytical_observable(observable = 'position')
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeExpectations_potential():
"""Can MBAR calculate E(u_kn)??"""
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)
mbar = MBAR(u_kn, N_k)
mu, sigma = mbar.computeExpectations(u_kn, state_dependent=True)
mu0 = test.analytical_observable(observable='potential energy')
print(mu)
print(mu0)
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_exponential_mbar_xkn_squared():
"""Harmonic Oscillators Test: can MBAR calculate E(x_kn^2)??"""
test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
x_kn, u_kln, N_k_output = test.sample(N_k)
eq(N_k, N_k_output)
mbar = MBAR(u_kln, N_k)
mu, sigma = mbar.computeExpectations(x_kn ** 2)
mu0 = test.analytical_x_squared()
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeEffectiveSampleNumber():
""" testing computeEffectiveSampleNumber """
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)
mbar = MBAR(u_kn, N_k)
# one mathematical effective sample numbers should be between N_k and sum_k N_k
N_eff = mbar.computeEffectiveSampleNumber()
sumN = np.sum(N_k)
assert all(N_eff > N_k)
assert all(N_eff < sumN)
def test_mbar_computeEffectiveSampleNumber():
""" testing computeEffectiveSampleNumber """
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)
mbar = MBAR(u_kn, N_k)
# one mathematical effective sample numbers should be between N_k and sum_k N_k
N_eff = mbar.computeEffectiveSampleNumber()
sumN = np.sum(N_k)
assert all(N_eff > N_k)
assert all(N_eff < sumN)
def gather_dg(self, u_kln, nstates):
# Subsample data to extract uncorrelated equilibrium timeseries
N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
for k in range(nstates):
[_, g, __] = timeseries.detectEquilibration(u_kln[k, k, :])
indices = timeseries.subsampleCorrelatedData(u_kln[k, k, :], g=g)
N_k[k] = len(indices)
u_kln[k, :, 0:N_k[k]] = u_kln[k, :, indices].T
# Compute free energy differences and statistical uncertainties
mbar = MBAR(u_kln, N_k)
[DeltaF_ij, dDeltaF_ij, _] = mbar.getFreeEnergyDifferences()
print("Number of uncorrelated samples per state: {}".format(N_k))
return DeltaF_ij, dDeltaF_ij
def test_exponential_mbar_free_energies():
"""Exponential Distribution Test: can MBAR calculate correct free energy differences?"""
test = exponential_distributions.ExponentialTestCase(rates)
x_kn, u_kln, N_k_output = test.sample(N_k, mode='u_kln')
eq(N_k, N_k_output)
mbar = MBAR(u_kln, N_k)
fe, fe_sigma = mbar.getFreeEnergyDifferences()
fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]
fe0 = test.analytical_free_energies()
fe0 = fe0[1:] - fe0[0]
z = (fe - fe0) / fe_sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeExpectations_position2():
"""Can MBAR calculate E(x_n^2)??"""
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)
mbar = MBAR(u_kn, N_k)
results = mbar.computeExpectations(x_n**2, return_dict=True)
mu = results['mu']
sigma = results['sigma']
mu0 = test.analytical_observable(observable='position^2')
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeExpectations_potential():
"""Can MBAR calculate E(u_kn)??"""
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)
mbar = MBAR(u_kn, N_k)
mu, sigma = mbar.computeExpectations(u_kn, state_dependent = True)
mu0 = test.analytical_observable(observable = 'potential energy')
print(mu)
print(mu0)
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_harmonic_oscillators_mbar_free_energies():
"""Harmonic Oscillators Test: can MBAR calculate correct free energy differences?"""
test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
x_kn, u_kln, N_k_output = test.sample(N_k)
eq(N_k, N_k_output)
mbar = MBAR(u_kln, N_k)
fe, fe_sigma = mbar.getFreeEnergyDifferences()
fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]
fe0 = test.analytical_free_energies()
fe0 = fe0[1:] - fe0[0]
z = (fe - fe0) / fe_sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def initialize_MBAR(ncfile, u_kln=None, N_k=None):
"""
Initialize MBAR for Free Energy and Enthalpy estimates, this may take a while.
ncfile : NetCDF
Input YANK netcdf file
u_kln : array of numpy.float64, optional, default=None
Reduced potential energies of the replicas; if None, will be extracted from the ncfile
N_k : array of ints, optional, default=None
Number of samples drawn from each kth replica; if None, will be extracted from the ncfile
TODO
----
* Ensure that the u_kln and N_k are decorrelated if not provided in this function
"""
if u_kln is None or N_k is None:
(u_kln, N_k, u_n) = extract_ncfile_energies(ncfile)
# Initialize MBAR (computing free energy estimates, which may take a while)
logger.info("Computing free energy differences...")
mbar = MBAR(u_kln, N_k)
return mbar
def test_exponential_mbar_free_energies():
"""Exponential Distribution Test: can MBAR calculate correct free energy differences?"""
test = exponential_distributions.ExponentialTestCase(rates)
x_n, u_kn, origin = test.sample(N_k)
u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
eq(N_k, N_k_output.values)
mbar = MBAR(u_ijn.values, N_k)
fe, fe_sigma = mbar.getFreeEnergyDifferences()
fe, fe_sigma = fe[0], fe_sigma[0]
fe0 = test.analytical_free_energies()
z = (fe - fe0) / fe_sigma
z = z[1:] # First component is undetermined.
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def simulate_protocol(lambdas_k, n_samples_per_window=100, seed=None):
"""Generate samples from each lambda window, plug into MBAR"""
O_k, K_k = poorly_spaced_path(lambdas_k)
testsystem = HarmonicOscillatorsTestCase(O_k, K_k)
N_k = [n_samples_per_window] * len(O_k)
xs, u_kn, N_k, s_n = testsystem.sample(N_k, seed=seed)
mbar = MBAR(u_kn, N_k)
return mbar
def test_mbar_free_energies():
"""Can MBAR 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)
mbar = MBAR(u_kn, N_k)
fe, fe_sigma, Theta_ij = mbar.getFreeEnergyDifferences()
fe, fe_sigma = fe[0, 1:], fe_sigma[0, 1:]
fe0 = test.analytical_free_energies()
fe0 = fe0[1:] - fe0[0]
z = (fe - fe0) / fe_sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_harmonic_oscillators_mbar_free_energies():
"""Harmonic Oscillators Test: can MBAR calculate correct free energy differences?"""
test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
x_n, u_kn, origin = test.sample(N_k)
u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
eq(N_k, N_k_output.values)
mbar = MBAR(u_ijn.values, N_k)
fe, fe_sigma = mbar.getFreeEnergyDifferences()
fe, fe_sigma = fe[0], fe_sigma[0]
fe0 = test.analytical_free_energies()
z = (fe - fe0) / fe_sigma
z = z[1:] # First component is undetermined.
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_free_energies():
"""Can MBAR 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)
mbar = MBAR(u_kn, N_k)
fe, fe_sigma, Theta_ij = mbar.getFreeEnergyDifferences()
fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]
fe0 = test.analytical_free_energies()
fe0 = fe0[1:] - fe0[0]
z = (fe - fe0) / fe_sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computePerturbedFreeEnergeies():
""" testing computePerturbedFreeEnergies """
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')
numN = np.sum(N_k[:2])
mbar = MBAR(u_kn[:2,:numN], N_k[:2]) # only do MBAR with the first and last set
fe, fe_sigma = mbar.computePerturbedFreeEnergies(u_kn[2:,:numN])
fe, fe_sigma = fe[0,1:], fe_sigma[0,1:]
print(fe, fe_sigma)
fe0 = test.analytical_free_energies()[2:]
fe0 = fe0[1:] - fe0[0]
z = (fe - fe0) / fe_sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_exponential_mbar__xkn():
"""Harmonic Oscillators Test: can MBAR calculate E(x_kn)??"""
test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
x_n, u_kn, origin = test.sample(N_k)
u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
eq(N_k, N_k_output.values)
mbar = MBAR(u_ijn.values, N_k)
x_kn = convert_xn_to_x_kn(x_n)
x_kn = x_kn.values # Convert to numpy for MBAR
x_kn[np.isnan(x_kn)] = 0.0 # Convert nans to 0.0
mu, sigma = mbar.computeExpectations(x_kn)
mu0 = test.analytical_means()
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_mbar_computeMultipleExpectations():
"""Can MBAR calculate E(u_kn)??"""
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)
mbar = MBAR(u_kn, N_k)
A = np.zeros([2,len(x_n)])
A[0,:] = x_n
A[1,:] = x_n**2
state = 1
mu, sigma, covariances = mbar.computeMultipleExpectations(A,u_kn[state,:])
mu0 = test.analytical_observable(observable = 'position')[state]
mu1 = test.analytical_observable(observable = 'position^2')[state]
z = (mu0 - mu[0]) / sigma[0]
eq(z / z_scale_factor, 0*z, decimal=0)
z = (mu1 - mu[1]) / sigma[1]
eq(z / z_scale_factor, 0*z, decimal=0)
def test_exponential_mbar_xkn():
"""Exponential Distribution Test: can MBAR calculate E(x_kn)?"""
test = exponential_distributions.ExponentialTestCase(rates)
x_n, u_kn, origin = test.sample(N_k)
u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
eq(N_k, N_k_output.values)
mbar = MBAR(u_ijn.values, N_k)
x_kn = convert_xn_to_x_kn(x_n)
x_kn = x_kn.values # Convert to numpy for MBAR
x_kn[np.isnan(x_kn)] = 0.0 # Convert nans to 0.0
mu, sigma = mbar.computeExpectations(x_kn)
mu0 = test.analytical_means()
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_exponential_mbar_xkn_squared():
"""Exponential Distribution Test: can MBAR calculate E(x_kn^2)"""
test = exponential_distributions.ExponentialTestCase(rates)
x_n, u_kn, origin = test.sample(N_k)
u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
eq(N_k, N_k_output.values)
mbar = MBAR(u_ijn.values, N_k)
x_kn = convert_xn_to_x_kn(x_n) ** 2.0
x_kn = x_kn.values # Convert to numpy for MBAR
x_kn[np.isnan(x_kn)] = 0.0 # Convert nans to 0.0
mu, sigma = mbar.computeExpectations(x_kn)
mu0 = test.analytical_x_squared()
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def test_exponential_mbar_xkn_squared():
"""Harmonic Oscillators Test: can MBAR calculate E(x_kn^2)??"""
test = harmonic_oscillators.HarmonicOscillatorsTestCase(O_k, k_k)
x_n, u_kn, origin = test.sample(N_k)
u_ijn, N_k_output = convert_ukn_to_uijn(u_kn)
eq(N_k, N_k_output.values)
mbar = MBAR(u_ijn.values, N_k)
x_kn = convert_xn_to_x_kn(x_n) ** 2.
x_kn = x_kn.values # Convert to numpy for MBAR
x_kn[np.isnan(x_kn)] = 0.0 # Convert nans to 0.0
mu, sigma = mbar.computeExpectations(x_kn)
mu0 = test.analytical_x_squared()
z = (mu0 - mu) / sigma
eq(z / z_scale_factor, np.zeros(len(z)), decimal=0)
def gather_dg(self, u_kln, nstates):
u_kln = np.vstack(u_kln)
# Subsample data to extract uncorrelated equilibrium timeseries
N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
for k in range(nstates):
[_, g, __] = timeseries.detectEquilibration(u_kln[k, k, :])
indices = timeseries.subsampleCorrelatedData(u_kln[k, k, :], g=g)
N_k[k] = len(indices)
u_kln[k, :, 0:N_k[k]] = u_kln[k, :, indices].T
# Compute free energy differences and statistical uncertainties
mbar = MBAR(u_kln, N_k)
[DeltaF_ij, dDeltaF_ij, _] = mbar.getFreeEnergyDifferences()
logger.debug(
"Number of uncorrelated samples per state: {}".format(N_k))
logger.debug("Relative free energy change for {0} = {1} +- {2}".format(
self.name, DeltaF_ij[0, nstates - 1] * self.kTtokcal,
dDeltaF_ij[0, nstates - 1] * self.kTtokcal))
return DeltaF_ij[0, nstates -
1] * self.kTtokcal, dDeltaF_ij[0, nstates -
1] * self.kTtokcal
def test_mbar_computeOverlap():
# tests with identical states, which gives analytical results.
d = len(N_k)
even_O_k = 2.0 * np.ones(d)
even_K_k = 0.5 * np.ones(d)
even_N_k = 100 * np.ones(d)
name, test = generate_ho(O_k=even_O_k, K_k=even_K_k)
x_n, u_kn, N_k_output, s_n = test.sample(even_N_k, mode='u_kn')
mbar = MBAR(u_kn, even_N_k)
overlap_scalar, eigenval, O = mbar.computeOverlap()
reference_matrix = np.matrix((1.0 / d) * np.ones([d, d]))
reference_eigenvalues = np.zeros(d)
reference_eigenvalues[0] = 1.0
reference_scalar = np.float64(1.0)
eq(O, reference_matrix, decimal=precision)
eq(eigenval, reference_eigenvalues, decimal=precision)
eq(overlap_scalar, reference_scalar, decimal=precision)
# test of more straightforward examples
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')
mbar = MBAR(u_kn, N_k)
overlap_scalar, eigenval, O = mbar.computeOverlap()
# rows of matrix should sum to one
sumrows = np.array(np.sum(O, axis=1))
eq(sumrows, np.ones(np.shape(sumrows)), decimal=precision)
eq(eigenval[0], np.float64(1.0), decimal=precision)
def run_mbar(self, ndiscard=0, nuse=None):
"""Estimate free energies of all alchemical states.
Parameters
----------
ndiscard : int, optinoal, default=0
number of iterations to discard to equilibration
nuse : int, optional, default=None
maximum number of iterations to use (after discarding)
Returns
-------
Deltaf_ij : np.ndarray, shape=(n_states, n_states)
The statewise free energy differences
dDeltaf_ij : np.ndarray, shape=(n_states, n_states)
The statewise free energy difference uncertainties
"""
u_kln_replica, u_kln, u_n = self.get_u_kln()
u_kln_replica, u_kln, u_n, N_k, N = self.equilibrate_and_subsample(u_kln_replica, u_kln, u_n, ndiscard=ndiscard, nuse=nuse)
logger.info("Initialing MBAR and computing free energy differences...")
mbar = MBAR(u_kln, N_k, verbose = False, method = 'self-consistent-iteration', maximum_iterations = 50000) # use slow self-consistent-iteration (the default)
# Get matrix of dimensionless free energy differences and uncertainty estimate.
logger.info("Computing covariance matrix...")
(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='svd-ew')
logger.info("\n%-24s %16s\n%s" % ("Deltaf_ij", "current state", pd.DataFrame(Deltaf_ij).to_string()))
logger.info("\n%-24s %16s\n%s" % ("Deltaf_ij", "current state", pd.DataFrame(dDeltaf_ij).to_string()))
return (Deltaf_ij, dDeltaf_ij)
def MBAR_analysis(self, debug = False):
"""MBAR analysis for populations and BICePs score"""
# load necessary data first
self.load_data()
# Suppose the energies sampled from each simulation are u_kln, where u_kln[k,l,n] is the reduced potential energy
# of snapshot n \in 1,...,N_k of simulation k \in 1,...,K evaluated at reduced potential for state l.
self.K = self.nlambda # number of thermodynamic ensembles
# N_k[k] will denote the number of correlated snapshots from state k
N_k = np.array( [len(self.traj[i]['trajectory']) for i in range(self.nlambda)] )
nsnaps = N_k.max()
u_kln = np.zeros( (self.K, self.K, nsnaps) )
nstates = int(self.states)
print 'nstates', nstates
states_kn = np.zeros( (self.K, nsnaps) )
# Get snapshot energies rescored in the different ensembles
"""['step', 'E', 'accept', 'state', 'sigma_noe', 'sigma_J', 'sigma_cs', 'sigma_pf''gamma']
[int(step), float(self.E), int(accept), int(self.state), int(self.sigma_noe_index), int(self.sigma_J_index), int(self.sigma_cs_H_index), int(self.sigma_cs_Ha_index), int(self.sigma_cs_N_index), int(self.sigma_cs_Ca_index), int(self.sigma_pf_index), int(self.gamma_index)] """
for n in range(nsnaps):
for k in range(self.K):
for l in range(self.K):
if debug:
print 'step', self.traj[k]['trajectory'][n][0],
if k==l:
print 'E%d evaluated in model %d'%(k,k), self.traj[k]['trajectory'][n][1],
u_kln[k,k,n] = self.traj[k]['trajectory'][n][1]
state, sigma_noe_index, sigma_J_index, sigma_cs_H_index, sigma_cs_Ha_index, sigma_cs_N_index, sigma_cs_Ca_index, sigma_pf_index, gamma_index = self.traj[k]['trajectory'][n][3:] # IMPORTANT: make sure the order of these parameters is the same as the way they are saved in PosteriorSampler
print 'state, sigma_noe_index, sigma_J_index, sigma_cs_H_index, sigma_cs_Ha_index, sigma_cs_N_index, sigma_cs_Ca_index, sigma_pf_index, gamma_index', state, sigma_noe_index, sigma_J_index, sigma_cs_H_index, sigma_cs_Ha_index, sigma_cs_N_index, sigma_cs_Ca_index, sigma_pf_index, gamma_index
states_kn[k,n] = state
sigma_noe = self.traj[k]['allowed_sigma_noe'][sigma_noe_index]
sigma_J = self.traj[k]['allowed_sigma_J'][sigma_J_index]
sigma_cs_H = self.traj[k]['allowed_sigma_cs_H'][sigma_cs_H_index]
sigma_cs_Ha = self.traj[k]['allowed_sigma_cs_Ha'][sigma_cs_Ha_index]
sigma_cs_N = self.traj[k]['allowed_sigma_cs_N'][sigma_cs_N_index]
sigma_cs_Ca = self.traj[k]['allowed_sigma_cs_Ca'][sigma_cs_Ca_index]
sigma_pf = self.traj[k]['allowed_sigma_pf'][sigma_pf_index]
u_kln[k,l,n] = self.sampler[l].neglogP(0, state, sigma_noe, sigma_J, sigma_cs_H, sigma_cs_Ha, sigma_cs_N, sigma_cs_Ca, sigma_pf, gamma_index)
if debug:
print 'E_%d evaluated in model_%d'%(k,l), u_kln[k,l,n]
# Initialize MBAR with reduced energies u_kln and number of uncorrelated configurations from each state N_k.
# u_kln[k,l,n] is the reduced potential energy beta*U_l(x_kn), where U_l(x) is the potential energy function for state l,
# beta is the inverse temperature, and and x_kn denotes uncorrelated configuration n from state k.
# N_k[k] is the number of configurations from state k stored in u_knm
# Note that this step may take some time, as the relative dimensionless free energies f_k are determined at this point.
mbar = MBAR(u_kln, N_k)
# Extract dimensionless free energy differences and their statistical uncertainties.
# (Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
#(Deltaf_ij, dDeltaf_ij, Theta_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='svd-ew')
(Deltaf_ij, dDeltaf_ij, Theta_ij) = mbar.getFreeEnergyDifferences(uncertainty_method='approximate')
#print 'Deltaf_ij', Deltaf_ij
#print 'dDeltaf_ij', dDeltaf_ij
beta = 1.0 # keep in units kT
#print 'Unit-bearing (units kT) free energy difference f_1K = f_K - f_1: %f +- %f' % ( (1./beta) * Deltaf_ij[0,K-1], (1./beta) * dDeltaf_ij[0,K-1])
self.f_df = np.zeros( (self.nlambda, 2) ) # first column is Deltaf_ij[0,:], second column is dDeltaf_ij[0,:]
self.f_df[:,0] = Deltaf_ij[0,:]
self.f_df[:,1] = dDeltaf_ij[0,:]
# Compute the expectation of some observable A(x) at each state i, and associated uncertainty matrix.
# Here, A_kn[k,n] = A(x_{kn})
#(A_k, dA_k) = mbar.computeExpectations(A_kn)
self.P_dP = np.zeros( (nstates, 2*self.K) ) # left columns are P, right columns are dP
if debug:
print 'state\tP\tdP'
for i in range(nstates):
A_kn = np.where(states_kn==i,1,0)
(p_i, dp_i) = mbar.computeExpectations(A_kn, uncertainty_method='approximate')
self.P_dP[i,0:self.K] = p_i
self.P_dP[i,self.K:2*self.K] = dp_i
print i
for p in p_i: print p,
for dp in dp_i: print dp,
print
pops, dpops = self.P_dP[:,0:self.K], self.P_dP[:,self.K:2*self.K]
# save results
self.save_MBAR()
sigma_PF = traj[k]['allowed_sigma_PF'][sigma_PF_index] #GYH
u_kln[k,l,n] = sampler[l].neglogP(0, state, sigma_noe, sigma_J, sigma_cs_H, sigma_cs_Ha, sigma_cs_N, sigma_cs_Ca, sigma_PF, gamma_index, beta_c_index, beta_h_index, beta_0_index, xcs_index, xhs_index, bs_index) #GYH 03/2017
# u_kln[k,l,n] = sampler[l].neglogP(0, state, sigma_noe, sigma_J, sigma_cs_H, sigma_cs_Ha, sigma_cs_N, sigma_cs_Ca, sigma_PF, gamma_index) #GYH
print 'E_%d evaluated in model_%d'%(k,l), u_kln[k,l,n]
# Initialize MBAR with reduced energies u_kln and number of uncorrelated configurations from each state N_k.
#
# u_kln[k,l,n] is the reduced potential energy beta*U_l(x_kn), where U_l(x) is the potential energy function for state l,
# beta is the inverse temperature, and and x_kn denotes uncorrelated configuration n from state k.
#
# N_k[k] is the number of configurations from state k stored in u_knm
#
# Note that this step may take some time, as the relative dimensionless free energies f_k are determined at this point.
mbar = MBAR(u_kln, N_k)
# Extract dimensionless free energy differences and their statistical uncertainties.
(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
print 'Deltaf_ij', Deltaf_ij
print 'dDeltaf_ij', dDeltaf_ij
beta = 1.0 # keep in units kT
print 'Unit-bearing (units kT) free energy difference f_1K = f_K - f_1: %f +- %f' % ( (1./beta) * Deltaf_ij[0,K-1], (1./beta) * dDeltaf_ij[0,K-1])
f_df = np.zeros( (nlambda, 2) ) # first column is Deltaf_ij[0,:], second column is dDeltaf_ij[0,:]
f_df[:,0] = Deltaf_ij[0,:]
f_df[:,1] = dDeltaf_ij[0,:]
print 'Writing %s...'%args.bayesfactorfile
savetxt(args.bayesfactorfile, f_df)
print '...Done.'
def estimate_free_energies(ncfile, ndiscard=0, nuse=None, g=None):
"""
Estimate free energies of all alchemical states.
Parameters
----------
ncfile : NetCDF
Input YANK netcdf file
ndiscard : int, optional, default=0
Number of iterations to discard to equilibration
nuse : int, optional, default=None
Maximum number of iterations to use (after discarding)
g : int, optional, default=None
Statistical inefficiency to use if desired; if None, will be computed.
TODO
----
* Automatically determine 'ndiscard'.
"""
# Get current dimensions.
niterations = ncfile.variables['energies'].shape[0]
nstates = ncfile.variables['energies'].shape[1]
natoms = ncfile.variables['energies'].shape[2]
# Extract energies.
logger.info("Reading energies...")
energies = ncfile.variables['energies']
u_kln_replica = np.zeros([nstates, nstates, niterations], np.float64)
for n in range(niterations):
u_kln_replica[:,:,n] = energies[n,:,:]
logger.info("Done.")
# Deconvolute replicas
logger.info("Deconvoluting replicas...")
u_kln = np.zeros([nstates, nstates, niterations], np.float64)
for iteration in range(niterations):
state_indices = ncfile.variables['states'][iteration,:]
u_kln[state_indices,:,iteration] = energies[iteration,:,:]
logger.info("Done.")
# Compute total negative log probability over all iterations.
u_n = np.zeros([niterations], np.float64)
for iteration in range(niterations):
u_n[iteration] = np.sum(np.diagonal(u_kln[:,:,iteration]))
#logger.info(u_n
# DEBUG
outfile = open('u_n.out', 'w')
for iteration in range(niterations):
outfile.write("%8d %24.3f\n" % (iteration, u_n[iteration]))
outfile.close()
# Discard initial data to equilibration.
u_kln_replica = u_kln_replica[:,:,ndiscard:]
u_kln = u_kln[:,:,ndiscard:]
u_n = u_n[ndiscard:]
# Truncate to number of specified conforamtions to use
if (nuse):
u_kln_replica = u_kln_replica[:,:,0:nuse]
u_kln = u_kln[:,:,0:nuse]
u_n = u_n[0:nuse]
# Subsample data to obtain uncorrelated samples
N_k = np.zeros(nstates, np.int32)
indices = timeseries.subsampleCorrelatedData(u_n, g=g) # indices of uncorrelated samples
#print u_n # DEBUG
#indices = range(0,u_n.size) # DEBUG - assume samples are uncorrelated
N = len(indices) # number of uncorrelated samples
N_k[:] = N
u_kln[:,:,0:N] = u_kln[:,:,indices]
logger.info("number of uncorrelated samples:")
logger.info(N_k)
logger.info("")
#===================================================================================================
# Estimate free energy difference with MBAR.
#===================================================================================================
# Initialize MBAR (computing free energy estimates, which may take a while)
logger.info("Computing free energy differences...")
mbar = MBAR(u_kln, N_k)
# Get matrix of dimensionless free energy differences and uncertainty estimate.
logger.info("Computing covariance matrix...")
try:
# pymbar 2
(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
except ValueError:
# pymbar 3
(Deltaf_ij, dDeltaf_ij, theta_ij) = mbar.getFreeEnergyDifferences()
# # Matrix of free energy differences
logger.info("Deltaf_ij:")
for i in range(nstates):
str_row = ""
for j in range(nstates):
str_row += "%8.3f" % Deltaf_ij[i, j]
logger.info(str_row)
# print Deltaf_ij
# # Matrix of uncertainties in free energy difference (expectations standard deviations of the estimator about the true free energy)
logger.info("dDeltaf_ij:")
for i in range(nstates):
str_row = ""
for j in range(nstates):
str_row += "%8.3f" % dDeltaf_ij[i, j]
logger.info(str_row)
# Return free energy differences and an estimate of the covariance.
return (Deltaf_ij, dDeltaf_ij)
def write_pdb_resampled(ncfile, atoms, output_pdb_filename, reference_state=0, ndiscard=0, nsamples=1000):
"""
Resample configurations with probability in specified state, concatenating into PDB file.
ARGUMENTS
ncfile (NetCDF) - input YANK netcdf file
atoms (list) - atom records from reference PDB file
output_pdb_filename (string) - PDB file of resampled configurations to write
OPTIONAL ARGUMENTS
reference_state (int) - state to reweight to
ndiscard (int) - number of iterations to discard to equilibration
nsamples (int) - number of sampls to generate
"""
import numpy
# Get current dimensions.
niterations = ncfile.variables['energies'].shape[0]
nstates = ncfile.variables['energies'].shape[1]
natoms = ncfile.variables['energies'].shape[2]
# Extract energies.
print "Reading energies..."
energies = ncfile.variables['energies']
u_kln_replica = zeros([nstates, nstates, niterations], float64)
for n in range(niterations):
u_kln_replica[:,:,n] = energies[n,:,:]
print "Done."
# Deconvolute replicas to collect by thermodynamic state.
print "Deconvoluting replicas..."
u_kln = numpy.zeros([nstates, nstates, niterations], float64)
for iteration in range(niterations):
state_indices = ncfile.variables['states'][iteration,:]
u_kln[state_indices,:,iteration] = energies[iteration,:,:]
print "Done."
# Discard initial data to equilibration.
u_kln = u_kln[:,:,ndiscard:]
[K,L,N] = u_kln.shape
N_k = N * numpy.ones([K], numpy.int32)
# Compute snapshot energies.
print "Computing snapshot energies..."
u_kn = numpy.zeros([K,N], numpy.float32)
# Get temperature.
temperature = ncfile.groups['thermodynamic_states'].variables['temperatures'][reference_state] * units.kelvin
kB = units.BOLTZMANN_CONSTANT_kB * units.AVOGADRO_CONSTANT_NA # Boltzmann constant
kT = kB * temperature # thermal energy
# Deserialize system.
system = openmm.XmlSerializer.deserialize(str(ncfile.groups['thermodynamic_states'].variables['systems'][reference_state]))
# Create Context.
integrator = openmm.VerletIntegrator(1.0 * units.femtoseconds)
context = openmm.Context(system, integrator)
# Turn off restraints.
context.setParameter('restraint_lambda', 0.0)
for n in range(N):
iteration = ndiscard + n
for replica_index in range(K):
# Recompute energies.
state_index = ncfile.variables['states'][iteration,replica_index]
positions = ncfile.variables['positions'][iteration,replica_index]
context.setPositions(positions)
openmm_state = context.getState(getEnergy=True)
u_kn[state_index,n] = openmm_state.getPotentialEnergy() / kT
#===================================================================================================
# Initialize MBAR.
#===================================================================================================
# Initialize MBAR (computing free energy estimates, which may take a while)
print "Initializing MBAR (warning: using all correlated data)..."
mbar = MBAR(u_kln, N_k, verbose=False)
# Get snapshot weights.
#u_kn = numpy.squeeze(u_kln[:,state,:])
log_w_kn = mbar._computeUnnormalizedLogWeights(u_kn)
f = _logsum(log_w_kn[mbar.indices])
w_kn = numpy.exp(log_w_kn - f)
p_kn = w_kn / w_kn[mbar.indices].sum()
# Form linear list of potential snapshots to choose.
[K, N] = p_kn.shape
snapshots = list()
probabilities = list()
for replica in range(K):
for sample in range(N):
state = ncfile.variables['states'][ndiscard+sample, replica]
snapshots.append( (replica, ndiscard+sample) )
probabilities.append( p_kn[state,sample] )
probabilities = numpy.array(probabilities)
# Draw samples.
state_hist = numpy.zeros([K], numpy.int32)
import numpy.random
sample_indices = numpy.random.choice(range(len(snapshots)), size=[nsamples], p=probabilities)
title = 'generated by resampling'
outfile = open(output_pdb_filename, 'w')
for (sample_number, sample_index) in enumerate(sample_indices):
[replica, iteration] = snapshots[sample_index]
state = ncfile.variables['states'][iteration, replica]
state_hist[state] += 1
print "sample %8d : replica %3d iteration %6d state %3d" % (sample_number, replica, iteration, state)
outfile.write('MODEL %4d\n' % (sample_number+1))
write_pdb(atoms, outfile, iteration, replica, title, ncfile)
outfile.write('ENDMDL\n')
outfile.close()
print "Counts from each state:"
print "%5s %6s" % ('state', 'counts')
for k in range(K):
print "%5d %6d" % (k, state_hist[k])
print ""
return
def execute(nstates, q_samp_space, epsi_samp_space, sig_samp_space, useO=False, sliceset=None):
#Initilize limts
sig3_samp_space = sig_samp_space**3
#These are the limits used to compute the constant matricies
#They should match with LJ 0 and 5, ALWAYS
q_min = -2.0
q_max = +2.0
epsi_min = epsi_samp_space[0]
epsi_max = epsi_samp_space[5]
sig_min = sig_samp_space[0]
sig_max = sig_samp_space[5]
lamto_epsi = lambda lam: (epsi_max - epsi_min)*lam + epsi_min
lamto_sig3 = lambda lam: (sig_max**3 - sig_min**3)*lam + sig
lamto_sig = lambda lam: lamto_sig3(lam)**(1.0/3)
if spacing is logspace:
StartSpace = -5
EndSpace = 0
spacename='log'
PlotStart = 10**StartSpace
PlotEnd = 10**EndSpace
elif spacing is linspace:
sigStartSpace = sig_min
sigEndSpace = sig_max
epsiStartSpace = epsi_min
qStartSpace = q_min
qEndSpace = q_max
if alle:
epsiEndSpace = 3.6 #!!! Manual set
else:
epsiEndSpace = epsi_max
spacename='linear'
sigPlotStart = sigStartSpace**sig_factor
sigPlotEnd = sigEndSpace**sig_factor
epsiPlotStart = epsiStartSpace
epsiPlotEnd = epsiEndSpace
qPlotStart = qStartSpace
qPlotEnd = qEndSpace
#generate sample length
g_en_start = 19 #Row where data starts in g_energy output
g_en_energy = 1 #Column where the energy is located
niterations_max = 30001
#Min and max sigmas:
fC12 = lambda epsi,sig: 4*epsi*sig**12
fC6 = lambda epsi,sig: 4*epsi*sig**6
fsig = lambda C12, C6: (C12/C6)**(1.0/6)
fepsi = lambda C12, C6: C6**2/(4*C12)
C12_delta = fC12(epsi_max, sig_max) - fC12(epsi_min, sig_min)
C6_delta = fC6(epsi_max, sig_max) - fC6(epsi_min, sig_min)
C12_delta_sqrt = fC12(epsi_max, sig_max)**.5 - fC12(epsi_min, sig_min)**.5
C6_delta_sqrt = fC6(epsi_max, sig_max)**.5 - fC6(epsi_min, sig_min)**.5
#Set up lambda calculation equations
flamC6 = lambda epsi, sig: (fC6(epsi, sig) - fC6(epsi_min, sig_min))/C6_delta
flamC12 = lambda epsi, sig: (fC12(epsi, sig) - fC12(epsi_min, sig_min))/C12_delta
flamC6sqrt = lambda epsi, sig: (fC6(epsi, sig)**.5 - fC6(epsi_min, sig_min)**.5)/C6_delta_sqrt
flamC12sqrt = lambda epsi, sig: (fC12(epsi, sig)**.5 - fC12(epsi_min, sig_min)**.5)/C12_delta_sqrt
flamC1 = lambda q: q
lamC12 = flamC12sqrt(epsi_samp_space, sig_samp_space)
lamC6 = flamC6sqrt(epsi_samp_space, sig_samp_space)
lamC1 = flamC1(q_samp_space)
#Try to load u_kln
lam_range = linspace(0,1,nstates)
subsampled = numpy.zeros([nstates],dtype=numpy.bool)
if load_ukln and os.path.isfile('esq_ukln_consts_n%i.npz'%nstates):
energies = consts(nstates, file='esq_ukln_consts_n%i.npz'%nstates)
else:
#Initial u_kln
energies = consts(nstates)
g_t = numpy.zeros([nstates])
#Read in the data
for k in xrange(nstates):
print "Importing LJ = %02i" % k
energy_dic = {'full':{}, 'rep':{}}
#Try to load the subsampled filenames
try:
energy_dic['null'] = open('lj%s/prod/subenergy%s_null.xvg' %(k,k),'r').readlines()[g_en_start:] #Read in the null energies (unaffected) of the K states
for l in xrange(nstates):
energy_dic['full']['%s'%l] = open('lj%s/prod/subenergy%s_%s.xvg' %(k,k,l),'r').readlines()[g_en_start:] #Read in the full energies for each state at KxL
if l == 5 or l == 0:
energy_dic['rep']['%s'%l] = open('lj%s/prod/subenergy%s_%s_rep.xvg' %(k,k,l),'r').readlines()[g_en_start:] #Read in the repulsive energies at 0, nstates-1, and K
energy_dic['q'] = open('lj%s/prod/subenergy%s_q.xvg' %(k,k),'r').readlines()[g_en_start:] #Read in the charge potential energy
energy_dic['q2'] = open('lj%s/prod/subenergy%s_q2.xvg' %(k,k),'r').readlines()[g_en_start:] #Read in the charge potential energy
iter = len(energy_dic['null'])
subsampled[k] = True
# Set the object to iterate over, since we want every frame of the subsampled proces, we just use every frame
frames = xrange(iter)
except: #Load the normal way
energy_dic['null'] = open('lj%s/prod/energy%s_null.xvg' %(k,k),'r').readlines()[g_en_start:] #Read in the null energies (unaffected) of the K states
for l in xrange(nstates):
energy_dic['full']['%s'%l] = open('lj%s/prod/energy%s_%s.xvg' %(k,k,l),'r').readlines()[g_en_start:] #Read in the full energies for each state at KxL
if l == 5 or l == 0:
energy_dic['rep']['%s'%l] = open('lj%s/prod/energy%s_%s_rep.xvg' %(k,k,l),'r').readlines()[g_en_start:] #Read in the repulsive energies at 0, nstates-1, and K
energy_dic['q'] = open('lj%s/prod/energy%s_q.xvg' %(k,k),'r').readlines()[g_en_start:] #Read in the charge potential energy
energy_dic['q2'] = open('lj%s/prod/energy%s_q2.xvg' %(k,k),'r').readlines()[g_en_start:] #Read in the charge potential energy
iter = niterations_max
#Subsample
tempenergy = numpy.zeros(iter)
for frame in xrange(iter):
tempenergy[frame] = float(energy_dic['full']['%s'%k][frame].split()[g_en_energy])
frames, temp_series, g_t[k] = subsample_series(tempenergy, return_g_t=True)
print "State %i has g_t of %i" % (k, g_t[k])
iter = len(frames)
#Update iterations if need be
energies.updateiter(iter)
#Fill in matricies
n = 0
for frame in frames:
#Unaffected state
energies.const_Un_matrix[k,n] = float(energy_dic['null'][frame].split()[g_en_energy])
#Charge only state
VI = float(energy_dic['q'][frame].split()[g_en_energy]) - energies.const_Un_matrix[k,n]
VII = float(energy_dic['q2'][frame].split()[g_en_energy]) - energies.const_Un_matrix[k,n]
q1 = 1
q2 = 0.5
QB = (q1**2*VII - q2**2*VI)/(q1**2*q2 - q2**2*q1)
QA = (VI - q1*QB)/q1**2
energies.const_q_matrix[k,n] = QB
energies.const_q2_matrix[k,n] = QA
#const_q_matrix[k,n] = float(energy_dic['q'][n].split()[g_en_energy]) - const_Un_matrix[k,n]
#Isolate the data
for l in xrange(nstates):
energies.u_kln[k,l,n] = float(energy_dic['full']['%s'%l][frame].split()[g_en_energy]) #extract the kln energy, get the line, split the line, get the energy, convert to float, store
#Repulsive terms:
#R0 = U_rep[k,k,n] + dhdl[k,0,n] - Un[k,n]
energies.const_R0_matrix[k,n] = float(energy_dic['rep']['%s'%(0)][frame].split()[g_en_energy]) - energies.const_Un_matrix[k,n]
#R1 = U_rep[k,k,n] + dhdl[k,-1,n] - Un[k,n]
energies.const_R1_matrix[k,n] = float(energy_dic['rep']['%s'%(5)][frame].split()[g_en_energy]) - energies.const_Un_matrix[k,n]
energies.const_R_matrix[k,n] = energies.const_R1_matrix[k,n] - energies.const_R0_matrix[k,n]
#Finish the total unaffected term
#Total unaffected = const_Un + U0 = const_Un + (U_full[k,0,n] - const_Un) = U_full[k,0,n]
energies.const_unaffected_matrix[k,n] = energies.u_kln[k,0,n]
#Fill in the q matrix
#Attractive term
#u_A = U_full[k,n] - constR[k,n] - const_unaffected[k,n]
energies.const_A0_matrix[k,n] = energies.u_kln[k,0,n] - energies.const_R0_matrix[k,n] - energies.const_Un_matrix[k,n]
energies.const_A1_matrix[k,n] = energies.u_kln[k,5,n] - energies.const_R1_matrix[k,n] - energies.const_Un_matrix[k,n]
energies.const_A_matrix[k,n] = energies.const_A1_matrix[k,n] - energies.const_A0_matrix[k,n]
n += 1
energies.determine_all_N_k()
#write_g_t(g_t)
constname = 'esq_ukln_consts_n%i.npz'%nstates
if load_ukln and not os.path.isfile(constname):
energies.save_consts(constname)
#Sanity check
sanity_kln = numpy.zeros(energies.u_kln.shape)
#Round q to the the GROMACS precision, otherwise there is drift in q which can cause false positives
lamC1r = numpy.around(lamC1, decimals=4)
for l in xrange(nstates):
#sanity_kln[:,l,:] = lamC12[l]*energies.const_R_matrix + lamC6[l]*energies.const_A_matrix + lamC1[l]*energies.const_q_matrix + lamC1[l]**2*energies.const_q2_matrix + energies.const_unaffected_matrix
sanity_kln[:,l,:] = lamC12[l]*energies.const_R_matrix + lamC6[l]*energies.const_A_matrix + lamC1r[l]*energies.const_q_matrix + lamC1r[l]**2*energies.const_q2_matrix + energies.const_unaffected_matrix
del_kln = numpy.abs(energies.u_kln - sanity_kln)
del_tol = 1 #in kJ per mol
print "Max Delta: %f" % numpy.nanmax(del_kln)
if numpy.nanmax(del_kln) > 1: #Check for numeric error
#Double check to see if the weight is non-zero.
#Most common occurance is when small particle is tested in large particle properties
#Results in energies > 60,000 kj/mol, which carry 0 weight to machine precision
nonzero_weights = numpy.count_nonzero(numpy.exp(-energies.u_kln[numpy.where(del_kln > .2)] * kjpermolTokT))
if nonzero_weights != 0:
print "and there are %d nonzero weights! Stopping execution" % nonzero_weights
pdb.set_trace()
else:
print "but these carry no weight and so numeric error does not change the answer"
##################################################
############### END DATA INPUT ###################
##################################################
#Convert to dimless
energies.dimless()
#Set up regular grid
sig_range = (spacing(sigStartSpace**3,sigEndSpace**3,Nparm))**(1.0/3)
epsi_range = spacing(epsiStartSpace,epsiEndSpace,Nparm)
q_range = spacing(qStartSpace,qEndSpace,Nparm)
epsi_plot_range = spacing(epsiStartSpace,epsi_max,Nparm)
#Load subsequent f_ki
f_ki_loaded = False
state_counter = nstates
while not f_ki_loaded and state_counter != 0:
#Load the largest f_ki you can
try:
f_ki_load = numpy.load('esq_f_k_{myint:{width}}.npy'.format(myint=state_counter, width=len(str(state_counter))))
f_ki_loaded = True
f_ki_n = state_counter
except:
pass
state_counter -= 1
try:
if nstates >= f_ki_n:
draw_ki = f_ki_n
else:
draw_ki = nstates
#Copy the loaded data
f_ki = numpy.zeros(nstates)
f_ki[:draw_ki] = f_ki_load[:draw_ki]
mbar = MBAR(energies.u_kln, energies.N_k, verbose = True, initial_f_k=f_ki, subsampling_protocol=[{'method':'L-BFGS-B','options':{'disp':True}}], subsampling=1)
except:
mbar = MBAR(energies.u_kln, energies.N_k, verbose = True, subsampling_protocol=[{'method':'L-BFGS-B','options':{'disp':True}}], subsampling=1)
if not f_ki_loaded or f_ki_n != nstates:
try:
numpy.save_compressed('esq_f_k_{myint:{width}}.npy'.format(myint=nstates, width=len(str(nstates))), mbar.f_k)
except:
numpy.save('esq_f_k_{myint:{width}}.npy'.format(myint=nstates, width=len(str(nstates))), mbar.f_k)
######## Begin Computing RDF's #########
basebins = 800
plotx = linspace(genrdf.distmin, genrdf.distmax, basebins)/10.0 #put in nm
dr = plotx[1] - plotx[0]
#Generate blank
rdfs = numpy.zeros([basebins, nstates, energies.itermax], dtype=numpy.float64)
if useO:
pathstr = 'lj%s/prod/rdfOhist%sb%i.npz'
else:
pathstr = 'lj%s/prod/rdfhist%s.npz'
for k in xrange(nstates):
try: #Try to load in
if useO:
rdfk = numpy.load(pathstr%(k,k,basebins))['rdf']
else:
rdfk = numpy.load(pathstr%(k,k))['rdf']
except:
raise
nkbin, nkiter = rdfk.shape
rdfs[:,k,:nkiter] = rdfk
#Do halfs
if sliceset is 1:
darange = xrange(0,25)
nrange = 25
elif sliceset is 2:
darange = xrange(25,Nparm)
nrange = 25
#Do fifths
elif sliceset is 3:
darange = xrange(0,10)
nrange = 10
elif sliceset is 4:
darange = xrange(10,20)
nrange = 10
elif sliceset is 5:
darange = xrange(20,30)
nrange = 10
elif sliceset is 6:
darange = xrange(30,40)
nrange = 10
elif sliceset is 7:
darange = xrange(40,Nparm)
nrange = Nparm-40
#Do Everything
else:
darange = xrange(Nparm)
nrange = Nparm
#Generate the rmins for each lj combo
# epsi sig
rmins = numpy.zeros([Nparm, Nparm])
q = 0
run_start_time = time.time()
number_of_iterations = Nparm*nrange
iteration = 0
gr = (numpy.sqrt(5)-1)/2.0
#Set up units for closest point
Uplotx = units.Quantity(plotx, units.nanometer)
zeroR = numpy.zeros(plotx.shape)
OWsig = 3.15061 * units.angstrom
OWepsi = 0.6364 * units.kilojoules_per_mole
for iepsi in darange:
epsi = epsi_range[iepsi]
for isig in xrange(Nparm):
if not os.path.isfile('esq_%s/RDFns%iE%iS%i.npz' %(spacename, nstates, iepsi, isig)):
initial_time = time.time()
iteration += 1
print "Working on e %i and s %i" % (iepsi, isig)
sig = sig_range[isig]
u_kn_P = flamC12sqrt(epsi,sig)*energies.const_R_matrix + flamC6sqrt(epsi,sig)*energies.const_A_matrix + flamC1(q)*energies.const_q_matrix + flamC1(q)**2*energies.const_q2_matrix + energies.const_unaffected_matrix
#Estimate zero-th order RDF
guessepsi = geo(epsi*units.kilojoules_per_mole,OWepsi)
guesssig = geo(sig*units.nanometer, OWsig)
rdfguess = exp0(Uplotx, guessepsi, guesssig)
depart0 = plotx[numpy.isclose(rdfguess,zeroR)].max()
#Find nearest value in array
#rcurrent = numpy.argmin(numpy.abs(sig*0.75-plotx))
rcurrent = numpy.argmin(numpy.abs(depart0-plotx))
foundmin = False
#Estimate current min
Ecurrent = mbar.computePerturbedExpectation(u_kn_P, rdfs[rcurrent,:,:], compute_uncertainty=False)
if numpy.isclose(Ecurrent, 0):
#Case where sigma already is g=0
while foundmin is False:
rnew = rcurrent + 1
Enew = mbar.computePerturbedExpectation(u_kn_P, rdfs[rnew,:,:], compute_uncertainty=False)
if numpy.isclose(Enew, 0):
#Have not found where g(r)=0 and g(r+dr)>0
rcurrent = rnew
else:
foundmin = True
else:
#Case where sigma is not g=0
while foundmin is False:
rnew = rcurrent - 1
Enew = mbar.computePerturbedExpectation(u_kn_P, rdfs[rnew,:,:], compute_uncertainty=False)
#Reversed logic of previous since we are searching backwards
if not numpy.isclose(Enew, 0):
#Have not found where g(r)=0 and g(r+dr)>0
rcurrent = rnew
else:
foundmin = True
#rmins[iepsi,isig] = rcurrent
laptime = time.clock()
# Show timing statistics. copied from Repex.py, copywrite John Chodera
final_time = time.time()
elapsed_time = final_time - initial_time
estimated_time_remaining = (final_time - run_start_time) / (iteration) * (number_of_iterations - iteration)
estimated_total_time = (final_time - run_start_time) / (iteration) * (number_of_iterations)
estimated_finish_time = final_time + estimated_time_remaining
print "Iteration took %.3f s." % elapsed_time
print "Estimated completion in %s, at %s (consuming total wall clock time %s)." % (str(datetime.timedelta(seconds=estimated_time_remaining)), time.ctime(estimated_finish_time), str(datetime.timedelta(seconds=estimated_total_time)))
savez('esq_%s/RDFns%iE%iS%i.npz' %(spacename, nstates, iepsi, isig), rmins=numpy.array([plotx[rnew]]))
else:
rmins[iepsi,isig] = numpy.load('esq_%s/RDFns%iE%iS%i.npz' %(spacename, nstates, iepsi, isig))['rmins']
savez('LJEffHS.npz', rmins=rmins)
def compute_hydration_energy(entry, parameters, platform_name="CPU"):
"""
Compute hydration energy of a single molecule given a GBSA parameter set.
ARGUMENTS
molecule (OEMol) - molecule with GBSA atom types
parameters (dict) - parameters for GBSA atom types
RETURNS
energy (float) - hydration energy in kcal/mol
"""
platform = openmm.Platform.getPlatformByName('CPU')
from pymbar import MBAR
timestep = 2 * units.femtoseconds
molecule = entry['molecule']
iupac_name = entry['iupac']
cid = molecule.GetData('cid')
# Retrieve OpenMM System.
vacuum_system = entry['system']
solvent_system = copy.deepcopy(entry['solvated_system'])
# Get nonbonded force.
forces = { solvent_system.getForce(index).__class__.__name__ : solvent_system.getForce(index) for index in range(solvent_system.getNumForces()) }
nonbonded_force = forces['NonbondedForce']
gbsa_force = forces['CustomGBForce']
# Build indexable list of atoms.
atoms = [atom for atom in molecule.GetAtoms()]
natoms = len(atoms)
# Create context for solvent system.
timestep = 2.0 * units.femtosecond
solvent_integrator = openmm.VerletIntegrator(timestep)
# Create context for vacuum system.
vacuum_integrator = openmm.VerletIntegrator(timestep)
# Assign GBSA parameters.
for (atom_index, atom) in enumerate(atoms):
[charge, sigma, epsilon] = nonbonded_force.getParticleParameters(atom_index)
atomtype = atom.GetStringData("gbsa_type") # GBSA atomtype
radius = parameters['%s_%s' % (atomtype, 'radius')] * units.angstroms
scalingFactor = parameters['%s_%s' % (atomtype, 'scalingFactor')]
gbsa_force.setParticleParameters(atom_index, [charge, radius, scalingFactor])
solvent_context = openmm.Context(solvent_system, solvent_integrator,platform)
vacuum_context = openmm.Context(vacuum_system, vacuum_integrator, platform)
# Compute energy differences.
temperature = entry['temperature']
kT = kB * temperature
beta = 1.0 / kT
initial_time = time.time()
x_n = entry['x_n']
u_n = entry['u_n']
nsamples = len(u_n)
nstates = 3 # number of thermodynamic states
u_kln = np.zeros([3,3,nsamples], np.float64)
for sample in range(nsamples):
positions = units.Quantity(x_n[sample,:,:], units.nanometers)
u_kln[0,0,sample] = u_n[sample]
vacuum_context.setPositions(positions)
vacuum_state = vacuum_context.getState(getEnergy=True)
u_kln[0,1,sample] = beta * vacuum_state.getPotentialEnergy()
solvent_context.setPositions(positions)
solvent_state = solvent_context.getState(getEnergy=True)
u_kln[0,2,sample] = beta * solvent_state.getPotentialEnergy()
N_k = np.zeros([nstates], np.int32)
N_k[0] = nsamples
mbar = MBAR(u_kln, N_k)
try:
df_ij, ddf_ij, _ = mbar.getFreeEnergyDifferences()
except linalg.LinAlgError:
return np.inf
DeltaG_in_kT = df_ij[1,2]
dDeltaG_in_kT = ddf_ij[1,2]
final_time = time.time()
elapsed_time = final_time - initial_time
#print "%48s | %48s | reweighting took %.3f s" % (cid, iupac_name, elapsed_time)
# Clean up.
del solvent_context, solvent_integrator
del vacuum_context, vacuum_integrator
energy = kT * DeltaG_in_kT
#print "%48s | %48s | DeltaG = %.3f +- %.3f kT " % (cid, iupac_name, energy, dDeltaG_in_kT)
print(DeltaG_in_kT)
print(type(DeltaG_in_kT))
return DeltaG_in_kT
# get the unreduced energies U_kln = u_kln/beta #============================================================================================= # Estimate free energies and expectations. #============================================================================================= print "======================================" print " Initializing MBAR " print "======================================" # Estimate free energies from simulation using MBAR. print "Estimating relative free energies from simulation (this may take a while)..." # Initialize the MBAR class, determining the free energies. mbar = MBAR(u_kln, N_k, method = 'adaptive',relative_tolerance=1.0e-10,verbose=True) # Get matrix of dimensionless free energy differences and uncertainty estimate. print "=============================================" print " Testing getFreeEnergyDifferences " print "=============================================" (Delta_f_ij_estimated, dDelta_f_ij_estimated) = mbar.getFreeEnergyDifferences() # Compute error from analytical free energy differences. Delta_f_ij_error = Delta_f_ij_estimated - Delta_f_ij_analytical print "Error in free energies is:" print Delta_f_ij_error print "Standard deviations away is:"
#=============================================================================================
# Generate independent data samples from K one-dimensional harmonic oscillators centered at q = 0.
#=============================================================================================
randomsample = testsystems.harmonic_oscillators.HarmonicOscillatorsTestCase(O_k=O_k, K_k=K_k, beta=beta)
[x_kn,u_kln,N_k] = randomsample.sample(N_k,mode='u_kln')
# get the unreduced energies
U_kln = u_kln/beta
#=============================================================================================
# Estimate free energies and expectations.
#=============================================================================================
# Initialize the MBAR class, determining the free energies.
mbar = MBAR(u_kln, N_k, method = 'adaptive',relative_tolerance=1.0e-10,verbose=False) # use fast Newton-Raphson solver
(Deltaf_ij_estimated, dDeltaf_ij_estimated) = mbar.getFreeEnergyDifferences()
# Compute error from analytical free energy differences.
Deltaf_ij_error = Deltaf_ij_estimated - Deltaf_ij_analytical
# Estimate the expectation of the mean-squared displacement at each condition.
if observe == 'RMS displacement':
A_kn = numpy.zeros([K,K,N_max], dtype = numpy.float64);
for k in range(0,K):
for l in range(0,K):
A_kn[k,l,0:N_k[k]] = (x_kn[k,0:N_k[k]] - O_k[l])**2 # observable is the squared displacement
# observable is the potential energy, a 3D array since the potential energy is a function of
# thermodynamic state
elif observe == 'potential energy':
""" Reminder: ['step', 'E', 'accept', 'state', 'sigma_noe', 'sigma_J', 'gamma']"""
sigma_noe = traj[k]['allowed_sigma_noe'][sigma_noe_index]
sigma_J = traj[k]['allowed_sigma_J'][sigma_J_index]
u_kln[k,l,i] = samplers[l].neglogP(state, sigma_noe, sigma_J, gamma_index)
print u_kln
# Initialize MBAR with reduced energies u_kln and number of uncorrelated configurations from each state N_k.
#
# u_kln[k,l,n] is the reduced potential energy beta*U_l(x_kn), where U_l(x) is the potential energy function for state l,
# beta is the inverse temperature, and and x_kn denotes uncorrelated configuration n from state k.
#
# N_k[k] is the number of configurations from state k stored in u_knm
#
# Note that this step may take some time, as the relative dimensionless free energies f_k are determined at this point.
mbar = MBAR(u_kln, N_k, verbose=True)
# OPTIONS
# maximum_iterations=10000, relative_tolerance=1e-07, verbose=False,
# initial_f_k=None, method='adaptive', use_optimized=None, newton_first_gamma=0.1, newton_self_consistent=2, maxrange=100000.0, initialize='zeros'
# Extract dimensionless free energy differences and their statistical uncertainties.
(Deltaf_ij, dDeltaf_ij) = mbar.getFreeEnergyDifferences()
print 'Deltaf_ij', Deltaf_ij
print 'dDeltaf_ij', dDeltaf_ij
K = 2
beta = 1.0 # keep in units kT
print 'Unit-bearing (units kT) free energy difference between states 1 and K: %f +- %f' % ( (1./beta) * Deltaf_ij[0,K-1], (1./beta) * dDeltaf_ij[0,K-1])
# Compute the expectation of some observable A(x) at each state i, and associated uncertainty matrix.
# Here, A_kn[k,n] = A(x_{kn})
#(A_k, dA_k) = mbar.computeExpectations(A_kn)