inherit from torch.utils.data.Dataset in DataSet.py

deeprank/learn/DataSet.py

@@ -15,13 +15,14 @@
 from deeprank.generate import NormParam, MinMaxParam, NormalizeData
 from deeprank.tools import sparse, pdb2sql
 from tqdm import tqdm
+from torch.utils.data import Dataset
 
 # import torch.utils.data as data_utils
 # The class used to subclass data_utils.Dataset
 # but that conflict with Sphinx that couldn't build the API
 # It's apparently not necessary though and works without subclassing
 
-class DataSet():
+class DataSet(Dataset):
 
 
     def __init__(self,train_database, valid_database=None, test_database = None,
@@ -255,15 +256,15 @@ def process_dataset(self):
         self.get_grid_shape()
 
         # get the input shape
-        self.get_input_shape()
+        self.get_input_shape() # called map_one_molecule()
 
         # get the target ordering
         #self._get_target_ordering()
 
         # get renormalization factor
         if self.normalize_features or self.normalize_targets:
             if self.mapfly:
-                self.compute_norm()
+                self.compute_norm() # called map_one_molecule() again
             else:
                 self.get_norm()
 
@@ -478,7 +479,7 @@ def filter(self,molgrp):
         Args:
             molgrp (str): group name of the molecule in the hdf5 file
         Returns:
-            bool: True if we keep the complex False otherwise
+            bool: True if we keep the complex, False otherwise
         Raises:
             ValueError: If an unsuported condition is provided
         '''
@@ -1068,7 +1069,10 @@ def map_one_molecule(self,fname,mol=None,angle=None,axis=None):
         t0 = time.time()
 
         outtype = 'float32'
-        fh5 = h5py.File(fname,'r')
+        try:
+            fh5 = h5py.File(fname,'r')
+        except:
+            raise IOError(f"{fname} cannot be opened for reading!")
 
         if mol is None:
             mol = list(fh5.keys())[0]
@@ -1095,7 +1099,6 @@ def map_one_molecule(self,fname,mol=None,angle=None,axis=None):
         # close
         fh5.close()
         print(' --> Map one molecule %f sec.' %(time.time()-t0))
-        #sys.exit()
 
         # make sure all the feature have exact same type
         # if they don't  collate_fn in the creation of the minibatch will fail.