Residues class creation

src/atom.py

@@ -18,3 +18,24 @@ class Atom:
         self.TEMP_FACT = TEMP_FACT
         self.ELEM_SYMBOL = ELEM_SYMBOL
         self.ATOM_CHARGE = ATOM_CHARGE
+
+class Residue:
+    def __init__(self, atomsList):
+        self.ATOMS = {}
+        for atom in atomsList:
+            self.ATOMS[atom.ATOM_NAME] = atom
+            if atom.ATOM_NAME == "N":
+                self.N = atom
+            if atom.ATOM_NAME == "CA":
+                self.CA = atom
+            if atom.ATOM_NAME == "C":
+                self.C = atom
+            if atom.ATOM_NAME == "O":
+                self.O = atom
+            if atom.ATOM_NAME == "CB":
+                self.CB = atom
+            if atom.ATOM_NAME == "CG1":
+                self.CG1 = atom
+            if atom.ATOM_NAME == "CG2":
+                self.CG2 = atom
+            

src/pdb.py

@@ -32,7 +32,8 @@ class PDBFile:
 
     def getAtoms(self):
         self.ATOMS = []
-        
+        self.RESIDUES = []
+        tempAtoms = []
         for line in self.rawLines:
             if line.startswith("ATOM" or "HETATM"):
                 atom = Atom(ATOM_ID = int(line[6:11].strip()),
@@ -50,7 +51,16 @@ class PDBFile:
                             ELEM_SYMBOL = line[76:78].strip(),
                             ATOM_CHARGE = line[78:80].strip())
                 self.ATOMS.append(atom)
-                      
+                # get the current indice of atom
+                i = self.ATOMS.index(atom)
+                # if this is a brand new residue
+                if(len(self.ATOMS)>1
+                   and atom.RES_SEQ_NB != self.ATOMS[i-1].RES_SEQ_NB):
+                    self.RESIDUES.append(Residue(tempAtoms))
+                    tempAtoms=[]
+                tempAtoms.append(atom)
+        # last residue
+        self.RESIDUES.append(Residue(tempAtoms))
        
     def __init__(self, filename):
         self.rawLines = self.getContent(filename)
@@ -66,3 +76,4 @@ if __name__ == "__main__":
               "call structure and parameters.")
     else:
         pdbFile = PDBFile(sys.argv[1])
+        print(pdbFile.RESIDUES[0].ATOMS["N"].COORD_X)