diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index e39e0a6a9e..8a34cbdf24 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -268,14 +268,18 @@ 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 @@ -283,7 +287,7 @@ def call( # 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 diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 84c229e1c2..e3437aa2f5 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/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." diff --git a/source/tests/common/dpmodel/test_loss_ener_spin.py b/source/tests/common/dpmodel/test_loss_ener_spin.py new file mode 100644 index 0000000000..5a3c37db39 --- /dev/null +++ b/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]) diff --git a/source/tests/consistent/loss/test_ener_spin.py b/source/tests/consistent/loss/test_ener_spin.py index 9f8137be73..b4aad2b4fe 100644 --- a/source/tests/consistent/loss/test_ener_spin.py +++ b/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. diff --git a/source/tests/pt_expt/loss/test_ener_spin.py b/source/tests/pt_expt/loss/test_ener_spin.py index c39f4d166b..2cd49f992a 100644 --- a/source/tests/pt_expt/loss/test_ener_spin.py +++ b/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"]) + )