feat(pt): Support denoise model #4644
Conversation
📝 WalkthroughWalkthroughThis pull request refactors the denoising components within the deepmd package. The Changes
Suggested reviewers
✨ Finishing Touches
🪧 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 (
|
for more information, see https://pre-commit.ci
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (7)
deepmd/pt/model/atomic_model/denoise_atomic_model.py (1)
35-61: Consider documenting or removing commented code.There's a large block of commented code that appears to implement a more sophisticated approach for handling the virial tensor. Either:
- Document why this code is commented out and when it might be used
- Remove it if it's not needed or implement it properly if it's ready
Keeping large blocks of commented code without explanation makes maintenance more difficult.
deepmd/pt/model/model/denoise_model.py (2)
27-28: Consider aligningmodel_typewith the denoise context.
Currently set to"property", which may cause confusion given this is a denoise model.- model_type = "property" + model_type = "denoise"
37-48: Remove or justify the large commented-out code blocks.
Long sections of commented-out logic can hamper readability. Retain them only if they serve as placeholders with clear documentation.Also applies to: 88-106
deepmd/pt/model/task/denoise.py (4)
65-107: Docstring clarity and consistency.
The updated parameter list in the docstring accurately reflects the constructor’s arguments and their usage. However, the mention of "Coordinate and cell fitting" might benefit from clarifying how these newly introduced parameters (e.g.,dim_case_embd,numb_aparam) fit into the broader denoising purpose.
328-328: Usefor kk in variables:instead offor kk in variables.keys():.
Invoking.keys()is unnecessary when iterating over a dictionary.- for kk in variables.keys(): + for kk in variables:🧰 Tools
🪛 Ruff (0.8.2)
328-328: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
445-456: Fix missing f-string usage in raise statements.
Placeholders{...}in string literals won’t interpolate unless you use an f-string.- raise ValueError( - "get an input fparam of dim {fparam.shape[-1]}, ", - "which is not consistent with {self.numb_fparam}.", - ) + raise ValueError( + f"get an input fparam of dim {fparam.shape[-1]}, " + f"which is not consistent with {self.numb_fparam}." + )Apply similarly at lines 445, 455, 480, etc.
Also applies to: 480-481
417-568: Docstring discrepancy and unused parameters inforward.
The docstring says “calculate total energy,” but the function returnsstrain_componentsandupdated_coord. Also,g2andh2are never used.Consider updating the docstring to match the actual outputs or clarifying if
g2andh2will be used later.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
deepmd/pt/loss/denoise.py(2 hunks)deepmd/pt/model/atomic_model/__init__.py(2 hunks)deepmd/pt/model/atomic_model/denoise_atomic_model.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/model/denoise_model.py(1 hunks)deepmd/pt/model/task/denoise.py(1 hunks)deepmd/utils/argcheck.py(2 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/loss/denoise.py
147-147: Undefined name get_cell_perturb_matrix_HEA
(F821)
deepmd/pt/model/task/denoise.py
59-59: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
230-230: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
245-245: Undefined name get_index_between_two_maps
(F821)
248-248: Undefined name map_atom_exclude_types
(F821)
301-301: Undefined name to_numpy_array
(F821)
302-302: Undefined name to_numpy_array
(F821)
303-303: Undefined name to_numpy_array
(F821)
304-304: Undefined name to_numpy_array
(F821)
305-305: Undefined name to_numpy_array
(F821)
306-306: Undefined name to_numpy_array
(F821)
323-323: Undefined name data
(F821)
327-327: Undefined name cls
(F821)
328-328: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
329-329: Undefined name to_torch_tensor
(F821)
⏰ Context from checks skipped due to timeout of 90000ms (19)
- 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: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (c-cpp)
- 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)
🔇 Additional comments (9)
deepmd/pt/model/atomic_model/__init__.py (2)
20-22: Import and class addition looks good.The addition of
DPDenoiseAtomicModelfrom the new module follows the existing import pattern used throughout the file.
53-53: LGTM!The
DPDenoiseAtomicModelclass is correctly added to the__all__list, making it part of the public API of the module.deepmd/utils/argcheck.py (2)
1865-1903: Implementation of fitting_denoise looks good.The function properly defines parameters with appropriate types, defaults, and documentation for the denoise fitting functionality, following the existing patterns in the codebase.
2873-2927: Parameters of loss_denoise are well defined.The remaining parameters for the denoise loss function are properly defined with appropriate types, defaults, and documentation, which follows the existing patterns in the codebase.
deepmd/pt/model/atomic_model/denoise_atomic_model.py (1)
18-25: Type checking in constructor looks good.The constructor properly checks that the fitting parameter is an instance of
DenoiseNet, which ensures proper compatibility between model components.deepmd/pt/model/model/__init__.py (3)
36-38: No issues with the added import.
100-101: Confirmed: Settingout_dimto embedding dimension
Aligningout_dimwithdescriptor.get_dim_emb()for "denoise" fitting ensures consistent dimensionality between descriptor embeddings and the fitting module.
274-275: Logic for selectingDenoiseModellooks good.deepmd/pt/model/task/denoise.py (1)
111-138:✅ Verification successful
Check partial usage of numpy and torch.
You are convertingbias_atom_efrom a NumPy array to a PyTorch tensor. This is correct, but ensure that all relevant data paths use consistent data types throughout the code base.Please confirm that no step inadvertently converts between NumPy and torch arrays without relevant device or dtype alignment in other modules or calls.
🏁 Script executed:
#!/bin/bash # Searching for unguarded conversions from numpy to torch or vice versa rg -A 2 'torch.tensor\(.*np\.'Length of output: 7489
Conversion Consistency Validated
- In
deepmd/pt/model/task/denoise.py,bias_atom_eis correctly converted to a torch tensor using explicitdeviceanddtypeparameters.- Across the repository, conversions from NumPy arrays to torch tensors are either performed via the
to_torch_tensorhelper (as seen in tests and defined indeepmd/pt/utils/utils.py) or use explicit conversion parameters. This ensures proper device and dtype alignment.- No unguarded or inconsistent conversions were detected.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (6)
deepmd/pt/loss/denoise.py (1)
180-186:⚠️ Potential issueFix variable naming mismatch and unnecessary initialization.
The code initializes
noise_on_coordto0.0and then immediately overwrites it. Also, there's a naming mismatch whereself.noiseis used but should beself.coord_noise.Apply this fix:
- noise_on_coord = 0.0 if self.noise_type == "uniform": noise_on_coord = np.random.uniform( - low=-self.noise, high=self.coord_noise, size=(mask_num, 3) + low=-self.coord_noise, high=self.coord_noise, size=(mask_num, 3) )deepmd/pt/model/task/denoise.py (5)
59-60: 🛠️ Refactor suggestionAvoid using mutable default argument.
Using an empty list as a default argument can lead to unexpected behavior because the same list object will be shared across all instances where the default is used.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None, trainable: Union[bool, list[bool]] = True,Then initialize the list inside the method:
def __init__(self, ...): # ... + if exclude_types is None: + exclude_types = [] + self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
59-59: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
253-253: 🛠️ Refactor suggestionSame mutable default issue in
reinit_exclude.Using a mutable default argument here has the same potential issues as in the constructor.
- def reinit_exclude( - self, - exclude_types: list[int] = [], - ) -> None: + def reinit_exclude( + self, + exclude_types: Optional[list[int]] = None, + ) -> None: + if exclude_types is None: + exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
253-253: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
268-271:⚠️ Potential issueUndefined functions causing potential errors.
The code calls
get_index_between_two_mapsandmap_atom_exclude_typeswhich don't appear to be defined or imported in this file.Make sure to import or define these functions:
+ from deepmd.pt.utils.type_mapping import get_index_between_two_maps, map_atom_exclude_typesOr replace them with proper implementations.
🧰 Tools
🪛 Ruff (0.8.2)
268-268: Undefined name
get_index_between_two_maps(F821)
271-271: Undefined name
map_atom_exclude_types(F821)
332-338:⚠️ Potential issueMissing utility function for numpy array conversion.
The serialization code calls
to_numpy_arraywhich isn't defined or imported.Ensure this function is imported or define it:
+ from deepmd.pt.utils.tensor_ops import to_numpy_arrayOr replace with native functionality:
- "bias_atom_e": to_numpy_array(self.bias_atom_e), + "bias_atom_e": self.bias_atom_e.detach().cpu().numpy() if self.bias_atom_e is not None else None,🧰 Tools
🪛 Ruff (0.8.2)
332-332: Undefined name
to_numpy_array(F821)
333-333: Undefined name
to_numpy_array(F821)
334-334: Undefined name
to_numpy_array(F821)
335-335: Undefined name
to_numpy_array(F821)
336-336: Undefined name
to_numpy_array(F821)
337-337: Undefined name
to_numpy_array(F821)
354-365:⚠️ Potential issueCritical issues in the deserialize method.
The deserialize method has several significant issues:
datais referenced but not definedclsis referenced but not definedto_torch_tensoris used but not definedFix the method signature and implementation:
- def deserialize(self) -> "DenoiseNet": - data = data.copy() - variables = data.pop("@variables") - cell_nets = data.pop("cell_nets") - coord_nets = data.pop("coord_nets") - token_nets = data.pop("token_nets") - obj = cls(**data) - for kk in variables.keys(): - obj[kk] = to_torch_tensor(variables[kk]) - obj.filter_layers_cell = NetworkCollection.deserialize(cell_nets) - obj.filter_layers_coord = NetworkCollection.deserialize(coord_nets) - obj.filter_layers_token = NetworkCollection.deserialize(token_nets) - return obj + @classmethod + def deserialize(cls, data: dict) -> "DenoiseNet": + data_copy = data.copy() + variables = data_copy.pop("@variables") + cell_nets = data_copy.pop("cell_nets") + coord_nets = data_copy.pop("coord_nets") + token_nets = data_copy.pop("token_nets") + obj = cls(**data_copy) + for kk in variables: + if variables[kk] is not None: + obj[kk] = torch.tensor(variables[kk], dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) + obj.filter_layers_cell = NetworkCollection.deserialize(cell_nets) + obj.filter_layers_coord = NetworkCollection.deserialize(coord_nets) + obj.filter_layers_token = NetworkCollection.deserialize(token_nets) + return obj🧰 Tools
🪛 Ruff (0.8.2)
354-354: Undefined name
data(F821)
359-359: Undefined name
cls(F821)
360-360: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
361-361: Undefined name
to_torch_tensor(F821)
🧹 Nitpick comments (2)
deepmd/pt/model/task/denoise.py (2)
510-514: Same unreachable code issue.Similar to the case above, this code block is unreachable because
xx_zerosis alwaysNone.- if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, fparam], - dim=-1, - )
534-538: Another unreachable code block.Same issue with unreachable code for
xx_zeros.- if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, aparam], - dim=-1, - )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (14)
deepmd/pt/loss/denoise.py(2 hunks)deepmd/pt/model/atomic_model/denoise_atomic_model.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/model/denoise_model.py(1 hunks)deepmd/pt/model/task/denoise.py(2 hunks)deepmd/pt/train/training.py(1 hunks)deepmd/utils/argcheck.py(2 hunks)examples/denoise/data/data_0/type.raw(1 hunks)examples/denoise/data/data_0/type_map.raw(1 hunks)examples/denoise/data/data_1/type.raw(1 hunks)examples/denoise/data/data_1/type_map.raw(1 hunks)examples/denoise/data/data_2/type.raw(1 hunks)examples/denoise/data/data_2/type_map.raw(1 hunks)examples/denoise/train/input.json(1 hunks)
✅ Files skipped from review due to trivial changes (6)
- examples/denoise/data/data_1/type.raw
- examples/denoise/data/data_0/type.raw
- examples/denoise/data/data_2/type.raw
- examples/denoise/data/data_0/type_map.raw
- examples/denoise/data/data_1/type_map.raw
- examples/denoise/data/data_2/type_map.raw
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/model/task/denoise.py
59-59: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
253-253: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
268-268: Undefined name get_index_between_two_maps
(F821)
271-271: Undefined name map_atom_exclude_types
(F821)
332-332: Undefined name to_numpy_array
(F821)
333-333: Undefined name to_numpy_array
(F821)
334-334: Undefined name to_numpy_array
(F821)
335-335: Undefined name to_numpy_array
(F821)
336-336: Undefined name to_numpy_array
(F821)
337-337: Undefined name to_numpy_array
(F821)
354-354: Undefined name data
(F821)
359-359: Undefined name cls
(F821)
360-360: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
361-361: Undefined name to_torch_tensor
(F821)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (clang, clang)
🔇 Additional comments (17)
deepmd/pt/train/training.py (1)
1293-1296: Properly passing denoising parameters from loss to modelThe added code correctly forwards the denoising-specific parameters from the loss configuration to the model parameters when a denoise loss type is detected. This ensures that both the model and loss function use consistent parameters for noise handling.
deepmd/pt/model/model/__init__.py (3)
36-38: Import of new DenoiseModel looks goodThe DenoiseModel is properly imported from the denoise_model module.
100-107: Well-implemented configuration for denoise fitting typeThe implementation correctly:
- Asserts that the last element in
type_mapmust be "MASKED_TOKEN" for denoise fitting- Sets the output dimension to match the embedding dimension
- Passes through the necessary noise-related parameters
This ensures proper configuration of the denoising model.
279-281: Appropriate model class selection for denoise fitting typeThe code correctly assigns the DenoiseModel class when the fitting type is "denoise", properly integrating it into the model selection logic.
deepmd/utils/argcheck.py (2)
1865-1903: Good addition of a new fitting type for the denoising model.The
fitting_denoisefunction defines necessary parameters for the denoising fitting network with appropriate documentation and default values. This new function is appropriately registered with thefitting_args_plugin.
2811-2928: Well-structured parameter definitions for the denoising loss function.The
loss_denoisefunction comprehensively defines all necessary parameters for configuring denoising loss, including masking options, noise settings, and loss preferences. The documentation is clear and the default values are sensible.deepmd/pt/model/atomic_model/denoise_atomic_model.py (2)
18-24: Good type checking in constructor.The constructor properly verifies that the
fittingparameter is an instance ofDenoiseNet, which ensures type safety and helps prevent runtime errors.
26-41: Clear implementation of noise-specific statistics application.The method appropriately handles different noise types and applies scaling based on statistical properties. The standard deviation calculation is correctly adjusted for different noise distributions (Gaussian vs. uniform).
examples/denoise/train/input.json (1)
1-99: Well-structured configuration file for denoising model training.This configuration example provides a comprehensive template for training a denoising model, with sensible defaults for all parameters across the model, learning rate, loss function, and training sections. The inclusion of comments adds clarity and the parameters align with the newly introduced functionality in the code.
deepmd/pt/loss/denoise.py (5)
23-38: Good implementation of the cell perturbation matrix function.This new helper function creates a random perturbation matrix for the cell based on the fraction parameter. The implementation looks solid with proper error handling for negative inputs.
44-109: Well-structured parameter set with comprehensive docstrings.The refactored constructor now supports a more flexible configuration with clear documentation for each parameter.
147-160: Possible missing import for cell perturbation function.The code uses
get_cell_perturb_matrixwhich is now defined in this file, replacing any previous usage ofget_cell_perturb_matrix_HEAthat was flagged in previous reviews.
231-234: Good error handling for masking requirements.The code properly validates that at least one masking type is enabled, ensuring the denoising process has something to work with.
241-256: Comprehensive loss calculation for coordinates and cell components.The implementation correctly handles both RMSE and MAE loss calculations with appropriate tracking in the
more_lossdictionary for debugging and analysis.deepmd/pt/model/task/denoise.py (3)
547-551: Unreachable code block.This code block will never be executed because
xx_zerosis alwaysNoneat this point in the code.You can safely remove this unreachable code:
- if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, case_embd], - dim=-1, - )
190-208: Well-structured network architecture for denoising.The implementation nicely separates the network components for coordinate, cell, and token fitting with consistent configuration. This design supports flexible denoising capabilities.
Also applies to: 209-226, 228-245
457-625: Comprehensive forward method implementation.The forward method is well-implemented with proper type checking, tensor handling, and efficient processing. It correctly uses the three network collections to generate predictions for strain components, updated coordinates, and logits.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (5)
deepmd/pt/model/model/denoise_model.py (2)
76-89:⚠️ Potential issueImplement the
forward_lowermethod or provide more context.The method is marked for export with
@torch.jit.exportbut only contains a TODO comment andpassstatement. This could cause issues when trying to use the model with TorchScript.@torch.jit.export def forward_lower( self, extended_coord, extended_atype, nlist, mapping: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, comm_dict: Optional[dict[str, torch.Tensor]] = None, ): - # TODO: implement forward_lower - pass + # TODO: Implement forward_lower for lower-level model operations + raise NotImplementedError( + "forward_lower is not yet implemented for DenoiseModel. " + "This method is required for TorchScript export and model deployment." + )
37-48:⚠️ Potential issueImplement the commented-out code in
translated_output_def.The method currently contains only a
passstatement followed by commented-out implementation. Either implement the method properly or document why it's not yet implemented.def translated_output_def(self): - pass - """ - out_def_data = self.model_output_def().get_data() - output_def = { - f"atom_{self.get_var_name()}": out_def_data[self.get_var_name()], - self.get_var_name(): out_def_data[f"{self.get_var_name()}_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - """ + out_def_data = self.model_output_def().get_data() + output_def = { + f"atom_{self.get_var_name()}": out_def_data[self.get_var_name()], + self.get_var_name(): out_def_data[f"{self.get_var_name()}_redu"], + } + if "mask" in out_def_data: + output_def["mask"] = out_def_data["mask"] + return output_defdeepmd/pt/loss/denoise.py (2)
190-192:⚠️ Potential issueVariable naming mismatch between self.noise and self.coord_noise.
The code uses
self.noisethoughcoord_noiseis the parameter name defined in the constructor.- noise_on_coord = np.random.uniform( - low=-self.noise, high=self.coord_noise, size=(mask_num, 3) - ) + noise_on_coord = np.random.uniform( + low=-self.coord_noise, high=self.coord_noise, size=(mask_num, 3) + )
305-310:⚠️ Potential issueImplement the serialize/deserialize methods.
These methods are currently just placeholders with
passstatements, but are needed for model serialization.def serialize(self) -> dict: - pass + return { + "ntypes": self.ntypes, + "mask_token": self.mask_token, + "mask_coord": self.mask_coord, + "mask_cell": self.mask_cell, + "token_loss": self.token_loss, + "coord_loss": self.coord_loss, + "cell_loss": self.cell_loss, + "noise_type": self.noise_type, + "coord_noise": self.coord_noise, + "cell_pert_fraction": self.cell_pert_fraction, + "noise_mode": self.noise_mode, + "mask_num": self.mask_num, + "mask_prob": self.mask_prob, + "same_mask": self.same_mask, + "loss_func": self.loss_func, + } @classmethod def deserialize(cls, data: dict) -> "TaskLoss": - pass + return cls(**data)deepmd/pt/model/task/denoise.py (1)
353-365:⚠️ Potential issueFix the deserialize method implementation.
The deserialize method has several issues:
- It doesn't take parameters but uses undefined variables
dataandcls- It uses undefined
to_torch_tensorfunction-def deserialize(self) -> "DenoiseNet": - data = data.copy() - variables = data.pop("@variables") - cell_nets = data.pop("cell_nets") - coord_nets = data.pop("coord_nets") - token_nets = data.pop("token_nets") - obj = cls(**data) - for kk in variables.keys(): - obj[kk] = to_torch_tensor(variables[kk]) - obj.filter_layers_cell = NetworkCollection.deserialize(cell_nets) - obj.filter_layers_coord = NetworkCollection.deserialize(coord_nets) - obj.filter_layers_token = NetworkCollection.deserialize(token_nets) - return obj +@classmethod +def deserialize(cls, data: dict) -> "DenoiseNet": + data = data.copy() + variables = data.pop("@variables") + cell_nets = data.pop("cell_nets") + coord_nets = data.pop("coord_nets") + token_nets = data.pop("token_nets") + obj = cls(**data) + for kk in variables: + obj[kk] = torch.tensor(variables[kk], device=device, dtype=dtype) + obj.filter_layers_cell = NetworkCollection.deserialize(cell_nets) + obj.filter_layers_coord = NetworkCollection.deserialize(coord_nets) + obj.filter_layers_token = NetworkCollection.deserialize(token_nets) + return obj🧰 Tools
🪛 Ruff (0.8.2)
354-354: Undefined name
data(F821)
359-359: Undefined name
cls(F821)
360-360: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
361-361: Undefined name
to_torch_tensor(F821)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deepmd/pt/loss/denoise.py(2 hunks)deepmd/pt/model/atomic_model/denoise_atomic_model.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/model/denoise_model.py(1 hunks)deepmd/pt/model/task/denoise.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/atomic_model/denoise_atomic_model.py
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/model/task/denoise.py
59-59: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
253-253: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
268-268: Undefined name get_index_between_two_maps
(F821)
271-271: Undefined name map_atom_exclude_types
(F821)
332-332: Undefined name to_numpy_array
(F821)
333-333: Undefined name to_numpy_array
(F821)
334-334: Undefined name to_numpy_array
(F821)
335-335: Undefined name to_numpy_array
(F821)
336-336: Undefined name to_numpy_array
(F821)
337-337: Undefined name to_numpy_array
(F821)
354-354: Undefined name data
(F821)
359-359: Undefined name cls
(F821)
360-360: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
361-361: Undefined name to_torch_tensor
(F821)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
🔇 Additional comments (6)
deepmd/pt/model/model/__init__.py (3)
36-38: LGTM: Clean import for DenoiseModel.The new import for DenoiseModel maintains the module's organization pattern, ensuring consistent structure.
100-107: DenoiseModel configuration looks good.The new condition for "denoise" fitting type correctly configures necessary parameters including the assertion to ensure MASKED_TOKEN is present in the type_map.
280-281: Fitting type integration is correct.The denoise model is properly integrated into the model factory method.
deepmd/pt/model/model/denoise_model.py (1)
50-74: Forward method implementation looks good.The forward method is well-structured, calling forward_common and constructing a clear dictionary of model predictions.
deepmd/pt/loss/denoise.py (2)
23-38: Well-implemented utility function for cell perturbation.The
get_cell_perturb_matrixfunction is cleanly implemented with proper error checking and return values.
44-109: Thoroughly documented parameter initialization.The constructor has excellent docstrings explaining each parameter's purpose and default values.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4644 +/- ##
==========================================
+ Coverage 84.80% 84.88% +0.08%
==========================================
Files 692 697 +5
Lines 66396 67009 +613
Branches 3539 3539
==========================================
+ Hits 56306 56883 +577
- Misses 8949 8985 +36
Partials 1141 1141 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (5)
deepmd/pt/model/model/denoise_model.py (1)
37-48:⚠️ Potential issueImplement or remove the commented-out code in
translated_output_def.The method currently contains only a
passstatement followed by commented-out implementation. Either implement the method properly or document why it's not yet implemented with a TODO comment.def translated_output_def(self): - pass - """ - out_def_data = self.model_output_def().get_data() - output_def = { - f"atom_{self.get_var_name()}": out_def_data[self.get_var_name()], - self.get_var_name(): out_def_data[f"{self.get_var_name()}_redu"], - } - if "mask" in out_def_data: - output_def["mask"] = out_def_data["mask"] - return output_def - """ + # TODO: Implement output definition translation for denoise model + # Will need to define how to map internal model outputs to external API + raise NotImplementedError("Method not yet implemented")deepmd/pt/loss/denoise.py (2)
190-192:⚠️ Potential issueFix variable name mismatch in noise generation.
There's a name mismatch between the class variable
self.coord_noisedefined in the constructor andself.noiseused here, which would cause a runtime error.noise_on_coord = rng.uniform( - low=-self.noise, high=self.coord_noise, size=(mask_num, 3) + low=-self.coord_noise, high=self.coord_noise, size=(mask_num, 3) )
300-310: 🛠️ Refactor suggestionImplement the serialize and deserialize methods.
The
serializeanddeserializemethods are marked withpass, which means they don't actually do anything. These should be properly implemented to ensure the class can be saved and loaded.def serialize(self) -> dict: - pass + """Serialize the loss to dict.""" + return { + "ntypes": self.ntypes, + "mask_token": self.mask_token, + "mask_coord": self.mask_coord, + "mask_cell": self.mask_cell, + "token_loss": self.token_loss, + "coord_loss": self.coord_loss, + "cell_loss": self.cell_loss, + "noise_type": self.noise_type, + "coord_noise": self.coord_noise, + "cell_pert_fraction": self.cell_pert_fraction, + "noise_mode": self.noise_mode, + "mask_num": self.mask_num, + "mask_prob": self.mask_prob, + "same_mask": self.same_mask, + "loss_func": self.loss_func, + } @classmethod def deserialize(cls, data: dict) -> "TaskLoss": - pass + """Deserialize to create an instance.""" + return cls(**data)deepmd/pt/model/task/denoise.py (2)
67-68:⚠️ Potential issueFix mutable default argument.
Using mutable objects (like empty lists) as default arguments can lead to unexpected behavior because the object is created once when the function is defined.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None,Then initialize it inside the method:
if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
67-67: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
260-261:⚠️ Potential issueFix the mutable default argument in reinit_exclude.
Same issue with mutable default argument.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None,Then initialize it inside the method:
if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
261-261: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
🧹 Nitpick comments (2)
deepmd/pt/loss/denoise.py (1)
215-216: Consider clearer variable naming in label assignment.The code applies coordinate updates but the comment and variable name don't clearly reflect what's happening. Consider renaming for clarity or adding a more descriptive comment.
- # label["updated_coord"] = (label["clean_frac_coord"] - frac_coord).clone().detach() + # Calculate coordinate differences in physical space by transforming fractional coordinates label["updated_coord"] = ( ( (label["clean_frac_coord"] - frac_coord) @ label["clean_box"].reshape(nbz, 3, 3) ) .clone() .detach() )deepmd/pt/model/task/denoise.py (1)
368-369: Improve dictionary iteration.Using
.keys()in a for loop is unnecessary in Python.-for kk in variables.keys(): +for kk in variables: obj[kk] = to_torch_tensor(variables[kk])🧰 Tools
🪛 Ruff (0.8.2)
368-368: 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 (3)
deepmd/pt/loss/denoise.py(2 hunks)deepmd/pt/model/model/denoise_model.py(1 hunks)deepmd/pt/model/task/denoise.py(2 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/model/task/denoise.py
67-67: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
261-261: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
368-368: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (26)
- 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: Build C++ (clang, clang)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (1, 3.9)
- 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: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test C++ (false)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test C++ (true)
- 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: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (5)
deepmd/pt/model/model/denoise_model.py (2)
76-105: Implement theforward_lowermethod properly.The method is marked for export with
@torch.jit.exportbut doesn't provide an informative error when not implemented, which could cause issues when using the model with TorchScript.@torch.jit.export def forward_lower( self, extended_coord, extended_atype, nlist, mapping: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, comm_dict: Optional[dict[str, torch.Tensor]] = None, ): model_ret = self.forward_common_lower( extended_coord, extended_atype, nlist, mapping, fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) model_predict = {} model_predict["updated_coord"] = model_ret["updated_coord"] model_predict["atom_strain_components"] = model_ret["strain_components"] model_predict["strain_components"] = model_ret["strain_components_redu"] model_predict["logits"] = model_ret["logits"] if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict
50-74: Model output structure is appropriate for denoising tasks.The forward method properly transforms the model's return values into a dictionary containing updated coordinates, strain components, and logits, which aligns with the PR objective of supporting denoise models for relaxing near-equilibrium structures.
deepmd/pt/loss/denoise.py (1)
253-286: The loss calculation logic is well-structured.The implementation handles both RMSE and MAE loss functions appropriately, with proper reporting of component losses in the
more_lossdictionary. This flexibility supports different training approaches as needed for the denoising tasks.deepmd/pt/model/task/denoise.py (2)
537-547: Good implementation of coordinate prediction.The code properly applies the network output to transform coordinates using the gradient rotations, ensuring that the model's predictions are applied in the correct geometric space.
71-74: Essential parameters for denoising tasks.The addition of
coord_noise,cell_pert_fraction, andnoise_typeparameters supports the PR's objective of incorporating noise for training, which will help with relaxing near-equilibrium structures for improving accuracy of the IS2RP task.
… into merge_denoise
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
deepmd/pt/model/model/__init__.py (1)
100-107: Consider using a more explicit exception for user input validation.
Usingasserthere causes anAssertionErrorwhich might be stripped out when Python is optimized. For production code, prefer raising a more explicit exception (e.g.,ValueError) to clearly indicate a user input error.-if fitting_net["type"] in ["denoise"]: - assert model_params["type_map"][-1] == "MASKED_TOKEN", ( - f"When using denoise fitting, the last element in `type_map` must be 'MASKED_TOKEN', but got '{model_params['type_map'][-1]}'" - ) +if fitting_net["type"] == "denoise": + if model_params["type_map"][-1] != "MASKED_TOKEN": + raise ValueError( + f"When using denoise fitting, the last element in `type_map` must be 'MASKED_TOKEN', " + f"but got '{model_params['type_map'][-1]}'" + )deepmd/pt/model/atomic_model/denoise_atomic_model.py (1)
26-47: Use mathematical constants instead of hardcoded approximations.
Lines involving1.732can be replaced bymath.sqrt(3)to avoid potential confusion. It helps maintain clarity and precision.+import math def apply_out_stat(self, ret: dict[str, torch.Tensor], atype: torch.Tensor): noise_type = self.fitting_net.get_noise_type() - cell_std = self.fitting_net.get_cell_pert_fraction() / 1.732 + cell_std = self.fitting_net.get_cell_pert_fraction() / math.sqrt(3) ... - coord_std = self.fitting_net.get_coord_noise() / 1.732 + coord_std = self.fitting_net.get_coord_noise() / math.sqrt(3) ...deepmd/pt/model/task/denoise.py (2)
361-374: Simplify dictionary iteration
Instead of iterating overvariables.keys(), you can iterate directly overvariablesfor clarity and efficiency.-for kk in variables.keys(): +for kk in variables: obj[kk] = to_torch_tensor(variables[kk])🧰 Tools
🪛 Ruff (0.8.2)
369-369: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
466-615: Refactor theforwardmethod for better maintainability
Theforwardmethod is lengthy and handles multiple responsibilities (coordinate, cell, token predictions, plus excluding atom types). Consider splitting it into smaller helper methods to improve readability, testability, and overall maintainability.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (12)
deepmd/pt/model/atomic_model/denoise_atomic_model.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/task/denoise.py(2 hunks)source/tests/pt/common.py(5 hunks)source/tests/pt/model/test_permutation.py(5 hunks)source/tests/pt/model/test_permutation_denoise.py(0 hunks)source/tests/pt/model/test_rot.py(7 hunks)source/tests/pt/model/test_rot_denoise.py(0 hunks)source/tests/pt/model/test_smooth.py(7 hunks)source/tests/pt/model/test_smooth_denoise.py(0 hunks)source/tests/pt/model/test_trans.py(5 hunks)source/tests/pt/model/test_trans_denoise.py(0 hunks)
💤 Files with no reviewable changes (4)
- source/tests/pt/model/test_rot_denoise.py
- source/tests/pt/model/test_smooth_denoise.py
- source/tests/pt/model/test_trans_denoise.py
- source/tests/pt/model/test_permutation_denoise.py
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/pt/common.py
295-295: Local variable atom_component_out is assigned to but never used
Remove assignment to unused variable atom_component_out
(F841)
deepmd/pt/model/task/denoise.py
67-67: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
261-261: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
369-369: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (30)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- 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: Analyze (python)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Analyze (c-cpp)
- 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 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 C++ (clang, clang)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cpu, cpu)
🔇 Additional comments (31)
deepmd/pt/model/model/__init__.py (2)
36-38: LGTM on the new import.
No issues found with these lines; the import statement correctly references the newDenoiseModel.
280-281: No issues with the denoise condition check.
AssigningmodelclstoDenoiseModelfor the "denoise" fitting type is correct and consistent with the rest of the codebase.source/tests/pt/common.py (5)
69-70: New arrays for strain data.
Defining new lists for strain components is fine. The naming is consistent with the rest of the code.
167-174: Properly collecting batch strain components.
These blocks correctly gather strain data frombatch_output. No issues noted.
195-199: Gathering strain components for tensor output.
Mirrors the logic of the previous block for NumPy arrays. Looks good.
229-238: Concatenating strain arrays in NumPy.
Creates properly shaped zero arrays when no data is found. Implementation is correct.
310-316: Validate inclusion of atom strain components in the denoising results.
Currently, the function returns only"strain_components","updated_coord", and"logits"underdenoise=True. Ifatom_strain_componentsis also needed at inference time, consider including it in the return dictionary.deepmd/pt/model/atomic_model/denoise_atomic_model.py (1)
19-25: Constructor checks for the correct fitting type.
Raising aTypeErroriffittingis not an instance ofDenoiseNetis appropriate for ensuring type safety.source/tests/pt/model/test_trans.py (5)
21-21: Correct import addition for model_denoise.The import of
model_denoiseis correctly added to the imports fromtest_permutation.py, which will be used in the new test class for denoise model testing.
58-60: Good implementation of denoise testing condition.The code correctly checks for the presence of
test_denoiseattribute and sets appropriate test keys for the denoise model. This follows the same pattern as the existingtest_spincondition.
71-71: Properly passing denoise flag to eval_model.The
denoiseparameter is correctly passed to botheval_modelcalls, ensuring consistent testing behavior.Also applies to: 80-80
85-92: Well-structured assertion conditions for new test keys.The assertion conditions have been appropriately extended to include the new denoise-related keys:
strain_components,updated_coord, andlogits. This ensures that the translational invariance properties are properly tested for these outputs.
171-176: Well-implemented test class for denoise model.The new test class
TestDenoiseModelDPA1follows the same pattern as the other test classes, properly initializing the model withmodel_denoiseconfiguration and settingtest_denoisetoTrue.source/tests/pt/model/test_rot.py (5)
21-21: Correct import addition for model_denoise.The import of
model_denoiseis correctly added to the imports fromtest_permutation.py, maintaining consistency with other test files.
55-58: Good implementation of denoise testing condition.The code correctly checks for the presence of
test_denoiseattribute and sets appropriate test keys for testing rotational invariance of denoise model outputs.
73-73: Properly passing denoise flag to eval_model.The
denoiseparameter is correctly passed to alleval_modelcalls, ensuring consistent testing behavior across all test scenarios.Also applies to: 82-82, 125-125, 134-134
86-87: Appropriate rotational invariance checks for denoise outputs.The assertion checks have been correctly modified to test:
- Scalar outputs (
energy,strain_components,logits) for equality- Vector outputs (
force,force_mag,updated_coord) for rotational transformationThis ensures that the rotational invariance properties are properly verified for all model outputs.
Also applies to: 88-91, 138-139, 140-143
224-229: Well-implemented test class for denoise model.The new test class
TestDenoiseModelDPA1correctly initializes the model withmodel_denoiseconfiguration and setstest_denoisetoTrue, maintaining consistency with other test files.source/tests/pt/model/test_smooth.py (5)
21-21: Correct import addition for model_denoise.The import of
model_denoiseis correctly added to the imports fromtest_permutation.py, maintaining consistency with other test files.
87-90: Good implementation of denoise testing condition.The code correctly checks for the presence of
test_denoiseattribute and sets appropriate test keys for testing smoothness of denoise model outputs.
101-101: Properly passing denoise flag to eval_model.The
denoiseparameter is correctly passed to alleval_modelcalls, ensuring consistent testing behavior across all test scenarios in the smoothness tests.Also applies to: 110-110, 119-119, 128-128
134-134: Appropriate smoothness checks for denoise outputs.The assertion checks in the
comparefunction have been correctly extended to include the new denoise-related keys:strain_components,updated_coord, andlogits. This ensures that the smoothness properties are properly tested for these outputs.
262-268: Well-implemented test class for denoise model.The new test class
TestDenoiseModelDPA1correctly initializes the model withmodel_denoiseconfiguration, setstest_denoisetoTrue, and initializesepsilonandaprectoNone, following the pattern of other test classes.source/tests/pt/model/test_permutation.py (5)
380-408: Well-defined model configuration for denoise model.The
model_denoiseconfiguration is comprehensive and properly structured:
- It includes a descriptor of type "se_atten" with appropriate attention-based embedding parameters
- The fitting_net is correctly set to type "denoise" with neural network parameters
- It includes a special token type "MASKED_TOKEN" in the type_map, which is likely used for masking tokens during denoising
This configuration aligns with the PR objective of incorporating noise into the equilibrium structure for denoising tasks.
429-432: Good implementation of denoise testing condition in permutation test.The code correctly checks for the presence of
test_denoiseattribute and sets appropriate test keys for testing permutation invariance of denoise model outputs.
442-442: Properly passing denoise flag to eval_model.The
denoiseparameter is correctly passed to botheval_modelcalls, ensuring consistent testing behavior for permutation invariance tests.Also applies to: 451-451
456-457: Appropriate permutation invariance checks for denoise outputs.The assertion checks have been correctly modified to:
- Check that scalar quantities (
energy,strain_components) are invariant to permutation- Verify that vector quantities (
force,force_mag,updated_coord,logits) transform correctly with permutationThis ensures that permutation invariance is properly tested for all model outputs.
Also applies to: 458-461
539-545: Well-implemented test class for denoise model.The new test class
TestDenoiseModelDPA1correctly initializes the model withmodel_denoiseconfiguration and setstest_denoisetoTrue, maintaining consistency with other test files.deepmd/pt/model/task/denoise.py (3)
405-407: Clarify usage whencoord_noiseis None
Ifcoord_noiseis not set, this method returnsNone, which might cause errors if the caller expects a float. Consider returning 0.0 or raising an appropriate warning.Are you sure all callers handle a
Nonereturn?
If required, I can generate a search script to verify usage across the codebase.
414-422: Validatecase_idxto avoid out-of-bounds
Indexingtorch.eye(self.dim_case_embd)[case_idx]may raise an error ifcase_idx >= self.dim_case_embd.Consider checking the index range before usage:
+if case_idx < 0 or case_idx >= self.dim_case_embd: + raise ValueError(f"Invalid case_idx={case_idx}") self.case_embd = torch.eye(self.dim_case_embd, dtype=self.prec, device=device)[case_idx]
607-611: Confirm zeroing out logic for excluded types
You're zeroing out thestrain_components,updated_coord, andlogitsfor excluded types using the mask. Verify that this is the desired behavior, as it might inadvertently hide data if used incorrectly.Would you like to run a quick script to identify usage patterns of the zeroed outputs across the codebase?
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
deepmd/pt/loss/denoise.py (1)
190-191:⚠️ Potential issueUndefined
self.noisereference.
Here,low=-self.noiseis used, but the class only definesself.coord_noisefor controlling coordinate noise. This likely causes runtime errors or unintended behavior. Replaceself.noisewithself.coord_noiseor another intended variable.-if self.noise_type == "uniform": - noise_on_coord = rng.uniform( - low=-self.noise, high=self.coord_noise, size=(mask_num, 3) - ) +if self.noise_type == "uniform": + noise_on_coord = rng.uniform( + low=-self.coord_noise, high=self.coord_noise, size=(mask_num, 3) + )deepmd/pt/model/task/denoise.py (1)
257-259: 🛠️ Refactor suggestionEliminate mutable default value for
exclude_types.
Similarly, using a mutable default here can lead to unexpected behaviors if the list is manipulated.-def reinit_exclude( - self, - exclude_types: list[int] = [], +def reinit_exclude( + self, + exclude_types: Optional[list[int]] = None, ) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
258-258: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
🧹 Nitpick comments (5)
source/tests/universal/dpmodel/fitting/test_fitting.py (1)
238-279: Avoid using a mutable list as a function default.
Python default arguments likeexclude_types=[]can lead to unexpected behavior. Consider usingNoneas the default and initializing the list inside the function to prevent potential side effects.-def FittingParamDenoise( - ntypes, - dim_descrpt, - mixed_types, - type_map, - exclude_types=[], - precision="float64", - embedding_width=None, - numb_param=0, -): +def FittingParamDenoise( + ntypes, + dim_descrpt, + mixed_types, + type_map, + exclude_types=None, + precision="float64", + embedding_width=None, + numb_param=0, +): + if exclude_types is None: + exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
243-243: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/dpmodel/fitting/denoise_fitting.py (2)
93-114: UseNonefor mutable default arguments.
In the constructor, mutable defaults such asneuron: list[int] = [120, 120, 120]andexclude_types: list[int] = []can unintentionally persist changes across instances. Prefer usingNoneand initializing within the method body to avoid hidden side effects.-def __init__( - self, - ntypes: int, - dim_descrpt: int, - embedding_width: int, - neuron: list[int] = [120, 120, 120], - bias_atom_e: Optional[np.ndarray] = None, - ... - exclude_types: list[int] = [], - ... -): +def __init__( + self, + ntypes: int, + dim_descrpt: int, + embedding_width: int, + neuron: Optional[list[int]] = None, + bias_atom_e: Optional[np.ndarray] = None, + ... + exclude_types: Optional[list[int]] = None, + ... +): + if neuron is None: + neuron = [120, 120, 120] + if exclude_types is None: + exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
97-97: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
107-107: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
501-503: Consider reshapinggronce outside the loop for efficiency.
gris repeatedly reshaped in this loop (for type_i in range(self.ntypes)) with no modifications in between. Reshaping it once beforehand avoids repetitive overhead.- for type_i in range(self.ntypes): - ... - gr = xp.reshape(gr, (nf * nloc, -1, 3)) + gr = xp.reshape(gr, (nf * nloc, -1, 3)) + for type_i in range(self.ntypes): ...deepmd/pt/model/task/denoise.py (2)
357-357: Usefor kk in variables:instead offor kk in variables.keys():
This is more idiomatic and aligns with static analysis suggestions.-for kk in variables.keys(): +for kk in variables: obj[kk] = to_torch_tensor(variables[kk])🧰 Tools
🪛 Ruff (0.8.2)
357-357: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
558-561: Reassess repeatedgr.view(...)usage within loops.
Performinggr = gr.view(-1, self.embedding_width, 3)multiple times inside the loop may introduce confusion. Consider doing the view operation once outside the loop or storing it in a separate variable for clarity.Also applies to: 600-603
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
deepmd/dpmodel/fitting/__init__.py(2 hunks)deepmd/dpmodel/fitting/denoise_fitting.py(1 hunks)deepmd/pt/loss/denoise.py(2 hunks)deepmd/pt/model/atomic_model/denoise_atomic_model.py(1 hunks)deepmd/pt/model/model/__init__.py(3 hunks)deepmd/pt/model/task/__init__.py(2 hunks)deepmd/pt/model/task/denoise.py(2 hunks)deepmd/utils/argcheck.py(2 hunks)source/tests/universal/common/cases/fitting/utils.py(2 hunks)source/tests/universal/dpmodel/fitting/test_fitting.py(2 hunks)source/tests/universal/pt/fitting/test_fitting.py(3 hunks)
✅ Files skipped from review due to trivial changes (1)
- deepmd/pt/model/task/init.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/pt/model/model/init.py
- deepmd/pt/model/atomic_model/denoise_atomic_model.py
🧰 Additional context used
🪛 Ruff (0.8.2)
source/tests/universal/dpmodel/fitting/test_fitting.py
243-243: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/dpmodel/fitting/denoise_fitting.py
97-97: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
107-107: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
311-311: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
359-359: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pt/model/task/denoise.py
56-56: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
66-66: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
258-258: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
357-357: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (28)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test C++ (true)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (clang, clang)
- 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: 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 (c-cpp)
🔇 Additional comments (15)
source/tests/universal/common/cases/fitting/utils.py (3)
79-82: Improved variable name handling for multiple outputsThis change enhances the test framework to support multiple output variables from the fitting module. Converting a string
var_nameto a list of variable names and handling the special case where "polar" needs to be mapped to "polarizability" makes the code more flexible for testing models that produce multiple outputs, like the new denoise model.
88-110: Well-structured result processing for multiple variablesThe updated result extraction and comparison logic properly handles multiple output variables from the fitting module. The code now:
- Stores the forward result in a dedicated variable
- Iterates through each output variable name
- Correctly applies the atom mask to each variable's output shape
- Makes appropriate assertions for each variable
This approach is more robust and aligns with testing requirements for the denoise model which likely needs to handle multiple types of outputs (e.g., token, coordinate, and cell noise).
176-198: Consistent handling of multiple variables in type map testsThe type map test has been updated consistently with the exclude types test to handle multiple variable names. This ensures that all testing paths properly support models that produce multiple outputs, maintaining test consistency across different configurations.
deepmd/dpmodel/fitting/__init__.py (2)
2-4: New import for DenoiseFitting properly integrated.The import of
DenoiseFittingfrom the denoise_fitting module follows the established pattern in this file, maintaining code consistency.
29-29: all list properly updated to expose the new class.The
DenoiseFittingclass has been correctly added to the__all__list, making it available through the package's public API. The addition maintains the alphabetical ordering (DenoiseFitting before DipoleFitting).source/tests/universal/pt/fitting/test_fitting.py (3)
5-5: New import for DenoiseFittingNet properly added.The import of
DenoiseFittingNetfrom the deepmd.pt.model.task module follows the established pattern in this file, maintaining code consistency.
20-20: FittingParamDenoise import properly added.The import of
FittingParamDenoisefrom the test_fitting module follows the established pattern in this file, maintaining code consistency.
39-39: Test case for DenoiseFittingNet properly integrated.The new test case tuple
(FittingParamDenoise, DenoiseFittingNet)has been correctly added to the parameterized test, ensuring that the new denoise fitting functionality will be tested alongside the existing fitting networks.deepmd/utils/argcheck.py (2)
1867-1905: Well-structured implementation of fitting_denoise functionThe implementation follows the same pattern as other fitting functions in the codebase, with proper documentation for each parameter and sensible default values. This function adds the necessary configuration options for the denoising model fitting process.
2812-2929: Comprehensive implementation of loss_denoise functionThe loss_denoise function properly defines all parameters needed for the denoising model's loss calculation, including mask controls (token, coordinate, cell), loss weights, noise configuration, and masking strategies. The previous issue with duplicate token_loss parameter definition has been fixed.
source/tests/universal/dpmodel/fitting/test_fitting.py (2)
8-8: Imports look good.
This import ofDenoiseFittingis clear and consistent with the newly introduced class in the codebase.
288-288: Addition to parameterized tests is appropriate.
Including the(FittingParamDenoise, DenoiseFitting)entry ensures that denoise fitting is tested alongside other fitting parameters. This helps maintain consistency and comprehensive coverage.deepmd/pt/loss/denoise.py (1)
306-310: Implement serialize & deserialize methods.
These methods are currently no-ops and will not function for saving/restoring theDenoiseLossobject. A proper implementation is essential for consistent training and inference workflows.Do you want a code snippet showing a minimal example for serialization/deserialization of the class?
deepmd/pt/model/task/denoise.py (2)
4-4: No concerns.
These import changes appear valid and do not raise any immediate functionality or style issues.Also applies to: 7-7, 16-21, 28-31
70-72: Add unit tests for noise parameters.
coord_noise,cell_pert_fraction, andnoise_typeare newly introduced parameters but are not apparently validated or tested. Ensure coverage in your test suite to verify their usage and expected impact on the model’s output.Do you want me to generate a minimal test harness for these parameters to confirm their effects on the model output?
Also applies to: 393-404
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
deepmd/dpmodel/model/__init__.py (1)
15-41: 💡 Verification agent🧩 Analysis chain
Missing implementation details for DenoiseModel integration.
The import and export of DenoiseModel are correctly implemented, but based on the PR summary, I would expect to see updates to
_get_standard_model_componentsandget_standard_modelfunctions to handle the "denoise" fitting type. These functions would need to be modified to properly initialize the DenoiseModel with appropriate parameters.
🏁 Script executed:
#!/bin/bash # Let's check if these functions exist in this file or related files echo "Checking for _get_standard_model_components and get_standard_model in __init__.py" cat deepmd/dpmodel/model/__init__.py | grep -n "def _get_standard_model_components\|def get_standard_model" echo -e "\nChecking for the functions in other files in the model directory" find deepmd/dpmodel/model -type f -name "*.py" | xargs grep -l "def _get_standard_model_components\|def get_standard_model"Length of output: 583
Action Required: Integrate “denoise” fitting support in model initialization functions
The
DenoiseModelis correctly imported, but the functions responsible for model initialization—namely_get_standard_model_componentsandget_standard_model(located indeepmd/dpmodel/model/model.py)—do not currently handle the "denoise" fitting type. Based on the PR summary, these functions should be updated to initializeDenoiseModelwith appropriate parameters when the "denoise" mode is specified.
- Ensure that
_get_standard_model_componentsincludes logic to recognize the "denoise" fitting type and properly instantiate theDenoiseModel.- Update
get_standard_modelsimilarly so that it supports "denoise" as a valid model configuration.
♻️ Duplicate comments (1)
deepmd/dpmodel/fitting/denoise_fitting.py (1)
448-451:⚠️ Potential issueFix missing f-string in the error message.
The error message is missing the 'f' prefix, so the variables won't be interpolated correctly.
- raise ValueError( - f"get an input descriptor of dim {nd}," - f"which is not consistent with {self.dim_descrpt}." - ) + raise ValueError( + f"get an input descriptor of dim {nd}, " + f"which is not consistent with {self.dim_descrpt}." + )
🧹 Nitpick comments (6)
deepmd/dpmodel/fitting/denoise_fitting.py (4)
98-98: Replace mutable default argument with None.Using mutable objects like lists as default parameters can lead to unexpected behavior if the default value is modified.
- neuron: list[int] = [120, 120, 120], + neuron: Optional[list[int]] = None,Then initialize in the method:
if neuron is None: self.neuron = [120, 120, 120] else: self.neuron = neuron🧰 Tools
🪛 Ruff (0.8.2)
98-98: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
108-108: Replace mutable default argument with None.Using mutable objects like lists as default parameters can lead to unexpected behavior if the default value is modified.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None,Then initialize in the method:
if exclude_types is None: self.exclude_types = [] else: self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
108-108: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
327-327: Replace mutable default argument with None.Using mutable objects like lists as default parameters can lead to unexpected behavior if the default value is modified.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None,Then initialize in the method:
if exclude_types is None: exclude_types = [] self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
327-327: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
375-376: Simplify dictionary key iteration.Using
.keys()is unnecessary when iterating over a dictionary.- for kk in variables.keys(): + for kk in variables:🧰 Tools
🪛 Ruff (0.8.2)
375-375: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/dpmodel/atomic_model/denoise_atomic_model.py (2)
39-54: Add default noise type handling.The code checks for "gaussian" and "uniform" noise types but doesn't handle the case when noise_type is None. Consider adding a default behavior for when noise_type is not specified.
noise_type = self.fitting_net.get_noise_type() + if noise_type is None: + # Return ret unchanged when no noise type is specified + return ret cell_std = self.fitting_net.get_cell_pert_fraction() / 1.732 if noise_type == "gaussian": coord_std = self.fitting_net.get_coord_noise() elif noise_type == "uniform": coord_std = self.fitting_net.get_coord_noise() / 1.732 else: raise RuntimeError(f"Unknown noise type {noise_type}")
47-54: Add comments explaining the scaling operations.The scaling by cell_std and coord_std is important, but it would be helpful to add comments explaining the mathematical reasoning behind these transformations, especially the 1.732 divisor (which is likely √3).
ret["strain_components"] = ( ret["strain_components"] * cell_std if cell_std > 0 else ret["strain_components"] ) + # Scale coordinates by standard deviation to match the noise magnitude + # 1.732 is approximately √3, used to convert between uniform and Gaussian distributions ret["updated_coord"] = ( ret["updated_coord"] * coord_std if coord_std > 0 else ret["updated_coord"] )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (11)
deepmd/dpmodel/atomic_model/__init__.py(2 hunks)deepmd/dpmodel/atomic_model/denoise_atomic_model.py(1 hunks)deepmd/dpmodel/fitting/__init__.py(2 hunks)deepmd/dpmodel/fitting/denoise_fitting.py(1 hunks)deepmd/dpmodel/model/__init__.py(2 hunks)deepmd/dpmodel/model/denoise_model.py(1 hunks)deepmd/dpmodel/model/model.py(3 hunks)deepmd/pt/model/atomic_model/denoise_atomic_model.py(1 hunks)source/tests/universal/common/cases/atomic_model/atomic_model.py(1 hunks)source/tests/universal/dpmodel/atomc_model/test_atomic_model.py(4 hunks)source/tests/universal/pt/atomc_model/test_atomic_model.py(4 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/dpmodel/fitting/init.py
- deepmd/pt/model/atomic_model/denoise_atomic_model.py
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/dpmodel/model/model.py
67-67: Undefined name model_params
(F821)
68-68: Undefined name model_params
(F821)
71-71: Undefined name model_params
(F821)
72-72: Undefined name model_params
(F821)
73-73: Undefined name model_params
(F821)
deepmd/dpmodel/fitting/denoise_fitting.py
98-98: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
108-108: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
327-327: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
375-375: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- 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: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (22)
source/tests/universal/common/cases/atomic_model/atomic_model.py (1)
132-153: Code follows established patterns and correctly defines expected outputs for the DenoiseAtomicModel.The new test class properly defines the expected outputs for the denoise model, including strain_components, updated_coord, logits, and mask. The model_output_equivariant list correctly includes only "updated_coord" which indicates it should be invariant to rotations and translations.
source/tests/universal/pt/atomc_model/test_atomic_model.py (4)
18-18: Import for DenoiseFittingNet looks good.Added import is necessary for the new denoise functionality.
30-30: Import for DenoiseAtomicModelTest looks good.This import brings in the base test class defined in the first file.
54-55: Fitting parameter imports look good.These imports are necessary for parameterizing the test class with denoise fitting parameters.
468-528: Test class implementation follows established patterns.The implementation of
TestDenoiseAtomicModelPTcorrectly follows the established pattern in the file, parameterizing the test with appropriate descriptor and fitting parameters. The class inherits from the right base classes and thesetUpClassmethod initializes all necessary attributes.source/tests/universal/dpmodel/atomc_model/test_atomic_model.py (4)
18-18: Import for DenoiseFitting looks good.Added import is necessary for the new denoise functionality in the deepmd implementation.
34-34: Import for DenoiseAtomicModelTest looks good.This import brings in the base test class for testing the denoise atomic model.
64-65: Fitting parameter imports look good.These imports are necessary for parameterizing the test class with denoise fitting parameters.
478-535: Test class implementation follows established patterns.The implementation of
TestDenoiseAtomicModelDPcorrectly follows the established pattern in the file. The class inherits from the right base classes, and the parameterization includes appropriate descriptor and fitting parameters. ThesetUpClassmethod correctly initializes all necessary attributes.Note that this class is properly skipped on non-CPU devices in CI environments, which is consistent with the other test classes in this file.
deepmd/dpmodel/model/__init__.py (2)
15-17: Good addition of DenoiseModel import.This import is necessary to make the new DenoiseModel class available from the module. It follows the existing pattern in the file of importing model classes from their respective modules.
36-36: Appropriate addition to__all__list.Adding DenoiseModel to the
__all__list correctly exposes it as part of the public API, making it accessible when users import from this package. This is consistent with how other model classes are handled.deepmd/dpmodel/model/denoise_model.py (1)
1-28: LGTM! The DenoiseModel class is well-structured and follows the established pattern.The implementation properly inherits from both DPModelCommon and the dynamically created DPDenoiseModel_. The class is correctly registered with BaseModel under the name "denoise" and the initialization properly calls both parent constructors. The code structure follows the project's patterns for model creation.
deepmd/dpmodel/fitting/denoise_fitting.py (5)
127-130: LGTM! Good addition of noise-related parameters.The parameters for coordinate noise, cell perturbation, and noise type are properly initialized and will help with the denoising functionality.
176-229: LGTM! Network collections are well-implemented.The implementation of coord_nets, cell_nets, and token_nets follows a consistent pattern and properly handles mixed and non-mixed types. The seed forwarding is properly done with child_seed to ensure reproducibility.
252-262: LGTM! Good accessor methods for noise parameters.The getter methods for coord_noise, cell_pert_fraction, and noise_type are well-implemented and provide a clean interface for accessing these parameters.
551-552: Ensure thenfvariable is used consistently.Previous static analysis flagged potential uses of undefined variable
nframesin this file. Ensure you're usingnf(defined on line 445) consistently throughout the code.#!/bin/bash # Check if nframes is referenced elsewhere in the code rg "\bnframes\b" --include "*.py" | grep -v "denoise_fitting.py"
516-520:❓ Verification inconclusive
Verify the matrix multiplication and reshaping operations.
This section handles reshaping the gradient calculations for the coordinate predictions. Make sure the dimensions are correct for the matrix multiplication operation between updated_coord_type and gr.
🏁 Script executed:
#!/bin/bash # Check for unit tests that verify this calculation rg -A 10 -B 10 "updated_coord_type.*@.*gr" --include "*.py" | grep -v "denoise_fitting.py"Length of output: 532
#!/bin/bash # Re-run the search using '--glob' instead of '--include' to find references to # the matrix multiplication between `updated_coord_type` and `gr`, excluding tests # in "denoise_fitting.py". rg -A 10 -B 10 "updated_coord_type\s*@\s*gr" --glob "*.py" | grep -v "denoise_fitting.py"Matrix multiplication and reshaping dimensions require manual verification
It seems our initial attempt to locate unit tests for the matrix multiplication between
updated_coord_typeandgrresulted in an error due to the use of an unsupported flag. Please verify that the reshaping ofgrto(nf * nloc, -1, 3)and the subsequent multiplication withupdated_coord_typeare dimensionally consistent. If no unit tests exist for this operation, consider adding tests to ensure that the output dimensions match the expectation—specifically that after the multiplication and reshaping,updated_coord_typetruly conforms to the shape(nf, nloc, 3).deepmd/dpmodel/atomic_model/denoise_atomic_model.py (1)
13-20: LGTM! Good type checking in the constructor.The constructor properly verifies that the fitting parameter is an instance of DenoiseFitting, which helps prevent incorrect usage.
deepmd/dpmodel/atomic_model/__init__.py (2)
20-22: LGTM! Proper module import.The DPDenoiseAtomicModel is correctly imported from the denoise_atomic_model module.
56-56: LGTM! Proper addition to module exports.DPDenoiseAtomicModel is correctly added to the all list, making it available when importing from this module.
deepmd/dpmodel/model/model.py (2)
22-24: Import of DenoiseModel looks good.This import adds the new DenoiseModel class that will be used later in the file, following the existing import pattern.
110-111: DenoiseModel integration looks correct.Adding the DenoiseModel to the model class selection follows the existing pattern for other model types and is correctly implemented.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
deepmd/dpmodel/model/model.py (1)
66-73:⚠️ Potential issueUndefined variable
model_paramscauses a NameError.Line 68 mistakenly references
model_paramsinstead ofdata. Apply this diff to fix the assertion message:- f"When using denoise fitting, the last element in 'type_map' must be 'MASKED_TOKEN', but got '{model_params['type_map'][-1]}'" + f"When using denoise fitting, the last element in 'type_map' must be 'MASKED_TOKEN', but got '{data['type_map'][-1]}'"🧰 Tools
🪛 Ruff (0.8.2)
68-68: Undefined name
model_params(F821)
🧹 Nitpick comments (2)
source/tests/universal/dpmodel/loss/test_loss.py (2)
209-234: Consider adding function parameters for customizationThe
LossParamDenoisefunction is well-structured but differs from other similar functions in this file by not accepting any parameters for customization. Consider adding parameters with default values (similar to other loss parameter functions) to allow more flexible testing.-def LossParamDenoise(): +def LossParamDenoise( + token_loss=1.0, + coord_loss=1.0, + cell_loss=1.0, + coord_noise=0.2, + mask_prob=0.2, + noise_type="gaussian", + noise_mode="prob" +):Also, consider adding a brief docstring explaining the purpose of this function and the meaning of key parameters in the context of the denoising model.
236-238: Consider using parameterize_func for consistencyThe other loss parameter functions in this file use
parameterize_functo create variations of parameter configurations for testing. Consider using the same pattern forLossParamDenoiseListto maintain consistency and enable testing with different parameter combinations.-LossParamDenoiseList = [LossParamDenoise] +LossParamDenoiseList = parameterize_func( + LossParamDenoise, + OrderedDict( + { + "token_loss": (1.0, 0.0), + "coord_loss": (1.0, 0.0), + "cell_loss": (1.0, 0.0), + "noise_type": ("gaussian", "uniform"), + } + ), +)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/dpmodel/model/model.py(3 hunks)deepmd/pt/loss/__init__.py(1 hunks)source/tests/universal/dpmodel/loss/test_loss.py(1 hunks)
✅ Files skipped from review due to trivial changes (1)
- deepmd/pt/loss/init.py
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/dpmodel/model/model.py
68-68: Undefined name model_params
(F821)
⏰ Context from checks skipped due to timeout of 90000ms (30)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- 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: Build C++ (clang, clang)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- 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)
🔇 Additional comments (3)
source/tests/universal/dpmodel/loss/test_loss.py (1)
224-232:❓ Verification inconclusive
Verify the supported values for string parameters
The parameters
noise_type,noise_mode, andloss_funcare set to specific string values. Please verify that these values match the options supported by theDenoiseLossclass implementation.For maintainability, consider using constants or enums for these string parameters to prevent typos and make it easier to update if supported values change.
🏁 Script executed:
#!/bin/bash # Check for valid string parameter values in the DenoiseLoss class rg -A 10 "noise_type|noise_mode|loss_func" --type py source/deepmd/pt/loss/denoise.pyLength of output: 160
Action Needed: Verify Parameter Values and File Path
It appears that the file at
source/deepmd/pt/loss/denoise.pycouldn't be located during the verification. Before finalizing any changes in the test parameters, please ensure that:
- The implementation for the
DenoiseLossclass (which should support thenoise_type,noise_mode, andloss_funcoptions) is present in the repository. If the file was moved or renamed, update the reference in your test or verification command accordingly.- The string values assigned in the test (e.g.,
"gaussian","prob","rmse") are consistent with the valid options defined in the actual implementation ofDenoiseLoss.- For better maintainability, consider using constants or enums to manage these parameters, thereby reducing the risk of typos and simplifying future updates.
Once you've confirmed the correct file location and verified the valid parameter options, the test and documentation can be updated accordingly.
deepmd/dpmodel/model/model.py (2)
22-24: Looks good: ImportingDenoiseModelis correct.
110-111: No issues with the new denoise model assignment.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (12)
deepmd/dpmodel/fitting/denoise_fitting.py (6)
98-98: Use None instead of a mutable list for the default value.Using a mutable data structure like a list as a default parameter value can lead to unexpected behavior if modified.
-neuron: list[int] = [120, 120, 120], +neuron: Optional[list[int]] = None,And then initialize it in the constructor:
if self.neuron is None: self.neuron = [120, 120, 120]🧰 Tools
🪛 Ruff (0.8.2)
98-98: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
108-108: Use None instead of a mutable list for the default value.Using an empty list as the default parameter value creates a shared mutable object across function calls.
-exclude_types: list[int] = [], +exclude_types: Optional[list[int]] = None,And then initialize it in the constructor:
if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
108-108: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
175-228: Consider refactoring duplicate network initialization code.The initialization code for the three networks (coord_nets, cell_nets, token_nets) is very similar with only minor differences. This could be refactored into a helper method to reduce code duplication.
-self.coord_nets = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim, - self.embedding_width, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - seed=child_seed(seed, ii), - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], -) -self.cell_nets = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim, - 6, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - seed=child_seed(self.seed, ii), - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], -) -self.token_nets = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim, - self.ntypes - 1, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - seed=child_seed(self.seed, ii), - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], -)A helper method could be added:
def _create_network_collection(self, in_dim, out_dim, seed): """Create a network collection with specified parameters.""" return NetworkCollection( 1 if not self.mixed_types else 0, self.ntypes, network_type="fitting_network", networks=[ FittingNet( in_dim, out_dim, self.neuron, self.activation_function, self.resnet_dt, self.precision, bias_out=True, seed=child_seed(seed, ii), ) for ii in range(self.ntypes if not self.mixed_types else 1) ], )Then use it like:
self.coord_nets = self._create_network_collection(in_dim, self.embedding_width, seed) self.cell_nets = self._create_network_collection(in_dim, 6, seed) self.token_nets = self._create_network_collection(in_dim, self.ntypes - 1, seed)
325-330: Use None instead of a mutable list for the default value.Using an empty list as a default parameter value can lead to shared state between function calls.
-def reinit_exclude( - self, - exclude_types: list[int] = [], -) -> None: +def reinit_exclude( + self, + exclude_types: Optional[list[int]] = None, +) -> None: + if exclude_types is None: + exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
327-327: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
378-379: Use simplifiedforloop without.keys().The
.keys()method is unnecessary here as dictionaries are directly iterable.-for kk in variables.keys(): +for kk in variables:🧰 Tools
🪛 Ruff (0.8.2)
378-378: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
497-546: Consider simplifying the nested loop structure.The code has multiple similar loops over atom types to apply fitting operations. This could potentially be refactored to reduce duplication and improve readability.
# For example, a helper method like: def _apply_network_to_type(self, network, xx, atype, output_shape, concat_dimension=None): result = xp.zeros(output_shape, dtype=get_xp_precision(xp, self.precision)) for type_i in range(self.ntypes): dims = [nf, nloc, 1] if concat_dimension is not None: dims[2] = concat_dimension mask = xp.tile(xp.reshape((atype == type_i), [nf, nloc, 1]), tuple(dims)) result_type = network[(type_i,)](xx) result_type = xp.where(mask, result_type, xp.zeros_like(result_type)) result = result + result_type return resultThen call it for each network operation. This would make the code more DRY and easier to maintain.
source/tests/universal/pt/model/test_model.py (1)
975-976: Remove commented-out code.There's commented-out code that appears to be related to setting special precision for DPA2 descriptors. Either uncomment and use it if needed or remove it entirely to keep the code clean.
-# if Descrpt in [DescrptDPA2]: -# cls.epsilon_dict["test_smooth"] = 1e-8🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 975-976: Commented-out code
This comment appears to contain commented-out code.source/tests/universal/dpmodel/fitting/test_fitting.py (1)
243-243: Fix mutable default argumentUsing mutable objects like an empty list as default arguments can lead to unexpected behavior if the list is modified elsewhere.
- exclude_types=[], + exclude_types=None,Then inside the function:
input_dict = { "ntypes": ntypes, "dim_descrpt": dim_descrpt, "mixed_types": mixed_types, "type_map": type_map, - "exclude_types": exclude_types, + "exclude_types": exclude_types if exclude_types is not None else [],🧰 Tools
🪛 Ruff (0.8.2)
243-243: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/pt/model/task/denoise.py (4)
56-56: Fix mutable default argumentUsing mutable objects like lists as default arguments can lead to unexpected behavior.
- neuron: list[int] = [128, 128, 128], + neuron: Optional[list[int]] = None,Then initialize inside the method:
if neuron is None: neuron = [128, 128, 128]🧰 Tools
🪛 Ruff (0.8.2)
56-56: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
66-66: Fix mutable default argumentSame issue with mutable default argument.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None,Then initialize inside the method:
if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
66-66: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
256-261: Fix mutable default argument inreinit_excludeSame issue with mutable default argument.
- exclude_types: list[int] = [], + exclude_types: Optional[list[int]] = None,Then initialize inside the method:
if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
258-258: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
360-361: Simplify dictionary keys iterationThe code can be simplified by using direct iteration over the dictionary.
-for kk in variables.keys(): +for kk in variables: obj[kk] = to_torch_tensor(variables[kk])🧰 Tools
🪛 Ruff (0.8.2)
360-360: 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 (7)
deepmd/dpmodel/fitting/denoise_fitting.py(1 hunks)deepmd/dpmodel/model/model.py(3 hunks)deepmd/pt/model/model/denoise_model.py(1 hunks)deepmd/pt/model/task/denoise.py(2 hunks)source/tests/universal/common/cases/model/model.py(1 hunks)source/tests/universal/dpmodel/fitting/test_fitting.py(2 hunks)source/tests/universal/pt/model/test_model.py(6 hunks)
🧰 Additional context used
🧬 Code Definitions (5)
source/tests/universal/common/cases/model/model.py (2)
source/tests/universal/common/cases/model/utils.py (1) (1)
ModelTestCase(30-979)source/tests/universal/common/cases/atomic_model/atomic_model.py (7) (7)
setUpClass(14-28)setUpClass(33-47)setUpClass(52-66)setUpClass(71-85)setUpClass(90-110)setUpClass(115-129)setUpClass(134-153)
source/tests/universal/dpmodel/fitting/test_fitting.py (1)
deepmd/dpmodel/fitting/denoise_fitting.py (1) (1)
DenoiseFitting(49-579)
source/tests/universal/pt/model/test_model.py (4)
deepmd/pt/model/model/denoise_model.py (2) (2)
DenoiseModel(26-104)translated_output_def(37-47)deepmd/pt/model/task/denoise.py (2) (2)
DenoiseFittingNet(50-637)output_def(287-313)source/tests/universal/common/cases/model/model.py (10) (10)
DenoiseModelTest(179-201)setUpClass(14-28)setUpClass(33-47)setUpClass(52-67)setUpClass(72-87)setUpClass(92-107)setUpClass(112-132)setUpClass(137-156)setUpClass(161-176)setUpClass(181-201)source/tests/universal/pt/backend.py (3) (3)
PTTestCase(19-114)modules_to_test(50-55)script_module(32-36)
deepmd/pt/model/task/denoise.py (1)
deepmd/dpmodel/fitting/denoise_fitting.py (7) (7)
reinit_exclude(325-330)change_type_map(271-291)output_def(385-411)get_dim_fparam(231-233)get_dim_aparam(235-237)get_sel_type(239-246)get_type_map(248-250)
deepmd/dpmodel/fitting/denoise_fitting.py (1)
deepmd/pt/model/task/denoise.py (13) (13)
output_def(287-313)reinit_exclude(256-261)get_dim_fparam(367-369)get_dim_aparam(371-373)get_sel_type(378-390)get_type_map(392-394)get_coord_noise(396-398)get_cell_pert_fraction(400-402)get_noise_type(404-406)set_case_embd(408-415)change_type_map(263-285)serialize(315-348)deserialize(351-365)
🪛 Ruff (0.8.2)
source/tests/universal/dpmodel/fitting/test_fitting.py
243-243: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
deepmd/pt/model/task/denoise.py
56-56: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
66-66: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
258-258: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
360-360: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/dpmodel/fitting/denoise_fitting.py
98-98: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
108-108: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
327-327: Do not use mutable data structures for argument defaults
Replace with None; initialize within function
(B006)
378-378: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
🪛 GitHub Check: CodeQL
source/tests/universal/pt/model/test_model.py
[notice] 975-976: Commented-out code
This comment appears to contain commented-out code.
⏰ Context from checks skipped due to timeout of 90000ms (27)
- 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 wheels for cp310-manylinux_aarch64
- 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: 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: Test C++ (false)
- GitHub Check: Test C++ (true)
- 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: Analyze (c-cpp)
🔇 Additional comments (23)
deepmd/dpmodel/fitting/denoise_fitting.py (1)
1-580: Overall assessment of DenoiseFitting implementation.The DenoiseFitting class is well-structured and follows the established patterns in the codebase. The implementation correctly handles the three types of fittings (coordinate, cell, and token) and properly manages different parameter types. The serialization and deserialization methods ensure the model can be saved and loaded consistently.
A few minor improvements could be made to address mutable default arguments and reduce code duplication, but the overall implementation is solid.
🧰 Tools
🪛 Ruff (0.8.2)
98-98: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
108-108: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
327-327: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
378-378: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/dpmodel/model/model.py (2)
66-73: Good implementation of denoise model configuration.The changes correctly handle the special case for denoise fitting by:
- Adding an assertion to ensure the last element in
type_mapis 'MASKED_TOKEN'- Setting appropriate default values for noise parameters
- Properly initializing the embedding width
This is consistent with how other model types are handled in this file.
110-111: Clean integration of DenoiseModel for denoise fitting.The addition of DenoiseModel as the model class for 'denoise' fitting type maintains the same pattern used for other fitting types in the function. This ensures consistent behavior and error handling.
source/tests/universal/common/cases/model/model.py (1)
179-201: Good test class implementation for Denoise model.The DenoiseModelTest class correctly follows the established pattern for test cases in this file. The expected outputs, especially
expected_model_output_typeandmodel_output_equivariant, align with the actual outputs of the DenoiseModel implementation.The class properly defines:
- Expected model output types (strain_components, updated_coord, logits, mask)
- Identification of equivariant outputs (updated_coord)
- Test parameters consistent with other test classes
The skip_test_autodiff flag is appropriately set due to the nature of the denoise model.
source/tests/universal/pt/model/test_model.py (2)
930-954: Good parameterized test setup for DenoiseModel.The parameterization includes appropriate descriptor and fitting parameter combinations for testing the denoise model. The test follows the same pattern established for other model tests in this file.
955-1019: Good implementation of TestDenoiseModelPT.The test class properly inherits from the base test classes and follows the same pattern as other model tests. It correctly:
- Sets up the model with the appropriate descriptor and fitting
- Handles JIT compilation for testing the scripted model
- Sets expected values based on the model configuration
- Provides proper inheritance for test methods
This ensures the denoise model will be thoroughly tested in the same way as other model types.
🧰 Tools
🪛 GitHub Check: CodeQL
[notice] 975-976: Commented-out code
This comment appears to contain commented-out code.source/tests/universal/dpmodel/fitting/test_fitting.py (3)
238-267: Add support for denoise model fitting parametersThe implementation adds a new function
FittingParamDenoisethat creates parameter dictionaries for testing the denoise fitting functionality, consistent with the existing pattern of fitting parameter functions in the file.🧰 Tools
🪛 Ruff (0.8.2)
243-243: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
270-281: Good usage of parameterize_func patternThe implementation properly follows the established pattern of creating a parameterized list and making the default function accessible through the first element of that list.
291-291: Update test class to include new denoise fitting typeThe test class is appropriately updated to include the new denoise fitting type in the parameterized test matrix.
deepmd/pt/model/model/denoise_model.py (5)
22-22: Good model creation patternUsing the
make_modelfunction to create a base model from the atomic model maintains consistency with the existing codebase architecture.
37-47: Complete implementation oftranslated_output_defThe method properly defines output mappings from the internal model to the external API, handling all required outputs including the optional "mask".
49-73: Well-structuredforwardmethodThe method correctly delegates to the parent class's
forward_commonand constructs the appropriate output dictionary, maintaining consistency with other model implementations.
75-104: Complete implementation offorward_lowermethodThe method properly handles the lower-level operations needed for TorchScript export, passing the necessary parameters to
forward_common_lowerand constructing the expected output dictionary.
25-36:✅ Verification successful
Class registration and initialization
The
DenoiseModelclass is properly registered with theBaseModelregistry and initializes both parent classes correctly.One minor oddity is that
model_typeis set to "property" even though the class is registered as "denoise". Verify if this is intentional.
🏁 Script executed:
#!/bin/bash # Check if other model classes follow the same pattern grep -r "model_type =" --include="*.py" deepmd/pt/model/model/Length of output: 863
Clarification on Model Type Assignment
After checking the similar pattern in other model classes (e.g., inproperty_model.pywheremodel_typeis also set to"property"), it appears that the assignment inDenoiseModelis intentional. If the denoise model is designed to function under a "property" classification, then no changes are needed. Otherwise, please clarify the design intent in the documentation to avoid confusion.deepmd/pt/model/task/denoise.py (9)
48-50: Registration with fitting type "denoise"The
DenoiseFittingNetclass is properly registered with theFittingregistry and is decorated withfitting_check_outputfor consistency with other fitting types.
51-74: Comprehensive constructor with denoise-specific parametersThe constructor includes all necessary parameters for denoising functionality including
coord_noise,cell_pert_fraction, andnoise_type.🧰 Tools
🪛 Ruff (0.8.2)
56-56: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
66-66: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
75-116: Detailed docstringThe docstring provides comprehensive information about all parameters, making the API clear and easy to use.
117-143: Proper initialization of class attributesThe initialization logic properly sets up all required attributes and includes error handling for parameter validation.
175-177: Safe handling of optional parametersThe code properly handles the case when parameters are not present by setting them to None.
195-250: Network initializationThe code creates three separate network collections for coordinate, cell, and token prediction, properly handling both mixed and non-mixed type configurations.
367-407: Good accessor methods implementationThe accessor methods are well-implemented and properly documented, providing a clean API for retrieving model parameters.
553-627: Comprehensive handling of mixed and non-mixed typesThe forward method properly handles both mixed and non-mixed types configurations, applying the appropriate networks and masks for each atom type.
628-637: Proper exclusion handlingThe code correctly applies the exclusion mask to zero out contributions from excluded atom types before returning the results.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (2)
deepmd/pt/loss/denoise.py (2)
191-193:⚠️ Potential issueVariable naming mismatch.
You're using
self.noisethoughself.coord_noiseis declared in the initialization, leading to a potentialAttributeError.- noise_on_coord = rng.uniform( - low=-self.noise, high=self.coord_noise, size=(mask_num, 3) - ) + noise_on_coord = rng.uniform( + low=-self.coord_noise, high=self.coord_noise, size=(mask_num, 3) + )
254-259:⚠️ Potential issueModel output mismatch risk.
The code assumes
model_predcontains "updated_coord" and "strain_components" keys, but these may not be provided if the respective masking is not enabled.Add conditional checks before accessing these fields:
- diff_coord = (label["updated_coord"] - model_pred["updated_coord"]).reshape(-1) - diff_cell = ( - label["strain_components"] - model_pred["strain_components"] - ).reshape(-1) + diff_coord = torch.zeros(0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) + diff_cell = torch.zeros(0, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) + + if self.mask_coord and "updated_coord" in model_pred and "updated_coord" in label: + diff_coord = (label["updated_coord"] - model_pred["updated_coord"]).reshape(-1) + + if self.mask_cell and "strain_components" in model_pred and "strain_components" in label: + diff_cell = ( + label["strain_components"] - model_pred["strain_components"] + ).reshape(-1)
🧹 Nitpick comments (4)
deepmd/pt/utils/stat.py (1)
284-286: Appropriate check to skip bias computation for denoise mode!The conditional check properly prevents bias computation for keys that are specific to the denoise model. Since labels are created in the loss function in denoise mode, this approach correctly avoids unnecessary computations.
Consider expanding the comment to explain briefly why these specific keys ("strain_components", "updated_coord", "logits") are special cases in the denoise model context. This would provide more context for future maintainers.
source/tests/universal/common/cases/loss/utils.py (1)
58-82: Well-structured conditional handling for denoise model testing!The conditional block correctly provides the specialized input data structure needed for the denoise model, while preserving backward compatibility for other tests.
Consider extracting the conditional logic into a separate helper method to improve readability, especially if more specialized test cases are added in the future. This would make the
test_forwardmethod cleaner and more maintainable.source/tests/universal/pt/model/test_model.py (1)
114-115: Good addition to the default fitting parameters list!The
FittingParamDenoiseis correctly added to the list of default fitting parameters.Note that there's a typo in the variable name
defalut_fit_param(should be "default"), but this is in the existing code and not part of your changes.deepmd/pt/loss/denoise.py (1)
301-327: Consider updating label requirements to match actual usage.The
label_requirementproperty marks all items as optional (must=False), but the forward method assumes these outputs are present when the corresponding masking is enabled.Consider modifying the requirements to reflect the actual dependencies based on the enabled masking options. For example, "logits" should have
must=Truewhenself.mask_tokenis True.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deepmd/pt/loss/denoise.py(2 hunks)deepmd/pt/utils/stat.py(1 hunks)source/tests/universal/common/cases/loss/utils.py(2 hunks)source/tests/universal/pt/loss/test_loss.py(3 hunks)source/tests/universal/pt/model/test_model.py(6 hunks)
🧰 Additional context used
🧬 Code Definitions (2)
source/tests/universal/pt/loss/test_loss.py (1)
deepmd/pt/loss/denoise.py (1)
DenoiseLoss(41-327)
source/tests/universal/pt/model/test_model.py (4)
deepmd/pt/model/model/denoise_model.py (2)
DenoiseModel(26-104)translated_output_def(37-47)deepmd/pt/model/task/denoise.py (2)
DenoiseFittingNet(50-637)output_def(287-313)source/tests/universal/common/cases/model/model.py (10)
DenoiseModelTest(179-201)setUpClass(14-28)setUpClass(33-47)setUpClass(52-67)setUpClass(72-87)setUpClass(92-107)setUpClass(112-132)setUpClass(137-156)setUpClass(161-176)setUpClass(181-201)source/tests/universal/dpmodel/fitting/test_fitting.py (1)
FittingParamDenoise(238-267)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- 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: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
🔇 Additional comments (12)
source/tests/universal/pt/loss/test_loss.py (3)
4-11: Good addition of the denoise loss import!The import of
DenoiseLossfollows the existing pattern in the file and integrates well with the existing testing structure.
19-26: Great inclusion of the parameter list import!The import of
LossParamDenoiseListis correctly placed alongside other testing parameter lists.
32-40: Well-structured parameterized test addition!The new test parameter pairing for
DenoiseLossfollows the existing pattern and ensures comprehensive test coverage for the denoising functionality.source/tests/universal/common/cases/loss/utils.py (1)
38-38: Good update to function signature to support kwargs!The addition of
**kwargsto the function signature allows for greater flexibility and future extensibility.source/tests/universal/pt/model/test_model.py (6)
21-22: Good addition of DenoiseModel import!The import is correctly placed alongside other model imports.
32-33: Good addition of DenoiseFittingNet import!The import is correctly placed alongside other fitting net imports.
47-48: Good addition of DenoiseModelTest import!The import is correctly placed alongside other model test imports.
77-79: Good addition of fitting parameter imports!The imports for
FittingParamDenoiseandFittingParamDenoiseListare correctly placed with other fitting parameter imports.
929-954: Well-structured parameterized test setup for DenoiseModel!The parameterized test setup follows the same pattern as other model tests, providing comprehensive test coverage for various descriptor types and fitting parameters.
955-1016: Well-implemented DenoiseModelPT test class!The
TestDenoiseModelPTclass properly inherits from the necessary base classes and follows the established pattern for model testing. The implementation is consistent with other model test classes in the file, ensuring proper test coverage for the denoise model functionality.deepmd/pt/loss/denoise.py (2)
23-38: Good implementation of cell perturbation matrix.The new
get_cell_perturb_matrixfunction properly handles the generation of a perturbation matrix for the unit cell with input validation. This replaces the previously undefinedget_cell_perturb_matrix_HEAfunction mentioned in a past review.Note: There's a TODO comment about allowing users to fix specific components. Consider documenting this as a future enhancement in your PR description.
41-110: LGTM: Comprehensive initialization with clear documentation.The class initialization has been significantly improved with more flexible and granular control over masking, noise types, and loss computation. The docstring thoroughly documents each parameter.
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (3)
deepmd/pt/loss/denoise.py (3)
181-181:⚠️ Potential issueMissing
raisekeyword for exception.The
NotImplementedErroris created but not raised, which means the error will be created but never thrown.- raise NotImplementedError(f"Unknown noise mode {self.noise_mode}!") + raise NotImplementedError(f"Unknown noise mode {self.noise_mode}!")
193-193:⚠️ Potential issueVariable naming mismatch.
The code uses
self.noisethoughself.coord_noiseis declared in the constructor, which will cause an error.- low=-self.noise, high=self.coord_noise, size=(mask_num, 3) + low=-self.coord_noise, high=self.coord_noise, size=(mask_num, 3)
290-299:⚠️ Potential issueConditional token loss calculation needed.
Token loss is computed regardless of whether
self.mask_tokenis enabled, which could cause errors if the model doesn't provide "logits" when token masking is disabled.-# token loss -type_mask = label["type_mask"] -masked_logits = model_pred["logits"][type_mask] -masked_target = label["clean_type"][type_mask] -token_loss = F.nll_loss( - F.log_softmax(masked_logits, dim=-1), - masked_target, - reduction="mean", -) -more_loss["token_loss"] = token_loss.detach() -loss += self.token_loss * token_loss.to(GLOBAL_PT_FLOAT_PRECISION) +# token loss +if self.mask_token and "logits" in model_pred and "type_mask" in label: + type_mask = label["type_mask"] + masked_logits = model_pred["logits"][type_mask] + masked_target = label["clean_type"][type_mask] + token_loss = F.nll_loss( + F.log_softmax(masked_logits, dim=-1), + masked_target, + reduction="mean", + ) + more_loss["token_loss"] = token_loss.detach() + loss += self.token_loss * token_loss.to(GLOBAL_PT_FLOAT_PRECISION)
🧹 Nitpick comments (6)
source/tests/universal/common/cases/loss/utils.py (4)
38-38: Consider using the passed keyword arguments in the fake_model function.The
fake_modelfunction has been updated to accept keyword arguments (**kwargs), but these arguments are not used within the function body. If they're meant to modify the behavior of the function, consider updating the implementation to utilize them.
80-84: Check the data type for atomic types tensor.The
atypetensor is created withdtype=env.GLOBAL_PT_FLOAT_PRECISION, but atomic types are typically represented as integers. Consider using an appropriate integer data type instead of floating-point.input_dict["atype"] = torch.tensor( [[0] * natoms, [0] * natoms], - dtype=env.GLOBAL_PT_FLOAT_PRECISION, + dtype=torch.long, device=env.DEVICE, )
58-104: Refactor module calls to reduce code duplication.The module function is called twice with very similar arguments, differing only in the
input_dictparameter. Consider refactoring to reduce duplication.+ input_dict = {} + if "updated_coord" in self.key_to_pref_map: + import torch + + from deepmd.pt.utils import ( + env, + ) + + labels.update( + { + "type_mask": torch.tensor( + [[False] * natoms, [False] * natoms], + dtype=torch.bool, + device=env.DEVICE, + ) + } + ) + input_dict["box"] = torch.tensor( + [[1, 0, 0, 0, 1, 0, 0, 0, 1]] * nframes, + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + input_dict["atype"] = torch.tensor( + [[0] * natoms, [0] * natoms], + dtype=torch.long, + device=env.DEVICE, + ) + input_dict["coord"] = torch.tensor( + [[[0] * 3] * natoms] * nframes, + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, + ) + + _, loss, more_loss = module( + input_dict, + fake_model, + labels, + natoms, + 1.0, + ) - if "updated_coord" in self.key_to_pref_map: - import torch - - from deepmd.pt.utils import ( - env, - ) - - labels.update( - { - "type_mask": torch.tensor( - [[False] * natoms, [False] * natoms], - dtype=torch.bool, - device=env.DEVICE, - ) - } - ) - input_dict = {} - input_dict["box"] = torch.tensor( - [[1, 0, 0, 0, 1, 0, 0, 0, 1]] * nframes, - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=env.DEVICE, - ) - input_dict["atype"] = torch.tensor( - [[0] * natoms, [0] * natoms], - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=env.DEVICE, - ) - input_dict["coord"] = torch.tensor( - [[[0] * 3] * natoms] * nframes, - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=env.DEVICE, - ) - _, loss, more_loss = module( - input_dict, - fake_model, - labels, - natoms, - 1.0, - ) - else: - _, loss, more_loss = module( - {}, - fake_model, - labels, - natoms, - 1.0, - )
58-104: Add documentation for the conditional behavior related to "updated_coord".The code introduces conditional handling when
"updated_coord"is inself.key_to_pref_map, but there's no documentation explaining this behavior. Consider adding a comment explaining why this special handling is needed and what it accomplishes.+ # When testing denoise models, we need to provide additional input tensors + # for coordinates, box, and atom types to simulate the structure data if "updated_coord" in self.key_to_pref_map: import torchdeepmd/pt/loss/denoise.py (2)
303-329: Implement serialize/deserialize methods.The
DenoiseLossclass would benefit from havingserializeanddeserializemethods to support model saving and loading, similar to other loss classes.Consider adding these methods to ensure proper model serialization:
def serialize(self): return { "mask_token": self.mask_token, "mask_coord": self.mask_coord, "mask_cell": self.mask_cell, "token_loss": self.token_loss, "coord_loss": self.coord_loss, "cell_loss": self.cell_loss, "noise_type": self.noise_type, "coord_noise": self.coord_noise, "cell_pert_fraction": self.cell_pert_fraction, "noise_mode": self.noise_mode, "mask_num": self.mask_num, "mask_prob": self.mask_prob, "same_mask": self.same_mask, "loss_func": self.loss_func, } @classmethod def deserialize(cls, data): return cls(**data)
112-249: Consider refactoring the forward method for clarity.The forward method is quite long and handles multiple responsibilities (cell perturbation, coordinate masking, token masking, and loss calculation). Consider refactoring it into smaller, focused methods.
You could break it down into methods like:
_apply_cell_perturbation(self, input_dict, label)_apply_coordinate_noise(self, input_dict, label)_apply_token_masking(self, input_dict, label)_calculate_loss(self, model_pred, label)This would make the code more maintainable and easier to test.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/loss/denoise.py(2 hunks)source/tests/universal/common/cases/loss/utils.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- 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: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test C++ (true)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (2)
deepmd/pt/loss/denoise.py (2)
257-260: Verify model output contains required fields.This code assumes
model_predcontainsupdated_coordandstrain_componentskeys, which might not be guaranteed by the model.Ensure that the model implementation returns both these fields in the forward pass:
#!/bin/bash # Search for DenoiseModel implementations to verify output fields ast-grep --pattern 'class DenoiseModel { $$$ forward($$$) { $$$ return $$$ } $$$ }' # Search for references to updated_coord and strain_components in model outputs rg -A 3 "def forward" --include="*.py" | grep -A 10 "return.*dict"
225-225: Verify coordinate loss calculation approach.The coordinate loss calculation involves transforming fractional coordinates and using matrix operations. Ensure this approach matches the physical meaning of the updated coordinates in your denoising model.
Could you verify this approach with experiments on small test cases? Alternatively, consider providing a brief comment explaining the physical meaning of these transformations and why they're appropriate for coordinate loss calculation.
| nframes, | ||
| natoms, | ||
| ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) | ||
| results = self.deep_eval.eval( |
Check notice
Code scanning / CodeQL
Unused local variable Note
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
deepmd/infer/deep_denoise.py (2)
122-122: Remove unused variable
The local variableresultsis assigned but not referenced, triggering a lint error (F841).- results = self.deep_eval.eval( + self.deep_eval.eval(🧰 Tools
🪛 Ruff (0.8.2)
122-122: Local variable
resultsis assigned to but never usedRemove assignment to unused variable
results(F841)
132-133: Unresolved TODO
A TODO comment indicates this method is partially implemented and no meaningful output is returned.Would you like me to propose a basic structure for returning the denoised outputs?
deepmd/entrypoints/test.py (1)
180-181: Alternative approach for unsupported test
RaisingNotImplementedErrorcan disrupt automated testing pipelines. Consider issuing a warning and skipping if it’s a planned future feature.- raise NotImplementedError("DeepDenoise is not supported in test mode.") + import warnings + warnings.warn("DeepDenoise is not yet supported in test mode. Skipping...") + continue
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/entrypoints/test.py(2 hunks)deepmd/infer/deep_denoise.py(1 hunks)deepmd/pt/infer/deep_eval.py(2 hunks)
🧰 Additional context used
🧬 Code Definitions (2)
deepmd/pt/infer/deep_eval.py (4)
deepmd/infer/deep_denoise.py (1)
DeepDenoise(21-133)deepmd/dpmodel/model/make_model.py (1)
model_output_type(180-189)deepmd/pt/model/model/make_model.py (1)
model_output_type(88-99)deepmd/pd/model/model/make_model.py (1)
model_output_type(87-98)
deepmd/entrypoints/test.py (1)
deepmd/infer/deep_denoise.py (1)
DeepDenoise(21-133)
🪛 Ruff (0.8.2)
deepmd/infer/deep_denoise.py
122-122: Local variable results is assigned to but never used
Remove assignment to unused variable results
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (25)
- 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: 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: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- 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)
🔇 Additional comments (5)
deepmd/infer/deep_denoise.py (2)
1-9: Clean import statements
No issues found; everything looks good and is well-organized.
21-38: Class docstring clarity
The docstring is well-structured and describes the purpose of this class thoroughly. Good job clarifying how noise is introduced to structures.deepmd/entrypoints/test.py (1)
18-20: New import verified
The added import forDeepDenoisematches its usage below. No concerns here.deepmd/pt/infer/deep_eval.py (2)
20-22: New import for DeepDenoise
No issues found, and it’s consistent with the newly introduced denoise functionality.
217-218: Model type extension
Mapping"updated_coord"toDeepDenoiseis coherent with your new denoising output. Nicely integrated.
| def eval( | ||
| self, | ||
| coords: np.ndarray, | ||
| cells: Optional[np.ndarray], | ||
| atom_types: Union[list[int], np.ndarray], | ||
| atomic: bool = False, | ||
| fparam: Optional[np.ndarray] = None, | ||
| aparam: Optional[np.ndarray] = None, | ||
| mixed_type: bool = False, | ||
| **kwargs: dict[str, Any], | ||
| ) -> tuple[np.ndarray, ...]: | ||
| """Evaluate properties. If atomic is True, also return atomic property. | ||
|
|
||
| Parameters | ||
| ---------- | ||
| coords : np.ndarray | ||
| The coordinates of the atoms, in shape (nframes, natoms, 3). | ||
| cells : np.ndarray | ||
| The cell vectors of the system, in shape (nframes, 9). If the system | ||
| is not periodic, set it to None. | ||
| atom_types : list[int] or np.ndarray | ||
| The types of the atoms. If mixed_type is False, the shape is (natoms,); | ||
| otherwise, the shape is (nframes, natoms). | ||
| atomic : bool, optional | ||
| Whether to return atomic property, by default False. | ||
| fparam : np.ndarray, optional | ||
| The frame parameters, by default None. | ||
| aparam : np.ndarray, optional | ||
| The atomic parameters, by default None. | ||
| mixed_type : bool, optional | ||
| Whether the atom_types is mixed type, by default False. | ||
| **kwargs : dict[str, Any] | ||
| Keyword arguments. | ||
|
|
There was a problem hiding this comment.
Incomplete return statement
The docstring indicates a tuple of properties should be returned, but the current implementation returns None. This discrepancy may cause confusion or errors downstream.
Training process:
Add token, coordinate and cell noise to the equilibrium structure. Then we use model to predict these noises.
Application:
Summary by CodeRabbit
New Features
DPDenoiseAtomicModel) and a dedicated denoising model (DenoiseModel) for improved predictive performance.input.json) for denoising model training with detailed parameters.type.rawandtype_map.raw) in multiple directories.DenoiseLoss) for enhanced denoising functionality.Enhancements
Utilities
fitting_denoiseandloss_denoise) to streamline the setup of denoising tasks.Bug Fixes
Tests
DenoiseModelandDenoiseFittingNetfunctionalities.