fix(loss): normalize global RMSE display per frame for mixed-size batches

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)

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)

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)

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
 

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
             )

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

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