pyg-nightly 2.7.0.dev20250701__py3-none-any.whl → 2.7.0.dev20250703__py3-none-any.whl

torch_geometric/nn/models/protein_mpnn.py

@@ -0,0 +1,304 @@
+from itertools import product
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+
+from torch_geometric.nn import knn_graph
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.utils import to_dense_adj, to_dense_batch
+
+
+class PositionWiseFeedForward(torch.nn.Module):
+    def __init__(self, in_channels: int, hidden_channels: int) -> None:
+        super().__init__()
+        self.out = torch.nn.Sequential(
+            torch.nn.Linear(in_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, in_channels),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.out(x)
+
+
+class PositionalEncoding(torch.nn.Module):
+    def __init__(self, hidden_channels: int,
+                 max_relative_feature: int = 32) -> None:
+        super().__init__()
+        self.max_relative_feature = max_relative_feature
+        self.emb = torch.nn.Embedding(2 * max_relative_feature + 2,
+                                      hidden_channels)
+
+    def forward(self, offset, mask) -> torch.Tensor:
+        d = torch.clip(offset + self.max_relative_feature, 0,
+                       2 * self.max_relative_feature) * mask + (1 - mask) * (
+                           2 * self.max_relative_feature + 1)  # noqa: E501
+        return self.emb(d.long())
+
+
+class Encoder(MessagePassing):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        dropout: float = 0.1,
+        scale: float = 30,
+    ) -> None:
+        super().__init__()
+        self.out_v = torch.nn.Sequential(
+            torch.nn.Linear(in_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, hidden_channels),
+        )
+        self.out_e = torch.nn.Sequential(
+            torch.nn.Linear(in_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, hidden_channels),
+        )
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+        self.dropout3 = torch.nn.Dropout(dropout)
+        self.norm1 = torch.nn.LayerNorm(hidden_channels)
+        self.norm2 = torch.nn.LayerNorm(hidden_channels)
+        self.norm3 = torch.nn.LayerNorm(hidden_channels)
+        self.scale = scale
+        self.dense = PositionWiseFeedForward(hidden_channels,
+                                             hidden_channels * 4)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: torch.Tensor,
+    ) -> torch.Tensor:
+        # x: [N, d_v]
+        # edge_index: [2, E]
+        # edge_attr: [E, d_e]
+        # update node features
+        h_message = self.propagate(x=x, edge_index=edge_index,
+                                   edge_attr=edge_attr)
+        dh = h_message / self.scale
+        x = self.norm1(x + self.dropout1(dh))
+        dh = self.dense(x)
+        x = self.norm2(x + self.dropout2(dh))
+        # update edge features
+        row, col = edge_index
+        x_i, x_j = x[row], x[col]
+        h_e = torch.cat([x_i, x_j, edge_attr], dim=-1)
+        h_e = self.out_e(h_e)
+        edge_attr = self.norm3(edge_attr + self.dropout3(h_e))
+        return x, edge_attr
+
+    def message(self, x_i: torch.Tensor, x_j: torch.Tensor,
+                edge_attr: torch.Tensor) -> torch.Tensor:
+        h = torch.cat([x_i, x_j, edge_attr], dim=-1)  # [E, 2*d_v + d_e]
+        h = self.out_e(h)  # [E, d_e]
+        return h
+
+
+class Decoder(MessagePassing):
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        dropout: float = 0.1,
+        scale: float = 30,
+    ) -> None:
+        super().__init__()
+        self.out_v = torch.nn.Sequential(
+            torch.nn.Linear(in_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, hidden_channels),
+            torch.nn.GELU(),
+            torch.nn.Linear(hidden_channels, hidden_channels),
+        )
+        self.dropout1 = torch.nn.Dropout(dropout)
+        self.dropout2 = torch.nn.Dropout(dropout)
+        self.norm1 = torch.nn.LayerNorm(hidden_channels)
+        self.norm2 = torch.nn.LayerNorm(hidden_channels)
+        self.scale = scale
+        self.dense = PositionWiseFeedForward(hidden_channels,
+                                             hidden_channels * 4)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: torch.Tensor,
+        x_label: torch.Tensor,
+        mask: torch.Tensor,
+    ) -> torch.Tensor:
+        # x: [N, d_v]
+        # edge_index: [2, E]
+        # edge_attr: [E, d_e]
+        h_message = self.propagate(x=x, x_label=x_label, edge_index=edge_index,
+                                   edge_attr=edge_attr, mask=mask)
+        dh = h_message / self.scale
+        x = self.norm1(x + self.dropout1(dh))
+        dh = self.dense(x)
+        x = self.norm2(x + self.dropout2(dh))
+        return x
+
+    def message(self, x_i: torch.Tensor, x_j: torch.Tensor,
+                x_label_j: torch.Tensor, edge_attr: torch.Tensor,
+                mask: torch.Tensor) -> torch.Tensor:
+        h_1 = torch.cat([x_j, edge_attr, x_label_j], dim=-1)
+        h_0 = torch.cat([x_j, edge_attr, torch.zeros_like(x_label_j)], dim=-1)
+        h = h_1 * mask + h_0 * (1 - mask)
+        h = torch.concat([x_i, h], dim=-1)
+        h = self.out_v(h)
+        return h
+
+
+class ProteinMPNN(torch.nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int = 128,
+        num_encoder_layers: int = 3,
+        num_decoder_layers: int = 3,
+        num_neighbors: int = 30,
+        num_rbf: int = 16,
+        dropout: float = 0.1,
+        augment_eps: float = 0.2,
+        num_positional_embedding: int = 16,
+        vocab_size: int = 21,
+    ) -> None:
+        super().__init__()
+        self.augment_eps = augment_eps
+        self.hidden_dim = hidden_dim
+        self.num_neighbors = num_neighbors
+        self.num_rbf = num_rbf
+        self.embedding = PositionalEncoding(num_positional_embedding)
+        self.edge_mlp = torch.nn.Sequential(
+            torch.nn.Linear(num_positional_embedding + 400, hidden_dim),
+            torch.nn.LayerNorm(hidden_dim),
+            torch.nn.Linear(hidden_dim, hidden_dim),
+        )
+        self.label_embedding = torch.nn.Embedding(vocab_size, hidden_dim)
+        self.encoder_layers = torch.nn.ModuleList([
+            Encoder(hidden_dim * 3, hidden_dim, dropout)
+            for _ in range(num_encoder_layers)
+        ])
+
+        self.decoder_layers = torch.nn.ModuleList([
+            Decoder(hidden_dim * 4, hidden_dim, dropout)
+            for _ in range(num_decoder_layers)
+        ])
+        self.output = torch.nn.Linear(hidden_dim, vocab_size)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                torch.nn.init.xavier_uniform_(p)
+
+    def _featurize(
+        self,
+        x: torch.Tensor,
+        mask: torch.Tensor,
+        batch: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        N, Ca, C, O = (x[:, i, :] for i in range(4))  # noqa: E741
+        b = Ca - N
+        c = C - Ca
+        a = torch.cross(b, c, dim=-1)
+        Cb = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + Ca
+
+        valid_mask = mask.bool()
+        valid_Ca = Ca[valid_mask]
+        valid_batch = batch[valid_mask]
+
+        edge_index = knn_graph(valid_Ca, k=self.num_neighbors,
+                               batch=valid_batch, loop=True)
+
+        row, col = edge_index
+        original_indices = torch.arange(Ca.size(0),
+                                        device=x.device)[valid_mask]
+        edge_index_original = torch.stack(
+            [original_indices[row], original_indices[col]], dim=0)
+        row, col = edge_index_original
+
+        rbf_all = []
+        for A, B in list(product([N, Ca, C, O, Cb], repeat=2)):
+            distances = torch.sqrt(torch.sum((A[row] - B[col])**2, 1) + 1e-6)
+            rbf = self._rbf(distances)
+            rbf_all.append(rbf)
+
+        return edge_index_original, torch.cat(rbf_all, dim=-1)
+
+    def _rbf(self, D: torch.Tensor) -> torch.Tensor:
+        D_min, D_max, D_count = 2., 22., self.num_rbf
+        D_mu = torch.linspace(D_min, D_max, D_count, device=D.device)
+        D_mu = D_mu.view([1, -1])
+        D_sigma = (D_max - D_min) / D_count
+        D_expand = torch.unsqueeze(D, -1)
+        RBF = torch.exp(-((D_expand - D_mu) / D_sigma)**2)
+        return RBF
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        chain_seq_label: torch.Tensor,
+        mask: torch.Tensor,
+        chain_mask_all: torch.Tensor,
+        residue_idx: torch.Tensor,
+        chain_encoding_all: torch.Tensor,
+        batch: torch.Tensor,
+    ) -> torch.Tensor:
+        device = x.device
+        if self.training and self.augment_eps > 0:
+            x = x + self.augment_eps * torch.randn_like(x)
+
+        edge_index, edge_attr = self._featurize(x, mask, batch)
+
+        row, col = edge_index
+        offset = residue_idx[row] - residue_idx[col]
+        # find self vs non-self interaction
+        e_chains = ((chain_encoding_all[row] -
+                     chain_encoding_all[col]) == 0).long()
+        e_pos = self.embedding(offset, e_chains)
+        h_e = self.edge_mlp(torch.cat([edge_attr, e_pos], dim=-1))
+        h_v = torch.zeros(x.size(0), self.hidden_dim, device=x.device)
+
+        # encoder
+        for encoder in self.encoder_layers:
+            h_v, h_e = encoder(h_v, edge_index, h_e)
+
+        # mask
+        h_label = self.label_embedding(chain_seq_label)
+        batch_chain_mask_all, _ = to_dense_batch(chain_mask_all * mask,
+                                                 batch)  # [B, N]
+        # 0 - visible - encoder, 1 - masked - decoder
+        decoding_order = torch.argsort(
+            (batch_chain_mask_all + 1e-4) * (torch.abs(
+                torch.randn(batch_chain_mask_all.shape, device=device))))
+        mask_size = batch_chain_mask_all.size(1)
+        permutation_matrix_reverse = F.one_hot(decoding_order,
+                                               num_classes=mask_size).float()
+        order_mask_backward = torch.einsum(
+            'ij, biq, bjp->bqp',
+            1 - torch.triu(torch.ones(mask_size, mask_size, device=device)),
+            permutation_matrix_reverse,
+            permutation_matrix_reverse,
+        )
+        adj = to_dense_adj(edge_index, batch)
+        mask_attend = order_mask_backward[adj.bool()].unsqueeze(-1)
+
+        # decoder
+        for decoder in self.decoder_layers:
+            h_v = decoder(
+                h_v,
+                edge_index,
+                h_e,
+                h_label,
+                mask_attend,
+            )
+
+        logits = self.output(h_v)
+        return F.log_softmax(logits, dim=-1)

