feat: add neb

Author: anyangml

lambench/models/ase_models.py

@@ -207,6 +207,17 @@ def evaluate(
 
                 assert task.test_data is not None
                 return {"metrics": run_torsionnet(self, task.test_data)}
+            elif task.task_name == "neb":
+                from lambench.tasks.calculator.neb.neb import run_inference
+                assert task.test_data is not None
+                return {
+                    "metrics": run_inference(
+                        self,
+                        task.test_data,
+                        task.calculator_params.get("fmax", 0.01),
+                        task.calculator_params.get("steps", 500),
+                    )
+                }
             else:
                 raise NotImplementedError(f"Task {task.task_name} is not implemented.")
 

lambench/models/models_config.yml

@@ -207,21 +207,36 @@
   show_finetune_task: True
   show_calculator_task: True
 
-- model_name: dpa3_250307_v3_0_0_3M
+- model_name: dpa3_250415_v3_1_0_3M
   model_type: DP
   model_family: DP
-  model_path: /bohr/lambench-model-55c1/v8/dpa3_250307_v3_0_0_3M/dp_dpa3_v3.0.0_0307_426w.pt
-  virtualenv: *bohrium_image
+  model_path: /bohr/lambench-model-55c1/v12/dpa3_250415_v3_1_0_3M/dp_dpa3_v3.1.0_0415_400w.pt
+  virtualenv: registry.dp.tech/dptech/dp/native/prod-375/lambench:v2.11
   model_metadata:
-    pretty_name: DPA-3.0-3M
-    date_added: 2025-04-10
-    model_description: DP 2025 Q1, 3M parameters
-    num_parameters: 2592390
+    pretty_name: DPA-3.1-3M
+    date_added: 2025-05-13
+    model_description: DP 2025 Q2, 16 layers with dynamic nnei.
+    num_parameters: 3268689
     packages:
-      deepmd-kit: P0131
-  show_finetune_task: False
+      deepmd-kit: P0414
+  show_finetune_task: True
   show_calculator_task: True
-  skip_tasks: [PropertyFinetuneTask]
+
+# - model_name: dpa3_250307_v3_0_0_3M
+#   model_type: DP
+#   model_family: DP
+#   model_path: /bohr/lambench-model-55c1/v8/dpa3_250307_v3_0_0_3M/dp_dpa3_v3.0.0_0307_426w.pt
+#   virtualenv: *bohrium_image
+#   model_metadata:
+#     pretty_name: DPA-3.0-3M
+#     date_added: 2025-04-10
+#     model_description: DP 2025 Q1, 3M parameters
+#     num_parameters: 2592390
+#     packages:
+#       deepmd-kit: P0131
+#   show_finetune_task: False
+#   show_calculator_task: True
+#   skip_tasks: [PropertyFinetuneTask]
 
 - model_name: mace_mp_0_medium
   model_type: ASE
@@ -275,18 +290,18 @@
       orb-models: 0.5.2
   show_calculator_task: True
 
-- model_name: 7net-0
-  model_type: ASE
-  model_family: SevenNet
-  virtualenv: *bohrium_image
-  model_metadata:
-    pretty_name: SevenNet-0
-    date_added: 2025-04-10
-    model_description: https://github.com/MDIL-SNU/SevenNet/blob/main/sevenn/pretrained_potentials/SevenNet_0__11Jul2024/checkpoint_sevennet_0.pth
-    num_parameters: 842440
-    packages:
-      sevenn: 0.11.0
-  show_calculator_task: True
+# - model_name: 7net-0
+#   model_type: ASE
+#   model_family: SevenNet
+#   virtualenv: *bohrium_image
+#   model_metadata:
+#     pretty_name: SevenNet-0
+#     date_added: 2025-04-10
+#     model_description: https://github.com/MDIL-SNU/SevenNet/blob/main/sevenn/pretrained_potentials/SevenNet_0__11Jul2024/checkpoint_sevennet_0.pth
+#     num_parameters: 842440
+#     packages:
+#       sevenn: 0.11.0
+#   show_calculator_task: True
 
 - model_name: 7net-l3i5
   model_type: ASE

lambench/tasks/calculator/calculator_tasks.yml

@@ -15,3 +15,8 @@ inference_efficiency:
 torsionnet:
   test_data: /bohr/lambench-torsionnet-e4sc/v2/torsionnet500_wB97m
   calculator_params: null
+neb:
+  test_data: /bohr/lambench-neb-b1ln/v1/OC20NEB-OOD
+  calculator_params:
+    fmax: 0.01
+    steps: 500

lambench/tasks/calculator/neb/neb.py

@@ -0,0 +1,77 @@
+from lambench.models.ase_models import ASEModel
+from pathlib import Path
+import pandas as pd
+import numpy as np
+from tqdm import tqdm
+from ase.io import read
+from sklearn.metrics import mean_absolute_error
+
+
+def run_inference(
+    model: ASEModel,
+    test_data: Path,
+    fmax: float,
+    steps: int,
+) -> dict[str, float]:
+    """ "
+    This tests perform energy barrier prediction for 460 OOD NEB trajs in the OC20NEB dataset.
+
+    The label.csv file contains the ground truth energy barrier and dE for each trajectory.
+    Each trajectory file contains 10 frames, and the first and last frames are the initial and final states.
+    Using the ground truth energy label, the index of the transition state is determined.
+
+    We then perform structure optimization for the initial, final, and transition states and calculate the energy barrier and dE.
+    """
+    result_df = pd.read_csv(Path(test_data, "label.csv"))
+    result_df.sort_values("traj", inplace=True)
+    result_df["pred_Ea"] = None
+    result_df["pred_dE"] = None
+    NUM_RECORDS = len(result_df)
+    ERROR_THRESHOLD = 0.1  # counting Ea predtion error > 0.1 eV as failure
+
+    for idx, row in tqdm(result_df.iterrows()):
+        traj_name = row["traj"]
+        data = read(test_data, f"{traj_name}.traj", ":")
+        energies = [frame.get_potential_energy() for frame in data]
+        barrier_idx = np.argmax(energies)
+
+        initial, transition, final = data[0], data[barrier_idx], data[-1]
+        relaxed_energy = []
+        for atoms in [initial, transition, final]:
+            atoms = ASEModel.run_ase_relaxation(
+                atoms=atoms,
+                calc=model.calc,
+                fmax=fmax,
+                steps=steps,
+                fix_symmetry=False,
+                relax_cell=False,
+            )
+            relaxed_energy.append(atoms.get_potential_energy())
+        e_a, de = (
+            relaxed_energy[1] - relaxed_energy[0],
+            relaxed_energy[-1] - relaxed_energy[0],
+        )
+        result_df.at[idx, "pred_Ea"] = e_a
+        result_df.at[idx, "pred_dE"] = de
+
+    result_df["type"] = result_df["traj"].apply(lambda x: x.split("_")[0])
+    result_df["error"] = result_df.apply(
+        lambda x: np.abs(x["pred_Ea"] - x["Ea"]), axis=1
+    )
+    type_percentages = (
+        result_df[result_df["error"] > ERROR_THRESHOLD].groupby("type").size()
+        / result_df.groupby("type").size()
+        * 100
+    )
+    type_percentages = type_percentages.round(2)
+    results = type_percentages.to_dict()
+    type_percentages.dropna(inplace=True)
+    results["success_rate"] = len(type_percentages) / NUM_RECORDS * 100
+    results["MAE_Ea"] = mean_absolute_error(
+        result_df["Ea"], result_df["pred_Ea"]
+    ).round(2)
+    results["MAE_dE"] = mean_absolute_error(
+        result_df["dE"], result_df["pred_dE"]
+    ).round(2)
+    
+    return results