only support energy models

Author: iProzd

deepmd/pt/infer/deep_eval.py

@@ -662,23 +662,10 @@ def get_observed_types(self) -> dict:
             - 'type_num': the total number of observed types in this model.
             - 'observed_type': a list of the observed types in this model.
         """
-        buffers_dict = dict(self.dp.named_buffers())
-        type_map = np.array(self.type_map)
-        out_bias = None
-        for k in buffers_dict:
-            if ".out_bias" in k:
-                # only use out_bias in the first fitting out_def
-                out_bias = buffers_dict[k].detach().cpu().numpy()[0]
-                break
-        assert out_bias is not None, "No out_bias found in the model buffers."
-        assert len(out_bias.shape) == 2, "The supported out_bias should be a 2D array."
-        assert out_bias.shape[0] == len(type_map), (
-            "The out_bias shape does not match the type map length."
-        )
-        bias_mask = (np.abs(out_bias) > 1e-6).any(-1)  # 1e-6 for stability
+        observed_type_list = self.dp.model["Default"].get_observed_type_list()
         return {
-            "type_num": bias_mask.sum(),
-            "observed_type": sort_element_type(type_map[bias_mask].tolist()),
+            "type_num": len(observed_type_list),
+            "observed_type": sort_element_type(observed_type_list),
         }
 
     def eval_descriptor(

deepmd/pt/model/model/ener_model.py

@@ -44,6 +44,32 @@ def enable_hessian(self):
         self.requires_hessian("energy")
         self._hessian_enabled = True
 
+    @torch.jit.export
+    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 = torch.gt(torch.abs(out_bias), 1e-6).any(
+            dim=-1
+        )  # 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()
         output_def = {

deepmd/pt/model/model/model.py

@@ -48,6 +48,16 @@ def compute_or_load_stat(
         """
         raise NotImplementedError
 
+    @torch.jit.export
+    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
+
     @torch.jit.export
     def get_model_def_script(self) -> str:
         """Get the model definition script."""

doc/model/show-model-info.md

@@ -17,7 +17,7 @@ dp --pt show <INPUT> <ATTRIBUTES...>
   - `descriptor`: Displays the model descriptor parameters.
   - `fitting-net`: Displays parameters of the fitting network.
   - `size`: (Supported Backends: PyTorch and PaddlePaddle) Shows the parameter counts for various components.
-  - `observed-type`: (Supported Backends: PyTorch) Shows the observed types (elements) of the model during data statistics.
+  - `observed-type`: (Supported Backends: PyTorch) Shows the observed types (elements) of the model during data statistics. Only energy models are supported now.
 
 ## Example Usage