torch_geometric/nn/norm/batch_norm.py

@@ -39,6 +39,8 @@ class BatchNorm(torch.nn.Module):
             with only a single element will work as during in evaluation.
             That is the running mean and variance will be used.
             Requires :obj:`track_running_stats=True`. (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -48,6 +50,7 @@ class BatchNorm(torch.nn.Module):
         affine: bool = True,
         track_running_stats: bool = True,
         allow_single_element: bool = False,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -56,7 +59,7 @@ class BatchNorm(torch.nn.Module):
                              "'track_running_stats' to be set to `True`")
 
         self.module = torch.nn.BatchNorm1d(in_channels, eps, momentum, affine,
-                                           track_running_stats)
+                                           track_running_stats, device=device)
         self.in_channels = in_channels
         self.allow_single_element = allow_single_element
 
@@ -114,6 +117,8 @@ class HeteroBatchNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -123,6 +128,7 @@ class HeteroBatchNorm(torch.nn.Module):
         momentum: Optional[float] = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -134,17 +140,21 @@ class HeteroBatchNorm(torch.nn.Module):
         self.track_running_stats = track_running_stats
 
         if self.affine:
-            self.weight = Parameter(torch.empty(num_types, in_channels))
-            self.bias = Parameter(torch.empty(num_types, in_channels))
+            self.weight = Parameter(
+                torch.empty(num_types, in_channels, device=device))
+            self.bias = Parameter(
+                torch.empty(num_types, in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
 
         if self.track_running_stats:
-            self.register_buffer('running_mean',
-                                 torch.empty(num_types, in_channels))
-            self.register_buffer('running_var',
-                                 torch.empty(num_types, in_channels))
+            self.register_buffer(
+                'running_mean',
+                torch.empty(num_types, in_channels, device=device))
+            self.register_buffer(
+                'running_var',
+                torch.empty(num_types, in_channels, device=device))
             self.register_buffer('num_batches_tracked', torch.tensor(0))
         else:
             self.register_buffer('running_mean', None)

torch_geometric/nn/norm/diff_group_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import Tensor
 from torch.nn import BatchNorm1d, Linear
@@ -39,6 +41,8 @@ class DiffGroupNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -49,6 +53,7 @@ class DiffGroupNorm(torch.nn.Module):
         momentum: float = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -56,9 +61,9 @@ class DiffGroupNorm(torch.nn.Module):
         self.groups = groups
         self.lamda = lamda
 
-        self.lin = Linear(in_channels, groups, bias=False)
+        self.lin = Linear(in_channels, groups, bias=False, device=device)
         self.norm = BatchNorm1d(groups * in_channels, eps, momentum, affine,
-                                track_running_stats)
+                                track_running_stats, device=device)
 
         self.reset_parameters()
 

torch_geometric/nn/norm/graph_norm.py

@@ -26,16 +26,21 @@ class GraphNorm(torch.nn.Module):
         in_channels (int): Size of each input sample.
         eps (float, optional): A value added to the denominator for numerical
             stability. (default: :obj:`1e-5`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, in_channels: int, eps: float = 1e-5):
+    def __init__(self, in_channels: int, eps: float = 1e-5,
+                 device: Optional[torch.device] = None):
         super().__init__()
 
         self.in_channels = in_channels
         self.eps = eps
 
-        self.weight = torch.nn.Parameter(torch.empty(in_channels))
-        self.bias = torch.nn.Parameter(torch.empty(in_channels))
-        self.mean_scale = torch.nn.Parameter(torch.empty(in_channels))
+        self.weight = torch.nn.Parameter(
+            torch.empty(in_channels, device=device))
+        self.bias = torch.nn.Parameter(torch.empty(in_channels, device=device))
+        self.mean_scale = torch.nn.Parameter(
+            torch.empty(in_channels, device=device))
 
         self.reset_parameters()
 

torch_geometric/nn/norm/instance_norm.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import torch
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.modules.instancenorm import _InstanceNorm
@@ -36,6 +37,8 @@ class InstanceNorm(_InstanceNorm):
             :obj:`False`, this module does not track such statistics and always
             uses instance statistics in both training and eval modes.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -44,9 +47,10 @@ class InstanceNorm(_InstanceNorm):
         momentum: float = 0.1,
         affine: bool = False,
         track_running_stats: bool = False,
+        device: Optional[torch.device] = None,
     ):
         super().__init__(in_channels, eps, momentum, affine,
-                         track_running_stats)
+                         track_running_stats, device=device)
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""

