breaking: seperate params in dpa2 (#3768)

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

- **New Features**
- Enhanced descriptor configuration with the introduction of
`RepinitArgs` and `RepformerArgs` classes for more flexible and
structured argument handling.
- Improved JSON configuration structure for descriptors, aligning
properties under `repinit` and `repformer` into nested objects for
clarity.

- **Refactor**
- Updated the initialization process in the model to utilize the new
argument classes, ensuring more robust setup.
- Refactored argument checking functions to support the new class-based
configuration.

- **Documentation**
- Streamlined and clarified documentation related to descriptor
arguments to better align with the new configuration approach.

- **Tests**
- Extended testing suite to include new argument classes, enhancing test
coverage and consistency checks.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

examples/water/dpa2/input_torch.json

@@ -7,36 +7,40 @@
     ],
     "descriptor": {
       "type": "dpa2",
-      "repinit_tebd_dim": 8,
-      "repinit_rcut": 9.0,
-      "repinit_rcut_smth": 8.0,
-      "repinit_nsel": 120,
-      "repformer_rcut": 4.0,
-      "repformer_rcut_smth": 3.5,
-      "repformer_nsel": 40,
-      "repinit_neuron": [
-        25,
-        50,
-        100
-      ],
-      "repinit_axis_neuron": 12,
-      "repinit_activation_function": "tanh",
-      "repformer_nlayers": 12,
-      "repformer_g1_dim": 128,
-      "repformer_g2_dim": 32,
-      "repformer_attn2_hidden": 32,
-      "repformer_attn2_nhead": 4,
-      "repformer_attn1_hidden": 128,
-      "repformer_attn1_nhead": 4,
-      "repformer_axis_neuron": 4,
-      "repformer_update_h2": false,
-      "repformer_update_g1_has_conv": true,
-      "repformer_update_g1_has_grrg": true,
-      "repformer_update_g1_has_drrd": true,
-      "repformer_update_g1_has_attn": true,
-      "repformer_update_g2_has_g1g1": true,
-      "repformer_update_g2_has_attn": true,
-      "repformer_attn2_has_gate": true,
+      "repinit": {
+        "tebd_dim": 8,
+        "rcut": 9.0,
+        "rcut_smth": 8.0,
+        "nsel": 120,
+        "neuron": [
+          25,
+          50,
+          100
+        ],
+        "axis_neuron": 12,
+        "activation_function": "tanh"
+      },
+      "repformer": {
+        "rcut": 4.0,
+        "rcut_smth": 3.5,
+        "nsel": 40,
+        "nlayers": 12,
+        "g1_dim": 128,
+        "g2_dim": 32,
+        "attn2_hidden": 32,
+        "attn2_nhead": 4,
+        "attn1_hidden": 128,
+        "attn1_nhead": 4,
+        "axis_neuron": 4,
+        "update_h2": false,
+        "update_g1_has_conv": true,
+        "update_g1_has_grrg": true,
+        "update_g1_has_drrd": true,
+        "update_g1_has_attn": true,
+        "update_g2_has_g1g1": true,
+        "update_g2_has_attn": true,
+        "attn2_has_gate": true
+      },
       "add_tebd_to_repinit_out": false
     },
     "fitting_net": {

source/tests/common/dpmodel/test_descriptor_dpa2.py

@@ -6,6 +6,10 @@
 from deepmd.dpmodel.descriptor import (
     DescrptDPA2,
 )
+from deepmd.dpmodel.descriptor.dpa2 import (
+    RepformerArgs,
+    RepinitArgs,
+)
 
 from .case_single_frame_with_nlist import (
     TestCaseSingleFrameWithNlist,
@@ -28,14 +32,21 @@ def test_self_consistency(
         dstd = 0.1 + np.abs(dstd)
         dstd_2 = 0.1 + np.abs(dstd_2)
 
+        repinit = RepinitArgs(
+            rcut=self.rcut,
+            rcut_smth=self.rcut_smth,
+            nsel=self.sel_mix,
+        )
+        repformer = RepformerArgs(
+            rcut=self.rcut / 2,
+            rcut_smth=self.rcut_smth,
+            nsel=nnei // 2,
+        )
+
         em0 = DescrptDPA2(
             ntypes=self.nt,
-            repinit_rcut=self.rcut,
-            repinit_rcut_smth=self.rcut_smth,
-            repinit_nsel=self.sel_mix,
-            repformer_rcut=self.rcut / 2,
-            repformer_rcut_smth=self.rcut_smth,
-            repformer_nsel=nnei // 2,
+            repinit=repinit,
+            repformer=repformer,
         )
 
         em0.repinit.mean = davg

source/tests/consistent/descriptor/test_dpa2.py

@@ -32,6 +32,10 @@
 # not implemented
 DescrptDPA2TF = None
 
+from deepmd.dpmodel.descriptor.dpa2 import (
+    RepformerArgs,
+    RepinitArgs,
+)
 from deepmd.utils.argcheck import (
     descrpt_se_atten_args,
 )
@@ -88,45 +92,53 @@ def data(self) -> dict:
         ) = self.param
         return {
             "ntypes": self.ntypes,
-            "repinit_rcut": 6.00,
-            "repinit_rcut_smth": 5.80,
-            "repinit_nsel": 10,
-            "repformer_rcut": 4.00,
-            "repformer_rcut_smth": 3.50,
-            "repformer_nsel": 8,
             # kwargs for repinit
-            "repinit_neuron": [6, 12, 24],
-            "repinit_axis_neuron": 3,
-            "repinit_tebd_dim": 4,
-            "repinit_tebd_input_mode": repinit_tebd_input_mode,
-            "repinit_set_davg_zero": repinit_set_davg_zero,
-            "repinit_activation_function": "tanh",
-            "repinit_type_one_side": repinit_type_one_side,
+            "repinit": RepinitArgs(
+                **{
+                    "rcut": 6.00,
+                    "rcut_smth": 5.80,
+                    "nsel": 10,
+                    "neuron": [6, 12, 24],
+                    "axis_neuron": 3,
+                    "tebd_dim": 4,
+                    "tebd_input_mode": repinit_tebd_input_mode,
+                    "set_davg_zero": repinit_set_davg_zero,
+                    "activation_function": "tanh",
+                    "type_one_side": repinit_type_one_side,
+                }
+            ),
             # kwargs for repformer
-            "repformer_nlayers": 3,
-            "repformer_g1_dim": 20,
-            "repformer_g2_dim": 10,
-            "repformer_axis_neuron": 3,
-            "repformer_direct_dist": repformer_direct_dist,
-            "repformer_update_g1_has_conv": repformer_update_g1_has_conv,
-            "repformer_update_g1_has_drrd": repformer_update_g1_has_drrd,
-            "repformer_update_g1_has_grrg": repformer_update_g1_has_grrg,
-            "repformer_update_g1_has_attn": repformer_update_g1_has_attn,
-            "repformer_update_g2_has_g1g1": repformer_update_g2_has_g1g1,
-            "repformer_update_g2_has_attn": repformer_update_g2_has_attn,
-            "repformer_update_h2": repformer_update_h2,
-            "repformer_attn1_hidden": 12,
-            "repformer_attn1_nhead": 2,
-            "repformer_attn2_hidden": 10,
-            "repformer_attn2_nhead": 2,
-            "repformer_attn2_has_gate": repformer_attn2_has_gate,
-            "repformer_activation_function": "tanh",
-            "repformer_update_style": repformer_update_style,
-            "repformer_update_residual": 0.001,
-            "repformer_update_residual_init": repformer_update_residual_init,
-            "repformer_set_davg_zero": True,
-            "repformer_trainable_ln": repformer_trainable_ln,
-            "repformer_ln_eps": repformer_ln_eps,
+            "repformer": RepformerArgs(
+                **{
+                    "rcut": 4.00,
+                    "rcut_smth": 3.50,
+                    "nsel": 8,
+                    "nlayers": 3,
+                    "g1_dim": 20,
+                    "g2_dim": 10,
+                    "axis_neuron": 3,
+                    "direct_dist": repformer_direct_dist,
+                    "update_g1_has_conv": repformer_update_g1_has_conv,
+                    "update_g1_has_drrd": repformer_update_g1_has_drrd,
+                    "update_g1_has_grrg": repformer_update_g1_has_grrg,
+                    "update_g1_has_attn": repformer_update_g1_has_attn,
+                    "update_g2_has_g1g1": repformer_update_g2_has_g1g1,
+                    "update_g2_has_attn": repformer_update_g2_has_attn,
+                    "update_h2": repformer_update_h2,
+                    "attn1_hidden": 12,
+                    "attn1_nhead": 2,
+                    "attn2_hidden": 10,
+                    "attn2_nhead": 2,
+                    "attn2_has_gate": repformer_attn2_has_gate,
+                    "activation_function": "tanh",
+                    "update_style": repformer_update_style,
+                    "update_residual": 0.001,
+                    "update_residual_init": repformer_update_residual_init,
+                    "set_davg_zero": True,
+                    "trainable_ln": repformer_trainable_ln,
+                    "ln_eps": repformer_ln_eps,
+                }
+            ),
             # kwargs for descriptor
             "concat_output_tebd": True,
             "precision": precision,

source/tests/pt/model/models/dpa2.json

@@ -5,35 +5,40 @@
   ],
   "descriptor": {
     "type": "dpa2",
-    "repinit_rcut": 6.0,
-    "repinit_rcut_smth": 2.0,
-    "repinit_nsel": 30,
-    "repformer_rcut": 4.0,
-    "repformer_rcut_smth": 0.5,
-    "repformer_nsel": 10,
-    "repinit_neuron": [
-      2,
-      4,
-      8
-    ],
-    "repinit_axis_neuron": 4,
-    "repinit_activation_function": "tanh",
-    "repformer_nlayers": 12,
-    "repformer_g1_dim": 8,
-    "repformer_g2_dim": 5,
-    "repformer_attn2_hidden": 3,
-    "repformer_attn2_nhead": 1,
-    "repformer_attn1_hidden": 5,
-    "repformer_attn1_nhead": 1,
-    "repformer_axis_neuron": 4,
-    "repformer_update_h2": false,
-    "repformer_update_g1_has_conv": true,
-    "repformer_update_g1_has_grrg": true,
-    "repformer_update_g1_has_drrd": true,
-    "repformer_update_g1_has_attn": true,
-    "repformer_update_g2_has_g1g1": true,
-    "repformer_update_g2_has_attn": true,
-    "repformer_attn2_has_gate": true,
+    "repinit": {
+      "rcut": 6.0,
+      "rcut_smth": 2.0,
+      "nsel": 30,
+      "neuron": [
+        2,
+        4,
+        8
+      ],
+      "axis_neuron": 4,
+      "activation_function": "tanh"
+
+    },
+    "repformer": {
+      "rcut": 4.0,
+      "rcut_smth": 0.5,
+      "nsel": 10,
+      "nlayers": 12,
+      "g1_dim": 8,
+      "g2_dim": 5,
+      "attn2_hidden": 3,
+      "attn2_nhead": 1,
+      "attn1_hidden": 5,
+      "attn1_nhead": 1,
+      "axis_neuron": 4,
+      "update_h2": false,
+      "update_g1_has_conv": true,
+      "update_g1_has_grrg": true,
+      "update_g1_has_drrd": true,
+      "update_g1_has_attn": true,
+      "update_g2_has_g1g1": true,
+      "update_g2_has_attn": true,
+      "attn2_has_gate": true
+    },
     "add_tebd_to_repinit_out": false
   },
   "fitting_net": {

source/tests/pt/model/test_autodiff.py

@@ -197,27 +197,13 @@ def setUp(self):
 
 class TestEnergyModelDPA2Force(unittest.TestCase, ForceTest):
     def setUp(self):
-        model_params_sample = copy.deepcopy(model_dpa2)
-        model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][
-            "repinit_rcut"
-        ]
-        model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][
-            "repinit_nsel"
-        ]
         model_params = copy.deepcopy(model_dpa2)
         self.type_split = True
         self.model = get_model(model_params).to(env.DEVICE)
 
 
 class TestEnergyModelDPAUniVirial(unittest.TestCase, VirialTest):
     def setUp(self):
-        model_params_sample = copy.deepcopy(model_dpa2)
-        model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][
-            "repinit_rcut"
-        ]
-        model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][
-            "repinit_nsel"
-        ]
         model_params = copy.deepcopy(model_dpa2)
         self.type_split = True
         self.model = get_model(model_params).to(env.DEVICE)