deepmd/entrypoints/test.py (1)

421-435: Weighted force MAE/RMSE look correct; add a shape guard for robustness

The weighting logic (sum of abs residuals / sum of weights and sqrt(weighted MSE)) is sound, and size_f=sum(weights) makes the system-level aggregation consistent. To tolerate atom_pref provided as (nframes, natoms) without depending on repeat=3 at load time, add a defensive broadcast/validation.

-    if find_atom_pref == 1:
-        atom_pref = test_data["atom_pref"][:numb_test]
-        diff_f = diff_f * atom_pref
+    if find_atom_pref == 1:
+        atom_pref = test_data["atom_pref"][:numb_test]
+        # Accept both (nframes, natoms) and (nframes, natoms*3)
+        if atom_pref.shape != diff_f.shape:
+            if atom_pref.shape[1] * 3 == diff_f.shape[1]:
+                atom_pref = np.repeat(atom_pref, 3, axis=1)
+            else:
+                raise ValueError(
+                    f"atom_pref shape {atom_pref.shape} is incompatible with force shape {diff_f.shape}"
+                )
+        diff_f = diff_f * atom_pref
         size_f = np.sum(atom_pref)
         if size_f > 0:
             mae_f = np.sum(np.abs(diff_f)) / size_f
             rmse_f = np.sqrt(np.sum(diff_f * diff_f) / size_f)
         else:
             mae_f = 0.0
             rmse_f = 0.0

If desired, we can also factor this into small helpers (weighted_mae/weighted_rmse) to keep this block concise.

source/tests/pt/test_dp_test.py (2)

168-181: Mask preparation is correct; tiny robustness note about copytree target

Copying the sample system into a mkdtemp dir and writing atom_pref.npy as (nframes, natoms) is consistent with repeat=3 in the loader. Setting the last atom’s weight to 0 validates masking well. One minor robustness nit: shutil.copytree(..., dirs_exist_ok=True) is fine here, but if src ever contains symlinks or permissions need preserving, consider copy_function/copy2. Not required for this test data.

---

182-230: Good end-to-end assertions; add tolerances and verify the reported sample size

The comparison path is sound: run dp_test_ener, independently compute the masked MAE/RMSE from dp.eval outputs, and compare. To avoid flaky failures from float32 arithmetic and to verify the new size_f behavior, consider:

• Use explicit tolerances in assert_allclose.
• Assert that the returned sizes equal the mask sum.

-        np.testing.assert_allclose(err["mae_f"][0], mae_expected)
-        np.testing.assert_allclose(err["rmse_f"][0], rmse_expected)
+        np.testing.assert_allclose(err["mae_f"][0], mae_expected, rtol=1e-6, atol=1e-8)
+        np.testing.assert_allclose(err["rmse_f"][0], rmse_expected, rtol=1e-6, atol=1e-8)
+        # Also verify the effective sample size matches the mask sum
+        np.testing.assert_equal(err["mae_f"][1], denom)
+        np.testing.assert_equal(err["rmse_f"][1], denom)

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

source/tests/pt/test_dp_test.py (1)

168-181: Define a named constant for the force offset instead of magic number

This improves readability and mirrors prior feedback.

Apply:

 class TestDPTestForceWeight(DPTest, unittest.TestCase):
     def setUp(self) -> None:
         self.detail_file = "test_dp_test_force_weight_detail"
@@
-    def _prepare_weighted_system(self) -> str:
+    def _prepare_weighted_system(self) -> str:
         src = Path(__file__).parent / "water/data/single"
         tmp_dir = tempfile.mkdtemp()
         shutil.copytree(src, tmp_dir, dirs_exist_ok=True)
         set_dir = Path(tmp_dir) / "set.000"
         forces = np.load(set_dir / "force.npy")
