[Feature Request] Migrate grid density model to dpmodel for multi-backend support

Open
  • Is something wrong?
#6,029 1 comment 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 3 days

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
35/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch, tensorflow
Domain
backend, machine-learning

Research direction

Start by reading the PyTorch implementations in deepmd/pt/, especially DensityFittingNet, DPDensityAtomicModel, GridDensityModel, GridDensityLoss, and DeepDensity, then compare them with the available dpmodel counterparts and PR #5999. Done means the reference model is implemented in dpmodel, PyTorch wraps it, and pt_expt support is added without losing the existing grid/density behavior.

Written by the indexing model from the issue text.

Description

enhancement
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
  • DensityFittingNet is a plain InvarFitting subclass — nothing
    backend-specific.
  • The grid/atom merge in DPDensityAtomicModel.forward_atomic and 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/density are
    declared with the loader's special_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
  1. Implement the reference version in dpmodel (descriptor-agnostic
    atomic model + density fitting + output definition).
  2. Wrap it for the PyTorch backend (reuse the existing pt classes as the
    wrapping layer where possible).
  3. Add pt_expt support once the dpmodel reference exists.
  4. (Optional) tf/jax/pd wrappers following the same pattern.
  5. 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 #5999 (this issue is linked from its description)
Further Information, Files, and Links

No response

Dominant language
Python
Stars
2.1k
Forks
663
Avg merge
3d 11h
Merged PRs (30d)
25

Getting set up

Open in Codespaces

Starts the project's dev container in your browser, under your own GitHub account.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from deepmodeling/deepmd-kit

All issues in deepmodeling/deepmd-kit

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.