Small fixes in conditions

Author: FilippoOlivo

pina/condition/condition.py

@@ -56,23 +56,23 @@ class Condition:
 
     Example::
 
-    >>> 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
-    ... )
+        >>> 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
+        ... )
     """
 
     __slots__ = list(
@@ -87,6 +87,7 @@ class Condition:
     def __new__(cls, *args, **kwargs):
         """
         Create a new condition object based on the keyword arguments passed.
+
         - ``input`` and ``target``: :class:`InputTargetCondition`
         - ``domain`` and ``equation``: :class:`DomainEquationCondition`
         - ``input`` and ``equation``: :class:`InputEquationCondition`
@@ -95,8 +96,8 @@ def __new__(cls, *args, **kwargs):
 
         :raises ValueError: No valid condition has been found.
         :return: A new condition instance belonging to the proper class.
-        :rtype: ConditionInputTarget | ConditionInputEquation |
-            ConditionDomainEquation | ConditionData
+        :rtype: InputTargetCondition | DomainEquationCondition |
+                InputEquationCondition | DataCondition
         """
         if len(args) != 0:
             raise ValueError(

pina/condition/condition_interface.py

@@ -41,12 +41,14 @@ def problem(self, value):
     @staticmethod
     def _check_graph_list_consistency(data_list):
         """
-        Check if the list of Data/Graph objects is consistent.
+        Check if the list of :class:`torch_geometric.data.Data`/:class:`Graph` 
+        objects is consistent.
 
-        :param data_list: list of Data/Graph objects.
-        :type data_list: list(Data) | list(Graph)
+        :param data_list: List of graph type objects.
+        :type data_list: list(torch_geometric.data.Data) | list(Graph)
 
-        :raises ValueError: Input data must be either Data or Graph objects.
+        :raises ValueError: Input data must be either torch_geometric.data.Data 
+            or Graph objects.
         :raises ValueError: All elements in the list must have the same keys.
         :raises ValueError: Type mismatch in data tensors.
         :raises ValueError: Label mismatch in LabelTensors.

pina/condition/data_condition.py

@@ -11,10 +11,9 @@
 
 class DataCondition(ConditionInterface):
     """
-    Condition for data. This condition must be used every
-    time a Unsupervised Loss is needed in the Solver. The conditionalvariable
-    can be passed as extra-input when the model learns a conditional
-    distribution
+    This condition must be used every time a Unsupervised Loss is needed in
+    the Solver. The conditionalvariable can be passed as extra-input when
+    the model learns a conditional distribution.
     """
 
     __slots__ = ["input", "conditional_variables"]
@@ -27,14 +26,16 @@ def __new__(cls, input, conditional_variables=None):
         types of input data.
 
         :param input: Input data for the condition.
-        :type input: torch.Tensor | LabelTensor | Graph | Data
+        :type input: torch.Tensor | LabelTensor | Graph | \
+            torch_geometric.data.Data
         :param conditional_variables: Conditional variables for the condition.
         :type conditional_variables: torch.Tensor | LabelTensor
         :return: Subclass of DataCondition.
         :rtype: TensorDataCondition | GraphDataCondition
 
         :raises ValueError: If input is not of type :class:`torch.Tensor`,
-            :class:`LabelTensor`, :class:`Graph`, or :class:`Data`.
+            :class:`LabelTensor`, :class:`Graph`, or
+            :class:`torch_geometric.data.Data`.
 
 
         """
@@ -51,7 +52,7 @@ def __new__(cls, input, conditional_variables=None):
 
         raise ValueError(
             "Invalid input types. "
-            "Please provide either Data or Graph objects."
+            "Please provide either torch_geometric.data.Data or Graph objects."
         )
 
     def __init__(self, input, conditional_variables=None):
@@ -60,14 +61,14 @@ def __init__(self, input, conditional_variables=None):
         variables (if any).
 
         :param input: Input data for the condition.
-        :type input: torch.Tensor or Graph or Data
+        :type input: torch.Tensor or Graph or torch_geometric.data.Data
         :param conditional_variables: Conditional variables for the condition.
         :type conditional_variables: torch.Tensor or LabelTensor
 
         .. note::
             If either `input` is composed by a list of :class:`Graph`/
-            :class:`Data` objects, all elements must have the same structure
-            (keys and data types)
+            :class:`torch_geometric.data.Data` objects, all elements must have
+            the same structure (keys and data types)
         """
         super().__init__()
         self.input = input

pina/condition/input_equation_condition.py

@@ -33,7 +33,8 @@ def __new__(cls, input, equation):
         :rtype: InputTensorEquationCondition | InputGraphEquationCondition
 
         :raises ValueError: If input is not of type :class:`torch.Tensor`,
-            :class:`LabelTensor`, :class:`Graph`, or :class:`Data`.
+            :class:`LabelTensor`, :class:`Graph`, or 
+            :class:`torch_geometric.data.Data`.
         """
 
         # If the class is already a subclass, return the instance
@@ -59,16 +60,16 @@ def __init__(self, input, equation):
         """
         Initialize the InputEquationCondition by storing the input and equation.
 