-        forces[0, -3:] += 10.0
+        # bump the last atom only
+        self.FORCE_OFFSET = getattr(self, "FORCE_OFFSET", 10.0)
+        forces[0, -3:] += self.FORCE_OFFSET
         np.save(set_dir / "force.npy", forces)

deepmd/entrypoints/test.py (2)

294-295: Adding atom_pref looks good; consider avoiding 3x memory via compute-time broadcasting

Registering atom_pref as atomic with repeat=3 works, but it inflates this array 3x in memory and I/O. Since the per-atom weight is axis-invariant across x/y/z, you can keep atom_pref as shape [nframes, natoms] and broadcast it at compute time. This reduces memory and simplifies data loading while preserving semantics. See the follow-up diff in the weighting block for a safe broadcast that handles both shapes (with and without repeat), so you can optionally drop repeat=3 here.

Example (paired with the weighting-block change below):

-    data.add("atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3)
+    # Keep per-atom weights; broadcast to 3 components at compute-time
+    data.add("atom_pref", 1, atomic=True, must=False, high_prec=False)

---

620-623: Nice: system average now includes weighted force metrics

The conditional print of weighted force MAE/RMSE when present matches reviewer feedback (“show both with and without atom_pref”). Consider also printing the effective sample size (sum of weights) for transparency.

Example:

         if "rmse_fw" in avg:
             log.info(f"Force weighted MAE : {avg['mae_fw']:e} eV/A")
             log.info(f"Force weighted RMSE: {avg['rmse_fw']:e} eV/A")
+            if "mae_fw__size" in avg:
+                log.info(f"Force weighted effN: {avg['mae_fw__size']:g}")

Note: This needs weighted_average to optionally retain sizes in the aggregated output (e.g., by exposing them under a “__size” suffix).

source/tests/pt/test_dp_test.py (2)

217-219: Don’t hardcode mixed_type=False; derive from data to avoid surprises

Future tests may use mixed-type systems; pass the actual flag from DeepmdData.

Apply:

-            mixed_type=False,
+            mixed_type=data.mixed_type,

---

220-233: Strengthen test: also assert sizes (denominators) used in metrics

Verifies that dp_test_ener reports effective counts consistently for both weighted and unweighted metrics.

Apply:

         mae_unweighted = np.sum(np.abs(diff)) / diff.size
         rmse_unweighted = np.sqrt(np.sum(diff * diff) / diff.size)
         np.testing.assert_allclose(err["mae_f"][0], mae_unweighted)
         np.testing.assert_allclose(err["rmse_f"][0], rmse_unweighted)
+        # Also check the denominators
+        np.testing.assert_equal(err["mae_f"][1], diff.size)
+        np.testing.assert_equal(err["rmse_f"][1], diff.size)
         np.testing.assert_allclose(err["mae_fw"][0], mae_expected)
         np.testing.assert_allclose(err["rmse_fw"][0], rmse_expected)
+        np.testing.assert_allclose(err["mae_fw"][1], denom)
+        np.testing.assert_allclose(err["rmse_fw"][1], denom)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

source/tests/pt/test_dp_test.py (2)

168-182: Replace magic numbers and make intent explicit in weighted system prep

Use the constants suggested above; avoids magic numbers and clarifies which atoms are emphasized or ignored. Also fine to keep dirs_exist_ok=True since mkdtemp() pre-creates the directory.

         forces = np.load(set_dir / "force.npy")
-        forces[0, :3] += 1.0
-        forces[0, -3:] += 10.0
+        forces[0, :3] += self.FORCE_OFFSET_FIRST
+        forces[0, -3:] += self.FORCE_OFFSET_LAST
         np.save(set_dir / "force.npy", forces)
         natoms = forces.shape[1] // 3
         atom_pref = np.ones((forces.shape[0], natoms), dtype=forces.dtype)
-        atom_pref[:, 0] = 2.0
-        atom_pref[:, -1] = 0.0
+        atom_pref[:, 0] = self.FIRST_ATOM_WEIGHT
+        atom_pref[:, -1] = self.LAST_ATOM_WEIGHT
         np.save(set_dir / "atom_pref.npy", atom_pref)

---

207-214: Slice coord to 1 frame before reshape; current code flattens all frames

Reshape without slicing can consume all frames. Slice [:1] to match box and numb_test=1.

-        coord = test_data["coord"].reshape([1, -1])
+        coord = test_data["coord"][:1].reshape([1, -1])

source/tests/pt/test_dp_test.py (3)

150-167: Good test scaffolding; consider defining explicit constants for readability

The setup is clear and mirrors existing patterns. To improve clarity and future maintenance, define the force offsets and atom weights as class-level constants and reuse them in _prepare_weighted_system. This also addresses prior nitpicks about magic numbers.

 class TestDPTestForceWeight(DPTest, unittest.TestCase):
+    # Test constants for reproducibility and readability
+    FORCE_OFFSET_FIRST = 1.0
+    FORCE_OFFSET_LAST = 10.0
+    FIRST_ATOM_WEIGHT = 2.0
+    LAST_ATOM_WEIGHT = 0.0

---

189-191: Ensure temp model file is always closed and cleaned up, even on failures

Close the NamedTemporaryFile (SIM115) and move cleanup to tearDown so files aren’t leaked on assertion failures. Store the path on self.

-        tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth")
-        torch.jit.save(model, tmp_model.name)
+        tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth")
+        torch.jit.save(model, tmp_model.name)
+        tmp_model.close()  # ensure the handle is closed before other consumers use the path
+        self._tmp_model_path = tmp_model.name
@@
-        os.unlink(tmp_model.name)
+        # cleanup moved to tearDown
@@
     def tearDown(self) -> None:
         super().tearDown()
         shutil.rmtree(self.system_dir)
+        # Ensure model file is removed even if the test fails early
+        if hasattr(self, "_tmp_model_path") and os.path.exists(self._tmp_model_path):
+            os.unlink(self._tmp_model_path)

Also applies to: 237-237, 239-241

---

226-236: Strengthen assertions; verify denominators and guard against zero sum

Validate the effective sample sizes returned by dp_test_ener and add a guard to avoid division by zero if all weights are zero (future-proofing). This also aligns with the reviewer suggestion to “show both the error with and without atom_pref” by making both paths explicit in assertions.

         masked_diff = diff[mask]
-        mae_unweighted = np.sum(np.abs(masked_diff)) / mask.sum()
-        rmse_unweighted = np.sqrt(np.sum(masked_diff * masked_diff) / mask.sum())
+        mae_unweighted = np.sum(np.abs(masked_diff)) / mask.sum()
+        rmse_unweighted = np.sqrt(np.sum(masked_diff * masked_diff) / mask.sum())
         denom = weight.sum()
+        assert denom > 0, "Weight sum is zero; cannot compute weighted errors."
         mae_weighted = np.sum(np.abs(diff) * weight) / denom
         rmse_weighted = np.sqrt(np.sum(diff * diff * weight) / denom)
-        np.testing.assert_allclose(err["mae_f"][0], mae_unweighted)
-        np.testing.assert_allclose(err["rmse_f"][0], rmse_unweighted)
-        np.testing.assert_allclose(err["mae_fw"][0], mae_weighted)
-        np.testing.assert_allclose(err["rmse_fw"][0], rmse_weighted)
+        # values and denominators (effective sample sizes)
+        np.testing.assert_equal(err["mae_f"][1], mask.sum())
+        np.testing.assert_equal(err["rmse_f"][1], mask.sum())
+        np.testing.assert_allclose(err["mae_f"][0], mae_unweighted)
+        np.testing.assert_allclose(err["rmse_f"][0], rmse_unweighted)
+        np.testing.assert_equal(err["mae_fw"][1], weight.sum())
+        np.testing.assert_equal(err["rmse_fw"][1], weight.sum())
+        np.testing.assert_allclose(err["mae_fw"][0], mae_weighted)
+        np.testing.assert_allclose(err["rmse_fw"][0], rmse_weighted)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

source/tests/pt/test_dp_test.py (2)

168-183: Replace magic numbers with named constants for clarity

Use named constants for the force offsets and atom weights to improve readability and intent.

Apply within this block:

-        forces[0, :3] += 1.0
-        forces[0, -3:] += 10.0
+        forces[0, :3] += self.FORCE_OFFSET_HEAD
+        forces[0, -3:] += self.FORCE_OFFSET_TAIL
-        atom_pref[:, 0] = 2.0
-        atom_pref[:, -1] = 0.0
+        atom_pref[:, 0] = self.HEAD_WEIGHT
+        atom_pref[:, -1] = self.TAIL_WEIGHT

Add these class attributes near the top of the class:

# class-level constants (add near class definition)
FORCE_OFFSET_HEAD = 1.0
FORCE_OFFSET_TAIL = 10.0
HEAD_WEIGHT = 2.0
TAIL_WEIGHT = 0.0

---

207-222: Fix two bugs: frame slicing and dict access for eval output

• coord must be sliced to the first frame before reshape; current code flattens all frames.
• DeepEval.eval returns a dict; ret[1] will KeyError. Use ret["force"].

-        coord = test_data["coord"].reshape([1, -1])
+        coord = test_data["coord"][:1].reshape([1, -1])
@@
-        ret = dp.eval(
+        ret = dp.eval(
             coord,
             box,
             atype,
             fparam=None,
             aparam=None,
             atomic=False,
             efield=None,
             mixed_type=False,
             spin=None,
         )
-        force_pred = ret[1].reshape([1, -1])
+        force_pred = ret["force"].reshape([1, -1])

source/tests/pt/test_dp_test.py (3)

189-191: Use a context manager for the temp model file (ruff SIM115) and reuse the path

Prevents resource leaks and aligns with tooling suggestions.

-        tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth")
-        torch.jit.save(model, tmp_model.name)
-        dp = DeepEval(tmp_model.name)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pth") as tmp_model:
+            torch.jit.save(model, tmp_model.name)
+            tmp_model_path = tmp_model.name
+        dp = DeepEval(tmp_model_path)
@@
-        os.unlink(tmp_model.name)
+        os.unlink(tmp_model_path)

Also applies to: 235-235

---

226-234: Good: validates both unweighted and weighted MAE/RMSE (addresses reviewer ask)

Asserting both metrics matches the behavior with and without atom_pref.

Optionally add tolerances to avoid flakiness from dtype/ordering differences:

-        np.testing.assert_allclose(err["mae_f"][0], mae_unweighted)
-        np.testing.assert_allclose(err["rmse_f"][0], rmse_unweighted)
-        np.testing.assert_allclose(err["mae_fw"][0], mae_weighted)
-        np.testing.assert_allclose(err["rmse_fw"][0], rmse_weighted)
+        np.testing.assert_allclose(err["mae_f"][0],  mae_unweighted,  rtol=1e-6, atol=1e-8)
+        np.testing.assert_allclose(err["rmse_f"][0], rmse_unweighted, rtol=1e-6, atol=1e-8)
+        np.testing.assert_allclose(err["mae_fw"][0], mae_weighted,    rtol=1e-6, atol=1e-8)
+        np.testing.assert_allclose(err["rmse_fw"][0], rmse_weighted,  rtol=1e-6, atol=1e-8)

---

237-240: Make teardown more robust if setup fails mid-way

Guard against missing directory to avoid errors on teardown in partial failures.

-        shutil.rmtree(self.system_dir)
+        if hasattr(self, "system_dir") and os.path.isdir(self.system_dir):
+            shutil.rmtree(self.system_dir, ignore_errors=True)

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.