initial commit

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

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()
+
+
+
+