pyg-nightly 2.6.0.dev20240704__py3-none-any.whl → 2.8.0.dev20251207__py3-none-any.whl

torch_geometric/utils/sparse.py

@@ -1,4 +1,3 @@
-import typing
 import warnings
 from typing import Any, List, Optional, Tuple, Union
 
@@ -71,8 +70,9 @@ def dense_to_sparse(
                          f"three-dimensional (got {adj.dim()} dimensions)")
 
     if mask is not None and adj.dim() == 2:
-        warnings.warn("Mask should not be provided in case the dense "
-                      "adjacency matrix is two-dimensional")
+        warnings.warn(
+            "Mask should not be provided in case the dense "
+            "adjacency matrix is two-dimensional", stacklevel=2)
         mask = None
 
     if mask is not None and mask.dim() != 2:
@@ -124,8 +124,7 @@ def is_torch_sparse_tensor(src: Any) -> bool:
             return True
         if src.layout == torch.sparse_csr:
             return True
-        if (torch_geometric.typing.WITH_PT112
-                and src.layout == torch.sparse_csc):
+        if src.layout == torch.sparse_csc:
             return True
     return False
 
@@ -198,15 +197,23 @@ def to_torch_coo_tensor(
         # edge_attr = edge_attr.expand(edge_index.size(1))
         edge_attr = torch.ones(edge_index.size(1), device=edge_index.device)
 
-    adj = torch.sparse_coo_tensor(
+    if not torch_geometric.typing.WITH_PT21:
+        adj = torch.sparse_coo_tensor(
+            indices=edge_index,
+            values=edge_attr,
+            size=tuple(size) + edge_attr.size()[1:],
+            device=edge_index.device,
+        )
+        adj = adj._coalesced_(True)
+        return adj
+
+    return torch.sparse_coo_tensor(
         indices=edge_index,
         values=edge_attr,
         size=tuple(size) + edge_attr.size()[1:],
         device=edge_index.device,
+        is_coalesced=True,
     )
-    adj = adj._coalesced_(True)
-
-    return adj
 
 
 def to_torch_csr_tensor(
@@ -312,12 +319,6 @@ def to_torch_csc_tensor(
                size=(4, 4), nnz=6, layout=torch.sparse_csc)
 
     """
-    if not torch_geometric.typing.WITH_PT112:
-        if typing.TYPE_CHECKING:
-            raise NotImplementedError
-        return torch_geometric.typing.MockTorchCSCTensor(
-            edge_index, edge_attr, size)
-
     if size is None:
         size = int(edge_index.max()) + 1
 
@@ -384,7 +385,7 @@ def to_torch_sparse_tensor(
         return to_torch_coo_tensor(edge_index, edge_attr, size, is_coalesced)
     if layout == torch.sparse_csr:
         return to_torch_csr_tensor(edge_index, edge_attr, size, is_coalesced)
-    if torch_geometric.typing.WITH_PT112 and layout == torch.sparse_csc:
+    if layout == torch.sparse_csc:
         return to_torch_csc_tensor(edge_index, edge_attr, size, is_coalesced)
 
     raise ValueError(f"Unexpected sparse tensor layout (got '{layout}')")
@@ -423,7 +424,7 @@ def to_edge_index(adj: Union[Tensor, SparseTensor]) -> Tuple[Tensor, Tensor]:
         col = adj.col_indices().detach()
         return torch.stack([row, col], dim=0).long(), adj.values()
 
-    if torch_geometric.typing.WITH_PT112 and adj.layout == torch.sparse_csc:
+    if adj.layout == torch.sparse_csc:
         col = ptr2index(adj.ccol_indices().detach())
         row = adj.row_indices().detach()
         return torch.stack([row, col], dim=0).long(), adj.values()
@@ -472,7 +473,7 @@ def set_sparse_value(adj: Tensor, value: Tensor) -> Tensor:
             device=value.device,
         )
 
-    if torch_geometric.typing.WITH_PT112 and adj.layout == torch.sparse_csc:
+    if adj.layout == torch.sparse_csc:
         return torch.sparse_csc_tensor(
             ccol_indices=adj.ccol_indices(),
             row_indices=adj.row_indices(),
@@ -531,18 +532,25 @@ def cat_coo(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
             if not tensor.is_coalesced():
                 is_coalesced = False
 
-    out = torch.sparse_coo_tensor(
+    if not torch_geometric.typing.WITH_PT21:
+        out = torch.sparse_coo_tensor(
+            indices=torch.cat(indices, dim=-1),
+            values=torch.cat(values),
+            size=(num_rows, num_cols) + values[-1].size()[1:],
+            device=tensor.device,
+        )
+        if is_coalesced:
+            out = out._coalesced_(True)
+        return out
+
+    return torch.sparse_coo_tensor(
         indices=torch.cat(indices, dim=-1),
         values=torch.cat(values),
         size=(num_rows, num_cols) + values[-1].size()[1:],
         device=tensor.device,
+        is_coalesced=True if is_coalesced else None,
     )
 
-    if is_coalesced:
-        out = out._coalesced_(True)
-
-    return out
-
 
 def cat_csr(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
     assert dim in {0, 1, (0, 1)}

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
@@ -132,7 +132,7 @@ def _visualize_graph_via_networkx(
             xy=pos[src],
             xytext=pos[dst],
             arrowprops=dict(
-                arrowstyle="->",
+                arrowstyle="<-",
                 alpha=data['alpha'],
                 shrinkA=sqrt(node_size) / 2.0,
                 shrinkB=sqrt(node_size) / 2.0,
@@ -140,9 +140,8 @@ def _visualize_graph_via_networkx(
             ),
         )
 
-    nodes = nx.draw_networkx_nodes(g, pos, node_size=node_size,
-                                   node_color='white', margins=0.1)
-    nodes.set_edgecolor('black')
+    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:
@@ -151,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] = (1.0, 1.0),
+        edge_width_range: Tuple[float, float] = (0.1, 2.0),
+        edge_opacity_range: Tuple[float, float] = (1.0, 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] = (1.0, 1.0),
+        edge_width_range: Tuple[float, float] = (0.1, 2.0),
+        edge_opacity_range: Tuple[float, float] = (1.0, 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()

torch_geometric/warnings.py

@@ -4,11 +4,11 @@ from typing import Literal
 import torch_geometric
 
 
-def warn(message: str) -> None:
+def warn(message: str, stacklevel: int = 5) -> None:
     if torch_geometric.is_compiling():
         return
 
-    warnings.warn(message)
+    warnings.warn(message, stacklevel=stacklevel)
 
 
 def filterwarnings(
@@ -19,3 +19,12 @@ def filterwarnings(
         return
 
     warnings.filterwarnings(action, message)
+
+
+class WarningCache(set):
+    """Cache for warnings."""
+    def warn(self, message: str, stacklevel: int = 5) -> None:
+        """Trigger warning message."""
+        if message not in self:
+            self.add(message)
+            warn(message, stacklevel=stacklevel)

torch_geometric/nn/nlp/__init__.py

@@ -1,7 +0,0 @@
-from .sentence_transformer import SentenceTransformer
-from .llm import LLM
-
-__all__ = classes = [
-    'SentenceTransformer',
-    'LLM',
-]