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

{pyg_nightly-2.7.0.dev20250702.dist-info → pyg_nightly-2.7.0.dev20250703.dist-info}/METADATA

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=ap-t4q8f9aTE0oAW_K5390u2Mlk8-S76rdeUEgPzglo,2250
+torch_geometric/__init__.py,sha256=8_AcpgPpDfVr7gwK3sYPD-pu9JYUmziOQzs9SPRvqcE,2250
 torch_geometric/_compile.py,sha256=9yqMTBKatZPr40WavJz9FjNi7pQj8YZAZOyZmmRGXgc,1351
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -55,7 +55,7 @@ torch_geometric/data/temporal.py,sha256=WOJ6gFrTLikaLhUvotyUF5ql14FkE5Ox3hNkdSp6
 torch_geometric/data/view.py,sha256=XjkVSc-UWZFCT4DlXLShZtO8duhFQkS9gq88zZXANsk,1089
 torch_geometric/data/lightning/__init__.py,sha256=w3En1tJfy3kSqe1MycpOyZpHFO3fxBCgNCUOznPA3YU,178
 torch_geometric/data/lightning/datamodule.py,sha256=IjucsIKRcNv16DIqILQnqa_sz72a4-yivoySmEllv2o,29353
-torch_geometric/datasets/__init__.py,sha256=vIraHnDqD40Num-XwwNivjHQDboK9tmMvlZHjTAuljM,6291
+torch_geometric/datasets/__init__.py,sha256=rgfUmjd9U3o8renKVl81Brscx4LOtwWmt6qAoaG41C4,6417
 torch_geometric/datasets/actor.py,sha256=oUxgJIX8bi5hJr1etWNYIFyVQNDDXi1nyVpHGGMEAGQ,4304
 torch_geometric/datasets/airfrans.py,sha256=8cCBmHPttrlKY_iwfyr-K-CUX_JEDjrIOg3r9dQSN7o,5439
 torch_geometric/datasets/airports.py,sha256=b3gkv3gY2JkUpmGiz36Z-g7EcnSfU8lBG1YsCOWdJ6k,3758
@@ -95,7 +95,7 @@ torch_geometric/datasets/gdelt_lite.py,sha256=zE1WagpgmsQARQhEgdCBtALRKyuQvIZqxT
 torch_geometric/datasets/ged_dataset.py,sha256=dtd-C6pCygNHLXgVfg3ZTWtTVHKT13Q3GlGrze1_rpo,9551
 torch_geometric/datasets/gemsec.py,sha256=oMTSryTgyed9z_4ydg3ql12KM-_35uqL1AoNls5nG8M,2820
 torch_geometric/datasets/geometry.py,sha256=-BxUMirZcUOf01c3avvF0b6wGPn-4S3Zj3Oau1RaJVk,4223
-torch_geometric/datasets/git_mol_dataset.py,sha256=LsS_dPYUpwhWXMBh17iT7IbjlLOP0fFzb-we9cuBDaQ,10681
+torch_geometric/datasets/git_mol_dataset.py,sha256=l5u4U86tfjJdHtQPN7SM3Yjv25LD1Idtm7VHaqJqNik,10665
 torch_geometric/datasets/github.py,sha256=Qhqhkvi6eZ8VF_HqP1rL2iYToZavFNsQh7J1WdeM9dA,2687
 torch_geometric/datasets/gnn_benchmark_dataset.py,sha256=4P8n7czF-gf1egLYlAcSSvfB0GXIKpAbH5UjsuFld1M,6976
 torch_geometric/datasets/heterophilous_graph_dataset.py,sha256=yHHtwl4uPrid0vPOxvPV3sIS8HWdswar8FJ0h0OQ9is,4224
@@ -139,8 +139,9 @@ torch_geometric/datasets/pcqm4m.py,sha256=7ID_xXXIAyuNzYLI2lBWygZl9wGos-dbaz1b6E
 torch_geometric/datasets/planetoid.py,sha256=RksfwR_PI7qGVphs-T-4jXDepYwQCweMXElLm096hgg,7201
 torch_geometric/datasets/polblogs.py,sha256=IYzsvd4R0OojmOOZUoOdCwQYfwcTfth1PNtcBK1yOGc,3045
 torch_geometric/datasets/ppi.py,sha256=zPtg-omC7WYvr9Tzwkb7zNjpXLODsvxKxKdGEUswp2E,5030
+torch_geometric/datasets/protein_mpnn_dataset.py,sha256=TTeTVJMo0Rlt2_h9bbZMKJe3rTJcjCgY5cXGyWteBfA,17756
 torch_geometric/datasets/qm7.py,sha256=bYyK8xlh9kTr5vqueNbLu9EAjIXkQH1KX1VWnjKfOJc,3323
-torch_geometric/datasets/qm9.py,sha256=XU2HTPbgJJ_6hT--X0J2xkXliCbt7_-hub9nuIUQlug,17213
+torch_geometric/datasets/qm9.py,sha256=Ub1t8KNeWFZvw50_Qk-80yNFeYFDwdAeyQtp3JHZs7o,17197
 torch_geometric/datasets/rcdd.py,sha256=gvOoM1tw_X5QMyBB4FkMUwNErMXAvImyjz5twktBAh8,5317
 torch_geometric/datasets/reddit.py,sha256=QUgiKTaj6YTOYbgWgqV8mPYsctOui2ujaM8f8qy81v0,3131
 torch_geometric/datasets/reddit2.py,sha256=WSdrhbDPcUEG37XWNUd0uKnqgI911MOcfjXmgjbTPoQ,4291
