fix(dpmodel): guard empty magnetic loss masks (#5798)

## 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:

```python
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


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **Bug Fixes**
- Fixed magnetic-force loss calculations when no magnetic atoms are
present.
  - Prevented NaN or invalid loss values in both MSE and MAE modes.
- Ensured magnetic-force metrics and overall loss report finite zero
values for empty magnetic masks.
  - Aligned results across supported computational backends.

- **Tests**
- Added regression coverage across supported precisions and loss
configurations.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz-bot@users.noreply.github.com>
Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>

Author: 3 people

deepmd/dpmodel/loss/ener_spin.py

@@ -284,22 +284,26 @@ def call(
             # zero out non-magnetic atoms
             diff_fm = (force_mag_label - force_mag_pred) * mask_float
             n_valid = xp.sum(mask_float)
+            # Guard the denominator itself because array backends may evaluate
+            # both branches of ``where``. This is safe under JAX tracing and
+            # makes an all-empty magnetic mask contribute exactly zero.
+            safe_n_valid = xp.where(n_valid > 0, n_valid, xp.ones_like(n_valid))
             if self.loss_func == "mse":
-                l2_force_mag_loss = xp.sum(xp.square(diff_fm)) / (n_valid * 3)
+                l2_force_mag_loss = xp.sum(xp.square(diff_fm)) / (safe_n_valid * 3)
                 loss += pref_fm * l2_force_mag_loss
                 more_loss["rmse_fm"] = self.display_if_exist(
                     xp.sqrt(l2_force_mag_loss), find_force_mag
                 )
                 if mae:
-                    mae_fm = xp.sum(xp.abs(diff_fm)) / (n_valid * 3)
+                    mae_fm = xp.sum(xp.abs(diff_fm)) / (safe_n_valid * 3)
                     more_loss["mae_fm"] = self.display_if_exist(mae_fm, find_force_mag)
             elif self.loss_func == "mae":
                 abs_diff_fm = xp.abs(diff_fm)  # [nf, na, 3], zeros for non-magnetic
                 # Mean over frames, magnetic atoms and xyz (same reduction as
                 # force_mag MSE, force_real MAE and the displayed mae_fm) so the
                 # loss is batch-size independent: a 2-frame batch equals the mean
                 # of the two single-frame losses.
-                l1_force_mag_loss = xp.sum(abs_diff_fm) / (n_valid * 3)
+                l1_force_mag_loss = xp.sum(abs_diff_fm) / (safe_n_valid * 3)
                 loss += pref_fm * l1_force_mag_loss
                 more_loss["mae_fm"] = self.display_if_exist(
                     l1_force_mag_loss, find_force_mag

deepmd/pt/loss/ener_spin.py

@@ -368,15 +368,14 @@ def forward(
                 # over an empty tensor (NaN), and ``0 * NaN`` is still NaN, so
                 # ``nan_to_num`` keeps such a batch's force_mag term at zero
                 # loss and zero gradient instead of poisoning the whole step.
-                loss += (pref_fm * torch.nan_to_num(l2_force_mag_loss)).to(
-                    GLOBAL_PT_FLOAT_PRECISION
-                )
-                rmse_fm = l2_force_mag_loss.sqrt()
+                safe_l2_force_mag_loss = torch.nan_to_num(l2_force_mag_loss)
+                loss += (pref_fm * safe_l2_force_mag_loss).to(GLOBAL_PT_FLOAT_PRECISION)
+                rmse_fm = safe_l2_force_mag_loss.sqrt()
                 more_loss["rmse_fm"] = self.display_if_exist(
                     rmse_fm.detach(), find_force_m
                 )
                 if mae:
-                    mae_fm = torch.mean(torch.abs(diff_fm))
+                    mae_fm = torch.nan_to_num(torch.mean(torch.abs(diff_fm)))
                     more_loss["mae_fm"] = self.display_if_exist(
                         mae_fm.detach(), find_force_m
                     )
@@ -385,13 +384,13 @@ def forward(
                 # force_mag MSE, force_real MAE and the displayed mae_fm) so the
                 # loss is batch-size independent: a 2-frame batch equals the mean
                 # of the two single-frame losses.
-                l1_force_mag_loss = torch.mean(torch.abs(label_fm - pred_fm))
+                l1_force_mag_loss = torch.nan_to_num(
+                    torch.mean(torch.abs(label_fm - pred_fm))
+                )
                 more_loss["mae_fm"] = self.display_if_exist(
                     l1_force_mag_loss.detach(), find_force_m
                 )
-                loss += (pref_fm * torch.nan_to_num(l1_force_mag_loss)).to(
-                    GLOBAL_PT_FLOAT_PRECISION
-                )
+                loss += (pref_fm * l1_force_mag_loss).to(GLOBAL_PT_FLOAT_PRECISION)
             else:
                 raise NotImplementedError(
                     f"Loss type {self.loss_func} is not implemented for magnetic force loss."

source/tests/common/dpmodel/test_loss_ener_spin.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Regression tests for backend-independent spin-energy losses."""
+
+import numpy as np
+import pytest
+
+from deepmd.dpmodel.loss.ener_spin import (
+    EnergySpinLoss,
+)
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("loss_func", ["mse", "mae"])
+@pytest.mark.parametrize("find_force_mag", [0.0, 1.0])
+def test_no_magnetic_atoms(dtype, loss_func, find_force_mag) -> None:
+    """An all-false magnetic mask must not poison the total loss.
+
+    The magnetic-force prefactor remains enabled even when the label is absent,
+    reproducing the case where ``0 * NaN`` previously contaminated the loss.
+    Display metrics retain the standard ``display_if_exist`` behavior: zero for
+    a present-but-empty label and NaN when the label itself is absent.
+    """
+    nframes, natoms = 2, 6
+    loss_fn = EnergySpinLoss(
+        starter_learning_rate=1.0,
+        start_pref_fm=1.0,
+        limit_pref_fm=1.0,
+        loss_func=loss_func,
+    )
+    model_pred = {
+        # Energy selects the NumPy Array API namespace even though its loss term
+        # is disabled.
+        "energy": np.zeros((nframes,), dtype=dtype),
+        "force_mag": np.ones((nframes, natoms, 3), dtype=dtype),
+        "mask_mag": np.zeros((nframes, natoms, 1), dtype=bool),
+    }
+    label = {
+        "force_mag": np.zeros((nframes, natoms, 3), dtype=dtype),
+        "find_force_mag": find_force_mag,
+    }
+
+    loss, more_loss = loss_fn(
+        1.0,
+        natoms,
+        model_pred,
+        label,
+        mae=True,
+    )
+
+    assert np.isfinite(loss)
+    np.testing.assert_equal(loss, 0.0)
+    np.testing.assert_equal(more_loss["rmse"], 0.0)
+    metric_names = ["mae_fm"]
+    if loss_func == "mse":
+        metric_names.append("rmse_fm")
+    for name in metric_names:
+        if find_force_mag:
+            np.testing.assert_equal(more_loss[name], 0.0)
+        else:
+            assert np.isnan(more_loss[name])

source/tests/consistent/loss/test_ener_spin.py

@@ -229,6 +229,58 @@ def atol(self) -> float:
         return 1e-10
 
 
+class TestEnerSpinEmptyMagneticMaskConsistency(unittest.TestCase):
+    """Keep PT display metrics aligned with dpmodel for empty spin masks."""
+
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_empty_mask_metrics_are_finite_zero(self) -> None:
+        nframes, natoms = 2, 4
+        predict = {
+            # The dpmodel loss uses energy to resolve the active array
+            # namespace even when the energy term is disabled.
+            "energy": np.zeros(nframes, dtype=np.float64),
+            "force_mag": np.ones((nframes, natoms, 3), dtype=np.float64),
+            "mask_mag": np.zeros((nframes, natoms, 1), dtype=bool),
+        }
+        label = {
+            "force_mag": np.zeros((nframes, natoms, 3), dtype=np.float64),
+            "find_force_mag": 1.0,
+        }
+
+        for loss_func, mae in (("mse", False), ("mse", True), ("mae", False)):
+            with self.subTest(loss_func=loss_func, mae=mae):
+                kwargs = {
+                    "starter_learning_rate": 1e-3,
+                    "start_pref_fm": 1.0,
+                    "limit_pref_fm": 1.0,
+                    "loss_func": loss_func,
+                }
+                dp_loss = EnerSpinLossDP(**kwargs)
+                pt_loss = EnerSpinLossPT(**kwargs)
+
+                dp_value, dp_more = dp_loss(1e-3, natoms, predict, label, mae=mae)
+                pt_predict = {
+                    key: numpy_to_torch(value) for key, value in predict.items()
+                }
+                pt_label = {
+                    key: numpy_to_torch(value)
+                    if isinstance(value, np.ndarray)
+                    else value
+                    for key, value in label.items()
+                }
+                _, pt_value, pt_more = pt_loss(
+                    {}, lambda: pt_predict, pt_label, natoms, 1e-3, mae=mae
+                )
+
+                np.testing.assert_allclose(torch_to_numpy(pt_value), dp_value)
+                expected_keys = {"rmse_fm"} if loss_func == "mse" else {"mae_fm"}
+                if mae:
+                    expected_keys.add("mae_fm")
+                for key in expected_keys:
+                    self.assertEqual(float(np.asarray(dp_more[key])), 0.0)
+                    self.assertEqual(float(torch_to_numpy(pt_more[key])), 0.0)
+
+
 class TestEnerSpinIntensiveScaling(unittest.TestCase):
     """Regression test for natoms-scaling behavior with intensive normalization.
 

source/tests/pt_expt/loss/test_ener_spin.py

@@ -253,3 +253,43 @@ def test_all_masked(self, prec) -> None:
             atol=atol,
             err_msg="pt_expt vs dpmodel (all masked)",
         )
+
+    @pytest.mark.parametrize("prec", ["float64", "float32"])
+    @pytest.mark.parametrize("loss_func", ["mse", "mae"])
+    def test_no_magnetic_atoms(self, prec, loss_func) -> None:
+        """An all-false magnetic mask has a finite zero contribution.
+
+        This complements the backend-neutral NumPy regression by exercising the
+        same dpmodel implementation through the Torch Array API namespace.
+        """
+        rng = np.random.default_rng(GLOBAL_SEED + 3)
+        nframes, natoms = 2, 6
+        dtype = PRECISION_DICT[prec]
+        learning_rate = 1e-3
+        loss_fn = EnergySpinLoss(
+            starter_learning_rate=learning_rate,
+            start_pref_fm=1.0,
+            limit_pref_fm=1.0,
+            loss_func=loss_func,
+        )
+        model_pred, label = _make_data(rng, nframes, natoms, 0, dtype, self.device)
+
+        loss, more_loss = loss_fn(
+            learning_rate,
+            natoms,
+            model_pred,
+            label,
+            mae=True,
+        )
+        assert torch.isfinite(loss)
+        torch.testing.assert_close(loss, torch.zeros_like(loss))
+        torch.testing.assert_close(
+            more_loss["mae_fm"], torch.zeros_like(more_loss["mae_fm"])
+        )
+        torch.testing.assert_close(
+            more_loss["rmse"], torch.zeros_like(more_loss["rmse"])
+        )
+        if loss_func == "mse":
+            torch.testing.assert_close(
+                more_loss["rmse_fm"], torch.zeros_like(more_loss["rmse_fm"])
+            )