deepmd/pt/utils/stat.py (1)

110-119: Avoid potential index-out-of-range errors when concatenating Tensors.
The code checks “isinstance(sys_stat[key], list) and isinstance(sys_stat[key][0], torch.Tensor)” without verifying that the list is non-empty. This might raise an IndexError if the list is empty. Consider adding a length check before accessing sys_stat[key][0].

Possible safe check:

-if isinstance(sys_stat[key], list) and isinstance(sys_stat[key][0], torch.Tensor):
+if isinstance(sys_stat[key], list) and len(sys_stat[key]) > 0 \
+   and isinstance(sys_stat[key][0], torch.Tensor):

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

#!/bin/bash
# Search for patterns related to 'atype' handling and array emptiness checks
rg -A 3 "atype" --type python

# Also look for the specific class/function containing this code
ast-grep --pattern $'def $_(self, $_):
    $$$
    collected_elements = set()
    for sys_stat in lst:
        if "atype" in sys_stat:
            collected_elements.update($$$)
    $$$'

#!/bin/bash
# Search for patterns related to 'atype' handling and array emptiness checks
rg -A 3 "atype" -t py

# Look for similar array handling patterns in the codebase
rg -t py "numpy\.unique|\.update.*numpy" -A 2

# Find the specific file we're analyzing
fd "stat.py" --type f

-        print('we want', all_elements)
+        log.info(f"We want elements: {all_elements}")

#!/bin/bash
# Let's first look at the broader context of this code
rg -B 10 -A 5 "for missing_element in missing_elements" deepmd/pt/utils/stat.py

# Also search for any related tests that might help understand the intended behavior
fd "test.*stat.*py" -x rg -l "missing_element"

# Look for the initialization of nbatches and lst variables
rg -B 5 "nbatches.*=|lst.*=" deepmd/pt/utils/stat.py

#!/bin/bash
# Search for the function definition and its parameters
ast-grep --pattern 'def $FUNC($PARAMS) {
  $$$
  nbatches
  $$$
}'  deepmd/pt/utils/stat.py

# Look for any documentation or comments about nbatches
rg -B 3 -A 3 "nbatches" deepmd/pt/utils/stat.py

# Check if there are any tests that verify this behavior
fd "test.*stat.*py" -x cat {} 

#!/bin/bash
# Look for the function definition and documentation
rg -B 10 "def make_stat_input" deepmd/pt/utils/stat.py

# Search for any tests that verify the behavior with missing elements
rg -B 5 -A 5 "missing_element" deepmd/pt/utils/stat.py

#!/bin/bash
# Look for the full function implementation and its docstring
rg -B 20 -A 50 "def make_stat_input" deepmd/pt/utils/stat.py

# Search for any tests that verify the statistics collection
fd "test.*stat.*py" -x rg -A 10 "make_stat_input"

# Look for the initialization of lst and how frames are collected
rg -B 5 "lst.*=.*\[\]" deepmd/pt/utils/stat.py

deepmd/pt/utils/stat.py (2)

89-89: Remove debug print statement.

The print statement print("we want", all_elements) appears to be debug code that should be removed or replaced with proper logging.

-        print("we want", all_elements)
+        log.debug(f"Required elements for statistics: {all_elements}")

---

97-111: Optimize nested loops and add error handling.

The nested loops for handling missing elements could be optimized, and error handling should be added for invalid frame indices.

         for missing_element in missing_elements:
             for i, dataset in enumerate(datasets):
                 if hasattr(dataset, "element_to_frames"):
-                    frame_indices = dataset.element_to_frames.get(
-                        missing_element, []
-                    )
+                    try:
+                        frame_indices = dataset.element_to_frames.get(missing_element, [])
+                        if not frame_indices:
+                            continue
+                            
+                        # Pre-check if we need more frames
+                        if len(lst[i]["atype"]) >= nbatches:
+                            break
+                            
+                        # Process frames in batch
+                        for frame_idx in frame_indices:
+                            frame_data = dataset[frame_idx]
+                            if any(key not in lst[i] for key in frame_data):
+                                lst[i].update({key: [] for key in frame_data if key not in lst[i]})
+                            for key in frame_data:
+                                lst[i][key].append(frame_data[key])
+                            if len(lst[i]["atype"]) >= nbatches:
+                                break
+                    except Exception as e:
+                        log.warning(f"Error processing frames for element {missing_element}: {e}")
+                        continue
-                    for frame_idx in frame_indices:
-                        if len(lst[i]["atype"]) >= nbatches:
-                            break
-                        frame_data = dataset[frame_idx]
-                        for key in frame_data:
-                            if key not in lst[i]:
-                                lst[i][key] = []
-                            lst[i][key].append(frame_data[key])

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

