diff --git a/adaptive_umbrella.py b/adaptive_umbrella.py index 2017d0c..e877933 100755 --- a/adaptive_umbrella.py +++ b/adaptive_umbrella.py @@ -5,6 +5,26 @@ from copy import deepcopy #%% class UmbrellaRunner(): + """ + Runner class for adaptive umbrella sampling + + Attributes: + cvs (numpy array): A multidimensional array of all collective variables (cvs) in the form [cv_min, cv_max, cv_delta] + Example: to sample a 2 dimensional grid where cv1 and cv2 go from 0 to 3 with stepsize 0.5, use + np.array([[0,3,1],[0,3,1]]) + cvs_init (tuple): Starting coordinates for umbrella sampling. Must be a tuple with the same dimension as cvs + E_min (float, default=0): Starting energy for exporative umbrella sampling + E_max (float, default=inf): Final energy. Umbrella sampling is stopped if no frames with E < E_max are found + E_incr (float, default=1): E_min is incremented by this until E_max is reached + max_iterations (int, default=-1): Max. number of iterations before umbrella sampling stops. -1 for infinite sampling + + """ + + def __init__(self): + self.max_iterations = -1 + self.E_min = 0 + self.E_max = np.inf + self.E_incr = 1 def _get_pmf_shape(self): """ returns the shape of the pmf according to the cvs """ @@ -131,7 +151,7 @@ class UmbrellaRunner(): def _main(self): # get the initial simulation and surrounding frames - root_frames = [self._get_index_for_lambdas(self.lambda_init)] + root_frames = [self._get_index_for_lambdas(self.cvs_init)] new_frames = self._get_new_frames(self.pmf, root_frames) self.num_iterations = 0 @@ -146,15 +166,15 @@ class UmbrellaRunner(): return self.E = self.E_min - print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~") - print("Iteration: {} (max={})".format(self.num_iterations, self.max_iterations)) - new_lambdas = [self._get_lambdas_for_index(x) for x in new_frames.keys()] - + print("~~~~~~~~~~~~~~~ Iteration {}/{} ~~~~~~~~~~~~~~~~".format(self.num_iterations, self.max_iterations)) + lambdas = dict([(self._get_lambdas_for_index(x), self._get_lambdas_for_index(y)) for x,y in new_frames.items()]) + + self.pre_run_hook() print("Running simulations") - self.simulate_frames(new_frames, new_lambdas) + self.simulate_frames(lambdas, new_frames) print("Calculating new PMF") - self.pmf = self.wham() + self.pmf = self.calculate_new_pmf() self.after_run_hook() while self.E <= self.E_max: @@ -169,20 +189,124 @@ class UmbrellaRunner(): def run(self): + # initialize the pmf self.pmf = self._init_pmf() - self._main() - print("Umbrella sampling finished.") - - def simulate_frames(self, new_frames, new_lambdas): - print("TODO Implement me") - def wham(self): - print("TODO Implement me") - - def after_run_hook(self): + # start the simulation/evaluation loop + self._main() + print("Finished.") + + def simulate_frames(self, frames, lambdas): + """ This should be implemented to run the simulation for each frame in ```frames``` + with ```lambdas```. Both variables are dictionaries where new values are keys and root values are values + """ + print("New lambda values to simulate: {}".format(lambdas)) + print("(You should implement `simulate_frames` method yourself.)") + + def calculate_new_pmf(self): + """ This should be overwritten to calculate the PMF based on the new simulations and + return a pmf of correct shape/spacing according to the cvs""" pass + def after_run_hook(self): + """ This can be implemented to hook into the simulation cycle after the reevaluation of the pmf """ + pass + + def pre_run_hook(self): + """ This can be implemented to hook into the simulation cycle before the simulation runs """ + pass + +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 = "simulations" + print("Collecting sampling data from simulations folder") + + # collect COLVARs + wham_dir = "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("WHAM/{}.xvg {} {} {} {}\n".format(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 diff --git a/example.py b/example.py new file mode 100644 index 0000000..027d93d --- /dev/null +++ b/example.py @@ -0,0 +1,51 @@ +from adaptive_umbrella import WHAM2DRunner + +from copy import deepcopy +import os +import numpy as np +import matplotlib.pyplot as plt + + +class MyUmbrellaRunner(WHAM2DRunner): + def after_run_hook(self): + filename = "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="lower", cmap='jet') + cb = plt.colorbar(pad=0.1) + cb.set_label("kJ/mol") + plt.savefig(filename) + os.system("cp {} {}".format(filename, "pmf_current.pdf")) + + def simulate_frames(self, lambdas, frames): + print("{} new simulations:".format(len(lambdas))) + counter = 0 + + threads = [] + for f in lambdas: + counter += 1 + if os.path.exists("sim/sim_{}_{}/COLVAR".format(*f)): + print("{}) Skipping lambdas={}/{}: COLVAR exists".format(counter, *f)) + continue + + print("{}) Simulate lambda1={}, lambda2={}".format(counter, *f)) + command = "bash sim.sh {} {} 2>&1 > 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.1, 3.1, 0.1), + (-3.1, 3.1, 0.1), +]) +runner.cvs_init = (1, -1.4) +runner.E_min = 10 +runner.E_max = 200 +runner.E_incr = 10 +runner.max_iterations = 25 + +runner.run()