Adding logic

FilippoOlivo · FilippoOlivo · commit cbadb34e0fdb · 2025-03-07T10:30:25.000+01:00
diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py
@@ -1,3 +1,6 @@
+"""
+Module for conditions.
+"""
 __all__ = [
     "Condition",
     "ConditionInterface",
@@ -16,6 +19,7 @@
 ]
 
 from .condition_interface import ConditionInterface
+from .condition import Condition
 from .domain_equation_condition import DomainEquationCondition
 from .input_target_condition import (
     InputTargetCondition,
diff --git a/pina/condition/condition.py b/pina/condition/condition.py
@@ -1,10 +1,10 @@
 """Condition module."""
 
+import warnings
+from .data_condition import DataCondition
 from .domain_equation_condition import DomainEquationCondition
 from .input_equation_condition import InputEquationCondition
 from .input_target_condition import InputTargetCondition
-from .data_condition import DataCondition
-import warnings
 from ..utils import custom_warning_format
 
 # Set the custom format for warnings
@@ -13,6 +13,13 @@
 
 
 def warning_function(new, old):
+    """Handle the deprecation warning.
+
+    :param new: Object to use instead of the old one.
+    :type new: str
+    :param old: Object to deprecate.
+    :type old: str
+    """
     warnings.warn(
         f"'{old}' is deprecated and will be removed "
         f"in future versions. Please use '{new}' instead.",
@@ -48,7 +55,23 @@ class Condition:
 
     Example::
 
-    >>> TODO
+    >>> from pina import Condition
+    >>> condition = Condition(
+    ...     input=input,
+    ...     target=target
+    ... )
+    >>> condition = Condition(
+    ...     domain=location,
+    ...     equation=equation
+    ... )
+    >>> condition = Condition(
+    ...     input=input,
+    ...     equation=equation
+    ... )
+    >>> condition = Condition(
+    ...     input=data,
+    ...     conditional_variables=conditional_variables
+    ... )
 
     """
 
diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py
@@ -4,10 +4,9 @@
 
 import torch
 from torch_geometric.data import Data
-from . import ConditionInterface
+from .condition_interface import ConditionInterface
 from ..label_tensor import LabelTensor
 from ..graph import Graph
-from ..utils import check_consistency
 
 
 class DataCondition(ConditionInterface):
@@ -36,20 +35,21 @@ def __new__(cls, input, conditional_variables=None):
         :return: DataCondition subclass
         :rtype: TensorDataCondition or GraphDataCondition
         """
-        if cls == DataCondition and isinstance(
-            input, (torch.Tensor, LabelTensor)
-        ):
+        if cls != DataCondition:
+            return super().__new__(cls)
+        if isinstance(input, (torch.Tensor, LabelTensor)):
             subclass = TensorDataCondition
             return subclass.__new__(subclass, input, conditional_variables)
 
-        elif cls == DataCondition and isinstance(
-            input, (Graph, Data, list, tuple)
-        ):
+        if isinstance(input, (Graph, Data, list, tuple)):
             cls._check_graph_list_consistency(input)
             subclass = GraphDataCondition
             return subclass.__new__(subclass, input, conditional_variables)
 
-        return super().__new__(cls)
+        raise ValueError(
+            "Invalid input types. "
+            "Please provide either Data or Graph objects."
+        )
 
     def __init__(self, input, conditional_variables=None):
         """
@@ -67,17 +67,6 @@ def __init__(self, input, conditional_variables=None):
         self.input = input
         self.conditional_variables = conditional_variables
 
-    def __setattr__(self, key, value):
-        if key == "input":
-            check_consistency(value, self._avail_input_cls)
-            DataCondition.__dict__[key].__set__(self, value)
-        elif key == "conditional_variables":
-            if value is not None:
-                check_consistency(value, self._avail_conditional_variables_cls)
-            DataCondition.__dict__[key].__set__(self, value)
-        elif key in ("_problem"):
-            super().__setattr__(key, value)
-
 
 class TensorDataCondition(DataCondition):
     """
diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py
@@ -2,8 +2,6 @@
 DomainEquationCondition class definition.
 """
 
-import torch
-
 from .condition_interface import ConditionInterface
 from ..utils import check_consistency
 from ..domain import DomainInterface
diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py
@@ -2,7 +2,6 @@
 Module to define InputEquationCondition class and its subclasses.
 """
 
-import torch
 from torch_geometric.data import Data
 from .condition_interface import ConditionInterface
 from ..label_tensor import LabelTensor
@@ -33,17 +32,25 @@ def __new__(cls, input, equation):
         :return: InputEquationCondition subclass
         :rtype: InputTensorEquationCondition or InputGraphEquationCondition
         """
-        if cls == InputEquationCondition and isinstance(input, LabelTensor):
-            subclass = InputTensorEquationCondition
-            return subclass.__new__(subclass, input, equation)
-        elif cls == InputEquationCondition and isinstance(
-            input, (Graph, Data, list, tuple)
-        ):
-            cls._check_graph_list_consistency(input)
+
+        # If the class is already a subclass, return the instance
+        if cls != InputEquationCondition:
+            return super().__new__(cls)
+
+        # Instanciate the correct subclass
+        if isinstance(input, (Graph, Data, list, tuple)):
             subclass = InputGraphEquationCondition
+            cls._check_graph_list_consistency(input)
             subclass._check_label_tensor(input)
             return subclass.__new__(subclass, input, equation)
-        return super().__new__(cls)
+        if isinstance(input, LabelTensor):
+            subclass = InputTensorEquationCondition
+            return subclass.__new__(subclass, input, equation)
+
+        # If the input is not a LabelTensor or a Graph object raise an error
+        raise ValueError(
+            "The input data object must be a LabelTensor or a Graph object."
+        )
 
     def __init__(self, input, equation):
         """
diff --git a/pina/condition/input_target_condition.py b/pina/condition/input_target_condition.py
@@ -4,10 +4,9 @@
 
 import torch
 from torch_geometric.data import Data
-from .condition_interface import ConditionInterface
 from ..label_tensor import LabelTensor
 from ..graph import Graph
-from ..utils import check_consistency
+from .condition_interface import ConditionInterface
 
 
 class InputTargetCondition(ConditionInterface):
@@ -34,43 +33,41 @@ def __new__(cls, input, target):
             TensorInputGraphTargetCondition or GraphInputTensorTargetCondition
             or GraphInputGraphTargetCondition
         """
-        if (
-            cls == InputTargetCondition
-            and isinstance(input, (torch.Tensor, LabelTensor))
-            and isinstance(target, (torch.Tensor, LabelTensor))
+        if cls != InputTargetCondition:
+            return super().__new__(cls)
+
+        if isinstance(input, (torch.Tensor, LabelTensor)) and isinstance(
+            target, (torch.Tensor, LabelTensor)
         ):
             subclass = TensorInputTensorTargetCondition
             return subclass.__new__(subclass, input, target)
-
-        elif (
-            cls == InputTargetCondition
-            and isinstance(input, (torch.Tensor, LabelTensor))
-            and isinstance(target, (Graph, Data, list, tuple))
+        if isinstance(input, (torch.Tensor, LabelTensor)) and isinstance(
+            target, (Graph, Data, list, tuple)
         ):
             cls._check_graph_list_consistency(target)
             subclass = TensorInputGraphTargetCondition
             return subclass.__new__(subclass, input, target)
 
-        elif (
-            cls == InputTargetCondition
-            and isinstance(input, (Graph, Data, list, tuple))
-            and isinstance(target, (torch.Tensor, LabelTensor))
+        if isinstance(input, (Graph, Data, list, tuple)) and isinstance(
+            target, (torch.Tensor, LabelTensor)
         ):
             cls._check_graph_list_consistency(input)
             subclass = GraphInputTensorTargetCondition
             return subclass.__new__(subclass, input, target)
 
-        elif (
-            cls == InputTargetCondition
-            and isinstance(input, (Graph, Data, list, tuple))
-            and isinstance(target, (Graph, Data, list, tuple))
+        if isinstance(input, (Graph, Data, list, tuple)) and isinstance(
+            target, (Graph, Data, list, tuple)
         ):
             cls._check_graph_list_consistency(input)
             cls._check_graph_list_consistency(target)
             subclass = GraphInputGraphTargetCondition
             return subclass.__new__(subclass, input, target)
 
-        return super().__new__(cls)
+        raise ValueError(
+            "Invalid input/target types. "
+            "Please provide either Data, Graph, LabelTensor or torch.Tensor "
+            "objects."
+        )
 
     def __init__(self, input, target):
         """
@@ -82,28 +79,16 @@ def __init__(self, input, target):
         :type target: torch.Tensor or Graph or Data
         """
         super().__init__()
-        if isinstance(input, (list, tuple)) or isinstance(
-            target, (list, tuple)
-        ):
-            self._check_input_target_len(input, target)
+        self._check_input_target_len(input, target)
         self.input = input
         self.target = target
 
-    def __setattr__(self, key, value):
-        if key == "input":
-            check_consistency(value, (torch.Tensor, LabelTensor, Data, Graph))
-            InputTargetCondition.__dict__[key].__set__(self, value)
-        elif key == "target":
-            if value is not None:
-                check_consistency(
-                    value, (torch.Tensor, LabelTensor, Data, Graph)
-                )
-            InputTargetCondition.__dict__[key].__set__(self, value)
-        elif key in ("_problem"):
-            super().__setattr__(key, value)
-
     @staticmethod
     def _check_input_target_len(input, target):
+        if isinstance(input, (Graph, Data)) or isinstance(
+            target, (Graph, Data)
+        ):
+            return
         if len(input) != len(target):
             raise ValueError(
                 "The input and target lists must have the same length."
