Amber/sqm and Gromacs/gro Enhancement

dpdata/amber/sqm.py

@@ -1,11 +1,13 @@
 import numpy as np
+from ..periodic_table import ELEMENTS
 
 
 START = 0
-READ_CHARGE = 1
-READ_CHARGE_SUCCESS = 2
+READ_ENERGY = 1
+READ_CHARGE = 2
 READ_COORDS_START = 3
 READ_COORDS = 6
+READ_FORCES = 7
 
 def parse_sqm_out(fname):
     '''
@@ -14,41 +16,87 @@ def parse_sqm_out(fname):
     atom_symbols = []
     coords = []
     charges = []
+    forces = []
+    energies = []
+
     with open(fname) as f:
-        flag = 0
+        flag = START
         for line in f:
-            if line.startswith("  Atom    Element       Mulliken Charge"):
+            if line.startswith(" Total SCF energy"):
+                flag = READ_ENERGY
+            elif line.startswith("  Atom    Element       Mulliken Charge"):
                 flag = READ_CHARGE
             elif line.startswith(" Total Mulliken Charge"):
-                flag = READ_CHARGE_SUCCESS
+                flag = START
             elif line.startswith(" Final Structure"):
                 flag = READ_COORDS_START
+            elif line.startswith("QMMM: Forces on QM atoms"):
+                flag = READ_FORCES
+            elif flag == READ_ENERGY:
+                energy = float(line.strip().split()[-2])
+                energies.append(energy)
+                flag = START
             elif flag == READ_CHARGE:
                 ls = line.strip().split()
                 atom_symbols.append(ls[-2])
                 charges.append(float(ls[-1]))
             elif READ_COORDS_START <= flag < READ_COORDS:
                 flag += 1
             elif flag == READ_COORDS:
-                ls = line.strip()
-                if not ls:
-                    break
-                else:
-                    symbol = line.strip().split()[-4]
-                    coord = list(map(float, line.strip().split()[-3:]))
-                    coords.append(coord)
-    return atom_symbols, charges, np.array(coords)
-
-
-def to_system_data(fname):
+                coords.append([float(x) for x in line.strip().split()[-3:]])
+                if len(coords) == len(charges):
+                    flag = START
+            elif flag == READ_FORCES:
+                forces.append([float(x) for x in line.strip().split()[-3:]])
+                if len(forces) == len(charges):
+                    flag = START
+
     data = {}
-    atom_symbols, charges, coords = parse_sqm_out(fname)
     atom_names, data['atom_types'], atom_numbs = np.unique(atom_symbols, return_inverse=True, return_counts=True)
+    data['charges'] = np.array(charges)
     data['atom_names'] = list(atom_names)
     data['atom_numbs'] = list(atom_numbs)
-    data['charges'] = np.array([charges])
-    data['coords'] = np.array([coords])
     data['orig'] = np.array([0, 0, 0])
     data['cells'] = np.array([[[100., 0., 0.], [0., 100., 0.], [0., 0., 100.]]])
     data['nopbc'] = True
+    data['coords'] = np.array([coords])
+
+    energies = np.array(energies)
+    forces = np.array([forces], dtype=np.float32)
+    if len(forces) > 0:
+        data['energies'] = energies
+        data['forces'] = forces
+
     return data
+
+def make_sqm_in(data, fname=None, frame_idx=0, **kwargs):
+    symbols = [data['atom_names'][ii] for ii in data['atom_types']]
+    atomic_numbers = [ELEMENTS.index(ss) + 1 for ss in symbols]
+    charge = kwargs.get("charge", 0)
+
+    # multiplicity
+    mult = kwargs.get("mult", 1)
+    if mult != 1 :
+        raise RuntimeError("Multiplicity is not 1, which is not supported by sqm")
+
+    maxcyc = kwargs.get("maxcyc", 0)    # 0 represents a single-point calculation
+    theory = kwargs.get("qm_theory", "DFTB3")
+    ret = "Run semi-emperical minimization\n"
+    ret += " &qmmm\n"
+    ret += f"     qm_theory='{theory}'\n"
+    ret += f"     qmcharge={charge}\n"
+    ret += f"     maxcyc={maxcyc}\n"
+    ret += "     verbosity=4\n"
+    ret += " /\n"
+    for ii in range(len(data['atom_types'])):
+        ret += "{:>4s}{:>6s}{:>14s}{:>14s}{:>14s}\n".format(
+            str(atomic_numbers[ii]),
+            str(symbols[ii]),
+            f"{data['coords'][frame_idx][ii, 0]:.4f}",
+            f"{data['coords'][frame_idx][ii, 1]:.4f}",
+            f"{data['coords'][frame_idx][ii, 2]:.4f}"
+        )
+    if fname is not None:
+        with open(fname, 'w') as fp:
+            fp.write(ret)
+    return ret

dpdata/gromacs/gro.py

@@ -41,7 +41,7 @@ def _get_cell(line):
     cell = cell * nm2ang
     return cell
 
-def file_to_system_data(fname, format_atom_name=True):
+def file_to_system_data(fname, format_atom_name=True, **kwargs):
     system = {'coords': [], 'cells': []}
     with open(fname) as fp:
         frame = 0
@@ -74,7 +74,9 @@ def file_to_system_data(fname, format_atom_name=True):
     system['cells'] = np.array(system['cells'])
     return system
 
-def from_system_data(system, f_idx=0):
+def from_system_data(system, f_idx=0, **kwargs):
+    resname = kwargs.get("resname", "MOL")
+    shift = kwargs.get("shift", 0)
     ret = ""
     ret += " molecule" + "\n"
     n_atoms = sum(system["atom_numbs"])
@@ -83,8 +85,8 @@ def from_system_data(system, f_idx=0):
         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]))
+        ret += "{:>5d}{:<5s}{:>5s}{:5d}{:8.3f}{:8.3f}{:8.3f}\n".format(1, resname, atom_name, i+shift+1, *tuple(coords[i]))
     cell = (system["cells"][f_idx].flatten() * ang2nm)[cell_idx_gmx2dp]
-    ret += " " + " ".join([str(x) for x in cell])
+    ret += " " + " ".join([f"{x:.3f}" for x in cell])
 
     return ret

dpdata/plugins/amber.py

@@ -34,4 +34,20 @@ def from_system(self, fname, **kwargs):
         '''
         Read from ambertools sqm.out
         '''
-        return dpdata.amber.sqm.to_system_data(fname)
+        return dpdata.amber.sqm.parse_sqm_out(fname)
+
+    def from_labeled_system(self, fname, **kwargs):
+        '''
+        Read from ambertools sqm.out
+        '''
+        data = dpdata.amber.sqm.parse_sqm_out(fname)
+        assert "forces" in list(data.keys()), f"No forces in {fname}"
+        return data
+
+@Format.register("sqm/in")
+class SQMINFormat(Format):
+    def to_system(self, data, fname=None, frame_idx=0, **kwargs):
+        """
+        Generate input files for semi-emperical calculation in sqm software
+        """
+        return dpdata.amber.sqm.make_sqm_in(data, fname, frame_idx, **kwargs)

dpdata/plugins/gromacs.py

@@ -14,7 +14,7 @@ def from_system(self, file_name, format_atom_name=True, **kwargs):
         file_name : str
             The input file name
         """
-        return dpdata.gromacs.gro.file_to_system_data(file_name, format_atom_name=format_atom_name)
+        return dpdata.gromacs.gro.file_to_system_data(file_name, format_atom_name=format_atom_name, **kwargs)
 
     def to_system(self, data, file_name=None, frame_idx=-1, **kwargs):
         """
@@ -31,12 +31,13 @@ def to_system(self, data, file_name=None, frame_idx=-1, **kwargs):
         if frame_idx == -1:
             strs = []
             for idx in range(data['coords'].shape[0]):
-                gro_str = dpdata.gromacs.gro.from_system_data(data, f_idx=idx)
+                gro_str = dpdata.gromacs.gro.from_system_data(data, f_idx=idx,
+                                                              **kwargs)
                 strs.append(gro_str)
             gro_str = "\n".join(strs)
         else:
             gro_str = dpdata.gromacs.gro.from_system_data(
-                data, f_idx=frame_idx)
+                data, f_idx=frame_idx, **kwargs)
 
         if file_name is None:
             return gro_str

tests/amber/methane.mol

@@ -0,0 +1,14 @@
+
+     RDKit          3D
+
+  5  4  0  0  0  0  0  0  0  0999 V2000
+   -0.0221    0.0032    0.0165 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.6690    0.8894   -0.1009 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.3778   -0.8578   -0.5883 H   0  0  0  0  0  0  0  0  0  0  0  0
+    0.0964   -0.3151    1.0638 H   0  0  0  0  0  0  0  0  0  0  0  0
+    0.9725    0.2803   -0.3911 H   0  0  0  0  0  0  0  0  0  0  0  0
+  1  2  1  0
+  1  3  1  0
+  1  4  1  0
+  1  5  1  0
+M  END

tests/amber/sqm.in

@@ -0,0 +1,12 @@
+Run semi-emperical minimization
+ &qmmm
+     qm_theory='DFTB3'
+     qmcharge=0
+     maxcyc=0
+     verbosity=4
+ /
+   6     C       -0.0221        0.0032        0.0165
+   1     H       -0.6690        0.8894       -0.1009
+   1     H       -0.3778       -0.8578       -0.5883
+   1     H        0.0964       -0.3151        1.0638
+   1     H        0.9725        0.2803       -0.3911