refactor(loss): extract masked per-frame reduction idioms to cut nested branching

[wanghan-iapcm]

Resolves #5768. Follow-up to #5738.

What this does

The mixed_type padding fix (#5738) added a masked (maskf is not None) branch to every term of the shared losses. Combined with the existing has_* / mse|mae / use_huber conditions, that produced deeply nested, duplicated per-frame arithmetic repeated across energy, force, virial, atom_ener, atom_pref, dos, tensor, and again across ener_spin, in both the dpmodel and pt backends.

This PR extracts the three recurring masked reduction "idioms" into a single shared array_api_compat module, deepmd/dpmodel/loss/reduction.py:

• masked_atom_mean(elem, maskf, ncomp) — idiom 1: per-atom masked mean over ncomp components (force, atom_ener, atom_pref, local dos/tensor, spin real-force).
• per_frame_component_mean(err) — idiom 2 primitive: per-frame mean over components; the caller applies the extensive 1/real_natoms**norm_exp weighting (energy, virial).
• masked_atom_num(mask, natoms, dtype) — idiom 3 companion: the display-only atom-count divisor for already-reduced global quantities (global dos/tensor).

Both backends import the same module; the pt backend calls it with torch tensors, and array_api_compat's torch-namespace dispatch preserves autograd and produces bit-identical results.

Numerics are preserved

This is a pure readability/maintainability refactor with no intended behavior change. Each helper implements only the masked branch; every call site keeps its original non-masked else expression verbatim, so the "bit-identical for non-mixed batches" guarantee from #5738 holds (defaulting the mask to all-ones would change the reduction order at the ULP level and is deliberately avoided). Operand order and the mask-application point are preserved in every conversion.

Testing

The safety net is the grad-accumulation-invariant and no-op-for-non-mixed tests added in #5738 plus the cross-backend consistency loss tests, all of which stay green:

• source/tests/common/dpmodel/test_loss_padding.py, source/tests/pt/test_loss_padding.py
• source/tests/consistent/loss (dpmodel <-> pt <-> tf): 244 passed, 176 skipped — unchanged, the load-bearing bit-identity guard
• new direct unit tests for the three helpers (numpy + torch, autograd, bit-identity vs the torch-native inline form): source/tests/common/dpmodel/test_loss_reduction.py

Known limitations

• The shared module lives under deepmd/dpmodel/loss and is imported by deepmd/pt/loss — the first such cross-backend import in the loss layer. This is intentional (single source of truth for the reduction math, consistent with deepmd/dpmodel being the designated shared-math layer), but it does couple the pt loss backend to array_api_compat.
• The Huber-masked sub-branches, the f_use_norm force variant, the generalized-force (gf) input masking, and the spin magnetic-force (mask_mag) block are intentionally NOT routed through the helpers — they are not one of the three idioms.
• The TensorFlow (deepmd/tf) loss backend is out of scope; it does not use the mixed_type padding mask.

Summary by CodeRabbit

• Improvements

  • Improved masked loss reductions across energy, forces, virial, tensor, DOS, and spin training by switching to shared masked/per-frame reduction helpers.
  • Standardized masked atom normalization and per-frame aggregation for consistent RMSE/MAE reporting in local and global DOS/tensor computations.
  • Preserved behavior and control flow across NumPy/JAX and PyTorch execution paths.

• Tests

  • Added coverage for masked mean, per-frame component mean, and atom-count normalization (including None masks), plus PyTorch autograd checks and no-NaN behavior on all-padding frames.

[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: 90538b6a-098e-4adc-80c7-0570f1af71ab

📥 Commits

Reviewing files that changed from the base of the PR and between bc7fc31 and 97c26f5.

📒 Files selected for processing (1)

• source/tests/common/dpmodel/test_loss_reduction.py

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

• source/tests/common/dpmodel/test_loss_reduction.py

---

📝 Walkthrough

Walkthrough

The change adds shared masked-reduction helpers, tests their NumPy and PyTorch behavior, and applies them across dpmodel and PyTorch DOS, energy, spin-energy, and tensor losses.

Changes

Masked loss reduction refactor

Layer / File(s)	Summary
Reduction helpers and validation deepmd/dpmodel/loss/reduction.py, source/tests/common/dpmodel/test_loss_reduction.py	Adds shared masked atom, per-frame component, and masked atom-count reductions with backend-aware implementations and correctness, autograd, and all-padding tests.
dpmodel loss integration deepmd/dpmodel/loss/{dos,ener,ener_spin,tensor}.py	Replaces duplicated masked reshaping, summation, and normalization logic with shared reduction helpers.
PyTorch loss integration deepmd/pt/loss/{dos,ener,ener_spin,tensor}.py	Applies the shared reductions to masked DOS, energy, spin-energy, tensor, force, virial, atomic-energy, and atom-prefactor calculations.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related issues

• TF backend loss: exclude mixed_type padding atoms (per-frame normalization) #5760 — Covers the corresponding shared reduction helpers and their application across loss backends.
• fix(loss): normalize global RMSE display metrics per frame for mixed-size batches #5767 — Relates to mask-aware atom-count normalization in DOS and tensor loss paths.

Possibly related PRs

• deepmodeling/deepmd-kit#5738 — Introduced the masked per-frame normalization behavior consolidated by these helpers.
• deepmodeling/deepmd-kit#5345 — Added related loss implementations now being refactored.
• deepmodeling/deepmd-kit#5325 — Modifies overlapping energy-loss computation paths.

Suggested reviewers: iprozd, njzjz

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.77% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately summarizes the main refactor: extracting shared masked reduction helpers to simplify loss code.
Linked Issues check	✅ Passed	The PR implements the requested helpers and applies them across dpmodel and PyTorch loss paths with matching test coverage for the refactor.
Out of Scope Changes check	✅ Passed	The changes are limited to the shared reduction refactor and supporting tests, with no clear unrelated additions.

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

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

source/tests/common/dpmodel/test_loss_reduction.py (1)

19-33: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

_ref duplicates the implementation formula rather than independently verifying it.

_ref mirrors masked_atom_mean's exact computation, so test_numpy_matches_reference mainly confirms the numpy dispatch path matches itself rather than validating correctness against an independently derived expected value. Consider adding at least one case with a manually computed expected scalar to catch a shared logic bug in both.

🤖 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/common/dpmodel/test_loss_reduction.py` around lines 19 - 33,
Replace or supplement the formula-based _ref oracle in
test_numpy_matches_reference with at least one test case whose expected
reduction scalar is manually computed and asserted directly. Keep coverage for
both ncomp values as appropriate, but ensure the test can detect a shared error
in masked_atom_mean and _ref rather than comparing two identical
implementations.

🤖 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/dpmodel/loss/reduction.py`:
- Around line 30-55: Update masked_atom_mean to guard zero values in
per_frame_dof before dividing per_frame_sum by it, assigning a finite neutral
per-frame result for all-padding frames while preserving the existing
calculation for positive degrees of freedom and the final frame mean.

---

Nitpick comments:
In `@source/tests/common/dpmodel/test_loss_reduction.py`:
- Around line 19-33: Replace or supplement the formula-based _ref oracle in
test_numpy_matches_reference with at least one test case whose expected
reduction scalar is manually computed and asserted directly. Keep coverage for
both ncomp values as appropriate, but ensure the test can detect a shared error
in masked_atom_mean and _ref rather than comparing two identical
implementations.

🪄 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: 611b6d93-b490-4ef1-b0e6-52b9b72acb1c

📥 Commits

Reviewing files that changed from the base of the PR and between 0601c79 and 13b74c9.

📒 Files selected for processing (10)

• deepmd/dpmodel/loss/dos.py
• deepmd/dpmodel/loss/ener.py
• deepmd/dpmodel/loss/ener_spin.py
• deepmd/dpmodel/loss/reduction.py
• deepmd/dpmodel/loss/tensor.py
• deepmd/pt/loss/dos.py
• deepmd/pt/loss/ener.py
• deepmd/pt/loss/ener_spin.py
• deepmd/pt/loss/tensor.py
• source/tests/common/dpmodel/test_loss_reduction.py

[codecov]

Codecov Report

❌ Patch coverage is 87.09677% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.32%. Comparing base (0601c79) to head (97c26f5).
⚠️ Report is 11 commits behind head on master.

Files with missing lines	Patch %	Lines
deepmd/dpmodel/loss/ener.py	78.57%	3 Missing ⚠️
deepmd/dpmodel/loss/ener_spin.py	75.00%	3 Missing ⚠️
deepmd/pt/loss/ener.py	80.00%	3 Missing ⚠️
deepmd/pt/loss/ener_spin.py	75.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5783      +/-   ##
==========================================
- Coverage   79.77%   78.32%   -1.46%     
==========================================
  Files        1020     1050      +30     
  Lines      116872   120545    +3673     
  Branches     4308     4348      +40     
==========================================
+ Hits        93239    94413    +1174     
- Misses      22095    24574    +2479     
- Partials     1538     1558      +20     

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

[njzjz-bot]

Reviewed commit 13b74c9. I found no blocking correctness issues in the shared reduction helpers or their dpmodel and PyTorch call sites. The refactor preserves the existing shape, dtype, masking, normalization, and reduction order.

Local validation:

• common helper and padding tests: 65 passed
• PyTorch padding tests: 58 passed
• cross-backend loss consistency tests: 232 passed, 188 skipped
• Ruff check and format check: passed

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