pyg-nightly 2.7.0.dev20250428__tar.gz → 2.7.0.dev20250429__tar.gz

{pyg_nightly-2.7.0.dev20250428 → pyg_nightly-2.7.0.dev20250429}/PKG-INFO

@@ -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 → pyg_nightly-2.7.0.dev20250429}/pyproject.toml

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

{pyg_nightly-2.7.0.dev20250428 → pyg_nightly-2.7.0.dev20250429}/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',

{pyg_nightly-2.7.0.dev20250428 → pyg_nightly-2.7.0.dev20250429}/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,

{pyg_nightly-2.7.0.dev20250428 → pyg_nightly-2.7.0.dev20250429}/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',
 ]

pyg_nightly-2.7.0.dev20250429/torch_geometric/visualization/graph.py

@@ -0,0 +1,398 @@
+from math import sqrt
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+import torch
+from torch import Tensor
+
+BACKENDS = {'graphviz', 'networkx'}
+
+
+def has_graphviz() -> bool:
+    try:
+        import graphviz
+    except ImportError:
+        return False
+
+    try:
+        graphviz.Digraph().pipe()
+    except graphviz.backend.ExecutableNotFound:
+        return False
+
+    return True
+
+
+def visualize_graph(
+    edge_index: Tensor,
+    edge_weight: Optional[Tensor] = None,
+    path: Optional[str] = None,
+    backend: Optional[str] = None,
+    node_labels: Optional[List[str]] = None,
+) -> Any:
+    r"""Visualizes the graph given via :obj:`edge_index` and (optional)
+    :obj:`edge_weight`.
+
+    Args:
+        edge_index (torch.Tensor): The edge indices.
+        edge_weight (torch.Tensor, optional): The edge weights.
+        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`)
+        backend (str, optional): The graph drawing backend to use for
+            visualization (:obj:`"graphviz"`, :obj:`"networkx"`).
+            If set to :obj:`None`, will use the most appropriate
+            visualization backend based on available system packages.
+            (default: :obj:`None`)
+        node_labels (List[str], optional): The labels/IDs of nodes.
+            (default: :obj:`None`)
+    """
+    if edge_weight is not None:  # Normalize edge weights.
+        edge_weight = edge_weight - edge_weight.min()
+        edge_weight = edge_weight / edge_weight.max()
+
+    if edge_weight is not None:  # Discard any edges with zero edge weight:
+        mask = edge_weight > 1e-7
+        edge_index = edge_index[:, mask]
+        edge_weight = edge_weight[mask]
+
+    if edge_weight is None:
+        edge_weight = torch.ones(edge_index.size(1))
+
+    if backend is None:
+        backend = 'graphviz' if has_graphviz() else 'networkx'
+
+    if backend.lower() == 'networkx':
+        return _visualize_graph_via_networkx(edge_index, edge_weight, path,
+                                             node_labels)
+    elif backend.lower() == 'graphviz':
+        return _visualize_graph_via_graphviz(edge_index, edge_weight, path,
+                                             node_labels)
+
+    raise ValueError(f"Expected graph drawing backend to be in "
+                     f"{BACKENDS} (got '{backend}')")
+
+
+def _visualize_graph_via_graphviz(
+    edge_index: Tensor,
+    edge_weight: Tensor,
+    path: Optional[str] = None,
+    node_labels: Optional[List[str]] = None,
+) -> Any:
+    import graphviz
+
+    suffix = path.split('.')[-1] if path is not None else None
+    g = graphviz.Digraph('graph', format=suffix)
+    g.attr('node', shape='circle', fontsize='11pt')
+
+    for node in edge_index.view(-1).unique().tolist():
+        g.node(str(node) if node_labels is None else node_labels[node])
+
+    for (src, dst), w in zip(edge_index.t().tolist(), edge_weight.tolist()):
+        hex_color = hex(255 - round(255 * w))[2:]
+        hex_color = f'{hex_color}0' if len(hex_color) == 1 else hex_color
+        if node_labels is not None:
+            src = node_labels[src]
+            dst = node_labels[dst]
+        g.edge(str(src), str(dst), color=f'#{hex_color}{hex_color}{hex_color}')
+
+    if path is not None:
+        path = '.'.join(path.split('.')[:-1])
+        g.render(path, cleanup=True)
+    else:
+        g.view()
+
+    return g
+
+
+def _visualize_graph_via_networkx(
+    edge_index: Tensor,
+    edge_weight: Tensor,
+    path: Optional[str] = None,
+    node_labels: Optional[List[str]] = None,
+) -> Any:
+    import matplotlib.pyplot as plt
+    import networkx as nx
+
+    g = nx.DiGraph()
+    node_size = 800
+
+    for node in edge_index.view(-1).unique().tolist():
+        g.add_node(node if node_labels is None else node_labels[node])
+
+    for (src, dst), w in zip(edge_index.t().tolist(), edge_weight.tolist()):
+        if node_labels is not None:
+            src = node_labels[src]
+            dst = node_labels[dst]
+        g.add_edge(src, dst, alpha=w)
+
+    ax = plt.gca()
+    pos = nx.spring_layout(g)
+    for src, dst, data in g.edges(data=True):
+        ax.annotate(
+            '',
+            xy=pos[src],
+            xytext=pos[dst],
+            arrowprops=dict(
+                arrowstyle="<-",
+                alpha=data['alpha'],
+                shrinkA=sqrt(node_size) / 2.0,
+                shrinkB=sqrt(node_size) / 2.0,
+                connectionstyle="arc3,rad=0.1",
+            ),
+        )
+
+    nx.draw_networkx_nodes(g, pos, node_size=node_size, node_color='white',
+                           margins=0.1, edgecolors='black')
+    nx.draw_networkx_labels(g, pos, font_size=10)
+
+    if path is not None:
+        plt.savefig(path)
+    else:
+        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()