@@ -153,6 +154,7 @@ torch_geometric/datasets/snap_dataset.py,sha256=deJvB6cpIQ3bu_pcWoqgEo1-Kl_NcFi7
 torch_geometric/datasets/suite_sparse.py,sha256=eqjH4vAUq872qdk3YdLkZSwlu6r7HHpTgK0vEVGmY1s,3278
 torch_geometric/datasets/tag_dataset.py,sha256=qTnwr2N1tbWYeLGbItfv70UxQ3n1rKesjeVU3kcOCP8,14757
 torch_geometric/datasets/taobao.py,sha256=CUcZpbWsNTasevflO8zqP0YvENy89P7wpKS4MHaDJ6Q,4170
+torch_geometric/datasets/teeth3ds.py,sha256=hZvhcq9lsQENNFr5hk50w2T3CgxE_tlnQfrCgN6uIDQ,9919
 torch_geometric/datasets/tosca.py,sha256=nUSF8NQT1GlkwWQLshjWmr8xORsvRHzzIqhUyDCvABc,4632
 torch_geometric/datasets/tu_dataset.py,sha256=14OSaXBgVwT1dX2h1wZ3xVIwoo0GQBEfR3yWh6Q0VF0,7847
 torch_geometric/datasets/twitch.py,sha256=qfEerf-Uaojx2ZvegENowdG4E7RoUT_HUO9xtULadvo,3658
@@ -374,7 +376,7 @@ torch_geometric/nn/conv/hgt_conv.py,sha256=lUhTWUMovMtn9yR_b2-kLNLqHChGOUl2OtXBY
 torch_geometric/nn/conv/hypergraph_conv.py,sha256=4BosbbqJyprlI6QjPqIfMxCqnARU_0mUn1zcAQhbw90,8691
 torch_geometric/nn/conv/le_conv.py,sha256=DonmmYZOKk5wIlTZzzIfNKqBY6MO0MRxYhyr0YtNz-Q,3494
 torch_geometric/nn/conv/lg_conv.py,sha256=8jMa79iPsOUbXEfBIc3wmbvAD8T3d1j37LeIFTX3Yag,2369
-torch_geometric/nn/conv/meshcnn_conv.py,sha256=Z6p9KwGc_Kj4XQnTWqzbXQzbbpVlMv7ga0DuDB0jLSg,22279
+torch_geometric/nn/conv/meshcnn_conv.py,sha256=92zUcgfS0Fwv-MpddF4Ia1a65y7ddPAkazYf7D6kvwg,21951
 torch_geometric/nn/conv/message_passing.py,sha256=ZuTvSvodGy1GyAW4mHtuoMUuxclam-7opidYNY5IHm8,44377
 torch_geometric/nn/conv/mf_conv.py,sha256=SkOGMN1tFT9dcqy8xYowsB2ozw6QfkoArgR1BksZZaU,4340
 torch_geometric/nn/conv/mixhop_conv.py,sha256=qVDPWeWcnO7_eHM0ZnpKtr8SISjb4jp0xjgpoDrwjlk,4555
@@ -429,7 +431,7 @@ 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=4mZ5dyiZ9aa1NaBth1qYV-hZdnG_Np1XWvRLB4Qv6RM,2338
+torch_geometric/nn/models/__init__.py,sha256=fbHQauZw9Snvl2PuN5cjZoAW8SwUl6E-p2IOmwUKB3A,2395
 torch_geometric/nn/models/attentive_fp.py,sha256=1z3iTV2O5W9tqHFAdno8FeBFeXmuG-TDZk4lwwVh3Ac,6634
 torch_geometric/nn/models/attract_repel.py,sha256=h9OyogT0NY0xiT0DkpJHMxH6ZUmo8R-CmwZdKEwq8Ek,5277
 torch_geometric/nn/models/autoencoder.py,sha256=nGje-zty78Y3hxOJ9o0_6QziJjOvBlknk6z0_fDQwQU,10770
@@ -442,7 +444,7 @@ torch_geometric/nn/models/dimenet.py,sha256=O2rqEx5HWs_lMwRD8eq6WMkbqJaCLL5zgWUJ
 torch_geometric/nn/models/dimenet_utils.py,sha256=Eyn_EiJqwKvuYj6BtRpSxrzMG3v4Gk98X9MxZ7uvwm4,5069
 torch_geometric/nn/models/g_retriever.py,sha256=tVibbqM_r-1LnA3R3oVyzp0bpuN3qPoYqcU6LZ8dYEk,8260
 torch_geometric/nn/models/git_mol.py,sha256=Wc6Hx6RDDR7sDWRWHfA5eK9e9gFsrTZ9OLmpMfoj3pE,12676
-torch_geometric/nn/models/glem.py,sha256=PlLjfMM4lKLs7c7tRC4LVD8tj0jpUXNxcnGbYut7vBE,16624
+torch_geometric/nn/models/glem.py,sha256=GlL_I63g-_5eTycSGRj720YntldQ-CQ351RaDPc6XAU,16674
 torch_geometric/nn/models/gnnff.py,sha256=15dkiLgy0LmH1hnUrpeoHioIp4BPTfjpVATpnGRt9E0,7860
 torch_geometric/nn/models/gpse.py,sha256=acEAeeicLgzKRL54WhvIFxjA5XViHgXgMEH-NgbMdqI,41971
 torch_geometric/nn/models/graph_mixer.py,sha256=mthMeCOikR8gseEsu4oJ3Cd9C35zHSv1p32ROwnG-6s,9246
@@ -459,6 +461,7 @@ torch_geometric/nn/models/molecule_gpt.py,sha256=k-XULH6jaurj-R2EE4sIWTkqlNqa3Cz
 torch_geometric/nn/models/neural_fingerprint.py,sha256=pTLJgU9Uh2Lnf9bggLj4cKI8YdEFcMF-9MALuubqbuQ,2378
 torch_geometric/nn/models/node2vec.py,sha256=81Ku4Rp4IwLEAy06KEgJ2fYtXXVL_uv_Hb8lBr6YXrE,7664
 torch_geometric/nn/models/pmlp.py,sha256=dcAASVSyQMMhItSfEJWPeAFh0R3tNCwAHwdrShwQ8o4,3538
