123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596 |
- # custom imports
- from atom import *
- import sys
- import pymol
- #import collections
- class PDBFile:
- def getContent(self, filename):
- with open(filename) as f:
- return(f.readlines())
-
- def getHeader(self):
- #Metadata = collections.namedtuple("Metadata", ["header", "compound", "source", "author"])
- Metadata = {}
- for line in self.rawLines:
- # no need to continue if meta are complete
- if(len(Metadata) <4):
- if(line[0:10] == "HEADER "):
- Metadata['header']=line
- elif(line[0:10] == "COMPND 2"):
- Metadata['compound']=line
- elif(line[0:10] == "SOURCE 2"):
- Metadata['source']=line
- elif(line[0:10] == "AUTHOR "):
- Metadata['author']=line
- else:
- # if meta are complete, stop parsing
- break
- return(Metadata)
-
- def getAtoms(self, filename):
- self.atoms = []
- self.residues = []
- temp_atoms = []
- count_h = 0
- for line in self.rawLines:
- if line.startswith("ATOM" or "HETATM"):
- if(line[76:78].strip()=="H"):
- count_h+=1
- atom = Atom(atom_id = int(line[6:11].strip()),
- atom_name = line[12:16].strip(),
- res_name = line[17:20].strip(),
- chain_id = line[21:22].strip(),
- res_seq_nb = int(line[22:26].strip()),
- coordinates = [float(line[30:38].strip()),
- float(line[38:46].strip()),
- float(line[46:54].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(temp_atoms))
- temp_atoms=[]
- temp_atoms.append(atom)
- # last residue
- self.residues.append(Residue(temp_atoms))
- # hydrogens should represent in average 50% of total atoms... We use 30% threshold...
- if(count_h/len(temp_atoms)<0.30):
- #if(output_pdb==None):
- print("Need to add hydrogens ! If you want the modified PDB file, please use the -o output.pdb argument")
- self.add_hydrogens(filename)
-
- def check_hydrogens(self, atoms):
- print("ENTER CHECK HYDROGEN")
- return True
-
- def add_hydrogens(self, filename, output_pdb=None):
- pymol.finish_launching(['pymol', '-qc'])
- pymol.cmd.load(filename)
- pymol.cmd.select("nitrogens",'name n')
- pymol.cmd.h_add("nitrogens")
- pymol.stored.pos = []
- pymol.cmd.iterate_state(1, "hydrogens", 'stored.pos.append([name,resi,x,y,z])')
- if(output_pdb!=None):
- pymol.cmd.save(output_file)
- return(pymol.stored.pos)
-
- def __init__(self, filename, output_pdb=None):
- self.rawLines = self.getContent(filename)
- self.Metadata = self.getHeader()
- self.getAtoms(filename)
- # for elem in self.Metadata :
- # print(self.Metadata[elem], end="")
-
-
- if __name__ == "__main__":
- if(len(sys.argv)<2):
- print("Not enough arguments! Run with --help to learn more about proper"
- "call structure and parameters.")
- else:
- pdbFile = PDBFile(sys.argv[1])
- print(pdbFile.residues[15].atoms["C"].coord_x)
- print(pdbFile.add_hydrogens(sys.argv[1])[3])
|