From ddd874cd575416d60b727d5a282f7d4b09e1ecd0 Mon Sep 17 00:00:00 2001 From: Daniel Bauer Date: Sat, 16 Jun 2018 10:46:38 +0200 Subject: [PATCH] initial commit --- .idea/vcs.xml | 6 + adaptive_umbrella.py | 257 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 263 insertions(+) create mode 100644 .idea/vcs.xml create mode 100755 adaptive_umbrella.py diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/adaptive_umbrella.py b/adaptive_umbrella.py new file mode 100755 index 0000000..806e237 --- /dev/null +++ b/adaptive_umbrella.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +import numpy as np +from copy import deepcopy +from abc import ABC, abstractmethod + + +#%% +class UmbrellaRunner(ABC): + def init_pmf(self): + """ returns a NxM matrix filled with -1 where N and M are the total number + of lambda frames along the reaction coordinates as determined by + self.lambda_min, self.lambda_max and self.lambda_delta """ + ranges = [] + for dimen in range(len(self.lambda_delta)): + ranges.append(np.arange(self.lambda_min[dimen], self.lambda_max[dimen]+self.lambda_delta[dimen], self.lambda_delta[dimen])) + + mesh = np.meshgrid(*ranges) + pmf = np.zeros(np.dstack(mesh).shape[:-1]) # this is magic + pmf[:] = -1 + return pmf + + def get_lambdas_for_index(self, idx): + """ takes a coordinate tuple and returns corresponding lambda values """ + lambdas = self.lambda_min + idx*self.lambda_delta + return np.round(lambdas, 10) + + def get_index_for_lambdas(self, lambdas): + """ takes a lambda tuple and returns corresponding indeces of the pmf instance variable""" + idx = [] + for dimen in range(len(lambdas)): + r = np.arange(self.lambda_min[dimen], self.lambda_max[dimen]+self.lambda_delta[dimen], self.lambda_delta[dimen]) + + for i in range(len(r)): + if abs(r[i]-lambdas[dimen]) < 0.00001: + idx.append(i) + break + return idx + + # TODO make this work with more then 2 dimensions + def get_root_frames(self, pmf, E_max): + """ returns the index of all positions in the pmf where the energy is + smaller W_max and greater 0 """ + selection = np.where((pmf <= E_max) & (pmf >= 0)) + frames = [] + for i in range(len(selection[0])): + frames.append((selection[0][i], selection[1][i])) + return frames + + # TODO make this work with more then 2 dimensions + def get_new_frames(self, pmf, root_frames): + """ returns a dict of all frames surrounding the root_frames + that have not an assigned energy yet, as well as their corresponding root + frame in the format {new_frame1: root_frame1, new_frame2: root_frame2} """ + + # find all neighboring frames and create a dict that associates them to their + # root frames + new_frames = {} + for frame in root_frames: + for x in [-1, 0, 1]: + for y in [-1, 0, 1]: + new_frame = list(frame) + new_frame[0] += x + new_frame[1] += y + if new_frame[0] < 0 or new_frame[0]+1 > len(self.pmf[0]) or new_frame[1] < 0 or new_frame[1]+1 > len(self.pmf[1]): + continue + try: + old_root = new_frames[tuple(new_frame)] + if old_root is not None: + old_start_energy = pmf[old_root[0], old_root[1]] + new_start_energy = pmf[frame[0], frame[1]] + if old_start_energy > new_start_energy: + new_frames[tuple(new_frame)] = frame + except KeyError: + new_frames[tuple(new_frame)] = frame + + # remove already sampled frames (energy >= 0) + new_frames_list = list(new_frames.keys()) + for idx in range(len(new_frames_list)): + new_frame = new_frames_list[idx] + energy = pmf[new_frame] + if energy >= 0: + del(new_frames[new_frame]) + + return new_frames + + + def main(self): + # get the initial simulation and surrounding frames + root_frames = [self.get_index_for_lambdas(self.lambda_init)] + new_frames = self.get_new_frames(self.pmf, root_frames) + + self.num_iterations = 0 + + # outer main loop: increase E and calculate PMF until E > E_max + while True: + + # stop if max iterations is reached + self.num_iterations += 1 + if(self.max_iterations > 0 and self.num_iterations > self.max_iterations): + print("Max iterations reached ({})".format(self.max_iterations)) + 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("Running simulations") + self.simulate_frames(new_frames, new_lambdas) + + print("Calculating new PMF") + self.pmf = self.wham() + self.after_run_hook() + + while self.E <= self.E_max: + root_frames = self.get_root_frames(self.pmf, self.E) + new_frames = self.get_new_frames(self.pmf, root_frames) + + if len(new_frames) == 0: + self.E += self.E_incr + print("Max energy increased to {} (max={})".format(self.E, self.E_max)) + else: + break + + + def run(self): + self.pmf = self.init_pmf() + self.main() + print("Umbrella sampling finished.") + + @abstractmethod + def simulate_frames(new_frames, new_lambdas): + pass + + @abstractmethod + def wham(): + pass + + def after_run_hook(self): + pass + +if __name__ == "__main__": + from copy import deepcopy + import subprocess + import os + import pandas as pd + import matplotlib.pyplot as plt + class MyUmbrellaRunner(UmbrellaRunner): + + def wham(self): + # copy colvar files + os.system("mkdir -p WHAM") + for f in os.listdir("sim"): + with open("sim/{}/COLVAR".format(f), "r") as i: + with open("WHAM/{}.xvg".format(f), 'w') as o: + for line in i.readlines()[100:]: + o.write(line) + + + # generate metadata.dat + x_vals = [] + y_vals = [] + metadata_file = "WHAM/{}_metadata.dat".format(self.num_iterations) + with open(metadata_file, 'w') as o: + for f in os.listdir("sim"): + prefix, x, y = f.split("_") + x_vals.append(float(x)) + y_vals.append(float(y)) + o.write("WHAM/{}.xvg {} {} {} {}\n".format(f, x,y, 100, 100)) + x_vals = np.array(list(set(x_vals))) + y_vals = np.array(list(set(y_vals))) + + # min_x = x_vals.min() - self.lambda_delta[0] + # min_y = y_vals.min() - self.lambda_delta[1] + # max_x = x_vals.max() + self.lambda_delta[0] + # max_y = y_vals.max() + self.lambda_delta[1] + min_x, min_y = self.lambda_min + max_x, max_y = self.lambda_max + frames_x, frames_y = 1002, 1002 + print("Running WHAM-2d:") + wham_output = "WHAM/{}_freeenergy.dat".format(self.num_iterations) + cmd = "/opt/wham/wham-2d/wham-2d Px=pi {min_x} {max_x} {frames_x} Py=pi {min_y} {max_y} {frames_y} 0.1 298 0 {metafile} {outfile} 0".format( + min_x=min_x, + max_x=max_x, + frames_x=frames_x, + min_y=min_y, + max_y=max_y, + frames_y=frames_y, + metafile=metadata_file, + outfile=wham_output + ) + print(cmd) + os.system(cmd) + + print("Update pmf from wham") + df = pd.read_csv(wham_output, 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 + + def simulate_frames(self, new_frames, new_lambdas): + print("{} new simulations:".format(len(new_lambdas))) + counter = 0 + + threads = [] + for f in new_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) + + 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")) + + + + runner = MyUmbrellaRunner() + runner.lambda_max = np.array((3.1, 3.1)) + runner.lambda_min = -runner.lambda_max + runner.lambda_delta = np.array((0.1, 0.1)) + runner.lambda_init = np.array((1,-1.4)) + runner.E_min = 5 + runner.E_max = 100 + runner.E_incr = 10 + runner.max_iterations = 10 + + runner.run() + + + +