update code

Author: HydrogenSulfate

deepmd/pd/model/model/ener_model.py

@@ -33,6 +33,37 @@ def __init__(
     ) -> None:
         DPModelCommon.__init__(self)
         DPEnergyModel_.__init__(self, *args, **kwargs)
+        self._hessian_enabled = False
+
+    def enable_hessian(self):
+        raise NotImplementedError(
+            "Hessian calculation is not implemented yet on PaddlePaddle platform."
+        )
+
+    def get_observed_type_list(self) -> list[str]:
+        """Get observed types (elements) of the model during data statistics.
+
+        Returns
+        -------
+        observed_type_list: a list of the observed types in this model.
+        """
+        type_map = self.get_type_map()
+        out_bias = self.atomic_model.get_out_bias()[0]
+
+        assert out_bias is not None, "No out_bias found in the model."
+        assert out_bias.dim() == 2, "The supported out_bias should be a 2D tensor."
+        assert out_bias.size(0) == len(type_map), (
+            "The out_bias shape does not match the type_map length."
+        )
+        bias_mask = (
+            paddle.greater_than(paddle.abs(out_bias), 1e-6).any(axis=-1).detach().cpu()
+        )  # 1e-6 for stability
+
+        observed_type_list: list[str] = []
+        for i in range(len(type_map)):
+            if bias_mask[i]:
+                observed_type_list.append(type_map[i])
+        return observed_type_list
 
     def translated_output_def(self):
         out_def_data = self.model_output_def().get_data()
@@ -50,6 +81,8 @@ def translated_output_def(self):
             output_def["atom_virial"].squeeze(-3)
         if "mask" in out_def_data:
             output_def["mask"] = out_def_data["mask"]
+        if self._hessian_enabled:
+            output_def["hessian"] = out_def_data["energy_derv_r_derv_r"]
         return output_def
 
     def forward(
@@ -81,14 +114,12 @@ def forward(
                     model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(
                         -3
                     )
-                else:
-                    model_predict["atom_virial"] = paddle.zeros(
-                        [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64
-                    )
             else:
                 model_predict["force"] = model_ret["dforce"]
             if "mask" in model_ret:
                 model_predict["mask"] = model_ret["mask"]
+            if self._hessian_enabled:
+                model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-2)
         else:
             model_predict = model_ret
             model_predict["updated_coord"] += coord
@@ -128,10 +159,6 @@ def forward_lower(
                     model_predict["extended_virial"] = model_ret[
                         "energy_derv_c"
                     ].squeeze(-3)
-                else:
-                    model_predict["extended_virial"] = paddle.zeros(
-                        [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64
-                    )
             else:
                 assert model_ret["dforce"] is not None
                 model_predict["dforce"] = model_ret["dforce"]

deepmd/pd/model/model/model.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
+    NoReturn,
     Optional,
 )
 
@@ -8,6 +9,9 @@
 from deepmd.dpmodel.model.base_model import (
     make_base_model,
 )
+from deepmd.pd.utils import (
+    env,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -18,13 +22,16 @@ def __init__(self, *args, **kwargs):
         """Construct a basic model for different tasks."""
         paddle.nn.Layer.__init__(self)
         self.model_def_script = ""
-        self.min_nbor_dist = None
+        self.register_buffer(
+            "min_nbor_dist",
+            paddle.to_tensor(-1.0, dtype=paddle.float64, place=env.DEVICE),
+        )
 
     def compute_or_load_stat(
         self,
         sampled_func,
         stat_file_path: Optional[DPPath] = None,
-    ):
+    ) -> NoReturn:
         """
         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.
@@ -42,13 +49,24 @@ def compute_or_load_stat(
         """
         raise NotImplementedError
 
+    def get_observed_type_list(self) -> list[str]:
+        """Get observed types (elements) of the model during data statistics.
+
+        Returns
+        -------
+        observed_type_list: a list of the observed types in this model.
+        """
+        raise NotImplementedError
+
     def get_model_def_script(self) -> str:
         """Get the model definition script."""
         return self.model_def_script
 
     def get_min_nbor_dist(self) -> Optional[float]:
         """Get the minimum distance between two atoms."""
-        return self.min_nbor_dist
+        if self.min_nbor_dist.item() == -1.0:
+            return None
+        return self.min_nbor_dist.item()
 
     def get_ntypes(self):
         """Returns the number of element types."""