Add writing u.trajectory.ts.data['molecule_tag'] as molecule_tag atom attribute to LAMMPS datafile

package/CHANGELOG

@@ -39,6 +39,7 @@ Fixes
     (Issue #3336)
 
 Enhancements
+  * Add writing u.trajectory.ts.data['molecule_tag'] as molecule tags to LAMMPS data file (Issue #3548)
   * Improved speed of chi1_selection (PR #4109)
   * Add `progressbar_kwargs` parameter to `AnalysisBase.run` method, allowing to modify description, position etc of tqdm progressbars.
   * Add a nojump transformation, which unwraps trajectories so that particle

package/MDAnalysis/coordinates/LAMMPS.py

@@ -273,7 +273,7 @@ def __init__(self, filename, convert_units=True, **kwargs):
         self.units['velocity'] = kwargs.pop('velocityunit',
                                  self.units['length']+'/'+self.units['time'])
 
-    def _write_atoms(self, atoms):
+    def _write_atoms(self, atoms, data):
         self.f.write('\n')
         self.f.write('Atoms\n')
         self.f.write('\n')
@@ -288,20 +288,23 @@ def _write_atoms(self, atoms):
         indices = atoms.indices + 1
         types = atoms.types.astype(np.int32)
 
+        moltags = data.get("molecule_tag", np.zeros(len(atoms), dtype=int))
+
         if self.convert_units:
             coordinates = self.convert_pos_to_native(atoms.positions, inplace=False)
 
         if has_charges:
-            for index, atype, charge, coords in zip(indices, types, charges,
-                    coordinates):
-                self.f.write('{i:d} 0 {t:d} {c:f} {x:f} {y:f} {z:f}\n'.format(
-                             i=index, t=atype, c=charge, x=coords[0],
-                             y=coords[1], z=coords[2]))
+            for index, moltag, atype, charge, coords in zip(indices, moltags,
+                    types, charges, coordinates):
+                x, y, z = coords
+                self.f.write(f"{index:d} {moltag:d} {atype:d} {charge:f}"
+                             f" {x:f} {y:f} {z:f}\n")
         else:
-            for index, atype, coords in zip(indices, types, coordinates):
-                self.f.write('{i:d} 0 {t:d} {x:f} {y:f} {z:f}\n'.format(
-                             i=index, t=atype, x=coords[0], y=coords[1],
-                             z=coords[2]))
+            for index, moltag, atype, coords in zip(indices, moltags, types,
+                    coordinates):
+                x, y, z = coords
+                self.f.write(f"{index:d} {moltag:d} {atype:d}"
+                             f" {x:f} {y:f} {z:f}\n")
 
     def _write_velocities(self, atoms):
         self.f.write('\n')
@@ -441,7 +444,7 @@ def write(self, selection, frame=None):
             self._write_dimensions(atoms.dimensions)
 
             self._write_masses(atoms)
-            self._write_atoms(atoms)
+            self._write_atoms(atoms, u.trajectory.ts.data)
             for attr in features.values():
                 if attr is None or len(attr) == 0:
                     continue

testsuite/MDAnalysisTests/coordinates/test_lammps.py

@@ -115,6 +115,32 @@ def LAMMPSDATAWriter(request, tmpdir_factory):
     return u, u_new
 
 
+@pytest.fixture(params=[
+    [LAMMPSdata, True],
+    [LAMMPSdata_mini, True],
+    [LAMMPScnt, True],
+    [LAMMPShyd, True],
+    [LAMMPSdata, False]
+], scope='module')
+def LAMMPSDATAWriter_molecule_tag(request, tmpdir_factory):
+    filename, charges = request.param
+    u = mda.Universe(filename)
+    if charges == False:
+        u.del_TopologyAttr('charges')
+
+    u.trajectory.ts.data['molecule_tag'] = u.atoms.resids
+
+    fn = os.path.split(filename)[1]
+    outfile = str(tmpdir_factory.mktemp('data').join(fn))
+
+    with mda.Writer(outfile, n_atoms=u.atoms.n_atoms) as w:
+        w.write(u.atoms)
+
+    u_new = mda.Universe(outfile)
+
+    return u, u_new
+
+
 def test_unwrap_vel_force():
 
     u_wrapped = mda.Universe(LAMMPS_image_vf, [LAMMPSDUMP_image_vf], 
@@ -187,6 +213,13 @@ def test_Writer_numerical_attrs(self, attr, LAMMPSDATAWriter):
                                 atol=1e-6)
 
 
+class TestLAMMPSDATAWriter_molecule_tag(object):
+    def test_molecule_tag(self, LAMMPSDATAWriter_molecule_tag):
+        u_ref, u_new = LAMMPSDATAWriter_molecule_tag
+        assert_equal(u_ref.atoms.resids, u_new.atoms.resids,
+                     err_msg="resids different after writing",)
+
+
 def test_datawriter_universe(tmpdir):
     fn = str(tmpdir.join('out.data'))