[stability-pipeline] do proper cleanup of non-ligand pdbs (issue38)

software/rosetta_ddG_pipeline/helper.py

@@ -33,6 +33,17 @@
 from pdb_to_prism import rosetta_energy_to_prism
 
 
+def remove_hetatms(input, output):
+    with open(input) as fp, open(output, 'w') as fp2:
+        for line in fp:
+            if not line.startswith('HETATM'):
+                if line.startswith('ATOM'):
+                    if line[16] in [' ', 'A']:
+                        fp2.write(line)
+                else:
+                    fp2.write(line)
+    return output
+
 def create_symlinks(source_file, dist_folder, name=''):
     # WIP
     # import stat

software/rosetta_ddG_pipeline/structure_input.py

@@ -21,7 +21,7 @@
 import pdb_to_fasta_seq
 import rosetta_paths
 from AnalyseStruc import get_structure_parameters
-from helper import read_fasta
+from helper import read_fasta, remove_hetatms, create_copy
 
 class structure:
 
@@ -51,8 +51,11 @@ def clean_up_and_isolate(self, name='input',ligand=None):
         """This script is for cleaning the pdb file from unwanted entities"""
 
         #This cleans the protein and removes ligands
-        if  ligand == None:           
+        if  ligand == None:
             self.path_to_clean_pdb = rosetta_paths.path_to_clean_pdb
+            #remove all HETATM
+            self.tmp_prep_struc = create_copy(self.prep_struc, self.folder.prepare_cleaning, name='withHETATM.pdb')
+            self.prep_struc = remove_hetatms(self.tmp_prep_struc , self.prep_struc)
             #Runs shell script
             shell_command = f'python2 {self.path_to_clean_pdb} {self.prep_struc} {self.run_struc}  --keepzeroocc'
             self.logger.info('Running clean_pdb.py script')