fix(dpmodel): preserve padding when forced-sorting short neighbor lists

[njzjz-bot]

Fix #5629

Problem

format_nlist pads the input nlist to width nnei when the input is shorter, but the forced-sort branch (extra_nlist_sort=True) then re-read the shape from the original unpadded nlist instead of the padded ret. This caused:

• The sort to operate on the short tensor (width = original n_nnei, not nnei)
• The final ret[..., :nnei] slice to return a tensor with the original short width
• The -1 padding to be lost
• The assert ret.shape[-1] == nnei check to fail for static shapes

This is especially reachable for linear atomic models, which always request sorted lower neighbor lists via need_sorted_nlist_for_lower().

Fix

Use the padded ret tensor (instead of the original nlist) throughout the forced-sort branch. Added explanatory comments.

Test

Added source/tests/common/dpmodel/test_format_nlist_short_padding.py with three test cases:

1. Short nlist with sort: input width 2, target nnei=4, verifies output width is 4 and real neighbors are preserved.
2. All-padded input with sort: verifies output is all -1 with width nnei.
3. Short nlist without sort: sanity check that the non-sort path also works.

Attribution

Generated with opencode using model glm-5.2.

Recommended reviewers

@njzjz (maintainer of the dpmodel code)

Summary by CodeRabbit

• Improvements

  • Improved neighbor-list generation across supported array backends.
  • Enhanced compatibility with tensor devices, data types, and static-shape execution.
  • Neighbor lists now consistently produce the requested number of entries, using clear padding when needed.
  • Improved sorting, distance filtering, truncation, and handling of empty or fully padded neighbor lists.
  • Ghost-cell coordinate generation now provides more consistent results across execution environments.

• Tests

  • Added coverage for sorted neighbor lists, including partially populated and fully padded inputs.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 5a6c060c-5a41-4d41-8c71-1d632a351242

📥 Commits

Reviewing files that changed from the base of the PR and between eb52ebb and 8b30d37.

📒 Files selected for processing (2)

• deepmd/dpmodel/utils/nlist.py
• source/tests/common/dpmodel/test_nlist.py

🚧 Files skipped from review as they are similar to previous changes (1)

• deepmd/dpmodel/utils/nlist.py

---

📝 Walkthrough

Walkthrough

Neighbor-list utilities now support backend-specific sizing and ranges, validate static shapes, and normalize neighbor lists to fixed widths while preserving padding during optional distance sorting. Tests cover partially padded and fully padded forced-sort inputs.

Changes

Neighbor-list utilities

Layer / File(s)	Summary
Backend-compatible neighbor-list construction deepmd/dpmodel/utils/nlist.py	Backend detection, size calculation, static-shape checks, integer range generation, ndtensorflow handling, concatenation, and ghost-cell range construction are updated for Array API compatibility.
Fixed-width neighbor-list formatting deepmd/dpmodel/utils/nlist.py, source/tests/common/dpmodel/test_nlist.py	format_nlist pads short inputs with -1, preserves padding during forced sorting and truncation, masks invalid neighbors, and adds regression tests for partially and fully padded inputs.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• deepmodeling/deepmd-kit#5524: Both changes address neighbor-list ordering and sorting for lower-interface processing.
• deepmodeling/deepmd-kit#5529: Both update neighbor-list padding and cutoff masking behavior.

Suggested reviewers: OutisLi, njzjz

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Out of Scope Changes check	⚠️ Warning	The PR also changes build_neighbor_list and extend_coord_with_ghosts for Array API and static-shape handling, which are outside #5629's scope.	Move those compatibility refactors to a separate PR or explicitly tie them to this issue's acceptance criteria.
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and accurately summarizes the main fix in the PR.
Linked Issues check	✅ Passed	The patch fixes format_nlist to sort the padded tensor and keep -1 padding and width nnei as required by #5629, with regression tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[njzjz-bot]

The core fix is correct: using the padded ret throughout the forced-sort branch preserves the requested width and matches the corresponding PyTorch behavior. However, this PR is currently not mergeable: its parent is dd428dbb, it is 124 commits behind the current base, GitHub reports CONFLICTING, and pre-commit.ci stopped at the mergeability check. The stale base also makes the diff include many already-merged nlist.py changes. Please rebase onto current master, resolve the conflict, and retain the minimal three-source-line change plus the regression test before merging.

As a non-blocking test improvement after rebasing, an exact array assertion would verify that every remaining slot is -1; the current set-based check would also accept duplicate real neighbors.

Coding agent: Codex
Codex version: codex-cli 0.144.0-alpha.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.52%. Comparing base (7c362d7) to head (8b30d37).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5764      +/-   ##
==========================================
- Coverage   79.65%   79.52%   -0.13%     
==========================================
  Files        1015     1015              
  Lines      115781   115781              
  Branches     4272     4272              
==========================================
- Hits        92229    92079     -150     
- Misses      22008    22157     +149     
- Partials     1544     1545       +1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[wanghan-iapcm]

Non-blocking (test organization): consider folding these cases into the existing source/tests/common/dpmodel/test_nlist.py::TestNList, which already targets format_nlist with test_nlist_eq / test_nlist_st / test_nlist_lt and a shared fixture (self.md, self.coord_ext, self.atype_ext, self.expected_nlist), rather than a separate file that rebuilds its own nlist.

Root cause of why the existing UTs missed this: test_nlist_st already covers the short list (n_nnei < nnei), but every existing test calls format_nlist with the default extra_nlist_sort=False, so the buggy if extra_nlist_sort or n_nnei > nnei: branch was never crossed with a short input (test_nlist_lt hits the sort branch but with ret == nlist, so the padding bug is invisible). The untested cell is the intersection extra_nlist_sort=True × n_nnei < nnei. Adding extra_nlist_sort=True variants (e.g. test_nlist_st_sort) to the {eq, st, lt} set would make that matrix explicit and keep all format_nlist coverage in one place.