enhance and add several Amber-related functions (#135)

* add some useful functions: pick_atom_idx, correction

* pick_atom_idx for MultiSystems

* fix bugs

* move all constants to one file

* load amber energy from mdout file if mden file does not exist

* add remove_atom_names function; bugfix

* add nopbc.setter

* use_element_symbols for amber

* pick_by_amber_mask

* fix bugs

* make constants more accurate

* fix tests

* fix bug

* System from amber

* Revert "fix tests"

This reverts commit 9f65fac.

* Revert "make constants more accurate"

This reverts commit 43947eb.

* Revert "move all constants to one file"

This reverts commit 5a85a78.

* add tests and fix bugs

* add extras_require

* split unittests

Co-authored-by: Han Wang <[email protected]>

Author: njzjz

.github/workflows/test.yml

@@ -18,7 +18,7 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
-      run: pip install . coverage codecov
+      run: pip install .[amber] coverage codecov
     - name: Test
       run: cd tests && coverage run --source=../dpdata -m unittest && cd .. && coverage combine tests/.coverage && coverage report
     - run: codecov

dpdata/amber/mask.py

@@ -0,0 +1,35 @@
+"""Amber mask"""
+try:
+    import parmed
+except ImportError:
+    pass
+
+def pick_by_amber_mask(param, maskstr, coords=None):
+    """Pick atoms by amber masks
+    
+    Parameters
+    ----------
+    param: str or parmed.Structure
+      filename of Amber param file or parmed.Structure
+    maskstr: str
+      Amber masks
+    coords: np.ndarray (optional)
+      frame coordinates, shape: N*3
+    """
+    parm = load_param_file(param)
+    if coords is not None:
+        parm.initialize_topology(xyz=coords)
+    sele = []
+    if len(maskstr) > 0:
+        newmaskstr = maskstr.replace("@0", "!@*")
+        sele = [parm.atoms[i].idx for i in parmed.amber.mask.AmberMask(
+            parm, newmaskstr).Selected()]
+    return sele
+
+def load_param_file(param_file):
+    if isinstance(param_file, str):
+        return parmed.load_file(param_file)
+    elif isinstance(param_file, parmed.Structure):
+        return param_file
+    else:
+        raise RuntimeError("Unsupported structure")

dpdata/amber/md.py

@@ -1,28 +1,47 @@
 import re
+import os
 from scipy.io import netcdf
 import numpy as np
+from dpdata.amber.mask import pick_by_amber_mask
 
 kcalmol2eV= 0.04336410390059322
+symbols = ['X', 'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', 'Np', 'Pu', 'Am', 'Cm', 'Bk', 'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr', 'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc', 'Lv', 'Ts', 'Og']
 
 energy_convert = kcalmol2eV
 force_convert = energy_convert
 
 
-def read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file):
+def read_amber_traj(parm7_file, nc_file, mdfrc_file=None, mden_file = None, mdout_file = None,
+        use_element_symbols=None, labeled=True,
+    ):
     """The amber trajectory includes:
     * nc, NetCDF format, stores coordinates
     * mdfrc, NetCDF format, stores forces
-    * mden, text format, stores energies
+    * mden (optional), text format, stores energies
+    * mdout (optional), text format, may store energies if there is no mden_file
     * parm7, text format, stores types
+
+    Parameters
+    ----------
+    parm7_file, nc_file, mdfrc_file, mden_file, mdout_file:
+      filenames
+    use_element_symbols: None or list or str
+      If use_element_symbols is a list of atom indexes, these atoms will use element symbols
+      instead of amber types. For example, a ligand will use C, H, O, N, and so on
+      instead of h1, hc, o, os, and so on.
+      IF use_element_symbols is str, it will be considered as Amber mask.
     """
 
-    flag=False
+    flag_atom_type = False
+    flag_atom_numb = False
     amber_types = []
+    atomic_number = []
     with open(parm7_file) as f:
         for line in f:
             if line.startswith("%FLAG"):