+torch_geometric/nn/models/protein_mpnn.py,sha256=QXHfltiJPmakpzgJKw_1vwCGBlszv9nfY4r4F38Sg9k,11031
 torch_geometric/nn/models/re_net.py,sha256=pz66q5b5BoGDNVQvpEGS2RGoeKvpjkYAv9r3WAuvITk,8986
 torch_geometric/nn/models/rect.py,sha256=2F3XyyvHTAEuqfJpiNB5M8pSGy738LhPiom5I-SDWqM,2808
 torch_geometric/nn/models/rev_gnn.py,sha256=Bpme087Zs227lcB0ODOKWsxaly67q96wseaRt6bacjs,11796
@@ -613,7 +616,7 @@ torch_geometric/utils/_tree_decomposition.py,sha256=ZtpjPQJgXbQWtSWjo-Fmhrov0DGO
 torch_geometric/utils/_trim_to_layer.py,sha256=cauOEzMJJK4w9BC-Pg1bHVncBYqG9XxQex3rn10BFjc,8339
 torch_geometric/utils/_unbatch.py,sha256=B0vjKI96PtHvSBG8F_lqvsiJE134aVjUurPZsG6UZRI,2378
 torch_geometric/utils/augmentation.py,sha256=1F0YCuaklZ9ZbXxdFV0oOoemWvLd8p60WvFo2chzl7E,8600
-torch_geometric/utils/convert.py,sha256=j0t_87c-U_-15YKFfkOZfloEc5NbjgeLIk851zHG8WA,21665
+torch_geometric/utils/convert.py,sha256=RE5n5no74Xu39-QMWFE0-1RvTgykdK33ymyjF9WcuSs,21938
 torch_geometric/utils/cross_entropy.py,sha256=ZFS5bivtzv3EV9zqgKsekmuQyoZZggPSclhl_tRNHxo,3047
 torch_geometric/utils/dropout.py,sha256=gg0rDnD4FLvBaKSoLAkZwViAQflhLefJm6_Mju5dmQs,11416
 torch_geometric/utils/embedding.py,sha256=Ac_MPSrZGpw-e-gU6Yz-seUioC2WZxBSSzXFeclGwMk,5232
@@ -634,13 +637,13 @@ torch_geometric/utils/num_nodes.py,sha256=F15ciTFOe8AxjkUh1wKH7RLmJvQYYpz-l3pPPv
 torch_geometric/utils/ppr.py,sha256=ebiHbQqRJsQbGUI5xu-IkzQSQsgIaC71vgO0KcXIKAk,4055
 torch_geometric/utils/random.py,sha256=Rv5HlhG5310rytbT9EZ7xWLGKQfozfz1azvYi5nx2-U,5148
 torch_geometric/utils/repeat.py,sha256=RxCoRoEisaP6NouXPPW5tY1Rn-tIfrmpJPm0qGP6W8M,815
