Implement graph conditions and enhance Graph class

Author: FilippoOlivo

pina/collector.py

@@ -56,10 +56,6 @@ 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/condition/__init__.py

@@ -4,9 +4,15 @@
     "DomainEquationCondition",
     "InputPointsEquationCondition",
     "InputOutputPointsCondition",
+    "GraphInputOutputCondition",
+    "GraphDataCondition",
+    "GraphInputEquationCondition",
 ]
 
 from .condition_interface import ConditionInterface
 from .domain_equation_condition import DomainEquationCondition
 from .input_equation_condition import InputPointsEquationCondition
 from .input_output_condition import InputOutputPointsCondition
+from .graph_condition import GraphInputOutputCondition
+from .graph_condition import GraphDataCondition
+from .graph_condition import GraphInputEquationCondition

pina/condition/condition.py

@@ -4,6 +4,10 @@
 from .input_equation_condition import InputPointsEquationCondition
 from .input_output_condition import InputOutputPointsCondition
 from .data_condition import DataConditionInterface
+from .graph_condition import (
+    GraphInputOutputCondition,
+    GraphInputEquationCondition,
+)
 import warnings
 from ..utils import custom_warning_format
 
@@ -82,5 +86,9 @@ def __new__(cls, *args, **kwargs):
             return DataConditionInterface(**kwargs)
         elif sorted_keys == DataConditionInterface.__slots__[0]:
             return DataConditionInterface(**kwargs)
+        elif sorted_keys == sorted(GraphInputOutputCondition.__slots__):
+            return GraphInputOutputCondition(**kwargs)
+        elif sorted_keys == sorted(GraphInputEquationCondition.__slots__):
+            return GraphInputEquationCondition(**kwargs)
         else:
             raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")

pina/condition/condition_interface.py

@@ -3,10 +3,7 @@
 
 class ConditionInterface(metaclass=ABCMeta):
 
-    condition_types = ["physics", "supervised", "unsupervised"]
-
     def __init__(self, *args, **kwargs):
-        self._condition_type = None
         self._problem = None
 
     @property
@@ -16,19 +13,3 @@ def problem(self):
     @problem.setter
     def problem(self, value):
         self._problem = value
-
-    @property
-    def condition_type(self):
-        return self._condition_type
-
-    @condition_type.setter
-    def condition_type(self, values):
-        if not isinstance(values, (list, tuple)):
-            values = [values]
-        for value in values:
-            if value not in ConditionInterface.condition_types:
-                raise ValueError(
-                    "Unavailable type of condition, expected one of"
-                    f" {ConditionInterface.condition_types}."
-                )
-        self._condition_type = values

pina/condition/graph_condition.py

@@ -0,0 +1,68 @@
+from .condition_interface import ConditionInterface
+from ..graph import Graph
+from ..utils import check_consistency
+from torch_geometric.data import Data
+from ..equation.equation_interface import EquationInterface
+
+
+class GraphCondition(ConditionInterface):
+    """
+    TODO
+    """
+
+    __slots__ = ["graph"]
+
+    def __new__(cls, graph):
+        """
+        TODO : add docstring
+        """
+        check_consistency(graph, (Graph, Data))
+        graph = [graph] if isinstance(graph, Data) else graph
+
+        if all(g.y is not None for g in graph):
+            return super().__new__(GraphInputOutputCondition)
+        else:
+            return super().__new__(GraphDataCondition)
+
+    def __init__(self, graph):
+
+        super().__init__()
+        self.graph = graph
+
+    def __setattr__(self, key, value):
+        if key == "graph":
+            check_consistency(value, (Graph, Data))
+            GraphCondition.__dict__[key].__set__(self, value)
+        elif key in ("_problem", "_condition_type"):
+            super().__setattr__(key, value)
+
+
+class GraphInputEquationCondition(ConditionInterface):
+
+    __slots__ = ["graph", "equation"]
+
+    def __init__(self, graph, equation):
+        super().__init__()
+        self.graph = graph
+        self.equation = equation
+
+    def __setattr__(self, key, value):
+        if key == "graph":
+            check_consistency(value, (Graph, Data))
+            GraphInputEquationCondition.__dict__[key].__set__(self, value)
+        elif key == "equation":
+            check_consistency(value, (EquationInterface))
+            GraphInputEquationCondition.__dict__[key].__set__(self, value)
+        elif key in ("_problem", "_condition_type"):
+            super().__setattr__(key, value)
+
+
+# The split between GraphInputOutputCondition and GraphDataCondition
+# distinguishes different types of graph conditions passed to problems.
+# This separation simplifies consistency checks during problem creation.
+class GraphDataCondition(GraphCondition):
+    pass
+
+
+class GraphInputOutputCondition(GraphCondition):
+    pass

