From 488dfa74a19ceab7ea5b580692be85534e38ddde Mon Sep 17 00:00:00 2001
From: Thomas Forest <thomas.forest@etu.univ-paris-diderot.fr>
Date: Fri, 13 Sep 2019 00:30:32 +0200
Subject: [PATCH] Pymol silent mode and hydrogen_add threshold

---
 src/pdb.py | 32 +++++++++++++++++++-------------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/src/pdb.py b/src/pdb.py
index 9d598df..2ee6b6a 100644
--- a/src/pdb.py
+++ b/src/pdb.py
@@ -1,7 +1,7 @@
 # custom imports
 from atom import *
 import sys
-import pymol
+import pymol 
 #import collections
 class PDBFile:
     def getContent(self, filename):
@@ -27,12 +27,15 @@ class PDBFile:
                 break     
         return(Metadata)
 
-    def getAtoms(self):
+    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(),
@@ -53,17 +56,23 @@ class PDBFile:
                 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):
-        return true
+        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", selection='name n')  
+        pymol.cmd.select("nitrogens",'name n')  
         pymol.cmd.h_add("nitrogens")
-        pymol.cmd.select('hydrogens',selection='name h')
         pymol.stored.pos = []
-        pymol.cmd.iterate_state(1, 'hydrogens', 'stored.pos.append([name,resi,x,y,z])')
+        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)
@@ -71,13 +80,10 @@ class PDBFile:
     def __init__(self, filename, output_pdb=None):
         self.rawLines = self.getContent(filename)
         self.Metadata = self.getHeader()
-        self.getAtoms()
-        for elem in self.Metadata :    
-            print(self.Metadata[elem], end="")
-        if(self.check_hydrogens(self.atoms)==False):
-            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)
+        self.getAtoms(filename)
+        # for elem in self.Metadata : 
+        #     print(self.Metadata[elem], end="")
+        
 
 if __name__ == "__main__":
     if(len(sys.argv)<2):