fix consistent and bump version

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -575,7 +575,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "dpa3",
-            "@version": 1,
+            "@version": 2,
             "ntypes": self.ntypes,
             "repflow_args": self.repflow_args.serialize(),
             "concat_output_tebd": self.concat_output_tebd,
@@ -586,6 +586,7 @@ def serialize(self) -> dict:
             "trainable": self.trainable,
             "use_econf_tebd": self.use_econf_tebd,
             "use_tebd_bias": self.use_tebd_bias,
+            "use_loc_mapping": self.use_loc_mapping,
             "type_map": self.type_map,
             "type_embedding": self.type_embedding.serialize(),
         }
@@ -610,7 +611,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptDPA3":
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 1, 1)
+        check_version_compatibility(version, 2, 1)
         data.pop("@class")
         data.pop("type")
         repflow_variable = data.pop("repflow_variable").copy()

deepmd/dpmodel/descriptor/repflows.py

@@ -688,6 +688,7 @@ def serialize(self):
             "smooth_edge_update": self.smooth_edge_update,
             "use_dynamic_sel": self.use_dynamic_sel,
             "sel_reduce_factor": self.sel_reduce_factor,
+            "use_loc_mapping": self.use_loc_mapping,
             # variables
             "edge_embd": self.edge_embd.serialize(),
             "angle_embd": self.angle_embd.serialize(),

deepmd/pd/model/descriptor/dpa3.py

@@ -89,6 +89,9 @@ class DescrptDPA3(BaseDescriptor, paddle.nn.Layer):
         Whether to use electronic configuration type embedding.
     use_tebd_bias : bool, Optional
         Whether to use bias in the type embedding layer.
