fix wham2d sampling regions with inf energy

adaptiveumbrella/__init__.py

@@ -3,4 +3,4 @@ from .runner import UmbrellaRunner
 from .wham2d import WHAM2DRunner
 
 __all__ = ['AdaptiveUmbrella']
-__version__ = "0.3.12"
+__version__ = "0.3.13"

adaptiveumbrella/runner.py

@@ -85,9 +85,7 @@ class UmbrellaRunner():
 
     def _get_root_frames(self, pmf, frames, E_max):
         """ returns the index of all positions in the pmf where the energy is
-        smaller E_max"""
+        smaller E_max and that have already been sampled"""
-
-        # select positions of the pmf where E <= E_max and that have already been sampled (frames > 0)
         selection = np.where((pmf <= E_max) & (frames > 0))
         zipped = list(zip(*selection))
 
@@ -132,33 +130,42 @@ class UmbrellaRunner():
         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 the root frame with lowest energy
+        # find all neighboring frames and create a dict 
+        # that associates them to the root frame with lowest energy
+        # new_frame -> corresponding root_frame
         new_frames = {}
-        for frame in root_frames:
+        for root_frame in root_frames:
-            neighbors = self._generate_neighbor_list(frame)
 
-            # remove neighbors if they are not valid (i.e not part of the pmf)
+            # all frames surrounding this root
+            neighbors = self._generate_neighbor_list(root_frame)
+
+            # remove invalid (i.e not part of the pmf)
             neighbors = [n for n in neighbors if self.is_valid_frame(n)]
 
-            # for each neighbor, check if its already in the list and compare root frame energy
+            # for every neighbor, we try
-            for n in neighbors:
+            # to get the root frame energy
+            # see if its already with an associated root in new_frames
+            # if yes, we swap the root frame if the new one has lower energy
+            # if not, KeyError, so the current root is our best candidate so far 
+            for neighbor in neighbors:
                 try:
-                    root_energy = pmf[frame]
+                    root_energy = pmf[root_frame]
-                    old_root = new_frames[n]
+                    old_root = new_frames[neighbor]
                     old_root_energy = pmf[old_root]
                     if root_energy < old_root_energy:
-                        new_frames[n] = frame
+                        new_frames[neighbor] = root_frame
                 except KeyError:
-                    new_frames[n] = frame
+                    new_frames[neighbor] = root_frame
 
-
+        # remove already sampled frames (frames with energy > 0)
-        # remove already sampled frames (frames > 0)
         new_frames_list = list(new_frames.keys())
         for idx in range(len(new_frames_list)):
             new_frame = new_frames_list[idx]
             if frames[new_frame] > 0:
                 del(new_frames[new_frame])
 
+        # also remove all frames where the root has no energy
+                
         return new_frames
     
     

adaptiveumbrella/wham2d.py

@@ -101,6 +101,8 @@ class WHAM2DRunner(UmbrellaRunner):
     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)
+        # if e is inf, that means this region is unsampled. we cannot really 
+        # tell its energy and set it to NaN 
         df = df.replace([np.inf, -np.inf], np.nan).dropna(subset=['e'], how='all')
         return df
 
@@ -112,15 +114,13 @@ class WHAM2DRunner(UmbrellaRunner):
                 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
+                    self.pmf[x,y] = np.inf
-
+                else:
                     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):
         metafile_path = self.create_metadata_file()
         wham_pmf_file = self.get_wham_output_file()