refactor: clean up GraphDataset and Trainer class

CHANGELOG.md

@@ -66,7 +66,7 @@ Released on June 28, 2022
 * QueryDataset class #53
 * Unit tests for NeuralNet class #86
 * Error message if you pick the wrong metrics #110
-* Unit tests for HDF5DataSet class parameters #82
+* Unit tests for GraphDataset class parameters #82
 * Installation from PyPI in the readme #122
 
 ### Changed
@@ -76,7 +76,7 @@ Released on June 28, 2022
 * `multiprocessing.Queue` has been replaced with `multiprocessing.pool.map` in PreProcessor #56
 * `test_preprocess.py` does not fail anymore on Mac M1 #74
 * It's now possible to pass your own train/test split to NeuralNet class #81
-* HDF5DataSet class now is used in the UX #83
+* GraphDataset class now is used in the UX #83
 * IndexError running `NeuralNet.train()` has been fixed #89
 * pip installation has been fixed
 * Repository has been renamed deeprank-core, and the package deeprankcore #101

README.md

@@ -204,39 +204,39 @@ Data can be split in sets implementing custom splits according to the specific a
 Assuming that the training, validation and testing ids have been chosen (keys of the hdf5 file), then the corresponding graphs can be saved in hdf5 files containing only references (external links) to the original one. For example:
 
 ```python
-from deeprankcore.DataSet import save_hdf5_keys
+from deeprankcore.dataset import save_hdf5_keys
 
 save_hdf5_keys("<original_hdf5_path.hdf5>", train_ids, "<train_hdf5_path.hdf5>")
 save_hdf5_keys("<original_hdf5_path.hdf5>", valid_ids, "<val_hdf5_path.hdf5>")
 save_hdf5_keys("<original_hdf5_path.hdf5>", test_ids, "<test_hdf5_path.hdf5>")
 ```
 
-Now the HDF5DataSet objects can be defined:
+Now the GraphDataset objects can be defined:
 
 ```python
-from deeprankcore.DataSet import HDF5DataSet
+from deeprankcore.dataset import GraphDataset
 
 node_features = ["bsa", "res_depth", "hse", "info_content", "pssm"]
 edge_features = ["distance"]
 
-# Creating HDF5DataSet objects
-dataset_train = HDF5DataSet(
+# Creating GraphDataset objects
+dataset_train = GraphDataset(
     hdf5_path = "<train_hdf5_path.hdf5>",
-    node_feature = node_features,
-    edge_feature = edge_features,
+    node_features = node_features,
+    edge_features = edge_features,
     target = "binary"
 )
-dataset_val = HDF5DataSet(
+dataset_val = GraphDataset(
     hdf5_path = "<val_hdf5_path.hdf5>",
-    node_feature = node_features,
-    edge_feature = edge_features,
+    node_features = node_features,
+    edge_features = edge_features,
     target = "binary"
 
 )
-dataset_test = HDF5DataSet(
+dataset_test = GraphDataset(
     hdf5_path = "<test_hdf5_path.hdf5>",
-    node_feature = node_features,
-    edge_feature = edge_features,
+    node_features = node_features,
+    edge_features = edge_features,
     target = "binary"
 )
 ```
@@ -246,18 +246,18 @@ dataset_test = HDF5DataSet(
 Let's define a Trainer instance, using for example of the already existing GNNs, GINet:
 
 ```python
-from deeprankcore.Trainer import Trainer
+from deeprankcore.trainer import Trainer
 from deeprankcore.ginet import GINet
 from deeprankcore.utils.metrics import OutputExporter, ScatterPlotExporter
 
 metrics_output_directory = "./metrics"
 metrics_exporters = [OutputExporter(metrics_output_directory)]
 
 trainer = Trainer(
+    GINet,
     dataset_train,
     dataset_val,
     dataset_test,
-    GINet,
     batch_size = 64,
     metrics_exporters = metrics_exporters
 )
@@ -297,7 +297,7 @@ def normalized_cut_2d(edge_index, pos):
 
 class CustomNet(torch.nn.Module):
     def __init__(self):
-        super(Net, self).__init__()
+        super().__init__()
         self.conv1 = SplineConv(d.num_features, 32, dim=2, kernel_size=5)
         self.conv2 = SplineConv(32, 64, dim=2, kernel_size=5)
         self.fc1 = torch.nn.Linear(64, 128)
@@ -320,10 +320,10 @@ class CustomNet(torch.nn.Module):
         return F.log_softmax(self.fc2(x), dim=1)
 
 trainer = Trainer(
+    CustomNet,
     dataset_train,
     dataset_val,
     dataset_test,
-    CustomNet,
     batch_size = 64,
     metrics_exporters = metrics_exporters
 )