fix(dpmodel): guard empty magnetic loss masks

deepmd/dpmodel/loss/ener_spin.py

@@ -268,22 +268,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

@@ -370,15 +370,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
                     )
@@ -387,13 +386,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"])
+            )