Implement graph conditions

Author: FilippoOlivo

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
@@ -15,20 +12,4 @@ 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
+        self._problem = value

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/dataset.py

@@ -146,36 +146,31 @@ def __init__(
         self, conditions_dict, max_conditions_lengths, automatic_batching
     ):
         super().__init__(conditions_dict, max_conditions_lengths)
+        self.conditions_length = {
+            k: len(v["graph"]) for k, v in self.conditions_dict.items()
+        }
+        self.length = max(self.conditions_length.values())
+
         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_points"
-        ][0]
+        ex_data = conditions_dict[list(conditions_dict.keys())[0]]["graph"][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
-        )
+        if hasattr(ex_data, "y"):
+            self._divide_batch = self._extract_output(self._divide_batch)
 
     def fetch_from_idx_list(self, idx):
         to_return_dict = {}
@@ -184,34 +179,32 @@ 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: (
-                    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()
-            }
+            batch = self._create_graph_batch_from_list(
+                [data["graph"][i] for i in idx]
+            )
+            to_return_dict[condition] = self._divide_batch(batch=batch)
+        return to_return_dict
 
+    def _divide_batch(self, batch):
+        """
+        Divide the batch into input and output points
+        """
+        to_return_dict = {}
+        to_return_dict["input_points"] = batch
+        if hasattr(batch, "y"):
+            to_return_dict["output_points"] = batch.y
         return to_return_dict
 
     def _base_create_graph_batch_from_list(self, data):
         batch = PinaBatch.from_data_list(data)
         return batch
 
-    def _base_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_points"])]
-                for k_data in v.keys()
-            }
+            k: {"graph": v["graph"][idx % len(v["graph"])]}
             for k, v in self.conditions_dict.items()
         }
 
@@ -251,3 +244,23 @@ def create_graph_batch(self, data):
         if isinstance(data[0], Data):
             return self._create_graph_batch_from_list(data)
         return self._create_output_batch(data)
+
+    @staticmethod
+    def _extract_output(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            out = func(*args, **kwargs)
+            out["output_points"] = kwargs["batch"].y
+            return out
+
+        return wrapper
+
+    @staticmethod
+    def _extract_cond_vars(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            out = func(*args, **kwargs)
+            out["conditional_variables"] = kwargs["batch"].conditional_vars
+            return out
+
+        return wrapper

pina/problem/zoo/supervised_problem.py

@@ -1,9 +1,22 @@
+import torch
 from pina.problem import AbstractProblem
 from pina import Condition
 from pina import Graph
+from pina import LabelTensor
 
 
-class SupervisedProblem(AbstractProblem):
+class SupervisedProblem:
+
+    def __new__(cls, *args, **kwargs):
+
+        if sorted(list(kwargs.keys())) == sorted(["input_", "output_"]):
+            return SupervisedTensorProblem(**kwargs)
+        elif sorted(list(kwargs.keys())) == sorted(["graph_"]):
+            return SupervisedGraphProblem(**kwargs)
+        raise RuntimeError("Invalid arguments for SupervisedProblem")
+
+
+class SupervisedTensorProblem(AbstractProblem):
     """
     A problem definition for supervised learning in PINA.
 
@@ -29,9 +42,54 @@ def __init__(self, input_, output_):
         :param output_: Output data of the problem
         :type output_: torch.Tensor
         """
-        if isinstance(input_, Graph):
-            input_ = input_.data
+        if not isinstance(input_, (torch.Tensor, LabelTensor)):
+            raise ValueError(
+                "The input data must be a torch.Tensor or a " "LabelTensor"
+            )
+        if not isinstance(output_, (torch.Tensor, LabelTensor)):
+            raise ValueError(
+                "The output data must be a torch.Tensor or a " "LabelTensor"
+            )
+        if isinstance(output_, LabelTensor):
+            self.output_variables = output_.labels
+
         self.conditions["data"] = Condition(
             input_points=input_, output_points=output_
         )
         super().__init__()
+
+
+class SupervisedGraphProblem(AbstractProblem):
+    """
+    A problem definition for supervised learning in PINA.
+
+    This class allows an easy and straightforward definition of a Supervised problem,
+    based on a single condition of type `InputOutputPointsCondition`
+
+    :Example:
+        >>> import torch
+        >>> from pina.graph import RadiusGraph
+        >>> x = torch.rand((10, 100, 10))
+        >>> pos = torch.rand((10, 100, 2))
+        >>> y = torch.rand((10, 100, 2))
+        >>> input_data = RadiusGraph(x=x, pos=pos, r=.2, y=y)
+        >>> problem = SupervisedProblem(graph_=input_data)
+    """
+
+    conditions = dict()
+    output_variables = None
+
+    def __init__(self, graph_):
+        """
+        Initialize the SupervisedProblem class
+
+        :param graph_: Input data of the problem
+        :type graph_: Graph
+        """
+        if not isinstance(graph_, list) or not all(
+            isinstance(g, Graph) for g in graph_
+        ):
+            raise ValueError("The input data must be a Graph")
+
+        self.conditions["data"] = Condition(graph=graph_)
+        super().__init__()