torch_geometric/nn/norm/layer_norm.py

@@ -35,6 +35,8 @@ class LayerNorm(torch.nn.Module):
             is used, each graph will be considered as an element to be
             normalized. If `"node"` is used, each node will be considered as
             an element to be normalized. (default: :obj:`"graph"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -42,6 +44,7 @@ class LayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'graph',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -51,8 +54,8 @@ class LayerNorm(torch.nn.Module):
         self.mode = mode
 
         if affine:
-            self.weight = Parameter(torch.empty(in_channels))
-            self.bias = Parameter(torch.empty(in_channels))
+            self.weight = Parameter(torch.empty(in_channels, device=device))
+            self.bias = Parameter(torch.empty(in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
@@ -134,6 +137,8 @@ class HeteroLayerNorm(torch.nn.Module):
             normalization (:obj:`"node"`). If `"node"` is used, each node will
             be considered as an element to be normalized.
             (default: :obj:`"node"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -142,6 +147,7 @@ class HeteroLayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'node',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
         assert mode == 'node'
@@ -152,8 +158,10 @@ class HeteroLayerNorm(torch.nn.Module):
         self.affine = affine
 
         if affine:
-            self.weight = Parameter(torch.empty(num_types, in_channels))
-            self.bias = Parameter(torch.empty(num_types, in_channels))
+            self.weight = Parameter(
+                torch.empty(num_types, in_channels, device=device))
+            self.bias = Parameter(
+                torch.empty(num_types, in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)

torch_geometric/nn/norm/msg_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -19,10 +21,14 @@ class MessageNorm(torch.nn.Module):
         learn_scale (bool, optional): If set to :obj:`True`, will learn the
             scaling factor :math:`s` of message normalization.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, learn_scale: bool = False):
+    def __init__(self, learn_scale: bool = False,
+                 device: Optional[torch.device] = None):
         super().__init__()
-        self.scale = Parameter(torch.empty(1), requires_grad=learn_scale)
+        self.scale = Parameter(torch.empty(1, device=device),
+                               requires_grad=learn_scale)
         self.reset_parameters()
 
     def reset_parameters(self):

torch_geometric/utils/convert.py

@@ -452,15 +452,22 @@ def to_cugraph(
     g = cugraph.Graph(directed=directed)
     df = cudf.from_dlpack(to_dlpack(edge_index.t()))
 
+    df = cudf.DataFrame({
+        'source':
+        cudf.from_dlpack(to_dlpack(edge_index[0])),
+        'destination':
+        cudf.from_dlpack(to_dlpack(edge_index[1])),
+    })
+
     if edge_weight is not None:
         assert edge_weight.dim() == 1
-        df['2'] = cudf.from_dlpack(to_dlpack(edge_weight))
+        df['weight'] = cudf.from_dlpack(to_dlpack(edge_weight))
 
     g.from_cudf_edgelist(
         df,
-        source=0,
-        destination=1,
-        edge_attr='2' if edge_weight is not None else None,
+        source='source',
+        destination='destination',
+        edge_attr='weight' if edge_weight is not None else None,
         renumber=relabel_nodes,
     )
 
@@ -476,13 +483,13 @@ def from_cugraph(g: Any) -> Tuple[Tensor, Optional[Tensor]]:
     """
     df = g.view_edge_list()
 
-    src = from_dlpack(df[0].to_dlpack()).long()
-    dst = from_dlpack(df[1].to_dlpack()).long()
+    src = from_dlpack(df[g.source_columns].to_dlpack()).long()
+    dst = from_dlpack(df[g.destination_columns].to_dlpack()).long()
     edge_index = torch.stack([src, dst], dim=0)
 
     edge_weight = None
-    if '2' in df:
-        edge_weight = from_dlpack(df['2'].to_dlpack())
+    if g.weight_column is not None:
+        edge_weight = from_dlpack(df[g.weight_column].to_dlpack())
 
     return edge_index, edge_weight
 

torch_geometric/utils/smiles.py

@@ -148,7 +148,7 @@ def from_smiles(
     """
     from rdkit import Chem, RDLogger
 
-    RDLogger.DisableLog('rdApp.*')  # type: ignore
+    RDLogger.DisableLog('rdApp.*')
 
     mol = Chem.MolFromSmiles(smiles)