refactor(pt/dpa4): use edge schema for dpa4 and ignore sel#5562
Open
OutisLi wants to merge 4 commits into
Open
refactor(pt/dpa4): use edge schema for dpa4 and ignore sel#5562OutisLi wants to merge 4 commits into
OutisLi wants to merge 4 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughReplaces the SeZM/DPA4 extended-coordinate neighbor-list input contract with a compact edge-vector schema ( ChangesSeZM/DPA4 Edge-Vector Schema Migration
Sequence Diagram(s)sequenceDiagram
participant Caller
participant DeepEval
participant NeighborListBuilder
participant SeZMModel
participant core_compute
rect rgba(70, 130, 180, 0.5)
Note over DeepEval,core_compute: Edge-schema path (_uses_edge_schema=True)
Caller->>DeepEval: eval(coord, atype, box)
DeepEval->>NeighborListBuilder: build(return_mode="edges")
NeighborListBuilder-->>DeepEval: EdgeNeighborList
DeepEval->>SeZMModel: forward_common_lower(coord, atype, edge_index, edge_vec, edge_scatter_index, edge_mask)
SeZMModel->>core_compute: edge_vec.detach().requires_grad_(True)
core_compute-->>SeZMModel: energy, forces via edge_scatter_index
SeZMModel-->>DeepEval: predict dict
DeepEval-->>Caller: energy, forces, virial
end
rect rgba(180, 70, 70, 0.5)
Note over DeepEval,core_compute: Extended path (_uses_edge_schema=False)
Caller->>DeepEval: eval(coord, atype, box)
DeepEval->>NeighborListBuilder: build(return_mode="extended")
NeighborListBuilder-->>DeepEval: ext_coord, ext_atype, nlist, mapping
DeepEval->>SeZMModel: forward_common_lower(ext_coord, ext_atype, nlist, mapping)
SeZMModel-->>DeepEval: predict dict (extended)
DeepEval->>DeepEval: communicate_extended_output
DeepEval-->>Caller: energy, forces, virial
end
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![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: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
deepmd/pt_expt/infer/deep_eval.py (1)
1266-1357:⚠️ Potential issue | 🔴 CriticalFix coordinate tensor asymmetry in fallback path: use local coordinates, not extended.
The fallback path (lines 1347–1357) passes
ext_coord_t(extended coordinates) while slicingatypeto local atoms only, creating a shape mismatch. The model reshapescoordbased onatype.shape[1], which expects the two tensors to have compatible atom dimensions. The fast path correctly uses local coordinates; the fallback path should do the same. Replaceext_coord_twith a local coordinate slice in the fallback model_inputs tuple.🤖 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_expt/infer/deep_eval.py` around lines 1266 - 1357, In the fallback path where model_inputs is constructed (around lines 1350-1357 in the _build_edge_inputs_from_nlist fallback), there is an asymmetry where ext_coord_t (extended coordinates) is passed while ext_atype_t is sliced to local atoms only with [:, :natoms], causing a shape mismatch. To fix this, slice ext_coord_t to local atoms in the same way as the atype is being sliced so that both tensors have compatible atom dimensions, matching the pattern used in the fast path where edge_schema.coord (local coordinates) is used. Replace the ext_coord_t reference in the model_inputs tuple with the appropriately sliced coordinate tensor.deepmd/pt/utils/nlist.py (1)
132-149:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject uncapped mode when type splitting is enabled.
cap_neighbors=Falsesizesnselfrom all neighbors, butdistinguish_types=Truestill routes throughnlist_distinguish_types(..., sel), which slices each type back toseland can silently drop neighbors despite the “keep every neighbor” contract.Proposed guard
if isinstance(sel, int): sel = [sel] + if not cap_neighbors and distinguish_types: + raise ValueError( + "cap_neighbors=False is only supported when distinguish_types=False" + ) # Get the coordinates for the local atoms (first nloc atoms)🤖 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/utils/nlist.py` around lines 132 - 149, The issue is that when cap_neighbors is False, the code intends to keep every neighbor within rcut, but when distinguish_types is True, the downstream call to nlist_distinguish_types with the sel parameter still slices neighbors back to fixed sel values, silently dropping neighbors and violating the "keep every neighbor" contract. Add a guard at the beginning of the function (before the cap_neighbors conditional block) that raises an error or rejects the invalid combination when both cap_neighbors is False and distinguish_types is True, preventing this silent behavior mismatch.deepmd/pt/model/model/sezm_model.py (1)
2279-2294:⚠️ Potential issue | 🟠 Major | ⚡ Quick winPreserve the spin model’s extended-list contract after this return-type change.
build_neighbor_list()now returnsEdgeNeighborList, butSeZMSpinModel.build_neighbor_list()still delegates tosuper().build_neighbor_list()while advertising the historical tuple contract. Update that override to callbuild_extended_neighbor_list()or change its callers/annotation; otherwise spin callers that unpack(extended_coord, extended_atype, nlist, mapping)will break.Suggested downstream fix
# deepmd/pt/model/model/sezm_spin_model.py def build_neighbor_list( self, coord: torch.Tensor, atype: torch.Tensor, box: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Build the real-atom neighbor list before spin expansion.""" - return super().build_neighbor_list(coord, atype, box) + return super().build_extended_neighbor_list(coord, atype, box)🤖 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/sezm_model.py` around lines 2279 - 2294, The base class build_neighbor_list() method now returns EdgeNeighborList, but SeZMSpinModel.build_neighbor_list() override still delegates to super().build_neighbor_list() while maintaining a tuple return type contract that downstream callers depend on. Update the SeZMSpinModel.build_neighbor_list() override to call build_extended_neighbor_list() instead of super().build_neighbor_list() to preserve the historical tuple unpacking contract (extended_coord, extended_atype, nlist, mapping) that spin model callers expect.
🧹 Nitpick comments (1)
deepmd/pt_expt/infer/deep_eval.py (1)
1370-1414: 💤 Low valueConsider extracting shared fparam/aparam preparation logic.
This method duplicates the fparam/aparam preparation logic from
_prepare_nlist_inputs(lines 1201-1240). While not a bug, extracting this into a shared helper would reduce maintenance burden.🤖 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_expt/infer/deep_eval.py` around lines 1370 - 1414, The _prepare_optional_lower_inputs method contains fparam and aparam preparation logic that is duplicated from _prepare_nlist_inputs. Extract the shared tensor preparation logic for both fparam and aparam (including default value handling, validation, and tensor creation with reshape and expand operations) into a separate helper method. Then replace the duplicated code blocks in both _prepare_optional_lower_inputs and _prepare_nlist_inputs with calls to this new helper method to reduce maintenance burden and improve code reusability.
🤖 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/pt/infer/deep_eval.py`:
- Around line 257-259: The _uses_edge_schema switch in the assignment is
currently based on self.model_def_script, which remains the top-level model
dictionary in multitask scenarios. However, self.input_param is replaced with
the selected branch params when a specific head is instantiated. Change the
_is_sezm_model_params call to use self.input_param instead of
self.model_def_script so that the edge-schema switch correctly reflects the
selected head's parameters rather than the top-level model definition, ensuring
the correct forward_common_lower argument contract is chosen for the O(N) path.
In `@deepmd/pt/model/model/sezm_model.py`:
- Around line 1524-1531: Ruff F722 is failing because it's attempting to parse
the string literals in the jaxtyping shape annotations (such as those in the
forward_lower method parameters at lines 1524-1531 and
build_extended_neighbor_list method at lines 2298-2300) as forward annotations,
which contain invalid Python syntax with embedded spaces and underscores. Add a
file-level `# noqa: F722` directive at the top of the sezm_model.py file to
suppress this linting rule for the entire file, ensuring that ruff check passes
without modifying the jaxtyping annotations themselves.
In `@source/api_cc/include/commonPT.h`:
- Around line 202-225: The createEdgeTensors function incorrectly assumes that
loop index ii directly corresponds to the center atom index, but when nlist is
compacted via shuffle_exclude_empty, row ii actually corresponds to ilist[ii].
This causes incorrect scatter indices and edge vector calculations. Pass ilist
as a parameter to createEdgeTensors and update the center atom index calculation
to use const int center = ilist[ii] instead of ii, then use this center variable
in the center_offset calculation at line 203 and in the dst_ext.push_back call
at line 225.
In `@source/api_cc/src/DeepPotPT.cc`:
- Around line 220-224: The filter_by_distance call on nlist_data needs to be
executed on every timestep, not just when ago equals zero (during neighbor list
rebuilds). Move the filter_by_distance call that filters pairs by the rcut
cutoff distance outside of the ago == 0 conditional block, keeping the other
nlist operations (copy_from_nlist, shuffle_exclude_empty, and padding) inside
the conditional block where they belong. This ensures pairs are correctly
filtered based on the cutoff distance at every step, preventing missed pairs
that re-enter the cutoff radius between rebuilds.
---
Outside diff comments:
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 1266-1357: In the fallback path where model_inputs is constructed
(around lines 1350-1357 in the _build_edge_inputs_from_nlist fallback), there is
an asymmetry where ext_coord_t (extended coordinates) is passed while
ext_atype_t is sliced to local atoms only with [:, :natoms], causing a shape
mismatch. To fix this, slice ext_coord_t to local atoms in the same way as the
atype is being sliced so that both tensors have compatible atom dimensions,
matching the pattern used in the fast path where edge_schema.coord (local
coordinates) is used. Replace the ext_coord_t reference in the model_inputs
tuple with the appropriately sliced coordinate tensor.
In `@deepmd/pt/model/model/sezm_model.py`:
- Around line 2279-2294: The base class build_neighbor_list() method now returns
EdgeNeighborList, but SeZMSpinModel.build_neighbor_list() override still
delegates to super().build_neighbor_list() while maintaining a tuple return type
contract that downstream callers depend on. Update the
SeZMSpinModel.build_neighbor_list() override to call
build_extended_neighbor_list() instead of super().build_neighbor_list() to
preserve the historical tuple unpacking contract (extended_coord,
extended_atype, nlist, mapping) that spin model callers expect.
In `@deepmd/pt/utils/nlist.py`:
- Around line 132-149: The issue is that when cap_neighbors is False, the code
intends to keep every neighbor within rcut, but when distinguish_types is True,
the downstream call to nlist_distinguish_types with the sel parameter still
slices neighbors back to fixed sel values, silently dropping neighbors and
violating the "keep every neighbor" contract. Add a guard at the beginning of
the function (before the cap_neighbors conditional block) that raises an error
or rejects the invalid combination when both cap_neighbors is False and
distinguish_types is True, preventing this silent behavior mismatch.
---
Nitpick comments:
In `@deepmd/pt_expt/infer/deep_eval.py`:
- Around line 1370-1414: The _prepare_optional_lower_inputs method contains
fparam and aparam preparation logic that is duplicated from
_prepare_nlist_inputs. Extract the shared tensor preparation logic for both
fparam and aparam (including default value handling, validation, and tensor
creation with reshape and expand operations) into a separate helper method. Then
replace the duplicated code blocks in both _prepare_optional_lower_inputs and
_prepare_nlist_inputs with calls to this new helper method to reduce maintenance
burden and improve code reusability.
🪄 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: 64f127a8-03e5-4cdf-870f-77ed7c868ae1
📒 Files selected for processing (30)
deepmd/dpmodel/utils/default_neighbor_list.pydeepmd/dpmodel/utils/neighbor_list.pydeepmd/pt/entrypoints/freeze_pt2.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/model/sezm_spin_model.pydeepmd/pt/utils/compile_compat.pydeepmd/pt/utils/nlist.pydeepmd/pt/utils/nv_nlist.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/utils/edge_schema.pydeepmd/pt_expt/utils/vesin_neighbor_list.pydeepmd/utils/argcheck.pyexamples/water/dpa4/input.jsonsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/common.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPT.ccsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/src/DeepSpinPT.ccsource/api_cc/src/DeepSpinPTExpt.ccsource/api_cc/src/DeepTensorPT.ccsource/api_cc/src/common.ccsource/api_cc/tests/test_deeppot_pt.ccsource/api_cc/tests/test_deeppot_ptexpt.ccsource/api_cc/tests/test_neighbor_list_data.ccsource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_nv_nlist.pysource/tests/pt/model/test_sezm_export.pysource/tests/pt/model/test_sezm_model.py
💤 Files with no reviewable changes (1)
- examples/water/dpa4/input.json
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5562 +/- ##
==========================================
- Coverage 82.17% 82.15% -0.03%
==========================================
Files 898 899 +1
Lines 103576 104128 +552
Branches 4432 4470 +38
==========================================
+ Hits 85117 85544 +427
- Misses 17063 17176 +113
- Partials 1396 1408 +12 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt/model/test_nv_nlist.py (1)
286-301: ⚡ Quick winAssert scatter indices and low-
selinvariance.The new consistency check masks down to
(src, dst, edge_vec)and every case passessel=[64]. A builder could still scatter forces/virials to the wrong atoms, or cap edge-mode output whenselis small, while these tests pass. Includeedge_scatter_indexin the canonical comparison and repeat the cases with an intentionally smallsel.Proposed test strengthening
-def _canonical_edges(schema) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +def _canonical_edges( + schema, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: @@ edge_index = _to_numpy(schema.edge_index) src = edge_index[0][mask].astype(np.int64) dst = edge_index[1][mask].astype(np.int64) + edge_scatter = _to_numpy(schema.edge_scatter_index) + scatter_src = edge_scatter[0][mask].astype(np.int64) + scatter_dst = edge_scatter[1][mask].astype(np.int64) edge_vec = _to_numpy(schema.edge_vec)[mask].astype(np.float64) @@ - return src[order], dst[order], edge_vec[order] + return ( + src[order], + dst[order], + scatter_src[order], + scatter_dst[order], + edge_vec[order], + ) @@ - ref_src, ref_dst, ref_vec = _canonical_edges( + ref_src, ref_dst, ref_scatter_src, ref_scatter_dst, ref_vec = _canonical_edges( _dense_keepall_edges(coord, atype, box, rcut) ) @@ - src, dst, vec = _canonical_edges( + src, dst, scatter_src, scatter_dst, vec = _canonical_edges( builder.build(coord, atype, box, rcut, sel, return_mode="edges") ) np.testing.assert_array_equal(src, ref_src) np.testing.assert_array_equal(dst, ref_dst) + np.testing.assert_array_equal(scatter_src, ref_scatter_src) + np.testing.assert_array_equal(scatter_dst, ref_scatter_dst) np.testing.assert_allclose(vec, ref_vec, atol=1.0e-9, rtol=1.0e-9) @@ - self._assert_builders_match(coord, atype, box, 4.0, [64]) + for sel in ([64], [1]): + with self.subTest(sel=sel): + self._assert_builders_match(coord, atype, box, 4.0, sel) @@ - self._assert_builders_match(coord, atype, box, 4.0, [64]) + for sel in ([64], [1]): + with self.subTest(sel=sel): + self._assert_builders_match(coord, atype, box, 4.0, sel) @@ - self._assert_builders_match(coord, atype, None, 4.0, [64]) + for sel in ([64], [1]): + with self.subTest(sel=sel): + self._assert_builders_match(coord, atype, None, 4.0, sel)Also applies to: 356-390
🤖 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 `@source/tests/pt/model/test_nv_nlist.py` around lines 286 - 301, The _canonical_edges function only returns (src, dst, edge_vec) which doesn't catch potential bugs in scatter indices or edge-mode capping when sel is small. Modify the _canonical_edges function to also extract edge_scatter_index from the schema (similar to how edge_index and edge_vec are extracted using _to_numpy), apply the mask, and include it as a fourth return value in the tuple. Then update all test cases that call this function to unpack and validate the returned scatter index alongside the other edge properties. Finally, repeat the test cases mentioned in the comment (around lines 356-390) with an intentionally small sel parameter value (such as sel=2 or sel=1) to ensure proper edge-mode output behavior is tested.
🤖 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/pt/model/model/sezm_model.py`:
- Around line 1021-1022: The extended_coord_corr tensor is being reshaped but
not normalized to match the device and dtype of other geometry tensors like
coord. After reshaping extended_coord_corr when its ndim is 2, add code to move
it to coord.device and coord.dtype using the same normalization approach applied
to coord and the edge tensors. This ensures consistent device placement and data
types throughout the geometry tensor operations and prevents device/dtype
mismatches when reaching edge_energy_deriv().
---
Nitpick comments:
In `@source/tests/pt/model/test_nv_nlist.py`:
- Around line 286-301: The _canonical_edges function only returns (src, dst,
edge_vec) which doesn't catch potential bugs in scatter indices or edge-mode
capping when sel is small. Modify the _canonical_edges function to also extract
edge_scatter_index from the schema (similar to how edge_index and edge_vec are
extracted using _to_numpy), apply the mask, and include it as a fourth return
value in the tuple. Then update all test cases that call this function to unpack
and validate the returned scatter index alongside the other edge properties.
Finally, repeat the test cases mentioned in the comment (around lines 356-390)
with an intentionally small sel parameter value (such as sel=2 or sel=1) to
ensure proper edge-mode output behavior is tested.
🪄 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: 0a44c6b0-a928-4af2-87a2-0038bc305faa
📒 Files selected for processing (30)
deepmd/dpmodel/utils/default_neighbor_list.pydeepmd/dpmodel/utils/neighbor_list.pydeepmd/pt/entrypoints/freeze_pt2.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/model/sezm_spin_model.pydeepmd/pt/utils/compile_compat.pydeepmd/pt/utils/nlist.pydeepmd/pt/utils/nv_nlist.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/utils/edge_schema.pydeepmd/pt_expt/utils/vesin_neighbor_list.pydeepmd/utils/argcheck.pyexamples/water/dpa4/input.jsonsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/common.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPT.ccsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/src/DeepSpinPT.ccsource/api_cc/src/DeepSpinPTExpt.ccsource/api_cc/src/DeepTensorPT.ccsource/api_cc/src/common.ccsource/api_cc/tests/test_deeppot_pt.ccsource/api_cc/tests/test_deeppot_ptexpt.ccsource/api_cc/tests/test_neighbor_list_data.ccsource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_nv_nlist.pysource/tests/pt/model/test_sezm_export.pysource/tests/pt/model/test_sezm_model.py
✅ Files skipped from review due to trivial changes (1)
- examples/water/dpa4/input.json
🚧 Files skipped from review as they are similar to previous changes (25)
- source/api_cc/src/DeepPotPT.cc
- source/api_cc/include/common.h
- source/api_cc/src/DeepTensorPT.cc
- source/api_cc/src/DeepSpinPTExpt.cc
- source/api_cc/src/DeepSpinPT.cc
- source/api_cc/include/DeepPotPTExpt.h
- deepmd/dpmodel/utils/default_neighbor_list.py
- deepmd/dpmodel/utils/neighbor_list.py
- source/tests/pt/model/test_sezm_model.py
- source/api_cc/tests/test_deeppot_pt.cc
- deepmd/utils/argcheck.py
- source/api_cc/tests/test_deeppot_ptexpt.cc
- deepmd/pt_expt/utils/vesin_neighbor_list.py
- source/api_cc/include/commonPT.h
- deepmd/pt/utils/nv_nlist.py
- source/api_cc/src/common.cc
- deepmd/pt/utils/nlist.py
- source/api_cc/tests/test_neighbor_list_data.cc
- source/tests/pt/model/test_descriptor_sezm.py
- deepmd/pt/entrypoints/freeze_pt2.py
- source/api_cc/src/DeepPotPTExpt.cc
- deepmd/pt/model/model/sezm_spin_model.py
- source/tests/pt/model/test_sezm_export.py
- deepmd/pt_expt/utils/edge_schema.py
- deepmd/pt_expt/infer/deep_eval.py
feat(pt/ptexpt): Filter LAMMPS neighbor lists by model cutoff
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/pt/model/test_nv_nlist.py (1)
286-301: ⚡ Quick winCover scatter indices and undersized
selin edge-schema consistency tests.These tests can pass even if a builder corrupts
edge_scatter_indexor still truncates edge-mode output bysel, because_canonical_edges()drops scatter indices and every case uses[64]. Compareedge_scatter_indextoo, and include a deliberately smallselsuch as[1]to enforce the “ignore sel” contract.Proposed test strengthening
-def _canonical_edges(schema) -> tuple[np.ndarray, np.ndarray, np.ndarray]: +def _canonical_edges( + schema, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: @@ edge_index = _to_numpy(schema.edge_index) src = edge_index[0][mask].astype(np.int64) dst = edge_index[1][mask].astype(np.int64) + scatter = _to_numpy(schema.edge_scatter_index)[:, mask].astype(np.int64) edge_vec = _to_numpy(schema.edge_vec)[mask].astype(np.float64) keys = np.round(edge_vec * 1.0e6).astype(np.int64) order = np.lexsort((keys[:, 2], keys[:, 1], keys[:, 0], dst, src)) - return src[order], dst[order], edge_vec[order] + return src[order], dst[order], scatter[:, order], edge_vec[order] @@ - ref_src, ref_dst, ref_vec = _canonical_edges( + ref_src, ref_dst, ref_scatter, ref_vec = _canonical_edges( _dense_keepall_edges(coord, atype, box, rcut) ) @@ - src, dst, vec = _canonical_edges( + src, dst, scatter, vec = _canonical_edges( builder.build(coord, atype, box, rcut, sel, return_mode="edges") ) np.testing.assert_array_equal(src, ref_src) np.testing.assert_array_equal(dst, ref_dst) + np.testing.assert_array_equal(scatter, ref_scatter) np.testing.assert_allclose(vec, ref_vec, atol=1.0e-9, rtol=1.0e-9) @@ - self._assert_builders_match(coord, atype, box, 4.0, [64]) + for sel in ([1], [64]): + with self.subTest(device=str(device), sel=sel): + self._assert_builders_match(coord, atype, box, 4.0, sel)Also applies to: 356-372, 374-390
🤖 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 `@source/tests/pt/model/test_nv_nlist.py` around lines 286 - 301, The edge-schema consistency tests are not comprehensive enough to catch potential builder corruptions. The `_canonical_edges()` function currently drops scatter indices and all test cases use the same large `sel` value of `[64]`, which means the tests cannot detect if `edge_scatter_index` is corrupted or if edge-mode output is incorrectly truncated by the `sel` parameter. Enhance the consistency validation by adding checks for `edge_scatter_index` equality in the test assertions, and include test cases with deliberately small `sel` values such as `[1]` alongside the existing `[64]` cases to enforce the contract that edges should be returned regardless of the `sel` parameter. Apply these changes to all affected test functions mentioned in the comment range.
🤖 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.
Nitpick comments:
In `@source/tests/pt/model/test_nv_nlist.py`:
- Around line 286-301: The edge-schema consistency tests are not comprehensive
enough to catch potential builder corruptions. The `_canonical_edges()` function
currently drops scatter indices and all test cases use the same large `sel`
value of `[64]`, which means the tests cannot detect if `edge_scatter_index` is
corrupted or if edge-mode output is incorrectly truncated by the `sel`
parameter. Enhance the consistency validation by adding checks for
`edge_scatter_index` equality in the test assertions, and include test cases
with deliberately small `sel` values such as `[1]` alongside the existing `[64]`
cases to enforce the contract that edges should be returned regardless of the
`sel` parameter. Apply these changes to all affected test functions mentioned in
the comment range.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d67a780c-d5da-4f1b-bc3d-ef6e0fcc913d
📒 Files selected for processing (24)
deepmd/dpmodel/utils/default_neighbor_list.pydeepmd/dpmodel/utils/neighbor_list.pydeepmd/pt/entrypoints/freeze_pt2.pydeepmd/pt/infer/deep_eval.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/model/sezm_spin_model.pydeepmd/pt/utils/compile_compat.pydeepmd/pt/utils/nlist.pydeepmd/pt/utils/nv_nlist.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/utils/edge_schema.pydeepmd/pt_expt/utils/vesin_neighbor_list.pydeepmd/utils/argcheck.pyexamples/water/dpa4/input.jsonsource/api_cc/include/DeepPotPTExpt.hsource/api_cc/include/commonPT.hsource/api_cc/src/DeepPotPTExpt.ccsource/api_cc/tests/test_deeppot_pt.ccsource/api_cc/tests/test_deeppot_ptexpt.ccsource/api_cc/tests/test_neighbor_list_data.ccsource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_nv_nlist.pysource/tests/pt/model/test_sezm_export.pysource/tests/pt/model/test_sezm_model.py
💤 Files with no reviewable changes (1)
- examples/water/dpa4/input.json
🚧 Files skipped from review as they are similar to previous changes (19)
- deepmd/dpmodel/utils/default_neighbor_list.py
- source/tests/pt/model/test_sezm_model.py
- source/api_cc/include/DeepPotPTExpt.h
- deepmd/dpmodel/utils/neighbor_list.py
- source/api_cc/include/commonPT.h
- source/api_cc/tests/test_deeppot_pt.cc
- deepmd/pt/infer/deep_eval.py
- source/api_cc/tests/test_deeppot_ptexpt.cc
- deepmd/pt_expt/utils/vesin_neighbor_list.py
- deepmd/pt/utils/nlist.py
- deepmd/utils/argcheck.py
- source/tests/pt/model/test_descriptor_sezm.py
- deepmd/pt/entrypoints/freeze_pt2.py
- source/tests/pt/model/test_sezm_export.py
- deepmd/pt/utils/nv_nlist.py
- deepmd/pt/model/model/sezm_spin_model.py
- deepmd/pt_expt/infer/deep_eval.py
- source/api_cc/src/DeepPotPTExpt.cc
- deepmd/pt_expt/utils/edge_schema.py
wanghan-iapcm
approved these changes
Jun 20, 2026
Contributor
|
I reviewed the current head of this PR, focusing on the new edge-schema path for DPA4/SeZM and the pt_expt/C++ inference plumbing. The parts I was most worried about look OK in the current revision:
I also checked that the GitHub CI matrix is green. Locally I ran syntax checks for the touched Python files and Authored by OpenClaw 2026.6.8 (844f405) (model: custom-chat-jinzhezeng-group/gpt-5.5) |
njzjz
approved these changes
Jun 20, 2026
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
5 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.
Summary by CodeRabbit
New Features
Bug Fixes
Improvements