feat(pt): add observed-type option for dp show (#4820)

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

- **New Features**
- Added support for displaying "observed-type" information in the model
reporting tool, showing element types observed during training for
single-task and multi-task models.
- **Tests**
- Updated tests to verify correct reporting of observed types for both
single-task and multi-task models.
- **Documentation**
- Updated documentation to include the new "observed-type" attribute in
model information display, with examples and explanations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Duo <50307526+iProzd@users.noreply.github.com>
Co-authored-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

deepmd/entrypoints/show.py

@@ -4,6 +4,9 @@
 from deepmd.infer.deep_eval import (
     DeepEval,
 )
+from deepmd.utils.econf_embd import (
+    sort_element_type,
+)
 
 log = logging.getLogger(__name__)
 
@@ -69,3 +72,34 @@ def show(
         log.info(f"Parameter counts{log_prefix}:")
         for k in sorted(size_dict):
             log.info(f"Parameters in {k}: {size_dict[k]:,}")
+
+    if "observed-type" in ATTRIBUTES:
+        if model_is_multi_task:
+            log.info("The observed types for each branch: ")
+            total_observed_types_list = []
+            model_branches = list(model_params["model_dict"].keys())
+            for branch in model_branches:
+                tmp_model = DeepEval(INPUT, head=branch, no_jit=True)
+                observed_types = tmp_model.get_observed_types()
+                log.info(
+                    f"{branch}: Number of observed types: {observed_types['type_num']} "
+                )
+                log.info(
+                    f"{branch}: Observed types: {observed_types['observed_type']} "
+                )
+                total_observed_types_list += [
+                    tt
+                    for tt in observed_types["observed_type"]
+                    if tt not in total_observed_types_list
+                ]
+            log.info(
+                f"TOTAL number of observed types in the model: {len(total_observed_types_list)} "
+            )
+            log.info(
+                f"TOTAL observed types in the model: {sort_element_type(total_observed_types_list)} "
+            )
+        else:
+            log.info("The observed types for this model: ")
+            observed_types = model.get_observed_types()
+            log.info(f"Number of observed types: {observed_types['type_num']} ")
+            log.info(f"Observed types: {observed_types['observed_type']} ")

deepmd/infer/deep_eval.py

@@ -295,6 +295,10 @@ def get_model_size(self) -> dict:
         """Get model parameter count."""
         raise NotImplementedError("Not implemented in this backend.")
 
+    def get_observed_types(self) -> dict:
+        """Get observed types (elements) of the model during data statistics."""
+        raise NotImplementedError("Not implemented in this backend.")
+
 
 class DeepEval(ABC):
     """High-level Deep Evaluator interface.
@@ -568,3 +572,7 @@ def get_model_def_script(self) -> dict:
     def get_model_size(self) -> dict:
         """Get model parameter count."""
         return self.deep_eval.get_model_size()
+
+    def get_observed_types(self) -> dict:
+        """Get observed types (elements) of the model during data statistics."""
+        return self.deep_eval.get_observed_types()

deepmd/main.py

@@ -851,7 +851,14 @@ def main_parser() -> argparse.ArgumentParser:
     )
     parser_show.add_argument(
         "ATTRIBUTES",
-        choices=["model-branch", "type-map", "descriptor", "fitting-net", "size"],
+        choices=[
+            "model-branch",
+            "type-map",
+            "descriptor",
+            "fitting-net",
+            "size",
+            "observed-type",
+        ],
         nargs="+",
     )
     return parser

deepmd/pt/infer/deep_eval.py

@@ -64,6 +64,9 @@
     to_numpy_array,
     to_torch_tensor,
 )
+from deepmd.utils.econf_embd import (
+    sort_element_type,
+)
 
 if TYPE_CHECKING:
     import ase.neighborlist
@@ -98,6 +101,7 @@ def __init__(
         auto_batch_size: Union[bool, int, AutoBatchSize] = True,
         neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None,
         head: Optional[Union[str, int]] = None,
+        no_jit: bool = False,
         **kwargs: Any,
     ) -> None:
         self.output_def = output_def
@@ -130,7 +134,7 @@ def __init__(
                         ] = state_dict[item].clone()
                 state_dict = state_dict_head
             model = get_model(self.input_param).to(DEVICE)
-            if not self.input_param.get("hessian_mode"):
+            if not self.input_param.get("hessian_mode") and not no_jit:
                 model = torch.jit.script(model)
             self.dp = ModelWrapper(model)
             self.dp.load_state_dict(state_dict)
@@ -648,6 +652,22 @@ def get_model_size(self) -> dict:
             "total": sum_param_des + sum_param_fit,
         }
 
+    def get_observed_types(self) -> dict:
+        """Get observed types (elements) of the model during data statistics.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the information of observed type in the model:
+            - 'type_num': the total number of observed types in this model.
+            - 'observed_type': a list of the observed types in this model.
+        """
+        observed_type_list = self.dp.model["Default"].get_observed_type_list()
+        return {
+            "type_num": len(observed_type_list),
+            "observed_type": sort_element_type(observed_type_list),
+        }
+
     def eval_descriptor(
         self,
         coords: np.ndarray,

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).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()
         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."""

deepmd/utils/econf_embd.py

@@ -9,6 +9,7 @@
     "electronic_configuration_embedding",
     "make_econf_embedding",
     "normalized_electronic_configuration_embedding",
+    "sort_element_type",
     "transform_to_spin_rep",
 ]
 
@@ -263,3 +264,16 @@ def print_econf_embedding(res: dict[str, np.ndarray]) -> None:
         vvstr = ",".join([str(ii) for ii in vv])
         space = " " * (2 - len(kk))
         print(f'"{kk}"{space} : [{vvstr}],')  # noqa: T201
+
+
+def sort_element_type(elements: list[str]) -> list[str]:
+    """Sort element types based on their atomic number."""
+
+    def get_atomic_number(symbol):
+        try:
+            return element(symbol).atomic_number
+        except ValueError:
+            return float("inf")
+
+    sorted_elements = sorted(elements, key=lambda x: get_atomic_number(x))
+    return sorted_elements

doc/model/show-model-info.md

@@ -17,6 +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. Only energy models are supported now.
 
 ## Example Usage
 
@@ -60,6 +61,12 @@ Depending on the provided attributes and the model type, the output includes:
 
   - Prints the number of parameters for each component (`descriptor`, `fitting-net`, etc.), as well as the total parameter count.
 
+- **observed-type**
+
+  - Displays the count and list of observed element types of the model during data statistics.
+  - For multitask models, it shows the observed types for each branch.
+  - Note: This info shows the types observed during training data statistics, which may differ from the type map.
+
 ## Example Output
 
 For a singletask model, the output might look like:
@@ -73,6 +80,9 @@ Parameter counts:
 Parameters in descriptor: 19,350
 Parameters in fitting-net: 119,091
 Parameters in total: 138,441
+The observed types for this model:
+Number of observed types: 2
+Observed types: ['H', 'O']
 ```
 
 For a multitask model, if `model-branch` is selected, it will additionally display available branches: