Skip to content

feat: add charge spin to cpp runtime #5509

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anyangml wants to merge 24 commits into
base: master
Choose a base branch
from
Open

feat: add charge spin to cpp runtime #5509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
7f5540f
feat: add charge spin to cpp runtime
anyangml Jun 9, 2026
57df385
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 9, 2026
46a0164
fix: UT
anyangml Jun 9, 2026
982c711
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 9, 2026
7ae6be5
fix: param passing
anyangml Jun 9, 2026
d469444
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 9, 2026
dac5b4b
fix: UT
anyangml Jun 9, 2026
1b8e8b5
fix: export charge spin dim
anyangml Jun 9, 2026
f24390b
fix: ghost mapping
anyangml Jun 9, 2026
738f2e4
fix: address comments
anyangml Jun 10, 2026
66c471e
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 10, 2026
5c02d99
feat: add charge_spin to the .pth (TorchScript) cpp backend
anyangml Jun 10, 2026
b68dc9a
fix: address charge_spin review comments
anyangml Jun 10, 2026
f743f00
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 10, 2026
dce3c43
feat: wire charge_spin through the C API; drop charge_spin_from_compute
anyangml Jun 10, 2026
1c7252d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 10, 2026
f178f2b
feat: wire charge_spin through the C API model-deviation path
anyangml Jun 11, 2026
2ff72cb
fix: address CodeRabbit review on charge_spin (.pth / C API / tests)
anyangml Jun 11, 2026
ad66793
test: add LAMMPS integration test for the charge_spin pair_style keyword
anyangml Jun 11, 2026
55955c3
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 11, 2026
b667537
fix: drop get_dim_chg_spin() call on the .pth DeepEval in gen_chg_spin
anyangml Jun 11, 2026
34b441c
refactor: type charge_spin as FPTYPE in the C API (iProzd review)
anyangml Jun 11, 2026
1e07772
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jun 11, 2026
f3dd760
docs: document charge_spin argument (njzjz review)
anyangml Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions deepmd/pt_expt/utils/serialization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
7 changes: 6 additions & 1 deletion doc/third-party/lammps-command.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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_

<pre>
<i>out_file</i> value = filename
Expand All @@ -62,6 +62,8 @@ pair_style deepmd models ... keyword value ...
id = compute id used to update the frame parameter.
<i>aparam_from_compute</i> value = id
id = compute id used to update the atom parameter.
<i>charge_spin</i> 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.
<i>atomic</i> = no value is required.
If this keyword is set, the force model deviation of each atom will be output.
<i>relative</i> value = level
Expand All @@ -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
Expand All @@ -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.
Expand Down
411 changes: 410 additions & 1 deletion source/api_c/include/c_api.h
View file
Open in desktop

Large diffs are not rendered by default.

284 changes: 233 additions & 51 deletions source/api_c/include/deepmd.hpp
View file
Open in desktop

Large diffs are not rendered by default.

375 changes: 297 additions & 78 deletions source/api_c/src/c_api.cc
View file
Open in desktop

Large diffs are not rendered by default.

210 changes: 194 additions & 16 deletions source/api_cc/include/DeepPot.h
View file
Open in desktop

Large diffs are not rendered by default.