-                flag = line.startswith("%FLAG AMBER_ATOM_TYPE")
-            elif flag:
+                flag_atom_type = line.startswith("%FLAG AMBER_ATOM_TYPE")
+                flag_atom_numb = (use_element_symbols is not None) and line.startswith("%FLAG ATOMIC_NUMBER")
+            elif flag_atom_type or flag_atom_numb:
                 if line.startswith("%FORMAT"):
                     fmt = re.findall(r'\d+', line)
                     fmt0 = int(fmt[0])
@@ -33,7 +52,16 @@ def read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file):
                         end_index = (ii + 1) * fmt1
                         if end_index >= len(line):
                             continue
-                        amber_types.append(line[start_index:end_index].strip())
+                        content = line[start_index:end_index].strip()
+                        if flag_atom_type:
+                            amber_types.append(content)
+                        elif flag_atom_numb:
+                            atomic_number.append(int(content))    
+    if use_element_symbols is not None:
+        if isinstance(use_element_symbols, str):
+            use_element_symbols = pick_by_amber_mask(parm7_file, use_element_symbols)
+        for ii in use_element_symbols:
+            amber_types[ii] = symbols[atomic_number[ii]]
 
     with netcdf.netcdf_file(nc_file, 'r') as f:
         coords = np.array(f.variables["coordinates"][:])
@@ -49,26 +77,37 @@ def read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file):
         else:
             raise RuntimeError("Unsupported cells")
 
-    with netcdf.netcdf_file(mdfrc_file, 'r') as f:
-        forces = np.array(f.variables["forces"][:])
+    if labeled:
+        with netcdf.netcdf_file(mdfrc_file, 'r') as f:
+            forces = np.array(f.variables["forces"][:])
 
-    # energy
-    energies = []
-    with open(mden_file) as f:
-        for line in f:
-            if line.startswith("L6"):
-                s = line.split()
-                if s[2] != "E_pot":
-                    energies.append(float(s[2]))
+        # load energy from mden_file or mdout_file
+        energies = []
+        if mden_file is not None and os.path.isfile(mden_file):
+            with open(mden_file) as f:
+                for line in f:
+                    if line.startswith("L6"):
+                        s = line.split()
+                        if s[2] != "E_pot":
+                            energies.append(float(s[2]))
+        elif mdout_file is not None and os.path.isfile(mdout_file):
+            with open(mdout_file) as f:
+                for line in f:
+                    if "EPtot" in line:
+                        s = line.split()
+                        energies.append(float(s[-1]))
+        else:
+            raise RuntimeError("Please provide one of mden_file and mdout_file")
 
     atom_names, atom_types, atom_numbs = np.unique(amber_types, return_inverse=True, return_counts=True)
 
     data = {}
     data['atom_names'] = list(atom_names)
     data['atom_numbs'] = list(atom_numbs)
     data['atom_types'] = atom_types
-    data['forces'] = forces * force_convert
-    data['energies'] = np.array(energies) * energy_convert
+    if labeled:
+        data['forces'] = forces * force_convert
+        data['energies'] = np.array(energies) * energy_convert
     data['coords'] = coords
     data['cells'] = cells
     data['orig'] = np.array([0, 0, 0])

dpdata/system.py

@@ -28,6 +28,7 @@
 from monty.serialization import loadfn,dumpfn
 from dpdata.periodic_table import Element
 from dpdata.xyz.quip_gap_xyz import QuipGapxyzSystems
+from dpdata.amber.mask import pick_by_amber_mask, load_param_file
 
 
 class Register:
@@ -927,6 +928,10 @@ def nopbc(self):
             return True
         return False
 
+    @nopbc.setter
+    def nopbc(self, value):
+        self.data['nopbc'] = value
+
     def shuffle(self):
         """Shuffle frames randomly."""
         idx = np.random.permutation(self.get_nframes())
@@ -973,6 +978,93 @@ def predict(self, dp):
             labeled_sys.append(this_sys)
         return labeled_sys
 
