pyg-nightly 2.7.0.dev20250903__tar.gz → 2.7.0.dev20250904__tar.gz

{pyg_nightly-2.7.0.dev20250903 → pyg_nightly-2.7.0.dev20250904}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250903
+Version: 2.7.0.dev20250904
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.7.0.dev20250903 → pyg_nightly-2.7.0.dev20250904}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend="flit_core.buildapi"
 
 [project]
 name="pyg-nightly"
-version="2.7.0.dev20250903"
+version="2.7.0.dev20250904"
 authors=[
     {name="Matthias Fey", email="matthias@pyg.org"},
 ]

{pyg_nightly-2.7.0.dev20250903 → pyg_nightly-2.7.0.dev20250904}/torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250903'
+__version__ = '2.7.0.dev20250904'
 
 __all__ = [
     'Index',

{pyg_nightly-2.7.0.dev20250903 → pyg_nightly-2.7.0.dev20250904}/torch_geometric/data/data.py

@@ -1,5 +1,6 @@
 import copy
 import warnings
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from itertools import chain
@@ -904,6 +905,60 @@ class Data(BaseData, FeatureStore, GraphStore):
 
         return data
 
+    def connected_components(self) -> List[Self]:
+        r"""Extracts connected components of the graph using a union-find
+        algorithm. The components are returned as a list of
+        :class:`~torch_geometric.data.Data` objects, where each object
+        represents a connected component of the graph.
+
+        .. code-block::
+
+            data = Data()
+            data.x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
+            data.y = torch.tensor([[1.1], [2.1], [3.1], [4.1]])
+            data.edge_index = torch.tensor(
+                [[0, 1, 2, 3], [1, 0, 3, 2]], dtype=torch.long
+            )
+
+            components = data.connected_components()
+            print(len(components))
+            >>> 2
+
+            print(components[0].x)
+            >>> Data(x=[2, 1], y=[2, 1], edge_index=[2, 2])
+
+        Returns:
+            List[Data]: A list of disconnected components.
+        """
+        # Union-Find algorithm to find connected components
+        self._parents: Dict[int, int] = {}
+        self._ranks: Dict[int, int] = {}
+        for edge in self.edge_index.t().tolist():
+            self._union(edge[0], edge[1])
+
+        # Rerun _find_parent to ensure all nodes are covered correctly
+        for node in range(self.num_nodes):
+            self._find_parent(node)
+
+        # Group parents
+        grouped_parents = defaultdict(list)
+        for node, parent in self._parents.items():
+            grouped_parents[parent].append(node)
+        del self._ranks
+        del self._parents
+
+        # Create components based on the found parents (roots)
+        components: List[Self] = []
+        for nodes in grouped_parents.values():
+            # Convert the list of node IDs to a tensor
+            subset = torch.tensor(nodes, dtype=torch.long)
+
+            # Use the existing subgraph function
+            component_data = self.subgraph(subset)
+            components.append(component_data)
+
+        return components
+
     ###########################################################################
 
     @classmethod
@@ -1150,6 +1205,49 @@ class Data(BaseData, FeatureStore, GraphStore):
 
         return list(edge_attrs.values())
 
+    # Connected Components Helper Functions ###################################
+
+    def _find_parent(self, node: int) -> int:
+        r"""Finds and returns the representative parent of the given node in a
+        disjoint-set (union-find) data structure. Implements path compression
+        to optimize future queries.
+
+        Args:
+            node (int): The node for which to find the representative parent.
+
+        Returns:
+            int: The representative parent of the node.
+        """
+        if node not in self._parents:
+            self._parents[node] = node
+            self._ranks[node] = 0
+        if self._parents[node] != node:
+            self._parents[node] = self._find_parent(self._parents[node])
+        return self._parents[node]
+
+    def _union(self, node1: int, node2: int):
+        r"""Merges the sets containing node1 and node2 in the disjoint-set
+        data structure.
+
+        Finds the root parents of node1 and node2 using the _find_parent
+        method. If they belong to different sets, updates the parent of
+        root2 to be root1, effectively merging the two sets.
+
+        Args:
+            node1 (int): The index of the first node to union.
+            node2 (int): The index of the second node to union.
+        """
+        root1 = self._find_parent(node1)
+        root2 = self._find_parent(node2)
+        if root1 != root2:
+            if self._ranks[root1] < self._ranks[root2]:
+                self._parents[root1] = root2
+            elif self._ranks[root1] > self._ranks[root2]:
+                self._parents[root2] = root1
+            else:
+                self._parents[root2] = root1
+                self._ranks[root1] += 1
+
 
 ###############################################################################
 

{pyg_nightly-2.7.0.dev20250903 → pyg_nightly-2.7.0.dev20250904}/torch_geometric/data/hetero_data.py

@@ -487,6 +487,77 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
         return status
 
+    def connected_components(self) -> List[Self]:
+        r"""Extracts connected components of the heterogeneous graph using
+        a union-find algorithm. The components are returned as a list of
+        :class:`~torch_geometric.data.HeteroData` objects.
+
+        .. code-block::
+
+            data = HeteroData()
+            data["red"].x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
+            data["blue"].x = torch.tensor([[5.0], [6.0]])
+            data["red", "to", "red"].edge_index = torch.tensor(
+                [[0, 1, 2, 3], [1, 0, 3, 2]], dtype=torch.long
+            )
+
+            components = data.connected_components()
+            print(len(components))
+            >>> 4
+
+            print(components[0])
+            >>> HeteroData(
+                red={x: tensor([[1.], [2.]])},
+                blue={x: tensor([[]])},
+                red, to, red={edge_index: tensor([[0, 1], [1, 0]])}
+            )
+
+        Returns:
+            List[HeteroData]: A list of connected components.
+        """
+        # Initialize union-find structures
+        self._parents: Dict[Tuple[str, int], Tuple[str, int]] = {}
+        self._ranks: Dict[Tuple[str, int], int] = {}
+
+        # Union-Find algorithm to find connected components
+        for edge_type in self.edge_types:
+            src, _, dst = edge_type
+            edge_index = self[edge_type].edge_index
+            for src_node, dst_node in edge_index.t().tolist():
+                self._union((src, src_node), (dst, dst_node))
+
+        # Rerun _find_parent to ensure all nodes are covered correctly
+        for node_type in self.node_types:
+            for node_index in range(self[node_type].num_nodes):
+                self._find_parent((node_type, node_index))
+
+        # Group nodes by their representative parent
+        components_map = defaultdict(list)
+        for node, parent in self._parents.items():
+            components_map[parent].append(node)
+        del self._parents
+        del self._ranks
+
+        components: List[Self] = []
+        for nodes in components_map.values():
+            # Prefill subset_dict with all node types to ensure all are present
+            subset_dict = {node_type: [] for node_type in self.node_types}
+
+            # Convert the list of (node_type, node_id) tuples to a subset_dict
+            for node_type, node_id in nodes:
+                subset_dict[node_type].append(node_id)
+
+            # Convert lists to tensors
+            for node_type, node_ids in subset_dict.items():
+                subset_dict[node_type] = torch.tensor(node_ids,
+                                                      dtype=torch.long)
+
+            # Use the existing subgraph function to do all the heavy lifting
+            component_data = self.subgraph(subset_dict)
+            components.append(component_data)
+
+        return components
+
     def debug(self):
         pass  # TODO
 
@@ -1148,6 +1219,51 @@ class HeteroData(BaseData, FeatureStore, GraphStore):
 
         return list(edge_attrs.values())
 
+    # Connected Components Helper Functions ###################################
+
+    def _find_parent(self, node: Tuple[str, int]) -> Tuple[str, int]:
+        r"""Finds and returns the representative parent of the given node in a
+        disjoint-set (union-find) data structure. Implements path compression
+        to optimize future queries.
+
+        Args:
+            node (tuple[str, int]): The node for which to find the parent.
+            First element is the node type, second is the node index.
+
+        Returns:
+            tuple[str, int]: The representative parent of the node.
+        """
+        if node not in self._parents:
+            self._parents[node] = node
+            self._ranks[node] = 0
+        if self._parents[node] != node:
+            self._parents[node] = self._find_parent(self._parents[node])
+        return self._parents[node]
+
+    def _union(self, node1: Tuple[str, int], node2: Tuple[str, int]):
+        r"""Merges the node1 and node2 in the disjoint-set data structure.
+
+        Finds the root parents of node1 and node2 using the _find_parent
+        method. If they belong to different sets, updates the parent of
+        root2 to be root1, effectively merging the two sets.
+
+        Args:
+            node1 (Tuple[str, int]): The first node to union. First element is
+                the node type, second is the node index.
+            node2 (Tuple[str, int]): The second node to union. First element is
+                the node type, second is the node index.
+        """
+        root1 = self._find_parent(node1)
+        root2 = self._find_parent(node2)
+        if root1 != root2:
+            if self._ranks[root1] < self._ranks[root2]:
+                self._parents[root1] = root2
+            elif self._ranks[root1] > self._ranks[root2]:
+                self._parents[root2] = root1
+            else:
+                self._parents[root2] = root1
+                self._ranks[root1] += 1
+
 
 # Helper functions ############################################################