feat(modifier): add dpmodel dipole-charge adapters#5775
Open
njzjz-bot wants to merge 1 commit into
Open
Conversation
Contributor
📝 WalkthroughWalkthroughAdds a backend-neutral dipole-charge modifier with Ewald/DPLR utilities, JAX, exportable PyTorch, and TensorFlow implementations, public exports, serialization, and cross-backend consistency tests. ChangesDipole-charge modifier
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant BackendModifier
participant DipoleModel
participant DPLRUtilities
Caller->>BackendModifier: provide coordinates, atom types, and periodic box
BackendModifier->>DipoleModel: compute dipoles
BackendModifier->>DPLRUtilities: extend system and evaluate Ewald energy
DPLRUtilities-->>BackendModifier: return reciprocal energy
BackendModifier-->>Caller: return energy, force, and virial
Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
deepmd/jax/modifier/dipole_charge.py (1)
30-66: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the Coulomb constant to a shared module.
14.39964535475696995031is hard-coded here. Since PyTorch, exportable PyTorch, and TF2 implementations must reproduce the exact same value for cross-backend parity (validated by the consistency test), defining this once in a shared constants module and importing it in each backend would reduce the risk of a transcription mismatch between backends.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/jax/modifier/dipole_charge.py` around lines 30 - 66, Extract the Coulomb constant currently embedded in ewald_reciprocal_energy into the project’s shared constants module, then import and use that single symbol here. Update the PyTorch, exportable PyTorch, and TF2 implementations to reference the same shared constant, preserving the exact value required for cross-backend parity.deepmd/dpmodel/modifier/dipole_charge.py (1)
91-102: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
model_charge_mapis mapped tosel_typepurely by position, with no identity check.Only the length of
sel_typevsmodel_charge_mapis validated (Line 99). The actual charge assignment (Line 108-111, and mirrored in the JAX modifier at deepmd/jax/modifier/dipole_charge.py Lines 89-93/128-131) trusts thatmodel_charge_map[i]corresponds tosel_type[i]in the exact order returned by the embedded model'sget_sel_type(). If that ordering ever diverges from what the caller assumed (e.g. multiple selected types, or model retraining changes internal type ordering), charges would be silently swapped between atom types with no error — producing wrong energies/forces without any diagnostic.Consider validating identity (e.g. requiring
model_charge_mapas a{atom_type: charge}mapping, or asserting the caller-supplied ordering againstsel_typeexplicitly) rather than relying on a purely positional, undocumented contract.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/modifier/dipole_charge.py` around lines 91 - 102, Update make_charge_maps and its mirrored JAX implementation to validate selected atom-type identity, not only sel_type/model_charge_map length, before assigning charges. Replace or supplement the positional model_charge_map contract with an explicit atom-type-to-charge mapping or an equivalent ordering assertion against the embedded model’s get_sel_type(), and raise a clear error when identities diverge.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/jax/modifier/dipole_charge.py`:
- Around line 162-171: Update the energy-and-gradient calculation in the shown
method to use jax.value_and_grad with has_aux=True, returning the summed scalar
as the differentiated value and energy_fn’s per-frame result as auxiliary data.
Reuse that auxiliary result for "energy_by_frame" and remove the separate
energy_fn(coord, strain) call, while preserving the existing force and virial
gradient transformations.
- Around line 109-115: Update the grid construction in the surrounding
dipole/charge computation to derive reciprocal grid sizes from each frame’s box
rather than only box[0]. Preserve per-frame grid handling through the downstream
calculation, or explicitly validate that all batch boxes are identical before
reusing one grid.
In `@deepmd/pt/modifier/dipole_charge.py`:
- Around line 193-225: Update the self.dipole_model invocation in _energy to
forward the received charge_spin argument as charge_spin=charge_spin. Preserve
the existing coordinate, type, box, parameter, and virial arguments.
- Around line 261-263: Update the torch.jit.annotate type in the grad_outputs
initialization to use TorchScript-compatible List[Optional[torch.Tensor]]
instead of list[torch.Tensor | None], while preserving the existing tensor list
and jitable execution path.
In `@deepmd/tf2/modifier/dipole_charge.py`:
- Around line 138-144: Update the selection handling around nselected in the TF2
dipole-charge modifier to validate that every frame’s selected-atom count
matches frame 0 before reshaping selected_coord, selected_atype, and
selected_dipole. Raise a clear ValueError on mismatch, while preserving the
existing reshape path for consistent selections.
---
Nitpick comments:
In `@deepmd/dpmodel/modifier/dipole_charge.py`:
- Around line 91-102: Update make_charge_maps and its mirrored JAX
implementation to validate selected atom-type identity, not only
sel_type/model_charge_map length, before assigning charges. Replace or
supplement the positional model_charge_map contract with an explicit
atom-type-to-charge mapping or an equivalent ordering assertion against the
embedded model’s get_sel_type(), and raise a clear error when identities
diverge.
In `@deepmd/jax/modifier/dipole_charge.py`:
- Around line 30-66: Extract the Coulomb constant currently embedded in
ewald_reciprocal_energy into the project’s shared constants module, then import
and use that single symbol here. Update the PyTorch, exportable PyTorch, and TF2
implementations to reference the same shared constant, preserving the exact
value required for cross-backend parity.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 3316a559-9c64-4653-a985-b69785af8347
📒 Files selected for processing (12)
deepmd/dpmodel/modifier/__init__.pydeepmd/dpmodel/modifier/dipole_charge.pydeepmd/jax/modifier/__init__.pydeepmd/jax/modifier/dipole_charge.pydeepmd/pt/modifier/__init__.pydeepmd/pt/modifier/dipole_charge.pydeepmd/pt_expt/modifier/__init__.pydeepmd/pt_expt/modifier/dipole_charge.pydeepmd/tf2/modifier/__init__.pydeepmd/tf2/modifier/dipole_charge.pysource/tests/consistent/modifier/__init__.pysource/tests/consistent/modifier/test_dipole_charge.py
njzjz
force-pushed
the
feat/dplr-dpmodel-backends
branch
from
fc8a04c to
75b3660
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
deepmd/jax/modifier/dipole_charge.py (1)
71-109: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winStill recomputing
energy_fna second time (dipole model runs twice).
jax.value_and_gradat Line 101-103 already evaluatesenergy_fnonce to get the summed scalar, then Line 104 callsenergy_fn(coord, strain)again just to recover the per-frame values — this reruns the full dipole-model forward pass and Ewald sum a second time per invocation. This was flagged in a previous review round and does not appear to have been fixed in this version.♻️ Proposed fix using has_aux to avoid the duplicate forward pass
- strain = jnp.zeros((coord.shape[0], 3, 3), dtype=coord.dtype) - energy, gradients = jax.value_and_grad( - lambda cc, ss: jnp.sum(energy_fn(cc, ss)), argnums=(0, 1) - )(coord, strain) - energy_by_frame = energy_fn(coord, strain) + strain = jnp.zeros((coord.shape[0], 3, 3), dtype=coord.dtype) + + def summed_energy( + cc: jnp.ndarray, ss: jnp.ndarray + ) -> tuple[jnp.ndarray, jnp.ndarray]: + per_frame = energy_fn(cc, ss) + return jnp.sum(per_frame), per_frame + + (_, energy_by_frame), gradients = jax.value_and_grad( + summed_energy, argnums=(0, 1), has_aux=True + )(coord, strain) return { "energy": energy_by_frame, "force": -gradients[0], "virial": -jnp.swapaxes(gradients[1], -1, -2).reshape(coord.shape[0], 9), }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/jax/modifier/dipole_charge.py` around lines 71 - 109, Update the energy and gradient calculation around energy_fn and jax.value_and_grad to use has_aux, returning the per-frame energy as auxiliary data while differentiating the summed energy. Reuse the auxiliary energy output for energy_by_frame and remove the second energy_fn(coord, strain) call, preserving the existing force and virial calculations.Source: Linters/SAST tools
🧹 Nitpick comments (1)
deepmd/pt/modifier/dipole_charge.py (1)
25-30: 🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy liftPer-frame Python loop with repeated
.item()GPU syncs on a hot path.
ewald_reciprocal_energyloops over frames in Python (line 55), and_even_grid_sizecalls.item()(line 27) three times per frame. For batched training/MD steps this forces3 * nframesCPU/GPU synchronizations every forward call, serializing an otherwise vectorizable computation.At minimum, the host sync can be reduced to one call per batch by computing all cell-vector norms in a single batched op and moving to host once, rather than per-frame/per-axis.
♻️ Sketch of reducing host syncs
-def _even_grid_size(length: torch.Tensor, spacing: float) -> int: - """Return the even reciprocal grid extent used by the C++ Ewald kernel.""" - size = int(torch.ceil(length / spacing).item()) - if size % 2 != 0: - size += 1 - return size +def _even_grid_sizes(lengths: torch.Tensor, spacing: float) -> list[int]: + """Return even reciprocal grid extents for a batch of lengths in one sync.""" + sizes = torch.ceil(lengths / spacing).to(torch.int64).tolist() + return [s + 1 if s % 2 != 0 else s for s in sizes]Then compute
norms = torch.linalg.vector_norm(box, dim=-1)once outside the frame loop and index into the precomputedgrid_sizeslist inside it.Also applies to: 53-84
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/pt/modifier/dipole_charge.py` around lines 25 - 30, Refactor ewald_reciprocal_energy and _even_grid_size to compute all cell-vector norms and even grid extents in one batched operation before the per-frame loop, then perform a single host transfer to obtain grid sizes. Use the precomputed grid_sizes entries inside the loop and eliminate repeated per-axis .item() calls while preserving the existing even-size calculation.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/modifier/dipole_charge.py`:
- Around line 37-86: Validate charge-map coverage before the charge construction
in extend_dplr_system, matching make_charge_maps behavior. Reject any
non-virtual atype value outside sys_charge_map and any sel_type entry outside
the same map instead of silently assigning zero or ignoring it. Preserve support
for negative virtual atom types.
In `@deepmd/pt_expt/modifier/dipole_charge.py`:
- Around line 53-56: Ensure the modifier’s dipole model remains in evaluation
mode after later parent `.train()` calls, not only during initialization. Update
the relevant training-mode behavior around the registered `self.dipole_model` in
the modifier class so recursive mode changes cannot re-enable training for this
frozen correction model, while preserving normal mode handling for the modifier
itself.
In `@deepmd/pt/modifier/dipole_charge.py`:
- Around line 198-200: Update the scripted/exported method signatures for
_energy and forward to replace every torch.Tensor | None annotation with
typing.Optional[torch.Tensor]. Ensure typing.Optional is imported and preserve
the existing parameter and return behavior.
---
Duplicate comments:
In `@deepmd/jax/modifier/dipole_charge.py`:
- Around line 71-109: Update the energy and gradient calculation around
energy_fn and jax.value_and_grad to use has_aux, returning the per-frame energy
as auxiliary data while differentiating the summed energy. Reuse the auxiliary
energy output for energy_by_frame and remove the second energy_fn(coord, strain)
call, preserving the existing force and virial calculations.
---
Nitpick comments:
In `@deepmd/pt/modifier/dipole_charge.py`:
- Around line 25-30: Refactor ewald_reciprocal_energy and _even_grid_size to
compute all cell-vector norms and even grid extents in one batched operation
before the per-frame loop, then perform a single host transfer to obtain grid
sizes. Use the precomputed grid_sizes entries inside the loop and eliminate
repeated per-axis .item() calls while preserving the existing even-size
calculation.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: c491e28f-1438-4075-ba96-2fccb9eb92b7
📒 Files selected for processing (12)
deepmd/dpmodel/modifier/__init__.pydeepmd/dpmodel/modifier/dipole_charge.pydeepmd/jax/modifier/__init__.pydeepmd/jax/modifier/dipole_charge.pydeepmd/pt/modifier/__init__.pydeepmd/pt/modifier/dipole_charge.pydeepmd/pt_expt/modifier/__init__.pydeepmd/pt_expt/modifier/dipole_charge.pydeepmd/tf2/modifier/__init__.pydeepmd/tf2/modifier/dipole_charge.pysource/tests/consistent/modifier/__init__.pysource/tests/consistent/modifier/test_dipole_charge.py
🚧 Files skipped from review as they are similar to previous changes (6)
- deepmd/pt_expt/modifier/init.py
- deepmd/dpmodel/modifier/init.py
- source/tests/consistent/modifier/init.py
- deepmd/jax/modifier/init.py
- source/tests/consistent/modifier/test_dipole_charge.py
- deepmd/pt/modifier/init.py
Add the shared Array API dipole-charge core and eager adapters for PyTorch exportable, JAX, and TensorFlow 2. Cover energy, force, virial, and Array API Strict consistency. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
njzjz
force-pushed
the
feat/dplr-dpmodel-backends
branch
from
75b3660 to
9f6ab7c
Compare
njzjz-bot
changed the title
Contributor
Author
|
Also addressed the review-level nitpicks:
Coding agent: Codex |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5775 +/- ##
==========================================
- Coverage 79.69% 79.55% -0.15%
==========================================
Files 1020 1027 +7
Lines 116359 116606 +247
Branches 4303 4305 +2
==========================================
+ Hits 92736 92763 +27
- Misses 22076 22297 +221
+ Partials 1547 1546 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
dipole_chargeschema and Array API numerical core in dpmodelmodel_charge_mapcontract: entries follow the embedded modelget_sel_type()orderRefs #5754.
Validation
ruff check .ruff format .git diff --check1.5e-13, virial difference approximately9.3e-14The combined local consistent-test process hits an environment-level Triton/JAX/PyTorch import segfault during collection, before test execution. Each backend was therefore validated in an isolated process, matching the backend-isolated CI layout.
Scope
This PR adds eager/Python
dipole_chargeadapters for pt_expt, JAX, and TF2. It intentionally does not add a native PyTorch implementation. Trainer label preprocessing and compiled artifact packaging remain follow-up work.Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
New Features
Tests