+    def pick_atom_idx(self, idx, nopbc=None):
+        """Pick atom index
+        
+        Parameters
+        ----------
+        idx: int or list or slice
+            atom index
+        nopbc: Boolen (default: None)
+            If nopbc is True or False, set nopbc
+
+        Returns
+        -------
+        new_sys: System
+            new system
+        """
+        new_sys = self.copy()
+        new_sys.data['coords'] = self.data['coords'][:, idx, :]
+        new_sys.data['atom_types'] = self.data['atom_types'][idx]
+        # recalculate atom_numbs according to atom_types
+        atom_numbs = np.bincount(new_sys.data['atom_types'], minlength=len(self.get_atom_names()))
+        new_sys.data['atom_numbs'] = list(atom_numbs)
+        if nopbc is True or nopbc is False:
+            new_sys.nopbc = nopbc
+        return new_sys
+
+    def remove_atom_names(self, atom_names):
+        """Remove atom names and all such atoms.
+        For example, you may not remove EP atoms in TIP4P/Ew water, which
+        is not a real atom. 
+        """
+        if isinstance(atom_names, str):
+            atom_names = [atom_names]
+        removed_atom_idx = []
+        for an in atom_names:
+            # get atom name idx
+            idx = self.data['atom_names'].index(an)
+            atom_idx = self.data['atom_types'] == idx
+            removed_atom_idx.append(atom_idx)
+        picked_atom_idx = ~np.any(removed_atom_idx, axis=0)
+        new_sys = self.pick_atom_idx(picked_atom_idx)
+        # let's remove atom_names
+        # firstly, rearrange atom_names and put these atom_names in the end
+        new_atom_names = list([xx for xx in new_sys.data['atom_names'] if xx not in atom_names])
+        new_sys.sort_atom_names(type_map=new_atom_names + atom_names)
+        # remove atom_names and atom_numbs
+        new_sys.data['atom_names'] = new_atom_names
+        new_sys.data['atom_numbs'] = new_sys.data['atom_numbs'][:len(new_atom_names)]
+        return new_sys
+
+    def pick_by_amber_mask(self, param, maskstr, pass_coords=False, nopbc=None):
+        """Pick atoms by amber mask
+        
+        Parameters
+        ----------
+        param: str or parmed.Structure
+          filename of Amber param file or parmed.Structure
+        maskstr: str
+          Amber masks
+        pass_coords: Boolen (default: False)
+            If pass_coords is true, the function will pass coordinates and 
+            return a MultiSystem. Otherwise, the result is
+            coordinate-independent, and the function will return System or
+            LabeledSystem.
+        nopbc: Boolen (default: None)
+            If nopbc is True or False, set nopbc
+        """
+        parm = load_param_file(param)
+        if pass_coords:
+            ms = MultiSystems()
+            for sub_s in self:
+                # TODO: this can computed in pararrel
+                idx = pick_by_amber_mask(parm, maskstr, sub_s['coords'][0])
+                ms.append(sub_s.pick_atom_idx(idx, nopbc=nopbc))
+            return ms
+        else:
+            idx = pick_by_amber_mask(parm, maskstr)
+            return self.pick_atom_idx(idx, nopbc=nopbc)
+
+    @register_from_funcs.register_funcs('amber/md')
+    def from_amber_md(self, file_name=None, parm7_file=None, nc_file=None, use_element_symbols=None):
+        # assume the prefix is the same if the spefic name is not given
+        if parm7_file is None:
+            parm7_file = file_name + ".parm7"
+        if nc_file is None:
+            nc_file = file_name + ".nc"
+        self.data = dpdata.amber.md.read_amber_traj(parm7_file=parm7_file, nc_file=nc_file, use_element_symbols=use_element_symbols, labeled=False)
+
 def get_cell_perturb_matrix(cell_pert_fraction):
     if cell_pert_fraction<0:
         raise RuntimeError('cell_pert_fraction can not be negative')
@@ -1305,7 +1397,7 @@ def from_gaussian_md(self, file_name):
         self.from_gaussian_log(file_name, md=True)
 
     @register_from_funcs.register_funcs('amber/md')
