Merge pull request #1177 from deepmodeling/devel

Merge recent development on devel into master

.github/workflows/mirror_gitee.yml

@@ -8,11 +8,12 @@ concurrency:
 
 jobs:
   git-mirror:
+    if: github.repository_owner == 'deepmodeling'
     runs-on: ubuntu-latest
     steps:
       - uses: wearerequired/git-mirror-action@v1
         env:
           SSH_PRIVATE_KEY: ${{ secrets.SYNC_GITEE_PRIVATE_KEY }}
         with:
-          source-repo: "git@github.com:deepmodeling/deepmd-kit.git"
+          source-repo: "https://github.com/deepmodeling/deepmd-kit.git"
           destination-repo: "[email protected]:deepmodeling/deepmd-kit.git"

.gitignore

@@ -30,5 +30,6 @@ API_CC
 doc/api_py/
 dp/
 build_lammps/
+.idea/
 build_tests/
 build_cc_tests

.readthedocs.yml

@@ -0,0 +1,3 @@
+version: 2
+conda:
+  environment: doc/environment.yml

README.md

@@ -60,7 +60,7 @@ Please follow our [GitHub](https://github.com/deepmodeling/deepmd-kit) webpage t
 
 DeePMD-kit offers multiple installation methods. It is recommend using easily methods like [offline packages](doc/install/easy-install.md#offline-packages), [conda](doc/install/easy-install.md#with-conda) and [docker](doc/install/easy-install.md#with-docker). 
 
-One may manually install DeePMD-kit by following the instuctions on [installing the Python interface](doc/install/install-from-source.md#install-the-python-interface) and [installing the C++ interface](doc/install/install-from-source.md#install-the-c-interface). The C++ interface is necessary when using DeePMD-kit with LAMMPS and i-PI.
+One may manually install DeePMD-kit by following the instuctions on [installing the Python interface](doc/install/install-from-source.md#install-the-python-interface) and [installing the C++ interface](doc/install/install-from-source.md#install-the-c-interface). The C++ interface is necessary when using DeePMD-kit with LAMMPS, i-PI or GROMACS.
 
 
 # Use DeePMD-kit
@@ -71,7 +71,7 @@ A quick-start on using DeePMD-kit can be found as follows:
 - [Training a model](doc/train/training.md)
 - [Freeze a model](doc/freeze/freeze.md)
 - [Test a model](doc/test/test.md)
-- [Running MD with LAMMPS](doc/third-party/lammps.md)
+- [Run MD with LAMMPS](doc/third-party/lammps.md)
 
 A full [document](doc/train/train-input-auto.rst) on options in the training input script is available.
 
@@ -82,6 +82,7 @@ A full [document](doc/train/train-input-auto.rst) on options in the training inp
     - [Install from source code](doc/install/install-from-source.md)
     - [Install LAMMPS](doc/install/install-lammps.md)
     - [Install i-PI](doc/install/install-ipi.md)
+    - [Install GROMACS](doc/install/install-gromacs.md)
     - [Building conda packages](doc/install/build-conda.md)
 - [Data](doc/data/index.md)
     - [Data conversion](doc/data/data-conv.md)
@@ -113,10 +114,10 @@ A full [document](doc/train/train-input-auto.rst) on options in the training inp
     - [C++ interface](doc/inference/cxx.md)
 - [Integrate with third-party packages](doc/third-party/index.rst)
     - [Use deep potential with ASE](doc/third-party/ase.md)
-    - [Running MD with LAMMPS](doc/third-party/lammps.md)
+    - [Run MD with LAMMPS](doc/third-party/lammps.md)
     - [LAMMPS commands](doc/third-party/lammps-command.md)
     - [Run path-integral MD with i-PI](doc/third-party/ipi.md)
-
+    - [Run MD with GROMACS](doc/third-party/gromacs.md)
 
 # Code structure
 The code is organized as follows:
@@ -135,6 +136,8 @@ The code is organized as follows:
 
 * `source/lmp`: source code of Lammps module.
 
+* `source/gmx`: source code of Gromacs plugin.
+
 * `source/op`: tensorflow op implementation. working with library.
 
 

deepmd/__init__.py

@@ -1,5 +1,9 @@
 """Root of the deepmd package, exposes all public classes and submodules."""
 
+try:
+    from importlib import metadata
+except ImportError:  # for Python<3.8
+    import importlib_metadata as metadata
 import deepmd.utils.network as network
 
 from . import cluster, descriptor, fit, loss, utils
@@ -14,6 +18,10 @@
 except ImportError:
     from .__about__ import __version__
 
+# load third-party plugins
+for ep in metadata.entry_points().get('deepmd', []):
+    ep.load()
+
 __all__ = [
     "descriptor",
     "fit",

deepmd/common.py

@@ -23,6 +23,7 @@
 from deepmd.env import GLOBAL_TF_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION
 from deepmd.utils.sess import run_sess
 from deepmd.utils.errors import GraphWithoutTensorError
+from deepmd.utils.path import DPPath
 
 if TYPE_CHECKING:
     _DICT_VAL = TypeVar("_DICT_VAL")
@@ -429,9 +430,10 @@ def expand_sys_str(root_dir: Union[str, Path]) -> List[str]:
     List[str]
         list of string pointing to system directories
     """
-    matches = [str(d) for d in Path(root_dir).rglob("*") if (d / "type.raw").is_file()]
-    if (Path(root_dir) / "type.raw").is_file():
-        matches += [root_dir]
+    root_dir = DPPath(root_dir)
+    matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()]
+    if (root_dir / "type.raw").is_file():
+        matches.append(str(root_dir))
     return matches
 
 

deepmd/descriptor/__init__.py

@@ -1,7 +1,7 @@
+from .descriptor import Descriptor
 from .hybrid import DescrptHybrid
 from .se_a import DescrptSeA
 from .se_r import DescrptSeR
-from .se_ar import DescrptSeAR
 from .se_t import DescrptSeT
 from .se_a_ebd import DescrptSeAEbd
 from .se_a_ef import DescrptSeAEf

deepmd/descriptor/descriptor.py

@@ -3,21 +3,64 @@
 
 import numpy as np
 from deepmd.env import tf
+from deepmd.utils import Plugin, PluginVariant
 
 
-class Descriptor(ABC):
+class Descriptor(PluginVariant):
     r"""The abstract class for descriptors. All specific descriptors should
     be based on this class.
 
     The descriptor :math:`\mathcal{D}` describes the environment of an atom,
     which should be a function of coordinates and types of its neighbour atoms.
 
+    Examples
+    --------
+    >>> descript = Descriptor(type="se_e2_a", rcut=6., rcut_smth=0.5, sel=[50])
+    >>> type(descript)
+    <class 'deepmd.descriptor.se_a.DescrptSeA'>
+
     Notes
     -----
     Only methods and attributes defined in this class are generally public,
     that can be called by other classes.
     """
 
+    __plugins = Plugin()
+
+    @staticmethod
+    def register(key: str) -> "Descriptor":
+        """Regiester a descriptor plugin.
+
+        Parameters
+        ----------
+        key : str
+            the key of a descriptor
+
+        Returns
+        -------
+        Descriptor
+            the regiestered descriptor
+
+        Examples
+        --------
+        >>> @Descriptor.register("some_descrpt")
+            class SomeDescript(Descriptor):
+                pass
+        """
+        return Descriptor.__plugins.register(key)
+
+    def __new__(cls, *args, **kwargs):
+        if cls is Descriptor:
+            try:
+                descrpt_type = kwargs['type']
+            except KeyError:
+                raise KeyError('the type of descriptor should be set by `type`')
+            if descrpt_type in Descriptor.__plugins.plugins:
+                cls = Descriptor.__plugins.plugins[descrpt_type]
+            else:
+                raise RuntimeError('Unknown descriptor type: ' + descrpt_type)
+        return super().__new__(cls)
+
     @abstractmethod
     def get_rcut(self) -> float:
         """

deepmd/descriptor/hybrid.py

@@ -15,12 +15,12 @@
 from .descriptor import Descriptor
 from .se_a import DescrptSeA
 from .se_r import DescrptSeR
-from .se_ar import DescrptSeAR
 from .se_t import DescrptSeT
 from .se_a_ebd import DescrptSeAEbd
 from .se_a_ef import DescrptSeAEf
 from .loc_frame import DescrptLocFrame
 
+@Descriptor.register("hybrid")
 class DescrptHybrid (Descriptor):
     """Concate a list of descriptors to form a new descriptor.
 
@@ -37,11 +37,19 @@ def __init__ (self,
         """
         if descrpt_list == [] or descrpt_list is None:
             raise RuntimeError('cannot build descriptor from an empty list of descriptors.')
+        formatted_descript_list = []
+        for ii in descrpt_list:
+            if isinstance(ii, Descriptor):
+                formatted_descript_list.append(ii)
+            elif isinstance(ii, dict):
+                formatted_descript_list.append(Descriptor(**ii))
+            else:
+                raise NotImplementedError
         # args = ClassArg()\
         #        .add('list', list, must = True)
         # class_data = args.parse(jdata)
         # dict_list = class_data['list']
-        self.descrpt_list = descrpt_list
+        self.descrpt_list = formatted_descript_list
         self.numb_descrpt = len(self.descrpt_list)
         for ii in range(1, self.numb_descrpt):
             assert(self.descrpt_list[ii].get_ntypes() == 

deepmd/descriptor/se.py

@@ -0,0 +1,108 @@
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.utils.graph import get_embedding_net_variables
+from .descriptor import Descriptor
+
+
+class DescrptSe (Descriptor):
+    """A base class for smooth version of descriptors.
+    
+    Notes
+    -----
+    All of these descriptors have an environmental matrix and an
+    embedding network (:meth:`deepmd.utils.network.embedding_net`), so
+    they can share some similiar methods without defining them twice.
+
+    Attributes
+    ----------
+    embedding_net_variables : dict
+        initial embedding network variables
+    descrpt_reshape : tf.Tensor
+        the reshaped descriptor
+    descrpt_deriv : tf.Tensor
+        the descriptor derivative
+    rij : tf.Tensor
+        distances between two atoms
+    nlist : tf.Tensor
+        the neighbor list
+    
+    """
+    def _identity_tensors(self, suffix : str = "") -> None:
+        """Identify tensors which are expected to be stored and restored.
+        
+        Notes
+        -----
+        These tensors will be indentitied:
+            self.descrpt_reshape : o_rmat
+            self.descrpt_deriv : o_rmat_deriv
+            self.rij : o_rij
+            self.nlist : o_nlist
+        Thus, this method should be called during building the descriptor and
+        after these tensors are initialized.
+
+        Parameters
+        ----------
+        suffix : str
+            The suffix of the scope
+        """
+        self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat' + suffix)
+        self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv' + suffix)
+        self.rij = tf.identity(self.rij, name = 'o_rij' + suffix)
+        self.nlist = tf.identity(self.nlist, name = 'o_nlist' + suffix)
+
+    def get_tensor_names(self, suffix : str = "") -> Tuple[str]:
+        """Get names of tensors.
+        
+        Parameters
+        ----------
+        suffix : str
+            The suffix of the scope
+
+        Returns
+        -------
+        Tuple[str]
+            Names of tensors
+        """
+        return (f'o_rmat{suffix}:0', f'o_rmat_deriv{suffix}:0', f'o_rij{suffix}:0', f'o_nlist{suffix}:0')
+
+    def pass_tensors_from_frz_model(self,
+                                    descrpt_reshape : tf.Tensor,
+                                    descrpt_deriv   : tf.Tensor,
+                                    rij             : tf.Tensor,
+                                    nlist           : tf.Tensor
+    ):
+        """
+        Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def
+
+        Parameters
+        ----------
+        descrpt_reshape
+                The passed descrpt_reshape tensor
+        descrpt_deriv
+                The passed descrpt_deriv tensor
+        rij
+                The passed rij tensor
+        nlist
+                The passed nlist tensor
+        """
+        self.rij = rij
+        self.nlist = nlist
+        self.descrpt_deriv = descrpt_deriv
+        self.descrpt_reshape = descrpt_reshape
+
+    def init_variables(self,
+                       model_file : str,
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given dict
+
+        Parameters
+        ----------
+        model_file : str
+            The input frozen model file
+        suffix : str, optional
+            The suffix of the scope
+        """
+        self.embedding_net_variables = get_embedding_net_variables(model_file, suffix = suffix)