74 changes: 74 additions & 0 deletions source/api_cc/include/DeepPotPT.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ class DeepPotPT : public DeepPotBackend {
const std::vector<VALUETYPE>& box,
const std::vector<VALUETYPE>& fparam,
const std::vector<VALUETYPE>& aparam,
const std::vector<double>& charge_spin,
const bool atomic);
/**
* @brief Evaluate the energy, force, virial, atomic energy, and atomic virial
Expand Down Expand Up @@ -114,6 +115,7 @@ class DeepPotPT : public DeepPotBackend {
const int& ago,
const std::vector<VALUETYPE>& fparam,
const std::vector<VALUETYPE>& aparam,
const std::vector<double>& charge_spin,
const bool atomic);
/**
* @brief Evaluate the energy, force, and virial with the mixed type
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -304,6 +315,11 @@ class DeepPotPT : public DeepPotBackend {
const std::vector<float>& fparam,
const std::vector<float>& 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<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
Expand All @@ -329,15 +345,73 @@ class DeepPotPT : public DeepPotBackend {
const std::vector<float>& aparam,
const bool atomic);

// charge_spin overloads — pass runtime charge/spin per call
void computew(std::vector<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_energy,
std::vector<double>& atom_virial,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const std::vector<double>& fparam,
const std::vector<double>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew(std::vector<double>& ener,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_energy,
std::vector<float>& atom_virial,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const std::vector<float>& fparam,
const std::vector<float>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew(std::vector<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_energy,
std::vector<double>& atom_virial,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const int nghost,
const InputNlist& inlist,
const int& ago,
const std::vector<double>& fparam,
const std::vector<double>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew(std::vector<double>& ener,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_energy,
std::vector<float>& atom_virial,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const int nghost,
const InputNlist& inlist,
const int& ago,
const std::vector<float>& fparam,
const std::vector<float>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
Comment thread
coderabbitai[bot] marked this conversation as resolved.

private:
int num_intra_nthreads, num_inter_nthreads;
bool inited;
int ntypes;
int ntypes_spin;
int dfparam;
int daparam;
int dchgspin;
bool aparam_nall;
bool has_default_fparam_;
std::vector<double> default_chg_spin_;
// copy neighbor list info from host
torch::jit::script::Module module;
double rcut;
Expand Down
101 changes: 98 additions & 3 deletions source/api_cc/include/DeepPotPTExpt.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ class DeepPotPTExpt : public DeepPotBackend {
const int& ago,
const std::vector<VALUETYPE>& fparam,
const std::vector<VALUETYPE>& aparam,
const std::vector<double>& charge_spin,
const bool atomic);
/**
* @brief Evaluate without nlist (standalone — builds nlist, folds back).
Expand All @@ -92,6 +93,7 @@ class DeepPotPTExpt : public DeepPotBackend {
const std::vector<VALUETYPE>& box,
const std::vector<VALUETYPE>& fparam,
const std::vector<VALUETYPE>& aparam,
const std::vector<double>& charge_spin,
const bool atomic);

public:
Expand All @@ -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);
Expand All @@ -125,7 +131,8 @@ 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<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
Expand Down Expand Up @@ -201,12 +208,94 @@ class DeepPotPTExpt : public DeepPotBackend {
const std::vector<float>& aparam,
const bool atomic);

// charge_spin overloads — pass runtime charge/spin per call
void computew(std::vector<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_energy,
std::vector<double>& atom_virial,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const std::vector<double>& fparam,
const std::vector<double>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew(std::vector<double>& ener,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_energy,
std::vector<float>& atom_virial,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const std::vector<float>& fparam,
const std::vector<float>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew(std::vector<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_energy,
std::vector<double>& atom_virial,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const int nghost,
const InputNlist& inlist,
const int& ago,
const std::vector<double>& fparam,
const std::vector<double>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew(std::vector<double>& ener,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_energy,
std::vector<float>& atom_virial,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const int nghost,
const InputNlist& inlist,
const int& ago,
const std::vector<float>& fparam,
const std::vector<float>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew_mixed_type(std::vector<double>& ener,
std::vector<double>& force,
std::vector<double>& virial,
std::vector<double>& atom_energy,
std::vector<double>& atom_virial,
const int& nframes,
const std::vector<double>& coord,
const std::vector<int>& atype,
const std::vector<double>& box,
const std::vector<double>& fparam,
const std::vector<double>& aparam,
const std::vector<double>& charge_spin,
const bool atomic) override;
void computew_mixed_type(std::vector<double>& ener,
std::vector<float>& force,
std::vector<float>& virial,
std::vector<float>& atom_energy,
std::vector<float>& atom_virial,
const int& nframes,
const std::vector<float>& coord,
const std::vector<int>& atype,
const std::vector<float>& box,
const std::vector<float>& fparam,
const std::vector<float>& aparam,
const std::vector<double>& 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<double> default_fparam_;
Expand Down Expand Up @@ -255,6 +344,7 @@ class DeepPotPTExpt : public DeepPotBackend {
const std::vector<VALUETYPE>& box,
const std::vector<VALUETYPE>& fparam,
const std::vector<VALUETYPE>& aparam,
const std::vector<double>& charge_spin,
const bool atomic);

/**
Expand All @@ -272,6 +362,7 @@ class DeepPotPTExpt : public DeepPotBackend {
const std::vector<VALUETYPE>& box,
const std::vector<VALUETYPE>& fparam,
const std::vector<VALUETYPE>& aparam,
const std::vector<double>& charge_spin,
const bool atomic);

/**
Expand All @@ -282,20 +373,23 @@ 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<torch::Tensor> run_model(const torch::Tensor& coord,
const torch::Tensor& atype,
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.
Expand All @@ -307,6 +401,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<at::Tensor>& comm_tensors);

/**
Expand Down
Loading
Loading