Fix bug in Collector with Graph data (#456)

* Fix bug in Collector with Graph data
* Add comments in DataModule class and bug fix in collate

pina/data/dataset.py

@@ -1,10 +1,12 @@
 """
 This module provide basic data management functionalities
 """
+import functools
 import torch
 from torch.utils.data import Dataset
 from abc import abstractmethod
-from torch_geometric.data import Batch
+from torch_geometric.data import Batch, Data
+from pina import LabelTensor
 
 
 class PinaDatasetFactory:
@@ -62,7 +64,7 @@ class PinaTensorDataset(PinaDataset):
         if automatic_batching:
             self._getitem_func = self._getitem_int
         else:
-            self._getitem_func = self._getitem_list
+            self._getitem_func = self._getitem_dummy
 
     def _getitem_int(self, idx):
         return {
@@ -82,7 +84,7 @@ class PinaTensorDataset(PinaDataset):
         return to_return_dict
 
     @staticmethod
-    def _getitem_list(idx):
+    def _getitem_dummy(idx):
         return idx
 
     def get_all_data(self):
@@ -102,15 +104,56 @@ class PinaTensorDataset(PinaDataset):
         }
 
 
+class PinaBatch(Batch):
+    """
+    Add extract function to torch_geometric Batch object
+    """
+    def __init__(self):
+
+        super().__init__(self)
+
+    def extract(self, labels):
+        """
+        Perform extraction of labels on node features (x)
+
+        :param labels: Labels to extract 
+        :type labels: list[str] | tuple[str] | str 
+        :return: Batch object with extraction performed on x
+        :rtype: PinaBatch
+        """
+        self.x = self.x.extract(labels)
+        return self
+
+
 class PinaGraphDataset(PinaDataset):
 
     def __init__(self, conditions_dict, max_conditions_lengths,
                  automatic_batching):
         super().__init__(conditions_dict, max_conditions_lengths)
+        self.in_labels = {}
+        self.out_labels = None
         if automatic_batching:
             self._getitem_func = self._getitem_int
         else:
-            self._getitem_func = self._getitem_list
+            self._getitem_func = self._getitem_dummy
+
+        ex_data = conditions_dict[list(conditions_dict.keys())[
+            0]]['input_points'][0]
+        for name, attr in ex_data.items():
+            if isinstance(attr, LabelTensor):
+                self.in_labels[name] = attr.stored_labels
+        ex_data = conditions_dict[list(conditions_dict.keys())[
+            0]]['output_points'][0]
+        if isinstance(ex_data, LabelTensor):
+            self.out_labels = ex_data.labels
+
+        self._create_graph_batch_from_list = self._labelise_batch(
+            self._base_create_graph_batch_from_list) if self.in_labels \
+            else self._base_create_graph_batch_from_list
+
+        self._create_output_batch = self._labelise_tensor(
+            self._base_create_output_batch) if self.out_labels is not None \
+            else self._base_create_output_batch
 
     def fetch_from_idx_list(self, idx):
         to_return_dict = {}
@@ -119,17 +162,24 @@ class PinaGraphDataset(PinaDataset):
             condition_len = self.conditions_length[condition]
             if self.length > condition_len:
                 cond_idx = [idx % condition_len for idx in cond_idx]
-            to_return_dict[condition] = {k: Batch.from_data_list([
-                v[i] for i in cond_idx])
-                                         if isinstance(v, list)
-                                         else v[
-                                             cond_idx].reshape(
-                                                 -1, *v[cond_idx].shape[2:])
-                                         for k, v in data.items()
-                                         }
+            to_return_dict[condition] = {
+                k: self._create_graph_batch_from_list([v[i] for i in idx])
+                if isinstance(v, list)
+                else self._create_output_batch(v[idx])
+                for k, v in data.items()
+            }
+
         return to_return_dict
 
-    def _getitem_list(self, idx):
+    def _base_create_graph_batch_from_list(self, data):
+        batch = PinaBatch.from_data_list(data)
+        return batch
+
+    def _base_create_output_batch(self, data):
+        out = data.reshape(-1, *data.shape[2:])
+        return out
+
+    def _getitem_dummy(self, idx):
         return idx
 
     def _getitem_int(self, idx):
@@ -144,3 +194,31 @@ class PinaGraphDataset(PinaDataset):
 
     def __getitem__(self, idx):
         return self._getitem_func(idx)
+
+    def _labelise_batch(self, func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            batch = func(*args, **kwargs)
+            for k, v in self.in_labels.items():
+                tmp = batch[k]
+                tmp.labels = v
+                batch[k] = tmp
+            return batch
+        return wrapper
+
+    def _labelise_tensor(self, func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            out = func(*args, **kwargs)
+            if isinstance(out, LabelTensor):
+                out.labels = self.out_labels
+            return out
+        return wrapper
+
+    def create_graph_batch(self, data):
+        """
+        # TODO
+        """
+        if isinstance(data[0], Data):
+            return self._create_graph_batch_from_list(data)
+        return self._create_output_batch(data)