From 37015660a3e7c55dccc0698f1524ef8cf7c56b07 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Fri, 8 Nov 2024 15:37:50 +0800 Subject: [PATCH] fix(dp/pt): support auto sel for dpa2 (#4323) Fix #4314 . ## 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. --- deepmd/dpmodel/descriptor/dpa2.py | 8 ++++++ deepmd/pt/model/descriptor/dpa2.py | 8 ++++++ deepmd/utils/argcheck.py | 22 ++++++++++----- source/tests/pt/test_update_sel.py | 43 ++++++++++++++++++++++++++++++ 4 files changed, 75 insertions(+), 6 deletions(-) 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": {