(fix) Make bias statistics complete for all elements #4496
Conversation
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughThe pull request introduces modifications in the DeepMD-kit's PyTorch utility modules. A new public method Changes
Suggested labels
Suggested reviewers
Tip CodeRabbit's docstrings feature is now available as part of our Pro Plan! Simply use the command ✨ Finishing Touches
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
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):
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/utils/dataset.py(1 hunks)deepmd/pt/utils/stat.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/utils/dataset.py
44-44: SyntaxError: Expected an indented block after function definition
🔇 Additional comments (3)
deepmd/pt/utils/stat.py (3)
91-96: Validate 'atype' presence and handle empty arrays gracefully.
While updating the collected elements, consider explicitly checking that each 'atype' array is non-empty. If a dataset yields an empty array, it might bypass your missing-elements logic.
✅ Verification successful
Let me gather more information about how atype is handled in the codebase.
Let me try again with the correct file type flag for ripgrep.
Based on the code analysis, I can now generate the final response:
Empty arrays are handled safely by numpy.unique and set.update
The code uses numpy.unique() which safely handles empty arrays by returning an empty array, and set.update() which safely handles empty iterables. The logic for missing elements works correctly even with empty arrays since an empty array's unique values would simply not add any elements to the set.
🏁 Scripts executed
The following scripts were executed for the analysis:
Error: No shell scripts found
Script:
#!/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($$$)
$$$'
Length of output: 248
Script:
#!/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
Length of output: 89317
85-90: Consider using logging and verifying that all datasets are included.
You are only collecting the set of elements from the first dataset and simply printing them. If your intent is to gather possible elements across all datasets, you may want to iterate over all datasets. Also, using "log.info" would be more consistent than a raw "print" in a library setting.
Possible fix for the print statement:
- print('we want', all_elements)
+ log.info(f"We want elements: {all_elements}")97-109: Ensure partial coverage logic for missing elements is intended.
Here, frames are appended until the number of records reaches 'nbatches'. Verify that the short-circuiting with “break” actually meets your coverage needs for missing elements. If your goal is to gather as many frames as possible across all missing elements, you might need a more explicit limit or a different approach.
✅ Verification successful
Let me gather more context about this code to better understand the coverage logic for missing elements.
Let me search for more context about the initialization and usage of nbatches to better understand the coverage requirements.
Let me search for the initialization of nbatches to understand its purpose in the coverage logic.
Let me search for more context about how nbatches is used and its relationship with the dataset statistics.
Based on the gathered context, I can now provide a final response about the coverage logic for missing elements.
The current partial coverage logic is intentional and correct.
The code's behavior is appropriate because:
nbatchesis a parameter that defines the desired number of batches for collecting statistics, used consistently throughout themake_stat_inputfunction- The break condition
if len(lst[i]['atype']) >= nbatchesensures each system collects up to the specified number of batches, preventing over-sampling - This is consistent with how regular batches are collected earlier in the function where
numb_batches = min(nbatches, len(dataloaders[i]))is used - The goal is to gather a statistically significant sample (controlled by nbatches) rather than exhaustively processing all frames containing missing elements
The current implementation correctly balances between gathering sufficient data for statistics while maintaining consistent batch sizes across all elements.
🏁 Scripts executed
The following scripts were executed for the analysis:
Error: No shell scripts found
Script:
#!/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
Length of output: 1724
Script:
#!/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 {}
Length of output: 66294
Script:
#!/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
Length of output: 1079
Script:
#!/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
Length of output: 4539
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
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])
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/utils/dataset.py(1 hunks)deepmd/pt/utils/stat.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/utils/dataset.py
44-44: SyntaxError: Expected an indented block after function definition
🔇 Additional comments (2)
deepmd/pt/utils/dataset.py (2)
43-44:
Fix the indentation error in method definition.
The method definition has incorrect indentation which will cause a SyntaxError. It should be aligned with other class methods.
- def _build_element_to_frames(self):
- """Mapping element types to frame indexes"""
+ def _build_element_to_frames(self):
+ """Mapping element types to frame indexes"""Likely invalid or redundant comment.
🧰 Tools
🪛 Ruff (0.8.2)
44-44: SyntaxError: Expected an indented block after function definition
43-53: 🛠️ Refactor suggestion
Make frame limit configurable and enhance documentation.
- The hard-coded limit of 10 frames per element should be configurable.
- The docstring should be more descriptive about the method's purpose and return value.
- 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_framesLikely invalid or redundant comment.
🧰 Tools
🪛 Ruff (0.8.2)
44-44: SyntaxError: Expected an indented block after function definition
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (4)
source/tests/pt/test_make_stat_input.py (4)
14-23: Consider usingcollections.defaultdictforelement_to_frames.
You can simplify the nested checks for element presence in the dictionary by using adefaultdict(list), which would eliminate the need for the explicitif 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@propertybut naming itget_all_atypecan be confusing. Consider a more descriptive name likeall_atypes, since Python properties typically avoid "get_" prefixes.
53-59: Remove or use the assignedlstvariable.
The variablelstis 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
lstis assigned to but never usedRemove assignment to unused variable
lst(F841)
61-62: Optional test runner inclusion.
Having theif __name__ == "__main__": unittest.main()block is fine. You could remove it if tests are run by a dedicated test runner.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/pt/utils/dataset.py(4 hunks)deepmd/pt/utils/stat.py(1 hunks)source/tests/pt/test_make_stat_input.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/pt/test_make_stat_input.py
55-55: Local variable lst is assigned to but never used
Remove assignment to unused variable lst
(F841)
🔇 Additional comments (4)
source/tests/pt/test_make_stat_input.py (1)
40-52: Test setup looks good.
The dataset creation for testing is straightforward and clear. No issues found.
deepmd/pt/utils/dataset.py (2)
21-24: Docstring clarity is sufficient.
The docstring effectively describes constructor parameters. No corrections needed.
34-34: Initialization of element frames is a good approach.
Storing the result of _build_element_to_frames() in self.element_to_frames and self.get_all_atype reduces redundancy.
deepmd/pt/utils/stat.py (1)
86-94: No immediate issues with collection of atomic types.
Collecting and updating sets is correct.
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
deepmd/utils/data.py (1)
667-674: Consider providing more informative error details when the type map lookup fails.
While raisingIndexErroris 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 edeepmd/utils/argcheck.py (1)
2829-2834: Ensure user awareness of the new argument.The new argument
min_frames_per_element_forstatis 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
lstis assigned with the result ofmake_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
lstis assigned to but never usedRemove 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_elementsConfirm that the logic aligns with the earlier
missing_elementssets 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
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
deepmd/pt/train/training.py(2 hunks)deepmd/pt/utils/dataset.py(3 hunks)deepmd/pt/utils/stat.py(3 hunks)deepmd/utils/argcheck.py(1 hunks)deepmd/utils/data.py(1 hunks)source/tests/pt/test_make_stat_input.py(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/utils/dataset.py
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/utils/stat.py
134-134: SyntaxError: Expected a statement
134-134: SyntaxError: Expected a statement
134-134: SyntaxError: Expected a statement
134-134: SyntaxError: Expected a statement
135-135: SyntaxError: Unexpected indentation
144-144: SyntaxError: unindent does not match any outer indentation level
144-144: SyntaxError: Expected a statement
144-144: SyntaxError: Expected a statement
144-145: SyntaxError: Expected a statement
145-145: SyntaxError: Unexpected indentation
176-176: SyntaxError: Expected a statement
176-176: SyntaxError: Expected a statement
176-176: SyntaxError: Expected a statement
176-176: SyntaxError: Expected a statement
188-188: SyntaxError: unindent does not match any outer indentation level
189-189: SyntaxError: Unexpected indentation
231-231: SyntaxError: Expected a statement
231-231: SyntaxError: Expected a statement
231-231: SyntaxError: Expected a statement
231-231: SyntaxError: Expected a statement
source/tests/pt/test_make_stat_input.py
43-43: Loop control variable idx not used within loop body
Rename unused idx to _idx
(B007)
68-68: Local variable lst is assigned to but never used
Remove assignment to unused variable lst
(F841)
🔇 Additional comments (4)
deepmd/pt/train/training.py (2)
145-147: Add type check or validation for the statistic threshold.While setting
self.min_frames_per_element_forstat, consider ensuring it's a strictly positive integer. If a negative or zero value is passed, it may cause runtime issues or meaningless statistics.
232-232: Parameter usage looks correct.Passing
self.min_frames_per_element_forstattomake_stat_inputaligns with the newly introduced functionality. Just be sure to verify that all call sites expect this parameter and properly handle out-of-range values.source/tests/pt/test_make_stat_input.py (1)
41-50: Optimize loop variable usage.The static analysis hint suggests renaming
_idxif it were truly unused, but hereidxis used to assign frames. The accusations of an unused loop variable are false. You can safely ignore that static analysis hint.🧰 Tools
🪛 Ruff (0.8.2)
43-43: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
deepmd/pt/utils/stat.py (1)
52-110: Confirm correctness of cumulative element count logic.You aggregate
total_element_counts[elem]["count"], but also trackindices. Ensure you don’t exceed list boundaries when collecting indices for up tomin_frames_per_element_forstat. If more frames exist, consider whether you need them to fulfill certain statistics.
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (3)
source/tests/pt/test_make_stat_input.py (3)
62-72:⚠️ Potential issueTest with multiple batch sizes.
Using a fixed batch size of 1 may not catch batch-related issues.
- cls.dataloaders = [] - for dataset in cls.datasets: - dataloader = DataLoader( - dataset, - batch_size=1, - num_workers=0, - drop_last=False, - collate_fn=collate_fn, - pin_memory=False, - ) - cls.dataloaders.append(dataloader) + # Test with different batch sizes + cls.batch_sizes = [1, 4, 16] + cls.dataloaders = {} + for batch_size in cls.batch_sizes: + loaders = [] + for dataset in cls.datasets: + loader = DataLoader( + dataset, + batch_size=batch_size, + num_workers=0, + drop_last=False, + collate_fn=collate_fn, + pin_memory=False, + ) + loaders.append(loader) + cls.dataloaders[batch_size] = loaders
171-172:⚠️ Potential issueUse self.real_ntypes instead of hard-coded ntypes value.
The hard-coded
ntypes=57should use the derived value from the dataset.- bias_ori, _ = compute_output_stats(lst_ori, ntypes=57) - bias_all, _ = compute_output_stats(lst_all, ntypes=57) + bias_ori, _ = compute_output_stats(lst_ori, ntypes=self.real_ntypes) + bias_all, _ = compute_output_stats(lst_all, ntypes=self.real_ntypes)
192-193: 🛠️ Refactor suggestionDocument missing elements and improve test method.
The test method needs better documentation of missing elements and their impact.
- # missing element:13,31,37 - # only one frame would be count + """Test make_stat_input with element completion enabled/disabled. + + Missing elements in test data: + - Element 13: Missing due to [reason] + - Element 31: Missing due to [reason] + - Element 37: Missing due to [reason] + + When element completion is disabled: + - Only one frame is counted because [explain why] + - Missing elements should have zero bias + """
🧹 Nitpick comments (1)
source/tests/pt/test_make_stat_input.py (1)
25-44: Optimize dictionary key lookup in collate function.The dictionary key lookup in the loop can be optimized.
- for key in batch[0].keys(): + for key in batch[0]:🧰 Tools
🪛 Ruff (0.8.2)
30-30: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pt/test_make_stat_input.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/pt/test_make_stat_input.py
30-30: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (19)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (3)
source/tests/pt/test_make_stat_input.py (1)
54-81: 🛠️ Refactor suggestionTest setup needs significant improvements.
Current setup has several limitations:
- Uses single system when more than 10 were requested
- Hard-codes
real_ntypes = 6instead of deriving from dataset- Fixed batch size of 1 may not catch batch-related issues
- Missing CUDA availability check
@classmethod def setUpClass(cls): - with torch.device("cpu"): + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + with torch.device(device): - system_path = str(Path(__file__).parent / "mixed_type_data/sys.000000") - cls.real_ntypes = 6 - cls.datasets = DeepmdDataSetForLoader(system=system_path) + # Load multiple systems for comprehensive testing + base_path = Path(__file__).parent / "mixed_type_data" + system_paths = [str(base_path / f"sys.{i:06d}") for i in range(15)] + cls.datasets = [DeepmdDataSetForLoader(system=path) for path in system_paths] + + # Derive number of types from the datasets + cls.real_ntypes = max(max(dataset.get_types()) for dataset in cls.datasets) + 1 data_requirements = [ DataRequirementItem( "energy", ndof=1, atomic=False, ), ] - cls.datasets.add_data_requirement(data_requirements) - cls.datasets = [cls.datasets] + for dataset in cls.datasets: + dataset.add_data_requirement(data_requirements) - cls.dataloaders = [] - for dataset in cls.datasets: - dataloader = DataLoader( - dataset, - batch_size=1, - num_workers=0, - drop_last=False, - collate_fn=collate_fn, - pin_memory=False, - ) - cls.dataloaders.append(dataloader) + # Test with different batch sizes + cls.batch_sizes = [1, 4, 16] + cls.dataloaders = {} + for batch_size in cls.batch_sizes: + loaders = [] + for dataset in cls.datasets: + loader = DataLoader( + dataset, + batch_size=batch_size, + num_workers=0, + drop_last=False, + collate_fn=collate_fn, + pin_memory=False, + ) + loaders.append(loader) + cls.dataloaders[batch_size] = loadersdeepmd/pt/utils/stat.py (2)
38-44: 🛠️ Refactor suggestionAdd input validation for new parameters.
The function should validate the new parameters to ensure they meet requirements:
min_frames_per_element_forstatshould be positivedatasetsshould not be empty and should matchdataloaderslengthdef make_stat_input( datasets, dataloaders, nbatches, min_frames_per_element_forstat=10, enable_element_completion=True, ): + if not datasets: + raise ValueError("No datasets provided") + if len(datasets) != len(dataloaders): + raise ValueError("Number of datasets does not match number of dataloaders") + if min_frames_per_element_forstat < 1: + raise ValueError("min_frames_per_element_forstat must be positive")
195-225: 🛠️ Refactor suggestionReplace assertion with proper exception handling.
The assertion should be replaced with a proper exception that provides clear guidance.
def process_with_new_frame( sys_indices, newele_counter, min_frames, datasets, lst, collect_ele, miss ): """Process missing elements by adding new frames until the minimum is reached.""" + if not sys_indices: + log.warning(f"No system indices provided for element {miss}") + return for sys_info in sys_indices: sys_index = sys_info["sys_index"] frames = sys_info["frames"] + if not frames: + log.warning(f"No frames found for system {sys_index}") + continue sys = datasets if isinstance(datasets, list) else [datasets] sys = sys[sys_index] for frame in frames: newele_counter += 1 if newele_counter > min_frames: break - frame_data = sys.__getitem__(frame) + try: + frame_data = sys.__getitem__(frame) + if "atype" not in frame_data: + raise ValueError(f"Frame {frame} does not contain type information") + if miss not in frame_data["atype"]: + raise ValueError( + f"Element {miss} not found in frame {frame}. " + "To proceed without element completion:\n" + "1. For change-bias mode: Use '--skip-elementcheck'\n" + "2. For model training: Set 'enable_element_completion' to False" + ) + except Exception as e: + log.error(f"Failed to process frame {frame} from system {sys_index}: {e}") + raise
🧹 Nitpick comments (3)
source/tests/pt/test_make_stat_input.py (2)
33-51: Add error handling and type hints to collate_fn.The function could be improved with:
- Type hints for parameters and return value
- Better error handling for unsupported data types
- Documentation for supported data types
-def collate_fn(batch): +def collate_fn(batch: Union[dict, list[dict]]) -> dict: + """Collate batch data into a single dictionary. + + Args: + batch: A single dict or list of dicts containing batch data + + Returns: + dict: Collated batch with tensors stacked along first dimension + + Supported data types: + - torch.Tensor + - numpy.ndarray + - Numeric types convertible to tensor + """ if isinstance(batch, dict): batch = [batch] out = {} for key in batch[0].keys(): items = [sample[key] for sample in batch] if isinstance(items[0], torch.Tensor): out[key] = torch.stack(items, dim=0) elif isinstance(items[0], np.ndarray): out[key] = torch.from_numpy(np.stack(items, axis=0)) else: try: out[key] = torch.tensor(items) - except Exception: + except (TypeError, ValueError) as e: + log.warning(f"Could not convert {key} to tensor: {e}") out[key] = items return out🧰 Tools
🪛 Ruff (0.8.2)
38-38: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
82-83: Add docstring to helper method.The method needs documentation explaining its purpose and parameters.
def count_non_zero_elements(self, tensor, threshold=1e-8): + """Count elements in tensor with absolute value above threshold. + + Args: + tensor: Input tensor to analyze + threshold: Minimum absolute value to consider non-zero (default: 1e-8) + + Returns: + int: Number of elements with absolute value > threshold + """ return torch.sum(torch.abs(tensor) > threshold).item()deepmd/pt/utils/stat.py (1)
111-130: Improve error handling in process_batches.The function silently passes on unknown data types and may raise KeyError.
def process_batches(dataloader, sys_stat, nbatches): """Process batches from a dataloader and collect statistics.""" + def process_tensor_data(data): + """Process tensor data with proper error handling.""" + if data is None: + return None + try: + if isinstance(data, torch.Tensor): + return data + elif isinstance(data, np.float32): + return data + else: + log.warning(f"Unexpected data type {type(data)}") + return None + except Exception as e: + log.error(f"Error processing tensor data: {e}") + raise 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: + for dd, value in stat_data.items(): + result = process_tensor_data(value) + if result is None: sys_stat[dd] = None - elif isinstance(stat_data[dd], torch.Tensor): + elif isinstance(result, 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 + sys_stat[dd].append(result) + elif isinstance(result, np.float32): + sys_stat[dd] = result
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/utils/stat.py(1 hunks)source/tests/pt/test_make_stat_input.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/pt/test_make_stat_input.py
38-38: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
307-307: Redefinition of unused test_process_missing_elements from line 103
(F811)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
source/tests/pt/test_make_stat_input.py (1)
46-49: Improve error handling in collate_fn.The current broad exception handling could mask important errors. Consider catching specific exceptions and providing more descriptive error messages.
- try: - out[key] = torch.tensor(items) - except Exception: - out[key] = items + try: + out[key] = torch.tensor(items) + except (ValueError, TypeError) as e: + # If tensor conversion fails, keep original items + out[key] = items + print(f"Warning: Could not convert {key} to tensor: {str(e)}")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pt/test_make_stat_input.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/pt/test_make_stat_input.py
38-38: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (9)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
🔇 Additional comments (1)
source/tests/pt/test_make_stat_input.py (1)
82-83: LGTM! Well-implemented helper method.The
count_non_zero_elementsmethod correctly handles floating-point comparisons using a threshold parameter.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
source/tests/pt/test_make_stat_input.py (1)
85-101: 🛠️ Refactor suggestionTest methods need to handle multiple batch sizes.
The test methods currently use a fixed batch size, which may not catch batch-related issues.
def test_make_stat_input_with_element_counts(self): - lst = make_stat_input( - datasets=self.datasets, - dataloaders=self.dataloaders, - nbatches=1, - min_frames_per_element_forstat=1, - enable_element_completion=True, - ) - bias, _ = compute_output_stats(lst, ntypes=57) + for batch_size in self.batch_sizes: + lst = make_stat_input( + datasets=self.datasets, + dataloaders=self.dataloaders[batch_size], + nbatches=10, + min_frames_per_element_forstat=1, + enable_element_completion=True, + ) + bias, _ = compute_output_stats(lst, ntypes=self.real_ntypes)
🧹 Nitpick comments (2)
source/tests/pt/test_make_stat_input.py (2)
33-52: Optimize dictionary key lookup in collate_fn.The dictionary key lookup in the loop can be simplified.
- for key in batch[0].keys(): + for key in batch[0]:🧰 Tools
🪛 Ruff (0.8.2)
38-38: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
180-181: Document missing elements and their impact.The comment about missing elements needs better documentation to explain why these elements are missing and how it affects the test.
- # missing element:13,31,37 - # only one frame would be count + """Test make_stat_input with element completion disabled. + + Missing elements in test data: + - Element 13: Missing due to [reason] + - Element 31: Missing due to [reason] + - Element 37: Missing due to [reason] + + When element completion is disabled: + - Only one frame is counted because [explain why] + - Missing elements should have zero bias + """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
source/tests/pt/test_make_stat_input.py(1 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/pt/test_make_stat_input.py
38-38: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
source/tests/pt/test_make_stat_input.py (1)
268-303: LGTM!The test_process_element_count_and_missing_elements method is well-structured and includes proper assertions to verify the behavior.
There was a problem hiding this comment.
Actionable comments posted: 6
🧹 Nitpick comments (1)
deepmd/pt/utils/stat.py (1)
38-60: Improve docstring clarity.The docstring needs to be more descriptive about the function's purpose and return value.
- """Get list for statistics from datasets. + """Get a list of statistical data from datasets, with support for element completion. + + This function processes datasets to compute statistics, ensuring each element has sufficient + representation in the data. When element completion is enabled, it will attempt to find + additional frames containing underrepresented elements to meet the minimum frame requirement. Parameters ---------- nbatches : int The number of batches to process from each dataloader. min_frames_per_element_forstat : int, optional The minimum number of frames required for an element to be considered, default is 10. enable_element_completion : bool, optional If True, enables element completion for missing elements, default is True. Returns ------- list - A list containing the statistics for each dataset processed. + A list of dictionaries containing processed statistics for each dataset. + Each dictionary contains tensor data for various properties (e.g., coordinates, + energies, forces) aggregated from the dataset. """
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/utils/stat.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (5)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (c-cpp)
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
deepmd/pt/utils/stat.py (1)
64-64: 🛠️ Refactor suggestionRemove unused variable.
The variable
global_element_countsis initialized but never used within the function scope.total_element_types = set() - global_element_counts = {} global_type_name = {} collect_ele = defaultdict(int)
🧹 Nitpick comments (7)
deepmd/pt/utils/stat.py (7)
38-60: Input validation needed for new parameters.The function
make_stat_inputhas been updated with new parameters, but there's no validation to ensure they have valid values.def make_stat_input( datasets, dataloaders, nbatches, min_frames_per_element_forstat=10, enable_element_completion=True, ): + if not datasets: + raise ValueError("No datasets provided") + if len(datasets) != len(dataloaders): + raise ValueError("Number of datasets does not match number of dataloaders") + if min_frames_per_element_forstat < 1: + raise ValueError("min_frames_per_element_forstat must be positive") + if nbatches < 1: + raise ValueError("nbatches must be positive")
80-88: Add error handling to prevent runtime exceptions.The tensor operations could potentially fail if the data format is unexpected. Add appropriate try-except blocks to handle potential errors.
- process_batches(dataloader, sys_stat, nbatches) - if do_element_completion: - element_data = torch.cat(sys_stat["atype"], dim=0) - collect_values = torch.unique(element_data.flatten(), sorted=True) - for elem in collect_values.tolist(): - frames_with_elem = torch.any(element_data == elem, dim=1) - row_indices = torch.where(frames_with_elem)[0] - collect_ele[elem] += len(row_indices) + try: + process_batches(dataloader, sys_stat, nbatches) + if do_element_completion: + if "atype" not in sys_stat or not sys_stat["atype"]: + log.warning(f"No type information found in dataset {sys_index}") + continue + element_data = torch.cat(sys_stat["atype"], dim=0) + collect_values = torch.unique(element_data.flatten(), sorted=True) + for elem in collect_values.tolist(): + frames_with_elem = torch.any(element_data == elem, dim=1) + row_indices = torch.where(frames_with_elem)[0] + collect_ele[elem] += len(row_indices) + except Exception as e: + log.error(f"Error processing dataset {sys_index}: {e}") + continue
114-133: Add error handling for unexpected data types.The function silently skips processing for unknown data types. Add explicit logging to capture potential issues.
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: + log.warning(f"Unexpected data type {type(stat_data[dd])} for key {dd}")
135-147: Add error handling for tensor operations.The tensor concatenation operation could fail if tensors have incompatible shapes or if the list is empty. Add error handling for these cases.
def finalize_stats(sys_stat): """Finalize statistics based on the data type.""" 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) + try: + if isinstance(sys_stat[key], np.float32): + continue + 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) + else: + log.warning(f"Unexpected data type for key {key}") + except Exception as e: + log.error(f"Error finalizing stats for key {key}: {e}") + raise
150-177: Add validation for element counts data.The function should validate the element counts and data format to prevent potential issues with malformed data.
def process_element_counts( sys_index, dataset, min_frames, global_element_counts, global_type_name, total_element_types, ): """Count element occurrences in the dataset and update global statistics.""" element_counts, type_name = dataset.get_frame_index_for_elements() + if not isinstance(element_counts, dict): + log.warning(f"Invalid element counts for dataset {sys_index}") + return for new_idx, elem_name in type_name.items(): if new_idx not in global_type_name: global_type_name[new_idx] = elem_name 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
179-196: Add validation for missing elements processing.The function needs validation for empty inputs and better logging for insufficient frames.
def process_missing_elements( min_frames, global_element_counts, total_element_types, collect_ele, datasets, lst ): """Handle missing elements by adding them to the statistics.""" + if not total_element_types: + log.warning("No elements found in any dataset") + return + if not global_element_counts: + log.warning("No element counts available") + return collect_elements = collect_ele.keys() missing_elements = total_element_types - collect_elements collect_miss_element = set() for ele, count in collect_ele.items(): if count < min_frames: + log.info(f"Element {ele} has insufficient frames: {count}/{min_frames}") collect_miss_element.add(ele) missing_elements.add(ele) for miss in missing_elements: + if miss not in global_element_counts: + log.warning(f"No count data available for element {miss}") + continue
212-224: Extract common tensor processing logic.The tensor processing logic is duplicated in multiple places. Extract it into a helper function to reduce duplication and improve maintainability.
+ def process_tensor_data(data): + """Process tensor data with proper error handling.""" + if data is None: + return None + try: + if isinstance(data, np.ndarray): + tensor_data = torch.from_numpy(data).unsqueeze(0) + return tensor_data + elif isinstance(data, np.float32): + return data + return None + except Exception as e: + log.error(f"Error processing tensor data: {e}") + raise + 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] = [] - tensor_data = torch.from_numpy(frame_data[dd]).unsqueeze(0) - sys_stat_new[dd].append(tensor_data) - elif isinstance(frame_data[dd], np.float32): - sys_stat_new[dd] = frame_data[dd] + result = process_tensor_data(frame_data[dd]) + if result is None: + sys_stat_new[dd] = None + elif isinstance(result, torch.Tensor): + if dd not in sys_stat_new: + sys_stat_new[dd] = [] + sys_stat_new[dd].append(result) + elif isinstance(result, np.float32): + sys_stat_new[dd] = result
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/utils/stat.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (19)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
deepmd/pt/utils/stat.py (1)
68-76: Good user feedback with clear status logging.The added logging statements provide clear feedback on whether element completion is enabled, and the
do_element_completionvariable improves code readability by combining the conditions.
|
What is the bias statistics result before and after this PR? |
Summary by CodeRabbit
New Features
Bug Fixes
Tests
make_stat_inputfunction to ensure accurate processing of atomic types.