From 4f38e2a780c0400bed97b8c786850cc01c3e7fa6 Mon Sep 17 00:00:00 2001
From: daniel <bauer@cbs.tu-darmstadt.de>
Date: Mon, 6 Jul 2020 11:33:28 +0200
Subject: [PATCH] reproduceable tmp pmfs

---
 adaptiveumbrella/runner.py | 20 +++++++++++++++-----
 adaptiveumbrella/wham2d.py | 18 +++++++-----------
 2 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/adaptiveumbrella/runner.py b/adaptiveumbrella/runner.py
index fefe551..47b0b09 100755
--- a/adaptiveumbrella/runner.py
+++ b/adaptiveumbrella/runner.py
@@ -47,7 +47,6 @@ class UmbrellaRunner():
         pmf = np.empty(shape)
         pmf.fill(-1)
         return pmf
-
         
     def _get_lambdas_for_index(self, idx):
         """ takes a coordinate tuple of the pmf and returns corresponding lambda values """
@@ -68,10 +67,21 @@ class UmbrellaRunner():
                 if abs(r[i]-lambdas[dimen]) < 0.00001:
                     idx.append(i)
                     break
-        if len(idx) == len(lambdas):
-            return tuple(idx)
-        else: # if len differs, theres no index for every dimension
-            raise ValueError("{} has no index.".format(lambdas))
+            if not len(idx)-1 == dimen:
+                raise ValueError(f"{lambdas} has no index because {lambdas[dimen]} is not in {r}.")
+
+        return tuple(idx)
+        
+    def _get_sampled_lambdas(self, step=None):
+        """ returns an array of all sampled lambdas. If step is given, only return
+        lambdas for this step"""
+        if step is None:
+            sampled_coords = np.where(self.sample_list > 0)
+        else:
+            sampled_coords = np.where(self.sample_list == step)
+        sampled_coords = np.array(sampled_coords).T
+        
+        return np.array([self._get_lambdas_for_index(cvs) for cvs in sampled_coords])
 
     def _get_root_frames(self, pmf, frames, E_max):
         """ returns the index of all positions in the pmf where the energy is
diff --git a/adaptiveumbrella/wham2d.py b/adaptiveumbrella/wham2d.py
index 46de7c3..b112eac 100644
--- a/adaptiveumbrella/wham2d.py
+++ b/adaptiveumbrella/wham2d.py
@@ -33,18 +33,14 @@ class WHAM2DRunner(UmbrellaRunner):
         """ 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):
-                filepath = os.path.join(self.simulation_folder, file, "COLVAR")
-                if not os.path.exists(filepath) and self.verbose:
-                    print("Not found: {}".format(filepath))
+            for x, y in self._get_sampled_lambdas():
+                colvar_file = os.path.join(self.simulation_folder, f"umb_{x}_{y}", "COLVAR")
+                if not os.path.exists(colvar_file) and self.verbose:
+                    print("Not found: {}".format(colvar_file))
                     continue
-                prefix, x, y = file.split("_")
-                index = self._get_index_for_lambdas((float(x),float(y)))
-                if self.sample_list[index] != 0:
-                    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']
-                    ))
+                out.write("{file}\t{x}\t{y}\t{fc_x}\t{fc_y}\n".format(
+                    file=colvar_file, x=x, y=y, fc_x=self.whamconfig['fc_x'], fc_y=self.whamconfig['fc_y']
+                ))
         return path