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

tests/test_condition.py

@@ -3,32 +3,37 @@
 
 from pina import LabelTensor, Condition
 from pina.domain import CartesianDomain
+from pina.condition import (
+    GraphInputOutputCondition,
+    GraphInputEquationCondition,
+)
 from pina.equation.equation_factory import FixedValue
+from pina.graph import RadiusGraph
+from torch_geometric.data import Data
+from pina.operator import laplacian
+from pina.equation.equation import Equation
 
-example_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
-example_input_pts = LabelTensor(torch.tensor([[0, 0, 0]]), ['x', 'y', 'z'])
-example_output_pts = LabelTensor(torch.tensor([[1, 2]]), ['a', 'b'])
+example_domain = CartesianDomain({"x": [0, 1], "y": [0, 1]})
+example_input_pts = LabelTensor(torch.tensor([[0, 0, 0]]), ["x", "y", "z"])
+example_output_pts = LabelTensor(torch.tensor([[1, 2]]), ["a", "b"])
 
 
 def test_init_inputoutput():
     Condition(input_points=example_input_pts, output_points=example_output_pts)
     with pytest.raises(ValueError):
         Condition(example_input_pts, example_output_pts)
     with pytest.raises(ValueError):
-        Condition(input_points=3., output_points='example')
+        Condition(input_points=3.0, output_points="example")
     with pytest.raises(ValueError):
         Condition(input_points=example_domain, output_points=example_domain)
 
 
-test_init_inputoutput()
-
-
 def test_init_domainfunc():
     Condition(domain=example_domain, equation=FixedValue(0.0))
     with pytest.raises(ValueError):
         Condition(example_domain, FixedValue(0.0))
     with pytest.raises(ValueError):
-        Condition(domain=3., equation='example')
+        Condition(domain=3.0, equation="example")
     with pytest.raises(ValueError):
         Condition(domain=example_input_pts, equation=example_output_pts)
 
@@ -38,6 +43,78 @@ def test_init_inputfunc():
     with pytest.raises(ValueError):
         Condition(example_domain, FixedValue(0.0))
     with pytest.raises(ValueError):
-        Condition(input_points=3., equation='example')
+        Condition(input_points=3.0, equation="example")
     with pytest.raises(ValueError):
         Condition(input_points=example_domain, equation=example_output_pts)
+
+
+def test_graph_io_condition():
+    x = torch.rand(10, 10, 4)
+    pos = torch.rand(10, 10, 2)
+    y = torch.rand(10, 10, 2)
+    graph = [
+        RadiusGraph(x=x_, pos=pos_, radius=0.1, build_edge_attr=True, y=y_)
+        for x_, pos_, y_ in zip(x, pos, y)
+    ]
+    condition = Condition(graph=graph)
+    assert isinstance(condition, GraphInputOutputCondition)
+    assert isinstance(condition.graph, list)
+
+    x = x[0]
+    pos = pos[0]
+    y = y[0]
+    edge_index = graph[0].edge_index
+    graph = Data(x=x, pos=pos, edge_index=edge_index, y=y)
+    condition = Condition(graph=graph)
+    assert isinstance(condition, GraphInputOutputCondition)
+    assert isinstance(condition.graph, Data)
+
+
+def laplace_equation(input_, output_):
+    """
+    Implementation of the laplace equation.
+    """
+    force_term = torch.sin(input_.extract(["x"]) * torch.pi) * torch.sin(
+        input_.extract(["y"]) * torch.pi
+    )
+    delta_u = laplacian(output_.extract(["u"]), input_)
+    return delta_u - force_term
+
+
+def test_graph_eq_condition():
+    def laplace(input_, output_):
+        """
+        Implementation of the laplace equation.
+        """
+        force_term = torch.sin(input_.extract(["x"]) * torch.pi) * torch.sin(
+            input_.extract(["y"]) * torch.pi
+        )
+        delta_u = laplacian(output_.extract(["u"]), input_)
+        return delta_u - force_term
+
+    x = torch.rand(10, 10, 4)
+    pos = torch.rand(10, 10, 2)
+    graph = [
+        RadiusGraph(
+            x=x_,
+            pos=pos_,
+            radius=0.1,
+            build_edge_attr=True,
+        )
+        for x_, pos_, in zip(
+            x,
+            pos,
+        )
+    ]
+    laplace_equation = Equation(laplace)
+    condition = Condition(graph=graph, equation=laplace_equation)
+    assert isinstance(condition, GraphInputEquationCondition)
+    assert isinstance(condition.graph, list)
+
+    x = x[0]
+    pos = pos[0]
+    edge_index = graph[0].edge_index
+    graph = Data(x=x, pos=pos, edge_index=edge_index)
+    condition = Condition(graph=graph, equation=laplace_equation)
+    assert isinstance(condition, GraphInputEquationCondition)
+    assert isinstance(condition.graph, Data)