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

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

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

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=ap-t4q8f9aTE0oAW_K5390u2Mlk8-S76rdeUEgPzglo,2250
+torch_geometric/__init__.py,sha256=GOuL0XBOcsFqK-Q-c_STDpzZAG-vsctiDiU_Tg9W3t8,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
@@ -119,7 +119,7 @@ torch_geometric/datasets/medshapenet.py,sha256=eCBCXKpueweCwDSf_Q4_MwVA3IbJd04FS
 torch_geometric/datasets/mixhop_synthetic_dataset.py,sha256=4NNvTHUvvV6pcqQCyVDS5XhppXUeF2H9GTfFoc49eyU,3951
 torch_geometric/datasets/mnist_superpixels.py,sha256=o2ArbZ0_OE0u8VCaHmWwvngESlOFr9oM9dSEP_tjAS4,3340
 torch_geometric/datasets/modelnet.py,sha256=-qmLjlQiKVWmtHefAIIE97dQxEcaBfetMJnvgYZuwkg,5347
-torch_geometric/datasets/molecule_gpt_dataset.py,sha256=gVZv14PuZCanE4oxxHlqRNrvzGv6_KN318q5yFA3lS0,18797
+torch_geometric/datasets/molecule_gpt_dataset.py,sha256=TFBduE3_3xxTFSHL3tirV-OAlBjSi6iHPOHJGQ_-tug,18785
 torch_geometric/datasets/molecule_net.py,sha256=pMzaJzd-LbBncZ0VoC87HfA8d1F4NwCWTb5YKvLM890,7404
 torch_geometric/datasets/movie_lens.py,sha256=M4Bu0Xus8IkW8GYzjxPxSdPXNbcCCx9cu6cncxBvLx8,4033
 torch_geometric/datasets/movie_lens_100k.py,sha256=eTpBAteM3jqTEtiwLxmhVj4r8JvftvPx8Hvs-3ZIHlU,6057
@@ -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
@@ -333,8 +335,9 @@ torch_geometric/nn/aggr/set_transformer.py,sha256=FG7_JizpFX14M6VSCwLSjYXYdJ1ZiQ
 torch_geometric/nn/aggr/sort.py,sha256=bvOOWnFkNOBOZih4rqVZQsjfeDX3vmXo1bpPSFD846w,2507
 torch_geometric/nn/aggr/utils.py,sha256=SQvdc0g6p_E2j0prA14MW2ekjEDvV-g545N0Q85uc-o,8625
 torch_geometric/nn/aggr/variance_preserving.py,sha256=fu-U_aGYpVLpgSFvVg0ONMe6nqoyv8tZ6Y35qMYTf9w,1126
-torch_geometric/nn/attention/__init__.py,sha256=wLKTmlfP7qL9sZHy4cmDFHEtdwa-MEKE1dT51L1_w10,192
+torch_geometric/nn/attention/__init__.py,sha256=w-jDQFpVqARJKjttTgKkD9kkAqRJl4MpASCfiNYIfr0,263
 torch_geometric/nn/attention/performer.py,sha256=2PCDn4_-oNTao2-DkXIaoi18anP01OxRELF2pvp-jk8,7357
+torch_geometric/nn/attention/polynormer.py,sha256=uBxGs0nldp6oGlByqbxgEk23VeXLEd6B3myS5BOKDRs,3998
 torch_geometric/nn/attention/qformer.py,sha256=7J-pWm_vpumK38IC-iCBz4oqL-BEIofEIxJ0wfjWq9A,2338
 torch_geometric/nn/attention/sgformer.py,sha256=OBC5HQxbY289bPDtwN8UbPH46To2GRTeVN-najogD-o,3747
 torch_geometric/nn/conv/__init__.py,sha256=8CK-DFG2PEo2ZaFyg-IUlQH8ecQoDDi556uv3ugeQyc,3572
@@ -374,7 +377,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 +432,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=71Hqc-ZMfCKn9lelFYDjpHXapbEa0wqVAd2OXCb1y5o,2448
 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 +445,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 +462,8 @@ 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/polynormer.py,sha256=mayWdzdolT5PCt_Oo7UGG-JUripMHWB2lUWF1bh6goU,7640
+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 +618,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 +639,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.dev20250704.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250704.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250704.dist-info/METADATA,sha256=Nau44bIMI13OXEqYNOlI1hYfDz8FXpUoPydv-JxRW2Q,63005
+pyg_nightly-2.7.0.dev20250704.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.dev20250704'
 
 __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/molecule_gpt_dataset.py

@@ -438,7 +438,7 @@ class MoleculeGPTDataset(InMemoryDataset):
         for mol in tqdm(suppl):
             if mol.HasProp('PUBCHEM_COMPOUND_CID'):
                 CID = mol.GetProp("PUBCHEM_COMPOUND_CID")
-                CAN_SMILES = mol.GetProp("PUBCHEM_OPENEYE_CAN_SMILES")
+                CAN_SMILES = mol.GetProp("PUBCHEM_SMILES")
 
                 m: Chem.Mol = Chem.MolFromSmiles(CAN_SMILES)
                 if m is None:

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: