nvnmd: init-model feature and 256 neighbors (#3058)

1. Add init-modele feature. you can initialize the CNN model by:
`dp train-nvnmd train_cnn.json -s s1 -i nvnmd_cnn_bck/model.ckpt `.
2. The maximum number of neighbors has been increased to 256. 
3. Update the doc of nvnmd.

---------

Co-authored-by: LiuGroupHNU <liujie123@HNU>
Co-authored-by: MoPinghui <[email protected]>
Co-authored-by: Han Wang <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

deepmd/fit/ener.py

@@ -514,6 +514,11 @@ def build(
                     self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
             self.bias_atom_e = self.bias_atom_e[:ntypes_atom]
 
+        if nvnmd_cfg.enable:
+            # fix the bug: CNN and QNN have different t_bias_atom_e.
+            if "t_bias_atom_e" in nvnmd_cfg.weight.keys():
+                self.bias_atom_e = nvnmd_cfg.weight["t_bias_atom_e"]
+
         with tf.variable_scope("fitting_attr" + suffix, reuse=reuse):
             t_dfparam = tf.constant(self.numb_fparam, name="dfparam", dtype=tf.int32)
             t_daparam = tf.constant(self.numb_aparam, name="daparam", dtype=tf.int32)

deepmd/nvnmd/data/data.py

@@ -60,13 +60,15 @@
     },
     "ctrl": {
         # NSTDM
+        "MAX_NNEI": 128,
         "NSTDM": 64,
         "NSTDM_M1": 32,
         "NSTDM_M2": 2,
         "NSTDM_M1X": 8,
         "NSEL": "NSTDM*NTYPE_MAX",
         "NSADV": "NSTDM+1",
         "VERSION": 0,
+        "SUB_VERSION": 1,
     },
     "nbit": {
         # general
@@ -116,6 +118,22 @@
     "end": "",
 }
 
+# change the configuration accordng to the max_nnei
+jdata_config_v0_ni128 = jdata_config_v0.copy()
+jdata_config_v0_ni256 = jdata_config_v0.copy()
+jdata_config_v0_ni256["ctrl"] = {
+    "MAX_NNEI": 256,
+    "NSTDM": 128,
+    "NSTDM_M1": 32,
+    "NSTDM_M2": 4,
+    "NSTDM_M1X": 8,
+    "NSEL": "NSTDM*NTYPE_MAX",
+    "NSADV": "NSTDM+1",
+    "VERSION": 0,
+    "SUB_VERSION": 1,
+}
+jdata_config_v0_ni256["nbit"]["NBIT_NEIB"] = 9
+
 jdata_config_v1 = {
     "dscp": {
         # basic config from deepmd model
@@ -174,13 +192,15 @@
     },
     "ctrl": {
         # NSTDM
+        "MAX_NNEI": 128,
         "NSTDM": 64,
         "NSTDM_M1": 32,
         "NSTDM_M2": 2,
         "NSTDM_M1X": 8,
         "NSEL": "NSTDM",
         "NSADV": "NSTDM+1",
         "VERSION": 1,
+        "SUB_VERSION": 1,
     },
     "nbit": {
         # general
@@ -230,6 +250,22 @@
     "end": "",
 }
 
+# change the configuration accordng to the max_nnei
+jdata_config_v1_ni128 = jdata_config_v1.copy()
+jdata_config_v1_ni256 = jdata_config_v1.copy()
+jdata_config_v1_ni256["ctrl"] = {
+    "MAX_NNEI": 256,
+    "NSTDM": 128,
+    "NSTDM_M1": 32,
+    "NSTDM_M2": 4,
+    "NSTDM_M1X": 8,
+    "NSEL": "NSTDM",
+    "NSADV": "NSTDM+1",
+    "VERSION": 1,
+    "SUB_VERSION": 1,
+}
+jdata_config_v1_ni256["nbit"]["NBIT_NEIB"] = 9
+
 jdata_deepmd_input_v0 = {
     "model": {
         "descriptor": {
@@ -247,6 +283,7 @@
     },
     "nvnmd": {
         "version": 0,
+        "max_nnei": 128,  # 128 or 256
         "net_size": 128,
         "config_file": "none",
         "weight_file": "none",
@@ -286,6 +323,10 @@
     },
 }
 
+jdata_deepmd_input_v0_ni128 = jdata_deepmd_input_v0.copy()
+jdata_deepmd_input_v0_ni256 = jdata_deepmd_input_v0.copy()
+jdata_deepmd_input_v0_ni256["nvnmd"]["max_nnei"] = 256
+
 jdata_deepmd_input_v1 = {
     "model": {
         "descriptor": {
@@ -308,6 +349,7 @@
     },
     "nvnmd": {
         "version": 1,
+        "max_nnei": 128,  # 128 or 256
         "net_size": 128,
         "config_file": "none",
         "weight_file": "none",
@@ -347,6 +389,10 @@
     },
 }
 
+jdata_deepmd_input_v1_ni128 = jdata_deepmd_input_v1.copy()
+jdata_deepmd_input_v1_ni256 = jdata_deepmd_input_v1.copy()
+jdata_deepmd_input_v1_ni256["nvnmd"]["max_nnei"] = 256
+
 NVNMD_WELCOME = (
     r" _   _  __     __  _   _   __  __   ____  ",
     r"| \ | | \ \   / / | \ | | |  \/  | |  _ \ ",

deepmd/nvnmd/descriptor/se_a.py

@@ -50,12 +50,17 @@ def check_switch_range(davg, dstd):
         else:
             min_dist = nvnmd_cfg.weight["train_attr.min_nbor_dist"]
     else:
-        min_dist = rmin
+        min_dist = None
+
+    # fix the bug: if model initial mode is 'init_from_model',
+    # we need dmin to calculate smin and smax in mapt.py
+    if min_dist is not None:
+        nvnmd_cfg.dscp["dmin"] = min_dist
+        nvnmd_cfg.save()
 
     # if davg and dstd is None, the model initial mode is in
     #  'init_from_model', 'restart', 'init_from_frz_model', 'finetune'
     if (davg is not None) and (dstd is not None):
-        nvnmd_cfg.dscp["dmin"] = min_dist
         nvnmd_cfg.get_s_range(davg, dstd)
 
 

deepmd/nvnmd/descriptor/se_atten.py

@@ -49,7 +49,13 @@ def check_switch_range(davg, dstd):
         else:
             min_dist = nvnmd_cfg.weight["train_attr.min_nbor_dist"]
     else:
-        min_dist = rmin
+        min_dist = None
+
+    # fix the bug: if model initial mode is 'init_from_model',
+    # we need dmin to calculate smin and smax in mapt.py
+    if min_dist is not None:
+        nvnmd_cfg.dscp["dmin"] = min_dist
+        nvnmd_cfg.save()
 
     # if davg and dstd is None, the model initial mode is in
     #  'init_from_model', 'restart', 'init_from_frz_model', 'finetune'
@@ -58,7 +64,6 @@ def check_switch_range(davg, dstd):
             davg = np.zeros([ntype, ndescrpt])
         if dstd is None:
             dstd = np.ones([ntype, ndescrpt])
-        nvnmd_cfg.dscp["dmin"] = min_dist
         nvnmd_cfg.get_s_range(davg, dstd)
 
 

deepmd/nvnmd/entrypoints/freeze.py

@@ -52,6 +52,7 @@ def filter_tensorVariableList(tensorVariableList) -> dict:
         p1 = p1 or name.startswith("filter_type_")
         p1 = p1 or name.startswith("layer_")
         p1 = p1 or name.startswith("final_layer")
+        p1 = p1 or name.endswith("t_bias_atom_e")
         p2 = "Adam" not in name
         p3 = "XXX" not in name
         if p1 and p2 and p3:
@@ -75,4 +76,5 @@ def save_weight(sess, file_name: str = "nvnmd/weight.npy"):
     else:
         min_dist = 0.0
     dic_key_value["train_attr.min_nbor_dist"] = min_dist
+    dic_key_value["t_bias_atom_e"] = dic_key_value["fitting_attr.t_bias_atom_e"]
     FioDic().save(file_name, dic_key_value)

deepmd/nvnmd/entrypoints/mapt.py

@@ -87,9 +87,22 @@ def __init__(self, config_file: str, weight_file: str, map_file: str):
         jdata["weight_file"] = weight_file
         jdata["enable"] = True
 
+        # 0 : xyz_scatter = xyz_scatter * two_embd + xyz_scatter;
+        # Gs + 1, Gt + 0
+        # 1 : xyz_scatter = xyz_scatter * two_embd + two_embd   ;
+        # Gs + 0, Gt + 1
+        self.Gs_Gt_mode = 1
+
         nvnmd_cfg.init_from_jdata(jdata)
 
     def build_map(self):
+        if self.Gs_Gt_mode == 0:
+            self.shift_Gs = 1
+            self.shift_Gt = 0
+        if self.Gs_Gt_mode == 1:
+            self.shift_Gs = 0
+            self.shift_Gt = 1
+        #
         M = nvnmd_cfg.dscp["M1"]
         if nvnmd_cfg.version == 0:
             ndim = nvnmd_cfg.dscp["ntype"]
@@ -482,7 +495,7 @@ def build_s2g_grad(self):
             shift = 0
         if nvnmd_cfg.version == 1:
             ndim = 1
-            shift = 1
+            shift = self.shift_Gs
         #
         dic_ph = {}
         dic_ph["s"] = tf.placeholder(tf.float64, [None, 1], "t_s")
@@ -496,6 +509,13 @@ def run_s2g(self):
         r"""Build s-> graph and run it to get value of mapping table."""
         smin = nvnmd_cfg.dscp["smin"]
         smax = nvnmd_cfg.dscp["smax"]
+        # fix the bug: if model initial mode is 'init_from_model',
+        # we need dmin to calculate smin and smax in mapt.py
+        if smin == -2:
+            davg, dstd = get_normalize(nvnmd_cfg.weight)
+            nvnmd_cfg.get_s_range(davg, dstd)
+            smin = nvnmd_cfg.dscp["smin"]
+            smax = nvnmd_cfg.dscp["smax"]
 
         tf.reset_default_graph()
         dic_ph = self.build_s2g_grad()
@@ -567,9 +587,11 @@ def build_t2g(self):
             two_side_type_embedding,
             [-1, two_side_type_embedding.shape[-1]],
         )
-
+        # see se_atten.py in dp
         wbs = [get_filter_type_weight(nvnmd_cfg.weight, ll) for ll in range(1, 5)]
-        dic_ph["gt"] = self.build_embedding_net(two_side_type_embedding, wbs)
+        dic_ph["gt"] = (
+            self.build_embedding_net(two_side_type_embedding, wbs) + self.shift_Gt
+        )
         return dic_ph
 
     def run_t2g(self):

deepmd/nvnmd/entrypoints/train.py

@@ -100,6 +100,7 @@ def normalized_input_qnn(jdata, PATH_QNN, CONFIG_CNN, WEIGHT_CNN, MAP_CNN):
     jdata_nvnmd = jdata_deepmd_input_v0["nvnmd"]
     jdata_nvnmd["enable"] = True
     jdata_nvnmd["version"] = nvnmd_cfg.version
+    jdata_nvnmd["max_nnei"] = nvnmd_cfg.max_nnei
     jdata_nvnmd["config_file"] = CONFIG_CNN
     jdata_nvnmd["weight_file"] = WEIGHT_CNN
     jdata_nvnmd["map_file"] = MAP_CNN
@@ -117,6 +118,7 @@ def normalized_input_qnn(jdata, PATH_QNN, CONFIG_CNN, WEIGHT_CNN, MAP_CNN):
 def train_nvnmd(
     *,
     INPUT: str,
+    init_model: Optional[str],
     restart: Optional[str],
     step: str,
     skip_neighbor_stat: bool = False,
@@ -142,6 +144,7 @@ def train_nvnmd(
         jdata = jdata_cmd_train.copy()
         jdata["INPUT"] = INPUT_CNN
         jdata["log_path"] = LOG_CNN
+        jdata["init_model"] = init_model
         jdata["restart"] = restart
         jdata["skip_neighbor_stat"] = skip_neighbor_stat
         train(**jdata)

deepmd/nvnmd/entrypoints/wrap.py

@@ -156,33 +156,75 @@ def wrap_head(self, nhs, nws):
         r"""Wrap the head information.
 
         version
+        nhead
         nheight
-        nweight
-        rcut
+        nwidth
+        rcut       cut-off radius
+        ntype      number of atomic species
+        nnei       number of neighbors
+        atom_ener  atom bias energy
         """
         nbit = nvnmd_cfg.nbit
         ctrl = nvnmd_cfg.ctrl
+        dscp = nvnmd_cfg.dscp
+        fitn = nvnmd_cfg.fitn
+        weight = nvnmd_cfg.weight
         VERSION = ctrl["VERSION"]
+        SUB_VERSION = ctrl["SUB_VERSION"]
+        MAX_NNEI = ctrl["MAX_NNEI"]
+        nhead = 128
         NBIT_MODEL_HEAD = nbit["NBIT_MODEL_HEAD"]
         NBIT_FIXD_FL = nbit["NBIT_FIXD_FL"]
-        rcut = nvnmd_cfg.dscp["rcut"]
+        rcut = dscp["rcut"]
+        ntype = dscp["ntype"]
+        SEL = dscp["SEL"]
 
         bs = ""
         e = Encode()
         # version
-        bs = e.dec2bin(VERSION, NBIT_MODEL_HEAD)[0] + bs
+        vv = VERSION + 256 * SUB_VERSION + 256 * 256 * MAX_NNEI
+        bs = e.dec2bin(vv, NBIT_MODEL_HEAD)[0] + bs
+        # nhead
+        bs = e.dec2bin(nhead, NBIT_MODEL_HEAD)[0] + bs
         # height
         for n in nhs:
             bs = e.dec2bin(n, NBIT_MODEL_HEAD)[0] + bs
-        # weight
+        # width
         for n in nws:
             bs = e.dec2bin(n, NBIT_MODEL_HEAD)[0] + bs
-        # dscp
+        # rcut
         RCUT = e.qr(rcut, NBIT_FIXD_FL)
         bs = e.dec2bin(RCUT, NBIT_MODEL_HEAD)[0] + bs
+        # ntype
+        bs = e.dec2bin(ntype, NBIT_MODEL_HEAD)[0] + bs
+        # nnei
+        if VERSION == 0:
+            for tt in range(ntype):
+                bs = e.dec2bin(SEL[tt], NBIT_MODEL_HEAD)[0] + bs
+        if VERSION == 1:
+            bs = e.dec2bin(SEL, NBIT_MODEL_HEAD)[0] + bs
+        # atom_ener
+        # fix the bug: the different energy between qnn and lammps
+        if "t_bias_atom_e" in weight.keys():
+            atom_ener = weight["t_bias_atom_e"]
+        else:
+            atom_ener = [0] * 32
+        nlayer_fit = fitn["nlayer_fit"]
+        if VERSION == 0:
+            for tt in range(ntype):
+                w, b, _idt = get_fitnet_weight(weight, tt, nlayer_fit - 1, nlayer_fit)
+                shift = atom_ener[tt] + b[0]
+                SHIFT = e.qr(shift, NBIT_FIXD_FL)
+                bs = e.dec2bin(SHIFT, NBIT_MODEL_HEAD, signed=True)[0] + bs
+        if VERSION == 1:
+            for tt in range(ntype):
+                w, b, _idt = get_fitnet_weight(weight, 0, nlayer_fit - 1, nlayer_fit)
+                shift = atom_ener[tt] + b[0]
+                SHIFT = e.qr(shift, NBIT_FIXD_FL)
+                bs = e.dec2bin(SHIFT, NBIT_MODEL_HEAD, signed=True)[0] + bs
         # extend
         hs = e.bin2hex(bs)
-        hs = e.extend_hex(hs, NBIT_MODEL_HEAD * 32)
+        hs = e.extend_hex(hs, NBIT_MODEL_HEAD * nhead)
         return hs
 
     def wrap_dscp(self):

deepmd/nvnmd/utils/argcheck.py

@@ -8,6 +8,7 @@ def nvnmd_args():
     doc_version = (
         "configuration the nvnmd version (0 | 1), 0 for 4 types, 1 for 32 types"
     )
+    doc_max_nnei = "configuration the max number of neighbors, 128|256 for version 0, 128 for version 1"
     doc_net_size_file = (
         "configuration the number of nodes of fitting_net, just can be set as 128"
     )
@@ -25,6 +26,7 @@ def nvnmd_args():
     doc_quantize_fitting_net = "enable the quantizatioin of fitting_net"
     args = [
         Argument("version", int, optional=False, default=0, doc=doc_version),
+        Argument("max_nnei", int, optional=False, default=128, doc=doc_max_nnei),
         Argument("net_size", int, optional=False, default=128, doc=doc_net_size_file),
         Argument("map_file", str, optional=False, default="none", doc=doc_map_file),
         Argument(