pd: support dpa1

deepmd/pd/entrypoints/main.py

@@ -230,7 +230,7 @@
     use_pretrain_script: bool = False,
     force_load: bool = False,
     output: str = "out.json",
-):
+) -> None:
     log.info("Configuration path: %s", input_file)
     SummaryPrinter()()
     with open(input_file) as fin:
@@ -321,18 +321,26 @@
     # save min_nbor_dist
     if min_nbor_dist is not None:
         if not multi_task:
-            trainer.model.min_nbor_dist = min_nbor_dist
+            trainer.model.min_nbor_dist = paddle.to_tensor(
+                min_nbor_dist,
+                dtype=paddle.float64,
+                place=DEVICE,
+            )
         else:
             for model_item in min_nbor_dist:
-                trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item]
+                trainer.model[model_item].min_nbor_dist = paddle.to_tensor(
+                    min_nbor_dist[model_item],
+                    dtype=paddle.float64,
+                    place=DEVICE,
+                )
     trainer.run()
 
 
 def freeze(
     model: str,
     output: str = "frozen_model.json",
     head: Optional[str] = None,
-):
+) -> None:
     paddle.set_flags(
         {
             "FLAGS_save_cf_stack_op": 1,
@@ -383,7 +391,7 @@
     numb_batch: int = 0,
     model_branch: Optional[str] = None,
     output: Optional[str] = None,
-):
+) -> None:
     if input_file.endswith(".pd"):
         old_state_dict = paddle.load(input_file)
         model_state_dict = copy.deepcopy(old_state_dict.get("model", old_state_dict))

deepmd/pd/loss/ener.py

@@ -10,7 +10,6 @@
     TaskLoss,
 )
 from deepmd.pd.utils import (
-    decomp,
     env,
 )
 from deepmd.pd.utils.env import (
@@ -224,10 +223,7 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
 
             if self.relative_f is not None:
                 force_label_3 = force_label.reshape([-1, 3])
-                # norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
-                norm_f = (
-                    decomp.norm(force_label_3, axis=1, keepdim=True) + self.relative_f
-                )
+                norm_f = force_label_3.norm(axis=1, keepdim=True) + self.relative_f
                 diff_f_3 = diff_f.reshape([-1, 3])
                 diff_f_3 = diff_f_3 / norm_f
                 diff_f = diff_f_3.reshape([-1])

deepmd/pd/model/atomic_model/dp_atomic_model.py

@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-import copy
 import functools
 import logging
 from typing import (
@@ -52,7 +51,7 @@
         fitting,
         type_map: list[str],
         **kwargs,
-    ):
+    ) -> None:
         super().__init__(type_map, **kwargs)
         ntypes = len(type_map)
         self.type_map = type_map
@@ -201,7 +200,7 @@
 
     @classmethod
     def deserialize(cls, data) -> "DPAtomicModel":
-        data = copy.deepcopy(data)
+        data = data.copy()
         check_version_compatibility(data.pop("@version", 1), 2, 1)
         data.pop("@class", None)
         data.pop("type", None)
@@ -212,6 +211,37 @@
         obj = super().deserialize(data)
         return obj
 
+    def enable_compression(
+        self,
+        min_nbor_dist: float,
+        table_extrapolate: float = 5,
+        table_stride_1: float = 0.01,
+        table_stride_2: float = 0.1,
+        check_frequency: int = -1,
+    ) -> None:
+        """Call descriptor enable_compression().
+
+        Parameters
+        ----------
+        min_nbor_dist
+            The nearest distance between atoms
+        table_extrapolate
+            The scale of model extrapolation
+        table_stride_1
+            The uniform stride of the first table
+        table_stride_2
+            The uniform stride of the second table
+        check_frequency
+            The overflow check frequency
+        """
+        self.descriptor.enable_compression(
+            min_nbor_dist,
+            table_extrapolate,
+            table_stride_1,
+            table_stride_2,
+            check_frequency,
+        )
+
     def forward_atomic(
         self,
         extended_coord,
@@ -278,7 +308,7 @@
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
-    ):
+    ) -> None:
         """
         Compute or load the statistics parameters of the model,
         such as mean and standard deviation of descriptors or the energy bias of the fitting net.

deepmd/pd/model/descriptor/__init__.py

@@ -5,6 +5,10 @@
 from .descriptor import (
     DescriptorBlock,
 )
+from .dpa1 import (
+    DescrptBlockSeAtten,
+    DescrptDPA1,
+)
 from .env_mat import (
     prod_env_mat,
 )
@@ -17,6 +21,8 @@
     "BaseDescriptor",
     "DescriptorBlock",
     "DescrptBlockSeA",
+    "DescrptBlockSeAtten",
+    "DescrptDPA1",
     "DescrptSeA",
     "prod_env_mat",
 ]