XYZ elements

package/CHANGELOG

@@ -14,7 +14,7 @@ The rules for this file:
 
 ------------------------------------------------------------------------------
 mm/dd/yy richardjgowers, kain88-de, lilyminium, p-j-smith, bdice, joaomcteixeira,
-         PicoCentauri, davidercruz, jbarnoud
+         PicoCentauri, davidercruz, jbarnoud, RMeli
 
   * 0.21.0
 
@@ -35,6 +35,8 @@ Fixes
   * TXYZ parser uses strings for the atom types like other parsers (Issue #2435)
 
 Enhancements
+  * XYZ parser store elements attribute (#2420) and XYZ write uses the elements
+    attribute, if present (#2421).
   * Enhanges exception message when trajectory output file has no extension assigned.
   * Uniforms exception handling between Python 2.7 and Python 3: raised exceptions
     do not contain previous exceptions traceback. Uses six package to handle

package/MDAnalysis/coordinates/XYZ.py

@@ -108,6 +108,9 @@ class XYZWriter(base.WriterBase):
 
     .. _xyzplugin:
        http://www.ks.uiuc.edu/Research/vmd/plugins/molfile/xyzplugin.html
+
+    .. versionchanged: 0.21.0
+       Use elements attribute instead of names attribute, if present
     """
 
     format = 'XYZ'
@@ -149,15 +152,15 @@ def __init__(self, filename, n_atoms=None, atoms=None, convert_units=None,
             self.convert_units = convert_units
         else:
             self.convert_units = flags['convert_lengths']
-        self.atomnames = self._get_atomnames(atoms)
+        self.atomnames = self._get_atoms_elements_or_names(atoms)
         default_remark = "Written by {0} (release {1})".format(
             self.__class__.__name__, __version__)
         self.remark = default_remark if remark is None else remark
         # can also be gz, bz2
         self._xyz = util.anyopen(self.filename, 'wt')
 
-    def _get_atomnames(self, atoms):
-        """Return a list of atom names"""
+    def _get_atoms_elements_or_names(self, atoms):
+        """Return a list of atom elements (if present) or fallback to atom names"""
         # Default case
         if atoms is None:
             return itertools.cycle(('X',))
@@ -170,9 +173,12 @@ def _get_atomnames(self, atoms):
         # AtomGroup or Universe, grab the names else default
         # (AtomGroup.atoms just returns AtomGroup)
         try:
-            return atoms.atoms.names
+            return atoms.atoms.elements
         except (AttributeError, NoDataError):
-            return itertools.cycle(('X',))
+            try:
+                return atoms.atoms.names
+            except (AttributeError, NoDataError):
+                return itertools.cycle(('X',))
 
     def close(self):
         """Close the trajectory file and finalize the writing"""
@@ -184,7 +190,7 @@ def close(self):
     def write(self, obj):
         """Write object `obj` at current trajectory frame to file.
 
-        Atom names in the output are taken from the `obj` or default
+        Atom elements (or names) in the output are taken from the `obj` or default
         to the value of the `atoms` keyword supplied to the
         :class:`XYZWriter` constructor.
 
@@ -218,7 +224,7 @@ def write(self, obj):
                 # For Universe only --- get everything
                 ts = obj.trajectory.ts
             # update atom names
-            self.atomnames = self._get_atomnames(atoms)
+            self.atomnames = self._get_atoms_elements_or_names(atoms)
 
         self.write_next_timestep(ts)
 

package/MDAnalysis/topology/XYZParser.py

@@ -55,6 +55,7 @@
     Resids,
     Resnums,
     Segids,
+    Elements,
 )
 
 
@@ -69,6 +70,9 @@ class XYZParser(TopologyReaderBase):
      - Masses
 
     .. versionadded:: 0.9.1
+
+    .. versionchanged: 0.21.0
+       Store elements attribute, based on XYZ atom names
     """
     format = 'XYZ'
 
@@ -92,7 +96,7 @@ def parse(self, **kwargs):
 
         # Guessing time
         atomtypes = guessers.guess_types(names)
-        masses = guessers.guess_masses(atomtypes)
+        masses = guessers.guess_masses(names)
 
         attrs = [Atomnames(names),
                  Atomids(np.arange(natoms) + 1),
@@ -101,7 +105,7 @@ def parse(self, **kwargs):
                  Resids(np.array([1])),
                  Resnums(np.array([1])),
                  Segids(np.array(['SYSTEM'], dtype=object)),
-                 ]
+                 Elements(names)]
 
         top = Topology(natoms, 1, 1,
                        attrs=attrs)

testsuite/MDAnalysisTests/coordinates/test_xyz.py

@@ -161,3 +161,18 @@ def test_list_names(self, outfile):
 
         u2 = mda.Universe(outfile)
         assert all(u2.atoms.names == names)
+
+    @pytest.mark.parametrize("attr", ["elements", "names"])
+    def test_elements_and_names(self, outfile, attr):
+
+        u = mda.Universe.empty(n_atoms=5, trajectory=True)
+
+        u.add_TopologyAttr(attr, values=['Te', 'S', 'Ti', 'N', 'Ga'])
+
+        with mda.Writer(outfile) as w:
+            w.write(u)
+
+        with open(outfile, "r") as r:
+            names = ''.join(l.split()[0].strip() for l in r.readlines()[2:-1])
+
+        assert names[:-1].lower() == 'testing'

testsuite/MDAnalysisTests/topology/test_xyz.py

@@ -37,10 +37,9 @@ class XYZBase(ParserBase):
     parser = mda.topology.XYZParser.XYZParser
     expected_n_residues = 1
     expected_n_segments = 1
-    expected_attrs = ['names']
+    expected_attrs = ['names', "elements"]
     guessed_attrs = ['types', 'masses']
 
-
 class TestXYZMini(XYZBase):
     ref_filename = XYZ_mini
     expected_n_atoms = 3