-        def _build_element_to_frames(self):
-        """Mapping element types to frame indexes"""
+    def _build_element_to_frames(self):
+        """Mapping element types to frame indexes"""

-    def _build_element_to_frames(self):
-        """Mapping element types to frame indexes"""
+    def _build_element_to_frames(self, max_frames_per_element: int = 10) -> dict[int, list[int]]:
+        """Build a mapping of element types to their corresponding frame indices.
+        
+        Args:
+            max_frames_per_element: Maximum number of frames to store per element type.
+            
+        Returns:
+            A dictionary mapping element types (int) to lists of frame indices (list[int])
+            where each element type appears.
+        """
         element_to_frames = {element: [] for element in range(self._ntypes)}
         for frame_idx in range(len(self)):
             frame_data = self._data_system.get_item_torch(frame_idx)
 
             elements = frame_data["atype"]
             for element in set(elements):
-                if len(element_to_frames[element]) < 10:
+                if len(element_to_frames[element]) < max_frames_per_element:
                     element_to_frames[element].append(frame_idx)
         return element_to_frames

source/tests/pt/test_make_stat_input.py (4)

14-23: Consider using collections.defaultdict for element_to_frames.
You can simplify the nested checks for element presence in the dictionary by using a defaultdict(list), which would eliminate the need for the explicit if atype not in self.element_to_frames: condition.

-from collections import defaultdict

class TestDataset:
    def __init__(self, samples):
        self.samples = samples
-        self.element_to_frames = {}
+        from collections import defaultdict
+        self.element_to_frames = defaultdict(list)
        for idx, sample in enumerate(samples):
            atypes = sample["atype"]
            for atype in atypes:
-                if atype not in self.element_to_frames:
-                    self.element_to_frames[atype] = []
                self.element_to_frames[atype].append(idx)

---

25-28: Rename the property to better reflect usage.
Using @property but naming it get_all_atype can be confusing. Consider a more descriptive name like all_atypes, since Python properties typically avoid "get_" prefixes.

---

53-59: Remove or use the assigned lst variable.
The variable lst is assigned but never used, according to static analysis hints. Consider removing it or using it for additional assertions.

 def test_make_stat_input(self):
     nbatches = 1
-    lst = make_stat_input(self.datasets, self.dataloaders, nbatches=nbatches)
+    _ = make_stat_input(self.datasets, self.dataloaders, nbatches=nbatches)
     all_elements = self.system.get_all_atype
     unique_elements = {1, 2}
     self.assertEqual(unique_elements, all_elements, "make_stat_input miss elements")

🧰 Tools

🪛 Ruff (0.8.2)

55-55: Local variable lst is assigned to but never used

Remove assignment to unused variable lst

(F841)

---

61-62: Optional test runner inclusion.
Having the if __name__ == "__main__": unittest.main() block is fine. You could remove it if tests are run by a dedicated test runner.

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/utils/data.py (1)

667-674: Consider providing more informative error details when the type map lookup fails.
While raising IndexError is appropriate, developers might benefit from including the failing value(s). You could, for instance, collect and display the out-of-range types to provide immediate troubleshooting clues. A custom exception or a more descriptive error message can significantly improve clarity and debuggability.

Here's an example of how you might refine the exception:

            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
+                # Gather all invalid elements
+                invalid_types = np.unique(real_type[(real_type < 0) | (real_type >= len(self.type_idx_map))])
+                raise IndexError(
+                    f"One or more invalid types found in 'real_atom_types.npy' of set {set_name}: {invalid_types}. "
+                    f"Ensure all types are within [0, {self.get_ntypes()-1}]."
+                ) from e

deepmd/utils/argcheck.py (1)

2829-2834: Ensure user awareness of the new argument.

