def dA_endpoint_MBAR(polymorphs='p1 p2', Molecules=72, Independent=4, Temp=200):
# Setting constants
kJ_to_kcal = 1/4.184 # Converting kJ to kcal
kB = 0.0019872041 # boltzman constant in kcal/(mol*K)
# Getting the polymorph names
polymorphs = polymorphs.split()
# Place to store the free energy differences
dA = np.zeros(len(polymorphs))
ddA = np.zeros(len(polymorphs))
for i, poly in enumerate(polymorphs):
if os.path.isfile(poly + '/interactions/100/PROD.edr') and os.path.isfile(poly + '/interactions/100/END.edr'):
# Loading in the difference between the endpoint and production files
dU = panedr.edr_to_df(poly + '/interactions/100/END.edr')['Potential'].values - panedr.edr_to_df(poly + '/interactions/100/PROD.edr')['Potential'].values
# Converting the energy differences to go into pymbar
dW = dU / Molecules * kJ_to_kcal / (kB * Temp)
# Getting the energy differences with MBAR using Exponential Averaging
da = np.array(pymbar.EXP(dW)) * kB * Temp
dA[i] = -da[0]
ddA[i] = da[1]
else:
dA[i] = np.nan
ddA[i] = np.nan
# Check to see if there are any nan values
if np.any(dA == np.nan):
dA[:] = 0.
ddA[:] = 0.
return dA, ddA
def one_endpoint_perturbation(self):
hybrid_reduced_potentials = np.array(
self._lambda_one_reduced_potentials)
nonalchemical_reduced_potentials = np.array(
self._lambda_zero_reduced_potentials)
[df, ddf] = pymbar.EXP(nonalchemical_reduced_potentials -
hybrid_reduced_potentials)
return [df, ddf]
def calculate_cross_variance(all_results):
"""
Calculates the overlap (df and ddf) between the non-alchemical state at lambda=0 to the hybrid state at lambda=1 and visa versa
These ensembles are not expected to have good overlap, as they are of explicitly different system, but provides a benchmark of appropriate dissimilarity
"""
if len(all_results) != 4:
return
else:
non_a = all_results[0]
hybrid_a = all_results[1]
non_b = all_results[2]
hybrid_b = all_results[3]
print('CROSS VALIDATION')
[df, ddf] = pymbar.EXP(non_a - hybrid_b)
print('df: {}, ddf: {}'.format(df, ddf))
[df, ddf] = pymbar.EXP(non_b - hybrid_a)
print('df: {}, ddf: {}'.format(df, ddf))
return
def run_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0,
print_work=False,
write_system=False,
write_state=False,
write_trajectories=False):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional, default=0
The index, 0 or 1, at which to retrieve nonalchemical positions
print_work : bool, optional, default=False
If True, will print work values
write_system : bool, optional, default=False
If True, will write alchemical and nonalchemical System XML files
write_state : bool, optional, default=False
If True, write alchemical (hybrid) State XML files each iteration
write_trajectories : bool, optional, default=False
If True, will write trajectories
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
import mdtraj as md
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
if write_system:
with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
lambda_thermodynamic_state.system))
with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
nonalchemical_thermodynamic_state.system))
#initialize work array
w = np.zeros([n_iterations])
non_potential = np.zeros([n_iterations])
hybrid_potential = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
hybrid_trajectory = unit.Quantity(
np.zeros([
n_iterations,
lambda_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
nonalchemical_trajectory = unit.Quantity(
np.zeros([
n_iterations,
nonalchemical_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
for iteration in range(n_iterations):
# Generate a new sampler state for the hybrid system
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
# Compute the hybrid reduced potential at the new sampler state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
if write_state:
state = hybrid_context.getState(getPositions=True,
getParameters=True)
state_xml = openmm.XmlSerializer.serialize(state)
with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile:
outfile.write(state_xml)
# Construct a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
if write_trajectories:
state = hybrid_context.getState(getPositions=True)
hybrid_trajectory[iteration, :, :] = state.getPositions(
asNumpy=True)
nonalchemical_trajectory[iteration, :, :] = nonalchemical_positions
# Compute the nonalchemical reduced potential
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
# Compute and store the work
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
non_potential[iteration] = nonalchemical_reduced_potential
hybrid_potential[iteration] = hybrid_reduced_potential
if print_work:
print(
f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}'
)
if write_trajectories:
if lambda_index == 0:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.old_topology)
elif lambda_index == 1:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.new_topology)
md.Trajectory(
hybrid_trajectory / unit.nanometers,
factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb')
md.Trajectory(nonalchemical_trajectory / unit.nanometers,
nonalchemical_mdtraj_topology).save(
f'nonalchemical{lambda_index}.pdb')
# Analyze data and return results
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
results = [df, ddf_corrected, t0, Neff_max]
return results, non_potential, hybrid_potential
fname = 'graph' + str(i) + 'to' + str(j) + '.pdf'
plt.savefig(fname)
N_k = npoints*np.ones([nstates],int)
mbar = pymbar.MBAR(u_kln,N_k,relative_tolerance=1.0e-10,verbose=True)
(Delta_f_ij_estimated, dDelta_f_ij_estimated) = mbar.getFreeEnergyDifferences()
print Delta_f_ij_estimated
print dDelta_f_ij_estimated
# check these in the case of two.
# try exponential averaging, to see if there is a difference. Seems to work.
if len(dirnames) == 2:
wf = -(u_kln[0,1:,]-u_kln[0,0,:])
(df_forward,ddf_forward) = pymbar.EXP(wf)
print('EXP forward %10.4f +/- %7.4f' % (df_forward, ddf_forward))
wr = -(u_kln[1,1:,]-u_kln[1,0,:])
(df_rev, ddf_rev) = pymbar.EXP(wr)
print('EXP reverse %10.4f +/- %7.4f' % (df_rev, ddf_rev))
pdb.set_trace()
(df_bar, ddf_bar) = pymbar.BAR(-wf,wr)
print('BAR reverse %10.4f +/- %7.4f' % (-df_bar, ddf_bar))
# plots the overlap in energy betwen two. Looks good!
if plot:
plt.clf()
plt.hist(wf.T, facecolor='red')
plt.hist(wr.T, facecolor='blue')
def run_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional default 0
The index, 0 or 1, at which to retrieve nonalchemical positions
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
#initialize work array
w = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
for iteration in range(n_iterations):
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
#compute the reduced potential at the new state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
#generate a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
#compute the reduced potential at the nonalchemical system as well:
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
print(Neff_max)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
return [df, ddf_corrected, Neff_max]