diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 200747c0ef..097be2ef09 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/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, diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 277aa4917f..ad5167c572 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/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, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 916e4de1b0..3c4a4e6204 100644 --- a/deepmd/utils/argcheck.py +++ b/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,7 +974,9 @@ 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`." @@ -980,7 +984,7 @@ def dpa2_repinit_args(): # 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, diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py index f0b91cfe84..810351b74d 100644 --- a/source/tests/pt/test_update_sel.py +++ b/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": {