fix: honor loss schedules with a zero endpoint

deepmd/pd/loss/ener.py

@@ -171,12 +171,12 @@ def __init__(
             )
 
         self.starter_learning_rate = starter_learning_rate
-        self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference
-        self.has_f = (start_pref_f != 0.0 and limit_pref_f != 0.0) or inference
-        self.has_v = (start_pref_v != 0.0 and limit_pref_v != 0.0) or inference
-        self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference
-        self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference
-        self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0
+        self.has_e = (start_pref_e != 0.0 or limit_pref_e != 0.0) or inference
+        self.has_f = (start_pref_f != 0.0 or limit_pref_f != 0.0) or inference
+        self.has_v = (start_pref_v != 0.0 or limit_pref_v != 0.0) or inference
+        self.has_ae = (start_pref_ae != 0.0 or limit_pref_ae != 0.0) or inference
+        self.has_pf = (start_pref_pf != 0.0 or limit_pref_pf != 0.0) or inference
+        self.has_gf = start_pref_gf != 0.0 or limit_pref_gf != 0.0
 
         self.start_pref_e = start_pref_e
         self.limit_pref_e = limit_pref_e
@@ -304,16 +304,11 @@ def forward(
                 l1_ener_loss = F.l1_loss(
                     energy_pred.reshape([-1]),
                     energy_label.reshape([-1]),
-                    reduction="sum",
+                    reduction="mean",
                 )
-                loss += pref_e * l1_ener_loss
+                loss += atom_norm * (pref_e * l1_ener_loss)
                 more_loss["mae_e"] = self.display_if_exist(
-                    F.l1_loss(
-                        energy_pred.reshape([-1]),
-                        energy_label.reshape([-1]),
-                        reduction="mean",
-                    ).detach(),
-                    find_energy,
+                    l1_ener_loss.detach() * atom_norm, find_energy
                 )
                 # more_loss['log_keys'].append('rmse_e')
             else:
@@ -367,11 +362,10 @@ def forward(
                         rmse_f.detach(), find_force
                     )
                 elif self.loss_func == "mae":
-                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="mean")
+                    l1_force_loss = paddle.mean(paddle.abs(diff_f))
                     more_loss["mae_f"] = self.display_if_exist(
                         l1_force_loss.detach(), find_force
                     )
