feat(pt, dpmodel): use data stat for observed type

Author: iProzd

deepmd/dpmodel/model/base_model.py

@@ -142,6 +142,20 @@ def get_model_def_script(self) -> str:
             """Get the model definition script."""
             pass
 
+        def get_observed_type_list(self) -> list[str]:
+            """Get observed types from model metadata.
+
+            Returns empty list if not available.
+            """
+            if self.model_def_script:
+                import json
+
+                params = json.loads(self.model_def_script)
+                observed = params.get("info", {}).get("observed_type")
+                if observed is not None:
+                    return observed
+            return []
+
         def get_min_nbor_dist(self) -> float | None:
             """Get the minimum distance between two atoms."""
             return self.min_nbor_dist

deepmd/dpmodel/utils/stat.py

@@ -26,6 +26,62 @@
 log = logging.getLogger(__name__)
 
 
+def collect_observed_types(
+    sampled: list[dict], type_map: list[str]
+) -> list[str]:
+    """Collect observed element types from sampled training data.
+
+    Parameters
+    ----------
+    sampled : list[dict]
+        Sampled data from different data systems. Each dict must contain
+        ``"atype"`` with shape ``[nframes, natoms]``.
+    type_map : list[str]
+        Mapping from type index to element symbol.
+
+    Returns
+    -------
+    list[str]
+        Sorted list of observed element symbols.
+    """
+    from deepmd.utils.econf_embd import (
+        sort_element_type,
+    )
+
+    observed_indices: set[int] = set()
+    for system in sampled:
+        atype = to_numpy_array(system["atype"])  # shape: [nframes, natoms]
+        observed_indices.update(np.unique(atype).tolist())
+    observed_types = [type_map[i] for i in sorted(observed_indices) if i < len(type_map)]
+    return sort_element_type(observed_types)
+
+
+def _restore_observed_type_from_file(
+    stat_file_path: DPPath | None,
+) -> list[str] | None:
+    """Try to load observed_type from stat file."""
+    if stat_file_path is None:
+        return None
+    fp = stat_file_path / "observed_type"
+    if fp.is_file():
+        arr = fp.load_numpy()
+        # Decode bytes back to str if stored as bytes (for h5py compatibility)
+        return [x.decode() if isinstance(x, bytes) else x for x in arr.tolist()]
+    return None
+
+
+def _save_observed_type_to_file(
+    stat_file_path: DPPath | None, observed_type: list[str]
+) -> None:
+    """Save observed_type to stat file."""
+    if stat_file_path is None:
+        return
+    stat_file_path.mkdir(exist_ok=True, parents=True)
+    fp = stat_file_path / "observed_type"
+    # Use bytes dtype for h5py compatibility (h5py cannot store Unicode strings)
+    fp.save_numpy(np.array(observed_type, dtype="S"))
+
+
 def _restore_from_file(
     stat_file_path: DPPath,
     keys: list[str],

deepmd/entrypoints/show.py

@@ -126,6 +126,13 @@ def show(
             )
         else:
             log.info("The observed types for this model: ")
-            observed_types = model.get_observed_types()
+            observed_type_list = model_params.get("info", {}).get("observed_type")
+            if observed_type_list is not None:
+                observed_types = {
+                    "type_num": len(observed_type_list),
+                    "observed_type": observed_type_list,
+                }
+            else:
+                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/pt/infer/deep_eval.py

@@ -736,6 +736,16 @@ 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.
         """
+        # Try metadata first (from model_def_script, already a dict)
+        observed_type_list = (
+            self.model_def_script.get("info", {}).get("observed_type")
+        )
+        if observed_type_list is not None:
+            return {
+                "type_num": len(observed_type_list),
+                "observed_type": observed_type_list,
+            }
+        # Fallback: bias-based approach for old models
         observed_type_list = self.dp.model["Default"].get_observed_type_list()
         return {
             "type_num": len(observed_type_list),

deepmd/pt/model/atomic_model/base_atomic_model.py

@@ -90,6 +90,7 @@ def __init__(
         self.rcond = rcond
         self.preset_out_bias = preset_out_bias
         self.data_stat_protect = data_stat_protect
+        self._observed_type: list[str] | None = None
 
     def init_out_stat(self) -> None:
         """Initialize the output bias."""

deepmd/pt/model/atomic_model/dp_atomic_model.py

@@ -20,6 +20,11 @@
 from deepmd.pt.model.task.base_fitting import (
     BaseFitting,
 )
+from deepmd.pt.utils.stat import (
+    _restore_observed_type_from_file,
+    _save_observed_type_to_file,
+    collect_observed_types,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -307,6 +312,7 @@ def compute_or_load_stat(
         sampled_func: Callable[[], list[dict]],
         stat_file_path: DPPath | None = None,
         compute_or_load_out_stat: bool = True,
+        preset_observed_type: list[str] | None = None,
     ) -> None:
         """
         Compute or load the statistics parameters of the model,
@@ -358,6 +364,17 @@ def wrapped_sampler() -> list[dict]:
         if compute_or_load_out_stat:
             self.compute_or_load_out_stat(wrapped_sampler, stat_file_path)
 
+        # Collect observed types with priority: preset > stat_file > compute
+        if preset_observed_type is not None:
+            self._observed_type = preset_observed_type
+        else:
+            observed = _restore_observed_type_from_file(stat_file_path)
+            if observed is None:
+                sampled = wrapped_sampler()
+                observed = collect_observed_types(sampled, self.type_map)
+                _save_observed_type_to_file(stat_file_path, observed)
+            self._observed_type = observed
+
     def compute_fitting_input_stat(
         self,
         sample_merged: Callable[[], list[dict]] | list[dict],

deepmd/pt/model/model/make_model.py

@@ -587,9 +587,14 @@ def compute_or_load_stat(
             self,
             sampled_func: Callable[[], Any],
             stat_file_path: DPPath | None = None,
+            preset_observed_type: list[str] | None = None,
         ) -> None:
             """Compute or load the statistics."""
-            return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path)
+            return self.atomic_model.compute_or_load_stat(
+                sampled_func,
+                stat_file_path,
+                preset_observed_type=preset_observed_type,
+            )
 
         def get_sel(self) -> list[int]:
             """Returns the number of selected atoms for each type."""

deepmd/pt/model/model/model.py

@@ -30,6 +30,7 @@ def compute_or_load_stat(
         self,
         sampled_func: Any,
         stat_file_path: DPPath | None = None,
+        preset_observed_type: list[str] | None = None,
     ) -> NoReturn:
         """
         Compute or load the statistics parameters of the model,

deepmd/pt/train/training.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import functools
+import json
 import logging
 import time
 from collections.abc import (
@@ -288,6 +289,7 @@ def single_model_stat(
             _training_data: DpLoaderSet,
             _stat_file_path: str | None,
             finetune_has_new_type: bool = False,
+            preset_observed_type: list[str] | None = None,
         ) -> Callable[[], Any]:
             @functools.lru_cache
             def get_sample() -> Any:
@@ -302,6 +304,7 @@ def get_sample() -> Any:
                 _model.compute_or_load_stat(
                     sampled_func=get_sample,
                     stat_file_path=_stat_file_path,
+                    preset_observed_type=preset_observed_type,
                 )
                 if isinstance(_stat_file_path, DPH5Path):
                     _stat_file_path.root.close()
@@ -394,7 +397,16 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                 finetune_has_new_type=self.finetune_links["Default"].get_has_new_type()
                 if self.finetune_links is not None
                 else False,
+                preset_observed_type=model_params.get("info", {}).get("observed_type"),
             )
+            # Persist observed_type from stat into model_params and model_def_script
+            if not resuming and self.rank == 0:
+                observed = getattr(
+                    self.model.atomic_model, "_observed_type", None
+                )
+                if observed is not None:
+                    model_params.setdefault("info", {})["observed_type"] = observed
+                    self.model.model_def_script = json.dumps(model_params)
             (
                 self.training_dataloader,
                 self.training_data,
@@ -432,6 +444,11 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                 training_data[model_key].preload_and_modify_all_data_torch()
                 if validation_data[model_key] is not None:
                     validation_data[model_key].preload_and_modify_all_data_torch()
+                _mt_user_observed = (
+                    model_params["model_dict"][model_key]
+                    .get("info", {})
+                    .get("observed_type")
+                )
                 self.get_sample_func[model_key] = single_model_stat(
                     self.model[model_key],
                     model_params["model_dict"][model_key].get("data_stat_nbatch", 10),
@@ -442,7 +459,22 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                     ].get_has_new_type()
                     if self.finetune_links is not None
                     else False,
+                    preset_observed_type=_mt_user_observed,
                 )
+                # Persist observed_type into model_params and model_def_script
+                if not resuming and self.rank == 0:
+                    observed = getattr(
+                        self.model[model_key].atomic_model,
+                        "_observed_type",
+                        None,
+                    )
+                    if observed is not None:
+                        model_params["model_dict"][model_key].setdefault(
+                            "info", {}
+                        )["observed_type"] = observed
+                        self.model[model_key].model_def_script = json.dumps(
+                            model_params["model_dict"][model_key]
+                        )
 
                 (
                     self.training_dataloader[model_key],

deepmd/pt/utils/stat.py

@@ -35,6 +35,19 @@
 
 log = logging.getLogger(__name__)
 
+# Re-export from dpmodel (backend-agnostic implementations)
+from deepmd.dpmodel.utils.stat import (  # noqa: E402
+    _restore_observed_type_from_file,
+    _save_observed_type_to_file,
+    collect_observed_types,
+)
+
+__all__ = [
+    "collect_observed_types",
+    "_restore_observed_type_from_file",
+    "_save_observed_type_to_file",
+]
+
 
 def make_stat_input(
     datasets: list[Any], dataloaders: list[Any], nbatches: int