From 96c9afc3c2f01c802fc9b2cd949cbd273ee782ba Mon Sep 17 00:00:00 2001 From: Mikhail Glagolev Date: Wed, 5 Apr 2023 00:37:17 +0300 Subject: [PATCH 1/5] Add writing data['molecule_tag'] to LAMMPS datafile --- package/CHANGELOG | 1 + package/MDAnalysis/coordinates/LAMMPS.py | 29 +++++++++++------- .../coordinates/test_lammps.py | 30 +++++++++++++++++++ 3 files changed, 50 insertions(+), 10 deletions(-) diff --git a/package/CHANGELOG b/package/CHANGELOG index e228fdef8aa..f5d29a44bd8 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -37,6 +37,7 @@ Fixes (Issue #3336) Enhancements + * Add writing u.trajectory.ts.data['molecule_tag'] as molecule tags to LAMMPS data file (Issue #3548) * 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 paths are continuous. (Issue #3703, PR #4031) diff --git a/package/MDAnalysis/coordinates/LAMMPS.py b/package/MDAnalysis/coordinates/LAMMPS.py index d0d95b4ae61..da37d4f7f82 100644 --- a/package/MDAnalysis/coordinates/LAMMPS.py +++ b/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,27 @@ def _write_atoms(self, atoms): indices = atoms.indices + 1 types = atoms.types.astype(np.int32) + try: + moltags = data['molecule_tag'] + except KeyError: + moltags = [0] * len(atoms) + + 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], + for index, moltag, atype, charge, coords in zip(indices, moltags, + types, charges, coordinates): + self.f.write('{i:d} {m:d} {t:d} {c:f} {x:f} {y:f} {z:f}\n'.format( + i=index, m=moltag, t=atype, c=charge, x=coords[0], y=coords[1], z=coords[2])) 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): + self.f.write('{i:d} {m:d} {t:d} {x:f} {y:f} {z:f}\n'.format( + i=index, m=moltag, t=atype, x=coords[0], + y=coords[1], z=coords[2])) def _write_velocities(self, atoms): self.f.write('\n') @@ -397,6 +404,8 @@ def write(self, selection, frame=None): else: frame = u.trajectory.ts.frame + data = u.trajectory.ts.data + # make sure to use atoms (Issue 46) atoms = selection.atoms @@ -441,7 +450,7 @@ def write(self, selection, frame=None): self._write_dimensions(atoms.dimensions) self._write_masses(atoms) - self._write_atoms(atoms) + self._write_atoms(atoms, data) for attr in features.values(): if attr is None or len(attr) == 0: continue diff --git a/testsuite/MDAnalysisTests/coordinates/test_lammps.py b/testsuite/MDAnalysisTests/coordinates/test_lammps.py index b755c4aad02..6542be33d51 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_lammps.py +++ b/testsuite/MDAnalysisTests/coordinates/test_lammps.py @@ -115,6 +115,29 @@ def LAMMPSDATAWriter(request, tmpdir_factory): return u, u_new +@pytest.fixture(params=[ + LAMMPSdata, + LAMMPSdata_mini, + LAMMPScnt, + LAMMPShyd, +], scope='module') +def LAMMPSDATAWriter_molecule_tag(request, tmpdir_factory): + filename = request.param + u = mda.Universe(filename) + + 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 +210,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')) From c46853e3e77200f3829db8e09fa616a8fcdf439b Mon Sep 17 00:00:00 2001 From: Mikhail Glagolev Date: Tue, 11 Apr 2023 14:24:17 +0300 Subject: [PATCH 2/5] Update package/MDAnalysis/coordinates/LAMMPS.py Co-authored-by: Rocco Meli --- package/MDAnalysis/coordinates/LAMMPS.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/package/MDAnalysis/coordinates/LAMMPS.py b/package/MDAnalysis/coordinates/LAMMPS.py index da37d4f7f82..fc0e7188510 100644 --- a/package/MDAnalysis/coordinates/LAMMPS.py +++ b/package/MDAnalysis/coordinates/LAMMPS.py @@ -288,11 +288,7 @@ def _write_atoms(self, atoms, data): indices = atoms.indices + 1 types = atoms.types.astype(np.int32) - try: - moltags = data['molecule_tag'] - except KeyError: - moltags = [0] * len(atoms) - + 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) From c9344744d0fe31c674b92ec5374d6bfb5cd06387 Mon Sep 17 00:00:00 2001 From: Mikhail Glagolev Date: Thu, 27 Apr 2023 21:29:26 +0300 Subject: [PATCH 3/5] Improving code based on RMeli's suggestions --- package/MDAnalysis/coordinates/LAMMPS.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/package/MDAnalysis/coordinates/LAMMPS.py b/package/MDAnalysis/coordinates/LAMMPS.py index fc0e7188510..746d24497ee 100644 --- a/package/MDAnalysis/coordinates/LAMMPS.py +++ b/package/MDAnalysis/coordinates/LAMMPS.py @@ -296,15 +296,15 @@ def _write_atoms(self, atoms, data): if has_charges: for index, moltag, atype, charge, coords in zip(indices, moltags, types, charges, coordinates): - self.f.write('{i:d} {m:d} {t:d} {c:f} {x:f} {y:f} {z:f}\n'.format( - i=index, m=moltag, t=atype, c=charge, x=coords[0], - y=coords[1], z=coords[2])) + 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, moltag, atype, coords in zip(indices, moltags, types, coordinates): - self.f.write('{i:d} {m:d} {t:d} {x:f} {y:f} {z:f}\n'.format( - i=index, m=moltag, t=atype, x=coords[0], - y=coords[1], z=coords[2])) + 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') @@ -400,8 +400,6 @@ def write(self, selection, frame=None): else: frame = u.trajectory.ts.frame - data = u.trajectory.ts.data - # make sure to use atoms (Issue 46) atoms = selection.atoms @@ -446,7 +444,7 @@ def write(self, selection, frame=None): self._write_dimensions(atoms.dimensions) self._write_masses(atoms) - self._write_atoms(atoms, data) + self._write_atoms(atoms, u.trajectory.ts.data) for attr in features.values(): if attr is None or len(attr) == 0: continue From 62d97501f16ef6e654cbfe24bace27f35c9fb5bc Mon Sep 17 00:00:00 2001 From: Mikhail Glagolev Date: Tue, 2 May 2023 11:49:04 +0300 Subject: [PATCH 4/5] Adding a test on a system with deleted charges --- .../MDAnalysisTests/coordinates/test_lammps.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/testsuite/MDAnalysisTests/coordinates/test_lammps.py b/testsuite/MDAnalysisTests/coordinates/test_lammps.py index 6542be33d51..a97f3d87527 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_lammps.py +++ b/testsuite/MDAnalysisTests/coordinates/test_lammps.py @@ -116,14 +116,17 @@ def LAMMPSDATAWriter(request, tmpdir_factory): @pytest.fixture(params=[ - LAMMPSdata, - LAMMPSdata_mini, - LAMMPScnt, - LAMMPShyd, + [LAMMPSdata, True], + [LAMMPSdata_mini, True], + [LAMMPScnt, True], + [LAMMPShyd, True], + [LAMMPSdata, False] ], scope='module') def LAMMPSDATAWriter_molecule_tag(request, tmpdir_factory): - filename = request.param + filename, charges = request.param u = mda.Universe(filename) + if charges == False: + u.del_TopologyAttr('charges') u.trajectory.ts.data['molecule_tag'] = u.atoms.resids From eb7c177049d7b7b6437c428f4fc1a95464a9624c Mon Sep 17 00:00:00 2001 From: Mikhail Glagolev Date: Wed, 3 May 2023 13:39:55 +0300 Subject: [PATCH 5/5] Update testsuite/MDAnalysisTests/coordinates/test_lammps.py Co-authored-by: Rocco Meli --- testsuite/MDAnalysisTests/coordinates/test_lammps.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/testsuite/MDAnalysisTests/coordinates/test_lammps.py b/testsuite/MDAnalysisTests/coordinates/test_lammps.py index a97f3d87527..f3a9e53ab93 100644 --- a/testsuite/MDAnalysisTests/coordinates/test_lammps.py +++ b/testsuite/MDAnalysisTests/coordinates/test_lammps.py @@ -125,7 +125,7 @@ def LAMMPSDATAWriter(request, tmpdir_factory): def LAMMPSDATAWriter_molecule_tag(request, tmpdir_factory): filename, charges = request.param u = mda.Universe(filename) - if charges == False: + if not charges: u.del_TopologyAttr('charges') u.trajectory.ts.data['molecule_tag'] = u.atoms.resids