pyg-nightly 2.7.0.dev20250403__py3-none-any.whl → 2.7.0.dev20250404__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250403.dist-info → pyg_nightly-2.7.0.dev20250404.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=_qoLOGL1sQztFtP1NfwHKDJYsWreO5B0s-W4MHDYSVs,1978
+torch_geometric/__init__.py,sha256=FXQ6WZdtxVAyouslhtsPt0xXJ33VOV4_aeSKmDUWYrU,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
@@ -426,8 +426,9 @@ torch_geometric/nn/kge/distmult.py,sha256=dGQ0bVzjreZgFN1lXE23_IIidsiOq7ehPrMb-N
 torch_geometric/nn/kge/loader.py,sha256=5Uc1j3OUMQnBYSHDqL7pLCty1siFLzoPkztigYO2zP8,771
 torch_geometric/nn/kge/rotate.py,sha256=XLuO1AbyTt5cJxr97ZzoyAyIEsHKesgW5TvDmnGJAao,3208
 torch_geometric/nn/kge/transe.py,sha256=jlejq5BLMm-sb1wWcLDp7pZqCdelWBgjDIC8ctbjSdU,3088
-torch_geometric/nn/models/__init__.py,sha256=gmBRXrbjkxLv_g0hfI87bSXMgwFoIo3XgslpGGdMs3g,2197
+torch_geometric/nn/models/__init__.py,sha256=dwKH_fLrAMV9nda0nQHLyCzIOPezHCZa01pmTCdaJxU,2263
 torch_geometric/nn/models/attentive_fp.py,sha256=tkgvw28wg9-JqHIfBllfCwTHrZIUiv85yZJcDqjz3z0,6634
+torch_geometric/nn/models/attract_repel.py,sha256=h9OyogT0NY0xiT0DkpJHMxH6ZUmo8R-CmwZdKEwq8Ek,5277
 torch_geometric/nn/models/autoencoder.py,sha256=nGje-zty78Y3hxOJ9o0_6QziJjOvBlknk6z0_fDQwQU,10770
 torch_geometric/nn/models/basic_gnn.py,sha256=PGa0RUMyvrNy_5yRI2jX_zwPsmZXwOQWfsWvxOiHsSk,31225
 torch_geometric/nn/models/captum.py,sha256=kV3lsIM7QdP9MDXmBOTq0Sf1jTEcK-_9LnwLYm2aVzw,3971
