pyg-nightly 2.7.0.dev20250428__py3-none-any.whl → 2.7.0.dev20250429__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250428.dist-info → pyg_nightly-2.7.0.dev20250429.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250428
+Version: 2.7.0.dev20250429
 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.dev20250428.dist-info → pyg_nightly-2.7.0.dev20250429.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=Gbg4XAZMSi28NjwxD1wMjvY-M4y1bxpYHz0c6bjkLbs,1978
+torch_geometric/__init__.py,sha256=Fn0Tk4CEGmrh27S4w7JxkbV3OwwAKExri7NzXE8-dQE,1978
 torch_geometric/_compile.py,sha256=f-WQeH4VLi5Hn9lrgztFUCSrN_FImjhQa6BxFzcYC38,1338
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -193,7 +193,7 @@ torch_geometric/distributed/utils.py,sha256=FGrr3qw7hx7EQaIjjqasurloCFJ9q_0jt8jd
 torch_geometric/explain/__init__.py,sha256=pRxVB33zsxhED1StRWdHboQWh3e06__g9N298Hzi42Y,359
 torch_geometric/explain/config.py,sha256=_0j67NAwPwjrWHPncNywCT-oKyMiryJNxufxVN1BFlM,7834
 torch_geometric/explain/explainer.py,sha256=8_NZTmlT4WO9RgKxpSUQRt3rbVwFURF5bSWOPlfOLjA,10667
-torch_geometric/explain/explanation.py,sha256=Z2NlgavEnq0QadEr6p6pxAhV6lU7WrlcJLFWbTdsmvg,14903
+torch_geometric/explain/explanation.py,sha256=Bt8THLn-CSrvEFisdT9DX9fnOMaqficsChSCI9uhyQw,18873
 torch_geometric/explain/algorithm/__init__.py,sha256=fE29xbd0bPxg-EfrB2BDmmY9QnyO-7TgvYduGHofm5o,496
 torch_geometric/explain/algorithm/attention_explainer.py,sha256=65iGLmOt00ERtBDVxAoydIchykdWZU24aXzSzUGzQEI,11304
 torch_geometric/explain/algorithm/base.py,sha256=wwJcREUFKDLFUDjRa9o4X3DWqQgMvhS3Iciwb6Evtjc,6922
@@ -633,10 +633,10 @@ torch_geometric/utils/repeat.py,sha256=RxCoRoEisaP6NouXPPW5tY1Rn-tIfrmpJPm0qGP6W
 torch_geometric/utils/smiles.py,sha256=lGQ2BwJ49uBrQfIxxPz8ceTO9Jo-XCjlLxs1ql3xrsA,7130
 torch_geometric/utils/sparse.py,sha256=uYd0oPrp5XN0c2Zc15f-00rhhVMfLnRMqNcqcmILNKQ,25519
 torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5nUAUjw,6222
-torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
-torch_geometric/visualization/graph.py,sha256=ZuLPL92yGRi7lxlqsUPwL_EVVXF7P2kMcveTtW79vpA,4784
+torch_geometric/visualization/__init__.py,sha256=b-HnVesXjyJ_L1N-DnjiRiRVf7lhwKaBQF_2i5YMVSU,208
+torch_geometric/visualization/graph.py,sha256=PoI9tjbEXZVkMUg4CvTLbzqtEfzUwMUcsw57DNBEU0s,14311
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250428.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250428.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250428.dist-info/METADATA,sha256=FIDRBCg3wfpTqyNTcBx3DGEuYtTJ3pvyK7aoAnHFVIs,62979
-pyg_nightly-2.7.0.dev20250428.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250429.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250429.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250429.dist-info/METADATA,sha256=ywTA_H23sidl8wVuvJPmwZWso1slTN7-GwOKR6MSf04,62979
+pyg_nightly-2.7.0.dev20250429.dist-info/RECORD,,

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.dev20250428'
+__version__ = '2.7.0.dev20250429'
 
 __all__ = [
     'Index',

torch_geometric/explain/explanation.py

@@ -1,5 +1,5 @@
 import copy
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -8,7 +8,10 @@ from torch_geometric.data.data import Data, warn_or_raise
 from torch_geometric.data.hetero_data import HeteroData
 from torch_geometric.explain.config import ThresholdConfig, ThresholdType
 from torch_geometric.typing import EdgeType, NodeType
-from torch_geometric.visualization import visualize_graph
+from torch_geometric.visualization import (
+    visualize_graph,
+    visualize_hetero_graph,
+)
 
 
 class ExplanationMixin:
@@ -362,6 +365,87 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
 
         return _visualize_score(score, all_feat_labels, path, top_k)
 
+    def visualize_graph(
+            self,
+            path: Optional[str] = None,
+            node_labels: Optional[Dict[NodeType, List[str]]] = None,
+            node_size_range: Tuple[float, float] = (50, 500),
+            node_opacity_range: Tuple[float, float] = (0.2, 1.0),
+            edge_width_range: Tuple[float, float] = (0.1, 2.0),
+            edge_opacity_range: Tuple[float, float] = (0.2, 1.0),
+    ) -> None:
+        r"""Visualizes the explanation subgraph using networkx, with edge
+        opacity corresponding to edge importance and node colors
+        corresponding to node types.
+
+        Args:
+            path (str, optional): The path to where the plot is saved.
+                If set to :obj:`None`, will visualize the plot on-the-fly.
+                (default: :obj:`None`)
+            node_labels (Dict[NodeType, List[str]], optional): The display
+                names of nodes for each node type that will be shown in the
+                visualization. (default: :obj:`None`)
+            node_size_range (Tuple[float, float], optional): The minimum and
+                maximum node size in the visualization.
+                (default: :obj:`(50, 500)`)
+            node_opacity_range (Tuple[float, float], optional): The minimum and
+                maximum node opacity in the visualization.
+                (default: :obj:`(0.2, 1.0)`)
+            edge_width_range (Tuple[float, float], optional): The minimum and
+                maximum edge width in the visualization.
+                (default: :obj:`(0.1, 2.0)`)
+            edge_opacity_range (Tuple[float, float], optional): The minimum and
+                maximum edge opacity in the visualization.
+                (default: :obj:`(0.2, 1.0)`)
+        """
+        # Validate node labels if provided
+        if node_labels is not None:
+            for node_type, labels in node_labels.items():
+                if node_type not in self.node_types:
+                    raise ValueError(
+                        f"Node type '{node_type}' in node_labels "
+                        f"does not exist in the explanation graph")
+                if len(labels) != self[node_type].num_nodes:
+                    raise ValueError(f"Number of labels for node type "
+                                     f"'{node_type}' (got {len(labels)}) does "
+                                     f"not match the number of nodes "
+                                     f"(got {self[node_type].num_nodes})")
+        # Get the explanation subgraph
+        subgraph = self.get_explanation_subgraph()
+
+        # Prepare edge indices and weights for each edge type
+        edge_index_dict = {}
+        edge_weight_dict = {}
+        for edge_type in subgraph.edge_types:
+            if edge_type[0] == 'x' or edge_type[-1] == 'x':  # Skip edges
+                continue
+            edge_index_dict[edge_type] = subgraph[edge_type].edge_index
+            edge_weight_dict[edge_type] = subgraph[edge_type].get(
+                'edge_mask',
+                torch.ones(subgraph[edge_type].edge_index.size(1)))
+
+        # Prepare node weights for each node type
+        node_weight_dict = {}
+        for node_type in subgraph.node_types:
+            if node_type == 'x':  # Skip the global store
+                continue
+            node_weight_dict[node_type] = subgraph[node_type] \
+                .get('node_mask',
+                     torch.ones(subgraph[node_type].num_nodes)).squeeze(-1)
+
+        # Call the visualization function
+        visualize_hetero_graph(
+            edge_index_dict=edge_index_dict,
+            edge_weight_dict=edge_weight_dict,
+            path=path,
+            node_labels_dict=node_labels,
+            node_weight_dict=node_weight_dict,
+            node_size_range=node_size_range,
+            node_opacity_range=node_opacity_range,
+            edge_width_range=edge_width_range,
+            edge_opacity_range=edge_opacity_range,
+        )
+
 
 def _visualize_score(
     score: torch.Tensor,

torch_geometric/visualization/__init__.py

@@ -1,9 +1,10 @@
 r"""Visualization package."""
 
-from .graph import visualize_graph
+from .graph import visualize_graph, visualize_hetero_graph
 from .influence import influence
 
 __all__ = [
     'visualize_graph',
+    'visualize_hetero_graph',
     'influence',
 ]

torch_geometric/visualization/graph.py

@@ -1,5 +1,5 @@
 from math import sqrt
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional, Set, Tuple
 
 import torch
 from torch import Tensor
@@ -150,3 +150,249 @@ def _visualize_graph_via_networkx(
         plt.show()
 
     plt.close()
+
+
+def visualize_hetero_graph(
+        edge_index_dict: Dict[Tuple[str, str, str], Tensor],
+        edge_weight_dict: Dict[Tuple[str, str, str], Tensor],
+        path: Optional[str] = None,
+        backend: Optional[str] = None,
+        node_labels_dict: Optional[Dict[str, List[str]]] = None,
+        node_weight_dict: Optional[Dict[str, Tensor]] = None,
+        node_size_range: Tuple[float, float] = (50, 500),
+        node_opacity_range: Tuple[float, float] = (0.2, 1.0),
+        edge_width_range: Tuple[float, float] = (0.1, 2.0),
+        edge_opacity_range: Tuple[float, float] = (0.2, 1.0),
+) -> Any:
+    """Visualizes a heterogeneous graph using networkx."""
+    if backend is not None and backend != "networkx":
+        raise ValueError("Only 'networkx' backend is supported")
+
+    # Filter out edges with 0 weight
+    filtered_edge_index_dict = {}
+    filtered_edge_weight_dict = {}
+    for edge_type in edge_index_dict.keys():
+        mask = edge_weight_dict[edge_type] > 0
+        if mask.sum() > 0:
+            filtered_edge_index_dict[edge_type] = edge_index_dict[
+                edge_type][:, mask]
+            filtered_edge_weight_dict[edge_type] = edge_weight_dict[edge_type][
+                mask]
+
+    # Get all unique nodes that are still in the filtered edges
+    remaining_nodes: Dict[str, Set[int]] = {}
+    for edge_type, edge_index in filtered_edge_index_dict.items():
+        src_type, _, dst_type = edge_type
+        if src_type not in remaining_nodes:
+            remaining_nodes[src_type] = set()
+        if dst_type not in remaining_nodes:
+            remaining_nodes[dst_type] = set()
+        remaining_nodes[src_type].update(edge_index[0].tolist())
+        remaining_nodes[dst_type].update(edge_index[1].tolist())
+
+    # Filter node weights to only include remaining nodes
+    if node_weight_dict is not None:
+        filtered_node_weight_dict = {}
+        for node_type, weights in node_weight_dict.items():
+            if node_type in remaining_nodes:
+                mask = torch.zeros(len(weights), dtype=torch.bool)
+                mask[list(remaining_nodes[node_type])] = True
+                filtered_node_weight_dict[node_type] = weights[mask]
+        node_weight_dict = filtered_node_weight_dict
+
+    # Filter node labels to only include remaining nodes
+    if node_labels_dict is not None:
+        filtered_node_labels_dict = {}
+        for node_type, labels in node_labels_dict.items():
+            if node_type in remaining_nodes:
+                filtered_node_labels_dict[node_type] = [
+                    label for i, label in enumerate(labels)
+                    if i in remaining_nodes[node_type]
+                ]
+        node_labels_dict = filtered_node_labels_dict
+
+    return _visualize_hetero_graph_via_networkx(
+        filtered_edge_index_dict,
+        filtered_edge_weight_dict,
+        path,
+        node_labels_dict,
+        node_weight_dict,
+        node_size_range,
+        node_opacity_range,
+        edge_width_range,
+        edge_opacity_range,
+    )
+
+
+def _visualize_hetero_graph_via_networkx(
+        edge_index_dict: Dict[Tuple[str, str, str], Tensor],
+        edge_weight_dict: Dict[Tuple[str, str, str], Tensor],
+        path: Optional[str] = None,
+        node_labels_dict: Optional[Dict[str, List[str]]] = None,
+        node_weight_dict: Optional[Dict[str, Tensor]] = None,
+        node_size_range: Tuple[float, float] = (50, 500),
+        node_opacity_range: Tuple[float, float] = (0.2, 1.0),
+        edge_width_range: Tuple[float, float] = (0.1, 2.0),
+        edge_opacity_range: Tuple[float, float] = (0.2, 1.0),
+) -> Any:
+    import matplotlib.pyplot as plt
+    import networkx as nx
+
+    g = nx.DiGraph()
+    node_offsets: Dict[str, int] = {}
+    current_offset = 0
+
+    # First, collect all unique node types and their counts
+    node_types = set()
+    node_counts: Dict[str, int] = {}
+    remaining_nodes: Dict[str, Set[int]] = {
+    }  # Track which nodes are actually present in edges
+
+    # Get all unique nodes that are in the edges
+    for edge_type in edge_index_dict.keys():
+        src_type, _, dst_type = edge_type
+        node_types.add(src_type)
+        node_types.add(dst_type)
+
+        if src_type not in remaining_nodes:
+            remaining_nodes[src_type] = set()
+        if dst_type not in remaining_nodes:
+            remaining_nodes[dst_type] = set()
+
+        remaining_nodes[src_type].update(
+            edge_index_dict[edge_type][0].tolist())
+        remaining_nodes[dst_type].update(
+            edge_index_dict[edge_type][1].tolist())
+
+    # Set node counts based on remaining nodes
+    for node_type in node_types:
+        node_counts[node_type] = len(remaining_nodes[node_type])
+
+    # Add nodes for each node type
+    for node_type in node_types:
+        num_nodes = node_counts[node_type]
+        node_offsets[node_type] = current_offset
+
+        # Get node weights if provided
+        weights = None
+        if node_weight_dict is not None and node_type in node_weight_dict:
+            weights = node_weight_dict[node_type]
+            if len(weights) != num_nodes:
+                raise ValueError(f"Number of weights for node type "
+                                 f"{node_type} ({len(weights)}) does not "
+                                 f"match number of nodes ({num_nodes})")
+
+        for i in range(num_nodes):
+            node_id = current_offset + i
+            label = (node_labels_dict[node_type][i]
+                     if node_labels_dict is not None
+                     and node_type in node_labels_dict else "")
+
+            # Calculate node size and opacity if weights provided
+            size = node_size_range[1]
+            opacity = node_opacity_range[1]
+            if weights is not None:
+                w = weights[i].item()
+                size = node_size_range[0] + w * \
+                    (node_size_range[1] - node_size_range[0])
+                opacity = node_opacity_range[0] + w * \
+                    (node_opacity_range[1] - node_opacity_range[0])
+
+            g.add_node(node_id, label=label, type=node_type, size=size,
+                       alpha=opacity)
+
+        current_offset += num_nodes
+
+    # Add edges with remapped node indices
+    for edge_type, edge_index in edge_index_dict.items():
+        src_type, _, dst_type = edge_type
+        edge_weight = edge_weight_dict[edge_type]
+        src_offset = node_offsets[src_type]
+        dst_offset = node_offsets[dst_type]
+
+        # Create mappings for source and target nodes
+        src_mapping = {
+            old_idx: new_idx
+            for new_idx, old_idx in enumerate(sorted(
+                remaining_nodes[src_type]))
+        }
+        dst_mapping = {
+            old_idx: new_idx
+            for new_idx, old_idx in enumerate(sorted(
+                remaining_nodes[dst_type]))
+        }
+
+        for (src, dst), w in zip(edge_index.t().tolist(),
+                                 edge_weight.tolist()):
+            # Remap node indices
+            new_src = src_mapping[src] + src_offset
+            new_dst = dst_mapping[dst] + dst_offset
+
+            # Calculate edge width and opacity based on weight
+            width = edge_width_range[0] + w * \
+                (edge_width_range[1] - edge_width_range[0])
+            opacity = edge_opacity_range[0] + w * \
+                (edge_opacity_range[1] - edge_opacity_range[0])
+            g.add_edge(new_src, new_dst, width=width, alpha=opacity)
+
+    # Draw the graph
+    ax = plt.gca()
+    pos = nx.arf_layout(g)
+
+    # Draw edges with arrows
+    for src, dst, data in g.edges(data=True):
+        ax.annotate(
+            '',
+            xy=pos[src],
+            xytext=pos[dst],
+            arrowprops=dict(
+                arrowstyle="<-",
+                alpha=data['alpha'],
+                linewidth=data['width'],
+                shrinkA=sqrt(g.nodes[src]['size']) / 2.0,
+                shrinkB=sqrt(g.nodes[dst]['size']) / 2.0,
+                connectionstyle="arc3,rad=0.1",
+            ),
+        )
+
+    # Draw nodes colored by type
+    node_colors = []
+    node_sizes = []
+    node_alphas = []
+
+    # Use matplotlib tab20 colormap for consistent coloring
+    tab10_cmap = plt.cm.tab10  # type: ignore[attr-defined]
+    node_type_colors: Dict[str, Any] = {}  # Store color for each node type
+    for node in g.nodes():
+        node_type = g.nodes[node]['type']
+        # Assign a consistent color for each node type
+        if node_type not in node_type_colors:
+            color_idx = len(node_type_colors) % 10  # Cycle through colors
+            node_type_colors[node_type] = tab10_cmap(color_idx)
+        node_colors.append(node_type_colors[node_type])
+        node_sizes.append(g.nodes[node]['size'])
+        node_alphas.append(g.nodes[node]['alpha'])
+
+    nx.draw_networkx_nodes(g, pos, node_size=node_sizes,
+                           node_color=node_colors, margins=0.1,
+                           alpha=node_alphas)
+
+    # Draw labels
+    labels = nx.get_node_attributes(g, 'label')
+    nx.draw_networkx_labels(g, pos, labels, font_size=10)
+
+    # Add legend
+    legend_elements = []
+    for node_type, color in node_type_colors.items():
+        legend_elements.append(
+            plt.Line2D([0], [0], marker='o', color='w', label=node_type,
+                       markerfacecolor=color, markersize=10))
+    ax.legend(handles=legend_elements, loc='upper right',
+              bbox_to_anchor=(0.9, 1))
+
+    if path is not None:
+        plt.savefig(path, bbox_inches='tight')
+    else:
+        plt.show()
+
+    plt.close()