fixed a bug in read_epoch_data() in plot_utils.py

DeepRank · May 16, 2020 · 84810d9 · 84810d9
1 parent 70f2f66
commit 84810d9
Showing 1 changed file with 12 additions and 14 deletions.
diff --git a/deeprank/utils/plot_utils.py b/deeprank/utils/plot_utils.py
@@ -46,25 +46,23 @@ def read_epoch_data(DR_h5FL, epoch):
     0  Test  1AVX_ranair-it0_5286       0  0.503823  /home/lixue/DBs/BM5-haddock24/hdf5/000_1AVX.hdf5
     1  Test     1AVX_ti5-itw_354w       1  0.502845  /home/lixue/DBs/BM5-haddock24/hdf5/000_1AVX.hdf5
     """
-    print (f"-> Read epoch data from {DR_h5FL} into df")
+    print (f"-> Read epoch {epoch} data from {DR_h5FL} into df")
 
     # -- 1. read deeprank output data for the specific epoch
     h5 = h5py.File(DR_h5FL, 'r')
+
+    keys = list(h5.keys())
+    last_epoch_key = list(filter(lambda x: 'epoch_' in x, keys))[-1]
+
     if epoch is None:
-        print(f"epoch is not provided. Use the last epoch data.")
-        keys = list(h5.keys())
-        last_epoch_key = list(filter(lambda x: 'epoch_' in x, keys))[-1]
+        print(f"epoch is not provided. Use the last epoch data: {last_epoch_key}.")
+        epoch_key = last_epoch_key
     else:
-        last_epoch_key = 'epoch_%04d' % epoch
-        if last_epoch_key not in h5:
-            print(
-                'Incorrect epcoh name\n Possible options are: ' +
-                ' '.join(
-                    list(
-                        h5.keys())))
-            h5.close()
-            return
-    data = h5[last_epoch_key]
+        epoch_key = 'epoch_%04d' % epoch
+        if epoch_key not in h5:
+            print('Incorrect epoch name. Use the last epoch data: {last_epoch_key}.')
+            epoch_key = last_epoch_key
+    data = h5[epoch_key]
 
     # -- 2. convert into pd.DataFrame
     labels = list(data)  # labels = ['train', 'test', 'valid']