Conditions refactoring (#758)

Author: FilippoOlivo

pina/condition/__init__.py

@@ -5,13 +5,7 @@
     "ConditionInterface",
     "DomainEquationCondition",
     "InputTargetCondition",
-    "TensorInputTensorTargetCondition",
-    "TensorInputGraphTargetCondition",
-    "GraphInputTensorTargetCondition",
-    "GraphInputGraphTargetCondition",
     "InputEquationCondition",
-    "InputTensorEquationCondition",
-    "InputGraphEquationCondition",
     "DataCondition",
     "GraphDataCondition",
     "TensorDataCondition",
@@ -20,20 +14,6 @@
 from .condition_interface import ConditionInterface
 from .condition import Condition
 from .domain_equation_condition import DomainEquationCondition
-from .input_target_condition import (
-    InputTargetCondition,
-    TensorInputTensorTargetCondition,
-    TensorInputGraphTargetCondition,
-    GraphInputTensorTargetCondition,
-    GraphInputGraphTargetCondition,
-)
-from .input_equation_condition import (
-    InputEquationCondition,
-    InputTensorEquationCondition,
-    InputGraphEquationCondition,
-)
-from .data_condition import (
-    DataCondition,
-    GraphDataCondition,
-    TensorDataCondition,
-)
+from .input_target_condition import InputTargetCondition
+from .input_equation_condition import InputEquationCondition
+from .data_condition import DataCondition

pina/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/condition/condition.py

@@ -86,12 +86,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__
         )
     )
 
@@ -112,28 +112,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/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 .condition_interface import ConditionInterface
+from ..graph import LabelBatch
+from ..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
+            ),
+        )