deepmodeling / deepmodeling/deepmd-kit

[Feature Request] Implement DPLR dipole_charge modifier in dpmodel

Open
#5,754 1 comment 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

### 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 (#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:

- #1254: older DPLR training failure, later closed as duplicate/upstream.
- #2590: Python-interface DPLR failure for compressed DeepDipole models due to missing gradients.
- #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`

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.