-    def from_amber_md(self, file_name=None, parm7_file=None, nc_file=None, mdfrc_file=None, mden_file=None):
+    def from_amber_md(self, file_name=None, parm7_file=None, nc_file=None, mdfrc_file=None, mden_file=None, mdout_file=None, use_element_symbols=None):
         # assume the prefix is the same if the spefic name is not given
         if parm7_file is None:
             parm7_file = file_name + ".parm7"
@@ -1315,7 +1407,9 @@ def from_amber_md(self, file_name=None, parm7_file=None, nc_file=None, mdfrc_fil
             mdfrc_file = file_name + ".mdfrc"
         if mden_file is None:
             mden_file = file_name + ".mden"
-        self.data = dpdata.amber.md.read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file)
+        if mdout_file is None:
+            mdout_file = file_name + ".mdout"
+        self.data = dpdata.amber.md.read_amber_traj(parm7_file, nc_file, mdfrc_file, mden_file, mdout_file, use_element_symbols)
 
     @register_from_funcs.register_funcs('cp2k/output')
     def from_cp2k_output(self, file_name) :
@@ -1475,6 +1569,53 @@ def to_pymatgen_ComputedStructureEntry(self):
             entries.append(entry)
         return entries
 
+    def correction(self, hl_sys):
+        """Get energy and force correction between self and a high-level LabeledSystem.
+        The self's coordinates will be kept, but energy and forces will be replaced by
+        the correction between these two systems.
+
+        Note: The function will not check whether coordinates and elements of two systems
+        are the same. The user should make sure by itself.
+
+        Parameters
+        ----------
+        hl_sys: LabeledSystem
+            high-level LabeledSystem
+        Returns
+        ----------
+        corrected_sys: LabeledSystem
+            Corrected LabeledSystem
+        """
+        if not isinstance(hl_sys, LabeledSystem):
+            raise RuntimeError("high_sys should be LabeledSystem")
+        corrected_sys = self.copy()
+        corrected_sys.data['energies'] = hl_sys.data['energies'] - self.data['energies']
+        corrected_sys.data['forces'] = hl_sys.data['forces'] - self.data['forces']
+        if 'virials' in self.data and 'virials' in hl_sys.data:
+            corrected_sys.data['virials'] = hl_sys.data['virials'] - self.data['virials']
+        return corrected_sys
+
+    def pick_atom_idx(self, idx, nopbc=None):
+        """Pick atom index
+        
+        Parameters
+        ----------
+        idx: int or list or slice
+            atom index
+        nopbc: Boolen (default: None)
+            If nopbc is True or False, set nopbc
+
+        Returns
+        -------
+        new_sys: LabeledSystem
+            new system
+        """
+        new_sys = System.pick_atom_idx(self, idx, nopbc=nopbc)
+        # forces
+        new_sys.data['forces'] = self.data['forces'][:, idx, :]
+        return new_sys
+
+
 class MultiSystems:
     '''A set containing several systems.'''
 
@@ -1650,6 +1791,26 @@ def predict(self, dp):
         for ss in self:
             new_multisystems.append(ss.predict(dp))
         return new_multisystems
+    
+    def pick_atom_idx(self, idx, nopbc=None):
+        """Pick atom index
+        
+        Parameters
+        ----------
+        idx: int or list or slice
+            atom index
+        nopbc: Boolen (default: None)
+            If nopbc is True or False, set nopbc
+
+        Returns
+        -------
+        new_sys: MultiSystems
+            new system
+        """
+        new_sys = MultiSystems()
+        for ss in self:
+            new_sys.append(ss.pick_atom_idx(idx, nopbc=nopbc))
+        return new_sys
 
 
 def check_System(data):

setup.py

@@ -46,5 +46,9 @@
     ],
     keywords='lammps vasp deepmd-kit',
     install_requires=install_requires,
+    extras_require={
+        'ase': ['ase'],
+        'amber': ['parmed'],
+    }
 )