HETATM record type written out

package/AUTHORS

@@ -147,6 +147,7 @@ Chronological list of authors
   - Andrea Rizzi
   - William Glass
   - Marcello Sega
+  - Mieczyslaw Torchala
 
 External code
 -------------

package/CHANGELOG

@@ -19,6 +19,8 @@ The rules for this file:
   * 2.0.0
 
 Fixes
+  * Instead of using ATOM for both ATOM and HETATM, HETATM record type
+    keyword is properly written out by PDB writer (Issue #1753, PR #2880)
   * Bond attribute is no longer set if PDB file contains no CONECT records
     (Issue #2832)
   * ResidueAttrs now have Atom as a target class by default; ICodes and 

package/MDAnalysis/coordinates/PDB.py

@@ -475,6 +475,8 @@ class PDBWriter(base.WriterBase):
     .. _MODEL: http://www.wwpdb.org/documentation/file-format-content/format32/sect9.html#MODEL
     .. _ENDMDL: http://www.wwpdb.org/documentation/file-format-content/format32/sect9.html#ENDMDL
     .. _CONECT: http://www.wwpdb.org/documentation/file-format-content/format32/sect10.html#CONECT
+    .. _ATOM: http://www.wwpdb.org/documentation/file-format-content/format32/sect9.html#ATOM
+    .. _HETATM: http://www.wwpdb.org/documentation/file-format-content/format32/sect9.html#HETATM
 
 
     Note
@@ -493,6 +495,10 @@ class PDBWriter(base.WriterBase):
     are not set (:code:`None` or :code:`np.zeros(6)`,
     see Issue #2698).
 
+    When atom's record_type attribute is present (i.e. Universe object was
+    created by loading a PDB file), ATOM and HETATM record type keywords
+    are written out accordingly.
+
     See Also
     --------
     This class is identical to :class:`MultiPDBWriter` with the one
@@ -528,6 +534,11 @@ class PDBWriter(base.WriterBase):
             "{chainID:1s}{resSeq:4d}{iCode:1s}"
             "   {pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}{occupancy:6.2f}"
             "{tempFactor:6.2f}      {segID:<4s}{element:>2s}\n"),
+        'HETATM': (
+            "HETATM{serial:5d} {name:<4s}{altLoc:<1s}{resName:<4s}"
+            "{chainID:1s}{resSeq:4d}{iCode:1s}"
+            "   {pos[0]:8.3f}{pos[1]:8.3f}{pos[2]:8.3f}{occupancy:6.2f}"
+            "{tempFactor:6.2f}      {segID:<4s}{element:>2s}\n"),
         'REMARK': "REMARK     {0}\n",
         'COMPND': "COMPND    {0}\n",
         'HEADER': "HEADER    {0}\n",
@@ -995,6 +1006,10 @@ def _write_timestep(self, ts, multiframe=False):
            ChainID now comes from the last character of segid, as stated in the documentation.
            An indexing issue meant it previously used the first charater (Issue #2224)
 
+        .. versionchanged:: 2.0.0
+           when only record_type attribute is present, instead of using ATOM for both ATOM
+           and HETATM, HETATM record type keyword is properly written out (Issue #1753)
+
         """
         atoms = self.obj.atoms
         pos = atoms.positions
@@ -1028,6 +1043,7 @@ def get_attr(attrname, default):
         occupancies = get_attr('occupancies', 1.0)
         tempfactors = get_attr('tempfactors', 0.0)
         atomnames = get_attr('names', 'X')
+        record_types = get_attr('record_types', 'ATOM')
 
         for i, atom in enumerate(atoms):
             vals = {}
@@ -1044,8 +1060,14 @@ def get_attr(attrname, default):
             vals['segID'] = segids[i][:4]
             vals['element'] = guess_atom_element(atomnames[i].strip())[:2]
 
-            # .. _ATOM: http://www.wwpdb.org/documentation/file-format-content/format32/sect9.html#ATOM
-            self.pdbfile.write(self.fmt['ATOM'].format(**vals))
+            # record_type attribute, if exists, can be ATOM or HETATM
+            try:
+                self.pdbfile.write(self.fmt[record_types[i]].format(**vals))
+            except KeyError:
+                warnings.warn("Found '{}' for record type, but allowed types "
+                              "are ATOM or HETATM".format(record_types[i]))
+                raise ValueError
+
         if multiframe:
             self.ENDMDL()
         self.frames_written += 1

testsuite/MDAnalysisTests/coordinates/test_pdb.py

@@ -20,12 +20,12 @@
 # MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
 # J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
 #
-import pytest
-from io import StringIO
 import os
+from io import StringIO
 
 import MDAnalysis as mda
 import numpy as np
+import pytest
 from MDAnalysisTests import make_Universe
 from MDAnalysisTests.coordinates.base import _SingleFrameReader
 from MDAnalysisTests.coordinates.reference import (RefAdKSmall,
@@ -36,7 +36,8 @@
                                        INC_PDB, PDB_xlserial, ALIGN, ENT,
                                        PDB_cm, PDB_cm_gz, PDB_cm_bz2,
                                        PDB_mc, PDB_mc_gz, PDB_mc_bz2,
-                                       PDB_CRYOEM_BOX, MMTF_NOCRYST)
+                                       PDB_CRYOEM_BOX, MMTF_NOCRYST,
+                                       PDB_HOLE, mol2_molecule)
 from numpy.testing import (assert_equal,
                            assert_array_almost_equal,
                            assert_almost_equal)
@@ -188,6 +189,14 @@ def universe2(self):
     def universe3(self):
         return mda.Universe(PDB)
 
+    @pytest.fixture
+    def universe4(self):
+        return mda.Universe(PDB_HOLE, PDB_HOLE)
+
+    @pytest.fixture
+    def universe5(self):
+        return mda.Universe(mol2_molecule, mol2_molecule)
+
     @pytest.fixture(params=[
             [PDB_CRYOEM_BOX, np.zeros(6)],
             [MMTF_NOCRYST, None]
@@ -446,6 +455,62 @@ def test_stringio_outofrange(self, universe3):
             with mda.coordinates.PDB.PDBWriter(outstring) as writer:
                 writer.write(u.atoms)
 
+    def test_hetatm_written(self, universe4, tmpdir):
+        """
+        Checks that HETATM record types are written.
+        """
+
+        u = universe4
+        u_atoms = u.select_atoms("resname ETA and record_type HETATM")
+        assert_equal(len(u_atoms), 8)
+
+        outfile = str(tmpdir.join('test-hetatm.pdb'))
+        u.atoms.write(outfile)
+        written = mda.Universe(outfile, outfile)
+        written_atoms = written.select_atoms("resname ETA and "
+                                             "record_type HETATM")
+
+        assert_equal(len(u_atoms), len(written_atoms))
+
+    def test_default_atom_record_type_written(self, universe5, tmpdir):
+        """
+        Checks that ATOM record types are written when there is no
+        record_type attribute.
+        """
+
+        u = universe5
+        outfile = str(tmpdir.join('test-mol2-to-pdb.pdb'))
+
+        expected_msg = "Found no information for attr: " \
+                       "'record_types' Using default value of 'ATOM'"
+        with pytest.warns(UserWarning, match=expected_msg):
+            u.atoms.write(outfile)
+
+        written = mda.Universe(outfile, outfile)
+        assert_equal(len(u.atoms), len(written.atoms))
+
+        atms = written.select_atoms("record_type ATOM")
+        assert_equal(len(atms.atoms), len(u.atoms))
+
+        hetatms = written.select_atoms("record_type HETATM")
+        assert_equal(len(hetatms.atoms), 0)
+
+    def test_abnormal_record_type(self, universe5, tmpdir):
+        """
+        Checks whether KeyError is raised when record type is
+        neither ATOM or HETATM.
+        """
+        u = universe5
+        u.add_TopologyAttr('record_type', ['ABNORM']*len(u.atoms))
+        outfile = str(tmpdir.join('test-abnormal-record_type.pdb'))
+
+        expected_msg = "Found 'ABNORM' for record type, but allowed " \
+                       "types are ATOM or HETATM"
+
+        with pytest.raises(ValueError):
+            with pytest.warns(UserWarning, match=expected_msg):
+                u.atoms.write(outfile)
+
 
 class TestMultiPDBReader(object):
     @staticmethod