Add consistency check for graph data

Author: FilippoOlivo

pina/condition/condition_interface.py

@@ -1,15 +1,79 @@
+"""
+Module that defines the ConditionInterface class.
+"""
+
 from abc import ABCMeta
+from torch_geometric.data import Data
+from ..label_tensor import LabelTensor
+from ..graph import Graph
 
 
 class ConditionInterface(metaclass=ABCMeta):
+    """
+    Abstract class which defines a common interface for all the conditions.
+    """
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self):
         self._problem = None
 
     @property
     def problem(self):
+        """
+        Return the problem to which the condition is associated.
+
+        :return: Problem to which the condition is associated
+        :rtype: pina.problem.AbstractProblem
+        """
         return self._problem
 
     @problem.setter
     def problem(self, value):
         self._problem = value
+
+    @staticmethod
+    def _check_graph_list_consistency(data_list):
+
+        # If the data is a Graph or Data object, return (do not need to check
+        # anything)
+        if isinstance(data_list, (Graph, Data)):
+            return
+        data = data_list[0]
+        # Store the keys of the first element in the list
+        keys = sorted(list(data.keys()))
+
+        # Store the type of each tensor inside first element Data/Graph object
+        data_types = {name: tensor.__class__ for name, tensor in data.items()}
+
+        # Store the labels of each LabelTensor inside first element Data/Graph
+        # object
+        labels = {
+            name: tensor.labels
+            for name, tensor in data.items()
+            if isinstance(tensor, LabelTensor)
+        }
+        # Iterate over the list of Data/Graph objects
+        print(data_types)
+        for data in data_list[1:]:
+            # Check if the keys of the current element are the same as the first
+            # element
+            if sorted(list(data.keys())) != keys:
+                raise ValueError(
+                    "All elements in the list must have the same keys."
+                )
+
+            for name, tensor in data.items():
+                # Check if the type of each tensor inside the current element
+                # is the same as the first element
+                if tensor.__class__ is not data_types[name]:
+                    raise ValueError(
+                        f"Data {name} must be a {data_types[name]}, got "
+                        f"{tensor.__class__}"
+                    )
+
+                # If the tensor is a LabelTensor, check if the labels are the
+                # same as the first element
+                if isinstance(tensor, LabelTensor):
+                    if tensor.labels != labels[name]:
+                        raise ValueError(
+                            f"Data {name} must have the same labels"
+                        )

pina/condition/data_condition.py

@@ -88,3 +88,7 @@ class GraphDataCondition(DataCondition):
     """
     DataCondition for Graph/Data input data
     """
+
+    def __init__(self, input, conditional_variables=None):
+        self._check_graph_list_consistency(input)
+        super().__init__(input, conditional_variables)

pina/condition/input_equation_condition.py

@@ -78,3 +78,7 @@ class InputGraphEquationCondition(InputEquationCondition):
     """
     InputEquationCondition subclass for Graph input data.
     """
+
+    def __init__(self, input, equation):
+        super().__init__(input, equation)
+        self._check_graph_list_consistency(input)

pina/condition/input_target_condition.py

@@ -99,6 +99,10 @@ class TensorInputGraphTargetCondition(InputTargetCondition):
     data.
     """
 
+    def __init__(self, input, target):
+        self._check_graph_list_consistency(target)
+        super().__init__(input, target)
+
     @staticmethod
     def _check_input_target_consistency(input, target):
         if isinstance(target, (Graph, Data)):
@@ -115,6 +119,10 @@ class GraphInputTensorTargetCondition(InputTargetCondition):
     data.
     """
 
+    def __init__(self, input, target):
+        self._check_graph_list_consistency(input)
+        super().__init__(input, target)
+
     @staticmethod
     def _check_input_target_consistency(input, target):
         if isinstance(input, (Graph, Data)):
@@ -130,6 +138,11 @@ class GraphInputGraphTargetCondition(InputTargetCondition):
     InputTargetCondition subclass for Graph/Data input and target data.
     """
 
+    def __init__(self, input, target):
+        self._check_graph_list_consistency(input)
+        self._check_graph_list_consistency(target)
+        super().__init__(input, target)
+
     @staticmethod
     def _check_input_target_consistency(input, target):
         if isinstance(input, list) and isinstance(target, list):

tests/test_condition.py

@@ -48,6 +48,26 @@
     for x_, pos_ in zip(x, pos)
 ]
 
+x = LabelTensor(torch.rand(10, 20, 2), ["u", "v"])
+pos = LabelTensor(torch.rand(10, 20, 2), ["x", "y"])
+radius = 0.1
+input_graph_lt = [
+    RadiusGraph(
+        x=x[i],
+        pos=pos[i],
+        radius=radius,
+    )
+    for i in range(len(x))
+]
+target_graph_lt = [
+    RadiusGraph(
+        x=x[i],
+        pos=pos[i],
+        radius=radius,
+    )
+    for i in range(len(x))
+]
+
 input_single_graph = input_graph[0]
 target_single_graph = target_graph[0]
 
@@ -80,6 +100,16 @@ def test_init_input_target():
     with pytest.raises(ValueError):
         Condition(input=example_domain, target=example_domain)
 
+    # Test wrong graph condition initialisation
+    input = [input_graph[0], input_graph_lt[0]]
+    target = [target_graph[0], target_graph_lt[0]]
+    with pytest.raises(ValueError):
+        Condition(input=input, target=target)
+        
+    input_graph_lt[0].x.labels = ["a", "b"]
+    with pytest.raises(ValueError):
+        Condition(input=input_graph_lt, target=target_graph_lt)
+    
 
 def test_init_domain_equation():
     cond = Condition(domain=example_domain, equation=FixedValue(0.0))