Refactor datasets and implement LabelBatch

Author: FilippoOlivo

pina/data/dataset.py

@@ -2,12 +2,10 @@
 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, Data
-from pina import LabelTensor
+from torch.utils.data import Dataset
+from torch_geometric.data import Data
+from ..graph import Graph, LabelBatch
 
 
 class PinaDatasetFactory:
@@ -19,38 +17,53 @@ class PinaDatasetFactory:
     """
 
     def __new__(cls, conditions_dict, **kwargs):
+        # Check if conditions_dict is empty
         if len(conditions_dict) == 0:
             raise ValueError("No conditions provided")
-        if all(
-            [
-                isinstance(v["input"], torch.Tensor)
-                for v in conditions_dict.values()
-            ]
-        ):
-            return PinaTensorDataset(conditions_dict, **kwargs)
-        elif all(
-            [isinstance(v["input"], list) for v in conditions_dict.values()]
-        ):
+
+        # Check is a Graph is present in the conditions
+        is_graph = cls._is_graph_dataset(conditions_dict)
+        if is_graph:
+            # If a Graph is present, return a PinaGraphDataset
             return PinaGraphDataset(conditions_dict, **kwargs)
-        raise ValueError(
-            "Conditions must be either torch.Tensor or list of Data " "objects."
-        )
+        # If no Graph is present, return a PinaTensorDataset
+        return PinaTensorDataset(conditions_dict, **kwargs)
+
+    @staticmethod
+    def _is_graph_dataset(conditions_dict):
+        for v in conditions_dict.values():
+            for cond in v.values():
+                if isinstance(cond, (Data, Graph, list)):
+                    return True
+        return False
 
 
 class PinaDataset(Dataset):
     """
     Abstract class for the PINA dataset
     """
 
-    def __init__(self, conditions_dict, max_conditions_lengths):
+    def __init__(
+        self, conditions_dict, max_conditions_lengths, automatic_batching
+    ):
+        # Store the conditions dictionary
         self.conditions_dict = conditions_dict
+        # Store the maximum number of conditions to consider
         self.max_conditions_lengths = max_conditions_lengths
+        # Store length of each condition
         self.conditions_length = {
             k: len(v["input"]) for k, v in self.conditions_dict.items()
         }
+        # Store the maximum length of the dataset
         self.length = max(self.conditions_length.values())
+        # Dynamically set the getitem function based on automatic batching
+        if automatic_batching:
+            self._getitem_func = self._getitem_int
+        else:
+            self._getitem_func = self._getitem_dummy
 
     def _get_max_len(self):
+        """"""
         max_len = 0
         for condition in self.conditions_dict.values():
             max_len = max(max_len, len(condition["input"]))
@@ -59,50 +72,66 @@ def _get_max_len(self):
     def __len__(self):
         return self.length
 
-    @abstractmethod
-    def __getitem__(self, item):
-        pass
-
-
-class PinaTensorDataset(PinaDataset):
-    def __init__(
-        self, conditions_dict, max_conditions_lengths, automatic_batching
-    ):
-        super().__init__(conditions_dict, max_conditions_lengths)
+    def __getitem__(self, idx):
+        return self._getitem_func(idx)
 
-        if automatic_batching:
-            self._getitem_func = self._getitem_int
-        else:
-            self._getitem_func = self._getitem_dummy
+    def _getitem_dummy(self, idx):
+        # If automatic batching is disabled, return the data at the given index
+        return idx
 
     def _getitem_int(self, idx):
+        # If automatic batching is enabled, return the data at the given index
         return {
             k: {k_data: v[k_data][idx % len(v["input"])] for k_data in v.keys()}
             for k, v in self.conditions_dict.items()
         }
 
+    def get_all_data(self):
+        """
+        Return all data in the dataset
+
+        :return: All data in the dataset
+        :rtype: dict
+        """
+        index = list(range(len(self)))
+        return self.fetch_from_idx_list(index)
+
     def fetch_from_idx_list(self, idx):
+        """
+        Return data from the dataset given a list of indices
+
+        :param idx: List of indices
+        :type idx: list
+        :return: Data from the dataset
+        :rtype: dict
+        """
         to_return_dict = {}
         for condition, data in self.conditions_dict.items():
+            # Get the indices for the current condition
             cond_idx = idx[: self.max_conditions_lengths[condition]]
+            # Get the length of the current condition
             condition_len = self.conditions_length[condition]
+            # If the length of the dataset is greater than the length of the
+            # current condition, repeat the indices
             if self.length > condition_len:
                 cond_idx = [idx % condition_len for idx in cond_idx]
-            to_return_dict[condition] = {
-                k: v[cond_idx] for k, v in data.items()
-            }
+            # Retrieve the data from the current condition
+            to_return_dict[condition] = self._retrive_data(data, cond_idx)
         return to_return_dict
 
-    @staticmethod
-    def _getitem_dummy(idx):
-        return idx
+    @abstractmethod
+    def _retrive_data(self, data, idx_list):
+        pass
 
-    def get_all_data(self):
-        index = [i for i in range(len(self))]
-        return self.fetch_from_idx_list(index)
 
-    def __getitem__(self, idx):
-        return self._getitem_func(idx)
+class PinaTensorDataset(PinaDataset):
+    """
+    Class for the PINA dataset with torch.Tensor data
+    """
+
+    # Override _retrive_data method for torch.Tensor data
+    def _retrive_data(self, data, idx_list):
+        return {k: v[idx_list] for k, v in data.items()}
 
     @property
     def input(self):
@@ -112,129 +141,42 @@ def input(self):
         return {k: v["input"] for k, v in self.conditions_dict.items()}
 
 
-class PinaBatch(Batch):
+class PinaGraphDataset(PinaDataset):
     """
