fix(dpmodel): guard empty magnetic loss masks#5798
Conversation
Use Array API-safe nonzero denominators for magnetic-force loss reductions so an all-empty mask contributes finite zero instead of NaN. Add backend-neutral NumPy and Torch namespace regressions across loss functions, precisions, and label-presence flags. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Caution Review failedFailed to post review comments. GitHub was unavailable or timed out while CodeRabbit was posting the review. Please request a new review later if the pull request still needs one. Use ⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (5)
⏰ Context from checks skipped due to timeout. (2)
🧰 Additional context used📓 Path-based instructions (2)**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
**/tests/**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
🧠 Learnings (10)📚 Learning: 2026-04-11T08:01:08.364ZApplied to files:
📚 Learning: 2026-05-01T13:46:12.190ZApplied to files:
📚 Learning: 2026-05-02T12:55:24.164ZApplied to files:
📚 Learning: 2026-05-02T12:55:39.794ZApplied to files:
📚 Learning: 2026-05-28T23:48:28.388ZApplied to files:
📚 Learning: 2026-06-11T16:43:05.483ZApplied to files:
📚 Learning: 2026-06-11T16:43:04.825ZApplied to files:
📚 Learning: 2026-07-07T08:20:09.299ZApplied to files:
📚 Learning: 2026-06-13T01:31:17.863ZApplied to files:
📚 Learning: 2026-01-22T02:59:02.650ZApplied to files:
🪛 Ruff (0.15.21)source/tests/consistent/loss/test_ener_spin.py[warning] 272-272: Function definition does not bind loop variable (B023) 📝 WalkthroughWalkthrough
ChangesEmpty magnetic-mask handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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 |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5798 +/- ##
==========================================
- Coverage 79.85% 78.31% -1.54%
==========================================
Files 1022 1050 +28
Lines 117351 120601 +3250
Branches 4313 4354 +41
==========================================
+ Hits 93706 94453 +747
- Misses 22101 24586 +2485
- Partials 1544 1562 +18 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Reviewed via the code-review skill. Correct and behavior-preserving: for a non-empty magnetic mask safe_n_valid==n_valid (loss unchanged), and for an empty mask the numerator is already 0 so the contribution is a finite 0; guarding the denominator (not wrapping the result in where) is the right approach for JAX, which evaluates both where branches. All three magnetic-force reductions are covered. The regressions genuinely fail pre-fix (verified: all 8 numpy and 4 torch cases fail on isfinite, since 0*NaN=NaN regardless of find_force_mag), across mse/mae, float32/64, numpy+torch, with correct display_if_exist semantics (0 for a present-but-empty label, NaN for an absent label).
One non-blocking cross-backend note: the separate torch.jit backend deepmd/pt/loss/ener_spin.py already guards the accumulated loss via torch.nan_to_num, but its displayed more_loss['rmse_fm']/['mae_fm'] are NOT guarded and still evaluate to NaN for an empty magnetic mask. After this merges, dpmodel/pt_expt will report 0.0 for those display metrics while pt reports NaN for the same logical case, and source/tests/consistent/loss/test_ener_spin.py never exercises n_magnetic=0, so the divergence is untested. Worth aligning pt's display metrics (and/or adding an n_magnetic=0 consistency case) in a follow-up. Separately -- and clearly out of scope here -- the same NaN class (a mask-count of 0 used as a divisor) is still open in the energy/virial 1/real_natoms reciprocals across ener.py/ener_spin.py/pt equivalents for all-ghost frames; that belongs to the broader loss-padding-mask work. LGTM for the magnetic-mask fix.
Normalize empty magnetic-force reductions before publishing PyTorch RMSE and MAE metrics, and add a cross-backend consistency regression against dpmodel. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
I addressed the non-blocking PyTorch consistency note in The torch.jit loss now normalizes empty magnetic-force reductions before publishing The focused consistency test passed, including all three subcases; Ruff and formatting hooks also passed. Coding agent: Codex |
Summary
Root cause and fix
EnergySpinLossmasks non-magnetic atoms out of the force-magnitude residual, then divides its reductions byn_valid * 3. When a batch contains no magnetic atoms, both numerator and denominator are zero, producing NaN. Multiplying that value by a zero prefactor or a missing-label flag does not recover it because0 * NaNis still NaN.The implementation now constructs:
and uses it for every global magnetic-force MSE/MAE reduction. The masked numerator is already zero for an empty set, so the contribution becomes exactly zero. Guarding the denominator itself is important: array backends may evaluate both branches of a later
where, and an unselected divide-by-zero can still produce invalid values or gradients. The expression is compatible with Array API namespaces and JAX tracing; non-empty batches retain the exact previous formula.This PR intentionally does not change the native PyTorch loss or redefine its missing-label display metrics. When
find_force_mag == 0, the total loss is finite zero whiledisplay_if_existcontinues to report NaN for the unavailable magnetic metric.Why existing tests missed this
Existing spin-loss mask tests always selected at least one magnetic atom in each batch. They covered partial masks and all-magnetic batches, but never the global zero denominator.
The new backend-neutral regression parameterizes:
mseandmae;find_force_magpresent and absent.It verifies a finite zero total loss in all eight combinations, zero magnetic metrics for a present-but-empty label, and the existing NaN display behavior for an absent label. A separate pt_expt test exercises the same shared dpmodel implementation through the Torch Array API namespace.
On the previous implementation, all eight backend-neutral cases fail with NaN. With the fix, the combined new and related targeted tests pass.
Validation
array_api_strictMSE/MAE empty-mask scenarios passedruff format --check .ruff check .git diff --checkCloses #5637.
Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh
Summary by CodeRabbit
Bug Fixes
Tests