Bug fix and reimplement tests

Author: FilippoOlivo

pina/graph.py

@@ -2,14 +2,11 @@
 This module provides an interface to build torch_geometric.data.Data objects.
 """
 
-import warnings
-
 import torch
-
-from . import LabelTensor
-from .utils import check_consistency, is_function
 from torch_geometric.data import Data
 from torch_geometric.utils import to_undirected
+from . import LabelTensor
+from .utils import check_consistency, is_function
 
 
 class Graph(Data):
@@ -23,18 +20,18 @@ def __new__(
     ):
         """
         :param kwargs: Parameters to construct the Graph object.
-        :return: The Data object.
-        :rtype: torch_geometric.data.Data
+        :return: A new instance of the Graph class.
+        :rtype: Graph
         """
         # create class instance
         instance = Data.__new__(cls)
 
         # check the consistency of types defined in __init__, the others are not
         # checked (as in pyg Data object)
         instance._check_type_consistency(**kwargs)
-        
+
         return instance
-    
+
     def __init__(
         self,
         x=None,
@@ -46,18 +43,30 @@ def __init__(
     ):
         """
         Initialize the Graph object.
-        :param torch.Tensor pos: The position tensor.
-        :param torch.Tensor edge_index: The edge index tensor.
-        :param torch.Tensor edge_attr: The edge attribute tensor.
-        :param bool build_edge_attr: Whether to build the edge attributes.
-        :param kwargs: Additional parameters.
+
+        :param x: Optional tensor of node features (N, F) where F is the number
+            of features per node.
+        :type x: torch.Tensor, LabelTensor
+        :param torch.Tensor edge_index: A tensor of shape (2, E) representing
+            the indices of the graph's edges.
+        :param pos: A tensor of shape (N, D) representing the positions of N
+            points in D-dimensional space.
+        :type pos: torch.Tensor | LabelTensor
+        :param edge_attr: Optional tensor of edge_featured (E, F') where F' is
+            the number of edge features
+        :param bool undirected: Whether to make the graph undirected
+        :param kwargs: Additional keyword arguments passed to the
+            `torch_geometric.data.Data` class constructor. If the argument
+            is a `torch.Tensor` or `LabelTensor`, it is included in the Data
+            object as a graph parameter.
         """
         # preprocessing
         self._preprocess_edge_index(edge_index, undirected)
 
         # calling init
-        super().__init__(x=x, edge_index=edge_index, edge_attr=edge_attr,
-                         pos=pos, **kwargs)
+        super().__init__(
+            x=x, edge_index=edge_index, edge_attr=edge_attr, pos=pos, **kwargs
+        )
 
     def _check_type_consistency(self, **kwargs):
         # default types, specified in cls.__new__, by default they are Nont
@@ -85,9 +94,10 @@ def _check_pos_consistency(pos):
         Check if the position tensor is consistent.
         :param torch.Tensor pos: The position tensor.
         """
-        check_consistency(pos, (torch.Tensor, LabelTensor))
-        if pos.ndim != 2:
-            raise ValueError("pos must be a 2D tensor.")
+        if pos is not None:
+            check_consistency(pos, (torch.Tensor, LabelTensor))
+            if pos.ndim != 2:
+                raise ValueError("pos must be a 2D tensor.")
 
     @staticmethod
     def _check_edge_index_consistency(edge_index):
@@ -109,16 +119,17 @@ def _check_edge_attr_consistency(edge_attr, edge_index):
 
         :param torch.Tensor edge_index: The edge index tensor.
         """
-        check_consistency(edge_attr, (torch.Tensor, LabelTensor))
-        if edge_attr.ndim != 2:
-            raise ValueError("edge_attr must be a 2D tensor.")
-        if edge_attr.size(1) != edge_index.size(0):
-            raise ValueError(
-                "edge_attr must have shape "
-                "[num_edges, num_edge_features], expected "
-                f"num_edges {edge_index.size(0)} "
-                f"got {edge_attr.size(1)}."
-            )
+        if edge_attr is not None:
+            check_consistency(edge_attr, (torch.Tensor, LabelTensor))
+            if edge_attr.ndim != 2:
+                raise ValueError("edge_attr must be a 2D tensor.")
+            if edge_attr.size(0) != edge_index.size(1):
+                raise ValueError(
+                    "edge_attr must have shape "
+                    "[num_edges, num_edge_features], expected "
+                    f"num_edges {edge_index.size(1)} "
+                    f"got {edge_attr.size(0)}."
+                )
 
     @staticmethod
     def _check_x_consistency(x, pos=None):
@@ -134,6 +145,9 @@ def _check_x_consistency(x, pos=None):
             if pos is not None:
                 if x.size(0) != pos.size(0):
                     raise ValueError("Inconsistent number of nodes.")
+            if pos is not None:
+                if x.size(0) != pos.size(0):
+                    raise ValueError("Inconsistent number of nodes.")
 
     @staticmethod
     def _preprocess_edge_index(edge_index, undirected):
@@ -148,73 +162,158 @@ def _preprocess_edge_index(edge_index, undirected):
             edge_index = to_undirected(edge_index)
         return edge_index
 
-class RadiusGraph(Graph):
-    def __init__(
-        self,
-        radius,
+
+class GraphBuilder:
+    """
+    A class that allows the simple definition of Graph instances.
+    """
+
+    def __new__(
+        cls,
+        pos,
+        edge_index,
         x=None,
-        pos=None,
-        edge_attr=None,
-        undirected=False,
+        edge_attr=False,
+        custom_edge_func=None,
         **kwargs,
     ):
-        super().__init__(x=x, edge_index=None, edge_attr=edge_attr,
-                         pos=pos, undirected=undirected, **kwargs)
-        edge_index = self._radius_graph(pos, radius)
-        self.radius = radius
-        self.edge_index = edge_index
-    
+        """
+        Creates a new instance of the Graph class.
+
+        :param pos: A tensor of shape (N, D) representing the positions of N
+            points in D-dimensional space.
+        :type pos: torch.Tensor | LabelTensor
+        :param edge_index: A tensor of shape (2, E) representing the indices of
+            the graph's edges.
+        :type edge_index: torch.Tensor
+        :param x: Optional tensor of node features (N, F) where F is the number
+            of features per node.
+        :type x: torch.Tensor, LabelTensor
+        :param bool edge_attr: Optional edge attributes (E, F) where F is the
+            number of features per edge.
+        :param callable custom_edge_func: A custom function to compute edge
+            attributes.
+        :param kwargs: Additional keyword arguments passed to the Graph class
+            constructor.
+        :return: A Graph instance constructed using the provided information.
+        :rtype: Graph
+        """
+        edge_attr = cls._create_edge_attr(
+            pos, edge_index, edge_attr, custom_edge_func or cls._build_edge_attr
+        )
+        return Graph(
+            x=x,
+            edge_index=edge_index,
+            edge_attr=edge_attr,
+            pos=pos,
+            **kwargs,
+        )
+
     @staticmethod
-    def _radius_graph(points, r):
-        """
-        Implementation of the radius graph construction.
-        :param points: The input points.
-        :type points: torch.Tensor
-        :param r: The radius.
-        :type r: float
-        :return: The edge index.
-        :rtype: torch.Tensor
+    def _create_edge_attr(pos, edge_index, edge_attr, func):
+        check_consistency(edge_attr, bool)
+        if edge_attr:
+            if is_function(func):
+                return func(pos, edge_index)
+            raise ValueError("custom_edge_func must be a function.")
+        return None
+
+    @staticmethod
+    def _build_edge_attr(pos, edge_index):
+        return (
+            (pos[edge_index[0]] - pos[edge_index[1]])
+            .abs()
+            .as_subclass(torch.Tensor)
+        )
+
+
+class RadiusGraph(GraphBuilder):
+    """
+    A class to build a radius graph.
+    """
+
+    def __new__(cls, pos, radius, **kwargs):
+        """
+        Creates a new instance of the Graph class using a radius-based graph
+        construction.
+
+        :param pos: A tensor of shape (N, D) representing the positions of N
+            points in D-dimensional space.
+        :type pos: torch.Tensor | LabelTensor
+        :param float radius: The radius within which points are connected.
+        :Keyword Arguments:
+            The additional keyword arguments to be passed to GraphBuilder
+            and Graph classes
+        :return: Graph instance containg the information passed in input and
+            the computed edge_index
+        :rtype: Graph
+        """
+        edge_index = cls.compute_radius_graph(pos, radius)
+        return super().__new__(cls, pos=pos, edge_index=edge_index, **kwargs)
+
+    @staticmethod
+    def compute_radius_graph(points, radius):
+        """
+        Computes a radius-based graph for a given set of points.
+
+        :param points: A tensor of shape (N, D) representing the positions of
+            N points in D-dimensional space.
+        :type points: torch.Tensor | LabelTensor
+        :param float radius: The number of nearest neighbors to find for each
+            point.
+        :rtype torch.Tensor: A tensor of shape (2, E), where E is the number of
+            edges, representing the edge indices of the KNN graph.
         """
         dist = torch.cdist(points, points, p=2)
-        edge_index = torch.nonzero(dist <= r, as_tuple=False).t()
-        if isinstance(edge_index, LabelTensor):
-            edge_index = edge_index.tensor
-        return edge_index
-    
-class KNNGraph(Graph):
-    def __init__(
-        self,
-        neighboors,
-        x=None,
-        pos=None,
-        edge_attr=None,
-        undirected=False,
-        **kwargs,
-    ):
-        super().__init__(x=x, edge_index=None, edge_attr=edge_attr,
-                         pos=pos, undirected=undirected, **kwargs)
-        edge_index = self._knn_graph(pos, neighboors)
-        self.neighboors = neighboors
-        self.edge_index = edge_index
-    
+        return (
+            torch.nonzero(dist <= radius, as_tuple=False)
+            .t()
+            .as_subclass(torch.Tensor)
+        )
+
+
+class KNNGraph(GraphBuilder):
+    """
+    A class to build a KNN graph.
+    """
+
+    def __new__(cls, pos, neighbours, **kwargs):
+        """
+        Creates a new instance of the Graph class using k-nearest neighbors
+        to compute edge_index.
+
+        :param pos: A tensor of shape (N, D) representing the positions of N
+            points in D-dimensional space.
+        :type pos: torch.Tensor | LabelTensor
+        :param int neighbours: The number of nearest neighbors to consider when
+            building the graph.
+        :Keyword Arguments:
+            The additional keyword arguments to be passed to GraphBuilder
+            and Graph classes
+
+        :return: Graph instance containg the information passed in input and
+            the computed edge_index
+        :rtype: Graph
+        """
+
+        edge_index = cls.compute_knn_graph(pos, neighbours)
+        return super().__new__(cls, pos=pos, edge_index=edge_index, **kwargs)
+
     @staticmethod
-    def _knn_graph(points, k):
-        """
-        Implementation of the k-nearest neighbors graph construction.
-        :param points: The input points.
-        :type points: torch.Tensor
-        :param k: The number of nearest neighbors.
-        :type k: int
-        :return: The edge index.
-        :rtype: torch.Tensor
+    def compute_knn_graph(points, k):
+        """
+        Computes the edge_index based k-nearest neighbors graph algorithm
+
+        :param points: A tensor of shape (N, D) representing the positions of
+            N points in D-dimensional space.
+        :type points: torch.Tensor | LabelTensor
+        :param int k: The number of nearest neighbors to find for each point.
+        :rtype torch.Tensor: A tensor of shape (2, E), where E is the number of
+            edges, representing the edge indices of the KNN graph.
         """
-        if isinstance(points, LabelTensor):
-            points = points.tensor
+
         dist = torch.cdist(points, points, p=2)
         knn_indices = torch.topk(dist, k=k + 1, largest=False).indices[:, 1:]
         row = torch.arange(points.size(0)).repeat_interleave(k)
         col = knn_indices.flatten()
-        edge_index = torch.stack([row, col], dim=0)
-        if isinstance(edge_index, LabelTensor):
-            edge_index = edge_index.tensor
-        return edge_index
+        return torch.stack([row, col], dim=0).as_subclass(torch.Tensor)

tests/test_collector.py

@@ -112,14 +112,12 @@ def poisson_sol(self, pts):
 def test_supervised_graph_collector():
     pos = torch.rand((100, 3))
     x = [torch.rand((100, 3)) for _ in range(10)]
-    graph = RadiusGraph(pos=pos, build_edge_attr=True, r=0.4)
-    graph_list_1 = [graph(x=x_) for x_ in x]
+    graph_list_1 = [RadiusGraph(pos=pos, radius=0.4, x=x_) for x_ in x]
     out_1 = torch.rand((10, 100, 3))
 
     pos = torch.rand((50, 3))
     x = [torch.rand((50, 3)) for _ in range(10)]
-    graph = RadiusGraph(pos=pos, build_edge_attr=True, r=0.4)
-    graph_list_2 = [graph(x=x_) for x_ in x]
+    graph_list_2 = [RadiusGraph(pos=pos, radius=0.4, x=x_) for x_ in x]
     out_2 = torch.rand((10, 50, 3))
 
     class SupervisedProblem(AbstractProblem):
@@ -135,6 +133,3 @@ class SupervisedProblem(AbstractProblem):
     # assert all(collector._is_conditions_ready.values())
     for v in collector.conditions_name.values():
         assert v in problem.conditions.keys()
-
-
-test_supervised_graph_collector()

tests/test_data/test_data_module.py

@@ -16,8 +16,7 @@
 x = torch.rand((100, 50, 10))
 pos = torch.rand((100, 50, 2))
 input_graph = [
-    RadiusGraph(x=x_, pos=pos_, r=0.2, build_edge_attr=True)
-    for x_, pos_, in zip(x, pos)
+    RadiusGraph(x=x_, pos=pos_, radius=0.2) for x_, pos_, in zip(x, pos)
 ]
 output_graph = torch.rand((100, 50, 10))
 
@@ -183,8 +182,7 @@ def test_dataloader(input_, output_, automatic_batching):
 x = LabelTensor(torch.rand((100, 50, 3)), ["u", "v", "w"])
 pos = LabelTensor(torch.rand((100, 50, 2)), ["x", "y"])
 input_graph = [
-    RadiusGraph(x=x[i], pos=pos[i], r=0.1, build_edge_attr=True)
-    for i in range(len(x))
+    RadiusGraph(x=x[i], pos=pos[i], radius=0.1) for i in range(len(x))
 ]
 output_graph = LabelTensor(torch.rand((100, 50, 3)), ["u", "v", "w"])