[Feature Request] Implement DPLR dipole_charge modifier in dpmodel

[njzjz]

Summary

DPLR (model.modifier.type = "dipole_charge") is currently implemented as a TensorFlow-only modifier plus a TensorFlow-only C++/LAMMPS runtime wrapper. This issue proposes adding a backend-agnostic dpmodel implementation so the same DPLR workflow can be supported by dpmodel-driven backends (especially PyTorch/JAX) while preserving the existing user input schema and the LAMMPS fix dplr contract.

Detailed Description

Current TensorFlow/Python behavior

The reference implementation is in deepmd/tf/modifier/dipole_charge.py.

The current workflow is:

• The energy model accepts a model.modifier block with type, model_name, model_charge_map, sys_charge_map, ewald_h, and ewald_beta.
• The modifier loads a frozen DeepDipole/DW model under the dipole_charge prefix.
• For each frame, it predicts WC displacement vectors on sel_type atoms, appends virtual WC sites with model_charge_map, assigns real atom charges from sys_charge_map, and evaluates the reciprocal Ewald contribution.
• During training/data loading, it subtracts this long-range energy/force/virial contribution from labels before the short-range DP model statistics/training.
• For forces and virials, it treats the Ewald force on WC sites as an external vector t_ef and back-propagates it through the DW model, then exports o_dm_force, o_dm_virial, and o_dm_av in the frozen graph.
• deepmd/tf/entrypoints/freeze.py keeps the modifier metadata, the imported dipole_charge/* DW graph nodes, and the o_dm_* nodes.

Relevant files:

• deepmd/tf/modifier/dipole_charge.py
• deepmd/tf/infer/ewald_recp.py
• source/op/tf/ewald_recp.cc
• source/lib/include/ewald.h
• source/lib/src/ewald.cc
• deepmd/tf/entrypoints/train.py
• deepmd/tf/entrypoints/freeze.py

Current C++/LAMMPS behavior

The LAMMPS path has a separate runtime contract:

• source/api_cc/include/DataModifier.h defines DipoleChargeModifierBase and DipoleChargeModifier.
• source/api_cc/src/DataModifier.cc currently dispatches only TensorFlow models; PyTorch, PyTorch exportable, Paddle, and JAX all throw not supported yet.
• source/api_cc/src/DataModifierTF.cc reads TF graph metadata such as cutoff, number of types, and selected types, builds the LAMMPS neighbor-list input tensors, feeds t_ef, and reads o_dm_force/o_dm_virial.
• source/lmp/fix_dplr.cpp uses two objects: a DeepTensor-style evaluator to predict WC positions and a DipoleChargeModifier to pull back electric forces from virtual WC sites to real atoms. pppm/dplr provides the reciprocal-space electric forces.

Relevant files:

• source/api_cc/include/DataModifier.h
• source/api_cc/include/DataModifierTF.h
• source/api_cc/src/DataModifier.cc
• source/api_cc/src/DataModifierTF.cc
• source/lmp/fix_dplr.cpp
• source/lmp/pppm_dplr.cpp
• source/lmp/tests/test_dplr.py

Proposed dpmodel design

1. Add a backend-agnostic modifier class in deepmd/dpmodel/modifier/dipole_charge.py registered as "dipole_charge".

   It should own the schema and pure data flow:

   • preserve the existing input keys: model_name, model_charge_map, sys_charge_map, ewald_h, ewald_beta;
   • serialize/deserialize with @class = "Modifier", type = "dipole_charge", and a compatible version;
   • resolve the DW/dipole model to a BaseModel object in backend-specific wrappers;
   • compute selected atom indices from the DW model sel_type;
   • build the extended charge system: real atoms plus WC coordinates R_wc = R_ref + d_wc;
   • return standard modifier outputs for energy-model workflows: energy, force, and virial corrections.

2. Split the implementation into a reusable numerical core and backend adapters.

   The core should define the DPLR data flow and shape conventions. Backend adapters should provide:

   • DW model execution;
   • reciprocal Ewald evaluation;
   • vector-Jacobian product/pullback from WC forces to atom forces and virials.

   Suggested backend order:

   • NumPy/dpmodel reference: first implementation for correctness tests, likely based on the existing C++ source/lib/src/ewald.cc formula.
   • PyTorch: implement deepmd/pt/modifier/dipole_charge.py using the current deepmd/pt/modifier/base_modifier.py hooks. It can use autograd/VJP for training and Python inference first. TorchScript/JIT support can be a follow-up if needed.
   • JAX: implement via jax.vjp once the model wrapper path can provide the DW model as a JAX/dpmodel object.

3. Keep the existing training semantics.

   The modifier should still be applied before data statistics and short-range training labels are consumed. The existing PyTorch data path already has modifier hooks in deepmd/pt/modifier/base_modifier.py, deepmd/pt/entrypoints/main.py, and deepmd/pt/train/training.py, but no concrete DPLR modifier exists today.

4. Define how the DW model is packaged.

   TF freezing embeds/imports the DW graph under dipole_charge. For dpmodel backends, we need to decide whether frozen/exported artifacts should:

   • embed the serialized DW model inside the modifier serialization; or
   • keep model_name as an external sidecar; or
   • support both, using model_name during training and embedded serialization after freeze/export.

   I prefer embedding the serialized DW model at freeze/export time so inference and LAMMPS do not depend on an external dw.pb path.

5. Preserve or replace the C++/LAMMPS contract deliberately.

   Minimal compatibility path:

   • keep DipoleChargeModifier and fix dplr unchanged;
   • add backend implementations for PyTorch/JAX/exported models that expose the same behavior as the TF o_dm_force/o_dm_virial path.

   Cleaner longer-term path:

   • define a backend-neutral DipoleChargeModifierBase::compute contract in terms of model-level operations instead of TF node names;
   • implement compute_wfcc and compute_force_pullback per backend;
   • keep the current TF implementation as one adapter.

Open questions

• Should dpmodel serialization embed the DW model by default, or should model_name remain a sidecar file for all backends?
• Which backend should be the first supported target: PyTorch native training/inference, JAX, or exported C++/LAMMPS?
• Should the first implementation support only reciprocal Ewald, matching the current DPLR training modifier, or should it also expose a reusable Ewald operator/API?
• How should we handle compressed DW/dipole models? The historical TF implementation had Python-interface failures when gradients through compressed ops were unavailable ([BUG] Inference with Python interface DPLR model raises error. #2590).
• What output-key convention should the backend modifiers use so they compose cleanly with existing energy model outputs and ModelWrapper?
• Should fix dplr continue to reject atomic virial/atomic energy for this path, matching current behavior?

Suggested acceptance criteria

• The current DPLR input schema remains valid under dpmodel-backed training.
• A small water DPLR case matches the TensorFlow modifier for reciprocal energy, force, and virial within numerical tolerance.
• Training data are modified before statistics for the PyTorch/JAX path, matching the TensorFlow training semantics.
• Frozen/exported models contain enough modifier metadata and DW model data to run without the original model_name file, or the sidecar requirement is explicitly documented and tested.
• C++/LAMMPS either supports the new backend through DipoleChargeModifier, or reports a clear tracked limitation with a follow-up plan.
• Tests cover the Ewald reference, the modifier data flow, Python inference, and at least one training smoke test.

Further Information, Files, and Links

Related historical issues:

• [BUG] Training long range model #1254: older DPLR training failure, later closed as duplicate/upstream.
• [BUG] Inference with Python interface DPLR model raises error. #2590: Python-interface DPLR failure for compressed DeepDipole models due to missing gradients.
• Undefined key "modifier" is not allowed in strict model for running DPLR model #4128: user confusion around the model.modifier schema.

Existing DPLR docs and example:

• doc/model/dplr.md
• examples/water/dplr/train/ener.json
• examples/water/dplr/train/dw.json
• examples/water/dplr/lmp/in.lammps

[njzjz]

One implementation risk from offline discussion: the numerical algorithm can mostly follow the TensorFlow DPLR path, but the mask semantics are quite different between the existing paths and should be designed explicitly before porting code.

Current TF/C++ DPLR path:

• fix dplr keeps real atoms and explicit virtual WC atoms in the LAMMPS system.
• DipoleChargeModifierTF::compute calls select_real_atoms, which filters atom types to [0, ntypes) and maps virtual/NULL atoms to -1 in the forward map.
• The DW model is evaluated only on the filtered real-atom system. WC forces are passed in through t_ef, and the returned correction is mapped back to the original real+virtual LAMMPS layout.
• Selected atoms are handled by sel_type/select_by_type, not by a generic model-output mask.

Current dpmodel path:

• Dense nlist uses atype < 0 as a virtual/padding atom mask and nlist == -1 as the neighbor-slot mask.
• BaseAtomicModel emits a mask output where valid atoms are 1 and virtual/excluded atoms are zeroed.
• Graph lower uses flat nodes plus edge_mask; masked dummy edges are part of the exported graph contract.
• Pair/atom exclusions are merged into different masks depending on dense vs graph lower.

So the dpmodel DPLR implementation should not just copy the formula and rely on existing masks implicitly. It probably needs an explicit mask contract, e.g.:

1. real_atom_mask: atoms participating in the DW model and receiving physical force correction; excludes LAMMPS virtual WC/NULL atoms and dpmodel padding atoms.
2. selected_wc_mask or selected-index map: atoms whose DW output creates a WC site; derived from DW sel_type.
3. neighbor_mask: dense nlist != -1, only for descriptor/DW evaluation.
4. edge_mask: graph lower equivalent when we support graph-exported backends.
5. output_mask: the existing dpmodel atomic output mask; should not be used as a substitute for the selected-WC mapping.

Suggested acceptance test addition: construct a small system with explicit virtual/padding atoms and verify that the dpmodel implementation matches the TF/C++ path after all map-back operations, including cases where selected WC atoms are a strict subset of real atoms.