feat: add charge spin to cpp runtime#5509
Draft
anyangml wants to merge 9 commits into
Draft
Conversation
for more information, see https://pre-commit.ci
Contributor
📝 WalkthroughWalkthroughThe PR extends DeepMD-kit's inference stack with optional per-call charge/spin conditioning: backend interfaces expose dimension query and overloads accepting charge_spin vectors; DeepPot public API adds charge_spin parameter to all compute variants; PyTorch models build runtime tensors from caller vectors or metadata defaults; and LAMMPS integration enables explicit or compute-sourced charge_spin inputs with mutual-exclusion validation. ChangesCharge/Spin Conditioning Feature Implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds support for passing a runtime charge/spin conditioning vector from the LAMMPS pair style and C++ API down into the .pt2 (PTExpt) inference backend, instead of relying solely on defaults embedded in model metadata.
Changes:
- Extend LAMMPS
pair_style deepmdparsing to acceptcharge_spinandcharge_spin_from_compute, and thread the resulting vector intoDeepPot::compute(...). - Extend the C++
DeepPot/DeepPotBackendinterfaces and the PTExpt backend to accept and consume a runtimecharge_spininput with fallback to metadata defaults. - Add PTExpt
.pt2runtime tensor construction forcharge_spin, plus new computew overloads that carry it through.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
| source/lmp/pair_deepmd.cpp | Adds new pair_style keys for charge/spin and passes charge_spin into runtime inference calls. |
| source/lmp/pair_base.h | Adds storage/flags and a helper to populate charge_spin from a LAMMPS compute. |
| source/lmp/pair_base.cpp | Implements make_charge_spin_from_compute and initializes the new compute flag. |
| source/api_cc/src/DeepPotPTExpt.cc | Threads charge_spin into .pt2 model invocation and adds runtime/default tensor construction. |
| source/api_cc/src/DeepPot.cc | Extends DeepPot::compute overloads to accept charge_spin; adds dim_chg_spin() forwarding. |
| source/api_cc/include/DeepPotPTExpt.h | Declares PTExpt charge_spin-aware overloads and exposes dim_chg_spin(). |
| source/api_cc/include/DeepPot.h | Adds dim_chg_spin() to backends and provides default charge_spin-aware computew overloads. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
599
to
602
| dim_fparam = deep_pot_model_devi.dim_fparam(); | ||
| dim_aparam = deep_pot_model_devi.dim_aparam(); | ||
| dim_chg_spin = deep_pot.dim_chg_spin(); | ||
| assert(cutoff == deep_pot.cutoff() * dist_unit_cvt_factor); |
Comment on lines
+171
to
+176
| int icompute = modify->find_compute(compute_charge_spin_id); | ||
| Compute* compute = modify->compute[icompute]; | ||
|
|
||
| if (!compute) { | ||
| error->all(FLERR, "compute id is not found: " + compute_charge_spin_id); | ||
| } |
Comment on lines
+515
to
+522
| auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); | ||
| if (!charge_spin.empty()) { | ||
| charge_spin_tensor = | ||
| torch::from_blob(const_cast<double*>(charge_spin.data()), | ||
| {1, static_cast<std::int64_t>(charge_spin.size())}, | ||
| dbl_options) | ||
| .clone() | ||
| .to(device); |
Comment on lines
+839
to
+846
| auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); | ||
| if (!charge_spin.empty()) { | ||
| charge_spin_tensor = | ||
| torch::from_blob(const_cast<double*>(charge_spin.data()), | ||
| {1, static_cast<std::int64_t>(charge_spin.size())}, | ||
| dbl_options) | ||
| .clone() | ||
| .to(device); |
Comment on lines
+962
to
+967
| std::vector<double> frame_chg_spin; | ||
| if (!charge_spin.empty()) { | ||
| size_t s_dcsp = static_cast<size_t>(dchgspin); | ||
| frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, | ||
| charge_spin.begin() + (s_ff + 1) * s_dcsp); | ||
| } |
Comment on lines
+1139
to
+1144
| std::vector<double> frame_chg_spin; | ||
| if (!charge_spin.empty()) { | ||
| size_t s_dcsp = static_cast<size_t>(dchgspin); | ||
| frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, | ||
| charge_spin.begin() + (s_ff + 1) * s_dcsp); | ||
| } |
Comment on lines
+511
to
+514
| // Build charge_spin tensor: use runtime value when provided, fall back to | ||
| // default_chg_spin_ stored in the .pt2 metadata. | ||
| at::Tensor charge_spin_tensor; | ||
| if (dchgspin > 0) { |
Comment on lines
+699
to
+704
| } else if (string(arg[iarg]) == string("charge_spin")) { | ||
| for (int ii = 0; ii < dim_chg_spin; ++ii) { | ||
| if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) { | ||
| char tmp[1024]; | ||
| sprintf(tmp, "Illegal charge_spin, the dimension should be %d", | ||
| dim_chg_spin); |
Comment on lines
251
to
+256
| if (single_model || multi_models_no_mod_devi) { | ||
| // cvflag_atom is the right flag for the cvatom matrix | ||
| if (!(eflag_atom || cvflag_atom)) { | ||
| try { | ||
| deep_pot.compute(dener, dforce, dvirial, dcoord, dtype, dbox, nghost, | ||
| lmp_list, ago, fparam, daparam); | ||
| lmp_list, ago, fparam, daparam, charge_spin); |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/lmp/pair_deepmd.cpp (1)
318-320:⚠️ Potential issue | 🟠 MajorEnsure multi-model deviation propagates
charge_spin(or warns when unsupported)In
source/lmp/pair_deepmd.cpp, the multi-model deviation callsdeep_pot_model_devi.compute(...)at lines 318-320 and 326-329 withoutcharge_spin, unlike the single-model path which passescharge_spintodeep_pot.compute(). TheDeepBaseModelDeviAPI insource/api_c/include/deepmd.hppdoes not provide acompute(...charge_spin...)-style overload (nocompute.*charge_spinoccurrences), so deviation output cannot be consistent for models withdim_chg_spin > 0unless the API/caller is extended. Add a warning/error whendim_chg_spin > 0and multi-model deviation mode is enabled, or implement a charge_spin-aware deviation compute.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/lmp/pair_deepmd.cpp` around lines 318 - 320, The multi-model deviation path calls deep_pot_model_devi.compute(...) without passing charge_spin while the single-model path uses deep_pot.compute(..., charge_spin...), so when dim_chg_spin > 0 the deviation results are inconsistent; update the caller in source/lmp/pair_deepmd.cpp to detect dim_chg_spin > 0 when multi-model deviation (DeepBaseModelDevi) is active and either (A) emit a clear warning/error that charge_spin is unsupported for deviation mode (include dim_chg_spin and mode name in message) or (B) extend the deviation codepath to accept and forward charge_spin to DeepBaseModelDevi (implement a new compute overload or adapter in DeepBaseModelDevi and call it from deep_pot_model_devi.compute); reference deep_pot_model_devi.compute, deep_pot.compute, DeepBaseModelDevi, dim_chg_spin, and charge_spin when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/lmp/pair_deepmd.cpp`:
- Line 601: In the multi-model initialization path, change the source of
dim_chg_spin to be read from deep_pot_model_devi (like cutoff, numb_types,
numb_types_spin, dim_fparam, dim_aparam) and add an assertion that
deep_pot.dim_chg_spin() equals the value from deep_pot_model_devi; specifically,
replace or supplement the current direct read dim_chg_spin =
deep_pot.dim_chg_spin() with reading dim_chg_spin from deep_pot_model_devi
(e.g., dim_chg_spin = deep_pot_model_devi.dim_chg_spin()) and add the same
consistency check/assert between deep_pot and deep_pot_model_devi.
---
Outside diff comments:
In `@source/lmp/pair_deepmd.cpp`:
- Around line 318-320: The multi-model deviation path calls
deep_pot_model_devi.compute(...) without passing charge_spin while the
single-model path uses deep_pot.compute(..., charge_spin...), so when
dim_chg_spin > 0 the deviation results are inconsistent; update the caller in
source/lmp/pair_deepmd.cpp to detect dim_chg_spin > 0 when multi-model deviation
(DeepBaseModelDevi) is active and either (A) emit a clear warning/error that
charge_spin is unsupported for deviation mode (include dim_chg_spin and mode
name in message) or (B) extend the deviation codepath to accept and forward
charge_spin to DeepBaseModelDevi (implement a new compute overload or adapter in
DeepBaseModelDevi and call it from deep_pot_model_devi.compute); reference
deep_pot_model_devi.compute, deep_pot.compute, DeepBaseModelDevi, dim_chg_spin,
and charge_spin when making the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 173dd256-e998-4ff9-9db5-ca8571420859
📒 Files selected for processing (7)
source/api_cc/include/DeepPot.hsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/src/DeepPot.ccsource/api_cc/src/DeepPotPTExpt.ccsource/lmp/pair_base.cppsource/lmp/pair_base.hsource/lmp/pair_deepmd.cpp
| numb_types_spin = deep_pot_model_devi.numb_types_spin(); | ||
| dim_fparam = deep_pot_model_devi.dim_fparam(); | ||
| dim_aparam = deep_pot_model_devi.dim_aparam(); | ||
| dim_chg_spin = deep_pot.dim_chg_spin(); |
Contributor
There was a problem hiding this comment.
Query dim_chg_spin from deep_pot_model_devi for consistency.
In the multi-model initialization path, all other parameters (cutoff, numb_types, numb_types_spin, dim_fparam, dim_aparam) are queried from deep_pot_model_devi and then asserted to match deep_pot (lines 602-606). However, dim_chg_spin is queried only from deep_pot without a corresponding assertion.
🔧 Proposed fix
dim_fparam = deep_pot_model_devi.dim_fparam();
dim_aparam = deep_pot_model_devi.dim_aparam();
- dim_chg_spin = deep_pot.dim_chg_spin();
+ dim_chg_spin = deep_pot_model_devi.dim_chg_spin();
assert(cutoff == deep_pot.cutoff() * dist_unit_cvt_factor);
assert(numb_types == deep_pot.numb_types());
assert(numb_types_spin == deep_pot.numb_types_spin());
assert(dim_fparam == deep_pot.dim_fparam());
assert(dim_aparam == deep_pot.dim_aparam());
+ assert(dim_chg_spin == deep_pot.dim_chg_spin());
}📝 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.
Suggested change
| dim_chg_spin = deep_pot.dim_chg_spin(); | |
| dim_fparam = deep_pot_model_devi.dim_fparam(); | |
| dim_aparam = deep_pot_model_devi.dim_aparam(); | |
| dim_chg_spin = deep_pot_model_devi.dim_chg_spin(); | |
| assert(cutoff == deep_pot.cutoff() * dist_unit_cvt_factor); | |
| assert(numb_types == deep_pot.numb_types()); | |
| assert(numb_types_spin == deep_pot.numb_types_spin()); | |
| assert(dim_fparam == deep_pot.dim_fparam()); | |
| assert(dim_aparam == deep_pot.dim_aparam()); | |
| assert(dim_chg_spin == deep_pot.dim_chg_spin()); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@source/lmp/pair_deepmd.cpp` at line 601, In the multi-model initialization
path, change the source of dim_chg_spin to be read from deep_pot_model_devi
(like cutoff, numb_types, numb_types_spin, dim_fparam, dim_aparam) and add an
assertion that deep_pot.dim_chg_spin() equals the value from
deep_pot_model_devi; specifically, replace or supplement the current direct read
dim_chg_spin = deep_pot.dim_chg_spin() with reading dim_chg_spin from
deep_pot_model_devi (e.g., dim_chg_spin = deep_pot_model_devi.dim_chg_spin())
and add the same consistency check/assert between deep_pot and
deep_pot_model_devi.
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5509 +/- ##
==========================================
- Coverage 81.42% 79.75% -1.68%
==========================================
Files 871 872 +1
Lines 96951 97201 +250
Branches 4241 4271 +30
==========================================
- Hits 78941 77521 -1420
- Misses 16708 18542 +1834
+ Partials 1302 1138 -164 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes