feat(pt): add type-coverage option for dp show

Author: iProzd

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,32 @@ 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 "type-coverage" in ATTRIBUTES:
+        if model_is_multi_task:
+            log.info("The type coverage for each branch: ")
+            total_type_coverage_list = []
+            model_branches = list(model_params["model_dict"].keys())
+            for branch in model_branches:
+                tmp_model = DeepEval(INPUT, head=branch, no_jit=True)
+                type_coverage = tmp_model.get_type_coverage()
+                log.info(
+                    f"{branch}: Number of covered types: {type_coverage['type_num']} "
+                )
+                log.info(f"{branch}: Covered types: {type_coverage['covered_type']} ")
+                total_type_coverage_list += [
+                    tt
+                    for tt in type_coverage["covered_type"]
+                    if tt not in total_type_coverage_list
+                ]
+            log.info(
+                f"TOTAL number of covered types in the model: {len(total_type_coverage_list)} "
+            )
+            log.info(
+                f"TOTAL covered types in the model: {sort_element_type(total_type_coverage_list)} "
+            )
+        else:
+            log.info("The type coverage for this model: ")
+            type_coverage = model.get_type_coverage()
+            log.info(f"Number of covered types: {type_coverage['type_num']} ")
+            log.info(f"Covered types: {type_coverage['covered_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_type_coverage(self) -> dict:
+        """Get model type (element) coverage."""
+        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_type_coverage(self) -> dict:
+        """Get model type (element) coverage."""
+        return self.deep_eval.get_type_coverage()

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",
+            "type-coverage",
+        ],
         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,35 @@ def get_model_size(self) -> dict:
             "total": sum_param_des + sum_param_fit,
         }
 
