Fix: dpdata WBM without force

dev/ase_test_interface.py

@@ -6,9 +6,10 @@
 from pathlib import Path
 from lamstare.infra.ood_database import OODRecord
 
+
 def run_ase_dptest(
-        calc,
-        testpath: str,
+    calc,
+    testpath: str,
 ):
     """
     Given a ASE calculator and a test path, run ASE dptest and return the results.
@@ -25,22 +26,16 @@ def run_ase_dptest(
     virial_err_per_atom = []
     max_ele_num = 120
 
-    systems = []
-    for path in testpath:
-        systems.extend(glob.glob(f"{path}/*"))
-    # check if the system is mixed type
-    if len(glob.glob(systems[0] + '/**/real_atom_types.npy', recursive=True)) == 0:
-        mix_type = False
-    else:
-        mix_type = True
-
+    systems = [i.parent for i in testpath.rglob("type_map.raw")]
+    assert systems, f"No systems found in the test data {testpath}."
+    mix_type = any(systems[0].rglob("real_atom_types.npy"))
 
     for filepth in systems:
         if mix_type:
             sys = dpdata.MultiSystems()
-            sys.load_systems_from_file(filepth, fmt='deepmd/npy/mixed')
+            sys.load_systems_from_file(filepth, fmt="deepmd/npy/mixed")
         else:
-            sys = dpdata.LabeledSystem(filepth, fmt='deepmd/npy')
+            sys = dpdata.LabeledSystem(filepth, fmt="deepmd/npy")
 
         for ls in sys:
             for frame in ls:
@@ -57,46 +52,58 @@ def run_ase_dptest(
                     energy_lab.append(frame.data["energies"])
                     energy_err.append(energy_predict - frame.data["energies"])
                     force_err.append(frame.data["forces"].squeeze(0) - np.array(ff))
-                    energy_err_per_atom.append(energy_err[-1]/force_err[-1].shape[0])
+                    energy_err_per_atom.append(energy_err[-1] / force_err[-1].shape[0])
                     try:
                         stress = atoms.get_stress()
-                        stress_tensor = - np.array(
-                        [[stress[0],stress[5],stress[4]],
-                        [stress[5],stress[1],stress[3]],
-                        [stress[4],stress[3], stress[2]]]
-                        )* atoms.get_volume()
-                        virial_err.append(frame.data['virials'] - stress_tensor)
-                        virial_err_per_atom.append(virial_err[-1]/force_err[-1].shape[0])
+                        stress_tensor = (
+                            -np.array(
+                                [
+                                    [stress[0], stress[5], stress[4]],
+                                    [stress[5], stress[1], stress[3]],
+                                    [stress[4], stress[3], stress[2]],
+                                ]
+                            )
+                            * atoms.get_volume()
+                        )
+                        virial_err.append(frame.data["virials"] - stress_tensor)
+                        virial_err_per_atom.append(
+                            virial_err[-1] / force_err[-1].shape[0]
+                        )
                     except:
                         pass
                 else:
                     pass
-
-
+
     atom_num = np.array(atom_num)
     energy_err = np.array(energy_err)
     energy_pre = np.array(energy_pre)
     energy_lab = np.array(energy_lab)
     shift_bias, _, _, _ = np.linalg.lstsq(atom_num, energy_err, rcond=1e-10)
-    unbiased_energy = energy_pre - (atom_num @ shift_bias.reshape(max_ele_num, -1)).reshape(-1) - energy_lab.squeeze()
+    unbiased_energy = (
+        energy_pre
+        - (atom_num @ shift_bias.reshape(max_ele_num, -1)).reshape(-1)
+        - energy_lab.squeeze()
+    )
     unbiased_energy_err_per_a = unbiased_energy / atom_num.sum(-1)
 
-
     res = {
-    "Energy MAE": [np.mean(np.abs(np.stack(unbiased_energy)))],
-    "Energy RMSE": [np.sqrt(np.mean(np.square(unbiased_energy)))],
-    "Energy MAE/Natoms": [np.mean(np.abs(np.stack(unbiased_energy_err_per_a)))],
-    "Energy RMSE/Natoms":  [np.sqrt(np.mean(np.square(unbiased_energy_err_per_a)))],
-    "Force MAE": [np.mean(np.abs(np.concatenate(force_err)))],
-    "Force RMSE": [np.sqrt(np.mean(np.square(np.concatenate(force_err))))],
+        "Energy MAE": [np.mean(np.abs(np.stack(unbiased_energy)))],
+        "Energy RMSE": [np.sqrt(np.mean(np.square(unbiased_energy)))],
+        "Energy MAE/Natoms": [np.mean(np.abs(np.stack(unbiased_energy_err_per_a)))],
+        "Energy RMSE/Natoms": [np.sqrt(np.mean(np.square(unbiased_energy_err_per_a)))],
+        "Force MAE": [np.mean(np.abs(np.concatenate(force_err)))],
+        "Force RMSE": [np.sqrt(np.mean(np.square(np.concatenate(force_err))))],
     }
     if virial_err_per_atom != []:
         res.update(
             {
-        "Virial MAE": [np.mean(np.abs(np.stack(virial_err)))],
-        "Virial RMSE": [np.sqrt(np.mean(np.square(np.stack(virial_err))))],
-        "Virial MAE/Natoms": [np.mean(np.abs(np.stack(virial_err_per_atom)))],
-        "Virial RMSE/Natoms": [np.sqrt(np.mean(np.square(np.stack(virial_err_per_atom))))]}
+                "Virial MAE": [np.mean(np.abs(np.stack(virial_err)))],
+                "Virial RMSE": [np.sqrt(np.mean(np.square(np.stack(virial_err))))],
+                "Virial MAE/Natoms": [np.mean(np.abs(np.stack(virial_err_per_atom)))],
+                "Virial RMSE/Natoms": [
+                    np.sqrt(np.mean(np.square(np.stack(virial_err_per_atom))))
+                ],
+            }
         )
     return res
 
@@ -107,9 +114,11 @@ def ase_test(model_name, testpath_mapping):
     """
     if model_name == "DP":
         from deepmd.calculator import DP
+
         CALC = DP("/path/to/your/model.pth")
     elif model_name == "MACE":
         from mace.calculators import mace_mp
+
         CALC = mace_mp(model="medium", device="cuda", default_dtype="float64")
     else:
         raise ValueError(f"Model {model_name} not supported.")
@@ -120,10 +129,10 @@ def ase_test(model_name, testpath_mapping):
         with open(f"{ood_name}_dptest.json", "w") as f:
             json.dump(head_dptest_res, f, indent=4)
 
-        
+
 if __name__ == "__main__":
     mapping = {
         "head_A": ["/path/to/your/test_data"],
         "head_B": ["/path/to/your/test_data"],
     }
-    ase_test("DP", mapping)
+    ase_test("DP", mapping)

lambench/metrics/direct_tasks_metrics.yml

@@ -2,59 +2,104 @@ ANI:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.19852160741116331
+  dataset_lstsq_std_force: 2.1346504447862467
+  dataset_lstsq_std_virial: null
 HEA25_S:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.4030134885402307
+  dataset_lstsq_std_force: 1.547935906005218
+  dataset_lstsq_std_virial: null
 HEA25_bulk:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.3915534935785154
+  dataset_lstsq_std_force: 2.0282929631352595
+  dataset_lstsq_std_virial: null
 HEMC_HEMB:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: 1.0
+  dataset_lstsq_std_energy: 0.4747620729960006
+  dataset_lstsq_std_force: 1.8364064815279475
+  dataset_lstsq_std_virial: 0.46794153820403517
 MD22:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.007933722948473606
+  dataset_lstsq_std_force: 1.1408107903449518
+  dataset_lstsq_std_virial: null
 Collision:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.46905055129776263
+  dataset_lstsq_std_force: 1.9436163958502828
+  dataset_lstsq_std_virial: null
 H_nature_2022:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.3866962525278814
+  dataset_lstsq_std_force: 2.038498926750271
+  dataset_lstsq_std_virial: null
 REANN_CO2_Ni100:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.007410380917403982
+  dataset_lstsq_std_force: 0.6333957228116258
+  dataset_lstsq_std_virial: null
 NequIP_NC_2022:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.006640637887476306
+  dataset_lstsq_std_force: 0.9421629142611048
+  dataset_lstsq_std_virial: null
 AIMD-Chig:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.012594134644011467
+  dataset_lstsq_std_force: 0.8599492186937476
+  dataset_lstsq_std_virial: null
 CGM_MLP_NC2023:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.9985195007897916
+  dataset_lstsq_std_force: 2.0217020563415375
+  dataset_lstsq_std_virial: null
 Cu_MgO_catalysts:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.014304435776268616
+  dataset_lstsq_std_force: 0.18909695700558055
+  dataset_lstsq_std_virial: null
 Subalex_9k:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: null
+  dataset_lstsq_std_energy: 0.7749643377228373
+  dataset_lstsq_std_force: 1.1503770816187873
+  dataset_lstsq_std_virial: 0.8678699239404155
 WBM_downsampled:
   energy_weight: 1.0
   force_weight: null
   virial_weight: null
+  dataset_lstsq_std_energy: 0.374831692072961
+  dataset_lstsq_std_force: null
+  dataset_lstsq_std_virial: null
 Torsionnet500:
   energy_weight: 1.0
   force_weight: 1.0
   virial_weight: 1.0
+  dataset_lstsq_std_energy: 0.03387935785428199
+  dataset_lstsq_std_force: 0.05088788030453556
+  dataset_lstsq_std_virial: null