[Q&A] DPLR with DPA4 (SeZM) energy model

[yjw0510]

Question

Hi, we're using DPLR for systems where long-range electrostatics matter, and we'd love to use DPA4/SeZM as the energy backbone while keeping the DPLR treatment.

I dug around v3.2.0b0 a bit and it looks like the pieces aren't quite connected yet. There's a new BaseModifier scaffold under deepmd/pt/modifier/, but dipole_charge itself still only lives in the TF backend. Also get_sezm_model() is hardwired to SeZMEnergyFittingNet, so I can't pair a dipole fitting with the SeZM descriptor either. Meanwhile DeepTensorPT.cc already exists on the C++ side, so fix dplr loading a PyTorch DW model seems fine.

My understanding is that none of this is a real architectural constraint though. The modifier just subtracts the real-space part of the Wannier-centroid electrostatics from the labels, the k-space part is handled by pppm_dplr at run time, and the DW model is trained separately anyway. So an SeZM energy model trained on modifier-corrected labels should be perfectly valid, and what's missing is plumbing rather than physics. Am I getting this right?

Concretely I'm wondering:

• Is a PT port of dipole_charge on the roadmap? (the BaseModifier scaffold made me hopeful) Happy to test a branch if one exists.
• Would the recommended combo then be DW = dpa1/2/3 or se_atten (PT) + energy = SeZM, or is dipole fitting for SeZM itself planned too?
• In the meantime, is it sound to just preprocess the dataset ourselves, i.e. subtract the real-space Coulomb of the Wannier-centroid charges from E/F/virial, train SeZM on that through the normal PT path, and run pair deepmd + fix dplr + pppm_dplr with a PT DW model? Or are there gotchas with the modifier's real-space kernel conventions (cutoff/smoothing, virial bookkeeping, data-statistics init on modified labels)?

DeePMD-kit Version

No response

Backend and its version

No response

Python Version, CUDA Version, GCC Version, LAMMPS Version, etc

No response

Details

No response

Reproducible Example, Input Files, and Commands

No response

Further Information, Files, and Links

No response

[dosubot[bot]]

Hi @yjw0510! I'm Dosu and I'm helping the DeepModeling team.

Your analysis is essentially correct — the gap is plumbing, not physics. Here's the current picture:

PT dipole_charge modifier status

There are two relevant PRs plus a tracking issue:

1. PR #5138 (feat(pt): DPLR in PyTorch backend) — adds deepmd/pt/modifier/dipole_charge.py with a native PyTorch DipoleChargeModifier (using torch_admp for Ewald), serialization through modifier_data, and TF/PT consistency tests. This PR is still open.

2. PR #5775 (feat(modifier): add dpmodel dipole-charge adapters) — already merged. It adds a backend-neutral DipoleChargeModifierBase in deepmd/dpmodel/modifier/dipole_charge.py plus adapters for JAX, pt_expt, and TF2, with a shared Array-API Ewald core. It explicitly notes: "This PR intentionally does not add a native PyTorch implementation. Trainer label preprocessing and compiled artifact packaging remain follow-up work."

3. Issue #5754 tracks the full dpmodel design and open questions (DW model packaging, LAMMPS contract, mask semantics, etc.).

On the C++/LAMMPS side, DataModifier.cc still throws "PyTorch backend is not supported yet" — so fix dplr + pppm_dplr with a PT DW model is blocked at the C++ layer even if training works.

Answers to your specific questions

• Is a PT port of dipole_charge on the roadmap? Yes — PR #5138 is the active branch. If you want to test it, that's the one to check with @ChiahsinChu about.

• Recommended DW + energy combo? The planned path supports any PT dipole model (DPA-1/2/3, se_atten, etc.) as the DW backbone, plus any PT energy model for the short-range part. SeZM/DPA4 as the energy backbone is architecturally fine — what's missing is only the plumbing you identified (modifier label preprocessing + LAMMPS integration). Dipole fitting for SeZM itself isn't blocked in principle either, but it hasn't been explicitly paired with the SeZM descriptor yet.

• Is manual preprocessing sound in the meantime? Your physics reasoning is right — the modifier only subtracts the real-space Ewald contribution from E/F/virial before training. The main gotchas:

  • Virial bookkeeping: the modifier back-propagates electric forces on WC sites through the DW model (the t_ef → o_dm_force/o_dm_virial path in TF) [1]. If you preprocess manually, you need to correctly attribute the Wannier-force virial contribution, not just the direct Ewald virial.
  • Data statistics: the modifier is applied before the energy model's mean/std statistics are computed [1]. If you preprocess externally and supply corrected data directly, that's fine — just make sure the statistics are computed on corrected labels.
  • Cutoff/smoothing: make sure your manual Ewald kernel uses the same rcut and smoothing as your DW model to avoid systematic force discontinuities.
  • LAMMPS runtime: even with a manually-corrected dataset and a trained PT energy model, running fix dplr + pppm_dplr still requires the DW model through DipoleChargeModifier, which currently rejects PT models at the C++ level. That wall is tracked but not yet removed.

The cleanest near-term path is probably to follow PR #5138's branch for training, and watch #5754 for LAMMPS integration progress.

To reply, just mention @dosu.

---

Docs are dead. Just use Dosu.

[njzjz-bot]

The high-level separation is right, but there are two important corrections to the Dosu answer.

1. There is not yet a supported end-to-end PyTorch DPLR workflow. PR #5138, which adds the native PyTorch dipole_charge modifier, is still open. Merged PR #5775 added the backend-neutral numerical core and eager adapters for pt_expt/JAX/TF2, but explicitly left native PT training and compiled-artifact packaging as follow-up work. The C++ DipoleChargeModifier still rejects both PyTorch and PyTorch-exportable models, so fix dplr cannot currently use a PT DW model.

2. Manual preprocessing must subtract the reciprocal-space/Fourier Ewald contribution, not a real-space Coulomb term. That is the E_ele used by DPLR and implemented as ewald_reciprocal_energy. The force and virial corrections must include the chain rule through the DW-predicted WC positions, and the charge maps, ewald_beta, reciprocal grid convention, and label statistics must match.

Using SeZM as the short-range energy backbone is physically reasonable once those corrected labels are available: the energy model is fitting the residual short-range contribution. However, current SeZM construction supports only energy/property fittings, not dipole fitting. A DW model therefore still needs a separately supported dipole/tensor model; SeZM itself cannot presently be used as the DW fitting model.

So the practical answer today is:

• TensorFlow remains the supported end-to-end DPLR training + LAMMPS route.
• Externally corrected labels can be used to train a SeZM energy model, provided the full reciprocal correction and derivatives are reproduced exactly.
• That preprocessing alone does not make pair deepmd + fix dplr + pppm/dplr work with a PT DW model; the C++ modifier backend is still the runtime blocker.

Coding agent: Codex
Codex version: codex-cli 0.149.0
Model: gpt-5.6-sol
Reasoning effort: xhigh