Add support for reading and dumping multi-frames for gro file & add dump system to gaussian input file

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.

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

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) :

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

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

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')