add the dataset class and data structure for proteins

Author: Oxer11

torchdrug/data/__init__.py

@@ -1,15 +1,19 @@
 from .dictionary import PerfectHash, Dictionary
 from .graph import Graph, PackedGraph, cat
 from .molecule import Molecule, PackedMolecule
-from .dataset import MoleculeDataset, ReactionDataset, NodeClassificationDataset, KnowledgeGraphDataset, \
-    SemiSupervised, semisupervised, key_split, scaffold_split, ordered_scaffold_split
+from .protein import Protein, PackedProtein
+from .dataset import MoleculeDataset, ReactionDataset, ProteinDataset, \
+    ProteinPairDataset, ProteinLigandDataset, \
+    NodeClassificationDataset, KnowledgeGraphDataset, SemiSupervised, \
+    semisupervised, key_split, scaffold_split, ordered_scaffold_split
 from .dataloader import DataLoader, graph_collate
 from . import constant
 from . import feature
 
 __all__ = [
-    "Graph", "PackedGraph", "Molecule", "PackedMolecule", "PerfectHash", "Dictionary",
+    "Graph", "PackedGraph", "Molecule", "PackedMolecule", "Protein", "PackedProtein", "PerfectHash", "Dictionary",
     "MoleculeDataset", "ReactionDataset", "NodeClassificationDataset", "KnowledgeGraphDataset", "SemiSupervised",
+    "ProteinDataset", "ProteinPairDataset", "ProteinLigandDataset",
     "semisupervised", "key_split", "scaffold_split", "ordered_scaffold_split",
     "DataLoader", "graph_collate", "feature", "constant",
 ]

torchdrug/data/constant.py

@@ -29,9 +29,24 @@
     "Pa", "U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", "Fm", "Md", "No", "Lr", "Rf", "Db", "Sg", "Bh", "Hs",
     "Mt", "Ds", "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts", "Og"]
 
+# ordered by molecular mass
+RESIDUE_NAME = ["Glycine", "Alanine", "Serine", "Proline", "Valine", "Threonine", "Cysteine", "Isoleucine",
+                "Leucine", "Asparagine", "Aspartic acid", "Glutamine", "Lysine", "Glutamic acid", "Methionine",
+                "Histidine", "Phenylalanine", "Arginine", "Tyrosine", "Tryptophan"]
+
+RESIDUE_INITIAL = ["G", "A", "S", "P", "V", "T", "C", "I", "L", "N", "D", "Q", "K", "E", "M", "H", "F", "R", "Y", "W"]
+
+RESIDUE_ATOM_NAME = ["C", "CA", "CB", "CD", "CD1", "CD2", "CE", "CE1", "CE2", "CE3", "CG", "CG1", "CG2", "CH2",
+                     "CZ", "CZ2", "CZ3", "N", "ND1", "ND2", "NE", "NE1", "NE2", "NH1", "NH2", "NZ", "O", "OD1",
+                     "OD2", "OE1", "OE2", "OG", "OG1", "OH", "OXT", "SD", "SG"]
+
 NUM_ATOM = len(ATOM_NAME)
+NUM_AMINO_ACID = len(RESIDUE_NAME)
 
 for i, name in enumerate(ATOM_NAME):
     if i == 0:
         continue
+    setattr(module, name.upper(), i)
+
+for i, name in enumerate(RESIDUE_NAME):
     setattr(module, name.upper(), i)