pyg_nightly-2.7.0.dev20250428/torch_geometric/visualization/graph.py

@@ -1,152 +0,0 @@
-from math import sqrt
-from typing import Any, List, Optional
-
-import torch
-from torch import Tensor
-
-BACKENDS = {'graphviz', 'networkx'}
-
-
-def has_graphviz() -> bool:
-    try:
-        import graphviz
-    except ImportError:
-        return False
-
-    try:
-        graphviz.Digraph().pipe()
-    except graphviz.backend.ExecutableNotFound:
-        return False
-
-    return True
-
-
-def visualize_graph(
-    edge_index: Tensor,
-    edge_weight: Optional[Tensor] = None,
-    path: Optional[str] = None,
-    backend: Optional[str] = None,
-    node_labels: Optional[List[str]] = None,
-) -> Any:
-    r"""Visualizes the graph given via :obj:`edge_index` and (optional)
-    :obj:`edge_weight`.
-
-    Args:
-        edge_index (torch.Tensor): The edge indices.
-        edge_weight (torch.Tensor, optional): The edge weights.
-        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`)
-        backend (str, optional): The graph drawing backend to use for
-            visualization (:obj:`"graphviz"`, :obj:`"networkx"`).
-            If set to :obj:`None`, will use the most appropriate
-            visualization backend based on available system packages.
-            (default: :obj:`None`)
-        node_labels (List[str], optional): The labels/IDs of nodes.
-            (default: :obj:`None`)
-    """
-    if edge_weight is not None:  # Normalize edge weights.
-        edge_weight = edge_weight - edge_weight.min()
-        edge_weight = edge_weight / edge_weight.max()
-
-    if edge_weight is not None:  # Discard any edges with zero edge weight:
-        mask = edge_weight > 1e-7
-        edge_index = edge_index[:, mask]
-        edge_weight = edge_weight[mask]
-
-    if edge_weight is None:
-        edge_weight = torch.ones(edge_index.size(1))
-
-    if backend is None:
-        backend = 'graphviz' if has_graphviz() else 'networkx'
-
-    if backend.lower() == 'networkx':
-        return _visualize_graph_via_networkx(edge_index, edge_weight, path,
-                                             node_labels)
-    elif backend.lower() == 'graphviz':
-        return _visualize_graph_via_graphviz(edge_index, edge_weight, path,
-                                             node_labels)
-
-    raise ValueError(f"Expected graph drawing backend to be in "
-                     f"{BACKENDS} (got '{backend}')")
-
-
-def _visualize_graph_via_graphviz(
-    edge_index: Tensor,
-    edge_weight: Tensor,
-    path: Optional[str] = None,
-    node_labels: Optional[List[str]] = None,
-) -> Any:
-    import graphviz
-
-    suffix = path.split('.')[-1] if path is not None else None
-    g = graphviz.Digraph('graph', format=suffix)
-    g.attr('node', shape='circle', fontsize='11pt')
-
-    for node in edge_index.view(-1).unique().tolist():
-        g.node(str(node) if node_labels is None else node_labels[node])
-
-    for (src, dst), w in zip(edge_index.t().tolist(), edge_weight.tolist()):
-        hex_color = hex(255 - round(255 * w))[2:]
-        hex_color = f'{hex_color}0' if len(hex_color) == 1 else hex_color
-        if node_labels is not None:
-            src = node_labels[src]
-            dst = node_labels[dst]
-        g.edge(str(src), str(dst), color=f'#{hex_color}{hex_color}{hex_color}')
-
-    if path is not None:
-        path = '.'.join(path.split('.')[:-1])
-        g.render(path, cleanup=True)
-    else:
-        g.view()
-
-    return g
-
-
-def _visualize_graph_via_networkx(
-    edge_index: Tensor,
-    edge_weight: Tensor,
-    path: Optional[str] = None,
-    node_labels: Optional[List[str]] = None,
-) -> Any:
-    import matplotlib.pyplot as plt
-    import networkx as nx
-
-    g = nx.DiGraph()
-    node_size = 800
-
-    for node in edge_index.view(-1).unique().tolist():
-        g.add_node(node if node_labels is None else node_labels[node])
-
-    for (src, dst), w in zip(edge_index.t().tolist(), edge_weight.tolist()):
-        if node_labels is not None:
-            src = node_labels[src]
-            dst = node_labels[dst]
-        g.add_edge(src, dst, alpha=w)
-
-    ax = plt.gca()
-    pos = nx.spring_layout(g)
-    for src, dst, data in g.edges(data=True):
-        ax.annotate(
-            '',
-            xy=pos[src],
-            xytext=pos[dst],
-            arrowprops=dict(
-                arrowstyle="<-",
-                alpha=data['alpha'],
-                shrinkA=sqrt(node_size) / 2.0,
-                shrinkB=sqrt(node_size) / 2.0,
-                connectionstyle="arc3,rad=0.1",
-            ),
-        )
-
-    nx.draw_networkx_nodes(g, pos, node_size=node_size, node_color='white',
-                           margins=0.1, edgecolors='black')
-    nx.draw_networkx_labels(g, pos, font_size=10)
-
-    if path is not None:
-        plt.savefig(path)
-    else:
-        plt.show()
-
-    plt.close()