pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__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
 
@@ -320,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
 
@@ -392,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}')")
@@ -431,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()
@@ -480,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(),

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',
-]

torch_geometric/nn/nlp/sentence_transformer.py

@@ -1,101 +0,0 @@
-from enum import Enum
-from typing import List, Optional, Union
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-
-
-class PoolingStrategy(Enum):
-    MEAN = 'mean'
-    LAST = 'last'
-    CLS = 'cls'
-
-
-class SentenceTransformer(torch.nn.Module):
-    def __init__(
-        self,
-        model_name: str,
-        pooling_strategy: Union[PoolingStrategy, str] = 'mean',
-    ) -> None:
-        super().__init__()
-
-        self.model_name = model_name
-        self.pooling_strategy = PoolingStrategy(pooling_strategy)
-
-        from transformers import AutoModel, AutoTokenizer
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModel.from_pretrained(model_name)
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-
-    def forward(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
-        out = self.model(input_ids=input_ids, attention_mask=attention_mask)
-
-        emb = out[0]  # First element contains all token embeddings.
-        if self.pooling_strategy == PoolingStrategy.MEAN:
-            emb = mean_pooling(emb, attention_mask)
-        elif self.pooling_strategy == PoolingStrategy.LAST:
-            emb = last_pooling(emb, attention_mask)
-        else:
-            assert self.pooling_strategy == PoolingStrategy.CLS
-            emb = emb[:, 0, :]
-
-        emb = F.normalize(emb, p=2, dim=1)
-        return emb
-
-    @property
-    def device(self) -> torch.device:
-        return next(iter(self.model.parameters())).device
-
-    @torch.no_grad()
-    def encode(
-        self,
-        text: List[str],
-        batch_size: Optional[int] = None,
-        output_device: Optional[Union[torch.device, str]] = None,
-    ) -> Tensor:
-        is_empty = len(text) == 0
-        text = ['dummy'] if is_empty else text
-
-        batch_size = len(text) if batch_size is None else batch_size
-
-        embs: List[Tensor] = []
-        for start in range(0, len(text), batch_size):
-            token = self.tokenizer(
-                text[start:start + batch_size],
-                padding=True,
-                truncation=True,
-                return_tensors='pt',
-            )
-
-            emb = self(
-                input_ids=token.input_ids.to(self.device),
-                attention_mask=token.attention_mask.to(self.device),
-            ).to(output_device)
-
-            embs.append(emb)
-
-        out = torch.cat(embs, dim=0) if len(embs) > 1 else embs[0]
-        out = out[:0] if is_empty else out
-        return out
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}(model_name={self.model_name})'
-
-
-def mean_pooling(emb: Tensor, attention_mask: Tensor) -> Tensor:
-    mask = attention_mask.unsqueeze(-1).expand(emb.size()).to(emb.dtype)
-    return (emb * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
-
-
-def last_pooling(emb: Tensor, attention_mask: Tensor) -> Tensor:
-    # Check whether language model uses left padding,
-    # which is always used for decoder LLMs
-    left_padding = attention_mask[:, -1].sum() == attention_mask.size(0)
-    if left_padding:
-        return emb[:, -1]
-
-    seq_indices = attention_mask.sum(dim=1) - 1
-    return emb[torch.arange(emb.size(0), device=emb.device), seq_indices]