Fix bug in Collector with Graph data

Author: FilippoOlivo

pina/collector.py

@@ -1,3 +1,4 @@
+from .graph import Graph
 from .utils import check_consistency
 
 
@@ -52,6 +53,8 @@ def store_fixed_data(self):
                 # get data
                 keys = condition.__slots__
                 values = [getattr(condition, name) for name in keys]
+                values = [value.data if isinstance(
+                    value, Graph) else value for value in values]
                 self.data_collections[condition_name] = dict(zip(keys, values))
                 # condition now is ready
                 self._is_conditions_ready[condition_name] = True

pina/data/data_module.py

@@ -3,11 +3,11 @@
 from lightning.pytorch import LightningDataModule
 import torch
 from ..label_tensor import LabelTensor
-from torch.utils.data import DataLoader, BatchSampler, SequentialSampler, \
-    RandomSampler
+from torch.utils.data import DataLoader, SequentialSampler, RandomSampler
 from torch.utils.data.distributed import DistributedSampler
 from .dataset import PinaDatasetFactory
 from ..collector import Collector
+from torch_geometric.data import Data, Batch
 
 
 class DummyDataloader:
@@ -86,9 +86,8 @@ def _collate_standard_dataloader(self, batch):
                     single_cond_dict[arg] = LabelTensor.stack(data_list)
                 elif isinstance(data_list[0], torch.Tensor):
                     single_cond_dict[arg] = torch.stack(data_list)
-                else:
-                    raise NotImplementedError(
-                        f"Data type {type(data_list[0])} not supported")
+                elif isinstance(data_list[0], Data):
+                    single_cond_dict[arg] = Batch.from_data_list(data_list)
             batch_dict[condition_name] = single_cond_dict
         return batch_dict
 
@@ -125,7 +124,7 @@ def __init__(self,
                  batch_size=None,
                  shuffle=True,
                  repeat=False,
-                 automatic_batching=False,
+                 automatic_batching=None,
                  num_workers=0,
                  pin_memory=False,
                  ):
