Density for pr#5465
Open
YuzhiLiu-ai wants to merge 11 commits into
Open
Conversation
- Add DensityFittingNet and GridDensityModel for grid-based charge density - Add DPDensityAtomicModel with grid neighbor list handling - Add GridDensityLoss for training on density labels - Support density data loading (grid.npy / density.npy) - Support model inference with grid input (DeepPot / DeepEval) - Add density model argument checking (fitting_density, loss_grid_density) - Skip change_out_bias for density models (grid-based output) - Add example configs and evaluation script for density models
- Add DensityFittingNet and GridDensityModel for grid-based charge density - Add DPDensityAtomicModel with grid neighbor list handling - Add GridDensityLoss for training on density labels - Support density data loading (grid.npy / density.npy) - Support model inference with grid input (DeepPot / DeepEval) - Add density model argument checking (fitting_density, loss_grid_density) - Skip change_out_bias for density models (grid-based output) - Add example configs and evaluation script for density models
…/deepmd-charge-density into chg_for_merge_branch
for more information, see https://pre-commit.ci
Contributor
📝 WalkthroughWalkthroughThis PR adds charge density prediction on grid points: new density fitting net, grid-aware atomic model and model factory, GridDensityModel registration, grid-aware loss and training integration, inference/data-loading fast-paths, argument schemas, and examples. ChangesGrid Density Model Implementation
Loss, Training Integration, and Distributed Setup
Inference and Data Loading
Configuration Schemas and Documentation
Sequence Diagram(s)sequenceDiagram
participant EvalScript
participant DeepEval
participant GridDensityModel
EvalScript->>DeepEval: eval(coord, atype, grid)
DeepEval->>GridDensityModel: forward(coord, atype, grid=grid)
GridDensityModel->>DeepEval: {"density": tensor, "mask": optional}
DeepEval->>EvalScript: return reshaped density array
🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 16
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
🟡 Minor comments (6)
deepmd/pt/model/task/density.py-40-40 (1)
40-40:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAvoid mutable default for
neuroninDensityFittingNet.__init__
deepmd/pt/model/task/density.py:40usesneuron: list[int] = [128, 128, 128], and Ruff flags this as B006.Proposed fix
- neuron: list[int] = [128, 128, 128], + neuron: list[int] | None = None, @@ - super().__init__( + neuron = [128, 128, 128] if neuron is None else neuron + super().__init__(As per coding guidelines
**/*.py: Install linter and runruff check .before committing changes or the CI will fail; also runruff format ..🤖 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/model/task/density.py` at line 40, The parameter default for neuron in DensityFittingNet.__init__ uses a mutable list which triggers B006; change the signature to use an immutable default (e.g., tuple) or None (neuron: Optional[list[int]] = None) and inside __init__ set self.neuron = list(neuron) or self.neuron = [128,128,128] when None; update any uses of neuron within DensityFittingNet to rely on self.neuron.examples/density/README.md-11-22 (1)
11-22:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAdd a language identifier to the fenced code block.
The directory-tree fence is missing a language tag, which triggers markdownlint MD040.
🤖 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 `@examples/density/README.md` around lines 11 - 22, In README.md fix the fenced code block for the directory tree by adding a language identifier (e.g., change the opening "```" to "```text" or "```bash") so markdownlint MD040 is satisfied; update the block that contains the tree starting with "." and the child lines (the directory-tree fence) to include the chosen language tag.examples/density/README.md-167-174 (1)
167-174:⚠️ Potential issue | 🟡 Minor | ⚡ Quick win“Full dataset” example currently samples only 10%.
This section title says full dataset, but
--ratio 0.1still subsamples. Use--ratio 1.0(or remove the flag) to match the heading.🤖 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 `@examples/density/README.md` around lines 167 - 174, The README example labeled "Evaluate the Full QM9 Dataset" incorrectly passes --ratio 0.1 (subsampling 10%); update the command in examples/density/README.md so it either removes the --ratio flag or sets --ratio 1.0 to actually run on the full dataset (the change affects the example invocation of dptest_density_script.py and the --ratio argument).examples/density/dptest_density_script.py-7-7 (1)
7-7:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate the usage example to the current script name.
The usage line references
test_density_new.py, but this file isdptest_density_script.py.🤖 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 `@examples/density/dptest_density_script.py` at line 7, Update the usage example to reference the current script name: replace the referenced script `test_density_new.py` with `dptest_density_script.py` in the usage line (the command shown near the top of dptest_density_script.py), keeping the same arguments (`model.pt /path/to/data --ratio 0.1 --output result.txt`) so the example reflects the actual script filename used.examples/density/dptest_density_script.py-152-152 (1)
152-152:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winGuard
epsilon_MAEagainst zero-label edge cases.When
Mean|label| == 0, this printsinf/nan. Add a zero guard for stable summary output.🤖 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 `@examples/density/dptest_density_script.py` at line 152, The print currently divides mae by label_mean_abs causing inf/nan when Mean|label| == 0; compute a guarded epsilon_MAE value (e.g., set epsilon_MAE = mae / label_mean_abs if label_mean_abs != 0 else 0.0 or float('nan') as preferred) before printing, and change the print line to reference epsilon_MAE instead of performing the division inline so the summary output remains stable; update the symbol epsilon_MAE and use the existing label_mean_abs and mae variables.examples/density/README.md-148-148 (1)
148-148:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winAvoid machine-specific absolute paths in usage docs.
The hardcoded local path is not portable; use a repo-relative path (
cd examples/density) instead.🤖 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 `@examples/density/README.md` at line 148, Replace the machine-specific absolute path string "cd /aisi/yuzhiLiu/deepmd-kit-charge/deepmd-kit-dpa3/examples/density" in examples/density/README.md with a repo-relative invocation such as "cd examples/density" so the usage doc is portable across different developer environments; update that single line in the README to the relative path.
🧹 Nitpick comments (1)
deepmd/infer/deep_pot.py (1)
138-149: ⚡ Quick winUpdate
DeepPot.evaltyping/docs for density mode return type.This method now returns a bare
np.ndarrayfor grid-density mode, but overloads/docstring still describe tuple-only returns. Please add/adjust an overload/doc note for the grid path to avoid API ambiguity.Also applies to: 173-189, 216-218
🤖 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/infer/deep_pot.py` around lines 138 - 149, DeepPot.eval currently documents/annotates returns as a tuple but in grid-density mode it returns a bare np.ndarray; update DeepPot.eval's type hints and docstring to reflect both possibilities by adding an explicit overload (or Union return type) that returns np.ndarray for the grid-density path and the existing tuple[np.ndarray, ...] for the normal path, and clarify in the docstring when each form is returned; apply the same typing/docstring fix to the other overloaded signatures and doc blocks in this module that correspond to the same behavior (the alternate eval signatures and their docs).
🤖 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/infer/deep_pot.py`:
- Around line 216-218: The condition that returns density currently uses if
"grid" in kwargs which also matches grid=None and can access missing
results["density"]; change the predicate to check the actual value (e.g.,
kwargs.get("grid") is not None) before reshaping results["density"] in the
function where nframes and results are used so the density branch only runs when
grid is not None.
In `@deepmd/pt/infer/deep_eval.py`:
- Around line 432-434: The default non-spin/non-grid evaluation call to
self._eval_func(...)(...) omits the required charge_spin argument for
_eval_model; update the call site (the invocation that constructs out via
self._eval_func(self._eval_model, numb_test, natoms)(...)) to include the
charge_spin parameter (e.g., add charge_spin to the argument list alongside
coords, cells, atom_types, fparam, aparam, request_defs), and apply the same
change to the other similar block around the 527-536 region so both calls pass
charge_spin into _eval_model.
- Line 430: The "density" payload is returning a tuple (from
_eval_model_density/tuple(results)) which breaks DeepPot.eval when it calls
.reshape; convert the tuple to a NumPy array before returning. Locate the return
sites (the dict with "density": out around the return in deep_eval.py and the
similar block at lines ~769-775) and replace the tuple with a NumPy ndarray
(e.g., wrap results with numpy.array or ensure tuple(results) is created as
np.array(results)); add or use the existing import for numpy (np) so callers
receive an ndarray instead of a tuple.
In `@deepmd/pt/loss/charge.py`:
- Line 44: The density-loss activation logic in charge.py sets self.has_d
incorrectly using AND between start_pref_d and limit_pref_d so density loss is
disabled when either endpoint is zero; change the condition in the initializer
that assigns self.has_d to use OR between the endpoint checks (i.e., treat
density active if either start_pref_d or limit_pref_d is nonzero) while still
preserving the existing inference flag check (keep the overall OR with
inference); locate the assignment to self.has_d in charge.py and update the
boolean expression accordingly.
- Around line 1-3: CI currently fails ruff formatting/linting; run "ruff format
." to reformat the repository and then "ruff check ." to resolve remaining lint
errors, fix any reported issues, and commit the changes (ensure the module
deepmd.pt.loss.charge and any generated scripts like implib-gen.py are
reformatted). After fixing, re-run the CI lint step locally to confirm passing
before pushing.
In `@deepmd/pt/model/atomic_model/density_atomic_model.py`:
- Around line 134-136: Replace creation of grid type tensors that use zeros with
tensors filled with the reserved virtual type index (ntypes - 1): where code
currently calls torch.zeros([...], device=grid_type.device,
dtype=grid_type.dtype) to build grid_type, allocate a tensor with the same
shape, device and dtype but filled with the value (ntypes - 1) instead of 0;
update both occurrences involving the grid_type construction (the shown
torch.zeros call and the similar block at lines 286-288) so grid_type uses the
reserved type index rather than 0.
- Around line 240-258: The forward method in DensityAtomicModel currently omits
the required grid inputs when delegating to forward_common_atomic, causing
assertion failures; update the forward signature to accept grid: torch.Tensor |
None, grid_type: torch.Tensor | None, and grid_nlist: torch.Tensor | None (or
the same types used elsewhere) and forward them through (i.e., add parameters
grid, grid_type, grid_nlist to def forward(...) and include grid=grid,
grid_type=grid_type, grid_nlist=grid_nlist in the call to
self.forward_common_atomic) so the assertions in forward_common_atomic (grid,
grid_type, grid_nlist) are satisfied.
In `@deepmd/pt/model/model/make_density_model.py`:
- Around line 178-180: Replace the zero-initialized grid_type with the reserved
virtual type value (ntypes - 1) instead of 0: locate the grid_type creation in
make_density_model.py where grid_type = torch.zeros(..., device=gg.device,
dtype=atype.dtype) and change it to initialize all entries to the virtual type
value (use the model/config ntypes or available ntypes variable) while keeping
device=gg.device and dtype=atype.dtype so all grid points use the reserved
virtual atom type during inference.
- Around line 629-646: The forward method is incorrectly passing box as the
third positional argument to forward_common (the third parameter is grid), so
update the call in forward (function forward) to pass box as a keyword argument
(box=box) or otherwise ensure the third positional arg is the intended grid;
modify the forward_common invocation in make_density_model.forward to use named
parameters like box=box (and keep fparam=, aparam=, do_atomic_virial=) so inputs
map correctly to forward_common's signature.
- Around line 92-97: The local variable named "vars" shadows the built-in and
triggers Ruff A001; rename it to a non-built-in identifier (e.g., out_vars or
var_list) inside the function that iterates over var_defs and checks vv.category
== OutputVariableCategory.OUT.value, update all references and the return to use
the new name (ensure you only change the local symbol in make_density_model.py
around the loop over var_defs), and re-run the linter (ruff check .) before
committing.
In `@deepmd/pt/utils/dataloader.py`:
- Line 225: The function DpLoaderSet.print_summary contains an accidental early
"return" that makes the method a no-op and prevents the rank-checked summary
block from running; remove that stray return so the remainder of
DpLoaderSet.print_summary (including the rank check and summary printing logic)
executes normally, ensuring the existing summary code path in the method is
reachable again.
In `@deepmd/utils/argcheck.py`:
- Around line 3499-3541: There are two identical registrations/definitions of
loss_grid_density using `@loss_args_plugin.register`("grid_density"); remove the
duplicate block so only one loss_grid_density function remains (keep a single
`@loss_args_plugin.register`("grid_density") decorator + its associated
loss_grid_density definition that returns the Argument list), ensuring you don't
change the Argument names/start_pref/limit_pref usage or defaults.
In `@deepmd/utils/data.py`:
- Around line 816-818: The code returns rank-3 arrays for keys "grid"/"density"
which breaks _shuffle_data (expects 2D), causing coord shuffle misalignment;
change the loader in data.py so after data = path.load_numpy().astype(dtype) you
detect 3D arrays and reshape them to 2D with frames as the first dimension
(e.g., data = data.reshape(data.shape[0], -1)) and ensure dtype is preserved,
then return np.float32(1.0), data; reference the "grid"/"density" key handling
and the _shuffle_data assumption when making the change.
In `@examples/density/dataset/qm9/C7H15NO_val/type_map.raw`:
- Around line 1-9: The last entry in the type_map (used as the virtual grid type
where atype = ntypes - 1) is currently "F", which reserves a real element;
replace that final entry with a dedicated placeholder (e.g., "X") so the virtual
grid is unambiguous, and update any corresponding indices in type.raw to remain
consistent with the new mapping; locate the mapping in
examples/density/dataset/qm9/C7H15NO_val/type_map.raw and ensure the reserved
slot and type.raw references use the new "X" placeholder.
In `@examples/density/dpa3/input.json`:
- Around line 79-106: Remove the invalid training.opt_type entry and instead
add/update the top-level optimizer block to declare the optimizer type;
specifically, take the value from training.opt_type ("AdamW"), delete the
training.opt_type key, and set optimizer.type = "AdamW" (and move any
optimizer-specific params into that top-level optimizer block) so schema
validation uses the proper optimizer configuration instead of training.opt_type.
In `@examples/density/dptest_density_script.py`:
- Around line 56-61: The --ratio argument must be validated to avoid n_sample >
n_frames when calling random.sample; add validation for the parsed ratio (0 <
ratio <= 1 or 0 <= ratio <=1 as intended) either by using an argparse
type/validator that raises argparse.ArgumentTypeError or by checking the parsed
args.ratio before sampling, and if out of range raise a clear error. Update the
code that computes n_sample (variables n_sample, n_frames and the call to
random.sample on lines around 116-117) to rely on the validated ratio and/or
clamp it to at most 1 to prevent ValueError from random.sample.
---
Minor comments:
In `@deepmd/pt/model/task/density.py`:
- Line 40: The parameter default for neuron in DensityFittingNet.__init__ uses a
mutable list which triggers B006; change the signature to use an immutable
default (e.g., tuple) or None (neuron: Optional[list[int]] = None) and inside
__init__ set self.neuron = list(neuron) or self.neuron = [128,128,128] when
None; update any uses of neuron within DensityFittingNet to rely on self.neuron.
In `@examples/density/dptest_density_script.py`:
- Line 7: Update the usage example to reference the current script name: replace
the referenced script `test_density_new.py` with `dptest_density_script.py` in
the usage line (the command shown near the top of dptest_density_script.py),
keeping the same arguments (`model.pt /path/to/data --ratio 0.1 --output
result.txt`) so the example reflects the actual script filename used.
- Line 152: The print currently divides mae by label_mean_abs causing inf/nan
when Mean|label| == 0; compute a guarded epsilon_MAE value (e.g., set
epsilon_MAE = mae / label_mean_abs if label_mean_abs != 0 else 0.0 or
float('nan') as preferred) before printing, and change the print line to
reference epsilon_MAE instead of performing the division inline so the summary
output remains stable; update the symbol epsilon_MAE and use the existing
label_mean_abs and mae variables.
In `@examples/density/README.md`:
- Around line 11-22: In README.md fix the fenced code block for the directory
tree by adding a language identifier (e.g., change the opening "```" to
"```text" or "```bash") so markdownlint MD040 is satisfied; update the block
that contains the tree starting with "." and the child lines (the directory-tree
fence) to include the chosen language tag.
- Around line 167-174: The README example labeled "Evaluate the Full QM9
Dataset" incorrectly passes --ratio 0.1 (subsampling 10%); update the command in
examples/density/README.md so it either removes the --ratio flag or sets --ratio
1.0 to actually run on the full dataset (the change affects the example
invocation of dptest_density_script.py and the --ratio argument).
- Line 148: Replace the machine-specific absolute path string "cd
/aisi/yuzhiLiu/deepmd-kit-charge/deepmd-kit-dpa3/examples/density" in
examples/density/README.md with a repo-relative invocation such as "cd
examples/density" so the usage doc is portable across different developer
environments; update that single line in the README to the relative path.
---
Nitpick comments:
In `@deepmd/infer/deep_pot.py`:
- Around line 138-149: DeepPot.eval currently documents/annotates returns as a
tuple but in grid-density mode it returns a bare np.ndarray; update
DeepPot.eval's type hints and docstring to reflect both possibilities by adding
an explicit overload (or Union return type) that returns np.ndarray for the
grid-density path and the existing tuple[np.ndarray, ...] for the normal path,
and clarify in the docstring when each form is returned; apply the same
typing/docstring fix to the other overloaded signatures and doc blocks in this
module that correspond to the same behavior (the alternate eval signatures and
their docs).
🪄 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: 7fd9357d-b6e7-4173-8319-a1474b2678eb
📒 Files selected for processing (34)
deepmd/infer/deep_pot.pydeepmd/pt/entrypoints/main.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/loss/__init__.pydeepmd/pt/loss/charge.pydeepmd/pt/model/atomic_model/__init__.pydeepmd/pt/model/atomic_model/density_atomic_model.pydeepmd/pt/model/model/__init__.pydeepmd/pt/model/model/density_model.pydeepmd/pt/model/model/make_density_model.pydeepmd/pt/model/task/__init__.pydeepmd/pt/model/task/density.pydeepmd/pt/train/training.pydeepmd/pt/train/wrapper.pydeepmd/pt/utils/dataloader.pydeepmd/pt/utils/stat.pydeepmd/utils/argcheck.pydeepmd/utils/data.pyexamples/density/README.mdexamples/density/dataset/qm9/C7H15NO_train/set.000/box.npyexamples/density/dataset/qm9/C7H15NO_train/set.000/coord.npyexamples/density/dataset/qm9/C7H15NO_train/set.000/density.npyexamples/density/dataset/qm9/C7H15NO_train/set.000/grid.npyexamples/density/dataset/qm9/C7H15NO_train/type.rawexamples/density/dataset/qm9/C7H15NO_train/type_map.rawexamples/density/dataset/qm9/C7H15NO_val/set.000/box.npyexamples/density/dataset/qm9/C7H15NO_val/set.000/coord.npyexamples/density/dataset/qm9/C7H15NO_val/set.000/density.npyexamples/density/dataset/qm9/C7H15NO_val/set.000/grid.npyexamples/density/dataset/qm9/C7H15NO_val/type.rawexamples/density/dataset/qm9/C7H15NO_val/type_map.rawexamples/density/dpa2/input.jsonexamples/density/dpa3/input.jsonexamples/density/dptest_density_script.py
Comment on lines
+216
to
+218
| if "grid" in kwargs: | ||
| result = results["density"].reshape(nframes, -1) | ||
| return result |
Contributor
There was a problem hiding this comment.
Guard against grid=None before density extraction.
if "grid" in kwargs also matches grid=None, but the backend density branch only runs when grid is not None; this can make results["density"] missing and raise at runtime. Use the same predicate here (kwargs.get("grid") is not None).
Suggested fix
- if "grid" in kwargs:
+ if kwargs.get("grid") is not None:
result = results["density"].reshape(nframes, -1)
return result🤖 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/infer/deep_pot.py` around lines 216 - 218, The condition that returns
density currently uses if "grid" in kwargs which also matches grid=None and can
access missing results["density"]; change the predicate to check the actual
value (e.g., kwargs.get("grid") is not None) before reshaping results["density"]
in the function where nframes and results are used so the density branch only
runs when grid is not None.
| aparam, | ||
| request_defs, | ||
| ) | ||
| return {"density": out} |
Contributor
There was a problem hiding this comment.
Return ndarray instead of tuple for "density" payload.
out here is a tuple from _eval_model_density (currently tuple(results)), so {"density": out} makes downstream callers receive a tuple, not an array. DeepPot.eval then calls .reshape(...) on this value and will fail.
Suggested fix
- return {"density": out}
+ return {"density": out[0]}Also applies to: 769-775
🤖 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/infer/deep_eval.py` at line 430, The "density" payload is returning
a tuple (from _eval_model_density/tuple(results)) which breaks DeepPot.eval when
it calls .reshape; convert the tuple to a NumPy array before returning. Locate
the return sites (the dict with "density": out around the return in deep_eval.py
and the similar block at lines ~769-775) and replace the tuple with a NumPy
ndarray (e.g., wrap results with numpy.array or ensure tuple(results) is created
as np.array(results)); add or use the existing import for numpy (np) so callers
receive an ndarray instead of a tuple.
Comment on lines
+432
to
+434
| out = self._eval_func(self._eval_model, numb_test, natoms)( | ||
| coords, cells, atom_types, fparam, aparam, request_defs | ||
| ) |
Contributor
There was a problem hiding this comment.
Pass charge_spin in the default eval path.
_eval_model(...) requires charge_spin, but the non-spin/non-grid call omits it. This will raise TypeError for normal inference.
Suggested fix
- out = self._eval_func(self._eval_model, numb_test, natoms)(
- coords, cells, atom_types, fparam, aparam, request_defs
- )
+ out = self._eval_func(self._eval_model, numb_test, natoms)(
+ coords, cells, atom_types, fparam, aparam, request_defs, charge_spin
+ )Also applies to: 527-536
🤖 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/infer/deep_eval.py` around lines 432 - 434, The default
non-spin/non-grid evaluation call to self._eval_func(...)(...) omits the
required charge_spin argument for _eval_model; update the call site (the
invocation that constructs out via self._eval_func(self._eval_model, numb_test,
natoms)(...)) to include the charge_spin parameter (e.g., add charge_spin to the
argument list alongside coords, cells, atom_types, fparam, aparam,
request_defs), and apply the same change to the other similar block around the
527-536 region so both calls pass charge_spin into _eval_model.
| """ | ||
| super().__init__() | ||
| self.starter_learning_rate = starter_learning_rate | ||
| self.has_d = (start_pref_d != 0.0 and limit_pref_d != 0.0) or inference |
Contributor
There was a problem hiding this comment.
Fix density-loss activation logic on Line 44.
Using and here disables density loss whenever only one prefactor endpoint is zero (common for ramp-up/ramp-down schedules), so training can silently skip this objective.
Suggested fix
- self.has_d = (start_pref_d != 0.0 and limit_pref_d != 0.0) or inference
+ self.has_d = (start_pref_d != 0.0 or limit_pref_d != 0.0) or inference📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| self.has_d = (start_pref_d != 0.0 and limit_pref_d != 0.0) or inference | |
| self.has_d = (start_pref_d != 0.0 or limit_pref_d != 0.0) or inference |
🤖 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/loss/charge.py` at line 44, The density-loss activation logic in
charge.py sets self.has_d incorrectly using AND between start_pref_d and
limit_pref_d so density loss is disabled when either endpoint is zero; change
the condition in the initializer that assigns self.has_d to use OR between the
endpoint checks (i.e., treat density active if either start_pref_d or
limit_pref_d is nonzero) while still preserving the existing inference flag
check (keep the overall OR with inference); locate the assignment to self.has_d
in charge.py and update the boolean expression accordingly.
Comment on lines
+816
to
+818
| if key in ["grid", "density"] and path.is_file(): | ||
| data = path.load_numpy().astype(dtype) | ||
| return np.float32(1.0), data |
Contributor
There was a problem hiding this comment.
Preserve shuffle alignment for 3D grid/density arrays.
Returning raw rank-3 arrays here breaks existing _shuffle_data assumptions (len(shape)==2), so these keys won’t be shuffled with coord, causing silent frame misalignment in batch/test set loading.
Suggested fix direction
- if (
- isinstance(data[kk], np.ndarray)
- and len(data[kk].shape) == 2
- and data[kk].shape[0] == nframes
- and "find_" not in kk
- ):
+ if (
+ isinstance(data[kk], np.ndarray)
+ and data[kk].ndim >= 1
+ and data[kk].shape[0] == nframes
+ and "find_" not in kk
+ ):
ret[kk] = data[kk][idx]🤖 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/utils/data.py` around lines 816 - 818, The code returns rank-3 arrays
for keys "grid"/"density" which breaks _shuffle_data (expects 2D), causing coord
shuffle misalignment; change the loader in data.py so after data =
path.load_numpy().astype(dtype) you detect 3D arrays and reshape them to 2D with
frames as the first dimension (e.g., data = data.reshape(data.shape[0], -1)) and
ensure dtype is preserved, then return np.float32(1.0), data; reference the
"grid"/"density" key handling and the _shuffle_data assumption when making the
change.
Comment on lines
+1
to
+9
| Li | ||
| Ni | ||
| Co | ||
| Mn | ||
| O | ||
| C | ||
| H | ||
| N | ||
| F |
Contributor
There was a problem hiding this comment.
Use an explicit virtual grid type in the last slot.
For density mode, the last type_map entry is treated as the virtual grid type (atype = ntypes - 1). Using F here makes the reserved slot a real element, which is ambiguous and can collide when reusing this example on fluorine-containing systems. Please switch the last entry to a dedicated placeholder (for example, X) and keep type.raw indices consistent with that mapping.
🤖 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 `@examples/density/dataset/qm9/C7H15NO_val/type_map.raw` around lines 1 - 9,
The last entry in the type_map (used as the virtual grid type where atype =
ntypes - 1) is currently "F", which reserves a real element; replace that final
entry with a dedicated placeholder (e.g., "X") so the virtual grid is
unambiguous, and update any corresponding indices in type.raw to remain
consistent with the new mapping; locate the mapping in
examples/density/dataset/qm9/C7H15NO_val/type_map.raw and ensure the reserved
slot and type.raw references use the new "X" placeholder.
Comment on lines
+79
to
+106
| "training": { | ||
| "stat_file": "./qm9_charge_density.hdf5", | ||
| "training_data": { | ||
| "systems": [ | ||
| "../dataset/qm9/C7H15NO_train" | ||
| ], | ||
| "batch_size": "auto:128", | ||
| "_comment": "that's all" | ||
| }, | ||
| "validation_data": { | ||
| "systems": [ | ||
| "../dataset/qm9/C7H15NO_val" | ||
| ], | ||
| "batch_size": 1, | ||
| "numb_btch": 3, | ||
| "_comment": "that's all" | ||
| }, | ||
| "numb_steps": 100, | ||
| "warmup_steps": 0, | ||
| "gradient_max_norm": 5.0, | ||
| "seed": 10, | ||
| "max_ckpt_keep": 100, | ||
| "opt_type": "AdamW", | ||
| "disp_file": "lcurve.out", | ||
| "disp_freq": 100, | ||
| "save_freq": 10000, | ||
| "_comment": "that's all" | ||
| } |
Contributor
There was a problem hiding this comment.
training.opt_type is not a valid config key.
Optimizer selection is configured via the top-level optimizer block. Keeping opt_type under training can fail strict schema validation.
Proposed fix
"learning_rate": {
"type": "exp",
"decay_steps": 5000,
"start_lr": 0.001,
"stop_lr": 1e-05,
"_comment": "that's all"
},
+ "optimizer": {
+ "type": "AdamW"
+ },
"loss": {
"type": "grid_density",
"start_pref_d": 1,
"limit_pref_d": 1,
"_comment": " that's all"
},
@@
- "opt_type": "AdamW",🤖 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 `@examples/density/dpa3/input.json` around lines 79 - 106, Remove the invalid
training.opt_type entry and instead add/update the top-level optimizer block to
declare the optimizer type; specifically, take the value from training.opt_type
("AdamW"), delete the training.opt_type key, and set optimizer.type = "AdamW"
(and move any optimizer-specific params into that top-level optimizer block) so
schema validation uses the proper optimizer configuration instead of
training.opt_type.
Comment on lines
+56
to
+61
| parser.add_argument( | ||
| "--ratio", | ||
| type=float, | ||
| default=0.1, | ||
| help="Fraction of frames to randomly sample from each system (default: 0.1).", | ||
| ) |
Contributor
There was a problem hiding this comment.
Validate --ratio to prevent sampling runtime failures.
ratio > 1 can produce n_sample > n_frames, which raises ValueError in random.sample(...).
Proposed fix
+def _ratio_arg(value: str) -> float:
+ v = float(value)
+ if not (0.0 < v <= 1.0):
+ raise argparse.ArgumentTypeError("--ratio must be in (0, 1].")
+ return v
+
def parse_args() -> argparse.Namespace:
@@
parser.add_argument(
"--ratio",
- type=float,
+ type=_ratio_arg,
default=0.1,
help="Fraction of frames to randomly sample from each system (default: 0.1).",
)Also applies to: 116-117
🤖 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 `@examples/density/dptest_density_script.py` around lines 56 - 61, The --ratio
argument must be validated to avoid n_sample > n_frames when calling
random.sample; add validation for the parsed ratio (0 < ratio <= 1 or 0 <= ratio
<=1 as intended) either by using an argparse type/validator that raises
argparse.ArgumentTypeError or by checking the parsed args.ratio before sampling,
and if out of range raise a clear error. Update the code that computes n_sample
(variables n_sample, n_frames and the call to random.sample on lines around
116-117) to rely on the validated ratio and/or clamp it to at most 1 to prevent
ValueError from random.sample.
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt/model/model/make_density_model.py (1)
196-207:⚠️ Potential issue | 🟠 Major | ⚡ Quick winForward
charge_spinthrough both delegation layers.This wrapper accepts
charge_spinbut drops it in both calls, so charge-spin-capable models will silently ignore caller input and produce the wrong result. If density is not supposed to use charge-spin, the safer fix is to remove the parameter entirely instead of accepting and ignoring it.Proposed fix
model_predict_lower = self.forward_common_lower( extended_coord, extended_atype, nlist, grid=gg, grid_type=grid_type, grid_nlist=grid_nlist, mapping=mapping, do_atomic_virial=do_atomic_virial, fparam=fp, aparam=ap, + charge_spin=charge_spin, ) @@ atomic_ret = self.atomic_model.forward_common_atomic( cc_ext, extended_atype, nlist, mapping=mapping, fparam=fp, aparam=ap, comm_dict=comm_dict, grid=gg, grid_type=grid_type, grid_nlist=grid_nlist, + charge_spin=charge_spin, )Also applies to: 308-319
🤖 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/model/model/make_density_model.py` around lines 196 - 207, The wrapper currently drops the charge_spin argument when delegating to the lower/upper forwards; update the calls that construct model_predict_lower and model_predict_upper to pass the charge_spin parameter through to forward_common_lower and forward_common_upper (and ensure those functions' signatures accept and propagate charge_spin), or if density truly must not use charge_spin, remove charge_spin from the wrapper signature to avoid silent ignoring; target the model_predict_lower/model_predict_upper construction sites and the forward_common_lower/forward_common_upper function signatures to make the propagation consistent.
🤖 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.
Outside diff comments:
In `@deepmd/pt/model/model/make_density_model.py`:
- Around line 196-207: The wrapper currently drops the charge_spin argument when
delegating to the lower/upper forwards; update the calls that construct
model_predict_lower and model_predict_upper to pass the charge_spin parameter
through to forward_common_lower and forward_common_upper (and ensure those
functions' signatures accept and propagate charge_spin), or if density truly
must not use charge_spin, remove charge_spin from the wrapper signature to avoid
silent ignoring; target the model_predict_lower/model_predict_upper construction
sites and the forward_common_lower/forward_common_upper function signatures to
make the propagation consistent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: db4d6ca3-08a8-4307-af11-ae12fdfde4a4
📒 Files selected for processing (3)
deepmd/pt/model/atomic_model/density_atomic_model.pydeepmd/pt/model/model/density_model.pydeepmd/pt/model/model/make_density_model.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/pt/model/model/density_model.py
- deepmd/pt/model/atomic_model/density_atomic_model.py
for more information, see https://pre-commit.ci
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/utils/argcheck.py (1)
1897-1897:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix typo in documentation string.
"inital" should be "initial" in the rcond documentation.
📝 Proposed fix
- doc_rcond = "The condition number used to determine the inital density shift for each type of atoms. See `rcond` in :py:meth:`numpy.linalg.lstsq` for more details." + doc_rcond = "The condition number used to determine the initial density shift for each type of atoms. See `rcond` in :py:meth:`numpy.linalg.lstsq` for more details."🤖 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/utils/argcheck.py` at line 1897, The doc string assigned to the variable doc_rcond in deepmd/utils/argcheck.py contains a typo ("inital"); update the string to read "initial" so it becomes: "The condition number used to determine the initial density shift for each type of atoms. See `rcond` in :py:meth:`numpy.linalg.lstsq` for more details." Locate doc_rcond and replace the misspelled word only, preserving the rest of the text and formatting.
🤖 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.
Outside diff comments:
In `@deepmd/utils/argcheck.py`:
- Line 1897: The doc string assigned to the variable doc_rcond in
deepmd/utils/argcheck.py contains a typo ("inital"); update the string to read
"initial" so it becomes: "The condition number used to determine the initial
density shift for each type of atoms. See `rcond` in
:py:meth:`numpy.linalg.lstsq` for more details." Locate doc_rcond and replace
the misspelled word only, preserving the rest of the text and formatting.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 39486ecb-4a4a-4253-ac28-cbbe49d1fb24
📒 Files selected for processing (5)
deepmd/pt/model/model/density_model.pydeepmd/pt/model/model/make_density_model.pydeepmd/pt/model/task/density.pydeepmd/utils/argcheck.pyexamples/density/dptest_density_script.py
🚧 Files skipped from review as they are similar to previous changes (4)
- deepmd/pt/model/model/density_model.py
- deepmd/pt/model/task/density.py
- examples/density/dptest_density_script.py
- deepmd/pt/model/model/make_density_model.py
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Comment on lines
+64
to
+76
| def forward_atomic( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| comm_dict: dict[str, torch.Tensor] | None = None, | ||
| grid: torch.Tensor | None = None, | ||
| grid_type: torch.Tensor | None = None, | ||
| grid_nlist: torch.Tensor | None = None, | ||
| ) -> dict[str, torch.Tensor]: |
Comment on lines
+159
to
+171
| def forward_common_atomic( | ||
| self, | ||
| extended_coord: torch.Tensor, | ||
| extended_atype: torch.Tensor, | ||
| nlist: torch.Tensor, | ||
| mapping: torch.Tensor | None = None, | ||
| fparam: torch.Tensor | None = None, | ||
| aparam: torch.Tensor | None = None, | ||
| comm_dict: dict[str, torch.Tensor] | None = None, | ||
| grid: torch.Tensor | None = None, | ||
| grid_type: torch.Tensor | None = None, | ||
| grid_nlist: torch.Tensor | None = None, | ||
| ) -> dict[str, torch.Tensor]: |
|
|
||
| """ | ||
| log.warning("Not implemented yet for density out stat!") | ||
| pass |
| more_loss = {} | ||
| # more_loss['log_keys'] = [] # showed when validation on the fly | ||
| # more_loss['test_keys'] = [] # showed when doing dp test | ||
| atom_norm = 1.0 / natoms |
|
|
||
| """ | ||
| nframes, nloc, nnei = nlist.shape | ||
| atype = extended_atype[:, :nloc] |
| assert grid_nlist is not None | ||
| nframes, nloc, _ = nlist.shape | ||
| _, ngrid, _ = grid_nlist.shape | ||
| atype = extended_atype[:, :nloc] |
| ) | ||
| ret_dict = self.apply_out_stat(ret_dict, grid_type) | ||
|
|
||
| ext_grid_mask = self.make_atom_mask(grid_type) |
Comment on lines
+351
to
+356
| # for vv, kk in zip([fparam, aparam], ["frame", "atomic"]): | ||
| # if vv is not None and self.reverse_precision_dict[vv.dtype] != input_prec: | ||
| # log.warning( | ||
| # f"type of {kk} parameter {self.reverse_precision_dict[vv.dtype]}" | ||
| # " does not match" | ||
| # f" that of the coordinate {input_prec}" |
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
2 participants
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.
Feature: Add charge density prediction support for DPA3
Summary
This PR adds charge density prediction support to the PyTorch backend of DeePMD-kit.
The implementation enables a DPA3-based model to predict charge density values on user-defined spatial grid points. It extends the existing atom-centered energy/force framework to grid-based scalar field prediction.
The core idea is to treat each grid point as a virtual central atom. For each grid point, the model constructs a directional neighbor list from the grid point to surrounding real atoms, computes DPA3 descriptors, and applies a density fitting network to predict the charge density value at that grid location.
Key Changes
1. New density model architecture
GridDensityModelandmake_density_modelas model wrappers for grid-based density inference.DensityAtomicModelas the underlying atomic model for descriptor evaluation and density fitting.2. Data pipeline support
deepmd/utils/data.pyto support two additional data entries:griddensitytype_sel) because they are grid-based rather than atom-centered quantities.Example
A complete working example is provided in:
Data format
The dataset follows the standard
deepmd/npyformat, with two additional files in eachset.000/directory:coord.npy[nframes, natoms * 3]box.npy[nframes, 9]type.raw[natoms]type_map.raw[ntypes]grid.npy[nframes, ngrid, 3]density.npy[nframes, ngrid, 1]Training
cd examples/density/dpa3 dp --pt train input.jsonKey configuration in
input.json:{ "model": { "type_map": ["Li", "Ni", "Co", "Mn", "O", "C", "H", "N", "F", "X"], "descriptor": { "type": "dpa3", .... "env_protection": 0.1 }, "fitting_net": { "type": "density", "neuron": [240, 240, 240] } }, "loss": { "type": "grid_density", "start_pref_d": 1, "limit_pref_d": 1 } }Notes:
type_map,"X", is reserved for virtual grid points.atype = ntypes - 1, so that the descriptor treats them as distinct from real atoms.env_protectionis set to0.1to avoid numerical instability when a grid point is very close to, or coincides with, a real atom position.Evaluation
cd examples/density python dptest_density_script.py \ dpa3/model.ckpt-100.pt \ dataset/qm9/C7H15NO_val \ --ratio 0.1 \ --output val_result.txtNotes for reviewers
This implementation is optional and does not change the standard energy/force training workflow. The density model is activated only when the corresponding density model and loss types are specified in the input configuration.
Summary by CodeRabbit