generalized wham2d impl
This commit is contained in:
Daniel Bauer
@@ -2,4 +2,4 @@ from __future__ import absolute_import
|
||||
from .runner import UmbrellaRunner
|
||||
|
||||
__all__ = ['AdaptiveUmbrella']
|
||||
__version__ = "0.2.0"
|
||||
__version__ = "0.3.0"
|
||||
|
||||
@@ -155,7 +155,7 @@ class UmbrellaRunner():
|
||||
# outer main loop: increase E and calculate PMF until E > E_max
|
||||
while True:
|
||||
self.num_iterations += 1
|
||||
if reset_E:
|
||||
if self.reset_E:
|
||||
self.E = self.E_min
|
||||
|
||||
print("~~~~~~~~~~~~~~~ Iteration {}/{} ~~~~~~~~~~~~~~~~".format(self.num_iterations, self.max_iterations))
|
||||
|
||||
103
adaptiveumbrella/wham2d.py
Normal file
103
adaptiveumbrella/wham2d.py
Normal file
@@ -0,0 +1,103 @@
|
||||
import os
|
||||
import subprocess
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
import math
|
||||
|
||||
# ignore SettingsWithCopyWarning
|
||||
pd.options.mode.chained_assignment = None
|
||||
|
||||
from adaptiveumbrella 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'
|
||||
self.tmp_folder = "tmp/WHAM"
|
||||
self.simulation_folder = "tmp/simulations"
|
||||
|
||||
if not os.path.exists(self.tmp_folder):
|
||||
os.makedirs(self.tmp_folder)
|
||||
|
||||
def create_metadata_file(self):
|
||||
""" create the metadata file for wham-2d """
|
||||
path = os.path.join(self.tmp_folder, "{}_metadata.dat".format(self.num_iterations))
|
||||
with open(path, 'w') as out:
|
||||
for file in os.listdir(self.simulation_folder):
|
||||
prefix, x, y = file.split("_")
|
||||
colvar_file = os.path.join(file, 'COLVAR')
|
||||
out.write("{file}\t{x}\t{y}\t{fc_x}\t{fc_y}\n".format(
|
||||
file=os.path.join(self.simulation_folder, colvar_file), x=x, y=y, fc_x=self.whamconfig['fc_x'], fc_y=self.whamconfig['fc_y']
|
||||
))
|
||||
return path
|
||||
|
||||
|
||||
def get_wham_output_file(self):
|
||||
""" Output file for wham-2d """
|
||||
return os.path.join(self.tmp_folder, 'freeenergy_tmp.dat')
|
||||
|
||||
def run_wham2d(self, metafile_path, output_path):
|
||||
""" Runs wham-2d with the given parameters. See http://membrane.urmc.rochester.edu/sites/default/files/wham/doc.html """
|
||||
|
||||
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=self.whamconfig['Px'],
|
||||
min_x=self.whamconfig['hist_min_x'],
|
||||
max_x=self.whamconfig['hist_max_x'],
|
||||
frames_x=self.whamconfig['num_bins_x'],
|
||||
py=self.whamconfig['Py'],
|
||||
min_y=self.whamconfig['hist_min_y'],
|
||||
max_y=self.whamconfig['hist_max_y'],
|
||||
frames_y=self.whamconfig['num_bins_y'],
|
||||
tol=self.whamconfig['tolerance'],
|
||||
metafile=metafile_path,
|
||||
outfile=output_path
|
||||
)
|
||||
print(cmd)
|
||||
try:
|
||||
subprocess.call(cmd, shell=True)
|
||||
except OSError:
|
||||
print("wham failed.")
|
||||
exit(1)
|
||||
|
||||
def load_wham_pmf(self, wham_file):
|
||||
""" Load the new pmf into a pandas dataframe """
|
||||
df = pd.read_csv(wham_file, 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')
|
||||
return df
|
||||
|
||||
def update_pmf(self, wham_pmf):
|
||||
""" Update the internal pmf representation from the pmf generated with wham """
|
||||
# map all points of self.pmf to one point of the wham_pmf and update self.pmf accordingly
|
||||
for x in range(self.pmf.shape[0]):
|
||||
for y in range(self.pmf.shape[1]):
|
||||
lambdax, lambday = self._get_lambdas_for_index((x, y))
|
||||
reduced = wham_pmf[ (abs(wham_pmf.x-lambdax) < self.cvs[0][2]) & (abs(wham_pmf.y-lambday) < self.cvs[1][2]) ]
|
||||
if len(reduced) == 0:
|
||||
continue
|
||||
|
||||
reduced['dist'] = reduced.apply(lambda row: np.linalg.norm((lambdax-row['x'], lambday-row['y'])), axis=1)
|
||||
min = reduced[reduced.dist == reduced.dist.min()]
|
||||
min_row = reduced[(reduced.x == min.x.iloc[0]) & (reduced.y == min.y.iloc[0])]
|
||||
self.pmf[x, y] = min_row.e.iloc[0]
|
||||
|
||||
|
||||
|
||||
def calculate_new_pmf(self):
|
||||
print("Running wham-2d")
|
||||
metafile_path = self.create_metadata_file()
|
||||
wham_pmf_file = self.get_wham_output_file()
|
||||
self.run_wham2d(metafile_path, wham_pmf_file)
|
||||
|
||||
wham_pmf = self.load_wham_pmf(wham_pmf_file)
|
||||
self.update_pmf(wham_pmf)
|
||||
return self.pmf
|
||||
@@ -6,108 +6,20 @@ import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
import sys
|
||||
|
||||
from adaptiveumbrella.wham2d import WHAM2DRunner
|
||||
|
||||
sys.path.append('..')
|
||||
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]
|
||||
min_x = -3
|
||||
max_x = 3
|
||||
min_y = -3
|
||||
max_y = 3
|
||||
|
||||
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 __init__(self):
|
||||
WHAM2DRunner.__init__(self)
|
||||
|
||||
cum_frames = [0]
|
||||
|
||||
def after_run_hook(self):
|
||||
@@ -163,14 +75,29 @@ class MyUmbrellaRunner(WHAM2DRunner):
|
||||
|
||||
runner = MyUmbrellaRunner()
|
||||
runner.WHAM_EXEC = "/opt/wham/wham-2d/wham-2d"
|
||||
runner.whamconfig = {
|
||||
'Px': 'pi',
|
||||
'hist_min_x': -3,
|
||||
'hist_max_x': 3,
|
||||
'num_bins_x': 100,
|
||||
'Py': 'pi',
|
||||
'hist_min_y': -3,
|
||||
'hist_max_y': 3,
|
||||
'num_bins_y': 100,
|
||||
'tolerance': 0.1,
|
||||
'fc_x': 100,
|
||||
'fc_y': 100
|
||||
}
|
||||
|
||||
|
||||
runner.cvs = np.array([
|
||||
(-3, 3, 0.3),
|
||||
(-3, 3, 0.3),
|
||||
(-3, 3, 0.2),
|
||||
(-3, 3, 0.2),
|
||||
])
|
||||
runner.cvs_init = (-1.8, 1.8)
|
||||
runner.cvs_init = (0, 0)
|
||||
runner.E_min = 5
|
||||
runner.E_max = 100
|
||||
runner.E_incr = 10
|
||||
runner.max_iterations = 100
|
||||
runner.max_iterations = 5
|
||||
|
||||
runner.run()
|
||||
|
||||
Reference in New Issue
Block a user