feat(pt): add atomic_weight parameter to dipole model#2
Conversation
📝 WalkthroughWalkthroughAn optional per-atom Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✨ Finishing touches
📜 Recent review detailsConfiguration used: defaults Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (19)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
deepmd/pt/model/model/make_model.py (1)
131-203: Inconsistency:atomic_weightnot passed throughinput_type_castin PT version.In the dpmodel version (
deepmd/dpmodel/model/make_model.py),atomic_weightis passed throughinput_type_castand gets dtype-casted along with other inputs. However, in this PT version,atomic_weightis passed directly without casting.This inconsistency could cause dtype mismatches if
atomic_weightis provided with a different precision than the model's global precision.Consider adding
atomic_weightto theinput_type_castmethod for consistency:def input_type_cast( self, coord: torch.Tensor, box: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, + atomic_weight: Optional[torch.Tensor] = None, ) -> tuple[ torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], + Optional[torch.Tensor], str, ]:
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
deepmd/dpmodel/atomic_model/base_atomic_model.py(5 hunks)deepmd/dpmodel/model/make_model.py(11 hunks)deepmd/pt/model/atomic_model/base_atomic_model.py(5 hunks)deepmd/pt/model/model/dipole_model.py(4 hunks)deepmd/pt/model/model/make_model.py(4 hunks)source/tests/pt/model/test_dp_atomic_model.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
source/tests/pt/model/test_dp_atomic_model.py (3)
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
forward_common_atomic(152-235)deepmd/dpmodel/model/make_model.py (1)
forward_common_atomic(352-378)deepmd/pt/utils/utils.py (3)
to_numpy_array(224-224)to_numpy_array(228-228)to_numpy_array(231-247)
deepmd/pt/model/atomic_model/base_atomic_model.py (1)
deepmd/pt/model/network/network.py (1)
Tensor(36-37)
deepmd/dpmodel/model/make_model.py (2)
deepmd/pt/model/model/make_model.py (1)
input_type_cast(315-358)deepmd/pd/model/model/make_model.py (1)
input_type_cast(306-349)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (20)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Analyze (python)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
🔇 Additional comments (14)
deepmd/pt/model/atomic_model/base_atomic_model.py (2)
282-285: LGTM on atomic_weight scaling logic.The implementation correctly applies atomic_weight after the atom_mask, ensuring that virtual atoms remain zeroed. The scaling is properly applied before the "mask" key is added to ret_dict.
299-310: LGTM on forward method propagation.The
atomic_weightparameter is correctly added to the signature and propagated toforward_common_atomic.deepmd/pt/model/model/dipole_model.py (2)
56-92: LGTM on forward method.The
atomic_weightparameter is correctly added and propagated toforward_common.
94-133: LGTM on forward_lower method.The
atomic_weightparameter is correctly propagated toforward_common_lower.deepmd/dpmodel/atomic_model/base_atomic_model.py (2)
228-232: LGTM on atomic_weight scaling.The implementation correctly applies
atomic_weightafter masking. The reshape with-1allows broadcasting when the weight dimension matches the flattened output dimensions.
237-255: LGTM on call method.The
atomic_weightparameter is correctly added and propagated toforward_common_atomic.source/tests/pt/model/test_dp_atomic_model.py (2)
76-82: LGTM on atomic_weight test in self-consistency.Good test coverage validating that output energy is correctly scaled by
atomic_weight.
111-118: LGTM on atomic_weight test in DP consistency.Good cross-implementation validation ensuring NumPy and PyTorch versions produce consistent results when using
atomic_weight.deepmd/pt/model/model/make_model.py (2)
240-313: LGTM on forward_common_lower propagation.The
atomic_weightparameter is correctly propagated toatomic_model.forward_common_atomic.
622-639: Note:forward()does not includeatomic_weight.The base
forward()method doesn't include theatomic_weightparameter, which means models using this base implementation won't expose atomic_weight through the standard PyTorch interface. This appears intentional sinceDipoleModelprovides its ownforward()override that includesatomic_weight.deepmd/dpmodel/model/make_model.py (4)
225-283: LGTM on call method with atomic_weight integration.The
atomic_weightparameter is correctly integrated withinput_type_castand propagated tomodel_call_from_call_lower.
285-350: LGTM on call_lower method.The
atomic_weightparameter is correctly cast viainput_type_castand propagated toforward_common_atomic.
382-418: LGTM on input_type_cast extension.The method properly handles
atomic_weightdtype casting along with other inputs, maintaining consistency.
352-378: LGTM on forward_common_atomic.The
atomic_weightparameter is correctly propagated to the underlying atomic model.
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt/infer/deep_eval.py (1)
306-371: atomic_weight is undocumented in eval() and ignored when spin is provided
DeepEval.evalnow acceptsatomic_weightbut the docstring doesn’t describe it (purpose, expected shape, or broadcasting rules). That will confuse users trying to adopt the new feature.- In the
spinbranch,eval()calls_eval_model_spinwithout passingatomic_weight, so anyatomic_weightsupplied by the caller is silently ignored for spinful models. This is a behavioral discrepancy compared to the non-spin path and can yield scientifically inconsistent results without warning.I recommend:
- Documenting
atomic_weightin theevaldocstring (including allowed shapes and units/semantics).- Either (a) extending
_eval_model_spinand the spin model path to accept and applyatomic_weightas well, or (b) explicitly rejectingatomic_weightwhenspinis provided (e.g.,ValueError) so users don’t assume it is applied when it is not.
🧹 Nitpick comments (1)
deepmd/pt/infer/deep_eval.py (1)
474-541: Clarify atomic_weight shape expectations and auto-batch behavior in _eval_modelThe basic plumbing of
atomic_weightin_eval_modellooks correct: you reshape to(nframes, natoms, -1), convert withto_torch_tensor, and passatomic_weight=atomic_weight_inputinto the model, consistent with howfparamandaparamare handled (seeto_torch_tensorindeepmd/pt/utils/utils.py).A few points worth tightening up:
- The reshape
atomic_weight.reshape(nframes, natoms, -1)implicitly constrains acceptable input shapes. It will work nicely for(nframes, natoms)or(nframes, natoms, k)arrays/tensors, but will fail for shapes like(natoms,)that one might reasonably expect to be broadcast across frames. If you intend to support only full(nframes, natoms, [dim])inputs, this should be documented. If not, you’ll need extra logic (similar to howaparamis documented) to handle per-atom or global weights explicitly before reshaping.atomic_weightis threaded into_eval_func(self._eval_model, numb_test, natoms)(...)as a trailing positional argument. That meansAutoBatchSize.execute_allmust be aware that there is now an additional per-frame/per-atom array in*argsand split it consistently across batches. It’s probably fine, but worth confirming via tests that auto-batching withatomic_weightbehaves as expected.No immediate correctness bug for the documented use-case, but clarifying the supported shapes and verifying the auto-batch path would make this much more robust.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
deepmd/pt/infer/deep_eval.py(5 hunks)deepmd/pt/train/wrapper.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/infer/deep_eval.py (1)
deepmd/pt/utils/utils.py (3)
to_torch_tensor(251-251)to_torch_tensor(255-255)to_torch_tensor(258-276)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (20)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test C++ (true, true, true, false)
- GitHub Check: Test C++ (false, true, true, false)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Test C++ (true, false, false, true)
- GitHub Check: Test C++ (false, false, false, true)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
deepmd/pt/train/wrapper.py (1)
154-168: atomic_weight plumbing in ModelWrapper.forward looks correctThe new
atomic_weightargument is added at the end of the signature with a default ofNoneand is consistently threaded throughinput_dictfor both inference (self.model[task_key](**input_dict)) and training (passed as part ofinput_dictto the loss). This should be backward-compatible with existing call sites and enables losses/models that understandatomic_weightto consume it without impacting others.Also applies to: 175-183
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (1)
deepmd/pt/infer/deep_eval.py (1)
657-699: Dead code: validation never reached.This validation block is identical to lines 534-576 in
_eval_model. However, it's unreachable dead code becauseeval()raises an error whenatomic_weightis provided with spin (lines 379-382).This duplication should be addressed by either:
- Removing atomic_weight support from
_eval_model_spinentirely (if spin + atomic_weight shouldn't be supported)- Removing the error check in
eval()and extracting the shared validation logic (if it should be supported)See the earlier comment on lines 379-382 for detailed recommendations.
🧹 Nitpick comments (1)
deepmd/pt/infer/deep_eval.py (1)
543-557: Document ambiguous shape handling.When
natoms == nframes, a 2D input is ambiguous—it could be interpreted as either(nframes, natoms)or(natoms, nframes). The code resolves this by preferring the(nframes, natoms)interpretation (since the condition on line 546 requiresshape[1] != nframes).This is a reasonable choice, but consider adding a comment or updating the docstring to clarify this edge case for users.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
deepmd/pt/infer/deep_eval.py(9 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
deepmd/pt/infer/deep_eval.py (1)
deepmd/pt/utils/utils.py (3)
to_torch_tensor(251-251)to_torch_tensor(255-255)to_torch_tensor(258-276)
🪛 Ruff (0.14.8)
deepmd/pt/infer/deep_eval.py
380-382: Avoid specifying long messages outside the exception class
(TRY003)
539-541: Avoid specifying long messages outside the exception class
(TRY003)
555-557: Avoid specifying long messages outside the exception class
(TRY003)
564-566: Avoid specifying long messages outside the exception class
(TRY003)
568-570: Avoid specifying long messages outside the exception class
(TRY003)
662-664: Avoid specifying long messages outside the exception class
(TRY003)
678-680: Avoid specifying long messages outside the exception class
(TRY003)
687-689: Avoid specifying long messages outside the exception class
(TRY003)
691-693: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (20)
- 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
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Build C library (2.14, >=2.5.0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- 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: Test Python (3, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (c-cpp)
🔇 Additional comments (1)
deepmd/pt/infer/deep_eval.py (1)
314-314: Well-documented parameter addition.The
atomic_weightparameter is properly added with clear documentation covering the supported shapes.Also applies to: 344-349
| if atomic_weight is not None: | ||
| # Validate atomic_weight shape | ||
| if atomic_weight.ndim == 1: | ||
| # (natoms,) - broadcast across frames | ||
| if atomic_weight.shape[0] != natoms: | ||
| raise ValueError( | ||
| f"atomic_weight shape {atomic_weight.shape} incompatible with number of atoms {natoms}" | ||
| ) | ||
| atomic_weight = np.tile(atomic_weight, (nframes, 1)) | ||
| elif atomic_weight.ndim == 2: | ||
| # (nframes, natoms) or (natoms, k) | ||
| if ( | ||
| atomic_weight.shape[0] == natoms | ||
| and atomic_weight.shape[1] != nframes | ||
| ): | ||
| # (natoms, k) - broadcast across frames | ||
| atomic_weight = np.tile(atomic_weight, (nframes, 1, 1)) | ||
| elif ( | ||
| atomic_weight.shape[0] != nframes | ||
| or atomic_weight.shape[1] != natoms | ||
| ): | ||
| raise ValueError( | ||
| f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}" | ||
| ) | ||
| elif atomic_weight.ndim == 3: | ||
| # (nframes, natoms, k) | ||
| if ( | ||
| atomic_weight.shape[0] != nframes | ||
| or atomic_weight.shape[1] != natoms | ||
| ): | ||
| raise ValueError( | ||
| f"atomic_weight shape {atomic_weight.shape} incompatible with nframes={nframes} and natoms={natoms}" | ||
| ) | ||
| else: | ||
| raise ValueError( | ||
| f"atomic_weight must have 1, 2, or 3 dimensions, got {atomic_weight.ndim}" | ||
| ) | ||
|
|
||
| atomic_weight_input = to_torch_tensor( | ||
| atomic_weight.reshape(nframes, natoms, -1) | ||
| ) | ||
| else: | ||
| atomic_weight_input = None |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Extract duplicated validation logic to a helper method.
This validation block is duplicated identically in _eval_model_spin (lines 657-699), violating the DRY principle. The logic is also quite complex and would benefit from being in a separate, testable method.
Extract to a helper method:
def _validate_and_reshape_atomic_weight(
self,
atomic_weight: Optional[np.ndarray],
nframes: int,
natoms: int,
) -> Optional[np.ndarray]:
"""Validate and reshape atomic_weight to (nframes, natoms, k).
Parameters
----------
atomic_weight : np.ndarray, optional
Input atomic weight array
nframes : int
Number of frames
natoms : int
Number of atoms
Returns
-------
np.ndarray, optional
Reshaped array of shape (nframes, natoms, k) or None
Raises
------
ValueError
If shape is incompatible with nframes and natoms
"""
if atomic_weight is None:
return None
# Validation logic here...
# (lines 536-574)
return atomic_weight.reshape(nframes, natoms, -1)Then use it in both _eval_model and _eval_model_spin:
- if atomic_weight is not None:
- # Validate atomic_weight shape
- if atomic_weight.ndim == 1:
- ...
- atomic_weight_input = to_torch_tensor(
- atomic_weight.reshape(nframes, natoms, -1)
- )
- else:
- atomic_weight_input = None
+ reshaped_weight = self._validate_and_reshape_atomic_weight(
+ atomic_weight, nframes, natoms
+ )
+ atomic_weight_input = to_torch_tensor(reshaped_weight) if reshaped_weight is not None else None🧰 Tools
🪛 Ruff (0.14.8)
539-541: Avoid specifying long messages outside the exception class
(TRY003)
555-557: Avoid specifying long messages outside the exception class
(TRY003)
564-566: Avoid specifying long messages outside the exception class
(TRY003)
568-570: Avoid specifying long messages outside the exception class
(TRY003)
🤖 Prompt for AI Agents
In deepmd/pt/infer/deep_eval.py around lines 534 to 576, the atomic_weight
validation/reshaping logic is duplicated (also in _eval_model_spin lines
~657-699); extract this into a private helper method (e.g.
_validate_and_reshape_atomic_weight(self, atomic_weight, nframes, natoms)) that
returns either None or a np.ndarray shaped (nframes, natoms, k), move all shape
checks and broadcasting logic into it (preserve all current ValueError
messages), and replace the duplicated blocks in both _eval_model and
_eval_model_spin with a call to this helper and then to_torch_tensor on its
result if not None.
ChiahsinChu
force-pushed
the
patch/v3.1.1/dipole_with_atomic_weight
branch
from
516f1bf to
0c55d75
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
deepmd/dpmodel/model/make_model.py (1)
5-7: MissingOptionalimport causes NameError.The
Optionaltype hint is used on lines 71, 234, 295, and 362, but it's not imported from thetypingmodule. This will cause aNameErrorat runtime.Apply this diff to fix the import:
from typing import ( Any, + Optional, )deepmd/pt/model/model/make_model.py (1)
5-7: MissingOptionalimport causes NameError.The
Optionaltype hint is used on lines 141, 258, and 645, but it's not imported from thetypingmodule. This will cause aNameErrorat runtime.Apply this diff to fix the import:
from typing import ( Any, + Optional, )deepmd/pt/infer/deep_eval.py (1)
607-714: Dead code:atomic_weighthandling in_eval_model_spinis unreachable.Since
eval()raises an error at lines 380-383 whenatomic_weightis provided with spin, the_eval_model_spinmethod will always receiveatomic_weight=None. This makes:
- The parameter at line 616 effectively unused
- The validation logic at lines 658-700 unreachable dead code
- The
atomic_weight_inputat line 713 alwaysNoneConsider one of these options:
Option A: Remove atomic_weight from _eval_model_spin entirely
def _eval_model_spin( self, coords: np.ndarray, cells: np.ndarray | None, atom_types: np.ndarray, spins: np.ndarray, fparam: np.ndarray | None, aparam: np.ndarray | None, request_defs: list[OutputVariableDef], - atomic_weight: np.ndarray | None = None, ) -> tuple[np.ndarray, ...]:And remove lines 658-700 and the
atomic_weight=atomic_weight_inputfrom line 713.Option B: Keep for future support but add a comment explaining it's reserved for future use.
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
221-232: Add validation or clarification for atomic_weight broadcasting.The atomic_weight scaling multiplies each output by
atomic_weight.reshape([_out_shape[0], _out_shape[1], -1]). This assumes the third dimension of atomic_weight matches the flattened trailing dimensions of the output. While tests show atomic_weight shape[nf, nloc, 1]for energy outputs (which have shape[nf, nloc, 1]after flattening), the code does not validate this constraint.For outputs with different dimensions—e.g., a tensor property with shape
[nf, nloc, 3, 3]flattened to[nf, nloc, 9]—broadcasting with[nf, nloc, 1]atomic_weight would fail or behave unexpectedly. Either validate thatatomic_weight.shape[2] == math.prod(out_shape[2:])for each output, or explicitly document this requirement in the docstring.Additionally, the "mask" key added at line 232 is not scaled, which is correct since it's added after the loop completes.
deepmd/pt/model/model/dipole_model.py (1)
55-73: Fix undefinedOptional— change to PEP 604 union syntaxOn line 63,
atomic_weight: Optional[torch.Tensor]referencesOptional, which is not imported. Since__future__ annotationsis not enabled, this will raise aNameErrorat runtime.The rest of the file already uses the PEP 604 union syntax (
torch.Tensor | Noneon lines 59–61), so update line 63 to match:- atomic_weight: Optional[torch.Tensor] = None, + atomic_weight: torch.Tensor | None = None,This fixes the undefined name error and keeps typing consistent throughout the file.
🧹 Nitpick comments (2)
deepmd/pt/infer/deep_eval.py (1)
535-577: Consider extracting validation logic to a helper method.This validation/reshaping logic for
atomic_weight(lines 535-577) is complex and handles multiple shape scenarios (1D, 2D, 3D). While the implementation is correct, extracting it to a helper method would:
- Improve readability
- Make the logic easier to test in isolation
- Enable reuse if needed elsewhere
Example helper signature:
def _validate_and_reshape_atomic_weight( self, atomic_weight: np.ndarray | None, nframes: int, natoms: int, ) -> np.ndarray | None: """Validate and reshape atomic_weight to (nframes, natoms, k).""" ...deepmd/pt/train/wrapper.py (1)
152-200: Consider addingatomic_weighttoinput_dictonly when notNoneThe new
atomic_weightparameter is correctly threaded through to the model and loss, but it’s always included ininput_dict, even whenNone. If any existing model/loss still has a strictforward(...)signature withoutatomic_weight, this can causeTypeError: got an unexpected keyword argument 'atomic_weight'.To keep this wrapper backward‑compatible while still supporting the new feature, only add the key when a non‑
Nonetensor is provided:- input_dict = { - "coord": coord, - "atype": atype, - "box": box, - "do_atomic_virial": do_atomic_virial, - "fparam": fparam, - "aparam": aparam, - "atomic_weight": atomic_weight, - } + input_dict = { + "coord": coord, + "atype": atype, + "box": box, + "do_atomic_virial": do_atomic_virial, + "fparam": fparam, + "aparam": aparam, + } + if atomic_weight is not None: + input_dict["atomic_weight"] = atomic_weight
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (8)
deepmd/dpmodel/atomic_model/base_atomic_model.py(5 hunks)deepmd/dpmodel/model/make_model.py(11 hunks)deepmd/pt/infer/deep_eval.py(9 hunks)deepmd/pt/model/atomic_model/base_atomic_model.py(5 hunks)deepmd/pt/model/model/dipole_model.py(4 hunks)deepmd/pt/model/model/make_model.py(12 hunks)deepmd/pt/train/wrapper.py(2 hunks)source/tests/pt/model/test_dp_atomic_model.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/atomic_model/base_atomic_model.py
🧰 Additional context used
🧬 Code graph analysis (3)
deepmd/pt/infer/deep_eval.py (1)
deepmd/pt/utils/utils.py (3)
to_torch_tensor(249-249)to_torch_tensor(253-253)to_torch_tensor(256-274)
deepmd/pt/train/wrapper.py (2)
deepmd/pt/model/network/network.py (1)
Tensor(34-35)deepmd/pd/model/network/network.py (1)
Tensor(26-29)
deepmd/pt/model/model/dipole_model.py (1)
deepmd/pt/model/network/network.py (1)
Tensor(34-35)
🪛 Ruff (0.14.8)
deepmd/pt/infer/deep_eval.py
381-383: Avoid specifying long messages outside the exception class
(TRY003)
540-542: Avoid specifying long messages outside the exception class
(TRY003)
556-558: Avoid specifying long messages outside the exception class
(TRY003)
565-567: Avoid specifying long messages outside the exception class
(TRY003)
569-571: Avoid specifying long messages outside the exception class
(TRY003)
663-665: Avoid specifying long messages outside the exception class
(TRY003)
679-681: Avoid specifying long messages outside the exception class
(TRY003)
688-690: Avoid specifying long messages outside the exception class
(TRY003)
692-694: Avoid specifying long messages outside the exception class
(TRY003)
deepmd/dpmodel/model/make_model.py
71-71: Undefined name Optional
(F821)
234-234: Undefined name Optional
(F821)
295-295: Undefined name Optional
(F821)
362-362: Undefined name Optional
(F821)
deepmd/pt/model/model/make_model.py
141-141: Undefined name Optional
(F821)
258-258: Undefined name Optional
(F821)
645-645: Undefined name Optional
(F821)
deepmd/pt/model/model/dipole_model.py
63-63: Undefined name Optional
(F821)
🔇 Additional comments (15)
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
236-254: LGTM!The
callmethod correctly accepts and propagatesatomic_weighttoforward_common_atomic.deepmd/dpmodel/model/make_model.py (5)
48-131: LGTM!The
atomic_weightparameter is correctly threaded throughmodel_call_from_call_lowertocall_lower. The implementation is consistent with the overall design.
226-284: LGTM!The
CM.callmethod correctly:
- Receives
atomic_weightfrominput_type_castasaw- Passes
awtomodel_call_from_call_lower
286-351: LGTM!The
CM.call_lowermethod correctly threadsatomic_weightthrough input type casting and passes it toforward_common_atomic.
383-414: LGTM!The
input_type_castmethod correctly handlesatomic_weight:
- Casts to coord dtype alongside other parameters
- Returns in the extended tuple
- Handles precision conversion consistently with other optional parameters
353-379: LGTM!The
forward_common_atomicmethod correctly propagatesatomic_weightto the underlying atomic model.deepmd/pt/model/model/make_model.py (4)
133-210: LGTM!The
forward_commonmethod correctly:
- Accepts
atomic_weightparameter with docstring- Receives
awfrominput_type_cast- Deletes
atomic_weightfrom local scope- Passes
awtoforward_common_lower
247-325: LGTM!The
forward_common_lowermethod correctly threadsatomic_weightthrough input type casting toatomic_model.forward_common_atomic.
327-373: LGTM!The
input_type_castmethod correctly handlesatomic_weightwith proper dtype conversion using PyTorch's.to()method.
637-656: LGTM!The
forwardmethod correctly accepts and passesatomic_weighttoforward_common.deepmd/pt/infer/deep_eval.py (2)
307-398: LGTM for the publicevalmethod.The
atomic_weightparameter is correctly added with appropriate docstring. The error raised whenatomic_weightis provided with spin (lines 380-383) makes the design intent clear.
485-589: LGTM for_eval_modelatomic_weight handling.The validation logic correctly handles various input shapes:
- 1D
(natoms,)→ broadcast across frames- 2D
(nframes, natoms)or(natoms, k)→ appropriate handling- 3D
(nframes, natoms, k)→ direct useThe final reshape to
(nframes, natoms, -1)ensures consistent shape for the model.source/tests/pt/model/test_dp_atomic_model.py (2)
76-82: LGTM! Good test coverage for atomic_weight scaling.The test correctly verifies that
forward_common_atomicwithatomic_weightproduces energy scaled by the weight. The assertionret0["energy"] * aw.reshape(nf, nloc, -1)matches the expected scaling behavior implemented in the atomic model.
111-118: LGTM! Good cross-implementation consistency test.The test verifies that atomic_weight produces consistent results between the DP (numpy) and PT (torch) implementations, which is important for maintaining compatibility across backends.
deepmd/pt/model/model/dipole_model.py (1)
93-117: atomic_weight propagation inforward_lowerlooks correct
forward_lowernow acceptsatomic_weight: torch.Tensor | Noneand passes it through toforward_common_lower(..., atomic_weight=atomic_weight, ...), which matches the pattern used in the higher‑levelforwardmethod and keeps the lower‑level path in sync with the main dipole path.
ChiahsinChu
force-pushed
the
patch/v3.1.1/dipole_with_atomic_weight
branch
from
2a95032 to
9359011
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
deepmd/pt/infer/deep_eval.py (1)
315-350: Fix eval docstring vs behavior for atomic_weight with spinIn
eval, whenspinis provided andatomic_weightis notNone, aValueErroris raised (Lines 380–383), but the docstring still says:Note: This parameter is currently ignored when spin is provided.
That’s misleading for users.
Consider updating both the doc and the error message to reflect the actual behavior, e.g.:
- atomic_weight - The atomic weight for weighted averaging of atomic contributions. + atomic_weight + The atomic weight for weighted averaging of atomic contributions + on non-spin models. @@ - - nframes x natoms x k. Different weights for each atom in each frame with k dimensions. - Note: This parameter is currently ignored when spin is provided. + - nframes x natoms x k. Different weights for each atom in each frame with k dimensions. + Note: This parameter is not supported together with ``spin``. + A :class:`ValueError` is raised if both are provided. @@ - if atomic_weight is not None: - raise ValueError( - "atomic_weight is not supported when spin is provided." - ) + if atomic_weight is not None: + raise ValueError( + "atomic_weight is not supported when spin is provided; " + "omit atomic_weight or disable spin." + )This keeps the API contract explicit and avoids surprises.
Also applies to: 370-383
deepmd/pt/model/model/dipole_model.py (1)
55-73: Avoid breaking forward_lower callers by reordering new atomic_weight parameterAdding
atomic_weighttoforwardis a pure extension and safe. Inforward_lower, though, the new signature is:def forward_lower( self, extended_coord, extended_atype, nlist, mapping=None, fparam=None, aparam=None, do_atomic_virial: bool = False, atomic_weight: torch.Tensor | None = None, # new comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]:Previously,
comm_dictwas the positional argument immediately afterdo_atomic_virial. Any existing callers that do:dipole.forward_lower(..., do_atomic_virial, comm_dict)will now have
comm_dictbound toatomic_weightinstead, which is a silent but serious behavior change (likely causing shape errors or mis-scaling).To keep backward compatibility, consider keeping
comm_dictin its original position and appendingatomic_weightat the end:- def forward_lower( + def forward_lower( self, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - atomic_weight: torch.Tensor | None = None, - comm_dict: dict[str, torch.Tensor] | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + atomic_weight: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: @@ - do_atomic_virial=do_atomic_virial, - comm_dict=comm_dict, - extra_nlist_sort=self.need_sorted_nlist_for_lower(), - atomic_weight=atomic_weight, + do_atomic_virial=do_atomic_virial, + comm_dict=comm_dict, + extra_nlist_sort=self.need_sorted_nlist_for_lower(), + atomic_weight=atomic_weight,Since
forward_common_loweris called with keyword arguments, this change won’t affect that internal call, but will preserve existing positional uses ofcomm_dict.Also applies to: 94-117
🧹 Nitpick comments (4)
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
151-160: Clarify atomic_weight shape semantics and validation behaviorThe per-key scaling logic looks correct (masking then multiplying by a reshaped
atomic_weight), but a couple of details are worth tightening:
- The code enforces only
atomic_weight.shape[2] == prod(out_shape[2:]), while the docstring says the full shape isnf x nloc x dim_aw. If you intend to require per-frame/per-atom weights (no broadcasting on the first two axes), consider also checkingatomic_weight.shape[:2] == _out_shape[:2]for clearer failure modes. If broadcasting over frames/atoms is intended, it would help to document that explicitly.- The long
ValueErrormessage is a bit unwieldy; consider shortening it or moving some detail into a helper/constant if you want to appease static analysis.Functionally this is fine; these are mostly API clarity and ergonomics points.
Also applies to: 181-183, 221-239
deepmd/pt/infer/deep_eval.py (1)
881-936: Document all accepted atomic_weight shapes and simplify error textThe new
_validate_and_reshape_atomic_weighthelper is a good consolidation of shape handling, but:
- It supports shapes beyond what
eval’s docstring mentions:(natoms,)and(natoms, k)(broadcast over frames) as well as(nframes, natoms)and(nframes, natoms, k). It would be helpful to either reference this helper from the public docs or expand the docstring to list these accepted patterns so users know what’s allowed.- The multiple
ValueErrormessages are quite long and repetitive; you might consider shortening them a bit or factoring shared wording into a small internal helper/constant if you want to quiet static-analysis complaints.Functionally the logic (including the tiling and nframes/natoms checks) looks sound.
deepmd/dpmodel/model/make_model.py (1)
71-72: atomic_weight propagation through dpmodel CM looks consistentThe dpmodel factory now:
- Accepts
atomic_weightinmodel_call_from_call_lower,CM.call,CM.call_lower, andCM.forward_common_atomic.- Normalizes its dtype in
input_type_castalongsidebox,fparam, andaparam.- Passes the casted
awdown toatomic_model.forward_common_atomic, which owns the actual scaling logic.This wiring is coherent and preserves backward compatibility (all new params are optional and default to
None). If you want to tidy things further, you might:
- Mention
atomic_weightexplicitly in themodel_call_from_call_lowerandinput_type_castdocstrings, and- Align the
input_type_castreturn type hints to useArray | Noneinstead ofnp.ndarray | Nonefor consistency with the rest of the file.These are style-only; behavior looks correct.
Also applies to: 226-283, 286-352, 383-414
deepmd/pt/model/model/make_model.py (1)
132-177: atomic_weight integration in PT factory is coherent and backward-compatibleOn the PyTorch side:
forward_commonandforwardnow take an optionalatomic_weightand pass it intoinput_type_cast.input_type_castcastsatomic_weightto the coordinate dtype and returns it asawalongsidecoord,box,fparam, andaparam.forward_common_lowerpicks upawand passes it asatomic_weightintoatomic_model.forward_common_atomic, then intofit_output_to_model_output.All new parameters are optional and appended, so existing call sites remain valid. The design mirrors the dpmodel path and matches how DeepEval prepares a
(nframes, natoms, k)tensor.Only minor polish you might consider: expand the docstrings for
forward_common/forward_common_lowerslightly to mention thatatomic_weightis expected to have shape(nf, nloc, dim_aw)after any broadcasting/reshaping performed at higher levels (DeepEval).Also applies to: 247-315, 327-373, 637-656
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
deepmd/dpmodel/atomic_model/base_atomic_model.py(5 hunks)deepmd/dpmodel/model/make_model.py(11 hunks)deepmd/pt/infer/deep_eval.py(9 hunks)deepmd/pt/model/model/dipole_model.py(4 hunks)deepmd/pt/model/model/make_model.py(12 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
deepmd/dpmodel/model/make_model.py (2)
deepmd/pt/model/model/make_model.py (1)
input_type_cast(327-373)deepmd/pd/model/model/make_model.py (1)
input_type_cast(303-346)
deepmd/pt/model/model/make_model.py (2)
deepmd/pt/model/network/network.py (1)
Tensor(34-35)deepmd/dpmodel/model/make_model.py (1)
input_type_cast(383-414)
🪛 Ruff (0.14.8)
deepmd/pt/infer/deep_eval.py
381-383: Avoid specifying long messages outside the exception class
(TRY003)
912-914: Avoid specifying long messages outside the exception class
(TRY003)
922-924: Avoid specifying long messages outside the exception class
(TRY003)
928-930: Avoid specifying long messages outside the exception class
(TRY003)
932-934: Avoid specifying long messages outside the exception class
(TRY003)
deepmd/dpmodel/atomic_model/base_atomic_model.py
233-236: Avoid specifying long messages outside the exception class
(TRY003)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (8)
- GitHub Check: Test C++ (false, false, false, true)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (true, false, false, true)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (false, false, false, true)
- GitHub Check: Test C++ (true, true, true, false)
- GitHub Check: Test C++ (false, true, true, false)
🔇 Additional comments (2)
deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
244-262: call wrapper correctly forwards atomic_weightThe
callmethod’s newatomic_weightparameter is threaded straight through toforward_common_atomicand preserves the previous signature as a strict extension, so existing callers remain compatible.deepmd/pt/infer/deep_eval.py (1)
485-555: Non‑spin and spin model paths handle atomic_weight consistentlyThe non‑spin
_eval_modelpath now:
- Accepts
atomic_weight,- Validates/reshapes it via
_validate_and_reshape_atomic_weight,- Converts to a
(nframes, natoms, k)torch tensor,- And passes it as
atomic_weight=...to the model.The spin path
_eval_model_spin:
- No longer takes
atomic_weight,- Always passes
atomic_weight=Noneinto the model,- And
evalforbids non‑Noneatomic_weightwhenspinis present.This separation looks coherent and should keep the PT models’ signature uniform (
atomic_weightkwarg always exists) while making spin behavior explicit at the wrapper level.Also applies to: 573-637
ChiahsinChu
force-pushed
the
patch/v3.1.1/dipole_with_atomic_weight
branch
from
9359011 to
02fa9d0
Compare
1 participant
This change is part of the dipole_with_atomic_weight feature branch.
Summary by CodeRabbit
New Features
Tests
Chores
✏️ Tip: You can customize this high-level summary in your review settings.