nvnmd-v1 with 31-type chemical species (#2676)

1. Add the nvnmd-v1 with `se_atten` descriptor and 31-type chemical
species. nvnmd-v0 is the old version with `se_a` descriptor and 4-type
chemical species.
2. Fix the error for map_flt_nvnmd.cpp for nvnmd-v1. An unexpected value
of floor(-0.0)=-1 occurred at the index of the mapping table.
3. Improve the documentation about runing nvnmd in bohrium.

---------

Co-authored-by: LiuGroupHNU <liujie123@HNU>
Co-authored-by: MoPinghui <mopinghui1020@gmail.com>
Co-authored-by: Han Wang <92130845+wanghan-iapcm@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.pre-commit-config.yaml

@@ -16,7 +16,6 @@ repos:
         exclude: |
             (?x)^(
                 source/tests/infer/dipolecharge_e.pbtxt|
-                source/tests/nvnmd/map.npy|
                 source/tests/infer/deeppolar_new.pbtxt
             )$
     -   id: check-merge-conflict

deepmd/descriptor/se_atten.py

@@ -24,6 +24,17 @@
     op_module,
     tf,
 )
+from deepmd.nvnmd.descriptor.se_atten import (
+    build_davg_dstd,
+    build_op_descriptor,
+    check_switch_range,
+    descrpt2r4,
+    filter_GR2D,
+    filter_lower_R42GR,
+)
+from deepmd.nvnmd.utils.config import (
+    nvnmd_cfg,
+)
 from deepmd.utils.graph import (
     get_attention_layer_variables_from_graph_def,
     get_pattern_nodes_from_graph_def,
@@ -150,9 +161,10 @@ def __init__(
         """
         Constructor
         """
-        assert Version(TF_VERSION) > Version(
-            "2"
-        ), "se_atten only support tensorflow version 2.0 or higher."
+        if not (nvnmd_cfg.enable and (nvnmd_cfg.version == 1)):
+            assert Version(TF_VERSION) > Version(
+                "2"
+            ), "se_atten only support tensorflow version 2.0 or higher."
         self.stripped_type_embedding = stripped_type_embedding
         self.ntypes = ntypes
         self.att_n = attn
@@ -525,6 +537,11 @@ def build(
         """
         davg = self.davg
         dstd = self.dstd
+        if nvnmd_cfg.enable:
+            nvnmd_cfg.set_ntype(self.ntypes)
+            if nvnmd_cfg.restore_descriptor:
+                davg, dstd = build_davg_dstd()
+            check_switch_range(davg, dstd)
         with tf.variable_scope("descrpt_attr" + suffix, reuse=reuse):
             if davg is None:
                 davg = np.zeros([self.ntypes, self.ndescrpt])
@@ -566,14 +583,17 @@ def build(
         self.angular_weight = [None for i in range(self.attn_layer)]
         self.attn_weight_final = [None for i in range(self.attn_layer)]
 
+        op_descriptor = (
+            build_op_descriptor() if nvnmd_cfg.enable else op_module.prod_env_mat_a_mix
+        )
         (
             self.descrpt,
             self.descrpt_deriv,
             self.rij,
             self.nlist,
             self.nei_type_vec,
             self.nmask,
-        ) = op_module.prod_env_mat_a_mix(
+        ) = op_descriptor(
             coord,
             atype,
             natoms,
@@ -645,7 +665,8 @@ def _pass_filter(
                 self.nei_type_vec,  # extra input for atten
             )
             inputs_i *= mask
-
+        if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor:
+            inputs_i = descrpt2r4(inputs_i, atype)
         layer, qmat = self._filter(
             inputs_i,
             type_i,
@@ -1036,6 +1057,20 @@ def _filter_lower(
                         log.info(
                             "use the non-compressible model with stripped type embedding"
                         )
+                    if nvnmd_cfg.enable:
+                        if nvnmd_cfg.quantize_descriptor:
+                            return filter_lower_R42GR(
+                                inputs_i,
+                                atype,
+                                self.nei_type_vec,
+                            )
+                        elif nvnmd_cfg.restore_descriptor:
+                            self.embedding_net_variables = (
+                                nvnmd_cfg.get_dp_init_weights()
+                            )
+                            self.two_side_embeeding_net_variables = (
+                                nvnmd_cfg.get_dp_init_weights()
+                            )
                     if not self.compress:
                         xyz_scatter = embedding_net(
                             xyz_scatter,
@@ -1202,6 +1237,8 @@ def _filter(
             reuse=reuse,
             atype=atype,
         )
+        if nvnmd_cfg.enable:
+            return filter_GR2D(xyz_scatter_1)
         # natom x nei x outputs_size
         # xyz_scatter = tf.concat(xyz_scatter_total, axis=1)
         # natom x nei x 4

deepmd/entrypoints/main.py

@@ -566,6 +566,11 @@ def main_parser() -> argparse.ArgumentParser:
         choices=["s1", "s2"],
         help="steps to train model of NVNMD: s1 (train CNN), s2 (train QNN)",
     )
+    parser_train_nvnmd.add_argument(
+        "--skip-neighbor-stat",
+        action="store_true",
+        help="Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled.",
+    )
     return parser
 
 

deepmd/fit/ener.py

@@ -602,7 +602,13 @@ def build(
             )
             atype_filter = tf.cast(self.atype_nloc >= 0, GLOBAL_TF_FLOAT_PRECISION)
             self.atype_nloc = tf.reshape(self.atype_nloc, [-1])
-
+        if (
+            nvnmd_cfg.enable
+            and nvnmd_cfg.quantize_descriptor
+            and nvnmd_cfg.restore_descriptor
+            and (nvnmd_cfg.version == 1)
+        ):
+            type_embedding = nvnmd_cfg.map["t_ebd"]
         if type_embedding is not None:
             atype_embed = tf.nn.embedding_lookup(type_embedding, self.atype_nloc)
         else: