Add check consistency in InputTargetCondition

Author: FilippoOlivo

pina/condition/data_condition.py

@@ -35,8 +35,8 @@ def __new__(cls, input, conditional_variables=None):
         :rtype: TensorDataCondition or GraphDataCondition
         """
 
-        subclass = cls._get_subclass(input, conditional_variables)
-        if subclass is not cls:
+        if cls == DataCondition:
+            subclass = cls._get_subclass(input, conditional_variables)
             return subclass.__new__(subclass, input, conditional_variables)
         return super().__new__(cls)
 

pina/condition/input_equation_condition.py

@@ -1,5 +1,5 @@
 """
-InputEquationCondition class definition.
+Module to define InputEquationCondition class and its subclasses.
 """
 
 import torch
@@ -31,8 +31,10 @@ def __new__(cls, input, equation):
         :return: InputEquationCondition subclass
         :rtype: InputTensorEquationCondition or InputGraphEquationCondition
         """
-        subclass = cls._get_subclass(input, equation)
-        if subclass is not cls:
+        check_consistency(equation, (EquationInterface))
+
+        if cls == InputEquationCondition:
+            subclass = cls._get_subclass(input)
             return subclass.__new__(subclass, input, equation)
         return super().__new__(cls)
 
@@ -50,8 +52,7 @@ def __init__(self, input, equation):
         self.equation = equation
 
     @staticmethod
-    def _get_subclass(input, equation):
-        check_consistency(equation, (EquationInterface))
+    def _get_subclass(input):
         is_tensor_input = isinstance(input, (LabelTensor, torch.Tensor))
         is_graph_input = isinstance(input, (Data, Graph)) or (
             isinstance(input, list)

pina/condition/input_target_condition.py

@@ -1,5 +1,5 @@
 """
-TODO docstring
+This module contains condition classes for supervised learning tasks.
 """
 
 import torch
@@ -31,8 +31,9 @@ def __new__(cls, input, target):
             TensorInputGraphTargetCondition or GraphInputTensorTargetCondition
             or GraphInputGraphTargetCondition
         """
-        subclass = cls._get_subclass(input, target)
-        if subclass is not cls:
+
+        if cls == InputTargetCondition:
+            subclass = cls._get_subclass(input, target)
             return subclass.__new__(subclass, input, target)
         return super().__new__(cls)
 
@@ -46,6 +47,8 @@ def __init__(self, input, target):
         :type target: torch.Tensor or Graph or Data
         """
         super().__init__()
+        if hasattr(self, "_check_input_target_consistency"):
+            self._check_input_target_consistency(input, target)
         self.input = input
         self.target = target
 
@@ -82,22 +85,65 @@ class TensorInputTensorTargetCondition(InputTargetCondition):
     InputTargetCondition subclass for torch.Tensor input and target data.
     """
 
+    @staticmethod
+    def _check_input_target_consistency(input, target):
+        if len(input) != len(target):
+            raise ValueError(
+                "The input and target lists must have the same length."
+            )
+
 
 class TensorInputGraphTargetCondition(InputTargetCondition):
     """
     InputTargetCondition subclass for torch.Tensor input and Graph/Data target
     data.
     """
 
+    @staticmethod
+    def _check_input_target_consistency(input, target):
+        if isinstance(target, (Graph, Data)):
+            return
+        if len(input) != len(target):
+            raise ValueError(
+                "The input and target lists must have the same length."
+            )
+
 
 class GraphInputTensorTargetCondition(InputTargetCondition):
     """
     InputTargetCondition subclass for Graph/Data input and torch.Tensor target
     data.
     """
 
+    @staticmethod
+    def _check_input_target_consistency(input, target):
+        if isinstance(input, (Graph, Data)):
+            return
+        if len(input) != len(target):
+            raise ValueError(
+                "The input and target lists must have the same length."
+            )
+
 
 class GraphInputGraphTargetCondition(InputTargetCondition):
     """
     InputTargetCondition subclass for Graph/Data input and target data.
     """
+
+    @staticmethod
+    def _check_input_target_consistency(input, target):
+        if isinstance(input, list) and isinstance(target, list):
+            if len(input) != len(target):
+                raise ValueError(
+                    "The input and target lists must have the same length."
+                )
+            return
+        if isinstance(target, (Graph, Data)) and isinstance(
+            input, (Graph, Data)
+        ):
+            return
+        raise ValueError(
+            "Invalid input and target types. "
+            "input and target must be either both lists or both Graph/Data "
+            "objects."
+        )

tests/test_condition.py

@@ -14,7 +14,6 @@
     DomainEquationCondition,
 )
 from pina.condition import (
-    DataCondition,
     TensorDataCondition,
     GraphDataCondition,
 )
@@ -24,10 +23,10 @@
 
 example_domain = CartesianDomain({"x": [0, 1], "y": [0, 1]})
 
-input_tensor = torch.tensor([[0, 0, 0]])
-target_tensor = torch.tensor([[1, 2]])
-input_lt = LabelTensor(torch.tensor([[0, 0, 0]]), ["x", "y", "z"])
-target_lt = LabelTensor(torch.tensor([[1, 2]]), ["a", "b"])
+input_tensor = torch.rand((10,3))
+target_tensor = torch.rand((10,2))
+input_lt = LabelTensor(torch.rand((10,3)), ["x", "y", "z"])
+target_lt = LabelTensor(torch.rand((10,2)), ["a", "b"])
 
 x = torch.rand(10, 20, 2)
 pos = torch.rand(10, 20, 2)
@@ -69,7 +68,9 @@ def test_init_input_target():
     assert isinstance(cond, TensorInputGraphTargetCondition)
     cond = Condition(input=input_single_graph, target=target_lt)
     assert isinstance(cond, GraphInputTensorTargetCondition)
-    cond = Condition(input=input_graph, target=target_single_graph)
+    cond = Condition(input=input_graph, target=target_graph)
+    assert isinstance(cond, GraphInputGraphTargetCondition)
+    cond = Condition(input=input_single_graph, target=target_single_graph)
     assert isinstance(cond, GraphInputGraphTargetCondition)
 
     with pytest.raises(ValueError):

tests/test_problem_zoo/test_supervised_problem.py

@@ -23,7 +23,7 @@ def test_constructor_graph():
         RadiusGraph(x=x_, pos=pos_, radius=0.2, edge_attr=True)
         for x_, pos_ in zip(x, pos)
     ]
-    output_ = torch.rand((100, 10))
+    output_ = torch.rand((20, 100, 10))
     problem = SupervisedProblem(input_=input_, output_=output_)
     assert isinstance(problem, AbstractProblem)
     assert hasattr(problem, "conditions")