-                    l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
                     loss += (pref_f * l1_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
                 else:
                     raise NotImplementedError(
@@ -386,7 +380,8 @@ def forward(
             if self.has_pf and "atom_pref" in label:
                 atom_pref = label["atom_pref"]
                 find_atom_pref = label.get("find_atom_pref", 0.0)
-                pref_pf = pref_pf * find_atom_pref
+                effective_find_pf = find_force * find_atom_pref
+                pref_pf = pref_pf * effective_find_pf
                 atom_pref_reshape = atom_pref.reshape([-1])
 
                 if self.loss_func == "mse":
@@ -395,18 +390,18 @@ def forward(
                     ).mean()
                     if not self.inference:
                         more_loss["l2_pref_force_loss"] = self.display_if_exist(
-                            l2_pref_force_loss.detach(), find_atom_pref
+                            l2_pref_force_loss.detach(), effective_find_pf
                         )
                     loss += (pref_pf * l2_pref_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
                     rmse_pf = l2_pref_force_loss.sqrt()
                     more_loss["rmse_pf"] = self.display_if_exist(
-                        rmse_pf.detach(), find_atom_pref
+                        rmse_pf.detach(), effective_find_pf
                     )
                 elif self.loss_func == "mae":
                     l1_pref_force_loss = (paddle.abs(diff_f) * atom_pref_reshape).mean()
                     loss += (pref_pf * l1_pref_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
                     more_loss["mae_pf"] = self.display_if_exist(
-                        l1_pref_force_loss.detach(), find_atom_pref
+                        l1_pref_force_loss.detach(), effective_find_pf
                     )
                 else:
                     raise NotImplementedError(
@@ -416,7 +411,8 @@ def forward(
             if self.has_gf and "drdq" in label:
                 drdq = label["drdq"]
                 find_drdq = label.get("find_drdq", 0.0)
-                pref_gf = pref_gf * find_drdq
+                effective_find_gf = find_force * find_drdq
+                pref_gf = pref_gf * effective_find_gf
                 force_reshape_nframes = force_pred.reshape([-1, natoms * 3])
                 force_label_reshape_nframes = force_label.reshape([-1, natoms * 3])
                 drdq_reshape = drdq.reshape(
@@ -440,34 +436,47 @@ def forward(
                 l2_gen_force_loss = paddle.square(diff_gen_force).mean()
                 if not self.inference:
                     more_loss["l2_gen_force_loss"] = self.display_if_exist(
-                        l2_gen_force_loss.detach(), find_drdq
+                        l2_gen_force_loss.detach(), effective_find_gf
                     )
                 loss += (pref_gf * l2_gen_force_loss).to(GLOBAL_PD_FLOAT_PRECISION)
                 rmse_gf = l2_gen_force_loss.sqrt()
                 more_loss["rmse_gf"] = self.display_if_exist(
-                    rmse_gf.detach(), find_drdq
+                    rmse_gf.detach(), effective_find_gf
                 )
 
         if self.has_v and "virial" in model_pred and "virial" in label:
             find_virial = label.get("find_virial", 0.0)
             pref_v = pref_v * find_virial
             diff_v = label["virial"] - model_pred["virial"].reshape([-1, 9])
-            l2_virial_loss = paddle.mean(paddle.square(diff_v))
-            if not self.inference:
-                more_loss["l2_virial_loss"] = self.display_if_exist(
-                    l2_virial_loss.detach(), find_virial
+            if self.loss_func == "mse" or self.use_huber:
+                l2_virial_loss = paddle.mean(paddle.square(diff_v))
+                if not self.inference:
+                    more_loss["l2_virial_loss"] = self.display_if_exist(
+                        l2_virial_loss.detach(), find_virial
+                    )
+                if not self.use_huber:
+                    loss += atom_norm**norm_exp * (pref_v * l2_virial_loss)
+                else:
+                    l_huber_loss = custom_huber_loss(
+                        atom_norm * model_pred["virial"].reshape([-1]),
+                        atom_norm * label["virial"].reshape([-1]),
+                        delta=self._huber_delta_virial,
+                    )
+                    loss += pref_v * l_huber_loss
+                rmse_v = l2_virial_loss.sqrt() * atom_norm
+                more_loss["rmse_v"] = self.display_if_exist(
+                    rmse_v.detach(), find_virial
+                )
+            elif self.loss_func == "mae":
+                l1_virial_loss = paddle.mean(paddle.abs(diff_v))
+                loss += atom_norm * (pref_v * l1_virial_loss)
+                more_loss["mae_v"] = self.display_if_exist(
+                    l1_virial_loss.detach() * atom_norm, find_virial
                 )
-            if not self.use_huber:
-                loss += atom_norm**norm_exp * (pref_v * l2_virial_loss)
             else:
-                l_huber_loss = custom_huber_loss(
-                    atom_norm * model_pred["virial"].reshape([-1]),
-                    atom_norm * label["virial"].reshape([-1]),
-                    delta=self._huber_delta_virial,
+                raise NotImplementedError(
+                    f"Loss type {self.loss_func} is not implemented for virial loss."
                 )
-                loss += pref_v * l_huber_loss
-            rmse_v = l2_virial_loss.sqrt() * atom_norm
-            more_loss["rmse_v"] = self.display_if_exist(rmse_v.detach(), find_virial)
             if mae:
                 mae_v = paddle.mean(paddle.abs(diff_v)) * atom_norm
                 more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial)
@@ -534,7 +543,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     high_prec=True,
                 )
             )
-        if self.has_f:
+        if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf:
             label_requirement.append(
                 DataRequirementItem(
                     "force",

deepmd/pt/loss/dos.py

@@ -84,10 +84,10 @@ def __init__(
             and self.limit_pref_acdf >= 0.0
         ), "Can not assign negative weight to `pref` and `pref_atomic`"
 
-        self.has_dos = (start_pref_dos != 0.0 and limit_pref_dos != 0.0) or inference
-        self.has_cdf = (start_pref_cdf != 0.0 and limit_pref_cdf != 0.0) or inference
-        self.has_ados = (start_pref_ados != 0.0 and limit_pref_ados != 0.0) or inference
-        self.has_acdf = (start_pref_acdf != 0.0 and limit_pref_acdf != 0.0) or inference
+        self.has_dos = (start_pref_dos != 0.0 or limit_pref_dos != 0.0) or inference
+        self.has_cdf = (start_pref_cdf != 0.0 or limit_pref_cdf != 0.0) or inference
+        self.has_ados = (start_pref_ados != 0.0 or limit_pref_ados != 0.0) or inference
+        self.has_acdf = (start_pref_acdf != 0.0 or limit_pref_acdf != 0.0) or inference
 
         assert self.has_dos or self.has_cdf or self.has_ados or self.has_acdf, (
             AssertionError(

deepmd/pt/loss/ener.py

@@ -160,12 +160,12 @@ def __init__(
 
         self.loss_func = loss_func
         self.starter_learning_rate = starter_learning_rate
-        self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference
-        self.has_f = (start_pref_f != 0.0 and limit_pref_f != 0.0) or inference
-        self.has_v = (start_pref_v != 0.0 and limit_pref_v != 0.0) or inference
-        self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference
-        self.has_pf = (start_pref_pf != 0.0 and limit_pref_pf != 0.0) or inference
-        self.has_gf = start_pref_gf != 0.0 and limit_pref_gf != 0.0
+        self.has_e = (start_pref_e != 0.0 or limit_pref_e != 0.0) or inference
+        self.has_f = (start_pref_f != 0.0 or limit_pref_f != 0.0) or inference
+        self.has_v = (start_pref_v != 0.0 or limit_pref_v != 0.0) or inference
+        self.has_ae = (start_pref_ae != 0.0 or limit_pref_ae != 0.0) or inference
+        self.has_pf = (start_pref_pf != 0.0 or limit_pref_pf != 0.0) or inference
+        self.has_gf = start_pref_gf != 0.0 or limit_pref_gf != 0.0
         # Hessian labels scale quadratically with the atom count. Unlike the
         # linear-size labels above, do not request them merely because a mock
         # inference loss is collecting requirements for ``dp change-bias``.

deepmd/pt/loss/ener_spin.py

@@ -122,11 +122,11 @@ def __init__(
 
         self.loss_func = loss_func
         self.starter_learning_rate = starter_learning_rate
-        self.has_e = (start_pref_e != 0.0 and limit_pref_e != 0.0) or inference
-        self.has_fr = (start_pref_fr != 0.0 and limit_pref_fr != 0.0) or inference
-        self.has_fm = (start_pref_fm != 0.0 and limit_pref_fm != 0.0) or inference
-        self.has_v = (start_pref_v != 0.0 and limit_pref_v != 0.0) or inference
-        self.has_ae = (start_pref_ae != 0.0 and limit_pref_ae != 0.0) or inference
+        self.has_e = (start_pref_e != 0.0 or limit_pref_e != 0.0) or inference
+        self.has_fr = (start_pref_fr != 0.0 or limit_pref_fr != 0.0) or inference
+        self.has_fm = (start_pref_fm != 0.0 or limit_pref_fm != 0.0) or inference
+        self.has_v = (start_pref_v != 0.0 or limit_pref_v != 0.0) or inference
+        self.has_ae = (start_pref_ae != 0.0 or limit_pref_ae != 0.0) or inference
 
         self.start_pref_e = start_pref_e
         self.limit_pref_e = limit_pref_e
@@ -597,6 +597,18 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     high_prec=False,
                 )
             )
+        if self.enable_atom_ener_coeff:
+            label_requirement.append(
+                DataRequirementItem(
+                    "atom_ener_coeff",
+                    ndof=1,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                    default=1.0,
+                    source_policy="default",
+                )
+            )
         return label_requirement
 
     def serialize(self) -> dict:

deepmd/utils/data.py

@@ -1428,6 +1428,8 @@ def _load_single_data(
             else:
                 # For non-atomic data, shape should be [ndof]
                 data = np.full([ndof], vv["default"], dtype=dtype)
+            if vv["repeat"] != 1 and vv.get("special_shape") != "frame_major":
+                data = np.repeat(data, vv["repeat"], axis=-1)
             return np.float32(0.0), data
 
         # Branch 2: Data loading
@@ -1526,6 +1528,8 @@ def _load_single_data(
             else:
                 data = data.reshape([ndof])
 
+            if vv["repeat"] != 1:
+                data = np.repeat(data, vv["repeat"], axis=-1)
             # Atomic: return [natoms, ndof] or flattened hessian above
             # Non-atomic: return [ndof]
             return np.float32(1.0), data

source/tests/common/test_single_frame_repeat.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import numpy as np
+import pytest
+
+from deepmd.utils.data import (
+    DeepmdData,
+)
+
+
+@pytest.mark.parametrize("present", [False, True])
+@pytest.mark.parametrize("nframes", [1, 2])
+@pytest.mark.parametrize("backend", ["torch", "paddle"])
+def test_single_frame_repeat_matches_batch(tmp_path, present, nframes, backend):
+    """Single-frame and batch loaders repeat atomic labels identically."""
+    np.savetxt(tmp_path / "type.raw", [1, 0], fmt="%d")
+    set_path = tmp_path / "set.000"
+    set_path.mkdir()
+    np.save(set_path / "coord.npy", np.zeros((nframes, 6)))
+    np.save(set_path / "box.npy", np.tile(np.eye(3).reshape(1, 9), (nframes, 1)))
+    if present:
+        np.save(set_path / "atom_pref.npy", np.tile([[2.0, 3.0]], (nframes, 1)))
+    data = DeepmdData(str(tmp_path))
+    data.add("atom_pref", 1, atomic=True, repeat=3, default=1.0)
+    expected = data._load_set(set_path)["atom_pref"][0]
+    frame = getattr(data, f"get_item_{backend}")(0)
+    np.testing.assert_array_equal(frame["atom_pref"].reshape(-1), expected)
+    assert frame["find_atom_pref"] == float(present)