-torch_geometric/utils/smiles.py,sha256=lGQ2BwJ49uBrQfIxxPz8ceTO9Jo-XCjlLxs1ql3xrsA,7130
+torch_geometric/utils/smiles.py,sha256=CFqeNtSBXQtY9Ex2gQzI0La490IpVVrm01QdRYEpV7w,7114
 torch_geometric/utils/sparse.py,sha256=1DbaEwdyvnzvg5qVjPlnWcEVDMkxrQLX1jJ0dr6P4js,25135
 torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5nUAUjw,6222
 torch_geometric/visualization/__init__.py,sha256=b-HnVesXjyJ_L1N-DnjiRiRVf7lhwKaBQF_2i5YMVSU,208
 torch_geometric/visualization/graph.py,sha256=mfZHXYfiU-CWMtfawYc80IxVwVmtK9hbIkSKhM_j7oI,14311
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250702.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250702.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250702.dist-info/METADATA,sha256=66AyTfnfJvD0er8ePN_vOUgj6tD76JJy4QPaIvkh8bw,63005
-pyg_nightly-2.7.0.dev20250702.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250703.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250703.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250703.dist-info/METADATA,sha256=lm-xVNswGkfSaczkhV_ANegM0oGwBEfhCv75gwj0X5Q,63005
+pyg_nightly-2.7.0.dev20250703.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.dev20250702'
+__version__ = '2.7.0.dev20250703'
 
 __all__ = [
     'Index',

torch_geometric/datasets/__init__.py

@@ -81,8 +81,10 @@ from .web_qsp_dataset import WebQSPDataset, CWQDataset
 from .git_mol_dataset import GitMolDataset
 from .molecule_gpt_dataset import MoleculeGPTDataset
 from .instruct_mol_dataset import InstructMolDataset
+from .protein_mpnn_dataset import ProteinMPNNDataset
 from .tag_dataset import TAGDataset
 from .city import CityNetwork
+from .teeth3ds import Teeth3DS
 
 from .dbp15k import DBP15K
 from .aminer import AMiner
@@ -201,8 +203,10 @@ homo_datasets = [
     'GitMolDataset',
     'MoleculeGPTDataset',
     'InstructMolDataset',
+    'ProteinMPNNDataset',
     'TAGDataset',
     'CityNetwork',
+    'Teeth3DS',
 ]
 
 hetero_datasets = [

torch_geometric/datasets/git_mol_dataset.py

@@ -102,7 +102,7 @@ class GitMolDataset(InMemoryDataset):
 
         try:
             from rdkit import Chem, RDLogger
-            RDLogger.DisableLog('rdApp.*')  # type: ignore
+            RDLogger.DisableLog('rdApp.*')
             WITH_RDKIT = True
 
         except ImportError:

torch_geometric/datasets/protein_mpnn_dataset.py

@@ -0,0 +1,451 @@
+import os
+import pickle
+import random
+from collections import defaultdict
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_tar,
+)
+
+
+class ProteinMPNNDataset(InMemoryDataset):
+    r"""The ProteinMPNN dataset from the `"Robust deep learning based protein
+    sequence design using ProteinMPNN"
+    <https://www.biorxiv.org/content/10.1101/2022.06.03.494563v1>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        size (str): Size of the PDB information to train the model.
+            If :obj:`"small"`, loads the small dataset (229.4 MB).
+            If :obj:`"large"`, loads the large dataset (64.1 GB).
+            (default: :obj:`"small"`)
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"valid"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset.
+            (default: :obj:`"train"`)
+        datacut (str, optional): Date cutoff to filter the dataset.
+            (default: :obj:`"2030-01-01"`)
+        rescut (float, optional): PDB resolution cutoff.
+            (default: :obj:`3.5`)
+        homo (float, optional): Homology cutoff.
+            (default: :obj:`0.70`)
+        max_length (int, optional): Maximum length of the protein complex.
+            (default: :obj:`10000`)
+        num_units (int, optional): Number of units of the protein complex.
+            (default: :obj:`150`)
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        pre_filter (callable, optional): A function that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a boolean
+            value, indicating whether the data object should be included in the
+            final dataset. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+
+    raw_url = {
+        'small':
+        'https://files.ipd.uw.edu/pub/training_sets/'
+        'pdb_2021aug02_sample.tar.gz',
+        'large':
+        'https://files.ipd.uw.edu/pub/training_sets/'
+        'pdb_2021aug02.tar.gz',
+    }
+
+    splits = {
+        'train': 1,
+        'valid': 2,
+        'test': 3,
+    }
+
+    def __init__(
+        self,
+        root: str,
+        size: str = 'small',
+        split: str = 'train',
+        datacut: str = '2030-01-01',
+        rescut: float = 3.5,
+        homo: float = 0.70,
+        max_length: int = 10000,
+        num_units: int = 150,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        pre_filter: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+        self.size = size
+        self.split = split
+        self.datacut = datacut
+        self.rescut = rescut
+        self.homo = homo
+        self.max_length = max_length
+        self.num_units = num_units
+
+        self.sub_folder = self.raw_url[self.size].split('/')[-1].split('.')[0]
+
+        super().__init__(root, transform, pre_transform, pre_filter,
+                         force_reload=force_reload)
+        self.load(self.processed_paths[self.splits[self.split]])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [
+            f'{self.sub_folder}/{f}'
+            for f in ['list.csv', 'valid_clusters.txt', 'test_clusters.txt']
+        ]
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        return ['splits.pkl', 'train.pt', 'valid.pt', 'test.pt']
+
+    def download(self) -> None:
+        file_path = download_url(self.raw_url[self.size], self.raw_dir)
+        extract_tar(file_path, self.raw_dir)
+        os.unlink(file_path)
+
+    def process(self) -> None:
+        alphabet_set = set(list('ACDEFGHIKLMNPQRSTVWYX'))
+        cluster_ids = self._process_split()
+        total_items = sum(len(items) for items in cluster_ids.values())
+        data_list = []
+
+        with tqdm(total=total_items, desc="Processing") as pbar:
+            for _, items in cluster_ids.items():
+                for chain_id, _ in items:
+                    item = self._process_pdb1(chain_id)
+
+                    if 'label' not in item:
+                        pbar.update(1)
+                        continue
+                    if len(list(np.unique(item['idx']))) >= 352:
+                        pbar.update(1)
+                        continue
+
+                    my_dict = self._process_pdb2(item)
+
+                    if len(my_dict['seq']) > self.max_length:
+                        pbar.update(1)
+                        continue
+                    bad_chars = set(list(
+                        my_dict['seq'])).difference(alphabet_set)
+                    if len(bad_chars) > 0:
+                        pbar.update(1)
+                        continue
+
+                    x_chain_all, chain_seq_label_all, mask, chain_mask_all, residue_idx, chain_encoding_all = self._process_pdb3(  # noqa: E501
+                        my_dict)
+
+                    data = Data(
+                        x=x_chain_all,  # [seq_len, 4, 3]
+                        chain_seq_label=chain_seq_label_all,  # [seq_len]
+                        mask=mask,  # [seq_len]
+                        chain_mask_all=chain_mask_all,  # [seq_len]
+                        residue_idx=residue_idx,  # [seq_len]
+                        chain_encoding_all=chain_encoding_all,  # [seq_len]
+                    )
+
+                    if self.pre_filter is not None and not self.pre_filter(
+                            data):
+                        continue
+                    if self.pre_transform is not None:
+                        data = self.pre_transform(data)
+
+                    data_list.append(data)
+
+                    if len(data_list) >= self.num_units:
+                        pbar.update(total_items - pbar.n)
+                        break
+                    pbar.update(1)
+                else:
+                    continue
+                break
+            self.save(data_list, self.processed_paths[self.splits[self.split]])
+
+    def _process_split(self) -> Dict[int, List[Tuple[str, int]]]:
+        import pandas as pd
+        save_path = self.processed_paths[0]
+
+        if os.path.exists(save_path):
+            print('Load split')
+            with open(save_path, 'rb') as f:
+                data = pickle.load(f)
+        else:
+            # CHAINID, DEPOSITION, RESOLUTION, HASH, CLUSTER, SEQUENCE
+            df = pd.read_csv(self.raw_paths[0])
+            df = df[(df['RESOLUTION'] <= self.rescut)
+                    & (df['DEPOSITION'] <= self.datacut)]
+
+            val_ids = pd.read_csv(self.raw_paths[1], header=None)[0].tolist()
+            test_ids = pd.read_csv(self.raw_paths[2], header=None)[0].tolist()
+
+            # compile training and validation sets
+            data = {
+                'train': defaultdict(list),
+                'valid': defaultdict(list),
+                'test': defaultdict(list),
+            }
+
+            for _, r in tqdm(df.iterrows(), desc='Processing split',
+                             total=len(df)):
+                cluster_id = r['CLUSTER']
+                hash_id = r['HASH']
+                chain_id = r['CHAINID']
+                if cluster_id in val_ids:
+                    data['valid'][cluster_id].append((chain_id, hash_id))
+                elif cluster_id in test_ids:
+                    data['test'][cluster_id].append((chain_id, hash_id))
+                else:
+                    data['train'][cluster_id].append((chain_id, hash_id))
+
+            with open(save_path, 'wb') as f:
+                pickle.dump(data, f)
+
+        return data[self.split]
+
+    def _process_pdb1(self, chain_id: str) -> Dict[str, Any]:
+        pdbid, chid = chain_id.split('_')
+        prefix = f'{self.raw_dir}/{self.sub_folder}/pdb/{pdbid[1:3]}/{pdbid}'
+        # load metadata
+        if not os.path.isfile(f'{prefix}.pt'):
+            return {'seq': np.zeros(5)}
+        meta = torch.load(f'{prefix}.pt')
+        asmb_ids = meta['asmb_ids']
+        asmb_chains = meta['asmb_chains']
+        chids = np.array(meta['chains'])
+
+        # find candidate assemblies which contain chid chain
+        asmb_candidates = {
+            a
+            for a, b in zip(asmb_ids, asmb_chains) if chid in b.split(',')
+        }
+
+        # if the chains is missing is missing from all the assemblies
+        # then return this chain alone
+        if len(asmb_candidates) < 1:
+            chain = torch.load(f'{prefix}_{chid}.pt')
+            L = len(chain['seq'])
+            return {
+                'seq': chain['seq'],
+                'xyz': chain['xyz'],
+                'idx': torch.zeros(L).int(),
+                'masked': torch.Tensor([0]).int(),
+                'label': chain_id,
+            }
+
+        # randomly pick one assembly from candidates
+        asmb_i = random.sample(list(asmb_candidates), 1)
+
+        # indices of selected transforms
+        idx = np.where(np.array(asmb_ids) == asmb_i)[0]
+
+        # load relevant chains
+        chains = {
+            c: torch.load(f'{prefix}_{c}.pt')
+            for i in idx
+            for c in asmb_chains[i] if c in meta['chains']
+        }
+
+        # generate assembly
+        asmb = {}
+        for k in idx:
+
+            # pick k-th xform
+            xform = meta[f'asmb_xform{k}']
+            u = xform[:, :3, :3]
+            r = xform[:, :3, 3]
+
+            # select chains which k-th xform should be applied to
+            s1 = set(meta['chains'])
+            s2 = set(asmb_chains[k].split(','))
+            chains_k = s1 & s2
+
+            # transform selected chains
+            for c in chains_k:
+                try:
+                    xyz = chains[c]['xyz']
+                    xyz_ru = torch.einsum('bij,raj->brai', u, xyz) + r[:, None,
+                                                                       None, :]
+                    asmb.update({
+                        (c, k, i): xyz_i
+                        for i, xyz_i in enumerate(xyz_ru)
+                    })
+                except KeyError:
+                    return {'seq': np.zeros(5)}
+
+        # select chains which share considerable similarity to chid
+        seqid = meta['tm'][chids == chid][0, :, 1]
+        homo = {
+            ch_j
+            for seqid_j, ch_j in zip(seqid, chids) if seqid_j > self.homo
+        }
+        # stack all chains in the assembly together
+        seq: str = ''
+        xyz_all: List[torch.Tensor] = []
+        idx_all: List[torch.Tensor] = []
+        masked: List[int] = []
+        seq_list = []
+        for counter, (k, v) in enumerate(asmb.items()):
+            seq += chains[k[0]]['seq']
+            seq_list.append(chains[k[0]]['seq'])
+            xyz_all.append(v)
+            idx_all.append(torch.full((v.shape[0], ), counter))
+            if k[0] in homo:
+                masked.append(counter)
+
+        return {
+            'seq': seq,
+            'xyz': torch.cat(xyz_all, dim=0),
+            'idx': torch.cat(idx_all, dim=0),
+            'masked': torch.Tensor(masked).int(),
+            'label': chain_id,
+        }
+
+    def _process_pdb2(self, t: Dict[str, Any]) -> Dict[str, Any]:
+        init_alphabet = list(
+            'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz')
+        extra_alphabet = [str(item) for item in list(np.arange(300))]
+        chain_alphabet = init_alphabet + extra_alphabet
+        my_dict: Dict[str, Union[str, int, Dict[str, Any], List[Any]]] = {}
+        concat_seq = ''
+        mask_list = []
+        visible_list = []
+        for idx in list(np.unique(t['idx'])):
+            letter = chain_alphabet[idx]
+            res = np.argwhere(t['idx'] == idx)
+            initial_sequence = "".join(list(
+                np.array(list(t['seq']))[res][
+                    0,
+                ]))
+            if initial_sequence[-6:] == "HHHHHH":
+                res = res[:, :-6]
+            if initial_sequence[0:6] == "HHHHHH":
+                res = res[:, 6:]
+            if initial_sequence[-7:-1] == "HHHHHH":
+                res = res[:, :-7]
+            if initial_sequence[-8:-2] == "HHHHHH":
+                res = res[:, :-8]
+            if initial_sequence[-9:-3] == "HHHHHH":
+                res = res[:, :-9]
+            if initial_sequence[-10:-4] == "HHHHHH":
+                res = res[:, :-10]
+            if initial_sequence[1:7] == "HHHHHH":
+                res = res[:, 7:]
+            if initial_sequence[2:8] == "HHHHHH":
+                res = res[:, 8:]
+            if initial_sequence[3:9] == "HHHHHH":
+                res = res[:, 9:]
+            if initial_sequence[4:10] == "HHHHHH":
+                res = res[:, 10:]
+            if res.shape[1] >= 4:
+                chain_seq = "".join(list(np.array(list(t['seq']))[res][0]))
+                my_dict[f'seq_chain_{letter}'] = chain_seq
+                concat_seq += chain_seq
+                if idx in t['masked']:
+                    mask_list.append(letter)
+                else:
+                    visible_list.append(letter)
+                coords_dict_chain = {}
+                all_atoms = np.array(t['xyz'][res])[0]  # [L, 14, 3]
+                for i, c in enumerate(['N', 'CA', 'C', 'O']):
+                    coords_dict_chain[
+                        f'{c}_chain_{letter}'] = all_atoms[:, i, :].tolist()
+                my_dict[f'coords_chain_{letter}'] = coords_dict_chain
+        my_dict['name'] = t['label']
+        my_dict['masked_list'] = mask_list
+        my_dict['visible_list'] = visible_list
+        my_dict['num_of_chains'] = len(mask_list) + len(visible_list)
+        my_dict['seq'] = concat_seq
+        return my_dict
+
+    def _process_pdb3(
+        self, b: Dict[str, Any]
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,
+               torch.Tensor, torch.Tensor]:
+        L = len(b['seq'])
+        # residue idx with jumps across chains
+        residue_idx = -100 * np.ones([L], dtype=np.int32)
+        # get the list of masked / visible chains
+        masked_chains, visible_chains = b['masked_list'], b['visible_list']
+        visible_temp_dict, masked_temp_dict = {}, {}
+        for letter in masked_chains + visible_chains:
+            chain_seq = b[f'seq_chain_{letter}']
+            if letter in visible_chains:
+                visible_temp_dict[letter] = chain_seq
+            elif letter in masked_chains:
+                masked_temp_dict[letter] = chain_seq
+        # check for duplicate chains (same sequence but different identity)
+        for _, vm in masked_temp_dict.items():
+            for kv, vv in visible_temp_dict.items():
+                if vm == vv:
+                    if kv not in masked_chains:
+                        masked_chains.append(kv)
+                    if kv in visible_chains:
+                        visible_chains.remove(kv)
+        # build protein data structures
+        all_chains = masked_chains + visible_chains
+        np.random.shuffle(all_chains)
+        x_chain_list = []
+        chain_mask_list = []
+        chain_seq_list = []
+        chain_encoding_list = []
+        c, l0, l1 = 1, 0, 0
+        for letter in all_chains:
+            chain_seq = b[f'seq_chain_{letter}']
+            chain_length = len(chain_seq)
+            chain_coords = b[f'coords_chain_{letter}']
+            x_chain = np.stack([
+                chain_coords[c] for c in [
+                    f'N_chain_{letter}', f'CA_chain_{letter}',
+                    f'C_chain_{letter}', f'O_chain_{letter}'
+                ]
+            ], 1)  # [chain_length, 4, 3]
+            x_chain_list.append(x_chain)
+            chain_seq_list.append(chain_seq)
+            if letter in visible_chains:
+                chain_mask = np.zeros(chain_length)  # 0 for visible chains
+            elif letter in masked_chains:
+                chain_mask = np.ones(chain_length)  # 1 for masked chains
+            chain_mask_list.append(chain_mask)
+            chain_encoding_list.append(c * np.ones(chain_length))
+            l1 += chain_length
+            residue_idx[l0:l1] = 100 * (c - 1) + np.arange(l0, l1)
+            l0 += chain_length
+            c += 1
+        x_chain_all = np.concatenate(x_chain_list, 0)  # [L, 4, 3]
+        chain_seq_all = "".join(chain_seq_list)
+        # [L,] 1.0 for places that need to be predicted
+        chain_mask_all = np.concatenate(chain_mask_list, 0)
+        chain_encoding_all = np.concatenate(chain_encoding_list, 0)
+
+        # Convert to labels
+        alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
+        chain_seq_label_all = np.asarray(
+            [alphabet.index(a) for a in chain_seq_all], dtype=np.int32)
+
+        isnan = np.isnan(x_chain_all)
+        mask = np.isfinite(np.sum(x_chain_all, (1, 2))).astype(np.float32)
+        x_chain_all[isnan] = 0.
+
+        # Conversion
+        return (
+            torch.from_numpy(x_chain_all).to(dtype=torch.float32),
+            torch.from_numpy(chain_seq_label_all).to(dtype=torch.long),
+            torch.from_numpy(mask).to(dtype=torch.float32),
+            torch.from_numpy(chain_mask_all).to(dtype=torch.float32),
+            torch.from_numpy(residue_idx).to(dtype=torch.long),
+            torch.from_numpy(chain_encoding_all).to(dtype=torch.long),
+        )

torch_geometric/datasets/qm9.py

@@ -202,7 +202,7 @@ class QM9(InMemoryDataset):
             from rdkit import Chem, RDLogger
             from rdkit.Chem.rdchem import BondType as BT
             from rdkit.Chem.rdchem import HybridizationType
-            RDLogger.DisableLog('rdApp.*')  # type: ignore
+            RDLogger.DisableLog('rdApp.*')
             WITH_RDKIT = True
 
         except ImportError:

torch_geometric/datasets/teeth3ds.py

@@ -0,0 +1,269 @@
+import json
+import os
+import os.path as osp
+from glob import glob
+from typing import Callable, Dict, List, Optional
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from torch_geometric.data import (
+    Data,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+
+
+class Teeth3DS(InMemoryDataset):
+    r"""The Teeth3DS+ dataset from the `"An Extended Benchmark for Intra-oral
+    3D Scans Analysis" <https://crns-smartvision.github.io/teeth3ds/>`_ paper.
+
+    This dataset is the first comprehensive public benchmark designed to
+    advance the field of intra-oral 3D scan analysis developed as part of the
+    3DTeethSeg 2022 and 3DTeethLand 2024 MICCAI challenges, aiming to drive
+    research in teeth identification, segmentation, labeling, 3D modeling,
+    and dental landmark identification.
+    The dataset includes at least 1,800 intra-oral scans (containing 23,999
+    annotated teeth) collected from 900 patients, covering both upper and lower
+    jaws separately.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str): The split name (one of :obj:`"Teeth3DS"`,
+            :obj:`"3DTeethSeg22_challenge"` or :obj:`"3DTeethLand_challenge"`).
+        train (bool, optional): If :obj:`True`, loads the training dataset,
+            otherwise the test dataset. (default: :obj:`True`)
+        num_samples (int, optional): Number of points to sample from each mesh.
+            (default: :obj:`30000`)
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.Data` object and returns a transformed
+            version. The data object will be transformed before every access.
+            (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.Data` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+    """
+    urls = {
+        'data_part_1.zip':
+        'https://osf.io/download/qhprs/',
+        'data_part_2.zip':
+        'https://osf.io/download/4pwnr/',
+        'data_part_3.zip':
+        'https://osf.io/download/frwdp/',
+        'data_part_4.zip':
+        'https://osf.io/download/2arn4/',
+        'data_part_5.zip':
+        'https://osf.io/download/xrz5f/',
+        'data_part_6.zip':
+        'https://osf.io/download/23hgq/',
+        'data_part_7.zip':
+        'https://osf.io/download/u83ad/',
+        'train_test_split':
+        'https://files.de-1.osf.io/v1/'
+        'resources/xctdy/providers/osfstorage/?zip='
+    }
+
+    sample_url = {
+        'teeth3ds_sample': 'https://osf.io/download/vr38s/',
+    }
+
+    landmarks_urls = {
+        '3DTeethLand_landmarks_train.zip': 'https://osf.io/download/k5hbj/',
+        '3DTeethLand_landmarks_test.zip': 'https://osf.io/download/sqw5e/',
+    }
+
+    def __init__(
+        self,
+        root: str,
+        split:
+        str = 'Teeth3DS',  # [3DTeethSeg22_challenge, 3DTeethLand_challenge]
+        train: bool = True,
+        num_samples: int = 30000,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        force_reload: bool = False,
+    ) -> None:
+
+        self.mode = 'training' if train else 'testing'
+        self.split = split
+        self.num_samples = num_samples
+
+        super().__init__(root, transform, pre_transform,
+                         force_reload=force_reload)
+
+    @property
+    def processed_dir(self) -> str:
+        return os.path.join(self.root, f'processed_{self.split}_{self.mode}')
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['license.txt']
+
+    @property
+    def processed_file_names(self) -> List[str]:
+        # Directory containing train/test split files:
+        split_subdir = 'teeth3ds_sample' if self.split == 'sample' else ''
+        split_dir = osp.join(
+            self.raw_dir,
+            split_subdir,
+            f'{self.split}_train_test_split',
+        )
+
+        split_files = glob(osp.join(split_dir, f'{self.mode}*.txt'))
+
+        # Collect all file names from the split files:
+        combined_list = []
+        for file_path in split_files:
+            with open(file_path) as file:
+                combined_list.extend(file.read().splitlines())
+
+        # Generate the list of processed file paths:
+        return [f'{file_name}.pt' for file_name in combined_list]
+
+    def download(self) -> None:
+        if self.split == 'sample':
+            for key, url in self.sample_url.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)
+                os.unlink(path)
+        else:
+            for key, url in self.urls.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)
+                os.unlink(path)
+            for key, url in self.landmarks_urls.items():
+                path = download_url(url, self.root, filename=key)
+                extract_zip(path, self.raw_dir)  # Extract each downloaded part
+                os.unlink(path)
+
+    def process_file(self, file_path: str) -> Optional[Data]:
+        """Processes the input file path to load mesh data, annotations,
+        and prepare the input features for a graph-based deep learning model.
+        """
+        import trimesh
+        from fpsample import bucket_fps_kdline_sampling
+
+        mesh = trimesh.load_mesh(file_path)
+
+        if isinstance(mesh, list):
+            # Handle the case where a list of Geometry objects is returned
+            mesh = mesh[0]
+
+        vertices = mesh.vertices
+        vertex_normals = mesh.vertex_normals
+
+        # Perform sampling on mesh vertices:
+        if len(vertices) < self.num_samples:
+            sampled_indices = np.random.choice(
+                len(vertices),
+                self.num_samples,
+                replace=True,
+            )
+        else:
+            sampled_indices = bucket_fps_kdline_sampling(
+                vertices,
+                self.num_samples,
+                h=5,
+                start_idx=0,
+            )
+
+        if len(sampled_indices) != self.num_samples:
+            raise RuntimeError(f"Sampled points mismatch, expected "
+                               f"{self.num_samples} points, but got "
+                               f"{len(sampled_indices)} for '{file_path}'")
+
+        # Extract features and annotations for the sampled points:
+        pos = torch.tensor(vertices[sampled_indices], dtype=torch.float)
+        x = torch.tensor(vertex_normals[sampled_indices], dtype=torch.float)
+
+        # Load segmentation annotations:
+        seg_annotation_path = file_path.replace('.obj', '.json')
+        if osp.exists(seg_annotation_path):
+            with open(seg_annotation_path) as f:
+                seg_annotations = json.load(f)
+            y = torch.tensor(
+                np.asarray(seg_annotations['labels'])[sampled_indices],
+                dtype=torch.float)
+            instances = torch.tensor(
+                np.asarray(seg_annotations['instances'])[sampled_indices],
+                dtype=torch.float)
+        else:
+            y = torch.empty(0, 3)
+            instances = torch.empty(0, 3)
+
+        # Load landmarks annotations:
+        landmarks_annotation_path = file_path.replace('.obj', '__kpt.json')
+
+        # Parse keypoint annotations into structured tensors:
+        keypoints_dict: Dict[str, List] = {
+            key: []
+            for key in [
+                'Mesial', 'Distal', 'Cusp', 'InnerPoint', 'OuterPoint',
+                'FacialPoint'
+            ]
+        }
+        keypoint_tensors: Dict[str, torch.Tensor] = {
+            key: torch.empty(0, 3)
+            for key in [
+                'Mesial', 'Distal', 'Cusp', 'InnerPoint', 'OuterPoint',
+                'FacialPoint'
+            ]
+        }
+        if osp.exists(landmarks_annotation_path):
+            with open(landmarks_annotation_path) as f:
+                landmarks_annotations = json.load(f)
+
+            for keypoint in landmarks_annotations['objects']:
+                keypoints_dict[keypoint['class']].extend(keypoint['coord'])
+
+            keypoint_tensors = {
+                k: torch.tensor(np.asarray(v),
+                                dtype=torch.float).reshape(-1, 3)
+                for k, v in keypoints_dict.items()
+            }
+
+        data = Data(
+            pos=pos,
+            x=x,
+            y=y,
+            instances=instances,
+            jaw=file_path.split('.obj')[0].split('_')[1],
+            mesial=keypoint_tensors['Mesial'],
+            distal=keypoint_tensors['Distal'],
+            cusp=keypoint_tensors['Cusp'],
+            inner_point=keypoint_tensors['InnerPoint'],
+            outer_point=keypoint_tensors['OuterPoint'],
+            facial_point=keypoint_tensors['FacialPoint'],
+        )
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        return data
+
+    def process(self) -> None:
+        for file in tqdm(self.processed_file_names):
+            name = file.split('.')[0]
+            path = osp.join(self.raw_dir, '**', '*', name + '.obj')
+            paths = glob(path)
+            if len(paths) == 1:
+                data = self.process_file(paths[0])
+                torch.save(data, osp.join(self.processed_dir, file))
+
+    def len(self) -> int:
+        return len(self.processed_file_names)
+
+    def get(self, idx: int) -> Data:
+        return torch.load(
+            osp.join(self.processed_dir, self.processed_file_names[idx]),
+            weights_only=False,
+        )
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({len(self)}, '
+                f'mode={self.mode}, split={self.split})')

torch_geometric/nn/conv/meshcnn_conv.py

@@ -1,7 +1,7 @@
 # The below is to suppress the warning on torch.nn.conv.MeshCNNConv::update
 # pyright: reportIncompatibleMethodOverride=false
+import warnings
 from typing import Optional
-from warnings import warn
 
 import torch
 from torch.nn import Linear, Module, ModuleList
@@ -456,13 +456,10 @@ class MeshCNNConv(MessagePassing):
                 {type(network)}"
             if not hasattr(network, "in_channels") and \
                     not hasattr(network, "in_features"):
-                warn(
-                    f"kernel[{i}] does not have attribute \
-                            'in_channels' nor 'out_features'. The \
-                            network must take as input a \
-                            {self.in_channels}-dimensional tensor. \
-                            Still, assuming user configured \
-                            correctly. Continuing..", stacklevel=2)
+                warnings.warn(
+                    f"kernel[{i}] does not have attribute 'in_channels' nor "
+                    f"'out_features'. The network must take as input a "
+                    f"{self.in_channels}-dimensional tensor.", stacklevel=2)
             else:
                 input_dimension = getattr(network, "in_channels",
                                           network.in_features)
@@ -475,13 +472,10 @@ class MeshCNNConv(MessagePassing):
 
             if not hasattr(network, "out_channels") and \
                     not hasattr(network, "out_features"):
-                warn(
-                    f"kernel[{i}] does not have attribute \
-                                'in_channels' nor 'out_features'. The \
-                                network must take as input a \
-                                {self.in_channels}-dimensional tensor. \
-                                Still, assuming user configured \
-                                correctly. Continuing..", stacklevel=2)
+                warnings.warn(
+                    f"kernel[{i}] does not have attribute 'in_channels' nor "
+                    f"'out_features'. The network must take as input a "
+                    f"{self.in_channels}-dimensional tensor.", stacklevel=2)
             else:
                 output_dimension = getattr(network, "out_channels",
                                            network.out_features)

torch_geometric/nn/models/__init__.py

@@ -32,6 +32,7 @@ from .visnet import ViSNet
 from .g_retriever import GRetriever
 from .git_mol import GITMol
 from .molecule_gpt import MoleculeGPT
+from .protein_mpnn import ProteinMPNN
 from .glem import GLEM
 from .sgformer import SGFormer
 # Deprecated:
@@ -86,6 +87,7 @@ __all__ = classes = [
     'GRetriever',
     'GITMol',
     'MoleculeGPT',
+    'ProteinMPNN',
     'GLEM',
     'SGFormer',
     'ARLinkPredictor',

torch_geometric/nn/models/glem.py

@@ -8,6 +8,13 @@ from torch_geometric.loader import DataLoader, NeighborLoader
 from torch_geometric.nn.models import GraphSAGE, basic_gnn
 
 
+def deal_nan(x):
+    if isinstance(x, torch.Tensor):
+        x = x.clone()
+        x[torch.isnan(x)] = 0.0
+    return x
+
+
 class GLEM(torch.nn.Module):
     r"""This GNN+LM co-training model is based on GLEM from the `"Learning on
     Large-scale Text-attributed Graphs via Variational Inference"
@@ -379,9 +386,6 @@ class GLEM(torch.nn.Module):
             is_augmented: use EM or just train GNN and LM with gold data
 
         """
-        def deal_nan(x):
-            return 0 if torch.isnan(x) else x
-
         if is_augmented and (sum(~is_gold) > 0):
             mle_loss = deal_nan(loss_func(logits[is_gold], labels[is_gold]))
             # all other labels beside from ground truth(gold labels)

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/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)