feat(pt): add atomic_weight to dipole model inference #4628
Conversation
|
Caution Review failedThe pull request is closed. 📝 WalkthroughWalkthroughThis pull request introduces an optional Changes
Sequence Diagram(s)sequenceDiagram
participant Caller
participant ModelLayer
participant AtomicModel
Caller->>ModelLayer: forward(..., atomic_weight)
ModelLayer->>ModelLayer: Process inputs and pass atomic_weight to common routine
ModelLayer->>AtomicModel: forward_common_atomic(..., atomic_weight)
AtomicModel-->>ModelLayer: Return atomic predictions
ModelLayer-->>Caller: Return final results (atomic weight applied)
Possibly related PRs
Suggested reviewers
Tip ⚡🧪 Multi-step agentic review comment chat (experimental)
📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
deepmd/pt/model/model/make_model.py (1)
616-634:⚠️ Potential issueUpdate the forward method to include atomic_weight parameter.
The
forwardmethod at the bottom of the class needs to be updated to include the newatomic_weightparameter to ensure consistency with theforward_commonmethod signature.Apply this change to update the method signature and parameter passing:
def forward( self, coord, atype, box: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, + atomic_weight: Optional[torch.Tensor] = None, ) -> dict[str, torch.Tensor]: # directly call the forward_common method when no specific transform rule return self.forward_common( coord, atype, box, fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + atomic_weight=atomic_weight, )
🧹 Nitpick comments (1)
deepmd/pt/model/model/make_model.py (1)
299-301: Clarify the atomic weight initialization and usage.The code uses
next(iter(self.atomic_output_def().var_defs.keys()))to get the first key from the atomic output definitions. Consider adding a comment explaining why this key is chosen and ensuring it's the correct key for applying the weight.Add a clarifying comment:
# add weight to atomic_output kw = next(iter(self.atomic_output_def().var_defs.keys())) +# Create default weight tensor using the shape of the first output variable atomic_weight = torch.ones_like(atomic_ret[kw])
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/model/model/dipole_model.py(4 hunks)deepmd/pt/model/model/make_model.py(4 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (20)
- 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++ (clang, clang)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (2, 3.9)
- 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: Analyze (javascript-typescript)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cpu, cpu)
- 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 (7)
deepmd/pt/model/model/dipole_model.py (4)
63-63: Function signature enhancement: New atomic weight parameter added.This change appropriately adds the optional
atomic_weightparameter to the forward method, which allows users to define weights for atomic dipoles.
72-72: Correctly passing atomic_weight to the common implementation.The parameter is properly forwarded to the underlying implementation.
103-103: Function signature enhancement: New atomic weight parameter added to lower interface.The atomic_weight parameter is appropriately added to the forward_lower method, maintaining consistency with the higher-level interface.
114-114: Correctly passing atomic_weight to the lower-level implementation.The parameter is properly forwarded to the underlying implementation.
deepmd/pt/model/model/make_model.py (3)
138-138: Function signature enhancement: New atomic weight parameter added.This change appropriately adds the optional
atomic_weightparameter to the forward_common method, allowing the dipole model to incorporate atomic weights in its predictions.
192-192: Correctly passing atomic_weight to the lower-level implementation.The parameter is properly forwarded to the forward_common_lower method.
247-247: Function signature enhancement: New atomic weight parameter added to lower interface.The atomic_weight parameter is appropriately added to the forward_common_lower method.
deepmd/pt/model/model/make_model.py
Outdated
| # add weight to atomic_output | ||
| kw = next(iter(self.atomic_output_def().var_defs.keys())) | ||
| atomic_weight = torch.ones_like(atomic_ret[kw]) | ||
| if atomic_weight is not None: | ||
| # atomic_weight: nf x nloc x dim | ||
| atomic_ret[kw] = atomic_ret[kw] * atomic_weight.reshape( | ||
| *atomic_ret[kw].shape[:-1], -1 | ||
| ) |
There was a problem hiding this comment.
Fix parameter shadowing issue in the atomic weight implementation.
There's a critical bug in the implementation. The function parameter atomic_weight is being shadowed by a local variable with the same name on line 301. This will prevent the externally provided weights from being used.
Apply this fix to correct the parameter shadowing:
# add weight to atomic_output
kw = next(iter(self.atomic_output_def().var_defs.keys()))
-atomic_weight = torch.ones_like(atomic_ret[kw])
+weight_tensor = torch.ones_like(atomic_ret[kw])
if atomic_weight is not None:
# atomic_weight: nf x nloc x dim
- atomic_ret[kw] = atomic_ret[kw] * atomic_weight.reshape(
+ atomic_ret[kw] = atomic_ret[kw] * atomic_weight.reshape(
*atomic_ret[kw].shape[:-1], -1
)
+else:
+ atomic_ret[kw] = atomic_ret[kw] * weight_tensorOr, for a cleaner implementation:
# add weight to atomic_output
kw = next(iter(self.atomic_output_def().var_defs.keys()))
-atomic_weight = torch.ones_like(atomic_ret[kw])
if atomic_weight is not None:
# atomic_weight: nf x nloc x dim
atomic_ret[kw] = atomic_ret[kw] * atomic_weight.reshape(
*atomic_ret[kw].shape[:-1], -1
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # add weight to atomic_output | |
| kw = next(iter(self.atomic_output_def().var_defs.keys())) | |
| atomic_weight = torch.ones_like(atomic_ret[kw]) | |
| if atomic_weight is not None: | |
| # atomic_weight: nf x nloc x dim | |
| atomic_ret[kw] = atomic_ret[kw] * atomic_weight.reshape( | |
| *atomic_ret[kw].shape[:-1], -1 | |
| ) | |
| # add weight to atomic_output | |
| kw = next(iter(self.atomic_output_def().var_defs.keys())) | |
| weight_tensor = torch.ones_like(atomic_ret[kw]) | |
| if atomic_weight is not None: | |
| # atomic_weight: nf x nloc x dim | |
| atomic_ret[kw] = atomic_ret[kw] * atomic_weight.reshape( | |
| *atomic_ret[kw].shape[:-1], -1 | |
| ) | |
| else: | |
| atomic_ret[kw] = atomic_ret[kw] * weight_tensor |
ChiahsinChu
force-pushed
the
devel-dipole_with_atomic_weight
branch
from
64fbde6 to
50a8a21
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
deepmd/pt/model/model/make_model.py (1)
617-634:⚠️ Potential issueUpdate forward method to pass the atomic_weight parameter.
The
forwardmethod needs to be updated to accept and pass theatomic_weightparameter toforward_common, otherwise this functionality won't be available when using the baseforwardmethod.def forward( self, coord, atype, box: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, + atomic_weight: Optional[torch.Tensor] = None, ) -> dict[str, torch.Tensor]: # directly call the forward_common method when no specific transform rule return self.forward_common( coord, atype, box, fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + atomic_weight=atomic_weight, )
♻️ Duplicate comments (1)
deepmd/pt/model/model/make_model.py (1)
300-307:⚠️ Potential issueFix parameter shadowing issue as identified in previous review.
The current implementation introduces a potential parameter shadowing issue if the atomic model fitting net's var_name is the same as "atomic_weight". This could lead to unexpected behavior, similar to the issue identified in a previous review comment.
It's best to use a different variable name for clarity:
if hasattr(self.atomic_model, "fitting_net"): if hasattr(self.atomic_model.fitting_net, "var_name"): kw = self.atomic_model.fitting_net.var_name if atomic_weight is not None: # atomic_weight: nf x nloc x dim - atomic_ret[kw] = atomic_ret[kw] * atomic_weight.view( + atomic_ret[kw] = atomic_ret[kw] * atomic_weight.view( [atomic_ret[kw].shape[0], atomic_ret[kw].shape[1], -1] )The variable name doesn't need to be changed in this case as it doesn't actually shadow the parameter (unlike the previous issue), but the concern about having the same name is valid.
🧰 Tools
🪛 Ruff (0.8.2)
300-301: Use a single
ifstatement instead of nestedifstatements(SIM102)
🧹 Nitpick comments (2)
deepmd/pt/model/model/make_model.py (2)
138-138: Add the new parameter to docstring.The parameter
atomic_weightis added to the method signature but not documented in the method docstring. For better code maintenance and usability, update the docstring to include a description of this parameter.@@ -152,6 +152,8 @@ frame parameter. nf x ndf aparam atomic parameter. nf x nloc x nda + atomic_weight + Optional weight for atomic output. nf x nloc x dim do_atomic_virial If calculate the atomic virial.
247-247: Add the new parameter to docstring.The parameter
atomic_weightis added to the method signature but not documented in the method docstring. Update the docstring to include a description of this parameter.@@ -271,6 +271,8 @@ The data needed for communication for parallel inference. extra_nlist_sort whether to forcibly sort the nlist. + atomic_weight + Optional weight for atomic output. nf x nloc x dim
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/model/model/make_model.py(4 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
deepmd/pt/model/model/make_model.py
300-301: Use a single if statement instead of nested if statements
(SIM102)
⏰ Context from checks skipped due to timeout of 90000ms (14)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Analyze (javascript-typescript)
- 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 C++ (false)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (true)
🔇 Additional comments (1)
deepmd/pt/model/model/make_model.py (1)
192-192: Appropriate parameter passing.The
atomic_weightparameter is correctly passed down to the lower-level implementation, ensuring consistent behavior across the API layers.
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## devel #4628 +/- ##
=======================================
Coverage 84.77% 84.78%
=======================================
Files 688 688
Lines 66097 66102 +5
Branches 3539 3539
=======================================
+ Hits 56036 56042 +6
Misses 8919 8919
+ Partials 1142 1141 -1 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
I do not get what do you mean by " If the charges associated with the winner centroids are not identical, there is no way for users to get the derivative of (real) total dipole w.r.t. coord from DW models. (I tried to calculate the gradient of output dipole w.r.t. the input coord, but the gradient chain seems to be cut in the model.)"
Does the type exclusion mechanism solves your issue? you may check
No. It is not what I mean. The An example for my case: |
ChiahsinChu
force-pushed
the
devel-dipole_with_atomic_weight
branch
from
d48e45f to
a47b1e0
Compare
deepmd/pt/model/model/make_model.py
Outdated
| if hasattr(self.atomic_model, "fitting_net"): | ||
| if hasattr(self.atomic_model.fitting_net, "var_name"): | ||
| kw = self.atomic_model.fitting_net.var_name |
There was a problem hiding this comment.
I feel that self.atomic_model.fitting_output_def could be used here
There was a problem hiding this comment.
It's a bit confusing since it breaks the consistency of APIs between different models and backends.
There was a problem hiding this comment.
Would it be better to add this keyword for all models and backends?
There was a problem hiding this comment.
I would suggest incorporating the atomic weight into the base atomic model, following a similar approach to how atom masks are managed. This would ensure that atomic weight is consistently supported across all atomic modes.
There was a problem hiding this comment.
Caution
Inline review comments failed to post. This is likely due to GitHub's limits when posting large numbers of comments. If you are seeing this consistently it is likely a permissions issue. Please check "Moderation" -> "Code review limits" under your organization settings.
Actionable comments posted: 5
🔭 Outside diff range comments (1)
deepmd/pd/model/descriptor/se_a.py (1)
745-761:⚠️ Potential issueInconsistent code structure detected
There appears to be redundant or conflicting code blocks in this section. The same conditional logic appears twice with different indentation (lines 748-755 and 754-761), which could lead to unpredictable behavior.
This appears to be a merge conflict or code duplication that wasn't properly resolved. Please consolidate these blocks into a single coherent code path:
- "Compressed environment is not implemented yet." - ) - else: - # NOTE: control flow with double backward is not supported well yet by paddle.jit - if not paddle.framework.in_dynamic_mode() or decomp.numel(rr) > 0: - rr = rr * mm.unsqueeze(2).astype(rr.dtype) - ss = rr[:, :, :1] - if self.compress: - else: - # NOTE: control flow with double backward is not supported well yet by paddle.jit - if not paddle.in_dynamic_mode() or decomp.numel(rr) > 0: - rr = rr * mm.unsqueeze(2).astype(rr.dtype) - ss = rr[:, :, :1] - if self.compress: - raise NotImplementedError( - "Compressed environment is not implemented yet." - ) - else: + "Compressed environment is not implemented yet." + ) + else: + # NOTE: control flow with double backward is not supported well yet by paddle.jit + if not paddle.in_dynamic_mode() or decomp.numel(rr) > 0: + rr = rr * mm.unsqueeze(2).astype(rr.dtype) + ss = rr[:, :, :1] + if self.compress: + raise NotImplementedError( + "Compressed environment is not implemented yet." + ) + else:
🧹 Nitpick comments (22)
deepmd/utils/data.py (1)
92-96: Added validation to ensure comprehensive type map coverageThis addition improves input validation by checking if all elements in
self.type_mapare present in the providedtype_map. It raises a descriptive error that helps users identify which elements are missing from their type map, preventing silent errors or unexpected behavior later in the execution pipeline.Consider adding a unit test to verify this validation works as expected with various input combinations.
source/tests/array_api_strict/descriptor/repflows.py (2)
22-38: Consider handling None values before serialization/deserialization.In this logic, some attributes (e.g.,
layers,edge_embd,angle_embd) are immediately serialized/deserialized. Ifvalueis unexpectedlyNone, it could lead to runtime errors. A small type check or safety guard can improve robustness.def __setattr__(self, name: str, value: Any) -> None: if name in {"mean", "stddev"}: value = to_array_api_strict_array(value) + elif name in {"layers"} and value is not None: + value = [RepFlowLayer.deserialize(layer.serialize()) for layer in value] + elif name in {"edge_embd", "angle_embd"} and value is not None: + value = NativeLayer.deserialize(value.serialize()) elif name in {"env_mat_edge", "env_mat_angle"}: pass elif name == "emask": value = PairExcludeMask(value.ntypes, value.exclude_types) return super().__setattr__(name, value)
41-61: Check for None before callingvalue.serialize().In the current implementation, if any of the named attributes are set to
None, callingvalue.serialize()will fail. Consider adding a guard to avoid potential errors.def __setattr__(self, name: str, value: Any) -> None: if name in { "node_self_mlp", "node_sym_linear", "node_edge_linear", "edge_self_linear", "a_compress_n_linear", "a_compress_e_linear", "edge_angle_linear1", "edge_angle_linear2", "angle_self_linear", }: - if value is not None: + if value is not None and hasattr(value, "serialize"): value = NativeLayer.deserialize(value.serialize()) elif name in {"n_residual", "e_residual", "a_residual"}: value = [to_array_api_strict_array(vv) for vv in value] return super().__setattr__(name, value)deepmd/jax/descriptor/repflows.py (2)
25-42: Add a guard for potential None values.Similar to the array_api_strict variant, deserializing attributes without checking for None can cause errors. Adding a simple check can increase code resilience.
def __setattr__(self, name: str, value: Any) -> None: if name in {"mean", "stddev"}: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif name in {"layers"} and value is not None: + value = [RepFlowLayer.deserialize(layer.serialize()) for layer in value] + elif name in {"edge_embd", "angle_embd"} and value is not None: + value = NativeLayer.deserialize(value.serialize()) elif name in {"env_mat_edge", "env_mat_angle"}: pass elif name == "emask": value = PairExcludeMask(value.ntypes, value.exclude_types) return super().__setattr__(name, value)
46-66: Check for None or incorrect types beforevalue.serialize().When deserializing network layers, confirm that they are valid, non-None objects with a
serialize()method. This helps avoid unhandled runtime exceptions.if name in { "node_self_mlp", "node_sym_linear", "node_edge_linear", "edge_self_linear", "a_compress_n_linear", "a_compress_e_linear", "edge_angle_linear1", "edge_angle_linear2", "angle_self_linear", }: - if value is not None: + if value is not None and hasattr(value, "serialize"): value = NativeLayer.deserialize(value.serialize()) elif name in {"n_residual", "e_residual", "a_residual"}: value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value] return super().__setattr__(name, value)source/tests/pt/model/test_dpa3.py (1)
209-209: Remove or utilize the unused variablemodel.Local variable
modelis assigned but never used. This can be safely removed unless intended for later use.- model = torch.jit.script(dd0)🧰 Tools
🪛 Ruff (0.8.2)
209-209: Local variable
modelis assigned to but never usedRemove assignment to unused variable
model(F841)
deepmd/dpmodel/descriptor/dpa3.py (3)
212-253: Constructor largely follows good practices.The overall constructor is consistent, properly leveraging helper functions for initialization. However, please review the default argument for
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)
253-253: Replace mutable default argument.Defaulting
exclude_typesto a mutable list can introduce unexpected behaviors if the list is modified in-place. Consider usingNoneinstead:- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = 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)
574-574: Unused local variable.The variable
env_matis assigned but never used. Removing it clarifies the code and avoids confusion:- env_mat = repflow_variable.pop("env_mat")🧰 Tools
🪛 Ruff (0.8.2)
574-574: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
source/tests/universal/dpmodel/descriptor/test_descriptor.py (1)
467-527: New DescriptorParamDPA3 function appears correct.However, please address the mutable default argument for
exclude_typesat line 474 to avoid potential unintended side effects.- exclude_types=[], + exclude_types: Optional[list] = None, ... + if exclude_types is None: + exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
474-474: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/dpmodel/descriptor/repflows.py (3)
440-440: Rename unused loop variableidxto_idxfor clarityThe variable
idxisn’t referenced in the loop body. Renaming it to_idxclarifies that it’s intentionally unused.-for idx, ll in enumerate(self.layers): +for _idx, ll in enumerate(self.layers): # node_ebd: nb x nloc x n_dim # node_ebd_ext: nb x nall x n_dim🧰 Tools
🪛 Ruff (0.8.2)
440-440: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
937-937: Remove unused variablenallThe variable
nallis assigned but never used, creating dead code. Removing it tidies up the logic.- nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :] assert (nb, nloc) == node_ebd.shape[:2]🧰 Tools
🪛 Ruff (0.8.2)
937-937: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
1199-1199: Remove unused variablenitemThe
nitemvariable is never used. Eliminating it prevents confusion and aligns with clean coding practices.- nitem = len(update_list) uu = update_list[0] if update_name == "node": for ii, vv in enumerate(self.n_residual):🧰 Tools
🪛 Ruff (0.8.2)
1199-1199: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
deepmd/pt/model/descriptor/repflows.py (2)
461-461: Rename unused variable to underscore.The loop index
idxis never referenced. Use_to indicate it’s unused:- for idx, ll in enumerate(self.layers): + for _, ll in enumerate(self.layers):🧰 Tools
🪛 Ruff (0.8.2)
461-461: Loop control variable
idxnot used within loop bodyRename unused
idxto_idx(B007)
571-575: Use a ternary operator for concise code.You can simplify:
- if callable(merged): - sampled = merged() - else: - sampled = merged + sampled = merged() if callable(merged) else merged🧰 Tools
🪛 Ruff (0.8.2)
571-575: Use ternary operator
sampled = merged() if callable(merged) else mergedinstead ofif-else-block(SIM108)
deepmd/pt/model/descriptor/repflow_layer.py (3)
306-307: Remove unused local variable
Line 307 declarese_dimbut it is never referenced afterward. Consider removing it to streamline the code and avoid potential confusion.nb, nloc, nnei, _ = edge_ebd.shape - e_dim = edge_ebd.shape[-1]🧰 Tools
🪛 Ruff (0.8.2)
307-307: Local variable
e_dimis assigned to but never usedRemove assignment to unused variable
e_dim(F841)
534-534: Remove unused local variable
Line 534 introducesnallbut it is not used anywhere in the code. Cleaning it up will improve clarity.- nall = node_ebd_ext.shape[1]🧰 Tools
🪛 Ruff (0.8.2)
534-534: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
788-788: Remove unused local variable
Inlist_update_res_residual, the variablenitemis assigned but never utilized. Consider deleting it.def list_update_res_residual( self, update_list: list[torch.Tensor], update_name: str = "node" ) -> torch.Tensor: - nitem = len(update_list) uu = update_list[0]🧰 Tools
🪛 Ruff (0.8.2)
788-788: Local variable
nitemis assigned to but never usedRemove assignment to unused variable
nitem(F841)
source/tests/consistent/descriptor/test_dpa3.py (1)
39-44: Simplify the conditional logic
The currentif/elsestatements always assignDescrptDPA3PD = None. Using a ternary expression or removing the condition entirely can make the code simpler.-if INSTALLED_PD: - # not supported yet - DescrptDPA3PD = None -else: - DescrptDPA3PD = None +DescrptDPA3PD = None🧰 Tools
🪛 Ruff (0.8.2)
39-43: Use ternary operator
DescrptDPA3PD = None if INSTALLED_PD else Noneinstead ofif-else-block(SIM108)
deepmd/pt/model/descriptor/dpa3.py (3)
339-339: Rename or remove unused loop variable
The loop variableiiis not utilized within the loop body. This can be misleading. Rename it to_to clarify that it is unused, or remove the enumeration if not needed.- for ii, descrpt in enumerate(descrpt_list): + for _, descrpt in enumerate(descrpt_list): descrpt.compute_input_stats(merged, path)🧰 Tools
🪛 Ruff (0.8.2)
339-339: Loop control variable
iinot used within loop bodyRename unused
iito_ii(B007)
414-415: Remove assignment to unused variable
env_matis assigned here but never referenced. Remove it to reduce unnecessary code.env_mat = repflow_variable.pop("env_mat") -env_mat🧰 Tools
🪛 Ruff (0.8.2)
415-415: Local variable
env_matis assigned to but never usedRemove assignment to unused variable
env_mat(F841)
471-471: Remove unused variable
nallis assigned on line 471 without subsequent usage. Deleting it will keep the code concise.- nall = extended_coord.view(nframes, -1).shape[1] // 3🧰 Tools
🪛 Ruff (0.8.2)
471-471: Local variable
nallis assigned to but never usedRemove assignment to unused variable
nall(F841)
🛑 Comments failed to post (5)
deepmd/dpmodel/descriptor/repflows.py (2)
156-156: 🛠️ Refactor suggestion
Avoid using a mutable default value for
exclude_typesUsing a list as a default argument may lead to unexpected side effects because it is shared across all calls. Consider using
Noneand initializing inside the constructor.- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None,Then, after the constructor starts:
+ if exclude_types is None: + exclude_types = [] self.e_rcut = float(e_rcut) self.e_rcut_smth = float(e_rcut_smth) ...📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.exclude_types: Optional[list[tuple[int, int]]] = None, other_param: int = 0, # other parameters... # Inside the constructor after the parameters are received: if exclude_types is None: exclude_types = [] self.e_rcut = float(e_rcut) self.e_rcut_smth = float(e_rcut_smth) # ... rest of the constructor code🧰 Tools
🪛 Ruff (0.8.2)
156-156: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
353-353: 🛠️ Refactor suggestion
Avoid using a mutable default value for
exclude_typesinreinit_excludeSame reasoning applies: a default list can cause unexpected behavior. Use
Noneand initialize inside the method.-def reinit_exclude( - self, - exclude_types: list[tuple[int, int]] = [], -) -> None: +def reinit_exclude( + self, + exclude_types: Optional[list[tuple[int, int]]] = None, +) -> None: + if exclude_types is None: + exclude_types = [] self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.def reinit_exclude( self, exclude_types: Optional[list[tuple[int, int]]] = None, ) -> None: if exclude_types is None: exclude_types = [] self.exclude_types = exclude_types self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types)🧰 Tools
🪛 Ruff (0.8.2)
353-353: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/pt/model/descriptor/repflows.py (2)
360-360: 🛠️ Refactor suggestion
Avoid mutable default argument in
reinit_exclude.Applying the same reasoning here:
- def reinit_exclude( - self, - exclude_types: list[tuple[int, int]] = [], - ) -> None: + def reinit_exclude( + self, + exclude_types: Optional[list[tuple[int, int]]] = None, + ) -> None: + if exclude_types is None: + exclude_types = []📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.def reinit_exclude( self, exclude_types: Optional[list[tuple[int, int]]] = None, ) -> None: if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
360-360: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
178-178: 🛠️ Refactor suggestion
Use
Noneinstead of a mutable default argument.Using a mutable list (
[]) as a default argument forexclude_typescan lead to unexpected shared state across instances. Consider this refactor:- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, ... ): + if exclude_types is None: + exclude_types = []📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.exclude_types: Optional[list[tuple[int, int]]] = None, ... ): if exclude_types is None: exclude_types = []🧰 Tools
🪛 Ruff (0.8.2)
178-178: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
deepmd/pt/model/descriptor/dpa3.py (1)
105-105: 🛠️ Refactor suggestion
Avoid mutable default argument
Definingexclude_types: list[tuple[int, int]] = []can lead to unintended shared state across calls. Consider usingNoneas a default and initializing inside the constructor.- exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, ... def __init__( self, ... exclude_types: Optional[list[tuple[int, int]]] = None, ... ): if exclude_types is None: exclude_types = [] self.exclude_types = exclude_types📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.from typing import Optional # Add this if not already imported class YourClassName: def __init__( self, ..., - exclude_types: list[tuple[int, int]] = [], + exclude_types: Optional[list[tuple[int, int]]] = None, ... ): - self.exclude_types = exclude_types + if exclude_types is None: + exclude_types = [] + self.exclude_types = exclude_types🧰 Tools
🪛 Ruff (0.8.2)
105-105: Do not use mutable data structures for argument defaults
Replace with
None; initialize within function(B006)
ChiahsinChu
force-pushed
the
devel-dipole_with_atomic_weight
branch
from
ab0f1f5 to
d14fa86
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
152-224: Missing documentation for the atomic_weight parameter.The docstring for the forward_common_atomic method hasn't been updated to include information about the new atomic_weight parameter and its usage.
Add documentation for the atomic_weight parameter in the method's docstring:
def forward_common_atomic( self, extended_coord: np.ndarray, extended_atype: np.ndarray, nlist: np.ndarray, mapping: Optional[np.ndarray] = None, fparam: Optional[np.ndarray] = None, aparam: Optional[np.ndarray] = None, atomic_weight: Optional[np.ndarray] = None, ) -> dict[str, np.ndarray]: """Common interface for atomic inference. This method accept extended coordinates, extended atom typs, neighbor list, and predict the atomic contribution of the fit property. Parameters ---------- extended_coord extended coordinates, shape: nf x (nall x 3) extended_atype extended atom typs, shape: nf x nall for a type < 0 indicating the atomic is virtual. nlist neighbor list, shape: nf x nloc x nsel mapping extended to local index mapping, shape: nf x nall fparam frame parameters, shape: nf x dim_fparam aparam atomic parameter, shape: nf x nloc x dim_aparam + atomic_weight + atomic weights for scaling outputs, shape: nf x nloc + if provided, all output values will be multiplied by this weight Returns ------- ret_dict dict of output atomic properties. should implement the definition of `fitting_output_def`. ret_dict["mask"] of shape nf x nloc will be provided. ret_dict["mask"][ff,ii] == 1 indicating the ii-th atom of the ff-th frame is real. ret_dict["mask"][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual. """🧰 Tools
🪛 Ruff (0.8.2)
212-212: 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 (4)
deepmd/jax/model/base_model.py(1 hunks)deepmd/dpmodel/atomic_model/base_atomic_model.py(5 hunks)deepmd/dpmodel/model/make_model.py(12 hunks)deepmd/jax/model/base_model.py(8 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/jax/model/base_model.py
⏰ Context from checks skipped due to timeout of 90000ms (5)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (20)
deepmd/jax/model/base_model.py (5)
30-30: Added atomic_weight parameter for enhanced dipole model.The addition of the optional
atomic_weightparameter aligns with the PR's goal to improve dipole model calculations by allowing atomic weight specification.
39-39: Properly passed atomic_weight to the atomic model.The parameter is correctly forwarded to the atomic model's forward_common_atomic method.
61-62: Updated eval_output signature and handling of atomic_weight.The function now properly accepts and processes the atomic_weight parameter, with appropriate None-checking when passing it to the atomic model.
Also applies to: 73-75
111-112: Updated eval_ce signature and handling of atomic_weight.Similar to eval_output, this function has been properly updated to handle the atomic_weight parameter with appropriate None-checking.
Also applies to: 124-126
88-89: Correctly propagated atomic_weight to derivative calculations.The parameter is correctly passed to the JAX autodiff functions, ensuring that derivatives properly account for atomic weights.
Also applies to: 144-145
deepmd/dpmodel/model/make_model.py (10)
66-66: Added atomic_weight parameter to model_call_from_call_lower function.The optional atomic_weight parameter has been added to support weighted dipole calculations.
125-126: Properly propagated atomic_weight to lower-level functions.The parameter is correctly passed to call_lower, ensuring proper propagation through the call chain.
229-229: Added atomic_weight parameter to the call method.This addition maintains consistency with the other interface methods and supports the enhanced dipole model functionality.
256-262: Updated input_type_cast to handle atomic_weight parameter.The function now properly unpacks and includes atomic_weight in its parameter list and return values.
290-290: Added atomic_weight parameter to call_lower method.This change ensures that the atomic_weight parameter is supported throughout the model's interface.
328-333: Updated input_type_cast handling in call_lower method.The parameter is now correctly unpacked from the input_type_cast return value.
357-357: Added atomic_weight parameter to forward_common_atomic method.This change provides consistent support for atomic_weight throughout the model interface.
383-389: Updated input_type_cast function signature and return type.The function's signature and return type have been properly updated to include the atomic_weight parameter.
399-403: Added atomic_weight handling in input_type_cast conversion logic.The parameter is correctly included in the type conversion list and unpacked appropriately.
411-411: Added type casting for atomic_weight in precision conversion.The atomic_weight parameter is properly type-cast when converting precision, maintaining consistency with other parameters.
deepmd/dpmodel/atomic_model/base_atomic_model.py (5)
201-202: Properly passed atomic_weight to forward_atomic method.The parameter is correctly forwarded to the forward_atomic method.
218-223: Implemented atomic_weight scaling for output calculation.This implementation correctly scales the atomic output by the provided weights when atomic_weight is not None. The reshaping ensures proper broadcasting regardless of output shape.
235-235: Added atomic_weight parameter to call method.The parameter has been added to maintain consistency across the API.
244-245: Propagated atomic_weight to forward_common_atomic in call method.The parameter is correctly passed through the call method to forward_common_atomic.
152-152:❓ Verification inconclusive
Added atomic_weight parameter to forward_common_atomic method.
The optional atomic_weight parameter has been added to support weighted dipole calculations.
🌐 Web query:
What are atomic weights in molecular dipole modeling?💡 Result:
Atomic weights, which represent the average mass of an element's isotopes, are not directly used in molecular dipole modeling. Instead, dipole modeling focuses on atomic charges, polarizabilities, and molecular geometry to predict how electron density shifts within a molecule. Here’s how atomic-level properties contribute to dipole moment calculations:
Key Factors in Dipole Modeling:
Atomic Partial Charges:
Dipole moments arise from uneven charge distribution. Models like MuML (Machine learning for μ) decompose molecular dipoles into atomic contributions, combining local polarization (atomic dipoles) and long-range charge transfer (partial charges)[1][2][15]. These charges are derived from quantum mechanical calculations or machine learning predictions, not atomic weights[10][12].Atomic Polarizabilities:
Polarizability—how easily an atom’s electron cloud distorts under an electric field—determines induced dipoles. Thole-style models assign polarizabilities to atoms based on their chemical environment, enabling predictions of molecular polarization[7][16]. For example, oxygen’s lone pairs significantly influence polarizability in molecules like water[7][13].Molecular Geometry:
The spatial arrangement of atoms dictates whether bond dipoles cancel or reinforce each other. Linear molecules (e.g., CO₂) may have zero net dipole, while bent shapes (e.g., H₂O) result in measurable dipoles[6][9][13].Role of Machine Learning:
- Models like MuML use Gaussian process regression and symmetry-adapted kernels to predict atomic dipoles and charges, achieving accuracy comparable to quantum chemistry methods at lower computational cost[1][2][12].
- Recent approaches train neural networks on Wannier function centers (representing electron distributions) to predict dipole moments in liquids like methanol, capturing polarization effects from intermolecular interactions[12].
Clarifying Terminology:
- Atomic Weight ≠ Atomic Charge: Atomic weight (mass) is unrelated to charge distribution. Confusion may arise from terms like "weighting" in model parameters, which refers to the influence of atomic contributions, not isotopic mass[3][10].
- Atomic Polarizability vs. Mass: Polarizability depends on electron cloud flexibility, not nuclear mass. For example, larger atoms like bromine exhibit higher polarizability due to diffuse electron shells[7][16].
In summary, molecular dipole models prioritize electronic properties (charges, polarizabilities) and geometry, while atomic weights play no direct role. Machine learning bridges quantum accuracy and scalability, enabling predictions for complex systems like biomolecules[1][12].
Citations:
- 1: https://ww2.aip.org/scilights/a-machine-learning-model-to-predict-molecular-dipole-moments
- 2: https://pubmed.ncbi.nlm.nih.gov/32668949/
- 3: https://byjus.com/chemistry/difference-between-atomic-mass-and-atomic-weight/
- 4: https://pubs.acs.org/doi/10.1021/acs.jctc.4c00582
- 5: https://ch301.cm.utexas.edu/section2.php?target=imfs%2Fvsepr%2Fshape-dipole.html
- 6: https://en.wikipedia.org/wiki/Dipole
- 7: https://pmc.ncbi.nlm.nih.gov/articles/PMC8791008/
- 8: https://ursula.chem.yale.edu/~chem220/chem220js/STUDYAIDS/ESPotential/ESPdipoles.html
- 9: https://pressbooks.online.ucf.edu/chemistryfundamentals/chapter/7-3-molecular-polarity-and-dipole-moments/
- 10: https://pmc.ncbi.nlm.nih.gov/articles/PMC6104469/
- 11: https://www.youtube.com/watch?v=squbYnAHMqA
- 12: https://link.aps.org/doi/10.1103/PhysRevB.110.165159
- 13: https://wisc.pb.unizin.org/chem109fall2021ver02/chapter/dipole-dipole-attractions/
- 14: https://pubs.aip.org/aip/jcp/article/153/2/024113/1061473/Predicting-molecular-dipole-moments-by-combining
- 15: https://www.youtube.com/watch?v=mnf_4Uzz5sQ
- 16: https://www.osti.gov/pages/biblio/1681219
- 17: https://pubs.acs.org/doi/abs/10.1021/acs.jpclett.1c03476
Clarification on the
atomic_weightParameter UsageWhile the new
atomic_weightparameter has been added to support weighted dipole calculations inforward_common_atomic, note that in typical molecular dipole modeling atomic weights (which represent isotopic masses) are not directly used. Dipole calculations usually focus on atomic partial charges, polarizabilities, and molecular geometry. Please confirm that this parameter is intended as a custom weighting factor and, if so, consider updating its naming or documentation to clearly distinguish its purpose from the conventional concept of atomic weight.
- File:
deepmd/dpmodel/atomic_model/base_atomic_model.py(line 152)- Action: Verify the intended role of
atomic_weightand adjust its name/documentation as needed to avoid confusion.
ChiahsinChu
force-pushed
the
devel-dipole_with_atomic_weight
branch
from
d14fa86 to
ad063bd
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
deepmd/pt/model/atomic_model/base_atomic_model.py (2)
275-278: Consider adding shape validation for atomic_weightThe current implementation assumes that atomic_weight can be reshaped to
[out_shape[0], out_shape[1], -1]. If the provided atomic_weight has an incompatible shape, this could lead to runtime errors.Consider adding validation to check that the atomic_weight tensor has a compatible shape before reshaping, or at least document the expected shape more explicitly in the docstring.
if atomic_weight is not None: + expected_batch_size = out_shape[0] + expected_atoms = out_shape[1] + if atomic_weight.shape[0] != expected_batch_size or atomic_weight.shape[1] != expected_atoms: + raise ValueError(f"atomic_weight shape {atomic_weight.shape} incompatible with output shape {out_shape}") ret_dict[kk] = ret_dict[kk] * atomic_weight.view( [out_shape[0], out_shape[1], -1] )
279-284: Remove duplicated lineThe line
ret_dict["mask"] = atom_maskappears twice (lines 279 and 282). This is redundant and one of them should be removed.ret_dict["mask"] = atom_mask return ret_dict - ret_dict["mask"] = atom_mask - - return ret_dict
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (9)
deepmd/jax/model/base_model.py(1 hunks)deepmd/pt/model/atomic_model/base_atomic_model.py(4 hunks)deepmd/dpmodel/atomic_model/base_atomic_model.py(4 hunks)deepmd/dpmodel/model/make_model.py(12 hunks)source/tests/pt/model/test_dp_atomic_model.py(2 hunks)deepmd/dpmodel/atomic_model/base_atomic_model.py(1 hunks)deepmd/pt/model/atomic_model/base_atomic_model.py(1 hunks)deepmd/jax/atomic_model/dp_atomic_model.py(2 hunks)deepmd/jax/model/base_model.py(8 hunks)
🚧 Files skipped from review as they are similar to previous changes (5)
- deepmd/pt/model/atomic_model/base_atomic_model.py
- deepmd/jax/model/base_model.py
- deepmd/dpmodel/atomic_model/base_atomic_model.py
- deepmd/dpmodel/atomic_model/base_atomic_model.py
- deepmd/dpmodel/model/make_model.py
⏰ Context from checks skipped due to timeout of 90000ms (17)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C 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 C++ (clang, clang)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test C++ (false)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test C++ (true)
🔇 Additional comments (6)
deepmd/jax/model/base_model.py (2)
50-50: Logic change from c_differentiable to r_differentiableI noticed the condition changed from checking
vdef.c_differentiabletovdef.r_differentiable. This seems to be an intentional logic change since there's an assertion at line 98 that assumesvdef.r_differentiableis true whenvdef.c_differentiableis true.Could you confirm this is the intended behavior and not a bug? This change affects when the derivative calculations are performed.
30-30: Good implementation of atomic_weight parameterThe atomic_weight parameter has been properly added to the function signature and correctly propagated to all relevant function calls within the implementation. This ensures consistent handling of atomic weights throughout the computation path.
Also applies to: 39-39, 61-62, 73-75, 88-88, 111-112, 124-126, 144-144
source/tests/pt/model/test_dp_atomic_model.py (2)
77-82: Well-designed test case for atomic_weight in self_consistencyThe test implementation correctly verifies that applying atomic weights scales the energy output as expected. Multiplying the original energy by the reshaped atomic weights should match the energy when the weights are passed to the forward function.
112-118: Good cross-model consistency test for atomic_weightThis test appropriately verifies that the atomic_weight parameter works consistently between different model implementations (DPModel and PyTorch). This is important to ensure consistent behavior across the different backends.
deepmd/pt/model/atomic_model/base_atomic_model.py (1)
206-206: Correct implementation of atomic_weight parameterThe atomic_weight parameter has been properly added to the function signatures with appropriate documentation. The implementation correctly applies the weights to scale the output values.
Also applies to: 230-232, 275-278, 292-292, 302-302
deepmd/jax/atomic_model/dp_atomic_model.py (1)
61-61: Correctly propagated atomic_weight parameterThe atomic_weight parameter has been properly added to the function signature and correctly passed to the superclass method.
Also applies to: 70-70
ChiahsinChu
force-pushed
the
devel-dipole_with_atomic_weight
branch
from
0e3647f to
088b252
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (1)
deepmd/jax/utils/serialization.py (1)
78-91:⚠️ Potential issueUpdate exported signature to include atomic_weight
The exported function uses
jax.ShapeDtypeStructfor various parameters includingaparam, but there is no correspondingjax.ShapeDtypeStructfor the newly addedatomic_weightparameter. This could lead to serialization issues.Add a shape-dtype struct for
atomic_weightsimilar to the one foraparam:jax.ShapeDtypeStruct((nf, nloc, model.get_dim_aparam()), jnp.float64) if model.get_dim_aparam() else None, # aparam + jax.ShapeDtypeStruct((nf, nloc), jnp.float64) + if atomic_weight is not None + else None, # atomic_weight )
🧹 Nitpick comments (4)
deepmd/jax/jax2tf/make_model.py (1)
50-72: Update documentation for atomic_weight parameterThe function docstring should be updated to include information about the new
atomic_weightparameter.aparam atomic parameter. nf x nloc x nda + atomic_weight + atomic weights for scaling. nf x nloc do_atomic_virial If calculate the atomic virial.deepmd/jax/jax2tf/tfmodel.py (1)
82-82: Update docstring to include the new parameter.The
atomic_weightparameter has been added to the method signature, but it's missing from the docstring documentation.deepmd/jax/jax2tf/serialization.py (2)
85-86: Address the hardcoded dimension in a future updateThe comment indicates that the "1" should be replaced with the fitting output dimension. Consider creating a ticket to track this requirement for a future update.
120-121: Ensure consistent handling of the dimension issueThe same comment about replacing "1" with fitting output dimension appears multiple times. Ensure this is consistently addressed across all occurrences when implementing the fix.
Also applies to: 203-204, 229-230
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
deepmd/jax/model/base_model.py(1 hunks)deepmd/jax/atomic_model/dp_atomic_model.py(2 hunks)deepmd/jax/jax2tf/make_model.py(4 hunks)deepmd/jax/jax2tf/serialization.py(11 hunks)deepmd/jax/jax2tf/tfmodel.py(6 hunks)deepmd/jax/model/base_model.py(8 hunks)deepmd/jax/model/dp_model.py(2 hunks)deepmd/jax/utils/serialization.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/jax/model/base_model.py
- deepmd/jax/atomic_model/dp_atomic_model.py
⏰ Context from checks skipped due to timeout of 90000ms (18)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- 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 C++ (true)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
🔇 Additional comments (16)
deepmd/jax/model/base_model.py (4)
29-40: LGTM: atomic_weight parameter appropriately addedThe
atomic_weightparameter has been correctly added to the function signature and passed to the atomic model's forward method.
73-89: Update atomic_weight in eval_output functionThe
atomic_weightparameter has been correctly added to the function signature and passed to the atomic model's forward method.
124-144: LGTM: atomic_weight parameter correctly handled in atomic virial calculationThe
atomic_weightparameter is properly passed to theforward_common_atomicmethod in this code section.🧰 Tools
🪛 Ruff (0.8.2)
130-130: Function definition does not bind loop variable
def_ndim(B023)
50-50:✅ Verification successful
Verify intent of conditional change from c_differentiable to r_differentiable
The condition was changed from checking
vdef.c_differentiabletovdef.r_differentiable. Line 97-98 confirms thatr_differentiableis a prerequisite forc_differentiable(assert vdef.r_differentiable). This change potentially broadens which properties will have their derivatives calculated.
🏁 Script executed:
#!/bin/bash # Check where these differentiability flags are defined and what they represent rg -A 3 "r_differentiable|c_differentiable" --glob "*.py" | grep -v testLength of output: 29705
Clarify and Confirm Differentiability Flag Usage in Base Model
The change from checking
vdef.c_differentiableto checkingvdef.r_differentiableis consistent with the overall design in the codebase. In many modules (e.g., injax2tf/transform_output.pyandtfmodel.py),r_differentiableis used as a prerequisite forc_differentiable—as seen by theassert vdef.r_differentiableaccompanying the latter. This indicates that the intent is for properties flagged asr_differentiableto always have their derivatives calculated, while the additionalc_differentiableflag further controls specific derivative computations (with the built-in safety check ensuring consistency).Action Items:
- Documentation/Comment Update: Consider adding a clarifying comment in
deepmd/jax/model/base_model.py(and possibly related modules) to explicitly state that checkingr_differentiableis intentional and that it broadens the derivative computation scope, whilec_differentiableremains a subset requirement.- Review Consistency: Verify that downstream processing and derivative computations align with this broader condition as part of the overall differentiation strategy.
deepmd/jax/jax2tf/make_model.py (3)
47-47: LGTM: atomic_weight parameter added to function signatureThe
atomic_weightparameter has been properly added to the function signature.
77-78: LGTM: atomic_weight handled in parameter unpackingThe
atomic_weightparameter is correctly included in variable unpacking and removed from thedelstatement.
107-107: LGTM: atomic_weight correctly passed to call_lowerThe
atomic_weightparameter is properly passed to thecall_lowerfunction.deepmd/jax/utils/serialization.py (2)
60-61: LGTM: atomic_weight parameter added to function signatureThe
atomic_weightparameter has been correctly added to thecall_lower_with_fixed_do_atomic_virialfunction signature.
69-70: LGTM: atomic_weight correctly passed to call_lowerThe
atomic_weightparameter is properly passed to thecall_lowerfunction.deepmd/jax/model/dp_model.py (2)
59-60: LGTM: atomic_weight parameter added to function signatureThe
atomic_weightparameter has been correctly added to theforward_common_atomicmethod signature with an appropriate type annotation and default value.
70-71: LGTM: atomic_weight correctly passed to forward_common_atomicThe
atomic_weightparameter is properly passed to theforward_common_atomicfunction.deepmd/jax/jax2tf/tfmodel.py (1)
164-167: LGTM! Empty initialization for atomic_weightThe pattern of initializing
atomic_weightto an empty tensor with proper shape when it's None is consistent with how other optional parameters are handled.deepmd/jax/jax2tf/serialization.py (4)
43-49: LGTM! Parameter added correctlyThe
atomic_weightparameter has been correctly added to the function signature.
78-79: LGTM! Consistent shape specificationThe shape specification "(nf, nloc, 1)" for
atomic_weightis appropriate and consistent with the parameter's expected structure.
215-225: LGTM! Function call updated correctlyThe function signature and call to
make_call_whether_do_atomic_virialhave been properly updated to include theatomic_weightparameter.
241-251: LGTM! Function call updated correctlyThe function signature and call to
make_call_whether_do_atomic_virialhave been properly updated to include theatomic_weightparameter.
3 participants
Currently, the pytorch dipole model returns dipole (atomic quantities), global dipole (sum of atomic quantities) and force/virial operation on global dipoles. I am interested in not only the atomic dipole but also the derivative of total dipole w.r.t. coord, which is relevant to the dielectric response of the systems and has been used to train field-response MLPs (https://www.nature.com/articles/s41467-024-52491-3). However, the "global dipole" is not always the real total dipole of systems due to the missing information on charges. If the charges associated with the winner centroids are not identical, there is no way for users to get the derivative of (real) total dipole w.r.t. coord from DW models. (I tried to calculate the gradient of output dipole w.r.t. the input coord, but the gradient chain seems to be cut in the model.) Therefore, I add the
atomic_weightparameter to themodel.forward, which allows a user-defined atomic weight.There are some potential advantages of this implementation:
This implementation has been only activated for dipole models yet, but making it work for other models should be straightforword.
Summary by CodeRabbit
New Features
Tests