152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
import os
|
|
from copy import deepcopy
|
|
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
from adaptiveumbrella.runner import UmbrellaRunner
|
|
|
|
|
|
class WHAM2DRunner(UmbrellaRunner):
|
|
""" Umbrella runner implementation that uses wham-2d to perform
|
|
the pmf calculation.
|
|
|
|
Attributes:
|
|
WHAM_EXEC: path to wham executeable
|
|
|
|
"""
|
|
|
|
def __init__(self):
|
|
UmbrellaRunner.__init__(self)
|
|
self.WHAM_EXEC = 'wham-2d'
|
|
|
|
def calculate_new_pmf(self):
|
|
import os
|
|
from shutil import copyfile
|
|
|
|
simulation_dir = "tmp/simulations"
|
|
print("Collecting sampling data from simulations folder")
|
|
|
|
# collect COLVARs
|
|
wham_dir = "tmp/WHAM/"
|
|
if not os.path.exists(wham_dir):
|
|
os.makedirs(wham_dir)
|
|
|
|
for folder in os.listdir(simulation_dir):
|
|
src = os.path.join(simulation_dir, folder, "COLVAR")
|
|
dst = os.path.join(wham_dir, folder + ".xvg")
|
|
copyfile(src, dst)
|
|
|
|
# create metadata file
|
|
metadata_file = os.path.join(wham_dir, "{}_metadata.dat".format(self.num_iterations))
|
|
fc_x = 100
|
|
fc_y = 100
|
|
with open(metadata_file, 'w') as out:
|
|
for f in os.listdir(simulation_dir):
|
|
prefix, x, y = f.split("_")
|
|
out.write("{}/{}.xvg {} {} {} {}\n".format(wham_dir, f, x, y, fc_x, fc_y))
|
|
|
|
# run WHAM2d
|
|
print("Running WHAM-2d")
|
|
wham_output = os.path.join(wham_dir, "{}_freeenergy.dat".format(self.num_iterations))
|
|
periodicity_x = "pi"
|
|
periodicity_y = "pi"
|
|
tolerance = 0.1
|
|
frames_x, frames_y = 1002, 1002
|
|
min_x = self.cvs[0][0]
|
|
max_x = self.cvs[0][1]
|
|
min_y = self.cvs[1][0]
|
|
max_y = self.cvs[1][1]
|
|
|
|
cmd = "{exec} Px={px} {min_x} {max_x} {frames_x} Py={py} {min_y} {max_y} {frames_y} {tol} 298 0 {metafile} {outfile} 0".format(
|
|
exec=self.WHAM_EXEC,
|
|
px=periodicity_x,
|
|
min_x=min_x,
|
|
max_x=max_x,
|
|
frames_x=frames_x,
|
|
py=periodicity_y,
|
|
min_y=min_y,
|
|
max_y=max_y,
|
|
frames_y=frames_y,
|
|
tol=tolerance,
|
|
metafile=metadata_file,
|
|
outfile=wham_output
|
|
)
|
|
print(cmd)
|
|
os.system(cmd)
|
|
|
|
# read wham to new pmf
|
|
return self.read_pmf(wham_output)
|
|
|
|
def read_pmf(self, pmf_path):
|
|
import pandas as pd
|
|
print("Update PMF from WHAM")
|
|
df = pd.read_csv(pmf_path, delim_whitespace=True, names=['x', 'y', 'e', 'pro'], skiprows=1,
|
|
index_col=None)
|
|
df = df.replace([np.inf, -np.inf], np.nan).dropna(subset=['e'], how='all')
|
|
new_pmf = deepcopy(self.pmf)
|
|
for x in range(new_pmf.shape[0]):
|
|
for y in range(new_pmf.shape[1]):
|
|
lambdax, lambday = self._get_lambdas_for_index((x, y))
|
|
x_selection = (df.x - lambdax).abs() < 0.01
|
|
y_selection = (df.y - lambday).abs() < 0.01
|
|
selected_energies = df[(x_selection) & (y_selection)].e
|
|
if len(selected_energies) == 0:
|
|
new_pmf[x, y] = -1
|
|
else:
|
|
new_pmf[x, y] = selected_energies.iloc[0]
|
|
|
|
return new_pmf
|
|
|
|
|
|
class MyUmbrellaRunner(WHAM2DRunner):
|
|
def after_run_hook(self):
|
|
filename = "tmp/pmf_{}.pdf".format(self.num_iterations)
|
|
print("Writing new pmf to {}".format(filename))
|
|
pmf_to_plot = deepcopy(self.pmf.T)
|
|
pmf_to_plot[pmf_to_plot < 0] = None
|
|
plt.figure()
|
|
plt.imshow(pmf_to_plot, origin="bottom", cmap='jet')
|
|
ticks = [(x,x) for x in [-3, -2, -1, 0, 1, 2, 3]]
|
|
tick_positions = [ self._get_index_for_lambdas(x)[0] for x in ticks ]
|
|
tick_labels = [ str(x[0]) for x in ticks ]
|
|
|
|
plt.xticks(tick_positions, tick_labels)
|
|
plt.yticks(tick_positions, tick_labels)
|
|
|
|
|
|
cb = plt.colorbar(pad=0.1)
|
|
cb.set_label("kJ/mol")
|
|
plt.savefig(filename)
|
|
os.system("cp {} {}".format(filename, "tmp/pmf_current.pdf"))
|
|
|
|
|
|
def simulate_frames(self, lambdas, frames):
|
|
print("{} new simulations:".format(len(lambdas)))
|
|
counter = 0
|
|
|
|
if not os.path.exists("tmp"):
|
|
os.mkdir('tmp')
|
|
|
|
threads = []
|
|
for f in lambdas:
|
|
counter += 1
|
|
print("{}) Simulate lambda1={}, lambda2={}".format(counter, *f))
|
|
command = "bash data/sim.sh {} {} 2>&1 > tmp/run.log".format(*f)
|
|
# print("Running {}".format(command))
|
|
os.system(command)
|
|
|
|
runner = MyUmbrellaRunner()
|
|
runner.WHAM_EXEC = "/opt/wham/wham-2d/wham-2d"
|
|
runner.cvs = np.array([
|
|
(-3, 3, 0.2),
|
|
(-3, 3, 0.2),
|
|
])
|
|
runner.cvs_init = (1.4, -1.4)
|
|
runner.E_min = 10
|
|
runner.E_max = 100
|
|
runner.E_incr = 10
|
|
runner.max_iterations = 30
|
|
|
|
runner.run()
|