generalized wham2d impl

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"

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

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

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.2),
-    (-3, 3, 0.3),
+    (-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()