feat: add charge spin to cpp runtime - #5509
Conversation
for more information, see https://pre-commit.ci
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThreads optional per-call charge_spin through DeepPot backend and public compute APIs, PT/PTExpt runtimes, C API and deepmd.hpp bridge, LAMMPS PairDeep wiring, metadata export, and adds generators/tests exercising explicit vs default charge_spin and nlist paths. ChangesCharge/Spin Conditioning Feature Implementation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
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 |
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.
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
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% 82.13% +0.71%
==========================================
Files 871 894 +23
Lines 96951 102064 +5113
Branches 4241 4308 +67
==========================================
+ Hits 78941 83833 +4892
- Misses 16708 16893 +185
- Partials 1302 1338 +36 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/tests/infer/gen_chg_spin.py`:
- Around line 3-9: The docstring's explicit charge_spin example is out of sync
with the actual value used in the script: update the docstring text that
currently shows "[0.5, 0.8]" to match the script's actual explicit charge_spin
"[1.0, 2.0]" so the described [explicit] section matches the value used when
writing chg_spin.expected (the variable/parameter named charge_spin in this test
generator).
🪄 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: 2de1ac02-e00d-42aa-93c7-625581682ce7
📒 Files selected for processing (7)
deepmd/pt_expt/utils/serialization.pysource/api_c/include/deepmd.hppsource/api_cc/tests/test_deeppot_chg_spin_ptexpt.ccsource/install/test_cc_local.shsource/lmp/pair_base.cppsource/tests/infer/gen_chg_spin.pysource/tests/pt_expt/infer/test_deep_eval.py
🚧 Files skipped from review as they are similar to previous changes (1)
- source/lmp/pair_base.cpp
for more information, see https://pre-commit.ci
iProzd
left a comment
There was a problem hiding this comment.
The following should be completed for the main runtime paths before merge:
-
Please wire
charge_spinthrough the C API /deepmd.hpppath.
In common LAMMPS builds withDP_USING_C_API,PairDeepMDusesdeepmd::hpp::DeepPotfromsource/api_c/include/deepmd.hpp. The newcharge_spinargument is currently accepted there but discarded via(void)charge_spin, socharge_spinhas no effect on this widely used path. Please add proper C API entry points insource/api_c/include/c_api.h/source/api_c/src/c_api.ccand havedeepmd.hppforwardcharge_spinthrough to the underlyingapi_cc::DeepPotimplementation. -
Please support the regular PyTorch
.pthbackend as well.
DeepPotPTcurrently does not override the newcharge_spin-awarecomputewoverloads, so.pthmodels fall back to the defaultDeepPotBackendimplementation and silently ignorecharge_spin. If runtimecharge_spinis a C++ runtime feature, it should also be implemented for the PT backend, not onlyDeepPotPTExpt/.pt2. -
Please also thread
charge_spinthrough model-deviation.
LAMMPS model-deviation currently callsdeep_pot_model_devi.compute(..., fparam, aparam)withoutcharge_spin, so the main model path and deviation path can use different effective conditioning.
Without these, charge_spin support is incomplete and can silently no-op in common usage.
DeepPotPT silently dropped runtime charge_spin: it never read the model's charge/spin dimension and never passed charge_spin to forward/forward_lower, so .pth charge_spin models (DPA3 add_chg_spin_ebd) fell back to the stored default_chg_spin. This also affected the model-deviation path when its sub-models are .pth. - DeepPotPT: override dim_chg_spin(); read dim_chg_spin / default_chg_spin from the jit-exported model methods (guarded with find_method for old .pth without them); build the charge_spin tensor (runtime value or default fallback) and thread it into forward (no-nlist) and forward_lower (nlist), only when the model has a charge/spin embedding so other models are unaffected. For the non-message-passing edge case a real None comm_dict is passed (an empty Dict would wrongly flip parallel_mode). - gen_chg_spin.py: also export chg_spin.pth and verify it reproduces the .pt2 reference (PBC + NoPbc); add NoPbc reference sections. - Add test_deeppot_chg_spin_pt.cc mirroring the DPA3 .pth layout (PBC standalone + NoPbc lmp_nlist), covering explicit charge_spin and the default fallback. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- DeepPotPT.h: add `using DeepPotBackend::computew[_mixed_type]` so the base charge_spin-aware overloads are not hidden by the non-charge_spin declarations. - DeepPotPT.cc: validate runtime/default charge_spin size against dim_chg_spin before building tensors (both nlist and no-nlist paths), avoiding out-of-range reads / opaque TorchScript failures. - deepmd.hpp: initialize dchgspin in DeepPot/DeepPotModelDevi constructors and guard dim_chg_spin() with assert(dp); add a shared validate_charge_spin() helper (rejects wrong size or charge_spin on a non-charge_spin model) and use it across all compute overloads. - test_deeppot_chg_spin_pt.cc: GTEST_SKIP when chg_spin.pth is absent instead of hard-failing SetUp. - gen_chg_spin.py: remove any stale chg_spin.pth before export and gate the parity check on a pth_generated flag, not os.path.exists. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Exercises parsing + wiring of the `charge_spin` keyword end-to-end against chg_spin.pt2 (DPA3, generated by gen_chg_spin.py): the default path (no keyword -> stored default_chg_spin) and an explicit `charge_spin 1.0 2.0`, both checked against the gen-script reference, plus a guard that an explicit charge_spin actually changes the energy. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
The .pth (pt) DeepEval does not expose get_dim_chg_spin (only the .pt2 / pt_expt one does), so the .pth parity assert raised AttributeError and crashed gen_chg_spin.py. Under test_cc_local.sh's `set -e`, that aborted the script before ctest, failing the whole C++ job. The energy/force parity checks already validate that charge_spin is threaded for the .pth backend. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Per review, charge_spin should follow the same convention as the other inputs (coord/fparam/aparam/spin): typed FPTYPE and grouped with the inputs (right after aparam), not a trailing double after the outputs. - c_api.h / c_api.cc: the version-3 entry points now take charge_spin as FPTYPE (double in the *3 funcs, float in the *f3 funcs), positioned after aparam. The internal *_variant helpers take const VALUETYPE* charge_spin and convert to std::vector<double> for the api_cc::DeepPot interface (which stores charge_spin as float64). api_cc and the compute backends are unchanged. - deepmd.hpp: the _DP_* helper templates take const FPTYPE* charge_spin after aparam; the public DeepPot/DeepPotModelDevi::compute overloads take const std::vector<VALUETYPE>& charge_spin; validate_charge_spin is now a template returning const FPTYPE*. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
|
LGTM, @njzjz plz also review this. |
- c_api.h: give full @PARAM documentation for all arguments of the new version-3 entry points (DP_DeepPot[ModelDevi]Compute[NList]{3,f3}), not only charge_spin. - deepmd.hpp / DeepPot.h: add the @PARAM[in] charge_spin doc to every charge_spin-aware compute / compute_mixed_type overload (DeepPot and DeepPotModelDevi). - doc/third-party/lammps-command.md: document the new `charge_spin` pair_style keyword (keyword list, syntax, description, example). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Summary by CodeRabbit
New Features
Tests