diff --git a/README.md b/README.md index 8dd9c8649..1f4b33387 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,7 @@ The `System` or `LabeledSystem` can be constructed from the following file forma | PWmat | movement | True | True | LabeledSystem | 'pwmat/movement' | | PWmat | OUT.MLMD | True | True | LabeledSystem | 'pwmat/out.mlmd' | | Amber | multi | True | True | LabeledSystem | 'amber/md' | -| Gromacs | gro | False | False | System | 'gromacs/gro' | +| Gromacs | gro | True | False | System | 'gromacs/gro' | The Class `dpdata.MultiSystems` can read data from a dir which may contains many files of different systems, or from single xyz file which contains different systems. diff --git a/dpdata/gromacs/gro.py b/dpdata/gromacs/gro.py index 6e12fb1e1..b2473a03b 100644 --- a/dpdata/gromacs/gro.py +++ b/dpdata/gromacs/gro.py @@ -1,11 +1,23 @@ #!/usr/bin/env python3 - +import re import numpy as np +import json +import os nm2ang = 10. +ang2nm = 1. / nm2ang +cell_idx_gmx2dp = [0, 4, 8, 1, 2, 3, 5, 6, 7] + +def _format_atom_name(atom_name): + patt = re.compile("[a-zA-Z]*") + match = re.search(patt, atom_name) + fmt_name = match.group().capitalize() + return fmt_name -def _get_line(line): +def _get_line(line, fmt_atom_name=True): atom_name = line[10:15].split()[0] + if fmt_atom_name: + atom_name = _format_atom_name(atom_name) atom_idx = int(line[15:20].split()[0]) posis = [float(line[ii:ii+8]) for ii in range(20,44,8)] posis = np.array(posis) * nm2ang @@ -29,27 +41,50 @@ def _get_cell(line): cell = cell * nm2ang return cell -def file_to_system_data(fname): - names = [] - idxs = [] - posis = [] +def file_to_system_data(fname, format_atom_name=True): + system = {'coords': [], 'cells': []} with open(fname) as fp: - fp.readline() - natoms = int(fp.readline()) - for ii in range(natoms): - n, i, p = _get_line(fp.readline()) - names.append(n) - idxs.append(i) - posis.append(p) - cell = _get_cell(fp.readline()) - posis = np.array(posis) - system = {} - system['orig'] = np.array([0, 0, 0]) - system['atom_names'] = list(set(names)) - system['atom_names'].sort() - system['atom_numbs'] = [names.count(ii) for ii in system['atom_names']] - system['atom_types'] = [system['atom_names'].index(ii) for ii in names] - system['atom_types'] = np.array(system['atom_types'], dtype = int) - system['coords'] = np.array([posis]) - system['cells'] = np.array([cell]) + frame = 0 + while True: + flag = fp.readline() + if not flag: + break + else: + frame += 1 + names = [] + idxs = [] + posis = [] + natoms = int(fp.readline()) + for ii in range(natoms): + n, i, p = _get_line(fp.readline(), fmt_atom_name=format_atom_name) + names.append(n) + idxs.append(i) + posis.append(p) + cell = _get_cell(fp.readline()) + posis = np.array(posis) + if frame == 1: + system['orig'] = np.zeros(3) + system['atom_names'] = list(set(names)) + system['atom_numbs'] = [names.count(ii) for ii in system['atom_names']] + system['atom_types'] = [system['atom_names'].index(ii) for ii in names] + system['atom_types'] = np.array(system['atom_types'], dtype = int) + system['coords'].append(posis) + system['cells'].append(cell) + system['coords'] = np.array(system['coords']) + system['cells'] = np.array(system['cells']) return system + +def from_system_data(system, f_idx=0): + ret = "" + ret += " molecule" + "\n" + n_atoms = sum(system["atom_numbs"]) + ret += " " + str(n_atoms) + "\n" + for i in range(n_atoms): + atom_type = system["atom_types"][i] + atom_name = system["atom_names"][atom_type] + coords = system["coords"][f_idx] * ang2nm + ret += "{:>5d}{:<5s}{:>5s}{:5d}{:8.3f}{:8.3f}{:8.3f}\n".format(1, "MOL", atom_name, i+1, *tuple(coords[i])) + cell = (system["cells"][f_idx].flatten() * ang2nm)[cell_idx_gmx2dp] + ret += " " + " ".join([str(x) for x in cell]) + + return ret diff --git a/dpdata/system.py b/dpdata/system.py index 0af601896..5db810396 100644 --- a/dpdata/system.py +++ b/dpdata/system.py @@ -622,7 +622,7 @@ def from_deepmd_raw(self, folder, type_map = None) : @register_from_funcs.register_funcs("gro") @register_from_funcs.register_funcs("gromacs/gro") - def from_gromacs_gro(self, file_name) : + def from_gromacs_gro(self, file_name, format_atom_name=True) : """ Load gromacs .gro file @@ -631,7 +631,35 @@ def from_gromacs_gro(self, file_name) : file_name : str The input file name """ - self.data = dpdata.gromacs.gro.file_to_system_data(file_name) + self.data = dpdata.gromacs.gro.file_to_system_data(file_name, format_atom_name=format_atom_name) + + @register_to_funcs.register_funcs("gromacs/gro") + def to_gromacs_gro(self, file_name=None, frame_idx=-1): + """ + Dump the system in gromacs .gro format + + Parameters + ---------- + file_name : str or None + The output file name. If None, return the file content as a string + frame_idx : int + The index of the frame to dump + """ + assert(frame_idx < len(self.data['coords'])) + if frame_idx == -1: + strs = [] + for idx in range(self.get_nframes()): + gro_str = dpdata.gromacs.gro.from_system_data(self.data, f_idx=idx) + strs.append(gro_str) + gro_str = "\n".join(strs) + else: + gro_str = dpdata.gromacs.gro.from_system_data(self.data, f_idx=frame_idx) + + if file_name is None: + return gro_str + else: + with open(file_name, 'w+') as fp: + fp.write(gro_str) @register_to_funcs.register_funcs("deepmd/npy") def to_deepmd_npy(self, folder, set_size = 5000, prec=np.float32) : diff --git a/tests/gromacs/case_for_format_atom_name.gro b/tests/gromacs/case_for_format_atom_name.gro new file mode 100644 index 000000000..29b248bd4 --- /dev/null +++ b/tests/gromacs/case_for_format_atom_name.gro @@ -0,0 +1,35 @@ +ligand in water t= 0.00000 step= 0 + 32 + 1MOL C1 1 2.763 2.029 0.966 + 1MOL C2 2 2.844 2.141 0.941 + 1MOL C3 3 2.788 2.268 0.958 + 1MOL C4 4 2.653 2.285 0.991 + 1MOL C5 5 2.576 2.171 1.011 + 1MOL C6 6 2.630 2.042 1.007 + 1MOL CL1 7 2.405 2.194 1.045 + 1MOL C7 8 2.588 2.419 0.999 + 1MOL O1 9 2.558 2.474 0.894 + 1MOL N1 10 2.572 2.471 1.125 + 1MOL C8 11 2.504 2.587 1.170 + 1MOL C9 12 2.542 2.648 1.289 + 1MOL C10 13 2.461 2.753 1.336 + 1MOL N2 14 2.359 2.804 1.266 + 1MOL C11 15 2.321 2.742 1.153 + 1MOL C12 16 2.389 2.631 1.103 + 1MOL N3 17 2.217 2.806 1.082 + 1MOL C13 18 2.118 2.742 1.011 + 1MOL O2 19 2.121 2.624 0.980 + 1MOL C14 20 2.003 2.827 0.960 + 1MOL CL2 21 2.875 2.415 0.920 + 1MOL H1 22 2.026 2.932 0.936 + 1MOL H2 23 2.810 1.932 0.951 + 1MOL H3 24 2.945 2.132 0.902 + 1MOL H4 25 2.577 1.954 1.041 + 1MOL H5 26 2.608 2.408 1.196 + 1MOL H6 27 2.624 2.605 1.346 + 1MOL H7 28 2.483 2.804 1.430 + 1MOL H8 29 2.357 2.582 1.012 + 1MOL H9 30 2.203 2.904 1.105 + 1MOL H10 31 1.920 2.816 1.031 + 1MOL H11 32 1.967 2.787 0.864 + 3.24290 3.24290 2.29308 0.00000 0.00000 0.00000 0.00000 1.62145 1.62145 \ No newline at end of file diff --git a/tests/gromacs/multi_frames.gro b/tests/gromacs/multi_frames.gro new file mode 100644 index 000000000..2f16e4d6a --- /dev/null +++ b/tests/gromacs/multi_frames.gro @@ -0,0 +1,24 @@ +frame_1 +9 + 1SOL O 1 0.135 0.183 0.341 + 1SOL H 2 0.177 0.149 0.262 + 1SOL H 3 0.046 0.149 0.339 + 2SOL O 4 0.520 0.447 0.111 + 2SOL H 5 0.567 0.481 0.035 + 2SOL H 6 0.568 0.481 0.186 + 3SOL O 7 0.651 0.539 0.335 + 3SOL H 8 0.653 0.634 0.336 + 3SOL H 9 0.743 0.512 0.336 + 0.7822838765564372 0.7353572647182051 0.9036518515423753 +frame_2 +9 + 1SOL O 1 0.135 0.183 0.341 + 1SOL H 2 0.177 0.149 0.262 + 1SOL H 3 0.046 0.149 0.339 + 2SOL O 4 0.520 0.447 0.111 + 2SOL H 5 0.567 0.481 0.035 + 2SOL H 6 0.568 0.481 0.186 + 3SOL O 7 0.651 0.539 0.335 + 3SOL H 8 0.653 0.634 0.336 + 3SOL H 9 0.743 0.512 0.336 + 0.7822838765564372 0.7353572647182051 0.9036518515423753 diff --git a/tests/test_gromacs_gro.py b/tests/test_gromacs_gro.py index 1acb4d27f..16070448b 100644 --- a/tests/test_gromacs_gro.py +++ b/tests/test_gromacs_gro.py @@ -5,8 +5,9 @@ class TestGromacsGro(unittest.TestCase): def test_read_file(self): - system = dpdata.System('gromacs/1h.gro') - self.assertEqual(system['atom_names'], ['H', 'O']) + system = dpdata.System('gromacs/1h.gro', type_map=['H', 'O']) + self.assertTrue('H' in system['atom_names']) + self.assertTrue('O' in system['atom_names']) self.assertEqual(system['atom_numbs'], [6, 3]) for cc,ii in enumerate([1, 0, 0, 1, 0, 0, 1, 0, 0]): self.assertEqual(system['atom_types'][cc], ii) @@ -24,8 +25,9 @@ def test_read_file(self): self.assertAlmostEqual(system['coords'][0][8][2], 3.36) def test_read_file_tri(self): - system = dpdata.System('gromacs/1h.tri.gro') - self.assertEqual(system['atom_names'], ['H', 'O']) + system = dpdata.System('gromacs/1h.tri.gro', type_map=['H', 'O']) + self.assertTrue('H' in system['atom_names']) + self.assertTrue('O' in system['atom_names']) self.assertEqual(system['atom_numbs'], [6, 3]) for cc,ii in enumerate([1, 0, 0, 1, 0, 0, 1, 0, 0]): self.assertEqual(system['atom_types'][cc], ii) @@ -44,3 +46,68 @@ def test_read_file_tri(self): self.assertAlmostEqual(system['coords'][0][8][1], 5.12) self.assertAlmostEqual(system['coords'][0][8][2], 3.36) system.to('vasp/poscar', 'POSCAR') + +class TestGromacsGroMultiFrames(unittest.TestCase): + def test_read_file(self): + system = dpdata.System('gromacs/multi_frames.gro', type_map=['H', 'O']) + self.assertTrue('H' in system['atom_names']) + self.assertTrue('O' in system['atom_names']) + self.assertEqual(system['atom_numbs'], [6, 3]) + for cc,ii in enumerate([1, 0, 0, 1, 0, 0, 1, 0, 0]): + self.assertEqual(system['atom_types'][cc], ii) + self.assertEqual(len(system['cells']), 2) + self.assertEqual(len(system['coords']), 2) + for ii in range(3): + for jj in range(3): + if ii != jj: + self.assertAlmostEqual(system['cells'][0][ii][jj], 0) # frame no.1 + self.assertAlmostEqual(system['cells'][1][ii][jj], 0) # frame no.2 + # frame no.1 + self.assertAlmostEqual(system['cells'][0][0][0], 7.822838765564372) + self.assertAlmostEqual(system['cells'][0][1][1], 7.353572647182051) + self.assertAlmostEqual(system['cells'][0][2][2], 9.036518515423753) + self.assertAlmostEqual(system['coords'][0][8][0], 7.43) + self.assertAlmostEqual(system['coords'][0][8][1], 5.12) + self.assertAlmostEqual(system['coords'][0][8][2], 3.36) + # frame no.2 + self.assertAlmostEqual(system['cells'][1][0][0], 7.822838765564372) + self.assertAlmostEqual(system['cells'][1][1][1], 7.353572647182051) + self.assertAlmostEqual(system['cells'][1][2][2], 9.036518515423753) + self.assertAlmostEqual(system['coords'][1][8][0], 7.43) + self.assertAlmostEqual(system['coords'][1][8][1], 5.12) + self.assertAlmostEqual(system['coords'][1][8][2], 3.36) + + +class TestFormatAtomName(unittest.TestCase): + def test_format_atom_name(self): + system = dpdata.System("gromacs/case_for_format_atom_name.gro", fmt='gromacs/gro', type_map=['H','C','N','O','Cl']) + self.assertEqual(system.formula, "H11C14N3O2Cl2") + + def test_no_format_atom_name(self): + system = dpdata.System("gromacs/case_for_format_atom_name.gro", fmt='gromacs/gro', format_atom_name=False) + atoms = ['CL1', 'H6', 'C4', 'C3', 'C6', 'C11', 'H10', 'C2', 'N3', 'C14', + 'H7', 'H8', 'C13', 'H2', 'H1', 'H4', 'O2', 'H9', 'O1', 'N2', 'C9', + 'H3', 'C5', 'H11', 'N1', 'C7', 'C10', 'CL2', 'H5', 'C1', 'C8','C12'] + for at in atoms: + self.assertTrue(at in system['atom_names']) + + +class TestDumpGromacsGro(unittest.TestCase): + def setUp(self): + self.system = dpdata.System('gromacs/multi_frames.gro', type_map=['H', 'O']) + + def test_dump_single_frame(self): + self.system.to_gromacs_gro('gromacs/tmp_1.gro', frame_idx=0) + tmp = dpdata.System('gromacs/tmp_1.gro', type_map=['H', 'O']) + self.assertEqual(tmp.get_nframes(), 1) + + def test_dump_multi_frames(self): + self.system.to_gromacs_gro('gromacs/tmp_2.gro') + tmp = dpdata.System('gromacs/tmp_2.gro', type_map=['H', 'O']) + self.assertEqual(tmp.get_nframes(), 2) + + def tearDown(self): + if os.path.exists('gromacs/tmp_1.gro'): + os.remove('gromacs/tmp_1.gro') + if os.path.exists('gromacs/tmp_2.gro'): + os.remove('gromacs/tmp_2.gro')