Repository navigation
[Feature Request] Migrate grid density model to dpmodel for multi-backend support #6029
Description
Summary
PR #5999 implements a charge density model on grid points for the PyTorch
backend only. It should eventually be implemented once in dpmodel and
wrapped per backend. This issue tracks the current state and the
migration plan.
Detailed Description
Background
PR #5999 introduces a charge density model that predicts scalar densities on
grid points. The implementation currently exists only for the PyTorch
backend (deepmd/pt/): DensityFittingNet, DPDensityAtomicModel,
GridDensityModel, GridDensityLoss, and the DeepDensity evaluator.
This does not follow the project's direction of implementing new
functionality once in dpmodel and wrapping it for each backend
(pt / pt_expt / tf / jax / pd).
Current state
DensityFittingNetis a plainInvarFittingsubclass — nothing
backend-specific.- The grid/atom merge in
DPDensityAtomicModel.forward_atomicand the
directional grid-to-atom neighbor list both have dpmodel counterparts
(or can be expressed with existing dpmodel/array-api utilities). - The data-loading side is already backend-agnostic:
grid/densityare
declared with the loader'sspecial_shape="frame_major"contract. - The reserved grid type is the last entry of
type_map; its descriptor
input statistics are computed from the grid data during stat.
Plan
- Implement the reference version in
dpmodel(descriptor-agnostic
atomic model + density fitting + output definition). - Wrap it for the PyTorch backend (reuse the existing pt classes as the
wrapping layer where possible). - Add
pt_exptsupport once the dpmodel reference exists. - (Optional) tf/jax/pd wrappers following the same pattern.
- Follow-up optimization noted during review: evaluate the descriptor
only on grid centers instead of all merged grid+atom points
(requires center-subset support in the descriptor API).
References
- PR feat(pt): add charge density prediction support #5999 (this issue is linked from its description)
Further Information, Files, and Links
No response
Activity
fix(pt): close review round on density eval, statistics, and loader c…
njzjz-bot commentedon Oct 2, 2026
The proposed dpmodel-first migration is a sensible follow-up to #5999. I checked its current atomic-model implementation: it merges grid centers and extended atoms, shifts valid neighbor indices while retaining
-1padding, and sends only the grid rows of the descriptor output to the fitting net. Those are useful invariants to preserve explicitly in the reference implementation.I suggest making the first migration PR cover the dpmodel fitting/atomic/model definitions, serialization, and PyTorch parity together. Compare the same serialized weights and inputs across the reference and PyTorch paths, including periodic ghosts, grid points coincident with atoms, padded neighbors, the reserved final grid type, input statistics, and frame parameters. The output and loss contracts should continue to use the number of grid points, not the number of atoms; document whether variable/empty grids are supported or rejected.
Keep the center-subset optimization separate from that parity change. For message-passing descriptors, atom rows whose final outputs are discarded may still provide intermediate features used by grid rows, so skipping them needs a descriptor-specific equivalence argument. A direct migration of the current evaluation semantics gives the pt_expt wrapper a much clearer reference. Other backends can then be enabled as their actual inference/training tests pass.
Coding agent: Codex
Codex version: codex-cli 0.160.0
Model: gpt-6-astra
Reasoning effort: xhigh