diff --git a/adaptiveumbrella/__init__.py b/adaptiveumbrella/__init__.py index 7ce4628..6516b48 100644 --- a/adaptiveumbrella/__init__.py +++ b/adaptiveumbrella/__init__.py @@ -2,4 +2,4 @@ from __future__ import absolute_import from .runner import UmbrellaRunner __all__ = ['AdaptiveUmbrella'] -__version__ = "0.2.0" +__version__ = "0.3.0" diff --git a/adaptiveumbrella/runner.py b/adaptiveumbrella/runner.py index 69fec9c..8b94d25 100755 --- a/adaptiveumbrella/runner.py +++ b/adaptiveumbrella/runner.py @@ -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)) diff --git a/adaptiveumbrella/wham2d.py b/adaptiveumbrella/wham2d.py new file mode 100644 index 0000000..849894b --- /dev/null +++ b/adaptiveumbrella/wham2d.py @@ -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 diff --git a/examples/example.py b/examples/example.py index 89ba72d..86a78a3 100755 --- a/examples/example.py +++ b/examples/example.py @@ -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()