fix(dp/pt): support auto sel for dpa2 (#4323)

Fix #4314 . 

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Enhanced selection update functionality to accommodate three-body
interactions.
- Expanded argument handling to allow flexible input types for neighbor
selection parameters.

- **Bug Fixes**
- Improved error handling and robustness in the selection update
process.

- **Documentation**
- Updated documentation strings for clarity on new input types and usage
constraints.

- **Tests**
- Added a new test method to validate the functionality of the updated
selection method for the `dpa2` model.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: iProzd

deepmd/dpmodel/descriptor/dpa2.py

@@ -1057,6 +1057,14 @@ def update_sel(
             True,
         )
         local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0]
+        min_nbor_dist, repinit_three_body_sel = update_sel.update_one_sel(
+            train_data,
+            type_map,
+            local_jdata_cpy["repinit"]["three_body_rcut"],
+            local_jdata_cpy["repinit"]["three_body_sel"],
+            True,
+        )
+        local_jdata_cpy["repinit"]["three_body_sel"] = repinit_three_body_sel[0]
         min_nbor_dist, repformer_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/pt/model/descriptor/dpa2.py

@@ -842,6 +842,14 @@ def update_sel(
             True,
         )
         local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0]
+        min_nbor_dist, repinit_three_body_sel = update_sel.update_one_sel(
+            train_data,
+            type_map,
+            local_jdata_cpy["repinit"]["three_body_rcut"],
+            local_jdata_cpy["repinit"]["three_body_sel"],
+            True,
+        )
+        local_jdata_cpy["repinit"]["three_body_sel"] = repinit_three_body_sel[0]
         min_nbor_dist, repformer_sel = update_sel.update_one_sel(
             train_data,
             type_map,

deepmd/utils/argcheck.py

@@ -940,7 +940,9 @@ def dpa2_repinit_args():
     # repinit args
     doc_rcut = "The cut-off radius."
     doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`."
-    doc_nsel = "Maximally possible number of selected neighbors."
+    doc_nsel = 'Maximally possible number of selected neighbors. It can be:\n\n\
+    - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\
+    - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".'
     doc_neuron = (
         "Number of neurons in each hidden layers of the embedding net."
         "When two layers are of the same size or one layer is twice as large as the previous layer, "
@@ -972,15 +974,17 @@ def dpa2_repinit_args():
         "When two layers are of the same size or one layer is twice as large as the previous layer, "
         "a skip connection is built."
     )
-    doc_three_body_sel = "Maximally possible number of selected neighbors in the three-body representation."
+    doc_three_body_sel = 'Maximally possible number of selected neighbors in the three-body representation. It can be:\n\n\
+    - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\
+    - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".'
     doc_three_body_rcut = "The cut-off radius in the three-body representation."
     doc_three_body_rcut_smth = "Where to start smoothing in the three-body representation. For example the 1/r term is smoothed from `three_body_rcut` to `three_body_rcut_smth`."
 
     return [
         # repinit args
         Argument("rcut", float, doc=doc_rcut),
         Argument("rcut_smth", float, doc=doc_rcut_smth),
-        Argument("nsel", int, doc=doc_nsel),
+        Argument("nsel", [int, str], doc=doc_nsel),
         Argument(
             "neuron",
             list,
@@ -1066,7 +1070,11 @@ def dpa2_repinit_args():
             doc=doc_three_body_rcut_smth,
         ),
         Argument(
-            "three_body_sel", int, optional=True, default=40, doc=doc_three_body_sel
+            "three_body_sel",
+            [int, str],
+            optional=True,
+            default=40,
+            doc=doc_three_body_sel,
         ),
     ]
 
@@ -1076,7 +1084,9 @@ def dpa2_repformer_args():
     # repformer args
     doc_rcut = "The cut-off radius."
     doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`."
-    doc_nsel = "Maximally possible number of selected neighbors."
+    doc_nsel = 'Maximally possible number of selected neighbors. It can be:\n\n\
+    - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\
+    - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".'
     doc_nlayers = "The number of repformer layers."
     doc_g1_dim = "The dimension of invariant single-atom representation."
     doc_g2_dim = "The dimension of invariant pair-atom representation."
@@ -1139,7 +1149,7 @@ def dpa2_repformer_args():
         # repformer args
         Argument("rcut", float, doc=doc_rcut),
         Argument("rcut_smth", float, doc=doc_rcut_smth),
-        Argument("nsel", int, doc=doc_nsel),
+        Argument("nsel", [int, str], doc=doc_nsel),
         Argument(
             "nlayers",
             int,

source/tests/pt/test_update_sel.py

@@ -170,6 +170,49 @@ def test_update_sel_atten_list(self, sel_mock):
         jdata = update_sel(jdata)
         self.assertEqual(jdata, expected_out)
 
+    @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat")
+    def test_update_sel_dpa2_auto(self, sel_mock):
+        sel_mock.return_value = self.mock_min_nbor_dist, [25]
+
+        jdata = {
+            "model": {
+                "descriptor": {
+                    "type": "dpa2",
+                    "repinit": {
+                        "rcut": 6.0,
+                        "nsel": "auto",
+                        "three_body_rcut": 4.0,
+                        "three_body_sel": "auto",
+                    },
+                    "repformer": {
+                        "rcut": 4.0,
+                        "nsel": "auto",
+                    },
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        expected_out = {
+            "model": {
+                "descriptor": {
+                    "type": "dpa2",
+                    "repinit": {
+                        "rcut": 6.0,
+                        "nsel": 28,
+                        "three_body_rcut": 4.0,
+                        "three_body_sel": 28,
+                    },
+                    "repformer": {
+                        "rcut": 4.0,
+                        "nsel": 28,
+                    },
+                }
+            },
+            "training": {"training_data": {}},
+        }
+        jdata = update_sel(jdata)
+        self.assertEqual(jdata, expected_out)
+
     def test_skip_frozen(self):
         jdata = {
             "model": {