diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py index 2222fb0e73..0c07e58fbf 100644 --- a/deepmd/dpmodel/loss/dos.py +++ b/deepmd/dpmodel/loss/dos.py @@ -13,7 +13,7 @@ ) from deepmd.dpmodel.loss.reduction import ( masked_atom_mean, - masked_atom_num, + masked_global_rmse, ) from deepmd.utils.data import ( DataRequirementItem, @@ -204,10 +204,9 @@ def call( diff = global_pred - global_label # idiom 3: global dos is already padding-invariant; plain mean suffices l2_global_loss_dos = xp.mean(xp.square(diff)) - atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype) loss += pref_dos * l2_global_loss_dos more_loss["rmse_global_dos"] = self.display_if_exist( - xp.sqrt(l2_global_loss_dos) / atom_num, find_global + masked_global_rmse(diff, model_dict.get("mask"), natoms), find_global ) if self.has_cdf and "dos" in model_dict and "dos" in label_dict: @@ -222,10 +221,9 @@ def call( diff = global_pred_cdf - global_label_cdf # idiom 3: global cdf is already padding-invariant; plain mean suffices l2_global_loss_cdf = xp.mean(xp.square(diff)) - atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype) loss += pref_cdf * l2_global_loss_cdf more_loss["rmse_global_cdf"] = self.display_if_exist( - xp.sqrt(l2_global_loss_cdf) / atom_num, find_global + masked_global_rmse(diff, model_dict.get("mask"), natoms), find_global ) more_loss["rmse"] = xp.sqrt(loss) diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py index 39980cc707..b4800948f7 100644 --- a/deepmd/dpmodel/loss/reduction.py +++ b/deepmd/dpmodel/loss/reduction.py @@ -204,3 +204,46 @@ def masked_atom_num(mask: Array | None, natoms: Any, dtype: Any) -> Any: return natoms xp = array_api_compat.array_namespace(mask) return xp.mean(xp.astype(xp.sum(mask, axis=-1), dtype)) + + +def masked_global_rmse(diff: Array, mask: Array | None, natoms: Any) -> Array: + """Display-only RMSE of an extensive global term, normalized per atom. + + Without a mask every frame has ``natoms`` atoms and the RMSE of the whole + batch is divided by it. With a mask, frames may hold different numbers of + real atoms, so each frame's mean squared error is normalized by its own + squared atom count before averaging over frames. A batch then reports the + same value as the mean of its frames' squared RMSEs taken one at a time. + + Parameters + ---------- + diff : Array + Global prediction error of shape ``[nf, ncomp]``. + mask : Array or None + Per-atom mask of shape ``[nf, nloc]``, or ``None`` when not mixed_type. + natoms + Atom count used when ``mask`` is ``None``. + + Returns + ------- + Array + ``sqrt(mean(diff**2)) / natoms`` without a mask, else + ``sqrt(mean_f(mean(diff_f**2) / n_f**2))`` with ``n_f`` the number of + real atoms of frame ``f``. Frames without any real atom are excluded + from the average; a batch with none reports ``0``. + """ + xp = array_api_compat.array_namespace(diff) + if mask is None: + return xp.sqrt(xp.mean(xp.square(diff))) / natoms + real_natoms = xp.astype(xp.sum(mask, axis=-1), diff.dtype) + frame_mse = xp.mean(xp.square(diff), axis=-1) + # A frame of nothing but padding has no atom to normalize by; it is left + # out of the average, and a batch without any real atom reports 0. + has_atoms = real_natoms > 0 + safe_natoms = xp.where(has_atoms, real_natoms, xp.ones_like(real_natoms)) + frame_val = xp.where( + has_atoms, frame_mse / xp.square(safe_natoms), xp.zeros_like(frame_mse) + ) + nvalid = xp.sum(xp.astype(has_atoms, diff.dtype)) + safe_nvalid = xp.where(nvalid > 0, nvalid, xp.ones_like(nvalid)) + return xp.sqrt(xp.sum(frame_val) / safe_nvalid) diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index 259cbf3057..c0a21e5463 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -13,7 +13,7 @@ ) from deepmd.dpmodel.loss.reduction import ( masked_atom_mean, - masked_atom_num, + masked_global_rmse, ) from deepmd.utils.data import ( DataRequirementItem, @@ -35,9 +35,9 @@ class TensorLoss(Loss): [a_i(T_i-\hat T_i)]^2\right\rangle +p_{\mathrm{global}}\left\langle(T-\hat T)^2\right\rangle. - Padded atoms are omitted from the local mean. Reported global RMSE is - divided by the number of real atoms, matching the convention for - extensive tensor labels. + Padded atoms are omitted from the local mean. Reported global RMSE + divides each frame's error by that frame's number of real atoms, + matching the convention for extensive tensor labels. Parameters ---------- @@ -155,10 +155,9 @@ def call( diff = global_pred - global_label # idiom 3: global tensor is already padding-invariant; plain mean suffices l2_global_loss = xp.mean(xp.square(diff)) - atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype) loss += global_weight * l2_global_loss more_loss[f"rmse_global_{self.tensor_name}"] = self.display_if_exist( - xp.sqrt(l2_global_loss) / atom_num, find_global + masked_global_rmse(diff, model_dict.get("mask"), natoms), find_global ) more_loss["rmse"] = xp.sqrt(loss) diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index b5ba880c93..b38da71c97 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -7,7 +7,7 @@ from deepmd.dpmodel.loss.reduction import ( masked_atom_mean, - masked_atom_num, + masked_global_rmse, ) from deepmd.pt.loss.loss import ( TaskLoss, @@ -210,13 +210,12 @@ def forward( diff = global_tensor_pred_dos - global_tensor_label_dos # idiom 3: global dos is already padding-invariant; plain mean suffices l2_global_loss_dos = torch.mean(torch.square(diff)) - atom_num = masked_atom_num(model_pred.get("mask"), natoms, torch.float32) if not self.inference: more_loss["l2_global_dos_loss"] = self.display_if_exist( l2_global_loss_dos.detach(), find_global ) loss += pref_dos * l2_global_loss_dos - rmse_global_dos = l2_global_loss_dos.sqrt() / atom_num + rmse_global_dos = masked_global_rmse(diff, model_pred.get("mask"), natoms) more_loss["rmse_global_dos"] = self.display_if_exist( rmse_global_dos.detach(), find_global ) @@ -232,15 +231,14 @@ def forward( diff = global_tensor_pred_cdf - global_tensor_label_cdf # idiom 3: global cdf is already padding-invariant; plain mean suffices l2_global_loss_cdf = torch.mean(torch.square(diff)) - atom_num = masked_atom_num(model_pred.get("mask"), natoms, torch.float32) if not self.inference: more_loss["l2_global_cdf_loss"] = self.display_if_exist( l2_global_loss_cdf.detach(), find_global ) loss += pref_cdf * l2_global_loss_cdf - rmse_global_dos = l2_global_loss_cdf.sqrt() / atom_num + rmse_global_cdf = masked_global_rmse(diff, model_pred.get("mask"), natoms) more_loss["rmse_global_cdf"] = self.display_if_exist( - rmse_global_dos.detach(), find_global + rmse_global_cdf.detach(), find_global ) return model_pred, loss, more_loss diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 2f08467db2..eea446e0aa 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -7,7 +7,7 @@ from deepmd.dpmodel.loss.reduction import ( masked_atom_mean, - masked_atom_num, + masked_global_rmse, ) from deepmd.pt.loss.loss import ( TaskLoss, @@ -166,13 +166,12 @@ def forward( diff = global_tensor_pred - global_tensor_label # idiom 3: global tensor is already padding-invariant; plain mean suffices l2_global_loss = torch.mean(torch.square(diff)) - atom_num = masked_atom_num(model_pred.get("mask"), natoms, torch.float32) if not self.inference: more_loss[f"l2_global_{self.tensor_name}_loss"] = self.display_if_exist( l2_global_loss.detach(), find_global ) loss += global_weight * l2_global_loss - rmse_global = l2_global_loss.sqrt() / atom_num + rmse_global = masked_global_rmse(diff, model_pred.get("mask"), natoms) more_loss[f"rmse_global_{self.tensor_name}"] = self.display_if_exist( rmse_global.detach(), find_global ) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 17016aaf36..5551d53a4a 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -467,6 +467,97 @@ def test_no_op_for_non_mixed(self): ) +class TestGlobalRMSEDisplayPerFrame: + """Global rmse displays normalize each frame by its own real atom count.""" + + @staticmethod + def _frames(ncomp): + # extensive errors [NA, NB]: a per-atom error of 1 in each frame + label = _rnd(2, ncomp) + pred = label + np.array([[NA], [NB]], dtype=np.float64) + mask = np.array([[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64) + return pred, label, mask + + @staticmethod + def _assert_matches_frames(call, pred, label, mask, keys): + _, batch = call(pred, label, mask, NP) + _, more_A = call(pred[:1], label[:1], np.ones((1, NA)), NA) + _, more_B = call(pred[1:], label[1:], np.ones((1, NB)), NB) + for key in keys: + expected = np.sqrt((float(more_A[key]) ** 2 + float(more_B[key]) ** 2) / 2) + np.testing.assert_allclose(float(batch[key]), expected, rtol=1e-12) + + def test_dos_and_cdf(self): + loss = DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=1.0, + limit_pref_dos=1.0, + start_pref_cdf=1.0, + limit_pref_cdf=1.0, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + def call(pred, label, mask, natoms): + return loss.call( + 1.0, + natoms, + {"dos": pred, "mask": mask}, + {"dos": label, "find_dos": 1.0}, + ) + + pred, label, mask = self._frames(NUMB_DOS) + _, more = call(pred, label, mask, NP) + np.testing.assert_allclose(float(more["rmse_global_dos"]), 1.0, rtol=1e-12) + cdf = np.sqrt(np.mean(np.square(np.arange(1, NUMB_DOS + 1)))) + np.testing.assert_allclose(float(more["rmse_global_cdf"]), cdf, rtol=1e-12) + pred, label = _rnd(2, NUMB_DOS), _rnd(2, NUMB_DOS) + self._assert_matches_frames( + call, pred, label, mask, ("rmse_global_dos", "rmse_global_cdf") + ) + + def test_tensor(self): + loss = TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + + def call(pred, label, mask, natoms): + return loss.call( + 1.0, + natoms, + {"global_dipole": pred, "mask": mask}, + {"dipole": label, "find_dipole": 1.0}, + ) + + pred, label, mask = self._frames(TENSOR_SIZE) + _, more = call(pred, label, mask, NP) + np.testing.assert_allclose(float(more["rmse_global_dipole"]), 1.0, rtol=1e-12) + pred, label = _rnd(2, TENSOR_SIZE), _rnd(2, TENSOR_SIZE) + self._assert_matches_frames(call, pred, label, mask, ("rmse_global_dipole",)) + + def test_unmasked_unchanged(self): + loss = TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + pred, label = _rnd(2, TENSOR_SIZE), _rnd(2, TENSOR_SIZE) + _, more = loss.call( + 1.0, NB, {"global_dipole": pred}, {"dipole": label, "find_dipole": 1.0} + ) + expected = np.sqrt(np.mean(np.square(pred - label))) / NB + assert float(more["rmse_global_dipole"]) == expected + + # --------------------------------------------------------------------------- # Task 3: EnergyLoss -- energy, force, virial, atom_ener, atom_pref # --------------------------------------------------------------------------- diff --git a/source/tests/common/dpmodel/test_loss_reduction.py b/source/tests/common/dpmodel/test_loss_reduction.py index ddb59993c3..a4dbe93d9f 100644 --- a/source/tests/common/dpmodel/test_loss_reduction.py +++ b/source/tests/common/dpmodel/test_loss_reduction.py @@ -7,6 +7,7 @@ from deepmd.dpmodel.loss.reduction import ( masked_atom_mean, masked_atom_num, + masked_global_rmse, per_frame_component_mean, ) @@ -133,3 +134,28 @@ def test_torch_bit_identical_float32(self) -> None: got = masked_atom_num(torch.tensor(mask_np), 3, torch.float32) ref = torch.tensor(mask_np).sum(-1).float().mean() assert got.item() == ref.item() + + +class TestMaskedGlobalRMSE: + """Display-only RMSE of an extensive global term, normalized per frame.""" + + def test_all_padding_frame_is_excluded(self) -> None: + diff = np.array([[2.0, 2.0], [5.0, 5.0]]) + mask = np.array([[1.0, 1.0, 0.0], [0.0, 0.0, 0.0]]) + got = masked_global_rmse(diff, mask, 3) + np.testing.assert_allclose(got, 1.0, rtol=1e-12) + + def test_all_padding_batch_is_zero(self) -> None: + diff = np.array([[2.0, 2.0], [5.0, 5.0]]) + mask = np.zeros((2, 3)) + got = masked_global_rmse(diff, mask, 3) + assert np.isfinite(got) + assert float(got) == 0.0 + + def test_torch_all_padding_frame_is_excluded(self) -> None: + diff = torch.tensor([[2.0, 2.0], [5.0, 5.0]], dtype=torch.float64) + mask = torch.tensor([[1.0, 1.0, 0.0], [0.0, 0.0, 0.0]], dtype=torch.float64) + got = masked_global_rmse(diff, mask, 3) + assert got.item() == pytest.approx(1.0, rel=1e-12) + got = masked_global_rmse(diff, torch.zeros_like(mask), 3) + assert got.item() == 0.0 diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index b0455d771e..5109da0519 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -524,6 +524,111 @@ def test_no_op_for_non_mixed(self): ) +class TestPTGlobalRMSEDisplayPerFrame: + """Global rmse displays normalize each frame by its own real atom count.""" + + @staticmethod + def _mask(): + return torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", + ) + + @staticmethod + def _more_loss(loss_obj, model_pred, label, natoms): + _, _, more = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return more + + def _assert_matches_frames(self, call, pred, label, keys): + batch = call(pred, label, self._mask(), NP) + more_A = call( + pred[:1], + label[:1], + torch.ones(1, NA, dtype=torch.float64, device="cpu"), + NA, + ) + more_B = call( + pred[1:], + label[1:], + torch.ones(1, NB, dtype=torch.float64, device="cpu"), + NB, + ) + for key in keys: + expected = np.sqrt((more_A[key].item() ** 2 + more_B[key].item() ** 2) / 2) + np.testing.assert_allclose(batch[key].item(), expected, rtol=1e-12) + + def test_dos_and_cdf(self): + loss_obj = DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=1.0, + limit_pref_dos=1.0, + start_pref_cdf=1.0, + limit_pref_cdf=1.0, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + def call(pred, label, mask, natoms): + return self._more_loss( + loss_obj, + {"dos": pred, "mask": mask}, + {"dos": label, "find_dos": 1.0}, + natoms, + ) + + # extensive errors [NA, NB]: a per-atom error of 1 in each frame + label = _rnd_t(2, NUMB_DOS) + pred = label + torch.tensor([[NA], [NB]], dtype=torch.float64, device="cpu") + more = call(pred, label, self._mask(), NP) + np.testing.assert_allclose(more["rmse_global_dos"].item(), 1.0, rtol=1e-12) + cdf = np.sqrt(np.mean(np.square(np.arange(1, NUMB_DOS + 1)))) + np.testing.assert_allclose(more["rmse_global_cdf"].item(), cdf, rtol=1e-12) + self._assert_matches_frames( + call, + _rnd_t(2, NUMB_DOS), + _rnd_t(2, NUMB_DOS), + ("rmse_global_dos", "rmse_global_cdf"), + ) + + def test_tensor(self): + loss_obj = TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + + def call(pred, label, mask, natoms): + return self._more_loss( + loss_obj, + {"global_dipole": pred, "mask": mask}, + {"dipole": label, "find_dipole": 1.0}, + natoms, + ) + + label = _rnd_t(2, TENSOR_SIZE) + pred = label + torch.tensor([[NA], [NB]], dtype=torch.float64, device="cpu") + more = call(pred, label, self._mask(), NP) + np.testing.assert_allclose(more["rmse_global_dipole"].item(), 1.0, rtol=1e-12) + self._assert_matches_frames( + call, + _rnd_t(2, TENSOR_SIZE), + _rnd_t(2, TENSOR_SIZE), + ("rmse_global_dipole",), + ) + + # --------------------------------------------------------------------------- # Task 1: unit tests for TaskLoss._inject_atom_mask # ---------------------------------------------------------------------------