-    Add extract function to torch_geometric Batch object
+    Class for the PINA dataset with torch_geometric.data.Data data
     """
 
-    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_dummy
-
-        ex_data = conditions_dict[list(conditions_dict.keys())[0]]["input"][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]]["target"][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 = {}
-        for condition, data in self.conditions_dict.items():
-            cond_idx = idx[: self.max_conditions_lengths[condition]]
-            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: (
-                    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 _base_create_graph_batch_from_list(self, data):
-        batch = PinaBatch.from_data_list(data)
+    def _create_graph_batch_from_list(self, data):
+        batch = LabelBatch.from_data_list(data)
         return batch
 
-    def _base_create_output_batch(self, data):
+    def _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):
-        return {
-            k: {k_data: v[k_data][idx % len(v["input"])] for k_data in v.keys()}
-            for k, v in self.conditions_dict.items()
-        }
-
-    def get_all_data(self):
-        index = [i for i in range(len(self))]
-        return self.fetch_from_idx_list(index)
-
-    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
+        Create a Batch object from a list of Data objects.
+
+        :param data: List of Data objects
+        :type data: list
+        :return: Batch object
+        :rtype: Batch or PinaBatch
         """
         if isinstance(data[0], Data):
             return self._create_graph_batch_from_list(data)
         return self._create_output_batch(data)
+
+    # Override _retrive_data method for graph handling
+    def _retrive_data(self, data, idx_list):
+        # Return the data from the current condition
+        # If the data is a list of Data objects, create a Batch object
+        # If the data is a list of torch.Tensor objects, create a torch.Tensor
+        return {
+            k: (
+                self._create_graph_batch_from_list([v[i] for i in idx_list])
+                if isinstance(v, list)
+                else self._create_output_batch(v[idx_list])
+            )
+            for k, v in data.items()
+        }

pina/graph.py

@@ -3,7 +3,7 @@
 """
 
 import torch
-from torch_geometric.data import Data
+from torch_geometric.data import Data, Batch
 from torch_geometric.utils import to_undirected
 from . import LabelTensor
 from .utils import check_consistency, is_function
@@ -162,6 +162,18 @@ def _preprocess_edge_index(edge_index, undirected):
             edge_index = to_undirected(edge_index)
         return edge_index
 
+    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 GraphBuilder:
     """
@@ -317,3 +329,31 @@ def compute_knn_graph(points, k):
         row = torch.arange(points.size(0)).repeat_interleave(k)
         col = knn_indices.flatten()
         return torch.stack([row, col], dim=0).as_subclass(torch.Tensor)
+
+
+class LabelBatch(Batch):
+    """
+    Add extract function to torch_geometric Batch object
+    """
+
+    @classmethod
+    def from_data_list(cls, data_list):
+        """
+        Create a Batch object from a list of Data objects.
+        """
+        # Store the labels of Data/Graph objects (all data have the same labels)
+        # If the data do not contain labels, labels is an empty dictionary,
+        # therefore the labels are not stored
+        labels = {
+            k: v.labels
+            for k, v in data_list[0].items()
+            if isinstance(v, LabelTensor)
+        }
+
+        # Create a Batch object from the list of Data objects
+        batch = super().from_data_list(data_list)
+
+        # Put the labels back in the Batch object
+        for k, v in labels.items():
+            batch[k].labels = v
+        return batch