(fix) Make bias statistics complete for all elements

deepmd/pt/train/training.py

@@ -142,6 +142,9 @@ def __init__(
         self.max_ckpt_keep = training_params.get("max_ckpt_keep", 5)
         self.display_in_training = training_params.get("disp_training", True)
         self.timing_in_training = training_params.get("time_training", True)
+        self.min_frames_per_element_forstat = training_params.get(
+            "min_frames_per_element_forstat", 10
+        )
         self.change_bias_after_training = training_params.get(
             "change_bias_after_training", False
         )
@@ -226,6 +229,7 @@ def get_sample():
                     _training_data.systems,
                     _training_data.dataloaders,
                     _data_stat_nbatch,
+                    self.min_frames_per_element_forstat,
                 )
                 return sampled
 

deepmd/pt/utils/dataset.py

@@ -1,10 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 
 
+import glob
+import os
+from collections import (
+    defaultdict,
+)
 from typing import (
     Optional,
 )
 
+import numpy as np
 from torch.utils.data import (
     Dataset,
 )
@@ -13,14 +19,17 @@
     DataRequirementItem,
     DeepmdData,
 )
+from deepmd.utils.path import (
+    DPPath,
+)
 
 
 class DeepmdDataSetForLoader(Dataset):
     def __init__(self, system: str, type_map: Optional[list[str]] = None) -> None:
-        """Construct DeePMD-style dataset containing frames cross different systems.
+        """Construct DeePMD-style dataset containing frames across different systems.
 
         Args:
-        - systems: Paths to systems.
+        - system: Path to the system.
         - type_map: Atom types.
         """
         self.system = system
@@ -40,6 +49,25 @@ def __getitem__(self, index):
         b_data["natoms"] = self._natoms_vec
         return b_data
 
+    def true_types(self):
+        """Identify and count unique element types present in the dataset,
+        and count the number of frames each element appears in.
+        """
+        element_counts = defaultdict(lambda: {"count": 0, "frames": 0})
+        set_pattern = os.path.join(self.system, "set.*")
+        set_files = sorted(glob.glob(set_pattern))
+        for set_file in set_files:
+            element_data = self._data_system._load_type_mix(DPPath(set_file))
+            unique_elements, counts = np.unique(element_data, return_counts=True)
+            for elem, cnt in zip(unique_elements, counts):
+                element_counts[elem]["count"] += cnt
+            for elem in unique_elements:
+                frames_with_elem = np.any(element_data == elem, axis=1)
+                row_count = np.sum(frames_with_elem)
+                element_counts[elem]["frames"] += row_count
+        element_counts = dict(element_counts)
+        return element_counts
+
     def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
         """Add data requirement for this data system."""
         for data_item in data_requirement:

deepmd/pt/utils/stat.py

@@ -36,8 +36,7 @@
 
 log = logging.getLogger(__name__)
 
-
-def make_stat_input(datasets, dataloaders, nbatches):
+def make_stat_input(datasets, dataloaders, nbatches, min_frames_per_element_forstat):
     """Pack data for statistics.
 
     Args:
