example

adaptiveumbrella/runner.py

@@ -145,31 +145,46 @@ class UmbrellaRunner():
     
     
     def _main(self):
-        # get the initial simulation and surrounding frames
-        root_frames = [self._get_index_for_lambdas(self.cvs_init)]
-        new_frames = self._get_new_frames(self.pmf, self.sample_list, root_frames)
         
         self.num_iterations = 0
 
+        self.E = self.E_min # TODO move this in the loop?
+        
         # outer main loop: increase E and calculate PMF until E > E_max
         while True:
-            self.E = self.E_min
-
-            # 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
-
-            # stop if the pmf is sampled 
-            if len(np.where(self.pmf < 0)) == 0:
-                print("Every window of the PMF appears to be sampled.")
-                return
 
             print("~~~~~~~~~~~~~~~ Iteration {}/{} ~~~~~~~~~~~~~~~~".format(self.num_iterations, self.max_iterations))
+            
+            # find frames to sample
+            if self.num_iterations == 1:
+                # get the initial simulation and surrounding frames
+                root_frames = [self._get_index_for_lambdas(self.cvs_init)]
+                new_frames = self._get_new_frames(self.pmf, self.sample_list, root_frames)
+            else:
+                while self.E <= self.E_max:
+                    # get new frames for energy level
+                    root_frames = self._get_root_frames(self.pmf, self.sample_list, self.E)
+                    new_frames = self._get_new_frames(self.pmf, self.sample_list, root_frames)
+                    
+                    # increase energy level if no frames are found
+                    if len(new_frames) == 0:
+                        self.E += self.E_incr
+                        print("Max energy increased to {} (max={})".format(self.E, self.E_max))
+                    else:
+                        break
+
+            # abort sampling if no new frames could be found
+            if len(new_frames) == 0:
+                print("Sampling aborted. No neighbors found for E_max={}".format(self.E))
+                return
+
+
+            # lambda states for new frames
             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(lambdas, new_frames)
 
@@ -179,17 +194,17 @@ class UmbrellaRunner():
             # update list of sampled windows
             for new_frame in new_frames.keys():
                 self.sample_list[new_frame] = self.num_iterations
+
             self.after_run_hook()
 
-            while self.E <= self.E_max:
+            if(self.max_iterations > 0 and self.num_iterations == self.max_iterations):
-                root_frames = self._get_root_frames(self.pmf, self.sample_list, self.E)
+                print("Max iterations reached ({})".format(self.max_iterations))
-                new_frames = self._get_new_frames(self.pmf, self.sample_list, root_frames)
+                return
 
-                if len(new_frames) == 0:
+            # stop if the pmf is sampled
-                    self.E += self.E_incr
+            if len(np.where(self.pmf < 0)) == 0:
-                    print("Max energy increased to {} (max={})".format(self.E, self.E_max))
+                print("Every window of the PMF appears to be sampled.")
-                else:
+                return
-                    break
 
 
     def run(self):

examples/data/sim.sh

@@ -18,6 +18,6 @@ sed -i -e "s/LAMBDA2/${lambda2}/g" $FOLDER/plumed.dat
 
 cd $FOLDER
 source /usr/local/gromacs/bin/GMXRC
-gmx --quiet mdrun -deffnm topol -plumed plumed.dat -nsteps 100000 2>&1
+gmx --quiet mdrun -deffnm topol -plumed plumed.dat -nsteps 250000 2>&1
 cd ../..
 

examples/example.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 import os
 from copy import deepcopy
 
@@ -110,7 +111,7 @@ class MyUmbrellaRunner(WHAM2DRunner):
     cum_frames = [0]
 
     def after_run_hook(self):
-        filename = "tmp/pmf_{}.pdf".format(self.num_iterations)
+        filename = "tmp/pmf_{}.pdf".format("%02d" % self.num_iterations)
         print("Writing new pmf to {}".format(filename))
         pmf_to_plot = deepcopy(self.pmf)
         pmf_to_plot[pmf_to_plot < 0] = None
@@ -128,13 +129,13 @@ class MyUmbrellaRunner(WHAM2DRunner):
         ax1.set_xlabel("Cycles")
         ax1.set_ylabel("Number of umbrella Windows")
 
-        ticks = [(x,x) for x in [-3, -2, -1, 0, 1, 2, 3]]
+        # ticks = [(x,x) for x in [-3, -2, -1, 0, 1, 2, 3]]
-        tick_positions = [ self._get_index_for_lambdas(x)[0] for x in ticks ]
+        # tick_positions = [ self._get_index_for_lambdas(x)[0] for x in ticks ]
-        tick_labels = [ str(x[0]) for x in ticks ]
+        # tick_labels = [ str(x[0]) for x in ticks ]
-        ax0.set_yticks(tick_positions)
+        # ax0.set_yticks(tick_positions)
-        ax0.set_yticklabels(tick_labels)
+        # ax0.set_yticklabels(tick_labels)
-        ax0.set_xticks(tick_positions)
+        # ax0.set_xticks(tick_positions)
-        ax0.set_xticklabels(tick_labels)
+        # ax0.set_xticklabels(tick_labels)
         ax0.set_ylabel("$\phi$")
         ax0.set_xlabel("$\psi$")
 
@@ -163,13 +164,13 @@ class MyUmbrellaRunner(WHAM2DRunner):
 runner = MyUmbrellaRunner()
 runner.WHAM_EXEC = "/opt/wham/wham-2d/wham-2d"
 runner.cvs = np.array([
-    (-3, 3, 0.2),
+    (-3, 3, 0.3),
-    (-3, 3, 0.2),
+    (-3, 3, 0.3),
 ])
-runner.cvs_init = (-1.4, 1.4)
+runner.cvs_init = (-1.8, 1.8)
-runner.E_min = 2
+runner.E_min = 5
-runner.E_max = 50
+runner.E_max = 100
-runner.E_incr = 5
+runner.E_incr = 10
-runner.max_iterations = 30
+runner.max_iterations = 100
 
 runner.run()