pina/data/data_module.py

@@ -84,39 +84,39 @@ def _collate_standard_dataloader(self, batch):
         conditions_names = batch[0].keys()
         # Condition names
         for condition_name in conditions_names:
-            single_cond_dict = {}
             condition_args = batch[0][condition_name].keys()
-            for arg in condition_args:
-                data_list = [
-                    batch[idx][condition_name][arg]
-                    for idx in range(
-                        min(
-                            len(batch),
-                            self.max_conditions_lengths[condition_name],
-                        )
-                    )
-                ]
-                single_cond_dict[arg] = self._collate(data_list)
-
-            batch_dict[condition_name] = single_cond_dict
+            batch_dict[condition_name] = self._collate(
+                condition_args, condition_name, batch
+            )
         return batch_dict
 
-    @staticmethod
-    def _collate_tensor_dataset(data_list):
-        if isinstance(data_list[0], LabelTensor):
-            return LabelTensor.stack(data_list)
-        if isinstance(data_list[0], torch.Tensor):
-            return torch.stack(data_list)
-        raise RuntimeError("Data must be Tensors or LabelTensor ")
-
-    def _collate_graph_dataset(self, data_list):
-        if isinstance(data_list[0], LabelTensor):
-            return LabelTensor.cat(data_list)
-        if isinstance(data_list[0], torch.Tensor):
-            return torch.cat(data_list)
-        if isinstance(data_list[0], Data):
-            return self.dataset.create_graph_batch(data_list)
-        raise RuntimeError("Data must be Tensors or LabelTensor or pyG Data")
+    def _collate_tensor_dataset(self, condition_args, condition_name, batch):
+        to_return_dict = {}
+        for arg in condition_args:
+            data_list = [
+                batch[idx][condition_name][arg]
+                for idx in range(
+                    min(len(batch), self.max_conditions_lengths[condition_name])
+                )
+            ]
+            if isinstance(data_list[0], LabelTensor):
+                data = LabelTensor.stack(data_list)
+            elif isinstance(data_list[0], torch.Tensor):
+                data = torch.stack(data_list)
+
+            to_return_dict[arg] = data
+        return to_return_dict
+
+    def _collate_graph_dataset(self, condition_args, condition_name, batch):
+        data_list = [
+            batch[idx][condition_name]
+            for idx in range(
+                min(len(batch), self.max_conditions_lengths[condition_name])
+            )
+        ]
+        return self.dataset._divide_batch(
+            batch=self.dataset.create_graph_batch(data_list)
+        )
 
     def __call__(self, batch):
         return self.callable_function(batch)
@@ -285,7 +285,7 @@ def setup(self, stage=None):
 
     @staticmethod
     def _split_condition(condition_dict, splits_dict):
-        len_condition = len(condition_dict["input_points"])
+        len_condition = len(list(condition_dict.values())[0])
 
         lengths = [
             int(len_condition * length) for length in splits_dict.values()
@@ -308,7 +308,7 @@ def _split_condition(condition_dict, splits_dict):
                 if k != "equation"
                 # Equations are NEVER dataloaded
             }
-            if offset + stage_len >= len_condition:
+            if offset + stage_len > len_condition:
                 offset = len_condition - 1
                 continue
             offset += stage_len
@@ -343,7 +343,7 @@ def _apply_shuffle(condition_dict, len_data):
             condition_name,
             condition_dict,
         ) in collector.data_collections.items():
-            len_data = len(condition_dict["input_points"])
+            len_data = len(list(condition_dict.values())[0])
             if self.shuffle:
                 _apply_shuffle(condition_dict, len_data)
             for key, data in self._split_condition(
@@ -390,12 +390,12 @@ def find_max_conditions_lengths(self, split):
         max_conditions_lengths = {}
         for k, v in self.collector_splits[split].items():
             if self.batch_size is None:
-                max_conditions_lengths[k] = len(v["input_points"])
+                max_conditions_lengths[k] = len(list(v.values())[0])
             elif self.repeat:
                 max_conditions_lengths[k] = self.batch_size
             else:
                 max_conditions_lengths[k] = min(
-                    len(v["input_points"]), self.batch_size
+                    len(list(v.values())[0]), self.batch_size
                 )
         return max_conditions_lengths