@@ -50,28 +49,131 @@ def make_stat_input(datasets, dataloaders, nbatches):
     """
     lst = []
     log.info(f"Packing data for statistics from {len(datasets)} systems")
-    for i in range(len(datasets)):
-        sys_stat = {}
-        with torch.device("cpu"):
-            iterator = iter(dataloaders[i])
-            numb_batches = min(nbatches, len(dataloaders[i]))
-            for _ in range(numb_batches):
-                try:
-                    stat_data = next(iterator)
-                except StopIteration:
-                    iterator = iter(dataloaders[i])
-                    stat_data = next(iterator)
-                for dd in stat_data:
-                    if stat_data[dd] is None:
-                        sys_stat[dd] = None
-                    elif isinstance(stat_data[dd], torch.Tensor):
-                        if dd not in sys_stat:
-                            sys_stat[dd] = []
-                        sys_stat[dd].append(stat_data[dd])
-                    elif isinstance(stat_data[dd], np.float32):
-                        sys_stat[dd] = stat_data[dd]
+    collect_elements = set()
+    total_element_types = set()
+    total_element_counts = {}
+    if datasets[0].mixed_type:
+        for sys_index, (dataset, dataloader) in enumerate(zip(datasets, dataloaders)):
+            sys_stat = {}
+            with torch.device("cpu"):
+                iterator = iter(dataloader)
+                numb_batches = min(nbatches, len(dataloader))
+                for _ in range(numb_batches):
+                    try:
+                        stat_data = next(iterator)
+                    except StopIteration:
+                        iterator = iter(dataloader)
+                        stat_data = next(iterator)
+                    for dd in stat_data:
+                        if stat_data[dd] is None:
+                            sys_stat[dd] = None
+                        elif isinstance(stat_data[dd], torch.Tensor):
+                            if dd not in sys_stat:
+                                sys_stat[dd] = []
+                            sys_stat[dd].append(stat_data[dd])
+                        elif isinstance(stat_data[dd], np.float32):
+                            sys_stat[dd] = stat_data[dd]
+                        else:
+                            pass
+                if 'atype' in sys_stat and isinstance(sys_stat['atype'], list):
+                    collect_values = np.unique(torch.cat(sys_stat['atype']).numpy())
+                    collect_elements.update(collect_values)
+
+            for key in sys_stat:
+                if isinstance(sys_stat[key], np.float32):
+                    pass
+                elif sys_stat[key] is None or (isinstance(sys_stat[key], list) and (len(sys_stat[key]) == 0 or sys_stat[key][0] is None)):
+                    sys_stat[key] = None
+                elif isinstance(sys_stat[key][0], torch.Tensor):
+                    sys_stat[key] = torch.cat(sys_stat[key], dim=0)
+            dict_to_device(sys_stat)
+            lst.append(sys_stat)
+
+            element_counts = dataset.true_types()
+            for elem, data in element_counts.items():
+                count = data["count"]
+                frames = data["frames"]
+                total_element_types.add(elem)
+                if elem not in total_element_counts:
+                    total_element_counts[elem] = {"count": 0, "frames": 0, "indices": []}
+                total_element_counts[elem]["count"] += count
+                if len(total_element_counts[elem]["indices"]) < min_frames_per_element_forstat:
+                    total_element_counts[elem]["indices"].append({
+                        "sys_index": sys_index,
+                        "frames": frames
+                    })
+        for elem, data in total_element_counts.items():
+            count = data["count"]
+            indices_count = len(data["indices"])
+            if indices_count < min_frames_per_element_forstat:
+                log.warning(f'The number of frame with element {elem} is {indices_count}, which is less than the expected maximum value {min_frames_per_element_forstat}')
+        missing_elements = total_element_types - collect_elements
+        for miss in missing_elements:
+            sys_indices = total_element_counts[miss].get('indices', [])
+            for sys_info in sys_indices:
+                sys_index = sys_info['sys_index']
+                frames = sys_info['frames']
+                sys = datasets[sys_index]
+                frame_data = sys.__getitem__(frames)
+                sys_stat_new = {}
+                for dd in frame_data:
+                    if dd == "type":
+                        continue
+                    if frame_data[dd] is None:
+                        sys_stat_new[dd] = None
+                    elif isinstance(frame_data[dd], np.ndarray):
+                        if dd not in sys_stat_new:
+                            sys_stat_new[dd] = []
+                        frame_data[dd] = torch.from_numpy(frame_data[dd])
+                        frame_data[dd] = frame_data[dd].unsqueeze(0)
+                        sys_stat_new[dd].append(frame_data[dd])
+                    elif isinstance(frame_data[dd], np.float32):
+                        sys_stat_new[dd] = frame_data[dd]
                     else:
                         pass
+<<<<<<< HEAD
+                for key in sys_stat_new:
+                    if isinstance(sys_stat_new[key], np.float32):
+                        pass
+                    elif sys_stat_new[key] is None or sys_stat_new[key][0] is None:
+                        sys_stat_new[key] = None
+                    elif isinstance(frame_data[dd], torch.Tensor):
+                        sys_stat_new[key] = torch.cat(sys_stat_new[key], dim=0)
+                dict_to_device(sys_stat_new)
+                lst.append(sys_stat_new)
+    else:
+        for i in range(len(datasets)):
+            sys_stat = {}
+            with torch.device("cpu"):
+                iterator = iter(dataloaders[i])
+                numb_batches = min(nbatches, len(dataloaders[i]))
+                for _ in range(numb_batches):
+                    try:
+                        stat_data = next(iterator)
+                    except StopIteration:
+                        iterator = iter(dataloaders[i])
+                        stat_data = next(iterator)
+                    for dd in stat_data:
+                        if stat_data[dd] is None:
+                            sys_stat[dd] = None
+                        elif isinstance(stat_data[dd], torch.Tensor):
+                            if dd not in sys_stat:
+                                sys_stat[dd] = []
+                            sys_stat[dd].append(stat_data[dd])
+                        elif isinstance(stat_data[dd], np.float32):
+                            sys_stat[dd] = stat_data[dd]
+                        else:
+                            pass
+            for key in sys_stat:
+                if isinstance(sys_stat[key], np.float32):
+                    pass
+                elif sys_stat[key] is None or sys_stat[key][0] is None:
+                    sys_stat[key] = None
+                elif isinstance(stat_data[dd], torch.Tensor):
+                    sys_stat[key] = torch.cat(sys_stat[key], dim=0)
+            dict_to_device(sys_stat)
+            lst.append(sys_stat)
+=======
 
         for key in sys_stat:
             if isinstance(sys_stat[key], np.float32):
@@ -82,6 +184,51 @@ def make_stat_input(datasets, dataloaders, nbatches):
                 sys_stat[key] = torch.cat(sys_stat[key], dim=0)
         dict_to_device(sys_stat)
         lst.append(sys_stat)
+
+    collect_elements = set()
+    all_element = set()
+    for i in lst:
+        collect_values = np.unique(i["atype"].cpu().numpy())
+        collect_elements.update(collect_values)
+    for i in datasets:
+        all_elements_in_dataset = i.get_all_atype
+        all_element.update(all_elements_in_dataset)
+    missing_element = all_element - collect_elements
+    for miss in missing_element:
+        for i in datasets:
+            if i.element_to_frames.get(miss, []) is not None:
+                frame_indices = i.element_to_frames.get(miss, [])
+                frame_data = i.__getitem__(frame_indices[0])
+                break
+            else:
+                pass
+        sys_stat_new = {}
+        for dd in frame_data:
+            if dd == "type":
+                continue
+            if frame_data[dd] is None:
+                sys_stat_new[dd] = None
+            elif isinstance(frame_data[dd], np.ndarray):
+                if dd not in sys_stat_new:
+                    sys_stat_new[dd] = []
+                frame_data[dd] = torch.from_numpy(frame_data[dd])
+                frame_data[dd] = frame_data[dd].unsqueeze(0)
+                sys_stat_new[dd].append(frame_data[dd])
+            elif isinstance(stat_data[dd], np.float32):
+                sys_stat_new[dd] = frame_data[dd]
+            else:
+                pass
+        for key in sys_stat_new:
+            if isinstance(sys_stat_new[key], np.float32):
+                pass
+            elif sys_stat_new[key] is None or sys_stat_new[key][0] is None:
+                sys_stat_new[key] = None
+            elif isinstance(stat_data[dd], torch.Tensor):
+                sys_stat_new[key] = torch.cat(sys_stat_new[key], dim=0)
+        dict_to_device(sys_stat_new)
+        lst.append(sys_stat_new)
+
+>>>>>>> dc6430730dd18087d0b6fafd98c5e4add795a57a
     return lst
 
 

deepmd/utils/argcheck.py

@@ -2826,6 +2826,12 @@ def training_args(
             optional=True,
             doc=doc_only_pt_supported + doc_gradient_max_norm,
         ),
+        Argument(
+            "min_frames_per_element_forstat",
+            int,
+            optional=True,
+            doc="The minimum number of frames per element used for statistics.",
+        ),
     ]
     variants = [
         Variant(

deepmd/utils/data.py

@@ -530,14 +530,6 @@ def _load_set(self, set_name: DPPath):
         if self.mixed_type:
             # nframes x natoms
             atom_type_mix = self._load_type_mix(set_name)
-            if self.enforce_type_map:
-                try:
-                    atom_type_mix_ = self.type_idx_map[atom_type_mix].astype(np.int32)
-                except IndexError as e:
-                    raise IndexError(
-                        f"some types in 'real_atom_types.npy' of set {set_name} are not contained in {self.get_ntypes()} types!"
-                    ) from e
-                atom_type_mix = atom_type_mix_
             real_type = atom_type_mix.reshape([nframes, self.natoms])
             data["type"] = real_type
             natoms = data["type"].shape[1]
@@ -672,6 +664,14 @@ def _load_type(self, sys_path: DPPath):
     def _load_type_mix(self, set_name: DPPath):
         type_path = set_name / "real_atom_types.npy"
         real_type = type_path.load_numpy().astype(np.int32).reshape([-1, self.natoms])
+        if self.enforce_type_map:
+            try:
+                atom_type_mix_ = self.type_idx_map[real_type].astype(np.int32)
+            except IndexError as e:
+                raise IndexError(
+                    f"some types in 'real_atom_types.npy' of set {set_name} are not contained in {self.get_ntypes()} types!"
+                ) from e
+            real_type = atom_type_mix_
         return real_type
 
     def _make_idx_map(self, atom_type):