Merge pull request #119

Correct `MODEL`, `END` and `ENDMDL` in `pdb_tidy`

pdbtools/pdb_tidy.py

@@ -143,16 +143,23 @@ def make_TER(prev_line):
     fmt_TER = "TER   {:>5d}      {:3s} {:1s}{:>4s}{:1s}" + " " * 53 + "\n"
 
     records = ('ATOM', 'HETATM')
-    ignored = ('TER', 'END ', 'END\n', 'CONECT', 'MASTER')
+    ignored = ('TER', 'END', 'CONECT', 'MASTER', 'ENDMDL')
     # Iterate up to the first ATOM/HETATM line
     prev_line = None
+    num_models = 1
+    in_model = False
     for line in fhandle:
 
+        line = line.strip()  # We will pad/add \n later to make uniform
+
+        if line.startswith('MODEL'):
+            line = "MODEL " + "    " + str(num_models).rjust(4)
+            num_models += 1
+            in_model = True
+
         if line.startswith(ignored):  # to avoid matching END _and_ ENDMDL
             continue
 
-        line = line.strip()  # We will pad/add \n later to make uniform
-
         # Check line length
         line = "{:<80}\n".format(line)
 
@@ -167,11 +174,11 @@ def make_TER(prev_line):
     serial_offset = 0  # To offset after adding TER records
     for line in fhandle:
 
+        line = line.strip()
+
         if line.startswith(ignored):
             continue
 
-        line = line.strip()
-
         # Treat ATOM/HETATM differently
         #   - no TER in HETATM
         if line.startswith('ATOM'):
@@ -206,10 +213,17 @@ def make_TER(prev_line):
             if atom_section:
                 atom_section = False
                 yield make_TER(prev_line)
+                if in_model:
+                    yield "{:<80}\n".format("ENDMDL")
+                    in_model = False
 
             if line.startswith('MODEL'):
+                line = "MODEL " + "    " + str(num_models).rjust(4)
+                num_models += 1
+                in_model = True
                 serial_offset = 0
 
+
         if serial > 99999:
             emsg = 'ERROR!! Structure contains more than 99.999 atoms.\n'
             sys.stderr.write(emsg)
@@ -226,6 +240,9 @@ def make_TER(prev_line):
             # Add last TER statement
             atom_section = False
             yield make_TER(prev_line)
+            if in_model:
+                yield "{:<80}\n".format("ENDMDL")
+                in_model = False
 
     # Add END statement
     yield "{:<80}\n".format("END")

tests/data/ensemble_error_MODEL.pdb

@@ -0,0 +1,14 @@
+HEADER    THIS A DUMMY PDB FOR PDB-TOOLS TESTING                                
+TITLE     A RANDOM PDB                                                          
+REMARK MODEL 1 HAS A BAD FORMATTED LINE
+REMARK MODEL 2 HAS NO NUMBER
+REMARK PDB_TIDY SHOULD ADD AND CORRECT THE MODEL NUMBERS
+MODEL 1                                                                  
+ATOM      1  N   ASN A   1      22.066  40.557   0.420  1.00  0.00           N  
+ATOM      2  H   ASN A   1      21.629  41.305  -0.098  1.00  0.00           H  
+ENDMDL                                                                          
+MODEL                                                                  
+ATOM      1  N   ASN A   1      22.066  40.557   0.420  1.00  0.00           N  
+ATOM      2  H   ASN A   1      21.629  41.305  -0.098  1.00  0.00           H  
+ENDMDL                                                                          
+END                                                                             

tests/test_pdb_tidy.py

@@ -181,6 +181,49 @@ def test_default_strict_stdin(self):
         # Check if we added END statements correctly
         self.assertTrue(self.stdout[-1].startswith('END'))
 
+    def test_corrects_model_format_and_numbers(self):
+        """Correct MODEL lines."""
+        fpath = os.path.join(data_dir, 'ensemble_error_MODEL.pdb')
+        sys.argv = ['', fpath]
+        self.exec_module()
+        self.assertEqual(self.retcode, 0)
+        self.assertEqual(len(self.stdout), 16)
+        self.assertEqual(len(self.stderr), 0)
+        self.assertEqual(len(self.stdout[5].strip()), 14)
+        self.assertEqual(len(self.stdout[5]), 80)
+        self.assertEqual(len(self.stdout[10].strip()), 14)
+        self.assertEqual(len(self.stdout[10]), 80)
+        self.assertEqual(self.stdout[5].strip(), "MODEL        1")
+        self.assertEqual(self.stdout[10].strip(), "MODEL        2")
+
+    def test_corrects_model_ENDMDL(self):
+        """Correct MODEL lines."""
+        fpath = os.path.join(data_dir, 'ensemble_error_4.pdb')
+        sys.argv = ['', fpath]
+        self.exec_module()
+        self.assertEqual(self.retcode, 0)
+        self.assertEqual(len(self.stdout), 14)
+        self.assertEqual(len(self.stderr), 0)
+
+        # MODEL lines
+        self.assertEqual(len(self.stdout[3].strip()), 14)
+        self.assertEqual(len(self.stdout[3]), 80)
+        self.assertEqual(len(self.stdout[3].strip()), 14)
+        self.assertEqual(len(self.stdout[8]), 80)
+        self.assertEqual(self.stdout[3].strip(), "MODEL        1")
+        self.assertEqual(self.stdout[8].strip(), "MODEL        2")
+
+        # ENDMDL LINES
+        self.assertEqual(len(self.stdout[7]), 80)
+        self.assertEqual(len(self.stdout[12]), 80)
+        self.assertEqual(self.stdout[7].strip(), "ENDMDL")
+        self.assertEqual(self.stdout[12].strip(), "ENDMDL")
+
+        # END LINES
+        self.assertEqual(self.stdout[-1].strip(), "END")
+        self.assertEqual(len(self.stdout[-1]), 80)
+
+
     def test_file_not_found(self):
         """$ pdb_tidy not_existing.pdb"""