Merge branch 'codex/update-test_ener-to-handle-masked-forces-tkzu2y' into codex/update-test_ener-to-handle-masked-forces

Author: caic99

deepmd/entrypoints/test.py

@@ -419,18 +419,28 @@ def test_ener(
     mae_e = mae(diff_e)
     rmse_e = rmse(diff_e)
     diff_f = force - test_data["force"][:numb_test]
-    mae_f = mae(diff_f)
-    rmse_f = rmse(diff_f)
-    size_f = force.size
     if find_atom_pref == 1:
         atom_weight = test_data["atom_pref"][:numb_test]
+        mask = atom_weight != 0
+        masked_diff = diff_f[mask]
+        if masked_diff.size > 0:
+            mae_f = np.mean(np.abs(masked_diff))
+            rmse_f = np.sqrt(np.mean(masked_diff * masked_diff))
+        else:
+            mae_f = 0.0
+            rmse_f = 0.0
+        size_f = mask.sum()
         weight_sum = np.sum(atom_weight)
         if weight_sum > 0:
             mae_fw = np.sum(np.abs(diff_f) * atom_weight) / weight_sum
             rmse_fw = np.sqrt(np.sum(diff_f * diff_f * atom_weight) / weight_sum)
         else:
             mae_fw = 0.0
             rmse_fw = 0.0
+    else:
+        mae_f = mae(diff_f)
+        rmse_f = rmse(diff_f)
+        size_f = diff_f.size
     diff_v = virial - test_data["virial"][:numb_test]
     mae_v = mae(diff_v)
     rmse_v = rmse(diff_v)

source/tests/pt/test_dp_test.py

@@ -171,10 +171,12 @@ def _prepare_weighted_system(self) -> str:
         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] += 1.0
         forces[0, -3:] += 10.0
         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
         np.save(set_dir / "atom_pref.npy", atom_pref)
         return tmp_dir
@@ -221,16 +223,17 @@ def test_force_weight(self) -> None:
         force_true = test_data["force"][:1]
         weight = test_data["atom_pref"][:1]
         diff = force_pred - force_true
-        diff_w = diff * weight
+        mask = weight != 0
+        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())
         denom = weight.sum()
-        mae_expected = np.sum(np.abs(diff_w)) / denom
-        rmse_expected = np.sqrt(np.sum(diff * diff * weight) / denom)
-        mae_unweighted = np.sum(np.abs(diff)) / diff.size
-        rmse_unweighted = np.sqrt(np.sum(diff * diff) / diff.size)
+        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_expected)
-        np.testing.assert_allclose(err["rmse_fw"][0], rmse_expected)
+        np.testing.assert_allclose(err["mae_fw"][0], mae_weighted)
+        np.testing.assert_allclose(err["rmse_fw"][0], rmse_weighted)
         os.unlink(tmp_model.name)
 
     def tearDown(self) -> None: