add pairwise DPRc

deepmd/descriptor/descriptor.py

@@ -324,48 +324,6 @@ def prod_force_virial(
             The atomic virial
         """
 
-    def get_feed_dict(
-        self,
-        coord_: tf.Tensor,
-        atype_: tf.Tensor,
-        natoms: tf.Tensor,
-        box: tf.Tensor,
-        mesh: tf.Tensor,
-    ) -> Dict[str, tf.Tensor]:
-        """Generate the feed_dict for current descriptor.
-
-        Parameters
-        ----------
-        coord_ : tf.Tensor
-            The coordinate of atoms
-        atype_ : tf.Tensor
-            The type of atoms
-        natoms : tf.Tensor
-            The number of atoms. This tensor has the length of Ntypes + 2
-            natoms[0]: number of local atoms
-            natoms[1]: total number of atoms held by this processor
-            natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
-        box : tf.Tensor
-            The box. Can be generated by deepmd.model.make_stat_input
-        mesh : tf.Tensor
-            For historical reasons, only the length of the Tensor matters.
-            if size of mesh == 6, pbc is assumed.
-            if size of mesh == 0, no-pbc is assumed.
-
-        Returns
-        -------
-        feed_dict : dict[str, tf.Tensor]
-            The output feed_dict of current descriptor
-        """
-        feed_dict = {
-            "t_coord:0": coord_,
-            "t_type:0": atype_,
-            "t_natoms:0": natoms,
-            "t_box:0": box,
-            "t_mesh:0": mesh,
-        }
-        return feed_dict
-
     def init_variables(
         self,
         graph: tf.Graph,

deepmd/entrypoints/freeze.py

@@ -211,6 +211,7 @@ def _make_node_names(
         "model_attr/model_version",
         "train_attr/min_nbor_dist",
         "train_attr/training_script",
+        "t_mesh",
     ]
 
     if model_type == "ener":

deepmd/entrypoints/train.py

@@ -354,6 +354,11 @@ def get_modifier(modi_data=None):
 
 
 def get_rcut(jdata):
+    if jdata["model"].get("type") == "pairwise_dprc":
+        return max(
+            jdata["model"]["qm_model"]["descriptor"]["rcut"],
+            jdata["model"]["qmmm_model"]["descriptor"]["rcut"],
+        )
     descrpt_data = jdata["model"]["descriptor"]
     rcut_list = []
     if descrpt_data["type"] == "hybrid":
@@ -499,6 +504,11 @@ def update_sel(jdata):
     log.info(
         "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)"
     )
+    if jdata["model"].get("type") == "pairwise_dprc":
+        # do not update sel; only find min distance
+        rcut = get_rcut(jdata)
+        get_min_nbor_dist(jdata, rcut)
+        return jdata
     descrpt_data = jdata["model"]["descriptor"]
     if descrpt_data["type"] == "hybrid":
         for ii in range(len(descrpt_data["list"])):

deepmd/model/ener.py

@@ -198,7 +198,7 @@ def build(
         input_dict["nframes"] = tf.shape(coord)[0]
 
         # type embedding if any
-        if self.typeebd is not None:
+        if self.typeebd is not None and "type_embedding" not in input_dict:
             type_embedding = self.typeebd.build(
                 self.ntypes,
                 reuse=reuse,
@@ -368,7 +368,10 @@ def init_variables(
             tf.constant("compressed_model", name="model_type", dtype=tf.string)
         else:
             raise RuntimeError("Unknown model type %s" % model_type)
-        if self.typeebd is not None:
+        if (
+            self.typeebd is not None
+            and self.typeebd.type_embedding_net_variables is None
+        ):
             self.typeebd.init_variables(
                 graph, graph_def, suffix=suffix, model_type=model_type
             )

deepmd/model/model.py

@@ -7,6 +7,7 @@
     Enum,
 )
 from typing import (
+    Dict,
     List,
     Optional,
     Union,
@@ -82,12 +83,17 @@ def __new__(cls, *args, **kwargs):
             from deepmd.model.multi import (
                 MultiModel,
             )
+            from deepmd.model.pairwise_dprc import (
+                PairwiseDPRc,
+            )
 
             model_type = kwargs.get("type", "standard")
             if model_type == "standard":
                 cls = StandardModel
             elif model_type == "multi":
                 cls = MultiModel
+            elif model_type == "pairwise_dprc":
+                cls = PairwiseDPRc
             else:
                 raise ValueError(f"unknown model type: {model_type}")
             return cls.__new__(cls, *args, **kwargs)
@@ -261,14 +267,31 @@ def build_descrpt(
                 suffix=suffix,
                 reuse=reuse,
             )
-            dout = tf.identity(dout, name="o_descriptor")
+            dout = tf.identity(dout, name="o_descriptor" + suffix)
         else:
             tf.constant(
-                self.rcut, name="descrpt_attr/rcut", dtype=GLOBAL_TF_FLOAT_PRECISION
+                self.rcut,
+                name="descrpt_attr%s/rcut" % suffix,
+                dtype=GLOBAL_TF_FLOAT_PRECISION,
+            )
+            tf.constant(
+                self.ntypes, name="descrpt_attr%s/ntypes" % suffix, dtype=tf.int32
             )
-            tf.constant(self.ntypes, name="descrpt_attr/ntypes", dtype=tf.int32)
-            feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
-            return_elements = [*self.descrpt.get_tensor_names(), "o_descriptor:0"]
+            if "global_feed_dict" in input_dict:
+                feed_dict = input_dict["global_feed_dict"]
+            else:
+                extra_feed_dict = {}
+                if "fparam" in input_dict:
+                    extra_feed_dict["fparam"] = input_dict["fparam"]
+                if "aparam" in input_dict:
+                    extra_feed_dict["aparam"] = input_dict["aparam"]
+                feed_dict = self.get_feed_dict(
+                    coord_, atype_, natoms, box, mesh, **extra_feed_dict
+                )
+            return_elements = [
+                *self.descrpt.get_tensor_names(suffix=suffix),
+                "o_descriptor%s:0" % suffix,
+            ]
             if frz_model is not None:
                 imported_tensors = self._import_graph_def_from_frz_model(
                     frz_model, feed_dict, return_elements
@@ -343,8 +366,14 @@ def change_energy_bias(
         """
         raise RuntimeError("Not supported")
 
-    def enable_compression(self):
-        """Enable compression."""
+    def enable_compression(self, suffix: str = ""):
+        """Enable compression.
+
+        Parameters
+        ----------
+        suffix : str
+            suffix to name scope
+        """
         raise RuntimeError("Not supported")
 
     def get_numb_fparam(self) -> Union[int, dict]:
@@ -379,6 +408,55 @@ def get_ntypes(self) -> int:
     def data_stat(self, data: dict):
         """Data staticis."""
 
+    def get_feed_dict(
+        self,
+        coord_: tf.Tensor,
+        atype_: tf.Tensor,
+        natoms: tf.Tensor,
+        box: tf.Tensor,
+        mesh: tf.Tensor,
+        **kwargs,
+    ) -> Dict[str, tf.Tensor]:
+        """Generate the feed_dict for current descriptor.
+
+        Parameters
+        ----------
+        coord_ : tf.Tensor
+            The coordinate of atoms
+        atype_ : tf.Tensor
+            The type of atoms
+        natoms : tf.Tensor
+            The number of atoms. This tensor has the length of Ntypes + 2
+            natoms[0]: number of local atoms
+            natoms[1]: total number of atoms held by this processor
+            natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
+        box : tf.Tensor
+            The box. Can be generated by deepmd.model.make_stat_input
+        mesh : tf.Tensor
+            For historical reasons, only the length of the Tensor matters.
+            if size of mesh == 6, pbc is assumed.
+            if size of mesh == 0, no-pbc is assumed.
+        **kwargs : dict
+            The additional arguments
+
+        Returns
+        -------
+        feed_dict : dict[str, tf.Tensor]
+            The output feed_dict of current descriptor
+        """
+        feed_dict = {
+            "t_coord:0": coord_,
+            "t_type:0": atype_,
+            "t_natoms:0": natoms,
+            "t_box:0": box,
+            "t_mesh:0": mesh,
+        }
+        if kwargs.get("fparam") is not None:
+            feed_dict["t_fparam:0"] = kwargs["fparam"]
+        if kwargs.get("aparam") is not None:
+            feed_dict["t_aparam:0"] = kwargs["aparam"]
+        return feed_dict
+
 
 class StandardModel(Model):
     """Standard model, which must contain a descriptor and a fitting.
@@ -479,8 +557,14 @@ def enable_mixed_precision(self, mixed_prec: dict):
         self.descrpt.enable_mixed_precision(mixed_prec)
         self.fitting.enable_mixed_precision(mixed_prec)
 
-    def enable_compression(self):
-        """Enable compression."""
+    def enable_compression(self, suffix: str = ""):
+        """Enable compression.
+
+        Parameters
+        ----------
+        suffix : str
+            suffix to name scope
+        """
         graph, graph_def = load_graph_def(self.compress["model_file"])
         self.descrpt.enable_compression(
             self.compress["min_nbor_dist"],
@@ -490,11 +574,15 @@ def enable_compression(self):
             self.compress["table_config"][1],
             self.compress["table_config"][2],
             self.compress["table_config"][3],
+            suffix=suffix,
         )
         # for fparam or aparam settings in 'ener' type fitting net
-        self.fitting.init_variables(graph, graph_def)
-        if self.typeebd is not None:
-            self.typeebd.init_variables(graph, graph_def)
+        self.fitting.init_variables(graph, graph_def, suffix=suffix)
+        if (
+            self.typeebd is not None
+            and self.typeebd.type_embedding_net_variables is None
+        ):
+            self.typeebd.init_variables(graph, graph_def, suffix=suffix)
 
     def get_fitting(self) -> Union[Fitting, dict]:
         """Get the fitting(s)."""