From 7f5540f5bf6c6aa7fa0802d807c54a506fe42098 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 9 Jun 2026 14:06:39 +0800 Subject: [PATCH 01/24] feat: add charge spin to cpp runtime --- source/api_cc/include/DeepPot.h | 148 +++++++++++++-- source/api_cc/include/DeepPotPTExpt.h | 100 ++++++++++- source/api_cc/src/DeepPot.cc | 138 +++++++++----- source/api_cc/src/DeepPotPTExpt.cc | 248 ++++++++++++++++++++++---- source/lmp/pair_base.cpp | 32 ++++ source/lmp/pair_base.h | 6 + source/lmp/pair_deepmd.cpp | 42 ++++- 7 files changed, 617 insertions(+), 97 deletions(-) diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 6a30eed7ea..93ea703915 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -198,6 +198,115 @@ class DeepPotBackend : public DeepBaseModelBackend { const std::vector& aparam, const bool atomic) = 0; /** @} */ + + /** + * @brief Get dimension of charge/spin condition inputs. + * Returns 0 for backends that do not support charge/spin conditioning. + **/ + virtual int dim_chg_spin() const { return 0; } + + // charge_spin-aware computew overloads. Default implementations call the + // existing pure-virtual overloads (ignoring charge_spin) so that backends + // that do not support charge/spin do not need any changes. DeepPotPTExpt + // overrides these to thread charge_spin through to the model. + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam, atomic); + } + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam, atomic); + } + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam, atomic); + } + virtual void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam, atomic); + } + virtual void computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes, + coord, atype, box, fparam, aparam, atomic); + } + virtual void computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + computew_mixed_type(ener, force, virial, atom_energy, atom_virial, nframes, + coord, atype, box, fparam, aparam, atomic); + } }; /** @@ -259,7 +368,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -268,7 +378,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** * @brief Evaluate the energy, force and virial by using this DP. @@ -304,7 +415,8 @@ class DeepPot : public DeepBaseModel { const InputNlist& inlist, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -316,7 +428,8 @@ class DeepPot : public DeepBaseModel { const InputNlist& inlist, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial @@ -351,7 +464,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -362,7 +476,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** @@ -404,7 +519,8 @@ class DeepPot : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute(std::vector& ener, std::vector& force, @@ -418,7 +534,8 @@ class DeepPot : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** * @brief Evaluate the energy, force, and virial with the mixed type @@ -453,7 +570,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute_mixed_type( std::vector& ener, @@ -464,7 +582,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ /** * @brief Evaluate the energy, force, and virial with the mixed type @@ -503,7 +622,8 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); template void compute_mixed_type( std::vector& ener, @@ -516,8 +636,12 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** @} */ + + int dim_chg_spin() const; + protected: std::shared_ptr dp; }; diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index 94631ab12a..dbe2e69620 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -77,6 +77,7 @@ class DeepPotPTExpt : public DeepPotBackend { const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); /** * @brief Evaluate without nlist (standalone — builds nlist, folds back). @@ -92,6 +93,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); public: @@ -115,6 +117,10 @@ class DeepPotPTExpt : public DeepPotBackend { assert(inited); return daparam; }; + int dim_chg_spin() const override { + assert(inited); + return dchgspin; + }; void get_type_map(std::string& type_map); bool is_aparam_nall() const { assert(inited); @@ -125,7 +131,7 @@ class DeepPotPTExpt : public DeepPotBackend { return has_default_fparam_; }; - // forward to template class + // forward to template class (no charge_spin — uses default_chg_spin_ fallback) void computew(std::vector& ener, std::vector& force, std::vector& virial, @@ -201,12 +207,94 @@ class DeepPotPTExpt : public DeepPotBackend { const std::vector& aparam, const bool atomic); + // charge_spin overloads — pass runtime charge/spin per call + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + private: bool inited; int ntypes; int dfparam; int daparam; - int dim_chg_spin; + int dchgspin; bool aparam_nall; bool has_default_fparam_; std::vector default_fparam_; @@ -255,6 +343,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); /** @@ -272,6 +361,7 @@ class DeepPotPTExpt : public DeepPotBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); /** @@ -282,6 +372,7 @@ class DeepPotPTExpt : public DeepPotBackend { * @param[in] mapping Mapping tensor. * @param[in] fparam Frame parameter tensor (or empty). * @param[in] aparam Atomic parameter tensor (or empty). + * @param[in] charge_spin Charge/spin tensor (or empty). * @return Vector of output tensors in sorted key order. */ std::vector run_model(const torch::Tensor& coord, @@ -289,13 +380,15 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& nlist, const torch::Tensor& mapping, const torch::Tensor& fparam, - const torch::Tensor& aparam); + const torch::Tensor& aparam, + const torch::Tensor& charge_spin); /** * @brief Run the with-comm .pt2 artifact with comm tensors appended. * * @param[in] base 4-6 base inputs (coord, atype, nlist, mapping, * fparam?, aparam?) — same as ``run_model``. + * @param[in] charge_spin Charge/spin tensor (or empty). * @param[in] comm_tensors 8 comm tensors in canonical positional * order: send_list, send_proc, recv_proc, send_num, * recv_num, communicator, nlocal, nghost. @@ -307,6 +400,7 @@ class DeepPotPTExpt : public DeepPotBackend { const torch::Tensor& mapping, const torch::Tensor& fparam, const torch::Tensor& aparam, + const torch::Tensor& charge_spin, const std::vector& comm_tensors); /** diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index ae543e1e31..8005cc469a 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -93,11 +93,12 @@ void DeepPot::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector dener_; std::vector datom_energy_, datom_virial_; dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, fparam_, aparam_, false); + datype_, dbox, fparam_, aparam_, charge_spin, false); dener = dener_[0]; } @@ -109,10 +110,11 @@ void DeepPot::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector datom_energy_, datom_virial_; dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, fparam_, aparam_, false); + datype_, dbox, fparam_, aparam_, charge_spin, false); } template void DeepPot::compute(ENERGYTYPE& dener, @@ -122,7 +124,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -131,7 +134,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -140,7 +144,8 @@ template void DeepPot::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -149,7 +154,8 @@ template void DeepPot::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, @@ -162,11 +168,13 @@ void DeepPot::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { std::vector dener_; std::vector datom_energy_, datom_virial_; dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, false); + datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, + charge_spin, false); dener = dener_[0]; } @@ -181,10 +189,12 @@ void DeepPot::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { std::vector datom_energy_, datom_virial_; dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, false); + datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, + charge_spin, false); } template void DeepPot::compute(ENERGYTYPE& dener, @@ -197,7 +207,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -209,7 +220,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -221,7 +233,8 @@ template void DeepPot::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -233,7 +246,8 @@ template void DeepPot::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, @@ -245,10 +259,11 @@ void DeepPot::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector dener_; dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, fparam_, aparam_, true); + datype_, dbox, fparam_, aparam_, charge_spin, true); dener = dener_[0]; } template @@ -261,9 +276,10 @@ void DeepPot::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, fparam_, aparam_, true); + datype_, dbox, fparam_, aparam_, charge_spin, true); } template void DeepPot::compute(ENERGYTYPE& dener, @@ -275,7 +291,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -286,7 +303,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -297,7 +315,8 @@ template void DeepPot::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -308,7 +327,8 @@ template void DeepPot::compute(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, @@ -323,10 +343,12 @@ void DeepPot::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { std::vector dener_; dp->computew(dener_, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, true); + datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, + charge_spin, true); dener = dener_[0]; } template @@ -342,9 +364,11 @@ void DeepPot::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam_, - const std::vector& aparam__) { + const std::vector& aparam__, + const std::vector& charge_spin) { dp->computew(dener, dforce_, dvirial, datom_energy_, datom_virial_, dcoord_, - datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, true); + datype_, dbox, nghost, lmp_list, ago, fparam_, aparam__, + charge_spin, true); } template void DeepPot::compute(ENERGYTYPE& dener, @@ -359,7 +383,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(ENERGYTYPE& dener, std::vector& dforce_, @@ -373,7 +398,8 @@ template void DeepPot::compute(ENERGYTYPE& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -387,7 +413,8 @@ template void DeepPot::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); template void DeepPot::compute(std::vector& dener, std::vector& dforce_, @@ -401,7 +428,8 @@ template void DeepPot::compute(std::vector& dener, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_); + const std::vector& aparam_, + const std::vector& charge_spin); // mixed type template @@ -413,12 +441,13 @@ void DeepPot::compute_mixed_type(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector dener_; std::vector datom_energy_, datom_virial_; dp->computew_mixed_type(dener_, dforce_, dvirial, datom_energy_, datom_virial_, nframes, dcoord_, datype_, dbox, - fparam_, aparam_, false); + fparam_, aparam_, charge_spin, false); dener = dener_[0]; } template @@ -430,11 +459,12 @@ void DeepPot::compute_mixed_type(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector datom_energy_, datom_virial_; dp->computew_mixed_type(dener, dforce_, dvirial, datom_energy_, datom_virial_, nframes, dcoord_, datype_, dbox, fparam_, aparam_, - false); + charge_spin, false); } template void DeepPot::compute_mixed_type( @@ -446,7 +476,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type( ENERGYTYPE& dener, @@ -457,7 +488,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type( std::vector& dener, @@ -468,7 +500,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type( std::vector& dener, @@ -479,7 +512,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type(ENERGYTYPE& dener, @@ -492,11 +526,12 @@ void DeepPot::compute_mixed_type(ENERGYTYPE& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { std::vector dener_; dp->computew_mixed_type(dener_, dforce_, dvirial, datom_energy_, datom_virial_, nframes, dcoord_, datype_, dbox, - fparam_, aparam_, true); + fparam_, aparam_, charge_spin, true); dener = dener_[0]; } template @@ -510,10 +545,11 @@ void DeepPot::compute_mixed_type(std::vector& dener, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam_, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { dp->computew_mixed_type(dener, dforce_, dvirial, datom_energy_, datom_virial_, nframes, dcoord_, datype_, dbox, fparam_, aparam_, - true); + charge_spin, true); } template void DeepPot::compute_mixed_type( @@ -527,7 +563,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type( ENERGYTYPE& dener, @@ -540,7 +577,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type( std::vector& dener, @@ -553,7 +591,8 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPot::compute_mixed_type( std::vector& dener, @@ -566,7 +605,10 @@ template void DeepPot::compute_mixed_type( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); + +int DeepPot::dim_chg_spin() const { return dp->dim_chg_spin(); } DeepPotModelDevi::DeepPotModelDevi() { inited = false; diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index c8c1bfcfad..c47a3ef28d 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -104,9 +104,9 @@ void DeepPotPTExpt::init(const std::string& model, : static_cast(metadata["type_map"].as_array().size()); dfparam = metadata["dim_fparam"].as_int(); daparam = metadata["dim_aparam"].as_int(); - dim_chg_spin = metadata.obj_val.count("dim_chg_spin") - ? metadata["dim_chg_spin"].as_int() - : 0; + dchgspin = metadata.obj_val.count("dim_chg_spin") + ? metadata["dim_chg_spin"].as_int() + : 0; aparam_nall = false; // pt_expt models use nloc for aparam if (metadata.obj_val.count("has_default_fparam")) { has_default_fparam_ = metadata["has_default_fparam"].as_bool(); @@ -132,7 +132,7 @@ void DeepPotPTExpt::init(const std::string& model, << std::endl; } } - default_chg_spin_ = read_default_chg_spin(metadata, dim_chg_spin); + default_chg_spin_ = read_default_chg_spin(metadata, dchgspin); if (metadata.obj_val.count("do_atomic_virial")) { do_atomic_virial = metadata["do_atomic_virial"].as_bool(); @@ -241,10 +241,10 @@ std::vector DeepPotPTExpt::run_model( const torch::Tensor& nlist, const torch::Tensor& mapping, const torch::Tensor& fparam, - const torch::Tensor& aparam) { - // Only include fparam/aparam if the model was exported with them. - // When fparam/aparam are None at export time, AOTInductor compiles - // the model with fewer inputs (e.g. 4 instead of 6). + const torch::Tensor& aparam, + const torch::Tensor& charge_spin) { + // Only include fparam/aparam/charge_spin if the model was exported with them. + // When they are None at export time, AOTInductor compiles with fewer inputs. std::vector inputs = {coord, atype, nlist, mapping}; if (dfparam > 0) { inputs.push_back(fparam); @@ -252,11 +252,7 @@ std::vector DeepPotPTExpt::run_model( if (daparam > 0) { inputs.push_back(aparam); } - if (dim_chg_spin > 0) { - auto charge_spin = torch::tensor(default_chg_spin_, coord.options()) - .view({1, dim_chg_spin}) - .expand({coord.size(0), dim_chg_spin}) - .contiguous(); + if (dchgspin > 0) { inputs.push_back(charge_spin); } return loader->run(inputs); @@ -269,6 +265,7 @@ std::vector DeepPotPTExpt::run_model_with_comm( const torch::Tensor& mapping, const torch::Tensor& fparam, const torch::Tensor& aparam, + const torch::Tensor& charge_spin, const std::vector& comm_tensors) { if (!with_comm_loader) { throw deepmd::deepmd_exception( @@ -293,11 +290,7 @@ std::vector DeepPotPTExpt::run_model_with_comm( if (daparam > 0) { inputs.push_back(aparam); } - if (dim_chg_spin > 0) { - auto charge_spin = torch::tensor(default_chg_spin_, coord.options()) - .view({1, dim_chg_spin}) - .expand({coord.size(0), dim_chg_spin}) - .contiguous(); + if (dchgspin > 0) { inputs.push_back(charge_spin); } for (const auto& t : comm_tensors) { @@ -334,6 +327,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { // Fail fast before allocating any tensors: refuse to run if the caller // asked for atomic virial but the .pt2 was exported without it. @@ -514,6 +508,32 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor = torch::zeros({0}, options).to(device); } + // 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) { + auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); + if (!charge_spin.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(charge_spin.data()), + {1, static_cast(charge_spin.size())}, + dbl_options) + .clone() + .to(device); + } else if (!default_chg_spin_.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(default_chg_spin_.data()), + {1, dchgspin}, dbl_options) + .clone() + .to(device); + } else { + throw deepmd::deepmd_exception( + "charge_spin is empty and no default_chg_spin is available in the " + ".pt2 metadata. Provide charge_spin explicitly or regenerate the " + "model with a default charge/spin value."); + } + } + // ``use_with_comm`` was computed earlier alongside the fail-fast // dispatch check. Use the with-comm artifact for the multi-rank case // (the regular artifact uses the mapping tensor to gather ghost @@ -545,10 +565,11 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, } flat_outputs = run_model_with_comm( coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, - fparam_tensor, aparam_tensor, comm_tensors); + fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); } else { flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor); + mapping_tensor, fparam_tensor, aparam_tensor, + charge_spin_tensor); } // Map flat outputs to internal keys @@ -624,6 +645,7 @@ template void DeepPotPTExpt::compute>( const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepPotPTExpt::compute>( std::vector& ener, @@ -639,6 +661,7 @@ template void DeepPotPTExpt::compute>( const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template @@ -652,6 +675,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { // Fail fast before allocating any tensors (same check as the nlist // overload — see its comment). @@ -668,7 +692,7 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (nframes > 1) { // Multi-frame: loop over frames and concatenate compute_nframes(ener, force, virial, atom_energy, atom_virial, nframes, - coord, atype, box, fparam, aparam, atomic); + coord, atype, box, fparam, aparam, charge_spin, atomic); return; } // The .pt2 model only contains forward_common_lower, which requires @@ -808,9 +832,36 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, aparam_tensor = torch::zeros({0}, options).to(device); } + // 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) { + auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); + if (!charge_spin.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(charge_spin.data()), + {1, static_cast(charge_spin.size())}, + dbl_options) + .clone() + .to(device); + } else if (!default_chg_spin_.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(default_chg_spin_.data()), + {1, dchgspin}, dbl_options) + .clone() + .to(device); + } else { + throw deepmd::deepmd_exception( + "charge_spin is empty and no default_chg_spin is available in the " + ".pt2 metadata. Provide charge_spin explicitly or regenerate the " + "model with a default charge/spin value."); + } + } + // 5. Run the .pt2 model - auto flat_outputs = run_model(coord_Tensor, atype_Tensor, nlist_tensor, - mapping_tensor, fparam_tensor, aparam_tensor); + auto flat_outputs = + run_model(coord_Tensor, atype_Tensor, nlist_tensor, mapping_tensor, + fparam_tensor, aparam_tensor, charge_spin_tensor); // 6. Map flat outputs to internal keys std::map output_map; @@ -874,6 +925,7 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { int natoms = atype.size(); int dap = aparam.empty() ? 0 : static_cast(aparam.size()) / nframes; @@ -907,10 +959,17 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, frame_aparam.assign(aparam.begin() + s_ff * s_dap, aparam.begin() + (s_ff + 1) * s_dap); } + std::vector frame_chg_spin; + if (!charge_spin.empty()) { + size_t s_dcsp = static_cast(dchgspin); + frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, + charge_spin.begin() + (s_ff + 1) * s_dcsp); + } std::vector frame_ener; std::vector frame_force, frame_virial, frame_ae, frame_av; compute(frame_ener, frame_force, frame_virial, frame_ae, frame_av, - frame_coord, atype, frame_box, frame_fparam, frame_aparam, atomic); + frame_coord, atype, frame_box, frame_fparam, frame_aparam, + frame_chg_spin, atomic); ener.insert(ener.end(), frame_ener.begin(), frame_ener.end()); force.insert(force.end(), frame_force.begin(), frame_force.end()); virial.insert(virial.end(), frame_virial.begin(), frame_virial.end()); @@ -932,6 +991,7 @@ template void DeepPotPTExpt::compute>( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepPotPTExpt::compute>( std::vector& ener, @@ -944,6 +1004,7 @@ template void DeepPotPTExpt::compute>( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); void DeepPotPTExpt::get_type_map(std::string& type_map_str) { @@ -953,7 +1014,7 @@ void DeepPotPTExpt::get_type_map(std::string& type_map_str) { } } -// forward to template method +// forward to template method (no runtime charge_spin — uses default_chg_spin_) void DeepPotPTExpt::computew(std::vector& ener, std::vector& force, std::vector& virial, @@ -967,7 +1028,7 @@ void DeepPotPTExpt::computew(std::vector& ener, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - fparam, aparam, atomic); + fparam, aparam, {}, atomic); }); } void DeepPotPTExpt::computew(std::vector& ener, @@ -983,7 +1044,7 @@ void DeepPotPTExpt::computew(std::vector& ener, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - fparam, aparam, atomic); + fparam, aparam, {}, atomic); }); } void DeepPotPTExpt::computew(std::vector& ener, @@ -1002,7 +1063,7 @@ void DeepPotPTExpt::computew(std::vector& ener, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - nghost, inlist, ago, fparam, aparam, atomic); + nghost, inlist, ago, fparam, aparam, {}, atomic); }); } void DeepPotPTExpt::computew(std::vector& ener, @@ -1021,7 +1082,7 @@ void DeepPotPTExpt::computew(std::vector& ener, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - nghost, inlist, ago, fparam, aparam, atomic); + nghost, inlist, ago, fparam, aparam, {}, atomic); }); } template @@ -1037,6 +1098,7 @@ void DeepPotPTExpt::compute_mixed_type_impl( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { // Mixed-type: atype has nframes * natoms elements. // Loop over frames, each with its own atype slice. @@ -1074,11 +1136,17 @@ void DeepPotPTExpt::compute_mixed_type_impl( frame_aparam.assign(aparam.begin() + s_ff * s_dap, aparam.begin() + (s_ff + 1) * s_dap); } + std::vector frame_chg_spin; + if (!charge_spin.empty()) { + size_t s_dcsp = static_cast(dchgspin); + frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, + charge_spin.begin() + (s_ff + 1) * s_dcsp); + } std::vector frame_ener; std::vector frame_force, frame_virial, frame_ae, frame_av; compute(frame_ener, frame_force, frame_virial, frame_ae, frame_av, frame_coord, frame_atype, frame_box, frame_fparam, frame_aparam, - atomic); + frame_chg_spin, atomic); ener.insert(ener.end(), frame_ener.begin(), frame_ener.end()); force.insert(force.end(), frame_force.begin(), frame_force.end()); virial.insert(virial.end(), frame_virial.begin(), frame_virial.end()); @@ -1103,7 +1171,121 @@ void DeepPotPTExpt::computew_mixed_type(std::vector& ener, const bool atomic) { translate_error([&] { compute_mixed_type_impl(ener, force, virial, atom_energy, atom_virial, - nframes, coord, atype, box, fparam, aparam, atomic); + nframes, coord, atype, box, fparam, aparam, {}, + atomic); + }); +} +void DeepPotPTExpt::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + compute_mixed_type_impl(ener, force, virial, atom_energy, atom_virial, + nframes, coord, atype, box, fparam, aparam, {}, + atomic); + }); +} + +// charge_spin overloads — pass runtime charge/spin value through to compute() +void DeepPotPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam, charge_spin, atomic); + }); +} +void DeepPotPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam, charge_spin, atomic); + }); +} +void DeepPotPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam, charge_spin, atomic); + }); +} +void DeepPotPTExpt::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam, charge_spin, atomic); + }); +} +void DeepPotPTExpt::computew_mixed_type(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute_mixed_type_impl(ener, force, virial, atom_energy, atom_virial, + nframes, coord, atype, box, fparam, aparam, + charge_spin, atomic); }); } void DeepPotPTExpt::computew_mixed_type(std::vector& ener, @@ -1117,10 +1299,12 @@ void DeepPotPTExpt::computew_mixed_type(std::vector& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { translate_error([&] { compute_mixed_type_impl(ener, force, virial, atom_energy, atom_virial, - nframes, coord, atype, box, fparam, aparam, atomic); + nframes, coord, atype, box, fparam, aparam, + charge_spin, atomic); }); } #endif diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index 63e5462590..f533acd941 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -164,6 +164,37 @@ void PairDeepBaseModel::make_fparam_from_compute(vector& fparam) { } } +void PairDeepBaseModel::make_charge_spin_from_compute( + vector& charge_spin) { + assert(do_compute_charge_spin); + + 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); + } + charge_spin.resize(dim_chg_spin); + + if (dim_chg_spin == 1) { + if (!(compute->invoked_flag & Compute::INVOKED_SCALAR)) { + compute->compute_scalar(); + compute->invoked_flag |= Compute::INVOKED_SCALAR; + } + charge_spin[0] = compute->scalar; + } else if (dim_chg_spin > 1) { + if (!(compute->invoked_flag & Compute::INVOKED_VECTOR)) { + compute->compute_vector(); + compute->invoked_flag |= Compute::INVOKED_VECTOR; + } + double* cvector = compute->vector; + for (int jj = 0; jj < dim_chg_spin; ++jj) { + charge_spin[jj] = cvector[jj]; + } + } +} + void PairDeepBaseModel::make_aparam_from_compute(vector& aparam) { assert(do_compute_aparam); @@ -339,6 +370,7 @@ PairDeepBaseModel::PairDeepBaseModel( do_ttm = false; do_compute_fparam = false; do_compute_aparam = false; + do_compute_charge_spin = false; single_model = false; multi_models_mod_devi = false; multi_models_no_mod_devi = false; diff --git a/source/lmp/pair_base.h b/source/lmp/pair_base.h index 1dd4b84041..22555cb07e 100644 --- a/source/lmp/pair_base.h +++ b/source/lmp/pair_base.h @@ -65,6 +65,7 @@ class PairDeepBaseModel : public Pair { std::string out_file; int dim_fparam; int dim_aparam; + int dim_chg_spin; int out_each; int out_rel; int out_rel_v; @@ -80,6 +81,7 @@ class PairDeepBaseModel : public Pair { std::map old_idx_map; std::vector fparam; std::vector aparam; + std::vector charge_spin; double eps; double eps_v; @@ -90,6 +92,10 @@ class PairDeepBaseModel : public Pair { bool do_compute_aparam; std::string compute_aparam_id; + void make_charge_spin_from_compute(std::vector& charge_spin); + bool do_compute_charge_spin; + std::string compute_charge_spin_id; + void make_ttm_fparam(std::vector& fparam); void make_ttm_aparam(std::vector& dparam); diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index fb27fa0ff6..193334ab2f 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -216,6 +216,10 @@ void PairDeepMD::compute(int eflag, int vflag) { make_fparam_from_compute(fparam); } + if (do_compute_charge_spin) { + make_charge_spin_from_compute(charge_spin); + } + // int ago = numb_models > 1 ? 0 : neighbor->ago; int ago = neighbor->ago; if (numb_models > 1) { @@ -249,7 +253,7 @@ void PairDeepMD::compute(int eflag, int vflag) { 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); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -260,7 +264,8 @@ void PairDeepMD::compute(int eflag, int vflag) { vector dvatom(nall * 9, 0); try { deep_pot.compute(dener, dforce, dvirial, deatom, dvatom, dcoord, - dtype, dbox, nghost, lmp_list, ago, fparam, daparam); + dtype, dbox, nghost, lmp_list, ago, fparam, daparam, + charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -534,6 +539,8 @@ static bool is_key(const string& input) { keys.push_back("aparam"); keys.push_back("fparam_from_compute"); keys.push_back("aparam_from_compute"); + keys.push_back("charge_spin"); + keys.push_back("charge_spin_from_compute"); keys.push_back("ttm"); keys.push_back("atomic"); keys.push_back("relative"); @@ -577,6 +584,7 @@ void PairDeepMD::settings(int narg, char** arg) { numb_types_spin = deep_pot.numb_types_spin(); dim_fparam = deep_pot.dim_fparam(); dim_aparam = deep_pot.dim_aparam(); + dim_chg_spin = deep_pot.dim_chg_spin(); } else { try { deep_pot.init(arg[0], get_node_rank(), get_file_content(arg[0])); @@ -590,6 +598,7 @@ void PairDeepMD::settings(int narg, char** arg) { 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(); 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()); @@ -604,6 +613,7 @@ void PairDeepMD::settings(int narg, char** arg) { eps = 0.; fparam.clear(); aparam.clear(); + charge_spin.clear(); while (iarg < narg) { if (!is_key(arg[iarg])) { error->all(FLERR, @@ -686,6 +696,29 @@ void PairDeepMD::settings(int narg, char** arg) { do_compute_aparam = true; compute_aparam_id = arg[iarg + 1]; iarg += 1 + 1; + } 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); + error->all(FLERR, tmp); + } + charge_spin.push_back(atof(arg[iarg + 1 + ii])); + } + iarg += 1 + dim_chg_spin; + } else if (string(arg[iarg]) == string("charge_spin_from_compute")) { + for (int ii = 0; ii < 1; ++ii) { + if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) { + error->all( + FLERR, + "invalid charge_spin_from_compute key: should be " + "charge_spin_from_compute compute_charge_spin_id(str)"); + } + } + do_compute_charge_spin = true; + compute_charge_spin_id = arg[iarg + 1]; + iarg += 1 + 1; } else if (string(arg[iarg]) == string("atomic")) { out_each = 1; iarg += 1; @@ -725,6 +758,11 @@ void PairDeepMD::settings(int narg, char** arg) { FLERR, "fparam and fparam_from_compute should NOT be set simultaneously"); } + if (do_compute_charge_spin && charge_spin.size() > 0) { + error->all(FLERR, + "charge_spin and charge_spin_from_compute should NOT be set " + "simultaneously"); + } if (comm->me == 0) { if (numb_models > 1 && out_freq > 0) { From 57df3853c1a82a708d77cd7754c55107d513bfac Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 9 Jun 2026 06:07:52 +0000 Subject: [PATCH 02/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_cc/include/DeepPotPTExpt.h | 3 ++- source/api_cc/src/DeepPotPTExpt.cc | 10 +++++----- source/lmp/pair_base.cpp | 3 +-- source/lmp/pair_deepmd.cpp | 7 +++---- 4 files changed, 11 insertions(+), 12 deletions(-) diff --git a/source/api_cc/include/DeepPotPTExpt.h b/source/api_cc/include/DeepPotPTExpt.h index dbe2e69620..461fc6f33c 100644 --- a/source/api_cc/include/DeepPotPTExpt.h +++ b/source/api_cc/include/DeepPotPTExpt.h @@ -131,7 +131,8 @@ class DeepPotPTExpt : public DeepPotBackend { return has_default_fparam_; }; - // forward to template class (no charge_spin — uses default_chg_spin_ fallback) + // forward to template class (no charge_spin — uses default_chg_spin_ + // fallback) void computew(std::vector& ener, std::vector& force, std::vector& virial, diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index c47a3ef28d..cae339a9c0 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -567,9 +567,9 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, charge_spin_tensor, comm_tensors); } else { - flat_outputs = run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor, - charge_spin_tensor); + flat_outputs = + run_model(coord_Tensor, atype_Tensor, firstneigh_tensor, mapping_tensor, + fparam_tensor, aparam_tensor, charge_spin_tensor); } // Map flat outputs to internal keys @@ -963,7 +963,7 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, if (!charge_spin.empty()) { size_t s_dcsp = static_cast(dchgspin); frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, - charge_spin.begin() + (s_ff + 1) * s_dcsp); + charge_spin.begin() + (s_ff + 1) * s_dcsp); } std::vector frame_ener; std::vector frame_force, frame_virial, frame_ae, frame_av; @@ -1140,7 +1140,7 @@ void DeepPotPTExpt::compute_mixed_type_impl( if (!charge_spin.empty()) { size_t s_dcsp = static_cast(dchgspin); frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, - charge_spin.begin() + (s_ff + 1) * s_dcsp); + charge_spin.begin() + (s_ff + 1) * s_dcsp); } std::vector frame_ener; std::vector frame_force, frame_virial, frame_ae, frame_av; diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index f533acd941..7350e70d6a 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -172,8 +172,7 @@ void PairDeepBaseModel::make_charge_spin_from_compute( Compute* compute = modify->compute[icompute]; if (!compute) { - error->all(FLERR, - "compute id is not found: " + compute_charge_spin_id); + error->all(FLERR, "compute id is not found: " + compute_charge_spin_id); } charge_spin.resize(dim_chg_spin); diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 193334ab2f..bdef6ee2c7 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -710,10 +710,9 @@ void PairDeepMD::settings(int narg, char** arg) { } else if (string(arg[iarg]) == string("charge_spin_from_compute")) { for (int ii = 0; ii < 1; ++ii) { if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) { - error->all( - FLERR, - "invalid charge_spin_from_compute key: should be " - "charge_spin_from_compute compute_charge_spin_id(str)"); + error->all(FLERR, + "invalid charge_spin_from_compute key: should be " + "charge_spin_from_compute compute_charge_spin_id(str)"); } } do_compute_charge_spin = true; From 46a0164eef864799f8330b2a545f96abbf043645 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 9 Jun 2026 14:43:01 +0800 Subject: [PATCH 03/24] fix: UT --- source/api_c/include/deepmd.hpp | 6 + .../tests/test_deeppot_chg_spin_ptexpt.cc | 255 ++++++++++++++++++ source/tests/infer/gen_chg_spin.py | 160 +++++++++++ source/tests/pt_expt/infer/test_deep_eval.py | 108 ++++++++ 4 files changed, 529 insertions(+) create mode 100644 source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc create mode 100644 source/tests/infer/gen_chg_spin.py diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index c3ca40b75f..35f91188e0 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1074,6 +1074,12 @@ class DeepPot : public DeepBaseModel { dpbase = (DP_DeepBaseModel*)dp; }; + /** + * @brief Get the dimension of the charge/spin embedding input. + * @return Always 0; charge_spin is not supported via the C-API path. + **/ + int dim_chg_spin() const { return 0; } + /** * @brief Evaluate the energy, force and virial by using this DP. * @param[out] ener The system energy. diff --git a/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc b/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc new file mode 100644 index 0000000000..4d448c1eea --- /dev/null +++ b/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc @@ -0,0 +1,255 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test C++ inference for pt_expt (.pt2) backend with charge_spin input (DPA3). +// Model generated by source/tests/infer/gen_chg_spin.py. +#include + +#include +#include +#include + +#include "DeepPot.h" +#include "DeepPotPTExpt.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +#undef EPSILON +#define EPSILON (std::is_same::value ? 1e-7 : 1e-4) + +namespace { +constexpr const char* kRefPath = "../../tests/infer/chg_spin.expected"; +constexpr const char* kModelPath = "../../tests/infer/chg_spin.pt2"; +} // namespace + +template +class TestInferDeepPotChgSpinPtExpt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + // charge_spin is always double regardless of VALUETYPE + std::vector charge_spin_explicit = {0.5, 0.8}; + + std::vector expected_e_explicit; + std::vector expected_f_explicit; + std::vector expected_v_explicit; + std::vector expected_e_default; + std::vector expected_f_default; + std::vector expected_v_default; + int natoms; + double expected_tot_e_explicit; + double expected_tot_e_default; + std::vector expected_tot_v_explicit; + std::vector expected_tot_v_default; + + static deepmd::DeepPot dp; + + static void SetUpTestSuite() { +#if defined(BUILD_PYTORCH) && BUILD_PT_EXPT + dp.init(kModelPath); +#endif + } + + void SetUp() override { +#if !defined(BUILD_PYTORCH) || !BUILD_PT_EXPT + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + expected_e_explicit = ref.get("explicit", "expected_e"); + expected_f_explicit = ref.get("explicit", "expected_f"); + expected_v_explicit = ref.get("explicit", "expected_v"); + expected_e_default = ref.get("default", "expected_e"); + expected_f_default = ref.get("default", "expected_f"); + expected_v_default = ref.get("default", "expected_v"); + + natoms = expected_e_explicit.size(); + EXPECT_EQ(natoms * 3, static_cast(expected_f_explicit.size())); + EXPECT_EQ(natoms * 9, static_cast(expected_v_explicit.size())); + EXPECT_EQ(natoms, static_cast(expected_e_default.size())); + + expected_tot_e_explicit = 0.; + expected_tot_e_default = 0.; + expected_tot_v_explicit.assign(9, 0.); + expected_tot_v_default.assign(9, 0.); + for (int ii = 0; ii < natoms; ++ii) { + expected_tot_e_explicit += expected_e_explicit[ii]; + expected_tot_e_default += expected_e_default[ii]; + } + for (int ii = 0; ii < natoms; ++ii) { + for (int dd = 0; dd < 9; ++dd) { + expected_tot_v_explicit[dd] += expected_v_explicit[ii * 9 + dd]; + expected_tot_v_default[dd] += expected_v_default[ii * 9 + dd]; + } + } + } + + void TearDown() override {} + + static void TearDownTestSuite() { dp = deepmd::DeepPot(); } +}; + +template +deepmd::DeepPot TestInferDeepPotChgSpinPtExpt::dp; + +TYPED_TEST_SUITE(TestInferDeepPotChgSpinPtExpt, ValueTypes); + +TYPED_TEST(TestInferDeepPotChgSpinPtExpt, dim_chg_spin) { + deepmd::DeepPot& dp = this->dp; + EXPECT_EQ(dp.dim_chg_spin(), 2); +} + +TYPED_TEST(TestInferDeepPotChgSpinPtExpt, cpu_build_nlist_explicit) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& charge_spin = this->charge_spin_explicit; + std::vector& expected_f = this->expected_f_explicit; + std::vector& expected_tot_v = this->expected_tot_v_explicit; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_explicit; + deepmd::DeepPot& dp = this->dp; + + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box, {}, {}, charge_spin); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotChgSpinPtExpt, cpu_build_nlist_default) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f_default; + std::vector& expected_tot_v = this->expected_tot_v_default; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_default; + deepmd::DeepPot& dp = this->dp; + + double ener; + std::vector force, virial; + // Empty charge_spin triggers default_chg_spin_ fallback ([0.0, 1.0]) + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotChgSpinPtExpt, cpu_lmp_nlist_explicit) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& charge_spin = this->charge_spin_explicit; + std::vector& expected_f = this->expected_f_explicit; + std::vector& expected_tot_v = this->expected_tot_v_explicit; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_explicit; + deepmd::DeepPot& dp = this->dp; + + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector> nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force_, virial; + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0, {}, {}, charge_spin); + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } + + // Verify ago=1 (cached nlist) gives same result + ener = 0.; + std::fill(force_.begin(), force_.end(), 0.0); + std::fill(virial.begin(), virial.end(), 0.0); + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 1, {}, {}, charge_spin); + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotChgSpinPtExpt, cpu_lmp_nlist_default) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f_default; + std::vector& expected_tot_v = this->expected_tot_v_default; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_default; + deepmd::DeepPot& dp = this->dp; + + float rc = dp.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector> nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force_, virial; + // Empty charge_spin triggers default_chg_spin_ fallback ([0.0, 1.0]) + dp.compute(ener, force_, virial, coord_cpy, atype_cpy, box, nall - nloc, + inlist, 0); + std::vector force; + _fold_back(force, force_, mapping, nloc, nall, 3); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} diff --git a/source/tests/infer/gen_chg_spin.py b/source/tests/infer/gen_chg_spin.py new file mode 100644 index 0000000000..d464e79828 --- /dev/null +++ b/source/tests/infer/gen_chg_spin.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate chg_spin.pt2 test model and reference values. + +Creates a DPA3 model with add_chg_spin_ebd=True and default_chg_spin=[0.0, 1.0], +exports to .pt2, and writes chg_spin.expected with two sections: + [default] -- eval with no charge_spin (uses stored default [0.0, 1.0]) + [explicit] -- eval with charge_spin=[0.5, 0.8] +""" + +import copy +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..")) + +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, + write_expected_ref, +) + + +def main(): + from deepmd.dpmodel.model.model import ( + get_model, + ) + + ensure_inductor_compiler() + + config = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 8, + "e_dim": 5, + "a_dim": 4, + "nlayers": 2, + "e_rcut": 6.0, + "e_rcut_smth": 2.0, + "e_sel": 30, + "a_rcut": 4.0, + "a_rcut_smth": 2.0, + "a_sel": 20, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual_init": "const", + "smooth_edge_update": True, + }, + "concat_output_tebd": True, + "precision": "float64", + "add_chg_spin_ebd": True, + "default_chg_spin": [0.0, 1.0], + "seed": 1, + }, + "fitting_net": {"neuron": [5, 5, 5], "resnet_dt": True, "seed": 1}, + } + + model = get_model(copy.deepcopy(config)) + model_dict = model.serialize() + + data = { + "model": model_dict, + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file as pt_expt_deserialize_to_file, + ) + + load_custom_ops() + + base_dir = os.path.dirname(__file__) + pt2_path = os.path.join(base_dir, "chg_spin.pt2") + print(f"Exporting to {pt2_path} ...") # noqa: T201 + pt_expt_deserialize_to_file(pt2_path, copy.deepcopy(data), do_atomic_virial=True) + print("Export done.") # noqa: T201 + + from deepmd.infer import ( + DeepPot, + ) + + dp = DeepPot(pt2_path) + + dim = dp.deep_eval.get_dim_chg_spin() + assert dim == 2, f"Expected dim_chg_spin == 2, got {dim}" + print(f"dim_chg_spin = {dim}") # noqa: T201 + + coord = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 0.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=np.float64, + ) + atype = [0, 1, 1, 0, 1, 1] + box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], dtype=np.float64) + + # Default charge_spin: no argument → uses stored [0.0, 1.0] + e_def, f_def, v_def, ae_def, av_def = dp.eval(coord, box, atype, atomic=True) + print(f"\n// Default charge_spin total energy: {e_def[0, 0]:.18e}") # noqa: T201 + + # Explicit charge_spin = [0.5, 0.8] + charge_spin_exp = np.array([[0.5, 0.8]], dtype=np.float64) + e_exp, f_exp, v_exp, ae_exp, av_exp = dp.eval( + coord, box, atype, atomic=True, charge_spin=charge_spin_exp + ) + print(f"\n// Explicit charge_spin [0.5, 0.8] total energy: {e_exp[0, 0]:.18e}") # noqa: T201 + + # Sanity: different charge_spin should yield different outputs + assert not np.allclose(e_def, e_exp), ( + "Default and explicit charge_spin gave identical energy — charge_spin may be ignored" + ) + + ref_path = os.path.join(base_dir, "chg_spin.expected") + write_expected_ref( + ref_path, + sections={ + "default": { + "expected_e": ae_def[0, :, 0], + "expected_f": f_def[0], + "expected_v": av_def[0], + }, + "explicit": { + "expected_e": ae_exp[0, :, 0], + "expected_f": f_exp[0], + "expected_v": av_exp[0], + }, + }, + source_script="source/tests/infer/gen_chg_spin.py", + ) + print(f"Wrote {ref_path}") # noqa: T201 + print("\nDone!") # noqa: T201 + + +if __name__ == "__main__": + main() diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index a921a94a7f..620b594de8 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -2297,5 +2297,113 @@ def test_eval_descriptor_vesin_matches_native(self) -> None: np.testing.assert_allclose(d_native, d_vesin, rtol=1e-10, atol=1e-10) +class TestDeepEvalEnerChgSpinPt2(unittest.TestCase): + """Test .pt2 inference with charge_spin (DPA3 with add_chg_spin_ebd=True).""" + + @classmethod + def setUpClass(cls) -> None: + import copy + + from deepmd.dpmodel.model.model import get_model as dp_get_model + + cls.nt = 2 + cls.default_chg_spin = [0.5, 0.8] + + config = { + "type_map": ["O", "H"], + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 8, + "e_dim": 5, + "a_dim": 4, + "nlayers": 2, + "e_rcut": 6.0, + "e_rcut_smth": 2.0, + "e_sel": 30, + "a_rcut": 4.0, + "a_rcut_smth": 2.0, + "a_sel": 20, + "axis_neuron": 4, + "update_angle": True, + "update_style": "res_residual", + "update_residual_init": "const", + "smooth_edge_update": True, + }, + "concat_output_tebd": True, + "precision": "float64", + "add_chg_spin_ebd": True, + "default_chg_spin": cls.default_chg_spin, + "seed": GLOBAL_SEED, + }, + "fitting_net": {"neuron": [5, 5, 5], "resnet_dt": True, "seed": GLOBAL_SEED}, + } + + model = dp_get_model(copy.deepcopy(config)) + cls.model_data = { + "model": model.serialize(), + "model_def_script": config, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + + cls.tmpfile = tempfile.NamedTemporaryFile(suffix=".pt2", delete=False) + cls.tmpfile.close() + torch.set_default_device(None) + try: + deserialize_to_file(cls.tmpfile.name, cls.model_data, do_atomic_virial=True) + finally: + torch.set_default_device("cuda:9999999") + + cls.dp = DeepPot(cls.tmpfile.name) + + @classmethod + def tearDownClass(cls) -> None: + import os + + os.unlink(cls.tmpfile.name) + + def test_dim_chg_spin(self) -> None: + self.assertEqual(self.dp.deep_eval.get_dim_chg_spin(), 2) + + def test_charge_spin_takes_effect(self) -> None: + """Different charge_spin values must produce different outputs.""" + rng = np.random.default_rng(GLOBAL_SEED) + natoms = 5 + coords = rng.random((1, natoms, 3)) * 8.0 + cells = np.eye(3).reshape(1, 9) * 10.0 + atom_types = np.array([i % self.nt for i in range(natoms)], dtype=np.int32) + + e0, f0, v0 = self.dp.eval( + coords, cells, atom_types, charge_spin=np.array([[0.0, 0.0]]) + ) + e1, f1, v1 = self.dp.eval( + coords, cells, atom_types, charge_spin=np.array([[1.0, 2.0]]) + ) + + assert not np.allclose(e0, e1), ( + "Changing charge_spin did not change output — charge_spin may be ignored" + ) + + def test_default_matches_explicit_default(self) -> None: + """Eval without charge_spin should use stored default and match explicit.""" + rng = np.random.default_rng(GLOBAL_SEED) + natoms = 5 + coords = rng.random((1, natoms, 3)) * 8.0 + cells = np.eye(3).reshape(1, 9) * 10.0 + atom_types = np.array([i % self.nt for i in range(natoms)], dtype=np.int32) + + e_no, f_no, v_no = self.dp.eval(coords, cells, atom_types) + e_ex, f_ex, v_ex = self.dp.eval( + coords, cells, atom_types, + charge_spin=np.array([self.default_chg_spin]), + ) + + np.testing.assert_allclose(e_no, e_ex, atol=1e-10) + np.testing.assert_allclose(f_no, f_ex, atol=1e-10) + np.testing.assert_allclose(v_no, v_ex, atol=1e-10) + + if __name__ == "__main__": unittest.main() From 982c71136aa5937640c52e5de48a904841479508 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 9 Jun 2026 06:44:31 +0000 Subject: [PATCH 04/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/pt_expt/infer/test_deep_eval.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/source/tests/pt_expt/infer/test_deep_eval.py b/source/tests/pt_expt/infer/test_deep_eval.py index 620b594de8..f831b365f6 100644 --- a/source/tests/pt_expt/infer/test_deep_eval.py +++ b/source/tests/pt_expt/infer/test_deep_eval.py @@ -2336,7 +2336,11 @@ def setUpClass(cls) -> None: "default_chg_spin": cls.default_chg_spin, "seed": GLOBAL_SEED, }, - "fitting_net": {"neuron": [5, 5, 5], "resnet_dt": True, "seed": GLOBAL_SEED}, + "fitting_net": { + "neuron": [5, 5, 5], + "resnet_dt": True, + "seed": GLOBAL_SEED, + }, } model = dp_get_model(copy.deepcopy(config)) @@ -2396,7 +2400,9 @@ def test_default_matches_explicit_default(self) -> None: e_no, f_no, v_no = self.dp.eval(coords, cells, atom_types) e_ex, f_ex, v_ex = self.dp.eval( - coords, cells, atom_types, + coords, + cells, + atom_types, charge_spin=np.array([self.default_chg_spin]), ) From 7ae6be583c12661e3b51d8e3091395dcde527cab Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 9 Jun 2026 15:10:25 +0800 Subject: [PATCH 05/24] fix: param passing --- source/api_c/include/deepmd.hpp | 16 ++++++++++++---- source/lmp/pair_base.cpp | 3 +++ 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 35f91188e0..d6c6a3f76a 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1109,7 +1109,9 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1170,7 +1172,9 @@ class DeepPot : public DeepBaseModel { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1239,7 +1243,9 @@ class DeepPot : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1309,7 +1315,9 @@ class DeepPot : public DeepBaseModel { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index 7350e70d6a..bc02a1f473 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -367,6 +367,9 @@ PairDeepBaseModel::PairDeepBaseModel( eps_v = 0.; scale = NULL; do_ttm = false; + dim_fparam = 0; + dim_aparam = 0; + dim_chg_spin = 0; do_compute_fparam = false; do_compute_aparam = false; do_compute_charge_spin = false; From d469444acef2308c8ee7f30c0676665af5533f69 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 9 Jun 2026 07:11:35 +0000 Subject: [PATCH 06/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_c/include/deepmd.hpp | 96 ++++++++++++++++----------------- 1 file changed, 46 insertions(+), 50 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index d6c6a3f76a..4669e2b0e3 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1101,16 +1101,15 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute( - ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; @@ -1162,18 +1161,17 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute( - ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - std::vector& atom_energy, - std::vector& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; @@ -1232,19 +1230,18 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute( - ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; @@ -1302,21 +1299,20 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute( - ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - std::vector& atom_energy, - std::vector& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; From dac5b4bb6cd0c3a71b8885a695d13ef723f61a18 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 9 Jun 2026 16:10:40 +0800 Subject: [PATCH 07/24] fix: UT --- source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc | 6 ++++-- source/install/test_cc_local.sh | 3 +++ source/tests/infer/gen_chg_spin.py | 7 ++++--- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc b/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc index 4d448c1eea..c30a17995b 100644 --- a/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc @@ -29,8 +29,10 @@ class TestInferDeepPotChgSpinPtExpt : public ::testing::Test { 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; std::vector atype = {0, 1, 1, 0, 1, 1}; std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; - // charge_spin is always double regardless of VALUETYPE - std::vector charge_spin_explicit = {0.5, 0.8}; + // charge_spin is always double regardless of VALUETYPE. + // [1.0, 2.0] maps to charge idx 101, spin idx 2 — both differ from the + // model's default_chg_spin [0.0, 1.0] (charge idx 100, spin idx 1). + std::vector charge_spin_explicit = {1.0, 2.0}; std::vector expected_e_explicit; std::vector expected_f_explicit; diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 5ddbf0eecc..d5b2f3afb0 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -87,9 +87,12 @@ else: PID5=$! env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_model_devi.py & PID6=$! + env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_chg_spin.py & + PID9=$! wait $PID4 wait $PID5 wait $PID6 + wait $PID9 env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_spin.py & PID7=$! diff --git a/source/tests/infer/gen_chg_spin.py b/source/tests/infer/gen_chg_spin.py index d464e79828..e849c78d78 100644 --- a/source/tests/infer/gen_chg_spin.py +++ b/source/tests/infer/gen_chg_spin.py @@ -123,12 +123,13 @@ def main(): e_def, f_def, v_def, ae_def, av_def = dp.eval(coord, box, atype, atomic=True) print(f"\n// Default charge_spin total energy: {e_def[0, 0]:.18e}") # noqa: T201 - # Explicit charge_spin = [0.5, 0.8] - charge_spin_exp = np.array([[0.5, 0.8]], dtype=np.float64) + # Explicit charge_spin = [1.0, 2.0] (charge idx 101, spin idx 2 — both + # differ from the default [0.0, 1.0] which maps to charge idx 100, spin 1) + charge_spin_exp = np.array([[1.0, 2.0]], dtype=np.float64) e_exp, f_exp, v_exp, ae_exp, av_exp = dp.eval( coord, box, atype, atomic=True, charge_spin=charge_spin_exp ) - print(f"\n// Explicit charge_spin [0.5, 0.8] total energy: {e_exp[0, 0]:.18e}") # noqa: T201 + print(f"\n// Explicit charge_spin [1.0, 2.0] total energy: {e_exp[0, 0]:.18e}") # noqa: T201 # Sanity: different charge_spin should yield different outputs assert not np.allclose(e_def, e_exp), ( From 1b8e8b53c4e6453e2820df128d4aa6bdf0db2900 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 9 Jun 2026 17:11:33 +0800 Subject: [PATCH 08/24] fix: export charge spin dim --- deepmd/pt_expt/utils/serialization.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/deepmd/pt_expt/utils/serialization.py b/deepmd/pt_expt/utils/serialization.py index 88bd839454..9c03783574 100644 --- a/deepmd/pt_expt/utils/serialization.py +++ b/deepmd/pt_expt/utils/serialization.py @@ -460,6 +460,9 @@ def _collect_metadata(model: torch.nn.Module, is_spin: bool = False) -> dict: "nnei": sum(model.get_sel()), "dim_fparam": model.get_dim_fparam(), "dim_aparam": model.get_dim_aparam(), + "dim_chg_spin": ( + model.get_dim_chg_spin() if hasattr(model, "get_dim_chg_spin") else 0 + ), "mixed_types": model.mixed_types(), "has_default_fparam": model.has_default_fparam(), "default_fparam": model.get_default_fparam(), From f24390b88915d4b08675eb54c033faf9a1d1ad61 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 9 Jun 2026 18:10:20 +0800 Subject: [PATCH 09/24] fix: ghost mapping --- source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc b/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc index c30a17995b..f9db1a0f2b 100644 --- a/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc +++ b/source/api_cc/tests/test_deeppot_chg_spin_ptexpt.cc @@ -181,6 +181,9 @@ TYPED_TEST(TestInferDeepPotChgSpinPtExpt, cpu_lmp_nlist_explicit) { std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); convert_nlist(inlist, nlist_data); + // DPA3 repflows uses mapping; pass it so the C++ backend can + // correctly map ghost atoms back to local atoms. + inlist.mapping = mapping.data(); double ener; std::vector force_, virial; @@ -236,6 +239,9 @@ TYPED_TEST(TestInferDeepPotChgSpinPtExpt, cpu_lmp_nlist_default) { std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); convert_nlist(inlist, nlist_data); + // DPA3 repflows uses mapping; pass it so the C++ backend can + // correctly map ghost atoms back to local atoms. + inlist.mapping = mapping.data(); double ener; std::vector force_, virial; From 738f2e484016d41147326deec217b08d6741c7dc Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 10 Jun 2026 15:43:15 +0800 Subject: [PATCH 10/24] fix: address comments --- source/api_c/include/deepmd.hpp | 20 ++++++++++++--- source/api_cc/include/DeepPot.h | 12 ++++++--- source/api_cc/src/DeepPot.cc | 45 +++++++++++++++++++++------------ source/lmp/pair_deepmd.cpp | 4 +-- 4 files changed, 55 insertions(+), 26 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 4669e2b0e3..ce81c9584a 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -2147,7 +2147,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + // charge_spin is not supported via the C-API model-deviation path. + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2228,7 +2231,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + // charge_spin is not supported via the C-API model-deviation path. + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2326,7 +2332,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + // charge_spin is not supported via the C-API model-deviation path. + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2416,7 +2425,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()) { + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { + // charge_spin is not supported via the C-API model-deviation path. + (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 93ea703915..1f01396509 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -704,7 +704,8 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial @@ -739,7 +740,8 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& atype, const std::vector& box, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Evaluate the energy, force and virial by using these DP models. @@ -775,7 +777,8 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); /** * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial *by using these DP models. @@ -815,7 +818,8 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const InputNlist& lmp_list, const int& ago, const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector()); + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()); protected: std::vector> dps; diff --git a/source/api_cc/src/DeepPot.cc b/source/api_cc/src/DeepPot.cc index 8005cc469a..30c64ab83f 100644 --- a/source/api_cc/src/DeepPot.cc +++ b/source/api_cc/src/DeepPot.cc @@ -658,7 +658,8 @@ void DeepPotModelDevi::compute(std::vector& all_energy, const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { // without nlist if (numb_models == 0) { return; @@ -668,7 +669,7 @@ void DeepPotModelDevi::compute(std::vector& all_energy, all_virial.resize(numb_models); for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], dcoord_, - datype_, dbox, fparam, aparam_); + datype_, dbox, fparam, aparam_, charge_spin); } } @@ -680,7 +681,8 @@ template void DeepPotModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute( std::vector& all_energy, @@ -690,7 +692,8 @@ template void DeepPotModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute( @@ -703,7 +706,8 @@ void DeepPotModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { if (numb_models == 0) { return; } @@ -715,7 +719,7 @@ void DeepPotModelDevi::compute( for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], dcoord_, datype_, - dbox, fparam, aparam_); + dbox, fparam, aparam_, charge_spin); } } @@ -729,7 +733,8 @@ template void DeepPotModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute( std::vector& all_energy, @@ -741,7 +746,8 @@ template void DeepPotModelDevi::compute( const std::vector& datype_, const std::vector& dbox, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute(std::vector& all_energy, @@ -754,7 +760,8 @@ void DeepPotModelDevi::compute(std::vector& all_energy, const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { if (numb_models == 0) { return; } @@ -763,7 +770,8 @@ void DeepPotModelDevi::compute(std::vector& all_energy, all_virial.resize(numb_models); for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], dcoord_, - datype_, dbox, nghost, lmp_list, ago, fparam, aparam_); + datype_, dbox, nghost, lmp_list, ago, fparam, aparam_, + charge_spin); } } @@ -778,7 +786,8 @@ template void DeepPotModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute( std::vector& all_energy, @@ -791,7 +800,8 @@ template void DeepPotModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute( @@ -807,7 +817,8 @@ void DeepPotModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_) { + const std::vector& aparam_, + const std::vector& charge_spin) { if (numb_models == 0) { return; } @@ -819,7 +830,7 @@ void DeepPotModelDevi::compute( for (unsigned ii = 0; ii < numb_models; ++ii) { dps[ii]->compute(all_energy[ii], all_force[ii], all_virial[ii], all_atom_energy[ii], all_atom_virial[ii], dcoord_, datype_, - dbox, nghost, lmp_list, ago, fparam, aparam_); + dbox, nghost, lmp_list, ago, fparam, aparam_, charge_spin); } } @@ -836,7 +847,8 @@ template void DeepPotModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); template void DeepPotModelDevi::compute( std::vector& all_energy, @@ -851,4 +863,5 @@ template void DeepPotModelDevi::compute( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam); + const std::vector& aparam, + const std::vector& charge_spin); diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index bdef6ee2c7..3446a71048 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -317,7 +317,7 @@ void PairDeepMD::compute(int eflag, int vflag) { try { deep_pot_model_devi.compute(all_energy, all_force, all_virial, dcoord, dtype, dbox, nghost, lmp_list, ago, - fparam, daparam); + fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } @@ -326,7 +326,7 @@ void PairDeepMD::compute(int eflag, int vflag) { deep_pot_model_devi.compute(all_energy, all_force, all_virial, all_atom_energy, all_atom_virial, dcoord, dtype, dbox, nghost, lmp_list, ago, - fparam, daparam); + fparam, daparam, charge_spin); } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } From 66c471eda0ad1c30f378b3b2c7355696e3aea73c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 07:44:02 +0000 Subject: [PATCH 11/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_c/include/deepmd.hpp | 96 ++++++++++++++++----------------- 1 file changed, 46 insertions(+), 50 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index ce81c9584a..2eee017c28 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -2139,16 +2139,15 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute( - std::vector& ener, - std::vector>& force, - std::vector>& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(std::vector& ener, + std::vector>& force, + std::vector>& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. (void)charge_spin; unsigned int natoms = atype.size(); @@ -2221,18 +2220,17 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute( - std::vector& ener, - std::vector>& force, - std::vector>& virial, - std::vector>& atom_energy, - std::vector>& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(std::vector& ener, + std::vector>& force, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. (void)charge_spin; unsigned int natoms = atype.size(); @@ -2321,19 +2319,18 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute( - std::vector& ener, - std::vector>& force, - std::vector>& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(std::vector& ener, + std::vector>& force, + std::vector>& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. (void)charge_spin; unsigned int natoms = atype.size(); @@ -2412,21 +2409,20 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute( - std::vector& ener, - std::vector>& force, - std::vector>& virial, - std::vector>& atom_energy, - std::vector>& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute(std::vector& ener, + std::vector>& force, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. (void)charge_spin; unsigned int natoms = atype.size(); From 5c02d994f54825f6fe40e6603a43141d1ac74576 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 10 Jun 2026 17:48:08 +0800 Subject: [PATCH 12/24] feat: add charge_spin to the .pth (TorchScript) cpp backend 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 --- source/api_cc/include/DeepPotPT.h | 69 +++++ source/api_cc/src/DeepPotPT.cc | 204 ++++++++++++- .../api_cc/tests/test_deeppot_chg_spin_pt.cc | 280 ++++++++++++++++++ source/tests/infer/gen_chg_spin.py | 77 ++++- 4 files changed, 611 insertions(+), 19 deletions(-) create mode 100644 source/api_cc/tests/test_deeppot_chg_spin_pt.cc diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h index 14b545dce4..cd57fc79ac 100644 --- a/source/api_cc/include/DeepPotPT.h +++ b/source/api_cc/include/DeepPotPT.h @@ -73,6 +73,7 @@ class DeepPotPT : public DeepPotBackend { const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); /** * @brief Evaluate the energy, force, virial, atomic energy, and atomic virial @@ -114,6 +115,7 @@ class DeepPotPT : public DeepPotBackend { const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); /** * @brief Evaluate the energy, force, and virial with the mixed type @@ -229,6 +231,15 @@ class DeepPotPT : public DeepPotBackend { assert(inited); return daparam; }; + /** + * @brief Get the dimension of the charge/spin input. + * @return The dimension of the charge/spin input (0 if the model has no + *charge/spin embedding). + **/ + int dim_chg_spin() const override { + assert(inited); + return dchgspin; + }; /** * @brief Get the type map (element name of the atom types) of this model. * @param[out] type_map The type map of this model. @@ -329,6 +340,62 @@ class DeepPotPT : public DeepPotBackend { const std::vector& aparam, const bool atomic); + // charge_spin overloads — pass runtime charge/spin per call + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + void computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) override; + private: int num_intra_nthreads, num_inter_nthreads; bool inited; @@ -336,8 +403,10 @@ class DeepPotPT : public DeepPotBackend { int ntypes_spin; int dfparam; int daparam; + int dchgspin; bool aparam_nall; bool has_default_fparam_; + std::vector default_chg_spin_; // copy neighbor list info from host torch::jit::script::Module module; double rcut; diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index d69dbb8f82..a725da2439 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -135,6 +135,24 @@ void DeepPotPT::init(const std::string& model, } else { has_default_fparam_ = false; } + // Charge/spin embedding (e.g. DPA3 add_chg_spin_ebd). Guarded with + // find_method so .pth models exported before charge_spin support + // (which lack these jit-exported methods) keep loading with dchgspin=0. + if (module.find_method("get_dim_chg_spin")) { + dchgspin = module.run_method("get_dim_chg_spin").toInt(); + } else { + dchgspin = 0; + } + default_chg_spin_.clear(); + if (dchgspin > 0 && module.find_method("get_default_chg_spin")) { + auto cs = module.run_method("get_default_chg_spin"); + if (!cs.isNone()) { + torch::Tensor cs_t = + cs.toTensor().to(torch::kFloat64).to(torch::kCPU).view({-1}); + default_chg_spin_.assign(cs_t.data_ptr(), + cs_t.data_ptr() + cs_t.numel()); + } + } inited = true; } @@ -163,6 +181,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { torch::Device device(torch::kCUDA, gpu_id); if (!gpu_enabled) { @@ -242,18 +261,63 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, options) .to(device); } - auto outputs = - (do_message_passing) - ? module - .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, - firstneigh_tensor, mapping_tensor, fparam_tensor, - aparam_tensor, do_atom_virial_tensor, comm_dict) - .toGenericDict() - : module - .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, - firstneigh_tensor, mapping_tensor, fparam_tensor, - aparam_tensor, do_atom_virial_tensor) - .toGenericDict(); + // Build charge_spin tensor (always float64): use the runtime value when + // provided, otherwise fall back to the model's stored default_chg_spin. + // Only threaded into forward_lower when the model has a charge/spin + // embedding (dchgspin > 0), so .pth models without it are unaffected. + c10::optional charge_spin_tensor; + if (dchgspin > 0) { + auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); + if (!charge_spin.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(charge_spin.data()), + {1, static_cast(charge_spin.size())}, + dbl_options) + .clone() + .to(device); + } else if (!default_chg_spin_.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(default_chg_spin_.data()), + {1, dchgspin}, dbl_options) + .clone() + .to(device); + } else { + throw deepmd::deepmd_exception( + "charge_spin is empty and no default_chg_spin is available in the " + "model. Provide charge_spin explicitly or regenerate the model with " + "a default charge/spin value."); + } + } + c10::IValue outputs_ival; + if (dchgspin > 0) { + // charge_spin model. DPA3 (the only charge/spin descriptor) always uses + // message passing, so comm_dict is populated; for the non-message-passing + // edge case pass a real None for comm_dict (an empty Dict would wrongly + // flip the descriptor into parallel mode). + if (do_message_passing) { + outputs_ival = module.run_method( + "forward_lower", coord_wrapped_Tensor, atype_Tensor, firstneigh_tensor, + mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor, + comm_dict, charge_spin_tensor); + } else { + outputs_ival = module.run_method( + "forward_lower", coord_wrapped_Tensor, atype_Tensor, firstneigh_tensor, + mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor, + c10::IValue(), charge_spin_tensor); + } + } else { + outputs_ival = + (do_message_passing) + ? module.run_method("forward_lower", coord_wrapped_Tensor, + atype_Tensor, firstneigh_tensor, mapping_tensor, + fparam_tensor, aparam_tensor, + do_atom_virial_tensor, comm_dict) + : module.run_method("forward_lower", coord_wrapped_Tensor, + atype_Tensor, firstneigh_tensor, mapping_tensor, + fparam_tensor, aparam_tensor, + do_atom_virial_tensor); + } + auto outputs = outputs_ival.toGenericDict(); c10::IValue energy_ = outputs.at("energy"); c10::IValue force_ = outputs.at("extended_force"); c10::IValue virial_ = outputs.at("virial"); @@ -313,6 +377,7 @@ template void DeepPotPT::compute>( const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepPotPT::compute>( std::vector& ener, @@ -328,6 +393,7 @@ template void DeepPotPT::compute>( const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepPotPT::compute(ENERGYVTYPE& ener, @@ -340,6 +406,7 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { torch::Device device(torch::kCUDA, gpu_id); if (!gpu_enabled) { @@ -391,6 +458,34 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, inputs.push_back(aparam_tensor); bool do_atom_virial_tensor = atomic; inputs.push_back(do_atom_virial_tensor); + // Append charge_spin (always float64) only for models with a charge/spin + // embedding, so models without it (whose forward() lacks the parameter) + // keep working. Uses the runtime value when provided, else the model's + // stored default_chg_spin. + if (dchgspin > 0) { + auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); + c10::optional charge_spin_tensor; + if (!charge_spin.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(charge_spin.data()), + {1, static_cast(charge_spin.size())}, + dbl_options) + .clone() + .to(device); + } else if (!default_chg_spin_.empty()) { + charge_spin_tensor = + torch::from_blob(const_cast(default_chg_spin_.data()), + {1, dchgspin}, dbl_options) + .clone() + .to(device); + } else { + throw deepmd::deepmd_exception( + "charge_spin is empty and no default_chg_spin is available in the " + "model. Provide charge_spin explicitly or regenerate the model with " + "a default charge/spin value."); + } + inputs.push_back(charge_spin_tensor); + } c10::Dict outputs = module.forward(inputs).toGenericDict(); c10::IValue energy_ = outputs.at("energy"); @@ -437,6 +532,7 @@ template void DeepPotPT::compute>( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); template void DeepPotPT::compute>( std::vector& ener, @@ -449,6 +545,7 @@ template void DeepPotPT::compute>( const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic); void DeepPotPT::get_type_map(std::string& type_map) { auto ret = module.run_method("get_type_map").toList(); @@ -472,7 +569,79 @@ void DeepPotPT::computew(std::vector& ener, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - fparam, aparam, atomic); + fparam, aparam, {}, atomic); + }); +} +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam, {}, atomic); + }); +} +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam, {}, atomic); + }); +} +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const int& ago, + const std::vector& fparam, + const std::vector& aparam, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + nghost, inlist, ago, fparam, aparam, {}, atomic); + }); +} +// charge_spin overloads — thread runtime charge/spin through to compute() +void DeepPotPT::computew(std::vector& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam, + const std::vector& aparam, + const std::vector& charge_spin, + const bool atomic) { + translate_error([&] { + compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, + fparam, aparam, charge_spin, atomic); }); } void DeepPotPT::computew(std::vector& ener, @@ -485,10 +654,11 @@ void DeepPotPT::computew(std::vector& ener, const std::vector& box, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - fparam, aparam, atomic); + fparam, aparam, charge_spin, atomic); }); } void DeepPotPT::computew(std::vector& ener, @@ -504,10 +674,11 @@ void DeepPotPT::computew(std::vector& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - nghost, inlist, ago, fparam, aparam, atomic); + nghost, inlist, ago, fparam, aparam, charge_spin, atomic); }); } void DeepPotPT::computew(std::vector& ener, @@ -523,10 +694,11 @@ void DeepPotPT::computew(std::vector& ener, const int& ago, const std::vector& fparam, const std::vector& aparam, + const std::vector& charge_spin, const bool atomic) { translate_error([&] { compute(ener, force, virial, atom_energy, atom_virial, coord, atype, box, - nghost, inlist, ago, fparam, aparam, atomic); + nghost, inlist, ago, fparam, aparam, charge_spin, atomic); }); } void DeepPotPT::computew_mixed_type(std::vector& ener, diff --git a/source/api_cc/tests/test_deeppot_chg_spin_pt.cc b/source/api_cc/tests/test_deeppot_chg_spin_pt.cc new file mode 100644 index 0000000000..7bfce1b3dc --- /dev/null +++ b/source/api_cc/tests/test_deeppot_chg_spin_pt.cc @@ -0,0 +1,280 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +// Test C++ inference for the PT (.pth) backend with charge_spin input (DPA3). +// Model + reference generated by source/tests/infer/gen_chg_spin.py. +// +// Mirrors the DPA3 .pth test layout (test_deeppot_dpa3_pt.cc): the PBC case +// exercises the standalone path (forward), and the NoPbc case exercises the +// LAMMPS neighbor-list path (forward_lower) without ghost atoms. Both cases +// cover explicit charge_spin and the default_chg_spin fallback. +#include + +#include +#include +#include + +#include "DeepPot.h" +#include "expected_ref.h" +#include "neighbor_list.h" +#include "test_utils.h" + +#undef EPSILON +#define EPSILON (std::is_same::value ? 1e-10 : 1e-4) + +namespace { +constexpr const char* kRefPath = "../../tests/infer/chg_spin.expected"; +constexpr const char* kModelPath = "../../tests/infer/chg_spin.pth"; +} // namespace + +// --------------------------------------------------------------------------- +// PBC: standalone path (DeepPot::compute that builds its own nlist -> forward) +// --------------------------------------------------------------------------- +template +class TestInferDeepPotChgSpinPt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + // charge_spin is always double regardless of VALUETYPE. + // [1.0, 2.0] maps to charge idx 101, spin idx 2 — both differ from the + // model's default_chg_spin [0.0, 1.0] (charge idx 100, spin idx 1). + std::vector charge_spin_explicit = {1.0, 2.0}; + + std::vector expected_f_explicit, expected_f_default; + int natoms; + double expected_tot_e_explicit, expected_tot_e_default; + std::vector expected_tot_v_explicit, expected_tot_v_default; + + deepmd::DeepPot dp; + + static void accumulate(const std::vector& e, + const std::vector& v, + double& tot_e, + std::vector& tot_v) { + int n = e.size(); + tot_e = 0.; + tot_v.assign(9, 0.); + for (int ii = 0; ii < n; ++ii) { + tot_e += e[ii]; + for (int dd = 0; dd < 9; ++dd) { + tot_v[dd] += v[ii * 9 + dd]; + } + } + } + + void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + std::vector e_exp = ref.get("explicit", "expected_e"); + std::vector e_def = ref.get("default", "expected_e"); + expected_f_explicit = ref.get("explicit", "expected_f"); + expected_f_default = ref.get("default", "expected_f"); + accumulate(e_exp, ref.get("explicit", "expected_v"), + expected_tot_e_explicit, expected_tot_v_explicit); + accumulate(e_def, ref.get("default", "expected_v"), + expected_tot_e_default, expected_tot_v_default); + natoms = e_exp.size(); + + dp.init(kModelPath); + } + + void TearDown() override {} +}; + +TYPED_TEST_SUITE(TestInferDeepPotChgSpinPt, ValueTypes); + +TYPED_TEST(TestInferDeepPotChgSpinPt, dim_chg_spin) { + deepmd::DeepPot& dp = this->dp; + EXPECT_EQ(dp.dim_chg_spin(), 2); +} + +TYPED_TEST(TestInferDeepPotChgSpinPt, cpu_build_nlist_explicit) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& charge_spin = this->charge_spin_explicit; + std::vector& expected_f = this->expected_f_explicit; + std::vector& expected_tot_v = this->expected_tot_v_explicit; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_explicit; + deepmd::DeepPot& dp = this->dp; + + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box, {}, {}, charge_spin); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotChgSpinPt, cpu_build_nlist_default) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f_default; + std::vector& expected_tot_v = this->expected_tot_v_default; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_default; + deepmd::DeepPot& dp = this->dp; + + double ener; + std::vector force, virial; + // Empty charge_spin triggers the default_chg_spin fallback ([0.0, 1.0]). + dp.compute(ener, force, virial, coord, atype, box); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +// --------------------------------------------------------------------------- +// NoPbc: LAMMPS neighbor-list path (forward_lower), no ghost atoms. +// Mirrors test_deeppot_dpa3_pt.cc::TestInferDeepPotDpa3PtNoPbc.cpu_lmp_nlist. +// --------------------------------------------------------------------------- +template +class TestInferDeepPotChgSpinPtNoPbc : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {}; + std::vector charge_spin_explicit = {1.0, 2.0}; + + std::vector expected_f_explicit, expected_f_default; + int natoms; + double expected_tot_e_explicit, expected_tot_e_default; + std::vector expected_tot_v_explicit, expected_tot_v_default; + + deepmd::DeepPot dp; + + static void accumulate(const std::vector& e, + const std::vector& v, + double& tot_e, + std::vector& tot_v) { + int n = e.size(); + tot_e = 0.; + tot_v.assign(9, 0.); + for (int ii = 0; ii < n; ++ii) { + tot_e += e[ii]; + for (int dd = 0; dd < 9; ++dd) { + tot_v[dd] += v[ii * 9 + dd]; + } + } + } + + void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + deepmd_test::ExpectedRef ref; + ref.load(kRefPath); + std::vector e_exp = + ref.get("nopbc_explicit", "expected_e"); + std::vector e_def = + ref.get("nopbc_default", "expected_e"); + expected_f_explicit = ref.get("nopbc_explicit", "expected_f"); + expected_f_default = ref.get("nopbc_default", "expected_f"); + accumulate(e_exp, ref.get("nopbc_explicit", "expected_v"), + expected_tot_e_explicit, expected_tot_v_explicit); + accumulate(e_def, ref.get("nopbc_default", "expected_v"), + expected_tot_e_default, expected_tot_v_default); + natoms = e_exp.size(); + + dp.init(kModelPath); + } + + void TearDown() override {} +}; + +TYPED_TEST_SUITE(TestInferDeepPotChgSpinPtNoPbc, ValueTypes); + +TYPED_TEST(TestInferDeepPotChgSpinPtNoPbc, cpu_lmp_nlist_explicit) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& charge_spin = this->charge_spin_explicit; + std::vector& expected_f = this->expected_f_explicit; + std::vector& expected_tot_v = this->expected_tot_v_explicit; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_explicit; + deepmd::DeepPot& dp = this->dp; + + std::vector> nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force, virial; + dp.compute(ener, force, virial, coord, atype, box, 0, inlist, 0, {}, {}, + charge_spin); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} + +TYPED_TEST(TestInferDeepPotChgSpinPtNoPbc, cpu_lmp_nlist_default) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_f = this->expected_f_default; + std::vector& expected_tot_v = this->expected_tot_v_default; + int& natoms = this->natoms; + double& expected_tot_e = this->expected_tot_e_default; + deepmd::DeepPot& dp = this->dp; + + std::vector> nlist_data = { + {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, + {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector ilist(natoms), numneigh(natoms); + std::vector firstneigh(natoms); + deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + double ener; + std::vector force, virial; + // Empty charge_spin triggers the default_chg_spin fallback ([0.0, 1.0]). + dp.compute(ener, force, virial, coord, atype, box, 0, inlist, 0); + + EXPECT_EQ(force.size(), static_cast(natoms * 3)); + EXPECT_EQ(virial.size(), 9u); + EXPECT_LT(fabs(ener - expected_tot_e), EPSILON); + for (int ii = 0; ii < natoms * 3; ++ii) { + EXPECT_LT(fabs(force[ii] - expected_f[ii]), EPSILON); + } + for (int ii = 0; ii < 9; ++ii) { + EXPECT_LT(fabs(virial[ii] - expected_tot_v[ii]), EPSILON); + } +} diff --git a/source/tests/infer/gen_chg_spin.py b/source/tests/infer/gen_chg_spin.py index e849c78d78..ba43f00b12 100644 --- a/source/tests/infer/gen_chg_spin.py +++ b/source/tests/infer/gen_chg_spin.py @@ -1,11 +1,13 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: LGPL-3.0-or-later -"""Generate chg_spin.pt2 test model and reference values. +"""Generate chg_spin.pt2 / chg_spin.pth test models and reference values. Creates a DPA3 model with add_chg_spin_ebd=True and default_chg_spin=[0.0, 1.0], -exports to .pt2, and writes chg_spin.expected with two sections: +exports it to both .pt2 (AOTInductor) and .pth (TorchScript) from the same +weights, and writes chg_spin.expected with two sections: [default] -- eval with no charge_spin (uses stored default [0.0, 1.0]) - [explicit] -- eval with charge_spin=[0.5, 0.8] + [explicit] -- eval with charge_spin=[1.0, 2.0] +The .pth is verified to match the .pt2 reference for both sections. """ import copy @@ -71,6 +73,9 @@ def main(): "version": "3.0.0", } + from deepmd.pt.utils.serialization import ( + deserialize_to_file as pt_deserialize_to_file, + ) from deepmd.pt_expt.utils.serialization import ( deserialize_to_file as pt_expt_deserialize_to_file, ) @@ -81,6 +86,16 @@ def main(): pt2_path = os.path.join(base_dir, "chg_spin.pt2") print(f"Exporting to {pt2_path} ...") # noqa: T201 pt_expt_deserialize_to_file(pt2_path, copy.deepcopy(data), do_atomic_virial=True) + + pth_path = os.path.join(base_dir, "chg_spin.pth") + print(f"Exporting to {pth_path} ...") # noqa: T201 + try: + pt_deserialize_to_file(pth_path, copy.deepcopy(data)) + except RuntimeError as e: + # Custom ops may not be available in all build environments; .pth + # generation is not critical (the .pth test skips if the file is + # missing / PyTorch support is off). + print(f"WARNING: .pth export failed ({e}), skipping.") # noqa: T201 print("Export done.") # noqa: T201 from deepmd.infer import ( @@ -136,6 +151,19 @@ def main(): "Default and explicit charge_spin gave identical energy — charge_spin may be ignored" ) + # NoPbc variants (box=None) — used by the .pth nlist test, which mirrors the + # established DPA3 .pth pattern (NoPbc lmp_nlist, nghost=0) and so needs its + # own reference values. + e_np_def, f_np_def, v_np_def, ae_np_def, av_np_def = dp.eval( + coord, None, atype, atomic=True + ) + e_np_exp, f_np_exp, v_np_exp, ae_np_exp, av_np_exp = dp.eval( + coord, None, atype, atomic=True, charge_spin=charge_spin_exp + ) + assert not np.allclose(e_np_def, e_np_exp), ( + "NoPbc: default and explicit charge_spin gave identical energy" + ) + ref_path = os.path.join(base_dir, "chg_spin.expected") write_expected_ref( ref_path, @@ -150,10 +178,53 @@ def main(): "expected_f": f_exp[0], "expected_v": av_exp[0], }, + "nopbc_default": { + "expected_e": ae_np_def[0, :, 0], + "expected_f": f_np_def[0], + "expected_v": av_np_def[0], + }, + "nopbc_explicit": { + "expected_e": ae_np_exp[0, :, 0], + "expected_f": f_np_exp[0], + "expected_v": av_np_exp[0], + }, }, source_script="source/tests/infer/gen_chg_spin.py", ) print(f"Wrote {ref_path}") # noqa: T201 + + # ---- Verify .pth reproduces the .pt2 reference (PBC + NoPbc) ---- + if os.path.exists(pth_path): + dp_pth = DeepPot(pth_path) + assert dp_pth.deep_eval.get_dim_chg_spin() == 2 + tol = 1e-10 + e_def_p, f_def_p, v_def_p = dp_pth.eval(coord, box, atype) + e_exp_p, f_exp_p, v_exp_p = dp_pth.eval( + coord, box, atype, charge_spin=charge_spin_exp + ) + np.testing.assert_allclose(e_def_p, e_def, atol=tol, err_msg="pth default e") + np.testing.assert_allclose(f_def_p, f_def, atol=tol, err_msg="pth default f") + np.testing.assert_allclose(e_exp_p, e_exp, atol=tol, err_msg="pth explicit e") + np.testing.assert_allclose(f_exp_p, f_exp, atol=tol, err_msg="pth explicit f") + # NoPbc parity (the .pth nlist test uses NoPbc) + e_np_def_p, f_np_def_p, _ = dp_pth.eval(coord, None, atype) + e_np_exp_p, f_np_exp_p, _ = dp_pth.eval( + coord, None, atype, charge_spin=charge_spin_exp + ) + np.testing.assert_allclose( + e_np_def_p, e_np_def, atol=tol, err_msg="pth nopbc default e" + ) + np.testing.assert_allclose( + e_np_exp_p, e_np_exp, atol=tol, err_msg="pth nopbc explicit e" + ) + # different charge_spin must change the .pth output too + assert not np.allclose(e_def_p, e_exp_p), ( + ".pth: default and explicit charge_spin gave identical energy" + ) + print("// .pth matches .pt2 reference (PBC + NoPbc).") # noqa: T201 + else: + print("\n// Skipping .pth verification (file not generated).") # noqa: T201 + print("\nDone!") # noqa: T201 From b68dc9a0823479e98bfb1bd4df068d6e5d0e7cb4 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 10 Jun 2026 17:59:57 +0800 Subject: [PATCH 13/24] fix: address charge_spin review comments - pair_base.cpp: guard make_charge_spin_from_compute against modify->find_compute() returning -1 (out-of-bounds read before the null check), matching the find_compute contract. - DeepPotModelDevi: add dim_chg_spin() (api_cc delegates to the first model; C-API returns 0). pair_deepmd now reads dim_chg_spin from deep_pot_model_devi in the multi-model branch and asserts it matches deep_pot, consistent with the other dims. - DeepPotPTExpt: validate runtime charge_spin size against dim_chg_spin. Single-frame paths require exactly dim_chg_spin values; the multi-frame (nframes / mixed_type) paths accept either a single dim_chg_spin vector (now broadcast to all frames) or nframes * dim_chg_spin, and throw on any other length instead of slicing out of range. Co-Authored-By: Claude Opus 4.8 --- source/api_c/include/deepmd.hpp | 5 +++ source/api_cc/include/DeepPot.h | 10 +++++ source/api_cc/src/DeepPotPTExpt.cc | 62 ++++++++++++++++++++++++++++-- source/lmp/pair_base.cpp | 3 ++ source/lmp/pair_deepmd.cpp | 3 +- 5 files changed, 78 insertions(+), 5 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 2eee017c28..cdcfc4d056 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1891,6 +1891,11 @@ class DeepBaseModelDevi { assert(dpbase); return has_default_fparam_; } + /** + * @brief Get the dimension of the charge/spin input. + * @return Always 0; charge_spin is not supported via the C-API path. + **/ + int dim_chg_spin() const { return 0; } /** * @param[out] avg The average of vectors. * @param[in] xx The vectors of all models. diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h index 1f01396509..6172f08c1c 100644 --- a/source/api_cc/include/DeepPot.h +++ b/source/api_cc/include/DeepPot.h @@ -676,6 +676,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& file_contents = std::vector()); + /** + * @brief Get the dimension of the charge/spin input. + * @return The dimension of the charge/spin input (0 if the models have no + *charge/spin embedding). Taken from the first model; all models are assumed + *to share the same value. + **/ + int dim_chg_spin() const { + return numb_models > 0 ? dps[0]->dim_chg_spin() : 0; + }; + /** * @brief Evaluate the energy, force and virial by using these DP models. * @param[out] all_ener The system energies of all models. diff --git a/source/api_cc/src/DeepPotPTExpt.cc b/source/api_cc/src/DeepPotPTExpt.cc index cae339a9c0..caf73416c8 100644 --- a/source/api_cc/src/DeepPotPTExpt.cc +++ b/source/api_cc/src/DeepPotPTExpt.cc @@ -514,6 +514,13 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (dchgspin > 0) { auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); if (!charge_spin.empty()) { + // Single-frame path: charge_spin must hold exactly dim_chg_spin values. + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } charge_spin_tensor = torch::from_blob(const_cast(charge_spin.data()), {1, static_cast(charge_spin.size())}, @@ -838,6 +845,13 @@ void DeepPotPTExpt::compute(ENERGYVTYPE& ener, if (dchgspin > 0) { auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); if (!charge_spin.empty()) { + // Single-frame path: charge_spin must hold exactly dim_chg_spin values. + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } charge_spin_tensor = torch::from_blob(const_cast(charge_spin.data()), {1, static_cast(charge_spin.size())}, @@ -930,6 +944,21 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, int natoms = atype.size(); int dap = aparam.empty() ? 0 : static_cast(aparam.size()) / nframes; int dfp = fparam.empty() ? 0 : static_cast(fparam.size()) / nframes; + // charge_spin may be empty (default fallback), a single dim_chg_spin vector + // (broadcast to all frames), or nframes * dim_chg_spin (per-frame). Reject + // anything else up-front to avoid out-of-range slicing in the loop. + if (!charge_spin.empty()) { + size_t s_dcsp = static_cast(dchgspin); + if (charge_spin.size() != s_dcsp && + charge_spin.size() != s_dcsp * static_cast(nframes)) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + " (per frame) or " + + std::to_string(dchgspin * nframes) + " (for " + + std::to_string(nframes) + " frames)."); + } + } ener.clear(); force.clear(); virial.clear(); @@ -962,8 +991,13 @@ void DeepPotPTExpt::compute_nframes(ENERGYVTYPE& ener, std::vector frame_chg_spin; if (!charge_spin.empty()) { size_t s_dcsp = static_cast(dchgspin); - frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, - charge_spin.begin() + (s_ff + 1) * s_dcsp); + if (charge_spin.size() == s_dcsp) { + // single charge/spin vector broadcast to every frame + frame_chg_spin = charge_spin; + } else { + frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, + charge_spin.begin() + (s_ff + 1) * s_dcsp); + } } std::vector frame_ener; std::vector frame_force, frame_virial, frame_ae, frame_av; @@ -1105,6 +1139,21 @@ void DeepPotPTExpt::compute_mixed_type_impl( int natoms = static_cast(atype.size()) / nframes; int dap = aparam.empty() ? 0 : static_cast(aparam.size()) / nframes; int dfp = fparam.empty() ? 0 : static_cast(fparam.size()) / nframes; + // charge_spin may be empty (default fallback), a single dim_chg_spin vector + // (broadcast to all frames), or nframes * dim_chg_spin (per-frame). Reject + // anything else up-front to avoid out-of-range slicing in the loop. + if (!charge_spin.empty()) { + size_t s_dcsp = static_cast(dchgspin); + if (charge_spin.size() != s_dcsp && + charge_spin.size() != s_dcsp * static_cast(nframes)) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + " (per frame) or " + + std::to_string(dchgspin * nframes) + " (for " + + std::to_string(nframes) + " frames)."); + } + } ener.clear(); force.clear(); virial.clear(); @@ -1139,8 +1188,13 @@ void DeepPotPTExpt::compute_mixed_type_impl( std::vector frame_chg_spin; if (!charge_spin.empty()) { size_t s_dcsp = static_cast(dchgspin); - frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, - charge_spin.begin() + (s_ff + 1) * s_dcsp); + if (charge_spin.size() == s_dcsp) { + // single charge/spin vector broadcast to every frame + frame_chg_spin = charge_spin; + } else { + frame_chg_spin.assign(charge_spin.begin() + s_ff * s_dcsp, + charge_spin.begin() + (s_ff + 1) * s_dcsp); + } } std::vector frame_ener; std::vector frame_force, frame_virial, frame_ae, frame_av; diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index bc02a1f473..1ea4da94b2 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -169,6 +169,9 @@ void PairDeepBaseModel::make_charge_spin_from_compute( assert(do_compute_charge_spin); int icompute = modify->find_compute(compute_charge_spin_id); + if (icompute < 0) { + error->all(FLERR, "compute id is not found: " + compute_charge_spin_id); + } Compute* compute = modify->compute[icompute]; if (!compute) { diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 3446a71048..3c4dcbdb7b 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -598,12 +598,13 @@ void PairDeepMD::settings(int narg, char** arg) { 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(); + 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()); } out_freq = 100; From f743f00624048e9ddcc633d2d857d98a9317cc9c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 10:01:40 +0000 Subject: [PATCH 14/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_cc/src/DeepPotPT.cc | 12 ++++++------ source/api_cc/tests/test_deeppot_chg_spin_pt.cc | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index a725da2439..4878491c69 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -296,14 +296,14 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, // flip the descriptor into parallel mode). if (do_message_passing) { outputs_ival = module.run_method( - "forward_lower", coord_wrapped_Tensor, atype_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor, - comm_dict, charge_spin_tensor); + "forward_lower", coord_wrapped_Tensor, atype_Tensor, + firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, + do_atom_virial_tensor, comm_dict, charge_spin_tensor); } else { outputs_ival = module.run_method( - "forward_lower", coord_wrapped_Tensor, atype_Tensor, firstneigh_tensor, - mapping_tensor, fparam_tensor, aparam_tensor, do_atom_virial_tensor, - c10::IValue(), charge_spin_tensor); + "forward_lower", coord_wrapped_Tensor, atype_Tensor, + firstneigh_tensor, mapping_tensor, fparam_tensor, aparam_tensor, + do_atom_virial_tensor, c10::IValue(), charge_spin_tensor); } } else { outputs_ival = diff --git a/source/api_cc/tests/test_deeppot_chg_spin_pt.cc b/source/api_cc/tests/test_deeppot_chg_spin_pt.cc index 7bfce1b3dc..efa39b3304 100644 --- a/source/api_cc/tests/test_deeppot_chg_spin_pt.cc +++ b/source/api_cc/tests/test_deeppot_chg_spin_pt.cc @@ -220,9 +220,9 @@ TYPED_TEST(TestInferDeepPotChgSpinPtNoPbc, cpu_lmp_nlist_explicit) { double& expected_tot_e = this->expected_tot_e_explicit; deepmd::DeepPot& dp = this->dp; - std::vector> nlist_data = { - {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, - {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector> nlist_data = {{1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, + {0, 1, 3, 4, 5}, {0, 1, 2, 4, 5}, + {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; std::vector ilist(natoms), numneigh(natoms); std::vector firstneigh(natoms); deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); @@ -255,9 +255,9 @@ TYPED_TEST(TestInferDeepPotChgSpinPtNoPbc, cpu_lmp_nlist_default) { double& expected_tot_e = this->expected_tot_e_default; deepmd::DeepPot& dp = this->dp; - std::vector> nlist_data = { - {1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, {0, 1, 3, 4, 5}, - {0, 1, 2, 4, 5}, {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; + std::vector> nlist_data = {{1, 2, 3, 4, 5}, {0, 2, 3, 4, 5}, + {0, 1, 3, 4, 5}, {0, 1, 2, 4, 5}, + {0, 1, 2, 3, 5}, {0, 1, 2, 3, 4}}; std::vector ilist(natoms), numneigh(natoms); std::vector firstneigh(natoms); deepmd::InputNlist inlist(natoms, &ilist[0], &numneigh[0], &firstneigh[0]); From dce3c43b3353ae813bf1c3dc149e87ec08e1855e Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 10 Jun 2026 18:27:19 +0800 Subject: [PATCH 15/24] feat: wire charge_spin through the C API; drop charge_spin_from_compute MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address iProzd's review: Point 1 — charge_spin had no effect on the common LAMMPS build (DP_USING_C_API), which uses deepmd::hpp::DeepPot. It was accepted there but discarded via (void)charge_spin. Now: - Add C API entry points (DP_C_API_VERSION 26->27): DP_DeepPotGetDimChgSpin and the version-3 compute functions DP_DeepPotCompute3 / *f3 / ComputeNList3 / *f3, which take a per-frame charge_spin (nframes x dim_chg_spin, always double) and forward it to api_cc::DeepPot::compute. The internal variant helpers gained an optional charge_spin arg, so the existing version-2 functions are unchanged. - deepmd.hpp DeepPot now reads dim_chg_spin via DP_DeepPotGetDimChgSpin (dim_chg_spin() returns it) and threads charge_spin through to the version-3 entry points. When charge_spin is empty / the model has no charge-spin embedding, it keeps the version-2 path so non-charge_spin models still work against an older libdeepmd_c. Inline comment — remove the charge_spin_from_compute pair_style keyword (and make_charge_spin_from_compute): charge_spin is a per-frame global physical input and is set statically via the charge_spin keyword; deriving it from a per-step compute was unnecessary. Co-Authored-By: Claude Opus 4.8 --- source/api_c/include/c_api.h | 128 ++++++++++++++++++++++++++++++- source/api_c/include/deepmd.hpp | 95 ++++++++++++++++------- source/api_c/src/c_api.cc | 130 +++++++++++++++++++++++++++++--- source/lmp/pair_base.cpp | 34 --------- source/lmp/pair_base.h | 4 - source/lmp/pair_deepmd.cpp | 21 ------ 6 files changed, 314 insertions(+), 98 deletions(-) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 358480b0ad..78d193de64 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -12,7 +12,7 @@ extern "C" { /** C API version. Bumped whenever the API is changed. * @since API version 22 */ -#define DP_C_API_VERSION 26 +#define DP_C_API_VERSION 27 /** * @brief Neighbor list. @@ -675,6 +675,123 @@ extern void DP_DeepPotComputeNListf2(DP_DeepPot* dp, float* atomic_energy, float* atomic_virial); +/** + * @brief Evaluate the energy, force and virial by using a DP. (double version, + *with charge_spin) + * @version 3 + * @param[in] dp The DP to use. + * @param[in] nframes The number of frames. + * @param[in] natoms The number of atoms. + * @param[in] coord The coordinates of atoms. The array should be of size natoms + *x 3. + * @param[in] atype The atom types. The array should contain natoms ints. + * @param[in] cell The cell of the region. The array should be of size 9. Pass + *NULL if pbc is not used. + * @param[in] fparam The frame parameters. The array can be of size nframes x + *dim_fparam. + * @param[in] aparam The atom parameters. The array can be of size nframes x + *natoms x dim_aparam. + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @param[out] energy Output energy. + * @param[out] force Output force. The array should be of size natoms x 3. + * @param[out] virial Output virial. The array should be of size 9. + * @param[out] atomic_energy Output atomic energy. The array should be of size + *natoms. + * @param[out] atomic_virial Output atomic virial. The array should be of size + *natoms x 9. + * @warning The output arrays should be allocated before calling this function. + *Pass NULL if not required. + * @since API version 27 + **/ +extern void DP_DeepPotCompute3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate the energy, force and virial by using a DP. (float version, + *with charge_spin) + * @version 3 + * @since API version 27 + **/ +extern void DP_DeepPotComputef3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial); + +/** + * @brief Evaluate the energy, force and virial by using a DP with the neighbor + *list. (double version, with charge_spin) + * @version 3 + * @param[in] charge_spin The per-frame charge/spin input. The array can be of + *size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @since API version 27 + **/ +extern void DP_DeepPotComputeNList3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate the energy, force and virial by using a DP with the neighbor + *list. (float version, with charge_spin) + * @version 3 + * @since API version 27 + **/ +extern void DP_DeepPotComputeNListf3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial); + /** * @brief Evaluate the energy, force, magnetic force and virial by using a DP *spin model with the neighbor list. (float version) @@ -1583,6 +1700,15 @@ int DP_DeepPotGetDimFParam(DP_DeepPot* dp); */ int DP_DeepPotGetDimAParam(DP_DeepPot* dp); +/** + * @brief Get the dimension of the charge/spin input of a DP. + * @param[in] dp The DP to use. + * @return The dimension of the charge/spin input (0 if the model has no + * charge/spin embedding). + * @since API version 27 + */ +int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index cdcfc4d056..9e3e7256bf 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -57,7 +57,8 @@ inline void _DP_DeepPotCompute(DP_DeepPot* dp, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial); + FPTYPE* atomic_virial, + const double* charge_spin); template <> inline void _DP_DeepPotCompute(DP_DeepPot* dp, @@ -72,9 +73,18 @@ inline void _DP_DeepPotCompute(DP_DeepPot* dp, double* force, double* virial, double* atomic_energy, - double* atomic_virial) { - DP_DeepPotCompute2(dp, nframes, natom, coord, atype, cell, fparam, aparam, - energy, force, virial, atomic_energy, atomic_virial); + double* atomic_virial, + const double* charge_spin) { + // charge_spin == nullptr keeps the version-2 entry point so models without a + // charge/spin embedding still work against an older libdeepmd_c. + if (charge_spin) { + DP_DeepPotCompute3(dp, nframes, natom, coord, atype, cell, fparam, aparam, + charge_spin, energy, force, virial, atomic_energy, + atomic_virial); + } else { + DP_DeepPotCompute2(dp, nframes, natom, coord, atype, cell, fparam, aparam, + energy, force, virial, atomic_energy, atomic_virial); + } } template <> @@ -90,9 +100,16 @@ inline void _DP_DeepPotCompute(DP_DeepPot* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial) { - DP_DeepPotComputef2(dp, nframes, natom, coord, atype, cell, fparam, aparam, - energy, force, virial, atomic_energy, atomic_virial); + float* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotComputef3(dp, nframes, natom, coord, atype, cell, fparam, aparam, + charge_spin, energy, force, virial, atomic_energy, + atomic_virial); + } else { + DP_DeepPotComputef2(dp, nframes, natom, coord, atype, cell, fparam, aparam, + energy, force, virial, atomic_energy, atomic_virial); + } } // support spin @@ -171,7 +188,8 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial); + FPTYPE* atomic_virial, + const double* charge_spin); template <> inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, @@ -189,10 +207,17 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, double* force, double* virial, double* atomic_energy, - double* atomic_virial) { - DP_DeepPotComputeNList2(dp, nframes, natom, coord, atype, cell, nghost, nlist, - ago, fparam, aparam, energy, force, virial, - atomic_energy, atomic_virial); + double* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotComputeNList3(dp, nframes, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, charge_spin, energy, + force, virial, atomic_energy, atomic_virial); + } else { + DP_DeepPotComputeNList2(dp, nframes, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, energy, force, virial, + atomic_energy, atomic_virial); + } } template <> @@ -211,10 +236,17 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial) { - DP_DeepPotComputeNListf2(dp, nframes, natom, coord, atype, cell, nghost, - nlist, ago, fparam, aparam, energy, force, virial, - atomic_energy, atomic_virial); + float* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotComputeNListf3(dp, nframes, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, charge_spin, energy, + force, virial, atomic_energy, atomic_virial); + } else { + DP_DeepPotComputeNListf2(dp, nframes, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, energy, force, virial, + atomic_energy, atomic_virial); + } } // support spin @@ -1069,6 +1101,7 @@ class DeepPot : public DeepBaseModel { DP_CHECK_OK(DP_DeepPotCheckOK, dp); dfparam = DP_DeepPotGetDimFParam(dp); daparam = DP_DeepPotGetDimAParam(dp); + dchgspin = DP_DeepPotGetDimChgSpin(dp); aparam_nall = DP_DeepPotIsAParamNAll(dp); has_default_fparam_ = DP_DeepPotHasDefaultFParam(dp); dpbase = (DP_DeepBaseModel*)dp; @@ -1076,9 +1109,10 @@ class DeepPot : public DeepBaseModel { /** * @brief Get the dimension of the charge/spin embedding input. - * @return Always 0; charge_spin is not supported via the C-API path. + * @return The dimension of the charge/spin input (0 if the model has no + *charge/spin embedding). **/ - int dim_chg_spin() const { return 0; } + int dim_chg_spin() const { return dchgspin; } /** * @brief Evaluate the energy, force and virial by using this DP. @@ -1110,7 +1144,6 @@ class DeepPot : public DeepBaseModel { const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1131,10 +1164,14 @@ class DeepPot : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + // charge_spin routes to the version-3 C API; nullptr keeps version-2 so + // non-charge_spin models still work against an older libdeepmd_c. + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, virial_, - nullptr, nullptr); + nullptr, nullptr, charge_spin__); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -1172,7 +1209,6 @@ class DeepPot : public DeepBaseModel { const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1198,10 +1234,12 @@ class DeepPot : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, virial_, - atomic_ener_, atomic_virial_); + atomic_ener_, atomic_virial_, charge_spin__); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; @@ -1242,7 +1280,6 @@ class DeepPot : public DeepBaseModel { const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1266,10 +1303,13 @@ class DeepPot : public DeepBaseModel { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotComputeNList( dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, ener_, force_, virial_, nullptr, nullptr); + fparam__, aparam__, ener_, force_, virial_, nullptr, nullptr, + charge_spin__); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -1313,7 +1353,6 @@ class DeepPot : public DeepBaseModel { const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1342,11 +1381,14 @@ class DeepPot : public DeepBaseModel { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, aparam__, ener_, force_, virial_, - atomic_ener_, atomic_virial_); + atomic_ener_, atomic_virial_, + charge_spin__); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -1475,6 +1517,7 @@ class DeepPot : public DeepBaseModel { private: DP_DeepPot* dp; + int dchgspin; }; class DeepSpin : public DeepBaseModel { diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index b0e789648e..b3ef76f3a2 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -251,7 +251,8 @@ inline void DP_DeepPotCompute_variant(DP_DeepPot* dp, VALUETYPE* force, VALUETYPE* virial, VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { + VALUETYPE* atomic_virial, + const double* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector atype_(atype, atype + natoms); @@ -268,15 +269,21 @@ inline void DP_DeepPotCompute_variant(DP_DeepPot* dp, if (aparam) { aparam_.assign(aparam, aparam + nframes * natoms * dp->daparam); } + // charge_spin is always double (it is a per-frame physical input) + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } std::vector e; std::vector f, v, ae, av; if (atomic_energy || atomic_virial) { DP_REQUIRES_OK(dp, dp->dp.compute(e, f, v, ae, av, coord_, atype_, cell_, - fparam_, aparam_)); + fparam_, aparam_, charge_spin_)); } else { - DP_REQUIRES_OK( - dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, fparam_, aparam_)); + DP_REQUIRES_OK(dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, fparam_, + aparam_, charge_spin_)); } // copy from C++ vectors to C arrays, if not NULL pointer if (energy) { @@ -308,7 +315,8 @@ template void DP_DeepPotCompute_variant(DP_DeepPot* dp, double* force, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepPotCompute_variant(DP_DeepPot* dp, const int nframes, @@ -322,7 +330,8 @@ template void DP_DeepPotCompute_variant(DP_DeepPot* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const double* charge_spin); // support spin template inline void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, @@ -431,7 +440,8 @@ inline void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, VALUETYPE* force, VALUETYPE* virial, VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { + VALUETYPE* atomic_virial, + const double* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector atype_(atype, atype + natoms); @@ -451,16 +461,24 @@ inline void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, (dp->aparam_nall ? natoms : (natoms - nghost)) * dp->daparam); } + // charge_spin is always double (it is a per-frame physical input) + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } std::vector e; std::vector f, v, ae, av; if (atomic_energy || atomic_virial) { DP_REQUIRES_OK( dp, dp->dp.compute(e, f, v, ae, av, coord_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_)); + nlist->nl, ago, fparam_, aparam_, charge_spin_)); } else { - DP_REQUIRES_OK(dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_)); + DP_REQUIRES_OK(dp, + dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, + nlist->nl, ago, fparam_, aparam_, + charge_spin_)); } // copy from C++ vectors to C arrays, if not NULL pointer if (energy) { @@ -495,7 +513,8 @@ template void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, double* force, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, const int nframes, @@ -512,7 +531,8 @@ template void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const double* charge_spin); // support spin template @@ -1702,6 +1722,90 @@ void DP_DeepPotComputeNListf2(DP_DeepPot* dp, aparam, energy, force, virial, atomic_energy, atomic_virial); } +// charge_spin-aware variants (version 3): same as the version-2 functions +// plus a per-frame charge_spin input (nframes x dim_chg_spin, always double). +void DP_DeepPotCompute3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepPotCompute_variant(dp, nframes, natoms, coord, atype, cell, + fparam, aparam, energy, force, virial, + atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepPotComputef3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepPotCompute_variant(dp, nframes, natoms, coord, atype, cell, + fparam, aparam, energy, force, virial, + atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepPotComputeNList3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepPotComputeNList_variant( + dp, nframes, natoms, coord, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepPotComputeNListf3(DP_DeepPot* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepPotComputeNList_variant( + dp, nframes, natoms, coord, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, virial, atomic_energy, atomic_virial, charge_spin); +} + void DP_DeepSpinComputeNListf2(DP_DeepSpin* dp, const int nframes, const int natoms, @@ -2089,6 +2193,8 @@ int DP_DeepPotGetDimAParam(DP_DeepPot* dp) { return DP_DeepBaseModelGetDimAParam(static_cast(dp)); } +int DP_DeepPotGetDimChgSpin(DP_DeepPot* dp) { return dp->dp.dim_chg_spin(); } + bool DP_DeepPotIsAParamNAll(DP_DeepPot* dp) { return DP_DeepBaseModelIsAParamNAll(static_cast(dp)); } diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index 1ea4da94b2..faadd086b5 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -164,39 +164,6 @@ void PairDeepBaseModel::make_fparam_from_compute(vector& fparam) { } } -void PairDeepBaseModel::make_charge_spin_from_compute( - vector& charge_spin) { - assert(do_compute_charge_spin); - - int icompute = modify->find_compute(compute_charge_spin_id); - if (icompute < 0) { - error->all(FLERR, "compute id is not found: " + compute_charge_spin_id); - } - Compute* compute = modify->compute[icompute]; - - if (!compute) { - error->all(FLERR, "compute id is not found: " + compute_charge_spin_id); - } - charge_spin.resize(dim_chg_spin); - - if (dim_chg_spin == 1) { - if (!(compute->invoked_flag & Compute::INVOKED_SCALAR)) { - compute->compute_scalar(); - compute->invoked_flag |= Compute::INVOKED_SCALAR; - } - charge_spin[0] = compute->scalar; - } else if (dim_chg_spin > 1) { - if (!(compute->invoked_flag & Compute::INVOKED_VECTOR)) { - compute->compute_vector(); - compute->invoked_flag |= Compute::INVOKED_VECTOR; - } - double* cvector = compute->vector; - for (int jj = 0; jj < dim_chg_spin; ++jj) { - charge_spin[jj] = cvector[jj]; - } - } -} - void PairDeepBaseModel::make_aparam_from_compute(vector& aparam) { assert(do_compute_aparam); @@ -375,7 +342,6 @@ PairDeepBaseModel::PairDeepBaseModel( dim_chg_spin = 0; do_compute_fparam = false; do_compute_aparam = false; - do_compute_charge_spin = false; single_model = false; multi_models_mod_devi = false; multi_models_no_mod_devi = false; diff --git a/source/lmp/pair_base.h b/source/lmp/pair_base.h index 22555cb07e..93c5cb7b39 100644 --- a/source/lmp/pair_base.h +++ b/source/lmp/pair_base.h @@ -92,10 +92,6 @@ class PairDeepBaseModel : public Pair { bool do_compute_aparam; std::string compute_aparam_id; - void make_charge_spin_from_compute(std::vector& charge_spin); - bool do_compute_charge_spin; - std::string compute_charge_spin_id; - void make_ttm_fparam(std::vector& fparam); void make_ttm_aparam(std::vector& dparam); diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index 3c4dcbdb7b..c1c4c29635 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -216,10 +216,6 @@ void PairDeepMD::compute(int eflag, int vflag) { make_fparam_from_compute(fparam); } - if (do_compute_charge_spin) { - make_charge_spin_from_compute(charge_spin); - } - // int ago = numb_models > 1 ? 0 : neighbor->ago; int ago = neighbor->ago; if (numb_models > 1) { @@ -540,7 +536,6 @@ static bool is_key(const string& input) { keys.push_back("fparam_from_compute"); keys.push_back("aparam_from_compute"); keys.push_back("charge_spin"); - keys.push_back("charge_spin_from_compute"); keys.push_back("ttm"); keys.push_back("atomic"); keys.push_back("relative"); @@ -708,17 +703,6 @@ void PairDeepMD::settings(int narg, char** arg) { charge_spin.push_back(atof(arg[iarg + 1 + ii])); } iarg += 1 + dim_chg_spin; - } else if (string(arg[iarg]) == string("charge_spin_from_compute")) { - for (int ii = 0; ii < 1; ++ii) { - if (iarg + 1 + ii >= narg || is_key(arg[iarg + 1 + ii])) { - error->all(FLERR, - "invalid charge_spin_from_compute key: should be " - "charge_spin_from_compute compute_charge_spin_id(str)"); - } - } - do_compute_charge_spin = true; - compute_charge_spin_id = arg[iarg + 1]; - iarg += 1 + 1; } else if (string(arg[iarg]) == string("atomic")) { out_each = 1; iarg += 1; @@ -758,11 +742,6 @@ void PairDeepMD::settings(int narg, char** arg) { FLERR, "fparam and fparam_from_compute should NOT be set simultaneously"); } - if (do_compute_charge_spin && charge_spin.size() > 0) { - error->all(FLERR, - "charge_spin and charge_spin_from_compute should NOT be set " - "simultaneously"); - } if (comm->me == 0) { if (numb_models > 1 && out_freq > 0) { From 1c7252d4d6fdb4e72cbb9b9308a84840221771f5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 10:29:25 +0000 Subject: [PATCH 16/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_c/include/deepmd.hpp | 17 +++++++------ source/api_c/src/c_api.cc | 42 ++++++++++++++++----------------- 2 files changed, 29 insertions(+), 30 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 9e3e7256bf..07278054b0 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -1306,10 +1306,10 @@ class DeepPot : public DeepBaseModel { const double* charge_spin__ = (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; - _DP_DeepPotComputeNList( - dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, ener_, force_, virial_, nullptr, nullptr, - charge_spin__); + _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, + box_, nghost, lmp_list.nl, ago, fparam__, + aparam__, ener_, force_, virial_, + nullptr, nullptr, charge_spin__); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -1384,11 +1384,10 @@ class DeepPot : public DeepBaseModel { const double* charge_spin__ = (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; - _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, - box_, nghost, lmp_list.nl, ago, fparam__, - aparam__, ener_, force_, virial_, - atomic_ener_, atomic_virial_, - charge_spin__); + _DP_DeepPotComputeNList( + dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, + fparam__, aparam__, ener_, force_, virial_, atomic_ener_, + atomic_virial_, charge_spin__); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index b3ef76f3a2..cc9bf3b146 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -425,23 +425,24 @@ template void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, float* atomic_virial); template -inline void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, - const int nframes, - const int natoms, - const VALUETYPE* coord, - const int* atype, - const VALUETYPE* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const VALUETYPE* fparam, - const VALUETYPE* aparam, - double* energy, - VALUETYPE* force, - VALUETYPE* virial, - VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial, - const double* charge_spin = nullptr) { +inline void DP_DeepPotComputeNList_variant( + DP_DeepPot* dp, + const int nframes, + const int natoms, + const VALUETYPE* coord, + const int* atype, + const VALUETYPE* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const VALUETYPE* fparam, + const VALUETYPE* aparam, + double* energy, + VALUETYPE* force, + VALUETYPE* virial, + VALUETYPE* atomic_energy, + VALUETYPE* atomic_virial, + const double* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector atype_(atype, atype + natoms); @@ -475,10 +476,9 @@ inline void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, dp, dp->dp.compute(e, f, v, ae, av, coord_, atype_, cell_, nghost, nlist->nl, ago, fparam_, aparam_, charge_spin_)); } else { - DP_REQUIRES_OK(dp, - dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_, - charge_spin_)); + DP_REQUIRES_OK( + dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, nlist->nl, + ago, fparam_, aparam_, charge_spin_)); } // copy from C++ vectors to C arrays, if not NULL pointer if (energy) { From f178f2b9c7b1e8a7ae10325facacb8284f60b238 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Jun 2026 10:05:45 +0800 Subject: [PATCH 17/24] feat: wire charge_spin through the C API model-deviation path Completes the C API charge_spin support so model-deviation in the common LAMMPS build (DP_USING_C_API) no longer silently drops charge_spin: - Add C API entry points: DP_DeepPotModelDeviGetDimChgSpin and the version-3 model-deviation compute functions DP_DeepPotModelDeviCompute3 / *f3 / ComputeNList3 / *f3, forwarding the per-frame charge_spin to api_cc::DeepPotModelDevi::compute. The internal model-devi variant helpers gained an optional charge_spin arg, so the existing version-2 functions are unchanged. - deepmd.hpp DeepPotModelDevi now reads dim_chg_spin via DP_DeepPotModelDeviGetDimChgSpin and threads charge_spin through to the version-3 entry points; empty charge_spin / non-charge_spin models keep the version-2 path for compatibility with an older libdeepmd_c. Co-Authored-By: Claude Opus 4.8 --- source/api_c/include/c_api.h | 104 +++++++++++++++++++++++++ source/api_c/include/deepmd.hpp | 101 +++++++++++++++++------- source/api_c/src/c_api.cc | 131 +++++++++++++++++++++++++++++--- 3 files changed, 296 insertions(+), 40 deletions(-) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 78d193de64..0d626b3725 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -1463,6 +1463,102 @@ void DP_DeepPotModelDeviComputeNListf2(DP_DeepPotModelDevi* dp, float* atomic_energy, float* atomic_virial); +/** + * @brief Evaluate energy, force and virial with a DP model deviation. + *(double version, with charge_spin) + * @version 3 + * @param[in] charge_spin The per-frame charge/spin input. The array should be + *of size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @since API version 27 + **/ +void DP_DeepPotModelDeviCompute3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate energy, force and virial with a DP model deviation. + *(float version, with charge_spin) + * @version 3 + * @since API version 27 + **/ +void DP_DeepPotModelDeviComputef3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial); + +/** + * @brief Evaluate energy, force and virial with a DP model deviation and a + *neighbor list. (double version, with charge_spin) + * @version 3 + * @param[in] charge_spin The per-frame charge/spin input. The array should be + *of size nframes x dim_chg_spin. Pass NULL to use the model's stored + *default_chg_spin. + * @since API version 27 + **/ +void DP_DeepPotModelDeviComputeNList3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial); + +/** + * @brief Evaluate energy, force and virial with a DP model deviation and a + *neighbor list. (float version, with charge_spin) + * @version 3 + * @since API version 27 + **/ +void DP_DeepPotModelDeviComputeNListf3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial); + /** * @brief Evaluate the energy, force, magnetic force and virial by using a DP *spin model deviation with neighbor list. (float version) @@ -1756,6 +1852,14 @@ int DP_DeepPotModelDeviGetDimFParam(DP_DeepPotModelDevi* dp); */ int DP_DeepPotModelDeviGetDimAParam(DP_DeepPotModelDevi* dp); +/** + * @brief Get the dimension of the charge/spin input of a DP Model Deviation. + * @param[in] dp The DP Model Deviation to use. + * @return The dimension of the charge/spin input (0 if none). + * @since API version 27 + */ +int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp); + /** * @brief Check whether the atomic dimension of atomic parameters is nall * instead of nloc. diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 07278054b0..e16572e12b 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -383,7 +383,8 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial); + FPTYPE* atomic_virial, + const double* charge_spin); template <> inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, @@ -397,10 +398,17 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, double* force, double* virial, double* atomic_energy, - double* atomic_virial) { - DP_DeepPotModelDeviCompute2(dp, 1, natom, coord, atype, cell, fparam, aparam, - energy, force, virial, atomic_energy, - atomic_virial); + double* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotModelDeviCompute3(dp, 1, natom, coord, atype, cell, fparam, aparam, + charge_spin, energy, force, virial, + atomic_energy, atomic_virial); + } else { + DP_DeepPotModelDeviCompute2(dp, 1, natom, coord, atype, cell, fparam, aparam, + energy, force, virial, atomic_energy, + atomic_virial); + } } template <> @@ -415,10 +423,17 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial) { - DP_DeepPotModelDeviComputef2(dp, 1, natom, coord, atype, cell, fparam, aparam, - energy, force, virial, atomic_energy, - atomic_virial); + float* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotModelDeviComputef3(dp, 1, natom, coord, atype, cell, fparam, + aparam, charge_spin, energy, force, virial, + atomic_energy, atomic_virial); + } else { + DP_DeepPotModelDeviComputef2(dp, 1, natom, coord, atype, cell, fparam, + aparam, energy, force, virial, atomic_energy, + atomic_virial); + } } template @@ -492,7 +507,8 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial); + FPTYPE* atomic_virial, + const double* charge_spin); template <> inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, @@ -509,10 +525,18 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, double* force, double* virial, double* atomic_energy, - double* atomic_virial) { - DP_DeepPotModelDeviComputeNList2(dp, 1, natom, coord, atype, cell, nghost, - nlist, ago, fparam, aparam, energy, force, - virial, atomic_energy, atomic_virial); + double* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotModelDeviComputeNList3(dp, 1, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, charge_spin, + energy, force, virial, atomic_energy, + atomic_virial); + } else { + DP_DeepPotModelDeviComputeNList2(dp, 1, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, energy, force, + virial, atomic_energy, atomic_virial); + } } template <> @@ -530,10 +554,18 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial) { - DP_DeepPotModelDeviComputeNListf2(dp, 1, natom, coord, atype, cell, nghost, - nlist, ago, fparam, aparam, energy, force, - virial, atomic_energy, atomic_virial); + float* atomic_virial, + const double* charge_spin) { + if (charge_spin) { + DP_DeepPotModelDeviComputeNListf3(dp, 1, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, charge_spin, + energy, force, virial, atomic_energy, + atomic_virial); + } else { + DP_DeepPotModelDeviComputeNListf2(dp, 1, natom, coord, atype, cell, nghost, + nlist, ago, fparam, aparam, energy, force, + virial, atomic_energy, atomic_virial); + } } template @@ -2160,11 +2192,19 @@ class DeepPotModelDevi : public DeepBaseModelDevi { numb_models = models.size(); dfparam = DP_DeepPotModelDeviGetDimFParam(dp); daparam = DP_DeepPotModelDeviGetDimAParam(dp); + dchgspin = DP_DeepPotModelDeviGetDimChgSpin(dp); aparam_nall = DP_DeepPotModelDeviIsAParamNAll(dp); has_default_fparam_ = DP_DeepPotModelDeviHasDefaultFParam(dp); dpbase = (DP_DeepBaseModelDevi*)dp; }; + /** + * @brief Get the dimension of the charge/spin input. + * @return The dimension of the charge/spin input (0 if the models have no + *charge/spin embedding). + **/ + int dim_chg_spin() const { return dchgspin; } + /** * @brief Evaluate the energy, force and virial by using this DP model *deviation. @@ -2196,7 +2236,6 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2222,10 +2261,13 @@ class DeepPotModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotModelDeviCompute(dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, - virial_, nullptr, nullptr); + virial_, nullptr, nullptr, + charge_spin__); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2279,7 +2321,6 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2309,10 +2350,12 @@ class DeepPotModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotModelDeviCompute( dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, - virial_, atomic_ener_, atomic_virial_); + virial_, atomic_ener_, atomic_virial_, charge_spin__); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2378,8 +2421,6 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { - // charge_spin is not supported via the C-API model-deviation path. - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2408,10 +2449,12 @@ class DeepPotModelDevi : public DeepBaseModelDevi { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, - aparam__, ener_, force_, virial_, nullptr, nullptr); + aparam__, ener_, force_, virial_, nullptr, nullptr, charge_spin__); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2470,8 +2513,6 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), const std::vector& charge_spin = std::vector()) { - // charge_spin is not supported via the C-API model-deviation path. - (void)charge_spin; unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2504,10 +2545,13 @@ class DeepPotModelDevi : public DeepBaseModelDevi { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const double* charge_spin__ = + (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, - aparam__, ener_, force_, virial_, atomic_ener_, atomic_virial_); + aparam__, ener_, force_, virial_, atomic_ener_, atomic_virial_, + charge_spin__); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2539,6 +2583,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { private: DP_DeepPotModelDevi* dp; + int dchgspin; }; class DeepSpinModelDevi : public DeepBaseModelDevi { diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index cc9bf3b146..b1dfc4a809 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -740,7 +740,8 @@ void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, VALUETYPE* force, VALUETYPE* virial, VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { + VALUETYPE* atomic_virial, + const double* charge_spin = nullptr) { if (nframes > 1) { throw std::runtime_error("nframes > 1 not supported yet"); } @@ -760,16 +761,22 @@ void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, if (aparam) { aparam_.assign(aparam, aparam + nframes * natoms * dp->daparam); } + // charge_spin is always double (it is a per-frame physical input) + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } // different from DeepPot std::vector e; std::vector> f, v, ae, av; if (atomic_energy || atomic_virial) { DP_REQUIRES_OK(dp, dp->dp.compute(e, f, v, ae, av, coord_, atype_, cell_, - fparam_, aparam_)); + fparam_, aparam_, charge_spin_)); } else { - DP_REQUIRES_OK( - dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, fparam_, aparam_)); + DP_REQUIRES_OK(dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, fparam_, + aparam_, charge_spin_)); } // 2D vector to 2D array, flatten first if (energy) { @@ -810,7 +817,8 @@ template void DP_DeepPotModelDeviCompute_variant( double* force, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, const int nframes, @@ -824,7 +832,8 @@ template void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, float* force, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const double* charge_spin); template void DP_DeepSpinModelDeviCompute_variant(DP_DeepSpinModelDevi* dp, @@ -954,7 +963,8 @@ void DP_DeepPotModelDeviComputeNList_variant(DP_DeepPotModelDevi* dp, VALUETYPE* force, VALUETYPE* virial, VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial) { + VALUETYPE* atomic_virial, + const double* charge_spin = nullptr) { if (nframes > 1) { throw std::runtime_error("nframes > 1 not supported yet"); } @@ -976,6 +986,12 @@ void DP_DeepPotModelDeviComputeNList_variant(DP_DeepPotModelDevi* dp, aparam, aparam + (dp->aparam_nall ? natoms : (natoms - nghost)) * dp->daparam); } + // charge_spin is always double (it is a per-frame physical input) + std::vector charge_spin_; + if (charge_spin) { + charge_spin_.assign(charge_spin, + charge_spin + nframes * dp->dp.dim_chg_spin()); + } // different from DeepPot std::vector e; std::vector> f, v, ae, av; @@ -983,10 +999,12 @@ void DP_DeepPotModelDeviComputeNList_variant(DP_DeepPotModelDevi* dp, if (atomic_energy || atomic_virial) { DP_REQUIRES_OK( dp, dp->dp.compute(e, f, v, ae, av, coord_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_)); + nlist->nl, ago, fparam_, aparam_, charge_spin_)); } else { - DP_REQUIRES_OK(dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_)); + DP_REQUIRES_OK(dp, + dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, + nlist->nl, ago, fparam_, aparam_, + charge_spin_)); } // 2D vector to 2D array, flatten first if (energy) { @@ -1030,7 +1048,8 @@ template void DP_DeepPotModelDeviComputeNList_variant( double* force, double* virial, double* atomic_energy, - double* atomic_virial); + double* atomic_virial, + const double* charge_spin); template void DP_DeepPotModelDeviComputeNList_variant( DP_DeepPotModelDevi* dp, @@ -1048,7 +1067,8 @@ template void DP_DeepPotModelDeviComputeNList_variant( float* force, float* virial, float* atomic_energy, - float* atomic_virial); + float* atomic_virial, + const double* charge_spin); // support spin multi model. template @@ -2074,6 +2094,89 @@ void DP_DeepPotModelDeviComputeNListf2(DP_DeepPotModelDevi* dp, aparam, energy, force, virial, atomic_energy, atomic_virial); } +// charge_spin-aware model-deviation variants (version 3). +void DP_DeepPotModelDeviCompute3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepPotModelDeviCompute_variant( + dp, nframes, natoms, coord, atype, cell, fparam, aparam, energy, force, + virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepPotModelDeviComputef3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepPotModelDeviCompute_variant( + dp, nframes, natoms, coord, atype, cell, fparam, aparam, energy, force, + virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepPotModelDeviComputeNList3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const double* coord, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { + DP_DeepPotModelDeviComputeNList_variant( + dp, nframes, natoms, coord, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, virial, atomic_energy, atomic_virial, charge_spin); +} + +void DP_DeepPotModelDeviComputeNListf3(DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const float* fparam, + const float* aparam, + const double* charge_spin, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { + DP_DeepPotModelDeviComputeNList_variant( + dp, nframes, natoms, coord, atype, cell, nghost, nlist, ago, fparam, + aparam, energy, force, virial, atomic_energy, atomic_virial, charge_spin); +} + void DP_DeepSpinModelDeviComputeNListf2(DP_DeepSpinModelDevi* dp, const int nframes, const int natoms, @@ -2231,6 +2334,10 @@ int DP_DeepPotModelDeviGetDimAParam(DP_DeepPotModelDevi* dp) { static_cast(dp)); } +int DP_DeepPotModelDeviGetDimChgSpin(DP_DeepPotModelDevi* dp) { + return dp->dp.dim_chg_spin(); +} + bool DP_DeepPotModelDeviIsAParamNAll(DP_DeepPotModelDevi* dp) { return DP_DeepBaseModelDeviIsAParamNAll( static_cast(dp)); From 2ff72cbd159f7efd35c40819330f568875746e9a Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Jun 2026 10:13:18 +0800 Subject: [PATCH 18/24] fix: address CodeRabbit review on charge_spin (.pth / C API / tests) - 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 --- source/api_c/include/deepmd.hpp | 62 ++++++++++++++----- source/api_cc/include/DeepPotPT.h | 5 ++ source/api_cc/src/DeepPotPT.cc | 26 ++++++++ .../api_cc/tests/test_deeppot_chg_spin_pt.cc | 15 +++++ source/tests/infer/gen_chg_spin.py | 8 ++- 5 files changed, 101 insertions(+), 15 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index e16572e12b..f3b9a4c5c5 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -864,6 +864,33 @@ inline double* _DP_Get_Energy_Pointer(double& vec, const int nframes) { namespace deepmd { namespace hpp { +/** + * @brief Validate a runtime charge_spin vector and return a raw pointer to it + * (or nullptr when empty, which falls back to the model's default_chg_spin). + * @param[in] charge_spin The per-frame charge/spin values. + * @param[in] dchgspin The model's charge/spin dimension (0 if unsupported). + * @param[in] nframes The number of frames. + * @return charge_spin.data() if non-empty and valid, otherwise nullptr. + */ +inline const double* validate_charge_spin( + const std::vector& charge_spin, + const int dchgspin, + const unsigned int nframes) { + if (charge_spin.empty()) { + return nullptr; + } + if (dchgspin == 0) { + throw deepmd::hpp::deepmd_exception( + "charge_spin was provided, but this model does not support " + "charge/spin conditioning"); + } + if (charge_spin.size() != static_cast(nframes) * dchgspin) { + throw deepmd::hpp::deepmd_exception( + "the dim of charge_spin provided is not consistent with what the " + "model uses"); + } + return charge_spin.data(); +} /** * @brief Neighbor list. **/ @@ -1091,7 +1118,7 @@ class DeepPot : public DeepBaseModel { /** * @brief DP constructor without initialization. **/ - DeepPot() : dp(nullptr) {}; + DeepPot() : dp(nullptr), dchgspin(0) {}; ~DeepPot() { DP_DeleteDeepPot(dp); }; /** * @brief DP constructor with initialization. @@ -1102,7 +1129,7 @@ class DeepPot : public DeepBaseModel { DeepPot(const std::string& model, const int& gpu_rank = 0, const std::string& file_content = "") - : dp(nullptr) { + : dp(nullptr), dchgspin(0) { try { init(model, gpu_rank, file_content); } catch (...) { @@ -1144,7 +1171,10 @@ class DeepPot : public DeepBaseModel { * @return The dimension of the charge/spin input (0 if the model has no *charge/spin embedding). **/ - int dim_chg_spin() const { return dchgspin; } + int dim_chg_spin() const { + assert(dp); + return dchgspin; + } /** * @brief Evaluate the energy, force and virial by using this DP. @@ -1199,7 +1229,7 @@ class DeepPot : public DeepBaseModel { // charge_spin routes to the version-3 C API; nullptr keeps version-2 so // non-charge_spin models still work against an older libdeepmd_c. const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, virial_, @@ -1267,7 +1297,7 @@ class DeepPot : public DeepBaseModel { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, virial_, @@ -1336,7 +1366,7 @@ class DeepPot : public DeepBaseModel { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, @@ -1414,7 +1444,7 @@ class DeepPot : public DeepBaseModel { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotComputeNList( dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, @@ -2137,13 +2167,14 @@ class DeepPotModelDevi : public DeepBaseModelDevi { /** * @brief DP model deviation constructor without initialization. **/ - DeepPotModelDevi() : dp(nullptr) {}; + DeepPotModelDevi() : dp(nullptr), dchgspin(0) {}; ~DeepPotModelDevi() { DP_DeleteDeepPotModelDevi(dp); }; /** * @brief DP model deviation constructor with initialization. * @param[in] models The names of the frozen model file. **/ - DeepPotModelDevi(const std::vector& models) : dp(nullptr) { + DeepPotModelDevi(const std::vector& models) + : dp(nullptr), dchgspin(0) { try { init(models); } catch (...) { @@ -2203,7 +2234,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi { * @return The dimension of the charge/spin input (0 if the models have no *charge/spin embedding). **/ - int dim_chg_spin() const { return dchgspin; } + int dim_chg_spin() const { + assert(dp); + return dchgspin; + } /** * @brief Evaluate the energy, force and virial by using this DP model @@ -2262,7 +2296,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviCompute(dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, @@ -2351,7 +2385,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviCompute( dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, @@ -2450,7 +2484,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, @@ -2546,7 +2580,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; const double* charge_spin__ = - (dchgspin > 0 && !charge_spin.empty()) ? charge_spin.data() : nullptr; + validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h index cd57fc79ac..deb487610f 100644 --- a/source/api_cc/include/DeepPotPT.h +++ b/source/api_cc/include/DeepPotPT.h @@ -315,6 +315,11 @@ class DeepPotPT : public DeepPotBackend { const std::vector& fparam, const std::vector& aparam, const bool atomic); + // Keep the base charge_spin-aware computew / computew_mixed_type overloads + // visible: the non-charge_spin declarations below would otherwise hide them + // for callers using the static type DeepPotPT. + using DeepPotBackend::computew; + using DeepPotBackend::computew_mixed_type; void computew_mixed_type(std::vector& ener, std::vector& force, std::vector& virial, diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 4878491c69..d29783089c 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -269,6 +269,13 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, if (dchgspin > 0) { auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); if (!charge_spin.empty()) { + // Single-frame path: charge_spin must hold exactly dim_chg_spin values. + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } charge_spin_tensor = torch::from_blob(const_cast(charge_spin.data()), {1, static_cast(charge_spin.size())}, @@ -276,6 +283,12 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, .clone() .to(device); } else if (!default_chg_spin_.empty()) { + if (static_cast(default_chg_spin_.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "default_chg_spin has " + std::to_string(default_chg_spin_.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } charge_spin_tensor = torch::from_blob(const_cast(default_chg_spin_.data()), {1, dchgspin}, dbl_options) @@ -466,6 +479,13 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, auto dbl_options = torch::TensorOptions().dtype(torch::kFloat64); c10::optional charge_spin_tensor; if (!charge_spin.empty()) { + // Single-frame path: charge_spin must hold exactly dim_chg_spin values. + if (static_cast(charge_spin.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "charge_spin has " + std::to_string(charge_spin.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } charge_spin_tensor = torch::from_blob(const_cast(charge_spin.data()), {1, static_cast(charge_spin.size())}, @@ -473,6 +493,12 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, .clone() .to(device); } else if (!default_chg_spin_.empty()) { + if (static_cast(default_chg_spin_.size()) != dchgspin) { + throw deepmd::deepmd_exception( + "default_chg_spin has " + std::to_string(default_chg_spin_.size()) + + " values but the model expects dim_chg_spin=" + + std::to_string(dchgspin) + "."); + } charge_spin_tensor = torch::from_blob(const_cast(default_chg_spin_.data()), {1, dchgspin}, dbl_options) diff --git a/source/api_cc/tests/test_deeppot_chg_spin_pt.cc b/source/api_cc/tests/test_deeppot_chg_spin_pt.cc index efa39b3304..296ee3023c 100644 --- a/source/api_cc/tests/test_deeppot_chg_spin_pt.cc +++ b/source/api_cc/tests/test_deeppot_chg_spin_pt.cc @@ -10,6 +10,7 @@ #include #include +#include #include #include "DeepPot.h" @@ -79,6 +80,13 @@ class TestInferDeepPotChgSpinPt : public ::testing::Test { expected_tot_e_default, expected_tot_v_default); natoms = e_exp.size(); + { + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (e.g. .pth export was skipped)."; + } + } dp.init(kModelPath); } @@ -200,6 +208,13 @@ class TestInferDeepPotChgSpinPtNoPbc : public ::testing::Test { expected_tot_e_default, expected_tot_v_default); natoms = e_exp.size(); + { + std::ifstream model_file(kModelPath); + if (!model_file.good()) { + GTEST_SKIP() << "Skip because " << kModelPath + << " was not generated (e.g. .pth export was skipped)."; + } + } dp.init(kModelPath); } diff --git a/source/tests/infer/gen_chg_spin.py b/source/tests/infer/gen_chg_spin.py index ba43f00b12..8db0659106 100644 --- a/source/tests/infer/gen_chg_spin.py +++ b/source/tests/infer/gen_chg_spin.py @@ -88,9 +88,15 @@ def main(): pt_expt_deserialize_to_file(pt2_path, copy.deepcopy(data), do_atomic_virial=True) pth_path = os.path.join(base_dir, "chg_spin.pth") + # Remove any stale .pth first so a failed export below cannot leave an old + # artifact that the parity check would then validate against. + if os.path.exists(pth_path): + os.remove(pth_path) + pth_generated = False print(f"Exporting to {pth_path} ...") # noqa: T201 try: pt_deserialize_to_file(pth_path, copy.deepcopy(data)) + pth_generated = True except RuntimeError as e: # Custom ops may not be available in all build environments; .pth # generation is not critical (the .pth test skips if the file is @@ -194,7 +200,7 @@ def main(): print(f"Wrote {ref_path}") # noqa: T201 # ---- Verify .pth reproduces the .pt2 reference (PBC + NoPbc) ---- - if os.path.exists(pth_path): + if pth_generated: dp_pth = DeepPot(pth_path) assert dp_pth.deep_eval.get_dim_chg_spin() == 2 tol = 1e-10 From ad66793512322857cd89c19cbb6da005fffefcb4 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Jun 2026 10:16:05 +0800 Subject: [PATCH 19/24] test: add LAMMPS integration test for the charge_spin pair_style keyword 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 --- source/lmp/tests/test_lammps_chg_spin_pt2.py | 131 +++++++++++++++++++ 1 file changed, 131 insertions(+) create mode 100644 source/lmp/tests/test_lammps_chg_spin_pt2.py diff --git a/source/lmp/tests/test_lammps_chg_spin_pt2.py b/source/lmp/tests/test_lammps_chg_spin_pt2.py new file mode 100644 index 0000000000..742e475cda --- /dev/null +++ b/source/lmp/tests/test_lammps_chg_spin_pt2.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Test the LAMMPS ``charge_spin`` pair_style keyword (.pt2 DPA3 model). + +Exercises the parsing and wiring of the ``charge_spin`` keyword through to the +backend: the model (chg_spin.pt2, DPA3 with add_chg_spin_ebd=True and +default_chg_spin=[0.0, 1.0]) is generated by source/tests/infer/gen_chg_spin.py, +together with the reference values in chg_spin.expected. +""" + +import os +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from expected_ref import ( + read_expected_ref, +) +from lammps import ( + PyLammps, +) +from write_lmp_data import ( + write_lmp_data, +) + +pt2_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "chg_spin.pt2" +ref_file = ( + Path(__file__).parent.parent.parent / "tests" / "infer" / "chg_spin.expected" +) +data_file = Path(__file__).parent / "data_chg_spin_pt2.lmp" + +# Reference values written by source/tests/infer/gen_chg_spin.py (PBC case): +# "default" -- no charge_spin keyword -> stored default_chg_spin [0.0, 1.0] +# "explicit" -- charge_spin 1.0 2.0 +# Guarded with try/except because gen_chg_spin.py only runs when PyTorch is +# built; matrices that disable PyTorch still load this file at collection time. +try: + _ref = read_expected_ref(ref_file) + expected_e_default = float(np.sum(_ref["default"]["expected_e"])) + expected_f_default = _ref["default"]["expected_f"].reshape(6, 3) + expected_e_explicit = float(np.sum(_ref["explicit"]["expected_e"])) + expected_f_explicit = _ref["explicit"]["expected_f"].reshape(6, 3) +except FileNotFoundError: + expected_e_default = expected_f_default = None + expected_e_explicit = expected_f_explicit = None + +box = np.array([0, 13, 0, 13, 0, 13, 0, 0, 0]) +coord = np.array( + [ + [12.83, 2.56, 2.18], + [12.09, 2.87, 2.74], + [0.25, 3.32, 1.68], + [3.36, 3.00, 1.81], + [3.51, 2.51, 2.60], + [4.27, 3.22, 1.56], + ] +) +# atype [0, 1, 1, 0, 1, 1] (O, H, H, O, H, H) -> LAMMPS types [1, 2, 2, 1, 2, 2] +type_OH = np.array([1, 2, 2, 1, 2, 2]) + + +def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("Skip test because PyTorch support is not enabled.") + write_lmp_data(box, coord, type_OH, data_file) + + +def teardown_module() -> None: + if data_file.exists(): + os.remove(data_file) + + +def _lammps(data_file, units="metal") -> PyLammps: + lammps = PyLammps() + lammps.units(units) + lammps.boundary("p p p") + lammps.atom_style("atomic") + # DPA3 (message passing) resolves ghost atoms via the LAMMPS atom-map. + lammps.atom_modify("map yes") + lammps.neighbor("2.0 bin") + lammps.neigh_modify("every 10 delay 0 check no") + lammps.read_data(data_file.resolve()) + lammps.mass("1 16") + lammps.mass("2 2") + lammps.timestep(0.0005) + lammps.fix("1 all nve") + return lammps + + +@pytest.fixture +def lammps(): + if not pt2_file.is_file(): + pytest.skip(f"Skip because {pt2_file} was not generated.") + lmp = _lammps(data_file=data_file) + yield lmp + lmp.close() + + +def test_pair_deepmd_charge_spin_default(lammps) -> None: + """No charge_spin keyword -> the model's stored default_chg_spin is used.""" + lammps.pair_style(f"deepmd {pt2_file.resolve()}") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_default) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_default[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +def test_pair_deepmd_charge_spin_explicit(lammps) -> None: + """Explicit ``charge_spin`` keyword is parsed and threaded to the model.""" + lammps.pair_style(f"deepmd {pt2_file.resolve()} charge_spin 1.0 2.0") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") == pytest.approx(expected_e_explicit) + for ii in range(6): + assert lammps.atoms[ii].force == pytest.approx( + expected_f_explicit[lammps.atoms[ii].id - 1] + ) + lammps.run(1) + + +def test_charge_spin_changes_result(lammps) -> None: + """Different charge_spin must give a different energy (keyword takes effect).""" + lammps.pair_style(f"deepmd {pt2_file.resolve()} charge_spin 1.0 2.0") + lammps.pair_coeff("* *") + lammps.run(0) + assert lammps.eval("pe") != pytest.approx(expected_e_default) From 55955c35021bd27a8fa76764b9fc2a29cab8982f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Jun 2026 02:29:25 +0000 Subject: [PATCH 20/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_c/include/deepmd.hpp | 62 +++++++++-------- source/api_c/src/c_api.cc | 71 ++++++++++---------- source/lmp/tests/test_lammps_chg_spin_pt2.py | 4 +- 3 files changed, 67 insertions(+), 70 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index f3b9a4c5c5..1dc0eae8a4 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -401,12 +401,12 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, double* atomic_virial, const double* charge_spin) { if (charge_spin) { - DP_DeepPotModelDeviCompute3(dp, 1, natom, coord, atype, cell, fparam, aparam, - charge_spin, energy, force, virial, + DP_DeepPotModelDeviCompute3(dp, 1, natom, coord, atype, cell, fparam, + aparam, charge_spin, energy, force, virial, atomic_energy, atomic_virial); } else { - DP_DeepPotModelDeviCompute2(dp, 1, natom, coord, atype, cell, fparam, aparam, - energy, force, virial, atomic_energy, + DP_DeepPotModelDeviCompute2(dp, 1, natom, coord, atype, cell, fparam, + aparam, energy, force, virial, atomic_energy, atomic_virial); } } @@ -511,27 +511,27 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, const double* charge_spin); template <> -inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, - const int natom, - const double* coord, - const int* atype, - const double* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const double* fparam, - const double* aparam, - double* energy, - double* force, - double* virial, - double* atomic_energy, - double* atomic_virial, - const double* charge_spin) { +inline void _DP_DeepPotModelDeviComputeNList( + DP_DeepPotModelDevi* dp, + const int natom, + const double* coord, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial, + const double* charge_spin) { if (charge_spin) { - DP_DeepPotModelDeviComputeNList3(dp, 1, natom, coord, atype, cell, nghost, - nlist, ago, fparam, aparam, charge_spin, - energy, force, virial, atomic_energy, - atomic_virial); + DP_DeepPotModelDeviComputeNList3( + dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, + charge_spin, energy, force, virial, atomic_energy, atomic_virial); } else { DP_DeepPotModelDeviComputeNList2(dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, @@ -557,10 +557,9 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, float* atomic_virial, const double* charge_spin) { if (charge_spin) { - DP_DeepPotModelDeviComputeNListf3(dp, 1, natom, coord, atype, cell, nghost, - nlist, ago, fparam, aparam, charge_spin, - energy, force, virial, atomic_energy, - atomic_virial); + DP_DeepPotModelDeviComputeNListf3( + dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, + charge_spin, energy, force, virial, atomic_energy, atomic_virial); } else { DP_DeepPotModelDeviComputeNListf2(dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, @@ -2298,10 +2297,9 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const double* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); - _DP_DeepPotModelDeviCompute(dp, natoms, coord_, atype_, box_, - fparam__, aparam__, ener_, force_, - virial_, nullptr, nullptr, - charge_spin__); + _DP_DeepPotModelDeviCompute( + dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, + virial_, nullptr, nullptr, charge_spin__); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index b1dfc4a809..3b9da06555 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -820,20 +820,21 @@ template void DP_DeepPotModelDeviCompute_variant( double* atomic_virial, const double* charge_spin); -template void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, - const int nframes, - const int natoms, - const float* coord, - const int* atype, - const float* cell, - const float* fparam, - const float* aparam, - double* energy, - float* force, - float* virial, - float* atomic_energy, - float* atomic_virial, - const double* charge_spin); +template void DP_DeepPotModelDeviCompute_variant( + DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial, + const double* charge_spin); template void DP_DeepSpinModelDeviCompute_variant(DP_DeepSpinModelDevi* dp, @@ -948,23 +949,24 @@ template void DP_DeepSpinModelDeviCompute_variant( float* atomic_virial); template -void DP_DeepPotModelDeviComputeNList_variant(DP_DeepPotModelDevi* dp, - const int nframes, - const int natoms, - const VALUETYPE* coord, - const int* atype, - const VALUETYPE* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const VALUETYPE* fparam, - const VALUETYPE* aparam, - double* energy, - VALUETYPE* force, - VALUETYPE* virial, - VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial, - const double* charge_spin = nullptr) { +void DP_DeepPotModelDeviComputeNList_variant( + DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const VALUETYPE* coord, + const int* atype, + const VALUETYPE* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const VALUETYPE* fparam, + const VALUETYPE* aparam, + double* energy, + VALUETYPE* force, + VALUETYPE* virial, + VALUETYPE* atomic_energy, + VALUETYPE* atomic_virial, + const double* charge_spin = nullptr) { if (nframes > 1) { throw std::runtime_error("nframes > 1 not supported yet"); } @@ -1001,10 +1003,9 @@ void DP_DeepPotModelDeviComputeNList_variant(DP_DeepPotModelDevi* dp, dp, dp->dp.compute(e, f, v, ae, av, coord_, atype_, cell_, nghost, nlist->nl, ago, fparam_, aparam_, charge_spin_)); } else { - DP_REQUIRES_OK(dp, - dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, - nlist->nl, ago, fparam_, aparam_, - charge_spin_)); + DP_REQUIRES_OK( + dp, dp->dp.compute(e, f, v, coord_, atype_, cell_, nghost, nlist->nl, + ago, fparam_, aparam_, charge_spin_)); } // 2D vector to 2D array, flatten first if (energy) { diff --git a/source/lmp/tests/test_lammps_chg_spin_pt2.py b/source/lmp/tests/test_lammps_chg_spin_pt2.py index 742e475cda..a995b25061 100644 --- a/source/lmp/tests/test_lammps_chg_spin_pt2.py +++ b/source/lmp/tests/test_lammps_chg_spin_pt2.py @@ -25,9 +25,7 @@ ) pt2_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "chg_spin.pt2" -ref_file = ( - Path(__file__).parent.parent.parent / "tests" / "infer" / "chg_spin.expected" -) +ref_file = Path(__file__).parent.parent.parent / "tests" / "infer" / "chg_spin.expected" data_file = Path(__file__).parent / "data_chg_spin_pt2.lmp" # Reference values written by source/tests/infer/gen_chg_spin.py (PBC case): From b667537daf53c53d8fff23042bb997a03df13420 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Jun 2026 11:08:07 +0800 Subject: [PATCH 21/24] fix: drop get_dim_chg_spin() call on the .pth DeepEval in gen_chg_spin 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 --- source/tests/infer/gen_chg_spin.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/source/tests/infer/gen_chg_spin.py b/source/tests/infer/gen_chg_spin.py index 8db0659106..327d967af4 100644 --- a/source/tests/infer/gen_chg_spin.py +++ b/source/tests/infer/gen_chg_spin.py @@ -202,7 +202,9 @@ def main(): # ---- Verify .pth reproduces the .pt2 reference (PBC + NoPbc) ---- if pth_generated: dp_pth = DeepPot(pth_path) - assert dp_pth.deep_eval.get_dim_chg_spin() == 2 + # Note: the .pth (pt) DeepEval does not expose get_dim_chg_spin (only + # the .pt2 / pt_expt one does); the energy/force parity below is the + # real check that charge_spin is threaded for the .pth backend. tol = 1e-10 e_def_p, f_def_p, v_def_p = dp_pth.eval(coord, box, atype) e_exp_p, f_exp_p, v_exp_p = dp_pth.eval( From 34b441c95e32aa60bd69ec26d44439e16cf3abe1 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 11 Jun 2026 17:52:59 +0800 Subject: [PATCH 22/24] refactor: type charge_spin as FPTYPE in the C API (iProzd review) 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 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& charge_spin; validate_charge_spin is now a template returning const FPTYPE*. Co-Authored-By: Claude Opus 4.8 --- source/api_c/include/c_api.h | 8 +-- source/api_c/include/deepmd.hpp | 119 ++++++++++++++++---------------- source/api_c/src/c_api.cc | 36 +++++----- 3 files changed, 84 insertions(+), 79 deletions(-) diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h index 0d626b3725..c0c48a55c6 100644 --- a/source/api_c/include/c_api.h +++ b/source/api_c/include/c_api.h @@ -734,7 +734,7 @@ extern void DP_DeepPotComputef3(DP_DeepPot* dp, const float* cell, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, @@ -785,7 +785,7 @@ extern void DP_DeepPotComputeNListf3(DP_DeepPot* dp, const int ago, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, @@ -1501,7 +1501,7 @@ void DP_DeepPotModelDeviComputef3(DP_DeepPotModelDevi* dp, const float* cell, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, @@ -1552,7 +1552,7 @@ void DP_DeepPotModelDeviComputeNListf3(DP_DeepPotModelDevi* dp, const int ago, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 1dc0eae8a4..e12ef4c65c 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -53,12 +53,12 @@ inline void _DP_DeepPotCompute(DP_DeepPot* dp, const FPTYPE* cell, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial, - const double* charge_spin); + FPTYPE* atomic_virial); template <> inline void _DP_DeepPotCompute(DP_DeepPot* dp, @@ -69,12 +69,12 @@ inline void _DP_DeepPotCompute(DP_DeepPot* dp, const double* cell, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* virial, double* atomic_energy, - double* atomic_virial, - const double* charge_spin) { + double* atomic_virial) { // charge_spin == nullptr keeps the version-2 entry point so models without a // charge/spin embedding still work against an older libdeepmd_c. if (charge_spin) { @@ -96,12 +96,12 @@ inline void _DP_DeepPotCompute(DP_DeepPot* dp, const float* cell, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* virial, float* atomic_energy, - float* atomic_virial, - const double* charge_spin) { + float* atomic_virial) { if (charge_spin) { DP_DeepPotComputef3(dp, nframes, natom, coord, atype, cell, fparam, aparam, charge_spin, energy, force, virial, atomic_energy, @@ -184,12 +184,12 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, const int ago, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial, - const double* charge_spin); + FPTYPE* atomic_virial); template <> inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, @@ -203,12 +203,12 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, const int ago, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* virial, double* atomic_energy, - double* atomic_virial, - const double* charge_spin) { + double* atomic_virial) { if (charge_spin) { DP_DeepPotComputeNList3(dp, nframes, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, charge_spin, energy, @@ -232,12 +232,12 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, const int ago, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* virial, float* atomic_energy, - float* atomic_virial, - const double* charge_spin) { + float* atomic_virial) { if (charge_spin) { DP_DeepPotComputeNListf3(dp, nframes, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, charge_spin, energy, @@ -379,12 +379,12 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, const FPTYPE* cell, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial, - const double* charge_spin); + FPTYPE* atomic_virial); template <> inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, @@ -394,12 +394,12 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, const double* cell, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* virial, double* atomic_energy, - double* atomic_virial, - const double* charge_spin) { + double* atomic_virial) { if (charge_spin) { DP_DeepPotModelDeviCompute3(dp, 1, natom, coord, atype, cell, fparam, aparam, charge_spin, energy, force, virial, @@ -419,12 +419,12 @@ inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, const float* cell, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* virial, float* atomic_energy, - float* atomic_virial, - const double* charge_spin) { + float* atomic_virial) { if (charge_spin) { DP_DeepPotModelDeviComputef3(dp, 1, natom, coord, atype, cell, fparam, aparam, charge_spin, energy, force, virial, @@ -503,12 +503,12 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, const int ago, const FPTYPE* fparam, const FPTYPE* aparam, + const FPTYPE* charge_spin, double* energy, FPTYPE* force, FPTYPE* virial, FPTYPE* atomic_energy, - FPTYPE* atomic_virial, - const double* charge_spin); + FPTYPE* atomic_virial); template <> inline void _DP_DeepPotModelDeviComputeNList( @@ -522,12 +522,12 @@ inline void _DP_DeepPotModelDeviComputeNList( const int ago, const double* fparam, const double* aparam, + const double* charge_spin, double* energy, double* force, double* virial, double* atomic_energy, - double* atomic_virial, - const double* charge_spin) { + double* atomic_virial) { if (charge_spin) { DP_DeepPotModelDeviComputeNList3( dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, @@ -550,12 +550,12 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, const int ago, const float* fparam, const float* aparam, + const float* charge_spin, double* energy, float* force, float* virial, float* atomic_energy, - float* atomic_virial, - const double* charge_spin) { + float* atomic_virial) { if (charge_spin) { DP_DeepPotModelDeviComputeNListf3( dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, @@ -871,10 +871,10 @@ namespace hpp { * @param[in] nframes The number of frames. * @return charge_spin.data() if non-empty and valid, otherwise nullptr. */ -inline const double* validate_charge_spin( - const std::vector& charge_spin, - const int dchgspin, - const unsigned int nframes) { +template +inline const FPTYPE* validate_charge_spin(const std::vector& charge_spin, + const int dchgspin, + const unsigned int nframes) { if (charge_spin.empty()) { return nullptr; } @@ -1204,7 +1204,7 @@ class DeepPot : public DeepBaseModel { const std::vector& box, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1227,12 +1227,12 @@ class DeepPot : public DeepBaseModel { const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; // charge_spin routes to the version-3 C API; nullptr keeps version-2 so // non-charge_spin models still work against an older libdeepmd_c. - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, - fparam__, aparam__, ener_, force_, virial_, - nullptr, nullptr, charge_spin__); + fparam__, aparam__, charge_spin__, ener_, + force_, virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -1269,7 +1269,7 @@ class DeepPot : public DeepBaseModel { const std::vector& box, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1295,12 +1295,13 @@ class DeepPot : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, - fparam__, aparam__, ener_, force_, virial_, - atomic_ener_, atomic_virial_, charge_spin__); + fparam__, aparam__, charge_spin__, ener_, + force_, virial_, atomic_ener_, + atomic_virial_); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; @@ -1340,7 +1341,7 @@ class DeepPot : public DeepBaseModel { const int& ago, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1364,13 +1365,13 @@ class DeepPot : public DeepBaseModel { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, - aparam__, ener_, force_, virial_, - nullptr, nullptr, charge_spin__); + aparam__, charge_spin__, ener_, force_, + virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -1413,7 +1414,7 @@ class DeepPot : public DeepBaseModel { const int& ago, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1442,13 +1443,13 @@ class DeepPot : public DeepBaseModel { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotComputeNList( dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, ener_, force_, virial_, atomic_ener_, - atomic_virial_, charge_spin__); + fparam__, aparam__, charge_spin__, ener_, force_, virial_, atomic_ener_, + atomic_virial_); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -2267,7 +2268,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& box, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. unsigned int natoms = atype.size(); unsigned int nframes = 1; @@ -2294,12 +2295,12 @@ class DeepPotModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviCompute( - dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, - virial_, nullptr, nullptr, charge_spin__); + dp, natoms, coord_, atype_, box_, fparam__, aparam__, charge_spin__, + ener_, force_, virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2351,7 +2352,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const std::vector& box, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. unsigned int natoms = atype.size(); unsigned int nframes = 1; @@ -2382,12 +2383,12 @@ class DeepPotModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviCompute( - dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, - virial_, atomic_ener_, atomic_virial_, charge_spin__); + dp, natoms, coord_, atype_, box_, fparam__, aparam__, charge_spin__, + ener_, force_, virial_, atomic_ener_, atomic_virial_); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2452,7 +2453,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const int& ago, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2481,12 +2482,12 @@ class DeepPotModelDevi : public DeepBaseModelDevi { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, - aparam__, ener_, force_, virial_, nullptr, nullptr, charge_spin__); + aparam__, charge_spin__, ener_, force_, virial_, nullptr, nullptr); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape @@ -2544,7 +2545,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { const int& ago, const std::vector& fparam = std::vector(), const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2577,13 +2578,13 @@ class DeepPotModelDevi : public DeepBaseModelDevi { aparam); const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; - const double* charge_spin__ = + const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, - aparam__, ener_, force_, virial_, atomic_ener_, atomic_virial_, - charge_spin__); + aparam__, charge_spin__, ener_, force_, virial_, atomic_ener_, + atomic_virial_); DP_CHECK_OK(DP_DeepPotModelDeviCheckOK, dp); // reshape diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index 3b9da06555..e648e7e72e 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -252,7 +252,7 @@ inline void DP_DeepPotCompute_variant(DP_DeepPot* dp, VALUETYPE* virial, VALUETYPE* atomic_energy, VALUETYPE* atomic_virial, - const double* charge_spin = nullptr) { + const VALUETYPE* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector atype_(atype, atype + natoms); @@ -269,7 +269,8 @@ inline void DP_DeepPotCompute_variant(DP_DeepPot* dp, if (aparam) { aparam_.assign(aparam, aparam + nframes * natoms * dp->daparam); } - // charge_spin is always double (it is a per-frame physical input) + // charge_spin is converted to double for the api_cc::DeepPot interface + // (the compute backend stores/uses charge_spin as float64). std::vector charge_spin_; if (charge_spin) { charge_spin_.assign(charge_spin, @@ -331,7 +332,7 @@ template void DP_DeepPotCompute_variant(DP_DeepPot* dp, float* virial, float* atomic_energy, float* atomic_virial, - const double* charge_spin); + const float* charge_spin); // support spin template inline void DP_DeepSpinCompute_variant(DP_DeepSpin* dp, @@ -442,7 +443,7 @@ inline void DP_DeepPotComputeNList_variant( VALUETYPE* virial, VALUETYPE* atomic_energy, VALUETYPE* atomic_virial, - const double* charge_spin = nullptr) { + const VALUETYPE* charge_spin = nullptr) { // init C++ vectors from C arrays std::vector coord_(coord, coord + nframes * natoms * 3); std::vector atype_(atype, atype + natoms); @@ -462,7 +463,8 @@ inline void DP_DeepPotComputeNList_variant( (dp->aparam_nall ? natoms : (natoms - nghost)) * dp->daparam); } - // charge_spin is always double (it is a per-frame physical input) + // charge_spin is converted to double for the api_cc::DeepPot interface + // (the compute backend stores/uses charge_spin as float64). std::vector charge_spin_; if (charge_spin) { charge_spin_.assign(charge_spin, @@ -532,7 +534,7 @@ template void DP_DeepPotComputeNList_variant(DP_DeepPot* dp, float* virial, float* atomic_energy, float* atomic_virial, - const double* charge_spin); + const float* charge_spin); // support spin template @@ -741,7 +743,7 @@ void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, VALUETYPE* virial, VALUETYPE* atomic_energy, VALUETYPE* atomic_virial, - const double* charge_spin = nullptr) { + const VALUETYPE* charge_spin = nullptr) { if (nframes > 1) { throw std::runtime_error("nframes > 1 not supported yet"); } @@ -761,7 +763,8 @@ void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, if (aparam) { aparam_.assign(aparam, aparam + nframes * natoms * dp->daparam); } - // charge_spin is always double (it is a per-frame physical input) + // charge_spin is converted to double for the api_cc::DeepPot interface + // (the compute backend stores/uses charge_spin as float64). std::vector charge_spin_; if (charge_spin) { charge_spin_.assign(charge_spin, @@ -834,7 +837,7 @@ template void DP_DeepPotModelDeviCompute_variant( float* virial, float* atomic_energy, float* atomic_virial, - const double* charge_spin); + const float* charge_spin); template void DP_DeepSpinModelDeviCompute_variant(DP_DeepSpinModelDevi* dp, @@ -966,7 +969,7 @@ void DP_DeepPotModelDeviComputeNList_variant( VALUETYPE* virial, VALUETYPE* atomic_energy, VALUETYPE* atomic_virial, - const double* charge_spin = nullptr) { + const VALUETYPE* charge_spin = nullptr) { if (nframes > 1) { throw std::runtime_error("nframes > 1 not supported yet"); } @@ -988,7 +991,8 @@ void DP_DeepPotModelDeviComputeNList_variant( aparam, aparam + (dp->aparam_nall ? natoms : (natoms - nghost)) * dp->daparam); } - // charge_spin is always double (it is a per-frame physical input) + // charge_spin is converted to double for the api_cc::DeepPot interface + // (the compute backend stores/uses charge_spin as float64). std::vector charge_spin_; if (charge_spin) { charge_spin_.assign(charge_spin, @@ -1069,7 +1073,7 @@ template void DP_DeepPotModelDeviComputeNList_variant( float* virial, float* atomic_energy, float* atomic_virial, - const double* charge_spin); + const float* charge_spin); // support spin multi model. template @@ -1772,7 +1776,7 @@ void DP_DeepPotComputef3(DP_DeepPot* dp, const float* cell, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, @@ -1816,7 +1820,7 @@ void DP_DeepPotComputeNListf3(DP_DeepPot* dp, const int ago, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, @@ -2123,7 +2127,7 @@ void DP_DeepPotModelDeviComputef3(DP_DeepPotModelDevi* dp, const float* cell, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, @@ -2167,7 +2171,7 @@ void DP_DeepPotModelDeviComputeNListf3(DP_DeepPotModelDevi* dp, const int ago, const float* fparam, const float* aparam, - const double* charge_spin, + const float* charge_spin, double* energy, float* force, float* virial, From 1e07772b638fcd50d5d10fa1774ecb76371c4898 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 11 Jun 2026 09:54:05 +0000 Subject: [PATCH 23/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/api_c/include/deepmd.hpp | 247 ++++++++++++++++---------------- source/api_c/src/c_api.cc | 29 ++-- 2 files changed, 142 insertions(+), 134 deletions(-) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index e12ef4c65c..e67becee72 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -511,23 +511,22 @@ inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, FPTYPE* atomic_virial); template <> -inline void _DP_DeepPotModelDeviComputeNList( - DP_DeepPotModelDevi* dp, - const int natom, - const double* coord, - const int* atype, - const double* cell, - const int nghost, - const DP_Nlist* nlist, - const int ago, - const double* fparam, - const double* aparam, - const double* charge_spin, - double* energy, - double* force, - double* virial, - double* atomic_energy, - double* atomic_virial) { +inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, + const int natom, + const double* coord, + const int* atype, + const double* cell, + const int nghost, + const DP_Nlist* nlist, + const int ago, + const double* fparam, + const double* aparam, + const double* charge_spin, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { if (charge_spin) { DP_DeepPotModelDeviComputeNList3( dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, @@ -872,9 +871,10 @@ namespace hpp { * @return charge_spin.data() if non-empty and valid, otherwise nullptr. */ template -inline const FPTYPE* validate_charge_spin(const std::vector& charge_spin, - const int dchgspin, - const unsigned int nframes) { +inline const FPTYPE* validate_charge_spin( + const std::vector& charge_spin, + const int dchgspin, + const unsigned int nframes) { if (charge_spin.empty()) { return nullptr; } @@ -1196,15 +1196,16 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute(ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1259,17 +1260,18 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute(ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - std::vector& atom_energy, - std::vector& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1298,10 +1300,9 @@ class DeepPot : public DeepBaseModel { const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); - _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, - fparam__, aparam__, charge_spin__, ener_, - force_, virial_, atomic_ener_, - atomic_virial_); + _DP_DeepPotCompute( + dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, + charge_spin__, ener_, force_, virial_, atomic_ener_, atomic_virial_); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; @@ -1330,18 +1331,19 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute(ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1401,20 +1403,21 @@ class DeepPot : public DeepBaseModel { * @warning Natoms should not be zero when computing multiple frames. **/ template - void compute(ENERGYVTYPE& ener, - std::vector& force, - std::vector& virial, - std::vector& atom_energy, - std::vector& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); @@ -1446,10 +1449,10 @@ class DeepPot : public DeepBaseModel { const VALUETYPE* charge_spin__ = validate_charge_spin(charge_spin, dchgspin, nframes); - _DP_DeepPotComputeNList( - dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, - fparam__, aparam__, charge_spin__, ener_, force_, virial_, atomic_ener_, - atomic_virial_); + _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, + box_, nghost, lmp_list.nl, ago, fparam__, + aparam__, charge_spin__, ener_, force_, + virial_, atomic_ener_, atomic_virial_); DP_CHECK_OK(DP_DeepPotCheckOK, dp); }; /** @@ -2260,15 +2263,16 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute(std::vector& ener, - std::vector>& force, - std::vector>& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + std::vector& ener, + std::vector>& force, + std::vector>& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. unsigned int natoms = atype.size(); unsigned int nframes = 1; @@ -2342,17 +2346,18 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute(std::vector& ener, - std::vector>& force, - std::vector>& virial, - std::vector>& atom_energy, - std::vector>& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + std::vector& ener, + std::vector>& force, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { // charge_spin is not supported via the C-API model-deviation path. unsigned int natoms = atype.size(); unsigned int nframes = 1; @@ -2442,18 +2447,19 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute(std::vector& ener, - std::vector>& force, - std::vector>& virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + std::vector& ener, + std::vector>& force, + std::vector>& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); @@ -2532,20 +2538,21 @@ class DeepPotModelDevi : public DeepBaseModelDevi { *same aparam. **/ template - void compute(std::vector& ener, - std::vector>& force, - std::vector>& virial, - std::vector>& atom_energy, - std::vector>& atom_virial, - const std::vector& coord, - const std::vector& atype, - const std::vector& box, - const int nghost, - const InputNlist& lmp_list, - const int& ago, - const std::vector& fparam = std::vector(), - const std::vector& aparam = std::vector(), - const std::vector& charge_spin = std::vector()) { + void compute( + std::vector& ener, + std::vector>& force, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector(), + const std::vector& charge_spin = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); diff --git a/source/api_c/src/c_api.cc b/source/api_c/src/c_api.cc index e648e7e72e..1346bfbf20 100644 --- a/source/api_c/src/c_api.cc +++ b/source/api_c/src/c_api.cc @@ -730,20 +730,21 @@ inline void flatten_vector(std::vector& onedv, } template -void DP_DeepPotModelDeviCompute_variant(DP_DeepPotModelDevi* dp, - const int nframes, - const int natoms, - const VALUETYPE* coord, - const int* atype, - const VALUETYPE* cell, - const VALUETYPE* fparam, - const VALUETYPE* aparam, - double* energy, - VALUETYPE* force, - VALUETYPE* virial, - VALUETYPE* atomic_energy, - VALUETYPE* atomic_virial, - const VALUETYPE* charge_spin = nullptr) { +void DP_DeepPotModelDeviCompute_variant( + DP_DeepPotModelDevi* dp, + const int nframes, + const int natoms, + const VALUETYPE* coord, + const int* atype, + const VALUETYPE* cell, + const VALUETYPE* fparam, + const VALUETYPE* aparam, + double* energy, + VALUETYPE* force, + VALUETYPE* virial, + VALUETYPE* atomic_energy, + VALUETYPE* atomic_virial, + const VALUETYPE* charge_spin = nullptr) { if (nframes > 1) { throw std::runtime_error("nframes > 1 not supported yet"); } From f3dd7600c64119b328980cd1a07ffa231f58100c Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Mon, 15 Jun 2026 10:38:25 +0800 Subject: [PATCH 24/24] docs: document charge_spin argument (njzjz review) - 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 --- doc/third-party/lammps-command.md | 7 +- source/api_c/include/c_api.h | 179 ++++++++++++++++++++++++++++++ source/api_c/include/deepmd.hpp | 32 ++++++ source/api_cc/include/DeepPot.h | 40 +++++++ 4 files changed, 257 insertions(+), 1 deletion(-) diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md index 70df75d22d..3cc5d955cc 100644 --- a/doc/third-party/lammps-command.md +++ b/doc/third-party/lammps-command.md @@ -49,7 +49,7 @@ pair_style deepmd models ... keyword value ... - models = frozen model(s) to compute the interaction. If multiple models are provided, then only the first model serves to provide energy and force prediction for each timestep of molecular dynamics, and the model deviation will be computed among all models every `out_freq` timesteps. -- keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _aparam_from_compute_ or _atomic_ or _relative_ or _relative_v_ or _aparam_ or _ttm_ +- keyword = _out_file_ or _out_freq_ or _fparam_ or _fparam_from_compute_ or _aparam_from_compute_ or _charge_spin_ or _atomic_ or _relative_ or _relative_v_ or _aparam_ or _ttm_ 
     out_file value = filename