@@ -158,7 +157,6 @@ def __init__(self,
         logging.debug('Start initialization of Pina DataModule')
         logging.info('Start initialization of Pina DataModule')
         super().__init__()
-        self.automatic_batching = automatic_batching
         self.batch_size = batch_size
         self.shuffle = shuffle
         self.repeat = repeat
@@ -192,6 +190,10 @@ def __init__(self,
         collector = Collector(problem)
         collector.store_fixed_data()
         collector.store_sample_domains()
+
+        self.automatic_batching = self._set_automatic_batching_option(
+            collector, automatic_batching)
+
         if batch_size is None and num_workers != 0:
             warnings.warn(
                 "Setting num_workers when batch_size is None has no effect on "
@@ -393,6 +395,27 @@ def _check_slit_sizes(train_size, test_size, val_size, predict_size):
         if abs(train_size + test_size + val_size + predict_size - 1) > 1e-6:
             raise ValueError("The sum of the splits must be 1")
 
+    @staticmethod
+    def _set_automatic_batching_option(collector, automatic_batching):
+        """
+        Determines whether automatic batching should be enabled.
+
+        If all 'input_points' in the collector's data collections are 
+        tensors (torch.Tensor or LabelTensor), it respects the provided
+        `automatic_batching` value; otherwise, mainly in the Graph scenario, 
+        it forces automatic batching on.
+
+        :param Collector collector: Collector object with contains all data 
+        retrieved from input conditions
+        :param bool automatic_batching : If the user wants to enable automatic
+        batching or not
+        """
+        if all(isinstance(v['input_points'], (torch.Tensor, LabelTensor))
+                for v in collector.data_collections.values()):
+            return automatic_batching if automatic_batching is not None \
+                else False
+        return True
+
     @property
     def input_points(self):
         """

pina/data/dataset.py

@@ -2,9 +2,11 @@
 This module provide basic data management functionalities
 """
 import torch
+import functools
 from torch.utils.data import Dataset
 from abc import abstractmethod
 from torch_geometric.data import Batch
+from pina import LabelTensor
 
 
 class PinaDatasetFactory:
@@ -107,10 +109,25 @@ class PinaGraphDataset(PinaDataset):
     def __init__(self, conditions_dict, max_conditions_lengths,
                  automatic_batching):
         super().__init__(conditions_dict, max_conditions_lengths)
-        if automatic_batching:
-            self._getitem_func = self._getitem_int
-        else:
-            self._getitem_func = self._getitem_list
+        self.in_labels = {}
+        self.out_labels = None
+        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,18 +136,22 @@ def fetch_from_idx_list(self, idx):
             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, cond_idx)
+                if isinstance(v, list)
+                else self._create_output_batch(v, cond_idx)
+                for k, v in data.items()
+            }
+
         return to_return_dict
 
-    def _getitem_list(self, idx):
-        return idx
+    def _base_create_graph_batch_from_list(self, data, idx):
+        batch = Batch.from_data_list([data[i] for i in idx])
+        return batch
+
+    def _base_create_output_batch(self, data, idx):
+        out = data[idx].reshape(-1, *data[idx].shape[2:])
+        return out
 
     def _getitem_int(self, idx):
         return {
@@ -143,4 +164,25 @@ def get_all_data(self):
         return self.fetch_from_idx_list(index)
 
     def __getitem__(self, idx):
-        return self._getitem_func(idx)
+        return self._getitem_int(idx) if isinstance(idx, int) else \
+            self.fetch_from_idx_list(idx=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

pina/graph.py

@@ -108,16 +108,14 @@ def __init__(
                                                     x)
 
         # Perform the graph construction
-        self._build_graph_list(x, pos, edge_index, edge_attr, additional_params)
+        self._build_graph_list(
+            x, pos, edge_index, edge_attr, additional_params)
 
     def _build_graph_list(self, x, pos, edge_index, edge_attr,
                           additional_params):
         for i, (x_, pos_, edge_index_) in enumerate(zip(x, pos, edge_index)):
-            if isinstance(x_, LabelTensor):
-                x_ = x_.tensor
             add_params_local = {k: v[i] for k, v in additional_params.items()}
             if edge_attr is not None:
-
                 self.data.append(Data(x=x_, pos=pos_, edge_index=edge_index_,
                                       edge_attr=edge_attr[i],
                                       **add_params_local))
@@ -165,7 +163,8 @@ def _check_input_consistency(x, pos, edge_index=None):
         # If edge_index is a 3D tensor, we split it into a list of 2D tensors
         if edge_index is not None:
             if isinstance(edge_index, torch.Tensor) and edge_index.ndim == 3:
-                edge_index = [edge_index[i] for i in range(edge_index.shape[0])]
+                edge_index = [edge_index[i]
+                              for i in range(edge_index.shape[0])]
             elif not (isinstance(edge_index, list) and all(
                     t.ndim == 2 for t in edge_index)) and not (
                     isinstance(edge_index,
@@ -219,7 +218,7 @@ def _check_and_build_edge_attr(self, edge_attr, build_edge_attr, data_len,
             if isinstance(edge_attr, list):
                 if len(edge_attr) != data_len:
                     raise TypeError("edge_attr must have the same length as x "
-                                     "and pos.")
+                                    "and pos.")
             return [edge_attr] * data_len
 
         if build_edge_attr:
@@ -258,6 +257,8 @@ def _radius_graph(points, r):
         """
         dist = torch.cdist(points, points, p=2)
         edge_index = torch.nonzero(dist <= r, as_tuple=False).t()
+        if isinstance(edge_index, LabelTensor):
+            edge_index = edge_index.tensor
         return edge_index
 
 
@@ -293,4 +294,6 @@ def _knn_graph(points, k):
         row = torch.arange(points.size(0)).repeat_interleave(k)
         col = knn_indices.flatten()
         edge_index = torch.stack([row, col], dim=0)
+        if isinstance(edge_index, LabelTensor):
+            edge_index = edge_index.tensor
         return edge_index

pina/trainer.py

@@ -105,8 +105,6 @@ def __init__(self,
         # checking compilation and automatic batching
         if compile is None or sys.platform == "win32":
             compile = False
-        if automatic_batching is None:
-            automatic_batching = False
 
         # set attributes
         self.compile = compile

pina/utils.py

@@ -48,7 +48,8 @@ def labelize_forward(forward, input_variables, output_variables):
     :type output_variables: list[str] | tuple[str]
     """
     def wrapper(x):
-        x = x.extract(input_variables)
+        if isinstance(x, LabelTensor):
+            x = x.extract(input_variables)
         output = forward(x)
         # keep it like this, directly using LabelTensor(...) raises errors
         # when compiling the code