-        :param input: torch.Tensor or Graph/Data object containing the input.
+        :param input: Input data for the condition.
         :type input: torch.Tensor | Graph
         :param EquationInterface equation: Equation object containing the
             equation function.
 
         .. note::
-            If ``input`` is composed by a list of :class:`Graph`/:class:`Data`
-            objects, all elements must have the same structure (keys and data
-            types). Moreover, at least one attribute must be a
-            :class:`LabelTensor`.
+            If ``input`` is composed by a list of :class:`Graph`/
+            :class:`torch_geometric.data.Data` objects, all elements must have 
+            the same structure (keys and data types). Moreover, at least one 
+            attribute must be a :class:`LabelTensor`.
         """
 
         super().__init__()
@@ -103,7 +104,7 @@ def _check_label_tensor(input):
         Check if at least one LabelTensor is present in the Graph object.
 
         :param input: Input data.
-        :type input: torch.Tensor or Graph or Data
+        :type input: torch.Tensor | Graph | torch_geometric.data.Data
 
         :raises ValueError: If the input data object does not contain at least
             one LabelTensor.

pina/condition/input_target_condition.py

@@ -25,19 +25,20 @@ def __new__(cls, input, target):
         the types of input and target data.
 
         :param input: Input data for the condition.
-        :type input: torch.Tensor | Graph | Data | list | tuple
+        :type input: torch.Tensor | Graph | torch_geometric.data.Data | list | \
+            tuple
         :param target: Target data for the condition.
-            Graph, Data, or list/tuple.
-        :type target: torch.Tensor | Graph | Data | list | tuple
+        :type target: torch.Tensor | Graph | torch_geometric.data.Data | list \
+            | tuple
         :return: Subclass of InputTargetCondition
-        :rtype: TensorInputTensorTargetCondition |
-            TensorInputGraphTargetCondition |
-            GraphInputTensorTargetCondition |
+        :rtype: TensorInputTensorTargetCondition | \
+            TensorInputGraphTargetCondition | \
+            GraphInputTensorTargetCondition | \
             GraphInputGraphTargetCondition
 
         :raises ValueError: If input and or target are not of type
             :class:`torch.Tensor`, :class:`LabelTensor`, :class:`Graph`, or
-            :class:`Data`.
+            :class:`torch_geometric.data.Data`.
         """
         if cls != InputTargetCondition:
             return super().__new__(cls)
@@ -71,23 +72,23 @@ def __new__(cls, input, target):
 
         raise ValueError(
             "Invalid input/target types. "
-            "Please provide either Data, Graph, LabelTensor or torch.Tensor "
-            "objects."
+            "Please provide either torch_geometric.data.Data, Graph, LabelTensor "
+            "or torch.Tensor objects."
         )
 
     def __init__(self, input, target):
         """
         Initialize the InputTargetCondition, storing the input and target data.
 
-        :param input: torch.Tensor or Graph/Data object containing the input
-        :type input: torch.Tensor | Graph or Data
-        :param target: torch.Tensor or Graph/Data object containing the target
-        :type target: torch.Tensor or Graph or Data
+        :param input: Input data for the condition.
+        :type input: torch.Tensor | Graph | torch_geometric.data.Data
+        :param target: Target data for the condition.
+        :type target: torch.Tensor | Graph | torch_geometric.data.Data
 
         .. note::
             If either ``input`` or ``target`` are composed by a list of
-            :class:`Graph`/:class:`Data` objects, all elements must have the
-            same structure (keys and data types)
+            :class:`Graph`/:class:`torch_geometric.data.Data` objects, all 
+            elements must have the same structure (keys and data types)
         """
 
         super().__init__()
@@ -117,19 +118,20 @@ class TensorInputTensorTargetCondition(InputTargetCondition):
 class TensorInputGraphTargetCondition(InputTargetCondition):
     """
     InputTargetCondition subclass for :class:`torch.Tensor`/:class:`LabelTensor`
-    input and :class:`Graph`/:class:`Data` target data.
+    input and :class:`Graph`/:class:`torch_geometric.data.Data` target data.
     """
 
 
 class GraphInputTensorTargetCondition(InputTargetCondition):
     """
-    InputTargetCondition subclass for :class:`Graph`/:class:`Data` input and
-    :class:`torch.Tensor`/:class:`LabelTensor` target data.
+    InputTargetCondition subclass for :class:`Graph`/
+    :class:`torch_geometric.data.Data` input and :class:`torch.Tensor`/
+    :class:`LabelTensor` target data.
     """
 
 
 class GraphInputGraphTargetCondition(InputTargetCondition):
     """
-    InputTargetCondition subclass for :class:`Graph`/:class:`Data` input and
-    target data.
+    InputTargetCondition subclass for :class:`Graph`/
+    :class:`torch_geometric.data.Data` input and target data.
     """