+    def get_type_coverage(self) -> dict:
+        """Get model type (element) coverage.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the information of type coverage in the model:
+            - 'type_num': the total number of covered types in this model.
+            - 'covered_type': a list of the covered 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
+        return {
+            "type_num": bias_mask.sum(),
+            "covered_type": sort_element_type(type_map[bias_mask].tolist()),
+        }
+
     def eval_descriptor(
         self,
         coords: np.ndarray,

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

source/tests/pt/test_dp_show.py

@@ -47,43 +47,49 @@ def setUp(self) -> None:
 
     def test_checkpoint(self) -> None:
         INPUT = "model.pt"
-        ATTRIBUTES = "type-map descriptor fitting-net size"
+        ATTRIBUTES = "type-map descriptor fitting-net size type-coverage"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-8]
-        assert "The type_map is ['O', 'H', 'Au']" in results[-7]
+        assert "This is a singletask model" in results[-11]
+        assert "The type_map is ['O', 'H', 'Au']" in results[-10]
         assert (
             "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[-6]
+        ) in results[-9]
         assert (
             "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}"
-            in results[-5]
+            in results[-8]
         )
-        assert "Parameter counts:" in results[-4]
-        assert "Parameters in descriptor: 19,350" in results[-3]
-        assert "Parameters in fitting-net: 119,091" in results[-2]
-        assert "Parameters in total: 138,441" in results[-1]
+        assert "Parameter counts:" in results[-7]
+        assert "Parameters in descriptor: 19,350" in results[-6]
+        assert "Parameters in fitting-net: 119,091" in results[-5]
+        assert "Parameters in total: 138,441" in results[-4]
+        assert "The type coverage for this model:" in results[-3]
+        assert "Number of covered types: 2" in results[-2]
+        assert "Covered types: ['H', 'O']" in results[-1]
 
     def test_frozen_model(self) -> None:
         INPUT = "frozen_model.pth"
-        ATTRIBUTES = "type-map descriptor fitting-net size"
+        ATTRIBUTES = "type-map descriptor fitting-net size type-coverage"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-8]
-        assert "The type_map is ['O', 'H', 'Au']" in results[-7]
+        assert "This is a singletask model" in results[-11]
+        assert "The type_map is ['O', 'H', 'Au']" in results[-10]
         assert (
             "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[-6]
+        ) in results[-9]
         assert (
             "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}"
-            in results[-5]
+            in results[-8]
         )
-        assert "Parameter counts:" in results[-4]
-        assert "Parameters in descriptor: 19,350" in results[-3]
-        assert "Parameters in fitting-net: 119,091" in results[-2]
-        assert "Parameters in total: 138,441" in results[-1]
+        assert "Parameter counts:" in results[-7]
+        assert "Parameters in descriptor: 19,350" in results[-6]
+        assert "Parameters in fitting-net: 119,091" in results[-5]
+        assert "Parameters in total: 138,441" in results[-4]
+        assert "The type coverage for this model:" in results[-3]
+        assert "Number of covered types: 2" in results[-2]
+        assert "Covered types: ['H', 'O']" in results[-1]  # only covers two elements
 
     def test_checkpoint_error(self) -> None:
         INPUT = "model.pt"
@@ -152,62 +158,72 @@ def setUp(self) -> None:
 
     def test_checkpoint(self) -> None:
         INPUT = "model.ckpt.pt"
-        ATTRIBUTES = "model-branch type-map descriptor fitting-net size"
+        ATTRIBUTES = "model-branch type-map descriptor fitting-net size type-coverage"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a multitask model" in results[-12]
+        assert "This is a multitask model" in results[-19]
         assert (
             "Available model branches are ['model_1', 'model_2', 'RANDOM'], "
             "where 'RANDOM' means using a randomly initialized fitting net."
-            in results[-11]
+            in results[-18]
         )
-        assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[-10]
-        assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[-9]
+        assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[-17]
+        assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[-16]
         assert (
             "model_1"
             and "'type': 'se_e2_a'"
             and "'sel': [46, 92, 4]"
             and "'rcut_smth': 0.5"
-        ) in results[-8]
+        ) in results[-15]
         assert (
             "model_2"
             and "'type': 'se_e2_a'"
             and "'sel': [46, 92, 4]"
             and "'rcut_smth': 0.5"
-        ) in results[-7]
+        ) in results[-14]
         assert (
             "The fitting_net parameter of branch model_1 is {'neuron': [1, 2, 3], 'seed': 678}"
-            in results[-6]
+            in results[-13]
         )
         assert (
             "The fitting_net parameter of branch model_2 is {'neuron': [9, 8, 7], 'seed': 1111}"
-            in results[-5]
+            in results[-12]
         )
-        assert "Parameter counts for a single branch model:" in results[-4]
-        assert "Parameters in descriptor: 19,350" in results[-3]
-        assert "Parameters in fitting-net: 4,860" in results[-2]
-        assert "Parameters in total: 24,210" in results[-1]
+        assert "Parameter counts for a single branch model:" in results[-11]
+        assert "Parameters in descriptor: 19,350" in results[-10]
+        assert "Parameters in fitting-net: 4,860" in results[-9]
+        assert "Parameters in total: 24,210" in results[-8]
+        assert "The type coverage for each branch:" in results[-7]
+        assert "model_1: Number of covered types: 2" in results[-6]
+        assert "model_1: Covered types: ['H', 'O']" in results[-5]
+        assert "model_2: Number of covered types: 2" in results[-4]
+        assert "model_2: Covered types: ['H', 'O']" in results[-3]
+        assert "TOTAL number of covered types in the model: 2" in results[-2]
+        assert "TOTAL covered types in the model: ['H', 'O']" in results[-1]
 
     def test_frozen_model(self) -> None:
         INPUT = "frozen_model.pth"
-        ATTRIBUTES = "type-map descriptor fitting-net size"
+        ATTRIBUTES = "type-map descriptor fitting-net size type-coverage"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-8]
-        assert "The type_map is ['O', 'H', 'B']" in results[-7]
+        assert "This is a singletask model" in results[-11]
+        assert "The type_map is ['O', 'H', 'B']" in results[-10]
         assert (
             "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5"
-        ) in results[-6]
+        ) in results[-9]
         assert (
             "The fitting_net parameter is {'neuron': [1, 2, 3], 'seed': 678}"
-            in results[-5]
+            in results[-8]
         )
-        assert "Parameter counts:" in results[-4]
-        assert "Parameters in descriptor: 19,350" in results[-3]
-        assert "Parameters in fitting-net: 4,860" in results[-2]
-        assert "Parameters in total: 24,210" in results[-1]
+        assert "Parameter counts:" in results[-7]
+        assert "Parameters in descriptor: 19,350" in results[-6]
+        assert "Parameters in fitting-net: 4,860" in results[-5]
+        assert "Parameters in total: 24,210" in results[-4]
+        assert "The type coverage for this model:" in results[-3]
+        assert "Number of covered types: 2" in results[-2]
+        assert "Covered types: ['H', 'O']" in results[-1]  # only covers two elements
 
     def tearDown(self) -> None:
         for f in os.listdir("."):