+    use_loc_mapping : bool, Optional
+        Whether to use local atom index mapping in training or non-parallel inference.
+        Not supported yet in Paddle.
     type_map : list[str], Optional
         A list of strings. Give the name to each type of atoms.
     """
@@ -108,6 +111,7 @@ def __init__(
         seed: Optional[Union[int, list[int]]] = None,
         use_econf_tebd: bool = False,
         use_tebd_bias: bool = False,
+        use_loc_mapping: bool = False,
         type_map: Optional[list[str]] = None,
     ) -> None:
         super().__init__()
@@ -152,6 +156,7 @@ def init_subclass_params(sub_data, sub_class):
             smooth_edge_update=self.repflow_args.smooth_edge_update,
             use_dynamic_sel=self.repflow_args.use_dynamic_sel,
             sel_reduce_factor=self.repflow_args.sel_reduce_factor,
+            use_loc_mapping=use_loc_mapping,
             exclude_types=exclude_types,
             env_protection=env_protection,
             precision=precision,
@@ -160,6 +165,7 @@ def init_subclass_params(sub_data, sub_class):
 
         self.use_econf_tebd = use_econf_tebd
         self.use_tebd_bias = use_tebd_bias
+        self.use_loc_mapping = use_loc_mapping
         self.type_map = type_map
         self.tebd_dim = self.repflow_args.n_dim
         self.type_embedding = TypeEmbedNet(
@@ -370,7 +376,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "dpa3",
-            "@version": 1,
+            "@version": 2,
             "ntypes": self.ntypes,
             "repflow_args": self.repflow_args.serialize(),
             "concat_output_tebd": self.concat_output_tebd,
@@ -381,6 +387,7 @@ def serialize(self) -> dict:
             "trainable": self.trainable,
             "use_econf_tebd": self.use_econf_tebd,
             "use_tebd_bias": self.use_tebd_bias,
+            "use_loc_mapping": self.use_loc_mapping,
             "type_map": self.type_map,
             "type_embedding": self.type_embedding.embedding.serialize(),
         }
@@ -405,7 +412,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptDPA3":
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 1, 1)
+        check_version_compatibility(version, 2, 1)
         data.pop("@class")
         data.pop("type")
         repflow_variable = data.pop("repflow_variable").copy()

deepmd/pd/model/descriptor/repflows.py

@@ -112,6 +112,9 @@ class DescrptBlockRepflows(DescriptorBlock):
     optim_update : bool, optional
         Whether to enable the optimized update method.
         Uses a more efficient process when enabled. Defaults to True
+    use_loc_mapping : bool, Optional
+        Whether to use local atom index mapping in training or non-parallel inference.
+        Not supported yet in Paddle.
     ntypes : int
         Number of element types
     activation_function : str, optional
@@ -161,6 +164,7 @@ def __init__(
         smooth_edge_update: bool = False,
         use_dynamic_sel: bool = False,
         sel_reduce_factor: float = 10.0,
+        use_loc_mapping: bool = False,
         optim_update: bool = True,
         seed: Optional[Union[int, list[int]]] = None,
     ) -> None:
@@ -196,6 +200,8 @@ def __init__(
         self.use_dynamic_sel = use_dynamic_sel  # not supported yet
         self.sel_reduce_factor = sel_reduce_factor
         assert not self.use_dynamic_sel, "Dynamic selection is not supported yet."
+        self.use_loc_mapping = use_loc_mapping
+        assert not self.use_loc_mapping, "Local mapping is not supported yet."
 
         self.n_dim = n_dim
         self.e_dim = e_dim

deepmd/pt/model/descriptor/dpa3.py

@@ -371,7 +371,7 @@ def serialize(self) -> dict:
         data = {
             "@class": "Descriptor",
             "type": "dpa3",
-            "@version": 1,
+            "@version": 2,
             "ntypes": self.ntypes,
             "repflow_args": self.repflow_args.serialize(),
             "concat_output_tebd": self.concat_output_tebd,
@@ -382,6 +382,7 @@ def serialize(self) -> dict:
             "trainable": self.trainable,
             "use_econf_tebd": self.use_econf_tebd,
             "use_tebd_bias": self.use_tebd_bias,
+            "use_loc_mapping": self.use_loc_mapping,
             "type_map": self.type_map,
             "type_embedding": self.type_embedding.embedding.serialize(),
         }
@@ -406,7 +407,7 @@ def serialize(self) -> dict:
     def deserialize(cls, data: dict) -> "DescrptDPA3":
         data = data.copy()
         version = data.pop("@version")
-        check_version_compatibility(version, 1, 1)
+        check_version_compatibility(version, 2, 1)
         data.pop("@class")
         data.pop("type")
         repflow_variable = data.pop("repflow_variable").copy()

deepmd/pt/model/descriptor/repflows.py

@@ -156,6 +156,9 @@ class DescrptBlockRepflows(DescriptorBlock):
         In the dynamic selection case, neighbor-scale normalization will use `e_sel / sel_reduce_factor`
         or `a_sel / sel_reduce_factor` instead of the raw `e_sel` or `a_sel` values,
         accommodating larger selection numbers.
+    use_loc_mapping : bool, Optional
+        Whether to use local atom index mapping in training or non-parallel inference.
+        When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.
     optim_update : bool, optional
         Whether to enable the optimized update method.
         Uses a more efficient process when enabled. Defaults to True

deepmd/utils/argcheck.py

@@ -1372,6 +1372,10 @@ def descrpt_dpa3_args():
     doc_seed = "Random seed for parameter initialization."
     doc_use_econf_tebd = "Whether to use electronic configuration type embedding."
     doc_use_tebd_bias = "Whether to use bias in the type embedding layer."
+    doc_use_loc_mapping = (
+        "Whether to use local atom index mapping in training or non-parallel inference. "
+        "When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation."
+    )
     return [
         # doc_repflow args
         Argument("repflow", dict, dpa3_repflow_args(), doc=doc_repflow),
@@ -1426,8 +1430,7 @@ def descrpt_dpa3_args():
             bool,
             optional=True,
             default=True,
-            doc="Whether to use local atom index mapping in training or non-parallel inference. "
-            "When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation.",
+            doc=doc_use_loc_mapping,
         ),
     ]
 

source/tests/consistent/descriptor/test_dpa3.py

@@ -67,6 +67,7 @@
     (True, False),  # optim_update
     (True, False),  # use_exp_switch
     (True, False),  # use_dynamic_sel
+    (True, False),  # use_loc_mapping
     (0.3, 0.0),  # fix_stat_std
     (1, 2),  # n_multi_edge_message
     ("float64",),  # precision
@@ -84,6 +85,7 @@ def data(self) -> dict:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -124,6 +126,7 @@ def data(self) -> dict:
             "precision": precision,
             "exclude_types": exclude_types,
             "env_protection": 0.0,
+            "use_loc_mapping": use_loc_mapping,
             "trainable": True,
         }
 
@@ -139,6 +142,7 @@ def skip_pt(self) -> bool:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -157,6 +161,7 @@ def skip_pd(self) -> bool:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -166,6 +171,7 @@ def skip_pd(self) -> bool:
             or precision == "bfloat16"
             or use_exp_switch
             or use_dynamic_sel
+            or use_loc_mapping
         )  # not supported yet
 
     @property
@@ -180,6 +186,7 @@ def skip_dp(self) -> bool:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -198,6 +205,7 @@ def skip_tf(self) -> bool:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -258,6 +266,7 @@ def setUp(self) -> None:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -339,6 +348,7 @@ def rtol(self) -> float:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,
@@ -363,6 +373,7 @@ def atol(self) -> float:
             optim_update,
             use_exp_switch,
             use_dynamic_sel,
+            use_loc_mapping,
             fix_stat_std,
             n_multi_edge_message,
             precision,