@@ -633,7 +634,7 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=ZuLPL92yGRi7lxlqsUPwL_EVVXF7P2kMcveTtW79vpA,4784
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250403.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250403.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250403.dist-info/METADATA,sha256=kIraJoPWCIT5AqAy_reyPy82-5ZLfzUfqqG3j30iYdg,63021
-pyg_nightly-2.7.0.dev20250403.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250404.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250404.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250404.dist-info/METADATA,sha256=AO-euV7W-eeQ5YKe9jRqjlvrpn4qZGiLmiYV_VNxKyI,63021
+pyg_nightly-2.7.0.dev20250404.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.dev20250403'
+__version__ = '2.7.0.dev20250404'
 
 __all__ = [
     'Index',

torch_geometric/nn/models/__init__.py

@@ -36,6 +36,7 @@ from .sgformer import SGFormer
 # Deprecated:
 from torch_geometric.explain.algorithm.captum import (to_captum_input,
                                                       captum_output_to_dicts)
+from .attract_repel import ARLinkPredictor
 
 __all__ = classes = [
     'MLP',
@@ -84,4 +85,5 @@ __all__ = classes = [
     'MoleculeGPT',
     'GLEM',
     'SGFormer',
+    'ARLinkPredictor',
 ]

torch_geometric/nn/models/attract_repel.py

@@ -0,0 +1,148 @@
+import torch
+import torch.nn.functional as F
+
+
+class ARLinkPredictor(torch.nn.Module):
+    r"""Link predictor using Attract-Repel embeddings from the paper
+    `"Pseudo-Euclidean Attract-Repel Embeddings for Undirected Graphs"
+    <https://arxiv.org/abs/2106.09671>`_.
+
+    This model splits node embeddings into: attract and
+    repel.
+    The edge prediction score is computed as the dot product of attract
+    components minus the dot product of repel components.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of hidden embeddings.
+        out_channels (int, optional): Size of output embeddings.
+            If set to :obj:`None`, will default to :obj:`hidden_channels`.
+            (default: :obj:`None`)
+        num_layers (int): Number of message passing layers.
+            (default: :obj:`2`)
+        dropout (float): Dropout probability. (default: :obj:`0.0`)
+        attract_ratio (float): Ratio to use for attract component.
+            Must be between 0 and 1. (default: :obj:`0.5`)
+    """
+    def __init__(self, in_channels, hidden_channels, out_channels=None,
+                 num_layers=2, dropout=0.0, attract_ratio=0.5):
+        super().__init__()
+
+        if out_channels is None:
+            out_channels = hidden_channels
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.dropout = dropout
+
+        if not 0 <= attract_ratio <= 1:
+            raise ValueError(
+                f"attract_ratio must be between 0 and 1, got {attract_ratio}")
+
+        self.attract_ratio = attract_ratio
+        self.attract_dim = int(out_channels * attract_ratio)
+        self.repel_dim = out_channels - self.attract_dim
+
+        # Create model layers
+        self.lins = torch.nn.ModuleList()
+        self.lins.append(torch.nn.Linear(in_channels, hidden_channels))
+
+        for _ in range(num_layers - 2):
+            self.lins.append(torch.nn.Linear(hidden_channels, hidden_channels))
+
+        # Final layer splits into attract and repel components
+        self.lin_attract = torch.nn.Linear(hidden_channels, self.attract_dim)
+        self.lin_repel = torch.nn.Linear(hidden_channels, self.repel_dim)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Reset all learnable parameters."""
+        for lin in self.lins:
+            lin.reset_parameters()
+        self.lin_attract.reset_parameters()
+        self.lin_repel.reset_parameters()
+
+    def encode(self, x, *args, **kwargs):
+        """Encode node features into attract-repel embeddings.
+
+        Args:
+            x (torch.Tensor): Node feature matrix of shape
+                :obj:`[num_nodes, in_channels]`.
+            *args: Variable length argument list
+            **kwargs: Arbitrary keyword arguments
+
+        """
+        for lin in self.lins:
+            x = lin(x)
+            x = F.relu(x)
+            x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # Split into attract and repel components
+        attract_x = self.lin_attract(x)
+        repel_x = self.lin_repel(x)
+
+        return attract_x, repel_x
+
+    def decode(self, attract_z, repel_z, edge_index):
+        """Decode edge scores from attract-repel embeddings.
+
+        Args:
+            attract_z (torch.Tensor): Attract embeddings of shape
+                :obj:`[num_nodes, attract_dim]`.
+            repel_z (torch.Tensor): Repel embeddings of shape
+                :obj:`[num_nodes, repel_dim]`.
+            edge_index (torch.Tensor): Edge indices of shape
+                :obj:`[2, num_edges]`.
+
+        Returns:
+            torch.Tensor: Edge prediction scores.
+        """
+        # Get node embeddings for edges
+        row, col = edge_index
+        attract_z_row = attract_z[row]
+        attract_z_col = attract_z[col]
+        repel_z_row = repel_z[row]
+        repel_z_col = repel_z[col]
+
+        # Compute attract-repel scores
+        attract_score = torch.sum(attract_z_row * attract_z_col, dim=1)
+        repel_score = torch.sum(repel_z_row * repel_z_col, dim=1)
+
+        return attract_score - repel_score
+
+    def forward(self, x, edge_index):
+        """Forward pass for link prediction.
+
+        Args:
+            x (torch.Tensor): Node feature matrix.
+            edge_index (torch.Tensor): Edge indices to predict.
+
+        Returns:
+            torch.Tensor: Predicted edge scores.
+        """
+        # Encode nodes into attract-repel embeddings
+        attract_z, repel_z = self.encode(x)
+
+        # Decode target edges
+        return torch.sigmoid(self.decode(attract_z, repel_z, edge_index))
+
+    def calculate_r_fraction(self, attract_z, repel_z):
+        """Calculate the R-fraction (proportion of energy in repel space).
+
+        Args:
+            attract_z (torch.Tensor): Attract embeddings.
+            repel_z (torch.Tensor): Repel embeddings.
+
+        Returns:
+            float: R-fraction value.
+        """
+        attract_norm_squared = torch.sum(attract_z**2)
+        repel_norm_squared = torch.sum(repel_z**2)
+
+        r_fraction = repel_norm_squared / (attract_norm_squared +
+                                           repel_norm_squared + 1e-10)
+
+        return r_fraction.item()