The new argument min_frames_per_element_forstat is useful for controlling statistic completeness. It might be helpful to specify the expected range (e.g., must be ≥ 1) and how large values impact memory or performance overhead.

source/tests/pt/test_make_stat_input.py (1)

68-68: Remove or utilize the unused variable.

The variable lst is assigned with the result of make_stat_input(...) but never used. If no further checks are applied, remove it to keep the code clean.

-        lst = make_stat_input(
+        make_stat_input(

🧰 Tools

🪛 Ruff (0.8.2)

68-68: Local variable lst is assigned to but never used

Remove assignment to unused variable lst

(F841)

deepmd/pt/utils/stat.py (1)

188-197: Double-check sets for collected vs. missing elements.

This code block re-checks missing elements with:

missing_element = all_element - collect_elements

Confirm that the logic aligns with the earlier missing_elements sets in lines 110–111 to avoid confusion or duplication.

🧰 Tools

🪛 Ruff (0.8.2)

188-188: SyntaxError: unindent does not match any outer indentation level

---

189-189: SyntaxError: Unexpected indentation

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/utils/argcheck.py (1)

2903-2909: Consider enhancing the documentation.

While the parameter implementation looks good, the documentation could be more descriptive about:

• What exactly is being checked when elements are validated
• The implications of enabling/disabling this feature
• Any potential impact on performance or behavior

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/utils/stat.py (4)

38-58: Enhance docstring with parameter types and examples.

The docstring could be improved by:

1. Adding type hints for parameters in the docstring
2. Including an example usage section

    """Pack data for statistics.
       Element checking is only enabled with mixed_type.

    Args:
-    - datasets: A list of datasets to analyze.
-    - dataloaders: Corresponding dataloaders for the datasets.
-    - nbatches: Batch count for collecting stats.
-    - min_frames_per_element_forstat: Minimum frames required for statistics.
-    - enable_element_completion: Whether to perform missing element completion (default: True).
+    - datasets: list[Dataset]
+        A list of datasets to analyze.
+    - dataloaders: list[DataLoader]
+        Corresponding dataloaders for the datasets.
+    - nbatches: int
+        Batch count for collecting stats.
+    - min_frames_per_element_forstat: int, optional
+        Minimum frames required for statistics. Defaults to 10.
+    - enable_element_completion: bool, optional
+        Whether to perform missing element completion. Defaults to True.

    Returns
    -------
    list[dict]
        A list of dicts, each of which contains data from a system.

+    Example
+    -------
+    >>> stats = make_stat_input(datasets, dataloaders, 5, 
+    ...                        min_frames_per_element_forstat=10,
+    ...                        enable_element_completion=True)
    """

---

64-74: Improve logging structure and clarity.

The logging messages could be more structured and informative.

     if datasets[0].mixed_type:
         if enable_element_completion:
             log.info(
-                f"Element check enabled. "
-                f"Verifying if frames with elements meet the set of {min_frames_per_element_forstat}."
+                f"Element completion: enabled (minimum frames per element: {min_frames_per_element_forstat})"
             )
         else:
-            log.info(
-                "Element completion is disabled. Skipping missing element handling."
-            )
+            log.info("Element completion: disabled")

---

130-143: Add type hints to helper function.

The function lacks type hints which would improve code maintainability.

-    def finalize_stats(sys_stat):
+    def finalize_stats(sys_stat: dict[str, Union[None, np.float32, list[torch.Tensor]]]) -> None:
         """Finalize statistics by concatenating tensors."""

---

144-157: Consider memory optimization for large datasets.

The current implementation loads all tensors into memory. Consider processing in chunks for large datasets.

Consider implementing a chunked processing approach:

def process_in_chunks(dataloader, chunk_size=1000):
    sys_stat = {}
    for i in range(0, len(dataloader), chunk_size):
        chunk = itertools.islice(dataloader, i, i + chunk_size)
        process_chunk(chunk, sys_stat)
    return sys_stat

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

source/tests/pt/test_make_stat_input.py (1)

27-28: Optimize dictionary key iteration.

Use direct dictionary key iteration instead of calling .keys().

-    for key in batch[0].keys():
-        data_list = [d[key] for d in batch]
+    for key in batch[0]:
+        data_list = [d[key] for d in batch]

🧰 Tools

🪛 Ruff (0.8.2)

27-27: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

source/tests/pt/test_make_stat_input.py (1)

31-31: Optimize dictionary key lookup.

Minor optimization opportunity in the key lookup.

Apply this diff to optimize the key lookup:

-    for key in batch[0].keys():
+    for key in batch[0]:

🧰 Tools

🪛 Ruff (0.8.2)

31-31: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

source/tests/pt/test_make_stat_input.py (2)

24-35: Add type hints and docstring to the collate function.

The function would benefit from type hints and documentation explaining its purpose and parameters.

-def collate_fn(batch):
+def collate_fn(batch: Union[dict, List[dict]]) -> dict:
+    """Collate a batch of data samples into a single batch for PyTorch DataLoader.
+    
+    Args:
+        batch: A single dict or list of dicts containing numpy arrays or scalar values
+        
+    Returns:
+        dict: Collated batch with torch tensors
+    """

Also, optimize the key lookup:

-    for key in batch[0].keys():
+    for key in batch[0]:

🧰 Tools

🪛 Ruff (0.8.2)

28-28: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

---

75-76: Add type hints and docstring to helper method.

The helper method would benefit from type hints and documentation.

-    def count_non_zero_elements(self, tensor, threshold=1e-8):
+    def count_non_zero_elements(self, tensor: torch.Tensor, threshold: float = 1e-8) -> int:
+        """Count elements in tensor with absolute value greater than threshold.
+        
+        Args:
+            tensor: Input tensor
+            threshold: Minimum absolute value to consider non-zero
+            
+        Returns:
+            int: Number of elements with absolute value > threshold
+        """

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

source/tests/pt/test_make_stat_input.py (1)

25-36: Minor optimization for dict key lookup.

The code is functionally correct, but there's a small optimization opportunity.

-    for key in batch[0].keys():
+    for key in batch[0]:

🧰 Tools

🪛 Ruff (0.8.2)

29-29: Use key in dict instead of key in dict.keys()

Remove .keys()

(SIM118)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

deepmd/pt/entrypoints/main.py (2)

389-389: Consider renaming the parameter for consistency.

The parameter name elem_check_stat doesn't match the CLI argument name skip-elementcheck. Additionally, their logic is inverted which could be confusing.

-def change_bias(..., elem_check_stat: bool = True):
+def change_bias(..., enable_element_completion: bool = True):

---

560-560: Fix inverted logic between function and CLI argument.

The function parameter defaults to True but is set from a CLI flag that defaults to False. This inversion makes the code harder to understand.

-            elem_check_stat=FLAGS.skip_elementcheck,
+            enable_element_completion=not FLAGS.skip_elementcheck,

deepmd/pt/utils/stat.py (4)

61-74: Remove unused variable.

The variable global_element_counts is initialized but never used.

 total_element_types = set()
-global_element_counts = {}
 collect_ele = defaultdict(int)

---

75-96: Add error handling for unknown data types.

The function silently passes on unknown data types, which could hide potential issues.

                 else:
-                    pass
+                    log.warning(f"Unexpected data type {type(stat_data[dd])} for key {dd}")

---

161-194: Add validation for element counts.

The element counting logic needs validation and could be simplified.

         if datasets[0].mixed_type and enable_element_completion:
+            if not isinstance(element_counts, dict):
+                log.warning(f"Invalid element counts for dataset {sys_index}")
+                continue
             for elem, data in element_counts.items():
+                if not isinstance(data, dict) or "indices" not in data or "frames" not in data:
+                    log.warning(f"Invalid data format for element {elem}")
+                    continue
                 indices = data["indices"]
                 count = data["frames"]

---

195-216: Improve error handling for missing elements.

The missing element handling could be improved with better error messages and validation.

     if datasets[0].mixed_type and enable_element_completion:
+        if not total_element_types:
+            log.warning("No elements found in any dataset")
+            return lst
         for elem, data in global_element_counts.items():
             indices_count = data["count"]
             if indices_count < min_frames_per_element_forstat:
                 log.warning(
-                    f"The number of frames in your datasets with element {elem} is {indices_count}, "
-                    f"which is less than the required {min_frames_per_element_forstat}"
+                    f"Insufficient frames for element {elem}: found {indices_count}, "
+                    f"required {min_frames_per_element_forstat}. This may affect model accuracy."
                 )

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro