(fix)Make the weighted avarange fit for all kinds of systems

deepmd/entrypoints/test.py

@@ -134,7 +134,7 @@ def test(
         )
 
         if isinstance(dp, DeepPot):
-            err = test_ener(
+            err, find_energy, find_force, find_virial = test_ener(
                 dp,
                 data,
                 system,
@@ -143,6 +143,29 @@ def test(
                 atomic,
                 append_detail=(cc != 0),
             )
+            err_part = {}
+
+            if find_energy == 1:
+                err_part['mae_e'] = err['mae_e']
+                err_part['mae_ea'] = err['mae_ea']
+                err_part['rmse_e'] = err['rmse_e']
+                err_part['rmse_ea'] = err['rmse_ea']
+
+            if find_force == 1:
+                if 'rmse_f' in err:
+                    err_part['mae_f'] = err['mae_f']
+                    err_part['rmse_f'] = err['rmse_f']
+                else:
+                    err_part['mae_fr'] = err['mae_fr']
+                    err_part['rmse_fr'] = err['rmse_fr']
+                    err_part['mae_fm'] = err['mae_fm']
+                    err_part['rmse_fm'] = err['rmse_fm']
+            if find_virial == 1:
+                err_part['mae_v'] = err['mae_v']
+                err_part['rmse_v'] = err['rmse_v']
+
+            err = err_part
+
         elif isinstance(dp, DeepDOS):
             err = test_dos(
                 dp,
@@ -303,10 +326,11 @@ def test_ener(
     if dp.has_spin:
         data.add("spin", 3, atomic=True, must=True, high_prec=False)
         data.add("force_mag", 3, atomic=True, must=False, high_prec=False)
-    if dp.has_hessian:
-        data.add("hessian", 1, atomic=True, must=True, high_prec=False)
 
     test_data = data.get_test()
+    find_energy = test_data.get('find_energy')
+    find_force = test_data.get('find_force')
+    find_virial = test_data.get('find_virial')
     mixed_type = data.mixed_type
     natoms = len(test_data["type"][0])
     nframes = test_data["box"].shape[0]
@@ -354,9 +378,6 @@ def test_ener(
     energy = energy.reshape([numb_test, 1])
     force = force.reshape([numb_test, -1])
     virial = virial.reshape([numb_test, 9])
-    if dp.has_hessian:
-        hessian = ret[3]
-        hessian = hessian.reshape([numb_test, -1])
     if has_atom_ener:
         ae = ret[3]
         av = ret[4]
@@ -420,10 +441,6 @@ def test_ener(
     rmse_ea = rmse_e / natoms
     mae_va = mae_v / natoms
     rmse_va = rmse_v / natoms
-    if dp.has_hessian:
-        diff_h = hessian - test_data["hessian"][:numb_test]
-        mae_h = mae(diff_h)
-        rmse_h = rmse(diff_h)
     if has_atom_ener:
         diff_ae = test_data["atom_ener"][:numb_test].reshape([-1]) - ae.reshape([-1])
         mae_ae = mae(diff_ae)
@@ -439,26 +456,24 @@ def test_ener(
     log.info(f"Energy RMSE        : {rmse_e:e} eV")
     log.info(f"Energy MAE/Natoms  : {mae_ea:e} eV")
     log.info(f"Energy RMSE/Natoms : {rmse_ea:e} eV")
-    if not out_put_spin:
-        log.info(f"Force  MAE         : {mae_f:e} eV/A")
-        log.info(f"Force  RMSE        : {rmse_f:e} eV/A")
-    else:
-        log.info(f"Force atom MAE      : {mae_fr:e} eV/A")
-        log.info(f"Force atom RMSE     : {rmse_fr:e} eV/A")
-        log.info(f"Force spin MAE      : {mae_fm:e} eV/uB")
-        log.info(f"Force spin RMSE     : {rmse_fm:e} eV/uB")
+    if find_force == 1:
+        if not out_put_spin:
+            log.info(f"Force  MAE         : {mae_f:e} eV/A")
+            log.info(f"Force  RMSE        : {rmse_f:e} eV/A")
+        else:
+            log.info(f"Force atom MAE      : {mae_fr:e} eV/A")
+            log.info(f"Force atom RMSE     : {rmse_fr:e} eV/A")
+            log.info(f"Force spin MAE      : {mae_fm:e} eV/uB")
+            log.info(f"Force spin RMSE     : {rmse_fm:e} eV/uB")
 
-    if data.pbc and not out_put_spin:
+    if data.pbc and not out_put_spin and find_virial == 1:
         log.info(f"Virial MAE         : {mae_v:e} eV")
         log.info(f"Virial RMSE        : {rmse_v:e} eV")
         log.info(f"Virial MAE/Natoms  : {mae_va:e} eV")
         log.info(f"Virial RMSE/Natoms : {rmse_va:e} eV")
     if has_atom_ener:
         log.info(f"Atomic ener MAE    : {mae_ae:e} eV")
         log.info(f"Atomic ener RMSE   : {rmse_ae:e} eV")
-    if dp.has_hessian:
-        log.info(f"Hessian MAE        : {mae_h:e} eV/A^2")
-        log.info(f"Hessian RMSE       : {rmse_h:e} eV/A^2")
 
     if detail_file is not None:
         detail_path = Path(detail_file)
@@ -542,24 +557,8 @@ def test_ener(
             "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz",
             append=append_detail,
         )
-        if dp.has_hessian:
-            data_h = test_data["hessian"][:numb_test].reshape(-1, 1)
-            pred_h = hessian.reshape(-1, 1)
-            h = np.concatenate(
-                (
-                    data_h,
-                    pred_h,
-                ),
-                axis=1,
-            )
-            save_txt_file(
-                detail_path.with_suffix(".h.out"),
-                h,
-                header=f"{system}: data_h pred_h (3Na*3Na matrix in row-major order)",
-                append=append_detail,
-            )
     if not out_put_spin:
-        dict_to_return = {
+        return {
             "mae_e": (mae_e, energy.size),
             "mae_ea": (mae_ea, energy.size),
             "mae_f": (mae_f, force.size),
@@ -570,9 +569,9 @@ def test_ener(
             "rmse_f": (rmse_f, force.size),
             "rmse_v": (rmse_v, virial.size),
             "rmse_va": (rmse_va, virial.size),
-        }
+        }, find_energy, find_force, find_virial
     else:
-        dict_to_return = {
+        return {
             "mae_e": (mae_e, energy.size),
             "mae_ea": (mae_ea, energy.size),
             "mae_fr": (mae_fr, force_r.size),
@@ -585,11 +584,7 @@ def test_ener(
             "rmse_fm": (rmse_fm, force_m.size),
             "rmse_v": (rmse_v, virial.size),
             "rmse_va": (rmse_va, virial.size),
-        }
-    if dp.has_hessian:
-        dict_to_return["mae_h"] = (mae_h, hessian.size)
-        dict_to_return["rmse_h"] = (rmse_h, hessian.size)
-    return dict_to_return
+        }, find_energy, find_force, find_virial
 
 
 def print_ener_sys_avg(avg: dict[str, float]) -> None:
@@ -616,9 +611,6 @@ def print_ener_sys_avg(avg: dict[str, float]) -> None:
     log.info(f"Virial RMSE        : {avg['rmse_v']:e} eV")
     log.info(f"Virial MAE/Natoms  : {avg['mae_va']:e} eV")
     log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV")
-    if "rmse_h" in avg.keys():
-        log.info(f"Hessian MAE         : {avg['mae_h']:e} eV/A^2")
-        log.info(f"Hessian RMSE        : {avg['rmse_h']:e} eV/A^2")
 
 
 def test_dos(
@@ -739,9 +731,9 @@ def test_dos(
             frame_output = np.hstack((test_out, pred_out))
 
             save_txt_file(
-                detail_path.with_suffix(f".dos.out.{ii}"),
+                detail_path.with_suffix(".dos.out.%.d" % ii),
                 frame_output,
-                header=f"{system} - {ii}: data_dos pred_dos",
+                header="%s - %.d: data_dos pred_dos" % (system, ii),
                 append=append_detail,
             )
 
@@ -753,9 +745,9 @@ def test_dos(
                 frame_output = np.hstack((test_out, pred_out))
 
                 save_txt_file(
-                    detail_path.with_suffix(f".ados.out.{ii}"),
+                    detail_path.with_suffix(".ados.out.%.d" % ii),
                     frame_output,
-                    header=f"{system} - {ii}: data_ados pred_ados",
+                    header="%s - %.d: data_ados pred_ados" % (system, ii),
                     append=append_detail,
                 )
 
@@ -814,17 +806,9 @@ def test_property(
     tuple[list[np.ndarray], list[int]]
         arrays with results and their shapes
     """
-    var_name = dp.get_var_name()
-    assert isinstance(var_name, str)
-    data.add(var_name, dp.task_dim, atomic=False, must=True, high_prec=True)
+    data.add("property", dp.task_dim, atomic=False, must=True, high_prec=True)
     if has_atom_property:
-        data.add(
-            f"atom_{var_name}",
-            dp.task_dim,
-            atomic=True,
-            must=False,
-            high_prec=True,
-        )
+        data.add("atom_property", dp.task_dim, atomic=True, must=False, high_prec=True)
 
     if dp.get_dim_fparam() > 0:
         data.add(
@@ -875,12 +859,12 @@ def test_property(
         aproperty = ret[1]
         aproperty = aproperty.reshape([numb_test, natoms * dp.task_dim])
 
-    diff_property = property - test_data[var_name][:numb_test]
+    diff_property = property - test_data["property"][:numb_test]
     mae_property = mae(diff_property)
     rmse_property = rmse(diff_property)
 
     if has_atom_property:
-        diff_aproperty = aproperty - test_data[f"atom_{var_name}"][:numb_test]
+        diff_aproperty = aproperty - test_data["atom_property"][:numb_test]
         mae_aproperty = mae(diff_aproperty)
         rmse_aproperty = rmse(diff_aproperty)
 
@@ -897,29 +881,29 @@ def test_property(
         detail_path = Path(detail_file)
 
         for ii in range(numb_test):
-            test_out = test_data[var_name][ii].reshape(-1, 1)
+            test_out = test_data["property"][ii].reshape(-1, 1)
             pred_out = property[ii].reshape(-1, 1)
 
             frame_output = np.hstack((test_out, pred_out))
 
             save_txt_file(
-                detail_path.with_suffix(f".property.out.{ii}"),
+                detail_path.with_suffix(".property.out.%.d" % ii),
                 frame_output,
-                header=f"{system} - {ii}: data_property pred_property",
+                header="%s - %.d: data_property pred_property" % (system, ii),
                 append=append_detail,
             )
 
         if has_atom_property:
             for ii in range(numb_test):
-                test_out = test_data[f"atom_{var_name}"][ii].reshape(-1, 1)
+                test_out = test_data["atom_property"][ii].reshape(-1, 1)
                 pred_out = aproperty[ii].reshape(-1, 1)
 
                 frame_output = np.hstack((test_out, pred_out))
 
                 save_txt_file(
-                    detail_path.with_suffix(f".aproperty.out.{ii}"),
+                    detail_path.with_suffix(".aproperty.out.%.d" % ii),
                     frame_output,
-                    header=f"{system} - {ii}: data_aproperty pred_aproperty",
+                    header="%s - %.d: data_aproperty pred_aproperty" % (system, ii),
                     append=append_detail,
                 )
 

source/tests/pt/test_weighted_avg.py

@@ -0,0 +1,103 @@
+from deepmd.utils.weight_avg import (
+    weighted_average,
+)
+import unittest
+import numpy as np
+
+def test(all_sys):
+    err_coll = []
+    for sys_data in all_sys:
+        err, find_energy, find_force, find_virial = sys_data
+        err_part = {}
+        if find_energy == 1:
+            err_part['mae_e'] = err['mae_e']
+            err_part['mae_ea'] = err['mae_ea']
+            err_part['rmse_e'] = err['rmse_e']
+            err_part['rmse_ea'] = err['rmse_ea']
+        if find_force == 1:
+            if 'rmse_f' in err:
+                err_part['mae_f'] = err['mae_f']
+                err_part['rmse_f'] = err['rmse_f']
+            else:
+                err_part['mae_fr'] = err['mae_fr']
+                err_part['rmse_fr'] = err['rmse_fr']
+                err_part['mae_fm'] = err['mae_fm']
+                err_part['rmse_fm'] = err['rmse_fm']
+        if find_virial == 1:
+            err_part['mae_v'] = err['mae_v']
+            err_part['rmse_v'] = err['rmse_v']
+        err_coll.append(err_part)
+    avg_err = weighted_average(err_coll)
+    return avg_err
+
+def test_ori(all_sys):
+    err_coll = []
+    for sys_data in all_sys:
+        err, _, _, _ = sys_data
+        err_coll.append(err)
+    avg_err = weighted_average(err_coll)
+    return avg_err
+
+class TestWeightedAverage(unittest.TestCase):
+
+    def test_case1_energy_only(self):
+
+        all_sys = [
+            ({'mae_e': (2,2), 'mae_ea': (4,2), 'rmse_e': (3,2), 'rmse_ea': (5,2),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,0,0),
+            ({'mae_e': (4,3), 'mae_ea': (6,3), 'rmse_e': (5,3), 'rmse_ea': (7,3),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,0,0),
+            ({'mae_e': (6,5), 'mae_ea': (8,5), 'rmse_e': (7,5), 'rmse_ea': (9,5),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,0,0),
+        ]
+        expected_mae_e = (2*2 +4*3 +6*5)/(2+3+5)
+        expected_rmse_e = np.sqrt((3**2*2 +5**2*3 +7**2*5)/(2+3+5))
+        expected_mae_ea = (4*2 +6*3 +8*5)/10
+        expected_rmse_ea = np.sqrt((5**2*2 +7**2*3 +9**2*5)/10)
+
+        avg_err = test(all_sys)
+        self.assertAlmostEqual(avg_err['mae_e'], expected_mae_e)
+        self.assertAlmostEqual(avg_err['rmse_e'], expected_rmse_e)
+        self.assertAlmostEqual(avg_err['mae_ea'], expected_mae_ea)
+        self.assertAlmostEqual(avg_err['rmse_ea'], expected_rmse_ea)
+        self.assertAlmostEqual(avg_err['mae_f'], 0)
+        self.assertAlmostEqual(avg_err['mae_v'], 0)
+
+        avg_err_ori = test_ori(all_sys)
+        self.assertAlmostEqual(avg_err['mae_e'], avg_err_ori['mae_e'])
+        self.assertNotEqual(avg_err['mae_f'], avg_err_ori['rmse_f'])
+        self.assertNotEqual(avg_err['mae_v'], avg_err_ori['rmse_v'])
+
+
+
+    def test_case2_energy_force(self):
+        all_sys = [
+            ({'mae_e': (2,2), 'mae_ea': (4,2), 'rmse_e': (3,2), 'rmse_ea': (5,2),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,1,0),
+            ({'mae_e': (4,3), 'mae_ea': (6,3), 'rmse_e': (5,3), 'rmse_ea': (7,3),'mae_f': (1,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,1,0),
+            ({'mae_e': (6,5), 'mae_ea': (8,5), 'rmse_e': (7,5), 'rmse_ea': (9,5),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,0,0),
+        ]
+
+        avg_err = test(all_sys)
+        expected_mae_f = (2*3 +1*3 )/(3+3)
+        self.assertAlmostEqual(avg_err['mae_f'], expected_mae_f)
+
+        avg_err_ori = test_ori(all_sys)
+        self.assertAlmostEqual(avg_err['mae_e'], avg_err_ori['mae_e'])
+        self.assertNotEqual(avg_err['mae_f'], avg_err_ori['mae_f'])
+        self.assertNotEqual(avg_err['mae_v'], avg_err_ori['mae_v'])
+
+    def test_case3_all_components(self):
+        all_sys = [
+            ({'mae_e': (2,2), 'mae_ea': (4,2), 'rmse_e': (3,2), 'rmse_ea': (5,2),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,1,1),
+            ({'mae_e': (4,3), 'mae_ea': (6,3), 'rmse_e': (5,3), 'rmse_ea': (7,3),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (1,5), 'rmse_v': (2,3)}, 1,1,1),
+            ({'mae_e': (6,5), 'mae_ea': (8,5), 'rmse_e': (7,5), 'rmse_ea': (9,5),'mae_f': (2,3), 'rmse_f': (1,3), 'mae_v': (3,5), 'rmse_v': (3,3)}, 1,1,0),
+        ]
+
+        avg_err = test(all_sys)
+        expected_mae_v = (3*5 +1*5 )/(5+5)
+        self.assertAlmostEqual(avg_err['mae_v'], expected_mae_v)
+
+        avg_err_ori = test_ori(all_sys)
+        self.assertAlmostEqual(avg_err['mae_e'], avg_err_ori['mae_e'])
+        self.assertAlmostEqual(avg_err['mae_f'], avg_err_ori['mae_f'])
+        self.assertNotEqual(avg_err['mae_v'], avg_err_ori['rmse_v'])
+
+if __name__ == '__main__':
+    unittest.main()