fix(dpmodel): guard empty magnetic loss masks

[njzjz]

Summary

• replace zero magnetic-mask denominators with Array API-safe nonzero denominators
• make all-empty magnetic-force masks contribute finite zero instead of NaN for both MSE and MAE losses
• preserve the existing display semantics when the magnetic-force label itself is absent
• add backend-neutral NumPy and Torch Array API regression coverage

Root cause and fix

EnergySpinLoss masks non-magnetic atoms out of the force-magnitude residual, then divides its reductions by n_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 because 0 * NaN is still NaN.

The implementation now constructs:

safe_n_valid = xp.where(n_valid > 0, n_valid, xp.ones_like(n_valid))

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 while display_if_exist continues 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:

• mse and mae;
• float32 and float64; and
• find_force_mag present 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

• old backend-neutral regression: 8 failed with NaN as expected
• fixed NumPy and pt_expt targeted tests: 20 passed
• array_api_strict MSE/MAE empty-mask scenarios passed
• JAX was not installed locally; the denominator expression was kept free of Python data-dependent branching for JIT compatibility
• a broader consistent-loss run produced 35 passed / 45 skipped plus 2 unrelated pre-existing MAE parity failures because the shared editable native-PT install came from another worktree; neither failure exercises the empty-mask path
• ruff format --check .
• ruff check .
• git diff --check

Closes #5637.

Coding agent: Codex
Codex version: codex-cli 0.144.1
Model: gpt-5.6-sol
Reasoning effort: xhigh

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[coderabbitai]

Warning

Review limit reached

@njzjz, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 25 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 24577d96-a949-4f17-b588-6fa79892ef1d

📥 Commits

Reviewing files that changed from the base of the PR and between 3122138 and f81b12b.

📒 Files selected for processing (3)

• deepmd/dpmodel/loss/ener_spin.py
• source/tests/common/dpmodel/test_loss_ener_spin.py
• source/tests/pt_expt/loss/test_ener_spin.py

✨ 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.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.30%. Comparing base (7d5ad38) to head (f81b12b).
⚠️ Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5798      +/-   ##
==========================================
- Coverage   79.85%   78.30%   -1.55%     
==========================================
  Files        1022     1029       +7     
  Lines      117351   117576     +225     
  Branches     4313     4312       -1     
==========================================
- Hits        93706    92064    -1642     
- Misses      22101    23956    +1855     
- Partials     1544     1556      +12     

☔ 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.