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,20 +419,20 @@ 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_pref = test_data["atom_pref"][:numb_test]
-        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)
+        atom_weight = test_data["atom_pref"][:numb_test]
+        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_f = 0.0
-            rmse_f = 0.0
-    else:
-        mae_f = mae(diff_f)
-        rmse_f = rmse(diff_f)
-        size_f = force.size
+            mae_fw = 0.0
+            rmse_fw = 0.0
     diff_v = virial - test_data["virial"][:numb_test]
     mae_v = mae(diff_v)
     rmse_v = rmse(diff_v)
@@ -469,6 +469,11 @@ def test_ener(
         log.info(f"Force  RMSE        : {rmse_f:e} eV/A")
         dict_to_return["mae_f"] = (mae_f, size_f)
         dict_to_return["rmse_f"] = (rmse_f, size_f)
+        if find_atom_pref == 1:
+            log.info(f"Force weighted MAE : {mae_fw:e} eV/A")
+            log.info(f"Force weighted RMSE: {rmse_fw:e} eV/A")
+            dict_to_return["mae_fw"] = (mae_fw, weight_sum)
+            dict_to_return["rmse_fw"] = (rmse_fw, weight_sum)
     if out_put_spin and find_force == 1:
         log.info(f"Force atom MAE      : {mae_fr:e} eV/A")
         log.info(f"Force atom RMSE     : {rmse_fr:e} eV/A")
@@ -614,6 +619,9 @@ def print_ener_sys_avg(avg: dict[str, float]) -> None:
     if "rmse_f" in avg:
         log.info(f"Force  MAE         : {avg['mae_f']:e} eV/A")
         log.info(f"Force  RMSE        : {avg['rmse_f']:e} eV/A")
+        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")
     else:
         log.info(f"Force atom MAE      : {avg['mae_fr']:e} eV/A")
         log.info(f"Force spin MAE      : {avg['mae_fm']:e} eV/uB")

source/tests/pt/test_dp_test.py

@@ -14,20 +14,15 @@
 import numpy as np
 import torch
 
-from deepmd.entrypoints.test import test as dp_test
-from deepmd.entrypoints.test import test_ener as dp_test_ener
-from deepmd.infer.deep_eval import (
-    DeepEval,
-)
+from deepmd.entrypoints.test import test as dp_test, test_ener as dp_test_ener
 from deepmd.pt.entrypoints.main import (
     get_trainer,
 )
 from deepmd.pt.utils.utils import (
     to_numpy_array,
 )
-from deepmd.utils.data import (
-    DeepmdData,
-)
+from deepmd.infer.deep_eval import DeepEval
+from deepmd.utils.data import DeepmdData
 
 from .model.test_permutation import (
     model_property,
@@ -147,25 +142,25 @@ def setUp(self) -> None:
             json.dump(self.config, fp, indent=4)
 
 
-class TestDPTestForceMask(DPTest, unittest.TestCase):
+class TestDPTestForceWeight(DPTest, unittest.TestCase):
     def setUp(self) -> None:
-        self.detail_file = "test_dp_test_force_mask_detail"
+        self.detail_file = "test_dp_test_force_weight_detail"
         input_json = str(Path(__file__).parent / "water/se_atten.json")
         with open(input_json) as f:
             self.config = json.load(f)
         self.config["training"]["numb_steps"] = 1
         self.config["training"]["save_freq"] = 1
-        system_dir = self._prepare_masked_system()
+        system_dir = self._prepare_weighted_system()
         data_file = [system_dir]
         self.config["training"]["training_data"]["systems"] = data_file
         self.config["training"]["validation_data"]["systems"] = data_file
         self.config["model"] = deepcopy(model_se_e2_a)
         self.system_dir = system_dir
-        self.input_json = "test_dp_test_force_mask.json"
+        self.input_json = "test_dp_test_force_weight.json"
         with open(self.input_json, "w") as fp:
             json.dump(self.config, fp, indent=4)
 
-    def _prepare_masked_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)
@@ -179,7 +174,7 @@ def _prepare_masked_system(self) -> str:
         np.save(set_dir / "atom_pref.npy", atom_pref)
         return tmp_dir
 
-    def test_force_mask(self) -> None:
+    def test_force_weight(self) -> None:
         trainer = get_trainer(deepcopy(self.config))
         with torch.device("cpu"):
             trainer.get_data(is_train=False)
@@ -219,13 +214,18 @@ def test_force_mask(self) -> None:
         )
         force_pred = ret[1].reshape([1, -1])
         force_true = test_data["force"][:1]
-        mask = test_data["atom_pref"][:1]
-        diff = (force_pred - force_true) * mask
-        denom = mask.sum()
-        mae_expected = np.sum(np.abs(diff)) / denom
-        rmse_expected = np.sqrt(np.sum(diff * diff) / denom)
-        np.testing.assert_allclose(err["mae_f"][0], mae_expected)
-        np.testing.assert_allclose(err["rmse_f"][0], rmse_expected)
+        weight = test_data["atom_pref"][:1]
+        diff = force_pred - force_true
+        diff_w = diff * weight
+        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)
+        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)
         os.unlink(tmp_model.name)
 
     def tearDown(self) -> None: