Conditions refactoring (#758)

Author: FilippoOlivo

pina/_src/condition/batch_manager.py

@@ -0,0 +1,43 @@
+"""
+Module for managing batches of data with device transfer capabilities.
+"""
+
+
+class _BatchManager(dict):
+    """
+    A dictionary-based batch manager that supports dot-notation
+    and moving tensors to devices.
+    """
+
+    def to(self, device):
+        """
+        Move all tensors in the batch to the specified device.
+
+        :param device: The target device.
+        :type device: torch.device | str
+        :return: The updated batch manager.
+        :rtype: _BatchManager
+        """
+        for key, value in self.items():
+            if hasattr(value, "to"):
+                moved_value = value.to(device)
+                self[key] = moved_value  # Updates both dict and attribute
+        return self
+
+    def __getattribute__(self, name):
+        """
+        Alias attribute access to dictionary keys.
+
+        :param str name: The name of the attribute to retrieve.
+        :return: The value associated with the attribute name.
+        :rtype: Any
+        """
+        try:
+            return super().__getattribute__(name)
+        except AttributeError:
+            try:
+                return self[name]
+            except KeyError:
+                raise AttributeError(
+                    f"'BatchManager' object has no attribute '{name}'"
+                )

pina/_src/condition/condition.py

@@ -88,12 +88,12 @@ class Condition:
     """
 
     # Combine all possible keyword arguments from the different Condition types
-    __slots__ = list(
+    available_kwargs = list(
         set(
-            InputTargetCondition.__slots__
-            + InputEquationCondition.__slots__
-            + DomainEquationCondition.__slots__
-            + DataCondition.__slots__
+            InputTargetCondition.__fields__
+            + InputEquationCondition.__fields__
+            + DomainEquationCondition.__fields__
+            + DataCondition.__fields__
         )
     )
 
@@ -114,28 +114,28 @@ def __new__(cls, *args, **kwargs):
         if len(args) != 0:
             raise ValueError(
                 "Condition takes only the following keyword "
-                f"arguments: {Condition.__slots__}."
+                f"arguments: {Condition.available_kwargs}."
             )
 
         # Class specialization based on keyword arguments
         sorted_keys = sorted(kwargs.keys())
 
         # Input - Target Condition
-        if sorted_keys == sorted(InputTargetCondition.__slots__):
+        if sorted_keys == sorted(InputTargetCondition.__fields__):
             return InputTargetCondition(**kwargs)
 
         # Input - Equation Condition
-        if sorted_keys == sorted(InputEquationCondition.__slots__):
+        if sorted_keys == sorted(InputEquationCondition.__fields__):
             return InputEquationCondition(**kwargs)
 
         # Domain - Equation Condition
-        if sorted_keys == sorted(DomainEquationCondition.__slots__):
+        if sorted_keys == sorted(DomainEquationCondition.__fields__):
             return DomainEquationCondition(**kwargs)
 
         # Data Condition
         if (
-            sorted_keys == sorted(DataCondition.__slots__)
-            or sorted_keys[0] == DataCondition.__slots__[0]
+            sorted_keys == sorted(DataCondition.__fields__)
+            or sorted_keys[0] == DataCondition.__fields__[0]
         ):
             return DataCondition(**kwargs)
 

pina/_src/condition/condition_base.py

@@ -0,0 +1,127 @@
+"""
+Base class for conditions.
+"""
+
+from functools import partial
+import torch
+from torch_geometric.data import Batch
+from torch.utils.data import DataLoader
+from pina._src.condition.condition_interface import ConditionInterface
+from pina._src.core.graph import LabelBatch
+from pina._src.core.label_tensor import LabelTensor
+
+
+class ConditionBase(ConditionInterface):
+    """
+    Base abstract class for all conditions in PINA.
+    This class provides common functionality for handling data storage,
+    batching, and interaction with the associated problem.
+    """
+
+    collate_fn_dict = {
+        "tensor": torch.stack,
+        "label_tensor": LabelTensor.stack,
+        "graph": LabelBatch.from_data_list,
+        "data": Batch.from_data_list,
+    }
+
+    def __init__(self, **kwargs):
+        """
+        Initialization of the :class:`ConditionBase` class.
+
+        :param kwargs: Keyword arguments representing the data to be stored.
+        """
+        super().__init__()
+        self.data = self.store_data(**kwargs)
+
+    @property
+    def problem(self):
+        """
+        Return the problem associated with this condition.
+
+        :return: Problem associated with this condition.
+        :rtype: ~pina.problem.abstract_problem.AbstractProblem
+        """
+        return self._problem
+
+    @problem.setter
+    def problem(self, value):
+        """
+        Set the problem associated with this condition.
+
+        :param pina.problem.abstract_problem.AbstractProblem value: The problem
+            to associate with this condition
+        """
+        self._problem = value
+
+    def __len__(self):
+        """
+        Return the number of data points in the condition.
+
+        :return: Number of data points.
+        :rtype: int
+        """
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        """
+        Return the data point(s) at the specified index.
+
+        :param idx: Index(es) of the data point(s) to retrieve.
+        :type idx: int | list[int]
+        :return: Data point(s) at the specified index.
+        """
+        return self.data[idx]
+
+    @classmethod
+    def automatic_batching_collate_fn(cls, batch):
+        """
+        Collate function for automatic batching to be used in DataLoader.
+        :param batch: A list of items from the dataset.
+        :type batch: list
+        :return: A collated batch.
+        :rtype: dict
+        """
+        if not batch:
+            return {}
+        instance_class = batch[0].__class__
+        return instance_class.create_batch(batch)
+
+    @staticmethod
+    def collate_fn(batch, condition):
+        """
+        Collate function for custom batching to be used in DataLoader.
+
+        :param batch: A list of items from the dataset.
+        :type batch: list
+        :param condition: The condition instance.
+        :type condition: ConditionBase
+        :return: A collated batch.
+        :rtype: dict
+        """
+        data = condition.data[batch].to_batch()
+        return data
+
+    def create_dataloader(
+        self, dataset, batch_size, shuffle, automatic_batching
+    ):
+        """
+        Create a DataLoader for the condition.
+
+        :param int batch_size: The batch size for the DataLoader.
+        :param bool shuffle: Whether to shuffle the data. Default is ``False``.
+        :return: The DataLoader for the condition.
+        :rtype: torch.utils.data.DataLoader
+        """
+        if batch_size == len(dataset):
+            pass  # will be updated in the near future
+        return DataLoader(
+            dataset=dataset,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            collate_fn=(
+                partial(self.collate_fn, condition=self)
+                if not automatic_batching
+                else self.automatic_batching_collate_fn
+            ),
+        )

pina/_src/condition/condition_interface.py

@@ -1,9 +1,6 @@
 """Module for the Condition interface."""
 
-from abc import ABCMeta
-from torch_geometric.data import Data
-from pina._src.core.label_tensor import LabelTensor
-from pina._src.core.graph import Graph
+from abc import ABCMeta, abstractmethod
 
 
 class ConditionInterface(metaclass=ABCMeta):
@@ -15,112 +12,46 @@ class ConditionInterface(metaclass=ABCMeta):
     description of all available conditions and how to instantiate them.
     """
 
-    def __init__(self):
+    @abstractmethod
+    def __init__(self, **kwargs):
         """
         Initialization of the :class:`ConditionInterface` class.
         """
-        self._problem = None
 
     @property
+    @abstractmethod
     def problem(self):
         """
         Return the problem associated with this condition.
 
         :return: Problem associated with this condition.
         :rtype: ~pina.problem.abstract_problem.AbstractProblem
         """
-        return self._problem
 
     @problem.setter
+    @abstractmethod
     def problem(self, value):
         """
         Set the problem associated with this condition.
 
         :param pina.problem.abstract_problem.AbstractProblem value: The problem
             to associate with this condition
         """
-        self._problem = value
 
-    @staticmethod
-    def _check_graph_list_consistency(data_list):
+    @abstractmethod
+    def __len__(self):
         """
-        Check the consistency of the list of Data | Graph objects.
-        The following checks are performed:
+        Return the number of data points in the condition.
 
-        - All elements in the list must be of the same type (either
-          :class:`~torch_geometric.data.Data` or :class:`~pina.graph.Graph`).
-
-        - All elements in the list must have the same keys.
-
-        - The data type of each tensor must be consistent across all elements.
-
-        - If a tensor is a :class:`~pina.label_tensor.LabelTensor`, its labels
-          must also be consistent across all elements.
-
-        :param data_list: The list of Data | Graph objects to check.
-        :type data_list: list[Data] | list[Graph] | tuple[Data] | tuple[Graph]
-        :raises ValueError: If the input types are invalid.
-        :raises ValueError: If all elements in the list do not have the same
-            keys.
-        :raises ValueError: If the type of each tensor is not consistent across
-            all elements in the list.
-        :raises ValueError: If the labels of the LabelTensors are not consistent
-            across all elements in the list.
+        :return: Number of data points.
+        :rtype: int
         """
-        # If the data is a Graph or Data object, perform no checks
-        if isinstance(data_list, (Graph, Data)):
-            return
-
-        # Check all elements in the list are of the same type
-        if not all(isinstance(i, (Graph, Data)) for i in data_list):
-            raise ValueError(
-                "Invalid input. Please, provide either Data or Graph objects."
-            )
-
-        # Store the keys, data types and labels of the first element
-        data = data_list[0]
-        keys = sorted(list(data.keys()))
-        data_types = {name: tensor.__class__ for name, tensor in data.items()}
-        labels = {
-            name: tensor.labels
-            for name, tensor in data.items()
-            if isinstance(tensor, LabelTensor)
-        }
-
-        # Iterate over the list of Data | Graph objects
-        for data in data_list[1:]:
-
-            # Check that all elements in the list have the same keys
-            if sorted(list(data.keys())) != keys:
-                raise ValueError(
-                    "All elements in the list must have the same keys."
-                )
-
-            # Iterate over the tensors in the current element
-            for name, tensor in data.items():
-                # Check that the type of each tensor is consistent
-                if tensor.__class__ is not data_types[name]:
-                    raise ValueError(
-                        f"Data {name} must be a {data_types[name]}, got "
-                        f"{tensor.__class__}"
-                    )
-
-                # Check that the labels of each LabelTensor are consistent
-                if isinstance(tensor, LabelTensor):
-                    if tensor.labels != labels[name]:
-                        raise ValueError(
-                            "LabelTensor must have the same labels"
-                        )
 
-    def __getattribute__(self, name):
+    @abstractmethod
+    def __getitem__(self, idx):
         """
-        Get an attribute from the object.
+        Return the data point(s) at the specified index.
 
-        :param str name: The name of the attribute to get.
-        :return: The requested attribute.
-        :rtype: Any
+        :param int idx: Index of the data point(s) to retrieve.
+        :return: Data point(s) at the specified index.
         """
-        to_return = super().__getattribute__(name)
-        if isinstance(to_return, (Graph, Data)):
-            to_return = [to_return]
-        return to_return