@@ -62,6 +62,8 @@ pair_style deepmd models ... keyword value ...
         id = compute id used to update the frame parameter.
     aparam_from_compute value = id
         id = compute id used to update the atom parameter.
+    charge_spin value = parameters
+        parameters = the per-frame charge/spin values (dim_chg_spin numbers) required by models trained with a charge/spin embedding (e.g. DPA-3 with add_chg_spin_ebd). If omitted, the model's stored default_chg_spin is used.
     atomic = no value is required.
         If this keyword is set, the force model deviation of each atom will be output.
     relative value = level
@@ -88,6 +90,8 @@ compute    TEMP all temp
 
 pair_style deepmd ener.pb aparam_from_compute 1
 compute    1 all ke/atom
+
+pair_style deepmd dpa3.pth charge_spin 1.0 2.0
 ```
 
 ### Description
@@ -112,6 +116,7 @@ If the keyword `fparam` is set, the given frame parameter(s) will be fed to the
 If the keyword `fparam_from_compute` is set, the global parameter(s) from compute command (e.g., temperature from [compute temp command](https://docs.lammps.org/compute_temp.html)) will be fed to the model as the frame parameter(s).
 If the keyword `aparam_from_compute` is set, the atomic parameter(s) from compute command (e.g., per-atom translational kinetic energy from [compute ke/atom command](https://docs.lammps.org/compute_ke_atom.html)) will be fed to the model as the atom parameter(s).
 If the keyword `aparam` is set, the given atomic parameter(s) will be fed to the model, where each atom is assumed to have the same atomic parameter(s).
+If the keyword `charge_spin` is set, the given per-frame charge/spin value(s) will be fed to models that were trained with a charge/spin embedding (e.g. DPA-3 with `add_chg_spin_ebd`). If the keyword is not set, the model's stored `default_chg_spin` (if any) is used.
 If the keyword `ttm` is set, electronic temperatures from [fix ttm command](https://docs.lammps.org/fix_ttm.html) will be fed to the model as the atomic parameters.
 
 Only a single `pair_coeff` command is used with the deepmd style which specifies atom names. These are mapped to LAMMPS atom types (integers from 1 to Ntypes) by specifying Ntypes additional arguments after `* *` in the `pair_coeff` command.
diff --git a/source/api_c/include/c_api.h b/source/api_c/include/c_api.h
index c0c48a55c6..31b9e0626c 100644
--- a/source/api_c/include/c_api.h
+++ b/source/api_c/include/c_api.h
@@ -724,6 +724,30 @@ extern void DP_DeepPotCompute3(DP_DeepPot* dp,
  * @brief Evaluate the energy, force and virial by using a DP. (float version,
  *with charge_spin)
  * @version 3
+ * @param[in] dp The DP to use.
+ * @param[in] nframes The number of frames.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
+ * @param[in] charge_spin The per-frame charge/spin input. The array can be of
+ *size nframes x dim_chg_spin. Pass NULL to use the model's stored
+ *default_chg_spin.
+ * @param[out] energy Output energy.
+ * @param[out] force Output force. The array should be of size natoms x 3.
+ * @param[out] virial Output virial. The array should be of size 9.
+ * @param[out] atomic_energy Output atomic energy. The array should be of size
+ *natoms.
+ * @param[out] atomic_virial Output atomic virial. The array should be of size
+ *natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 extern void DP_DeepPotComputef3(DP_DeepPot* dp,
@@ -745,9 +769,33 @@ extern void DP_DeepPotComputef3(DP_DeepPot* dp,
  * @brief Evaluate the energy, force and virial by using a DP with the neighbor
  *list. (double version, with charge_spin)
  * @version 3
+ * @param[in] dp The DP to use.
+ * @param[in] nframes The number of frames.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] nghost The number of ghost atoms.
+ * @param[in] nlist The neighbor list.
+ * @param[in] ago Update the internal neighbour list if ago is 0.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
  * @param[in] charge_spin The per-frame charge/spin input. The array can be of
  *size nframes x dim_chg_spin. Pass NULL to use the model's stored
  *default_chg_spin.
+ * @param[out] energy Output energy.
+ * @param[out] force Output force. The array should be of size natoms x 3.
+ * @param[out] virial Output virial. The array should be of size 9.
+ * @param[out] atomic_energy Output atomic energy. The array should be of size
+ *natoms.
+ * @param[out] atomic_virial Output atomic virial. The array should be of size
+ *natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 extern void DP_DeepPotComputeNList3(DP_DeepPot* dp,
@@ -772,6 +820,33 @@ extern void DP_DeepPotComputeNList3(DP_DeepPot* dp,
  * @brief Evaluate the energy, force and virial by using a DP with the neighbor
  *list. (float version, with charge_spin)
  * @version 3
+ * @param[in] dp The DP to use.
+ * @param[in] nframes The number of frames.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] nghost The number of ghost atoms.
+ * @param[in] nlist The neighbor list.
+ * @param[in] ago Update the internal neighbour list if ago is 0.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
+ * @param[in] charge_spin The per-frame charge/spin input. The array can be of
+ *size nframes x dim_chg_spin. Pass NULL to use the model's stored
+ *default_chg_spin.
+ * @param[out] energy Output energy.
+ * @param[out] force Output force. The array should be of size natoms x 3.
+ * @param[out] virial Output virial. The array should be of size 9.
+ * @param[out] atomic_energy Output atomic energy. The array should be of size
+ *natoms.
+ * @param[out] atomic_virial Output atomic virial. The array should be of size
+ *natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 extern void DP_DeepPotComputeNListf3(DP_DeepPot* dp,
@@ -1467,9 +1542,32 @@ void DP_DeepPotModelDeviComputeNListf2(DP_DeepPotModelDevi* dp,
  * @brief Evaluate energy, force and virial with a DP model deviation.
  *(double version, with charge_spin)
  * @version 3
+ * @param[in] dp The DP model deviation to use.
+ * @param[in] nframes The number of frames. Only 1 is supported.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
  * @param[in] charge_spin The per-frame charge/spin input. The array should be
  *of size nframes x dim_chg_spin. Pass NULL to use the model's stored
  *default_chg_spin.
+ * @param[out] energy Output energy of all models.
+ * @param[out] force Output force of all models. The array should be of size
+ *nmodels x natoms x 3.
+ * @param[out] virial Output virial of all models. The array should be of size
+ *nmodels x 9.
+ * @param[out] atomic_energy Output atomic energy of all models. The array
+ *should be of size nmodels x natoms.
+ * @param[out] atomic_virial Output atomic virial of all models. The array
+ *should be of size nmodels x natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 void DP_DeepPotModelDeviCompute3(DP_DeepPotModelDevi* dp,
@@ -1491,6 +1589,32 @@ void DP_DeepPotModelDeviCompute3(DP_DeepPotModelDevi* dp,
  * @brief Evaluate energy, force and virial with a DP model deviation.
  *(float version, with charge_spin)
  * @version 3
+ * @param[in] dp The DP model deviation to use.
+ * @param[in] nframes The number of frames. Only 1 is supported.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
+ * @param[in] charge_spin The per-frame charge/spin input. The array should be
+ *of size nframes x dim_chg_spin. Pass NULL to use the model's stored
+ *default_chg_spin.
+ * @param[out] energy Output energy of all models.
+ * @param[out] force Output force of all models. The array should be of size
+ *nmodels x natoms x 3.
+ * @param[out] virial Output virial of all models. The array should be of size
+ *nmodels x 9.
+ * @param[out] atomic_energy Output atomic energy of all models. The array
+ *should be of size nmodels x natoms.
+ * @param[out] atomic_virial Output atomic virial of all models. The array
+ *should be of size nmodels x natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 void DP_DeepPotModelDeviComputef3(DP_DeepPotModelDevi* dp,
@@ -1512,9 +1636,35 @@ void DP_DeepPotModelDeviComputef3(DP_DeepPotModelDevi* dp,
  * @brief Evaluate energy, force and virial with a DP model deviation and a
  *neighbor list. (double version, with charge_spin)
  * @version 3
+ * @param[in] dp The DP model deviation to use.
+ * @param[in] nframes The number of frames. Only 1 is supported.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] nghost The number of ghost atoms.
+ * @param[in] nlist The neighbor list.
+ * @param[in] ago Update the internal neighbour list if ago is 0.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
  * @param[in] charge_spin The per-frame charge/spin input. The array should be
  *of size nframes x dim_chg_spin. Pass NULL to use the model's stored
  *default_chg_spin.
+ * @param[out] energy Output energy of all models.
+ * @param[out] force Output force of all models. The array should be of size
+ *nmodels x natoms x 3.
+ * @param[out] virial Output virial of all models. The array should be of size
+ *nmodels x 9.
+ * @param[out] atomic_energy Output atomic energy of all models. The array
+ *should be of size nmodels x natoms.
+ * @param[out] atomic_virial Output atomic virial of all models. The array
+ *should be of size nmodels x natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 void DP_DeepPotModelDeviComputeNList3(DP_DeepPotModelDevi* dp,
@@ -1539,6 +1689,35 @@ void DP_DeepPotModelDeviComputeNList3(DP_DeepPotModelDevi* dp,
  * @brief Evaluate energy, force and virial with a DP model deviation and a
  *neighbor list. (float version, with charge_spin)
  * @version 3
+ * @param[in] dp The DP model deviation to use.
+ * @param[in] nframes The number of frames. Only 1 is supported.
+ * @param[in] natoms The number of atoms.
+ * @param[in] coord The coordinates of atoms. The array should be of size natoms
+ *x 3.
+ * @param[in] atype The atom types. The array should contain natoms ints.
+ * @param[in] cell The cell of the region. The array should be of size 9. Pass
+ *NULL if pbc is not used.
+ * @param[in] nghost The number of ghost atoms.
+ * @param[in] nlist The neighbor list.
+ * @param[in] ago Update the internal neighbour list if ago is 0.
+ * @param[in] fparam The frame parameters. The array can be of size nframes x
+ *dim_fparam.
+ * @param[in] aparam The atom parameters. The array can be of size nframes x
+ *natoms x dim_aparam.
+ * @param[in] charge_spin The per-frame charge/spin input. The array should be
+ *of size nframes x dim_chg_spin. Pass NULL to use the model's stored
+ *default_chg_spin.
+ * @param[out] energy Output energy of all models.
+ * @param[out] force Output force of all models. The array should be of size
+ *nmodels x natoms x 3.
+ * @param[out] virial Output virial of all models. The array should be of size
+ *nmodels x 9.
+ * @param[out] atomic_energy Output atomic energy of all models. The array
+ *should be of size nmodels x natoms.
+ * @param[out] atomic_virial Output atomic virial of all models. The array
+ *should be of size nmodels x natoms x 9.
+ * @warning The output arrays should be allocated before calling this function.
+ *Pass NULL if not required.
  * @since API version 27
  **/
 void DP_DeepPotModelDeviComputeNListf3(DP_DeepPotModelDevi* dp,
diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp
index e67becee72..1a68d0e227 100644
--- a/source/api_c/include/deepmd.hpp
+++ b/source/api_c/include/deepmd.hpp
@@ -1193,6 +1193,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    * @warning Natoms should not be zero when computing multiple frames.
    **/
   template 
@@ -1257,6 +1261,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    * @warning Natoms should not be zero when computing multiple frames.
    **/
   template 
@@ -1328,6 +1336,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    * @warning Natoms should not be zero when computing multiple frames.
    **/
   template 
@@ -1400,6 +1412,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    * @warning Natoms should not be zero when computing multiple frames.
    **/
   template 
@@ -2261,6 +2277,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    **/
   template 
   void compute(
@@ -2344,6 +2364,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    **/
   template 
   void compute(
@@ -2445,6 +2469,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    **/
   template 
   void compute(
@@ -2536,6 +2564,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The per-frame charge/spin input. The array can be
+   *of size nframes x dim_chg_spin. Then all frames are assumed to be provided
+   *with the same charge_spin. Leave it empty to use the model's stored
+   *default_chg_spin.
    **/
   template 
   void compute(
diff --git a/source/api_cc/include/DeepPot.h b/source/api_cc/include/DeepPot.h
index 6172f08c1c..177513ebb0 100644
--- a/source/api_cc/include/DeepPot.h
+++ b/source/api_cc/include/DeepPot.h
@@ -358,6 +358,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    * @{
    **/
   template 
@@ -402,6 +406,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    * @{
    **/
   template 
@@ -452,6 +460,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    * @{
    **/
   template 
@@ -504,6 +516,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    * @{
    **/
   template 
@@ -558,6 +574,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    * @{
    **/
   template 
@@ -608,6 +628,10 @@ class DeepPot : public DeepBaseModel {
    * nframes x natoms x dim_aparam.
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    * @{
    **/
   template 
@@ -705,6 +729,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam. dim_aparam. Then all frames and atoms are provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    **/
   template 
   void compute(std::vector& all_ener,
@@ -739,6 +767,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam. dim_aparam. Then all frames and atoms are provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    **/
   template 
   void compute(std::vector& all_ener,
@@ -775,6 +807,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam. dim_aparam. Then all frames and atoms are provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    **/
   template 
   void compute(std::vector& all_ener,
@@ -814,6 +850,10 @@ class DeepPotModelDevi : public DeepBaseModelDevi {
    * natoms x dim_aparam. Then all frames are assumed to be provided with the
    *same aparam. dim_aparam. Then all frames and atoms are provided with the
    *same aparam.
+   * @param[in] charge_spin The charge/spin parameter. The array can be of size
+   *nframes x dim_chg_spin.
+   * dim_chg_spin. Then all frames are assumed to be provided with the same
+   *charge_spin. Leave it empty to use the model's stored default_chg_spin.
    **/
   template 
   void compute(std::vector& all_ener,