diff --git a/package/CHANGELOG b/package/CHANGELOG index 3ce7cedc140..b6527b9ea23 100644 --- a/package/CHANGELOG +++ b/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 diff --git a/package/MDAnalysis/coordinates/XYZ.py b/package/MDAnalysis/coordinates/XYZ.py index 470f8fecb3b..bb2f9d18082 100644 --- a/package/MDAnalysis/coordinates/XYZ.py +++ b/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) diff --git a/package/MDAnalysis/topology/XYZParser.py b/package/MDAnalysis/topology/XYZParser.py index fc37a026c57..c1167cab275 100644 --- a/package/MDAnalysis/topology/XYZParser.py +++ b/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) diff --git a/testsuite/MDAnalysisTests/coordinates/test_xyz.py b/testsuite/MDAnalysisTests/coordinates/test_xyz.py index 2fa7336e03c..6436233fa1d 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_xyz.py +++ b/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' \ No newline at end of file diff --git a/testsuite/MDAnalysisTests/topology/test_xyz.py b/testsuite/MDAnalysisTests/topology/test_xyz.py index 8cc2e5857fc..21f1a1937d7 100644 --- a/testsuite/MDAnalysisTests/topology/test_xyz.py +++ b/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